Initial CB tests for batched aggs

Author: ivancea

server/src/main/java/org/elasticsearch/common/io/stream/DelayableWriteable.java

@@ -75,7 +75,7 @@ public static <T extends Writeable> DelayableWriteable<T> referencing(Writeable.
         }
     }
 
-    private DelayableWriteable() {}
+    protected DelayableWriteable() {}
 
     /**
      * Returns a {@linkplain DelayableWriteable} that stores its contents

server/src/test/java/org/elasticsearch/action/search/QueryPhaseResultConsumerTests.java

@@ -13,8 +13,14 @@
 import org.apache.lucene.search.TopDocs;
 import org.apache.lucene.search.TotalHits;
 import org.elasticsearch.common.breaker.CircuitBreaker;
+import org.elasticsearch.common.breaker.CircuitBreakingException;
 import org.elasticsearch.common.breaker.NoopCircuitBreaker;
+import org.elasticsearch.common.io.stream.DelayableWriteable;
+import org.elasticsearch.common.io.stream.NamedWriteableRegistry;
+import org.elasticsearch.common.io.stream.StreamOutput;
+import org.elasticsearch.common.io.stream.Writeable;
 import org.elasticsearch.common.lucene.search.TopDocsAndMaxScore;
+import org.elasticsearch.common.unit.ByteSizeValue;
 import org.elasticsearch.common.util.BigArrays;
 import org.elasticsearch.common.util.concurrent.EsExecutors;
 import org.elasticsearch.common.util.concurrent.EsExecutors.TaskTrackingConfig;
@@ -25,7 +31,9 @@
 import org.elasticsearch.search.aggregations.AggregationBuilder;
 import org.elasticsearch.search.aggregations.AggregationReduceContext;
 import org.elasticsearch.search.aggregations.InternalAggregations;
+import org.elasticsearch.search.aggregations.metrics.SumAggregationBuilder;
 import org.elasticsearch.search.aggregations.pipeline.PipelineAggregator;
+import org.elasticsearch.search.builder.SearchSourceBuilder;
 import org.elasticsearch.search.query.QuerySearchResult;
 import org.elasticsearch.test.ESTestCase;
 import org.elasticsearch.threadpool.TestThreadPool;
@@ -40,7 +48,10 @@
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.Supplier;
 
+import static org.hamcrest.Matchers.equalTo;
+import static org.hamcrest.Matchers.greaterThanOrEqualTo;
 import static org.mockito.Mockito.mock;
 
 public class QueryPhaseResultConsumerTests extends ESTestCase {
@@ -148,6 +159,122 @@ public void testProgressListenerExceptionsAreCaught() throws Exception {
         }
     }
 
+    public void testBatchedEstimateSizeTooBig() throws Exception {
+        SearchRequest searchRequest = new SearchRequest("index");
+        searchRequest.source(new SearchSourceBuilder().aggregation(new SumAggregationBuilder("sum")));
+
+        var circuitBreakerLimit = ByteSizeValue.ofMb(256);
+        var circuitBreaker = newLimitedBreaker(circuitBreakerLimit);
+        // More than what the CircuitBreaker should allow
+        long aggregationEstimatedSize = (long) (circuitBreakerLimit.getBytes() * 1.1);
+
+        try (
+            QueryPhaseResultConsumer queryPhaseResultConsumer = new QueryPhaseResultConsumer(
+                searchRequest,
+                executor,
+                circuitBreaker,
+                searchPhaseController,
+                () -> false,
+                new SearchProgressListener() {
+                },
+                10,
+                e -> {}
+            )
+        ) {
+            var mergeResult = new QueryPhaseResultConsumer.MergeResult(List.of(), null, new DelegatingDelayableWriteable<>(() -> {
+                fail("This shouldn't be called");
+                return null;
+            }), aggregationEstimatedSize);
+            queryPhaseResultConsumer.addBatchedPartialResult(new SearchPhaseController.TopDocsStats(0), mergeResult);
+
+            try {
+                queryPhaseResultConsumer.reduce();
+                fail("Expecting a circuit breaking exception to be thrown");
+            } catch (CircuitBreakingException e) {
+                assertThat(e.getBytesWanted(), equalTo(aggregationEstimatedSize));
+            }
+        }
+    }
+
+    public void testBatchedEstimateSizeTooBigAfterDeserialization() throws Exception {
+        SearchRequest searchRequest = new SearchRequest("index");
+        searchRequest.source(new SearchSourceBuilder().aggregation(new SumAggregationBuilder("sum")));
+
+        var circuitBreakerLimit = ByteSizeValue.ofMb(256);
+        var circuitBreaker = newLimitedBreaker(circuitBreakerLimit);
+        // Less than the CB, but more after the 1.5x
+        long aggregationEstimatedSize = (long) (circuitBreakerLimit.getBytes() * 0.75);
+
+        try (
+            QueryPhaseResultConsumer queryPhaseResultConsumer = new QueryPhaseResultConsumer(
+                searchRequest,
+                executor,
+                circuitBreaker,
+                searchPhaseController,
+                () -> false,
+                new SearchProgressListener() {
+                },
+                10,
+                e -> {}
+            )
+        ) {
+            var mergeResult = new QueryPhaseResultConsumer.MergeResult(List.of(), null, new DelegatingDelayableWriteable<>(() -> {
+                fail("This shouldn't be called");
+                return null;
+            }), aggregationEstimatedSize);
+            queryPhaseResultConsumer.addBatchedPartialResult(new SearchPhaseController.TopDocsStats(0), mergeResult);
+
+            try {
+                queryPhaseResultConsumer.reduce();
+                fail("Expecting a circuit breaking exception to be thrown");
+            } catch (CircuitBreakingException e) {
+                assertThat(circuitBreaker.getUsed(), greaterThanOrEqualTo(aggregationEstimatedSize));
+                assertThat(e.getBytesWanted(), equalTo((long) (aggregationEstimatedSize * 0.5)));
+            }
+        }
+    }
+
+    /**
+     * DelayableWriteable that delegates expansion to a supplier.
+     */
+    private static class DelegatingDelayableWriteable<T extends Writeable> extends DelayableWriteable<T> {
+        private final Supplier<T> supplier;
+
+        private DelegatingDelayableWriteable(Supplier<T> supplier) {
+            this.supplier = supplier;
+        }
+
+        @Override
+        public void writeTo(StreamOutput out) {
+            throw new UnsupportedOperationException("Not to be called");
+        }
+
+        @Override
+        public T expand() {
+            return supplier.get();
+        }
+
+        @Override
+        public Serialized<T> asSerialized(Reader<T> reader, NamedWriteableRegistry registry) {
+            throw new UnsupportedOperationException("Not to be called");
+        }
+
+        @Override
+        public boolean isSerialized() {
+            return true;
+        }
+
+        @Override
+        public long getSerializedSize() {
+            return 0;
+        }
+
+        @Override
+        public void close() {
+            // noop
+        }
+    }
+
     private static class ThrowingSearchProgressListener extends SearchProgressListener {
         private final AtomicInteger onQueryResult = new AtomicInteger(0);
         private final AtomicInteger onPartialReduce = new AtomicInteger(0);