Add CircuitBreaker to TDigest, Step 3: Connect with CB (elastic#113387) (elastic#113885)

Part of elastic#99815

## Steps 1. Migrate TDigest classes to use a custom Array
implementation. Temporarily use a simple array wrapper
(elastic#112810) 2. Implement
CircuitBreaking in the `WrapperTDigestArrays` class. Add
Releasable/AutoCloseable and ensure everything is closed
(elastic#113105) 3. Pass the
CircuitBreaker as a parameter to TDigestState from wherever it's being
used (This PR)     - ESQL: Pass a real CB     - Other aggs: Use the
deprecated methods on `TDigestState`, that will use a No-op CB instead
4. Account remaining TDigest classes size ("SHALLOW_SIZE")

Every step should be safely mergeable to main: - The first and second
steps should have no impact. - The third and fourth ones will start
increasing the CB count partially.

## Remarks TDigestStates are Releasable, and should be closed now.
However, old aggregations don't close them, as it's not trivial, and as
they are using the NoopCircuitBreaker, there's no need to close them

Author: ivancea

docs/changelog/113387.yaml

@@ -0,0 +1,5 @@
+pr: 113387
+summary: "Add `CircuitBreaker` to TDigest, Step 3: Connect with ESQL CB"
+area: ES|QL
+type: enhancement
+issues: []

server/src/main/java/org/elasticsearch/search/aggregations/metrics/AbstractTDigestPercentilesAggregator.java

@@ -87,7 +87,7 @@ private TDigestState getExistingOrNewHistogram(final BigArrays bigArrays, long b
         states = bigArrays.grow(states, bucket + 1);
         TDigestState state = states.get(bucket);
         if (state == null) {
-            state = TDigestState.create(compression, executionHint);
+            state = TDigestState.createWithoutCircuitBreaking(compression, executionHint);
             states.set(bucket, state);
         }
         return state;

server/src/main/java/org/elasticsearch/search/aggregations/metrics/EmptyTDigestState.java

@@ -9,10 +9,12 @@
 
 package org.elasticsearch.search.aggregations.metrics;
 
+import org.elasticsearch.common.breaker.NoopCircuitBreaker;
+
 public final class EmptyTDigestState extends TDigestState {
     public EmptyTDigestState() {
         // Use the sorting implementation to minimize memory allocation.
-        super(MemoryTrackingTDigestArrays.INSTANCE, Type.SORTING, 1.0D);
+        super(new NoopCircuitBreaker("empty-tdigest-state-noop-breaker"), Type.SORTING, 1.0D);
     }
 
     @Override

server/src/main/java/org/elasticsearch/search/aggregations/metrics/InternalMedianAbsoluteDeviation.java

@@ -30,13 +30,14 @@ public static double computeMedianAbsoluteDeviation(TDigestState valuesSketch) {
             return Double.NaN;
         } else {
             final double approximateMedian = valuesSketch.quantile(0.5);
-            final TDigestState approximatedDeviationsSketch = TDigestState.createUsingParamsFrom(valuesSketch);
-            valuesSketch.centroids().forEach(centroid -> {
-                final double deviation = Math.abs(approximateMedian - centroid.mean());
-                approximatedDeviationsSketch.add(deviation, centroid.count());
-            });
+            try (TDigestState approximatedDeviationsSketch = TDigestState.createUsingParamsFrom(valuesSketch)) {
+                valuesSketch.centroids().forEach(centroid -> {
+                    final double deviation = Math.abs(approximateMedian - centroid.mean());
+                    approximatedDeviationsSketch.add(deviation, centroid.count());
+                });
 
-            return approximatedDeviationsSketch.quantile(0.5);
+                return approximatedDeviationsSketch.quantile(0.5);
+            }
         }
     }
 
@@ -69,7 +70,12 @@ static InternalMedianAbsoluteDeviation empty(
         double compression,
         TDigestExecutionHint executionHint
     ) {
-        return new InternalMedianAbsoluteDeviation(name, metadata, format, TDigestState.create(compression, executionHint));
+        return new InternalMedianAbsoluteDeviation(
+            name,
+            metadata,
+            format,
+            TDigestState.createWithoutCircuitBreaking(compression, executionHint)
+        );
     }
 
     @Override

server/src/main/java/org/elasticsearch/search/aggregations/metrics/InternalTDigestPercentileRanks.java

@@ -51,7 +51,7 @@ public static InternalTDigestPercentileRanks empty(
         DocValueFormat format,
         Map<String, Object> metadata
     ) {
-        TDigestState state = TDigestState.create(compression, executionHint);
+        TDigestState state = TDigestState.createWithoutCircuitBreaking(compression, executionHint);
         return new InternalTDigestPercentileRanks(name, keys, state, keyed, format, metadata);
     }
 

server/src/main/java/org/elasticsearch/search/aggregations/metrics/MedianAbsoluteDeviationAggregator.java

@@ -101,7 +101,7 @@ private TDigestState getExistingOrNewHistogram(final BigArrays bigArrays, long b
         valueSketches = bigArrays.grow(valueSketches, bucket + 1);
         TDigestState state = valueSketches.get(bucket);
         if (state == null) {
-            state = TDigestState.create(compression, executionHint);
+            state = TDigestState.createWithoutCircuitBreaking(compression, executionHint);
             valueSketches.set(bucket, state);
         }
         return state;

server/src/main/java/org/elasticsearch/search/aggregations/metrics/MemoryTrackingTDigestArrays.java

@@ -13,7 +13,6 @@
 import org.apache.lucene.util.ArrayUtil;
 import org.apache.lucene.util.RamUsageEstimator;
 import org.elasticsearch.common.breaker.CircuitBreaker;
-import org.elasticsearch.common.breaker.NoopCircuitBreaker;
 import org.elasticsearch.core.Releasable;
 import org.elasticsearch.tdigest.arrays.TDigestArrays;
 import org.elasticsearch.tdigest.arrays.TDigestByteArray;
@@ -29,16 +28,6 @@
  */
 public class MemoryTrackingTDigestArrays implements TDigestArrays {
 
-    /**
-     * Default no-op CB instance of the wrapper.
-     *
-     * @deprecated This instance shouldn't be used, and will be removed after all usages are replaced.
-     */
-    @Deprecated
-    public static final MemoryTrackingTDigestArrays INSTANCE = new MemoryTrackingTDigestArrays(
-        new NoopCircuitBreaker("default-wrapper-tdigest-arrays")
-    );
-
     private final CircuitBreaker breaker;
 
     public MemoryTrackingTDigestArrays(CircuitBreaker breaker) {

server/src/main/java/org/elasticsearch/search/aggregations/metrics/TDigestState.java

@@ -9,13 +9,14 @@
 package org.elasticsearch.search.aggregations.metrics;
 
 import org.elasticsearch.TransportVersions;
+import org.elasticsearch.common.breaker.CircuitBreaker;
+import org.elasticsearch.common.breaker.NoopCircuitBreaker;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.core.Releasable;
 import org.elasticsearch.core.Releasables;
 import org.elasticsearch.tdigest.Centroid;
 import org.elasticsearch.tdigest.TDigest;
-import org.elasticsearch.tdigest.arrays.TDigestArrays;
 
 import java.io.IOException;
 import java.util.Collection;
@@ -28,6 +29,10 @@
  */
 public class TDigestState implements Releasable {
 
+    protected static final CircuitBreaker DEFAULT_NOOP_BREAKER = new NoopCircuitBreaker("default-tdigest-state-noop-breaker");
+
+    private final CircuitBreaker breaker;
+
     private final double compression;
 
     private final TDigest tdigest;
@@ -51,11 +56,12 @@ static Type valueForHighAccuracy() {
     private final Type type;
 
     /**
-     * @deprecated this method will be removed after all usages are replaced
+     * @deprecated No-op circuit-breaked factory for TDigestState. Used in _search aggregations.
+     *             Please use the {@link #create(CircuitBreaker, double)} method instead on new usages.
      */
     @Deprecated
-    public static TDigestState create(double compression) {
-        return create(MemoryTrackingTDigestArrays.INSTANCE, compression);
+    public static TDigestState createWithoutCircuitBreaking(double compression) {
+        return create(DEFAULT_NOOP_BREAKER, compression);
     }
 
     /**
@@ -64,25 +70,26 @@ public static TDigestState create(double compression) {
      * @param compression the compression factor for the underlying {@link org.elasticsearch.tdigest.TDigest} object
      * @return a TDigestState object that's optimized for performance
      */
-    public static TDigestState create(TDigestArrays arrays, double compression) {
-        return new TDigestState(arrays, Type.defaultValue(), compression);
+    public static TDigestState create(CircuitBreaker breaker, double compression) {
+        return new TDigestState(breaker, Type.defaultValue(), compression);
     }
 
     /**
      * Factory for TDigestState that's optimized for high accuracy. It's substantially slower than the default implementation.
      * @param compression the compression factor for the underlying {@link org.elasticsearch.tdigest.TDigest} object
      * @return a TDigestState object that's optimized for performance
      */
-    public static TDigestState createOptimizedForAccuracy(TDigestArrays arrays, double compression) {
-        return new TDigestState(arrays, Type.valueForHighAccuracy(), compression);
+    static TDigestState createOptimizedForAccuracy(CircuitBreaker breaker, double compression) {
+        return new TDigestState(breaker, Type.valueForHighAccuracy(), compression);
     }
 
     /**
-     * @deprecated this method will be removed after all usages are replaced
+     * @deprecated No-op circuit-breaked factory for TDigestState. Used in _search aggregations.
+     *             Please use the {@link #create(CircuitBreaker, double, TDigestExecutionHint)} method instead on new usages.
      */
     @Deprecated
-    public static TDigestState create(double compression, TDigestExecutionHint executionHint) {
-        return create(MemoryTrackingTDigestArrays.INSTANCE, compression, executionHint);
+    public static TDigestState createWithoutCircuitBreaking(double compression, TDigestExecutionHint executionHint) {
+        return create(DEFAULT_NOOP_BREAKER, compression, executionHint);
     }
 
     /**
@@ -93,10 +100,10 @@ public static TDigestState create(double compression, TDigestExecutionHint execu
      * @param executionHint controls which implementation is used; accepted values are 'high_accuracy' and '' (default)
      * @return a TDigestState object
      */
-    public static TDigestState create(TDigestArrays arrays, double compression, TDigestExecutionHint executionHint) {
+    public static TDigestState create(CircuitBreaker breaker, double compression, TDigestExecutionHint executionHint) {
         return switch (executionHint) {
-            case HIGH_ACCURACY -> createOptimizedForAccuracy(arrays, compression);
-            case DEFAULT -> create(arrays, compression);
+            case HIGH_ACCURACY -> createOptimizedForAccuracy(breaker, compression);
+            case DEFAULT -> create(breaker, compression);
         };
     }
 
@@ -107,10 +114,12 @@ public static TDigestState create(TDigestArrays arrays, double compression, TDig
      * @return a TDigestState object
      */
     public static TDigestState createUsingParamsFrom(TDigestState state) {
-        return new TDigestState(MemoryTrackingTDigestArrays.INSTANCE, state.type, state.compression);
+        return new TDigestState(state.breaker, state.type, state.compression);
     }
 
-    protected TDigestState(TDigestArrays arrays, Type type, double compression) {
+    protected TDigestState(CircuitBreaker breaker, Type type, double compression) {
+        this.breaker = breaker;
+        var arrays = new MemoryTrackingTDigestArrays(breaker);
         tdigest = switch (type) {
             case HYBRID -> TDigest.createHybridDigest(arrays, compression);
             case AVL_TREE -> TDigest.createAvlTreeDigest(arrays, compression);
@@ -140,22 +149,23 @@ public static void write(TDigestState state, StreamOutput out) throws IOExceptio
     }
 
     /**
-     * @deprecated this method will be removed after all usages are replaced
+     * @deprecated No-op circuit-breaked factory for TDigestState. Used in _search aggregations.
+     *             Please use the {@link #read(CircuitBreaker, StreamInput)} method instead on new usages.
      */
     @Deprecated
     public static TDigestState read(StreamInput in) throws IOException {
-        return read(MemoryTrackingTDigestArrays.INSTANCE, in);
+        return read(DEFAULT_NOOP_BREAKER, in);
     }
 
-    public static TDigestState read(TDigestArrays arrays, StreamInput in) throws IOException {
+    public static TDigestState read(CircuitBreaker breaker, StreamInput in) throws IOException {
         double compression = in.readDouble();
         TDigestState state;
         long size = 0;
         if (in.getTransportVersion().onOrAfter(TransportVersions.V_8_9_X)) {
-            state = new TDigestState(arrays, Type.valueOf(in.readString()), compression);
+            state = new TDigestState(breaker, Type.valueOf(in.readString()), compression);
             size = in.readVLong();
         } else {
-            state = new TDigestState(arrays, Type.valueForHighAccuracy(), compression);
+            state = new TDigestState(breaker, Type.valueForHighAccuracy(), compression);
         }
         int n = in.readVInt();
         if (size > 0) {

server/src/test/java/org/elasticsearch/search/aggregations/metrics/InternalMedianAbsoluteDeviationTests.java

@@ -23,7 +23,7 @@ public class InternalMedianAbsoluteDeviationTests extends InternalAggregationTes
 
     @Override
     protected InternalMedianAbsoluteDeviation createTestInstance(String name, Map<String, Object> metadata) {
-        final TDigestState valuesSketch = TDigestState.create(randomFrom(50.0, 100.0, 200.0, 500.0, 1000.0));
+        final TDigestState valuesSketch = TDigestState.createWithoutCircuitBreaking(randomFrom(50.0, 100.0, 200.0, 500.0, 1000.0));
         final int numberOfValues = frequently() ? randomIntBetween(0, 1000) : 0;
         for (int i = 0; i < numberOfValues; i++) {
             valuesSketch.add(randomDouble());

server/src/test/java/org/elasticsearch/search/aggregations/metrics/InternalTDigestPercentilesRanksTests.java

@@ -36,7 +36,7 @@ protected InternalTDigestPercentileRanks createTestInstance(
         if (empty) {
             return new InternalTDigestPercentileRanks(name, percents, null, keyed, format, metadata);
         }
-        final TDigestState state = TDigestState.create(100);
+        final TDigestState state = TDigestState.createWithoutCircuitBreaking(100);
         Arrays.stream(values).forEach(state::add);
 
         return new InternalTDigestPercentileRanks(name, percents, state, keyed, format, metadata);