EQL: Add circuit breaker to sampling

x-pack/plugin/eql/src/main/java/org/elasticsearch/xpack/eql/execution/assembler/ExecutionManager.java

@@ -211,7 +211,7 @@ public Executable assemble(List<List<Attribute>> listOfKeys, List<PhysicalPlan>
  }
  }
 
- return new SampleIterator(new PITAwareQueryClient(session), criteria, cfg.fetchSize());
+ return new SampleIterator(new PITAwareQueryClient(session), criteria, cfg.fetchSize(), session.circuitBreaker());
  }
 
  private HitExtractor timestampExtractor(HitExtractor hitExtractor) {

x-pack/plugin/eql/src/main/java/org/elasticsearch/xpack/eql/execution/sample/SampleIterator.java

@@ -9,9 +9,12 @@
 
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
+import org.apache.lucene.util.Accountable;
+import org.apache.lucene.util.RamUsageEstimator;
 import org.elasticsearch.action.ActionListener;
 import org.elasticsearch.action.search.MultiSearchResponse;
 import org.elasticsearch.action.search.SearchRequest;
+import org.elasticsearch.common.breaker.CircuitBreaker;
 import org.elasticsearch.core.TimeValue;
 import org.elasticsearch.search.SearchHit;
 import org.elasticsearch.search.aggregations.Aggregation;
@@ -30,6 +33,7 @@
 import org.elasticsearch.xpack.ql.util.ActionListeners;
 
 import java.util.ArrayList;
+import java.util.Collection;
 import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
@@ -47,19 +51,45 @@ public class SampleIterator implements Executable {
 
  private final QueryClient client;
  private final List<SampleCriterion> criteria;
- private final Stack<Page> stack = new Stack<>();
+ final Stack<Page> stack = new Stack<>();
  private final int maxCriteria;
- private final List<Sample> samples;
+ final List<Sample> samples;
  private final int fetchSize;
 
  private long startTime;
 
- public SampleIterator(QueryClient client, List<SampleCriterion> criteria, int fetchSize) {
+ // ---------- CIRCUIT BREAKER -----------
+
+ /**
+ * Memory consumption will be calculated every CB_STACK_SIZE_PRECISION hits added to the stack
+ * ie. the sum of sizes of pages added to the stack
+ * (not considering stack.pop(), so the number of hits added to the stack is different
+ * from the number of hits currently present in the stack)
+ */
+ protected static final int CB_STACK_SIZE_PRECISION = 1000;
+ private static final String CB_COMPLETED_LABEL = "sample_completed";
+ private static final String CB_INFLIGHT_LABEL = "sample_inflight";
+ private final CircuitBreaker circuitBreaker;
+ private long samplesRamBytesUsed = 0;
+ private long stackRamBytesUsed = 0;
+ private long totalRamBytesUsed = 0;
+ /**
+ * total number of hits (ie. sum of page sizes) added to the stack
+ * (not considering stack.pop(), so different from current stack size)
+ */
+ private long totalPageSize = 0;
+ /**
+ * total number of hits (ie. sum of page sizes) added to the stack when last memory check was executed
+ */
+ private long previousTotalPageSize = 0;
+
+ public SampleIterator(QueryClient client, List<SampleCriterion> criteria, int fetchSize, CircuitBreaker circuitBreaker) {
  this.client = client;
  this.criteria = criteria;
  this.maxCriteria = criteria.size();
  this.fetchSize = fetchSize;
  this.samples = new ArrayList<>();
+ this.circuitBreaker = circuitBreaker;
  }
 
  @Override
@@ -69,6 +99,7 @@ public void execute(ActionListener<Payload> listener) {
  advance(runAfter(listener, () -> {
  stack.clear();
  samples.clear();
+ clearCircuitBreaker();
  client.close(listener.delegateFailure((l, r) -> {}));
  }));
  }
@@ -114,8 +145,8 @@ private void queryForCompositeAggPage(ActionListener<Payload> listener, final Sa
  InternalComposite composite = (InternalComposite) a;
  log.trace("Found [{}] composite buckets", composite.getBuckets().size());
  Page nextPage = new Page(composite, request);
- if (nextPage.size() > 0) {
- stack.push(nextPage);
+ if (nextPage.size > 0) {
+ pushToStack(nextPage);
  advance(listener);
  } else {
  if (stack.size() > 0) {
@@ -127,6 +158,15 @@ private void queryForCompositeAggPage(ActionListener<Payload> listener, final Sa
  }, listener::onFailure));
  }
 
+ protected void pushToStack(Page nextPage) {
+ stack.push(nextPage);
+ totalPageSize += nextPage.size;
+ if (totalPageSize - previousTotalPageSize >= CB_STACK_SIZE_PRECISION) {
+ updateMemoryUsage();
+ previousTotalPageSize = totalPageSize;
+ }
+ }
+
  /*
  * Creates a _msearch request containing maxCriteria (number of filters in the query) * number_of_join_keys queries.
  * For a query with three filters
@@ -160,7 +200,6 @@ private void finalStep(ActionListener<Payload> listener) {
 
  int initialSize = samples.size();
  client.multiQuery(searches, ActionListener.wrap(r -> {
- List<List<SearchHit>> finalSamples = new ArrayList<>();
  List<List<SearchHit>> sample = new ArrayList<>(maxCriteria);
  MultiSearchResponse.Item[] response = r.getResponses();
  int docGroupsCounter = 1;
@@ -174,7 +213,6 @@ private void finalStep(ActionListener<Payload> listener) {
  if (docGroupsCounter == maxCriteria) {
  List<SearchHit> match = matchSample(sample, maxCriteria);
  if (match != null) {
- finalSamples.add(match);
  samples.add(new Sample(sampleKeys.get(responseIndex / maxCriteria), match));
  }
  docGroupsCounter = 1;
@@ -187,20 +225,30 @@ private void finalStep(ActionListener<Payload> listener) {
  log.trace("Final step... found [{}] new Samples", samples.size() - initialSize);
  // if this final page is max_page_size in size it means: either it's the last page and it happens to have max_page_size elements
  // or it's just not the last page and we should advance
- var next = page.size() == fetchSize ? page : stack.pop();
+ var next = page.size == fetchSize ? page : stack.pop();
  log.trace("Final step... getting next page of the " + (next == page ? "current" : "previous") + " page");
  nextPage(listener, next);
  }, listener::onFailure));
  }
 
+ private void updateMemoryUsage() {
+ long newSamplesRamSize = RamUsageEstimator.sizeOfCollection(samples);
+ addMemory(newSamplesRamSize - samplesRamBytesUsed, CB_COMPLETED_LABEL);
+ samplesRamBytesUsed = newSamplesRamSize;
+
+ long newStackRamSize = RamUsageEstimator.sizeOfCollection(stack);
+ addMemory(newStackRamSize - stackRamBytesUsed, CB_INFLIGHT_LABEL);
+ stackRamBytesUsed = newStackRamSize;
+ }
+
  /*
  * Finds the next set of results using the after_key of the previous set of buckets.
  * It can go back on previous page(s) until either there are no more results, or it finds a page with an after_key to use.
  */
  private void nextPage(ActionListener<Payload> listener, Page page) {
- page.request().nextAfter(page.afterKey());
- log.trace("Getting next page for page [{}] with afterkey [{}]", page, page.afterKey());
- queryForCompositeAggPage(listener, page.request());
+ page.request.nextAfter(page.afterKey);
+ log.trace("Getting next page for page [{}] with afterkey [{}]", page, page.afterKey);
+ queryForCompositeAggPage(listener, page.request);
  }
 
  /*
@@ -288,19 +336,66 @@ private static boolean match(int currentCriterion, List<List<SearchHit>> hits, L
  return false;
  }
 
+ private void addMemory(long bytes, String label) {
+ totalRamBytesUsed += bytes;
+ circuitBreaker.addEstimateBytesAndMaybeBreak(bytes, label);
+ }
+
+ private void clearCircuitBreaker() {
+ circuitBreaker.addWithoutBreaking(-totalRamBytesUsed);
+ stackRamBytesUsed = 0;
+ samplesRamBytesUsed = 0;
+ totalRamBytesUsed = 0;
+ totalPageSize = 0;
+ previousTotalPageSize = 0;
+ }
+
  private TimeValue timeTook() {
  return new TimeValue(System.currentTimeMillis() - startTime);
  }
 
- private record Page(
- List<InternalComposite.InternalBucket> hits,
- int size,
- Map<String, Object> afterKey,
- List<String> keys,
- SampleQueryRequest request
- ) {
- Page(InternalComposite compositeAgg, SampleQueryRequest request) {
- this(compositeAgg.getBuckets(), compositeAgg.getBuckets().size(), compositeAgg.afterKey(), request.keys(), request);
+ protected static class Page implements Accountable {
+ final List<InternalComposite.InternalBucket> hits;
+ final int size;
+ final Map<String, Object> afterKey;
+ final List<String> keys;
+ final SampleQueryRequest request;
+
+ long ramBytesUsed = 0;
+
+ private static final long SHALLOW_SIZE = RamUsageEstimator.shallowSizeOfInstance(Page.class);
+
+ // for test purposes only
+ protected Page(int size) {
+ hits = null;
+ this.size = size;
+ afterKey = null;
+ keys = null;
+ request = null;
+ }
+
+ protected Page(InternalComposite compositeAgg, SampleQueryRequest request) {
+ hits = compositeAgg.getBuckets();
+ size = compositeAgg.getBuckets().size();
+ afterKey = compositeAgg.afterKey();
+ keys = request.keys();
+ this.request = request;
+ }
+
+ @Override
+ public long ramBytesUsed() {
+ if (ramBytesUsed == 0) {
+ ramBytesUsed = SHALLOW_SIZE;
+ ramBytesUsed += RamUsageEstimator.sizeOfCollection(hits);
+ ramBytesUsed += RamUsageEstimator.sizeOfCollection(keys);
+ ramBytesUsed += RamUsageEstimator.sizeOfMap(afterKey);
+ }
+ return ramBytesUsed;
+ }
+
+ @Override
+ public Collection<Accountable> getChildResources() {
+ return Accountable.super.getChildResources();
  }
  }
 }

x-pack/plugin/eql/src/test/java/org/elasticsearch/xpack/eql/EqlTestUtils.java

@@ -9,7 +9,10 @@
 
 import org.elasticsearch.Version;
 import org.elasticsearch.action.support.IndicesOptions;
+import org.elasticsearch.common.breaker.CircuitBreaker;
+import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.core.TimeValue;
+import org.elasticsearch.indices.breaker.BreakerSettings;
 import org.elasticsearch.search.DocValueFormat;
 import org.elasticsearch.search.SearchSortValues;
 import org.elasticsearch.tasks.TaskId;
@@ -20,6 +23,7 @@
 import org.elasticsearch.xpack.eql.expression.predicate.operator.comparison.InsensitiveNotEquals;
 import org.elasticsearch.xpack.eql.expression.predicate.operator.comparison.InsensitiveWildcardEquals;
 import org.elasticsearch.xpack.eql.expression.predicate.operator.comparison.InsensitiveWildcardNotEquals;
+import org.elasticsearch.xpack.eql.plugin.EqlPlugin;
 import org.elasticsearch.xpack.eql.session.EqlConfiguration;
 import org.elasticsearch.xpack.ql.expression.Expression;
 
@@ -130,4 +134,17 @@ public static SearchSortValues randomSearchLongSortValues() {
  }
  return new SearchSortValues(values, sortValueFormats);
  }
+
+ public static BreakerSettings circuitBreakerSettings(Settings settings) {
+ return BreakerSettings.updateFromSettings(
+ new BreakerSettings(
+ EqlPlugin.CIRCUIT_BREAKER_NAME,
+ EqlPlugin.CIRCUIT_BREAKER_LIMIT,
+ EqlPlugin.CIRCUIT_BREAKER_OVERHEAD,
+ CircuitBreaker.Type.MEMORY,
+ CircuitBreaker.Durability.TRANSIENT
+ ),
+ settings
+ );
+ }
 }