Add CircuitBreaker to TDigest, Step 2: Add CB to array wrappers (#113105)

Part of #99815

## Steps
1. Migrate TDigest classes to use a custom Array implementation. Temporarily use a simple array wrapper (#112810)
2. Implement CircuitBreaking in the `MemoryTrackingTDigestArrays` class. Add `Releasable` and ensure it's always closed within TDigest (This PR)
3. Pass the CircuitBreaker as a parameter to TDigestState from wherever it's being used
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
To simplify testing the CircuitBreaker, added a helper method + `@After` to ESTestCase.

Right now CBs are usually tested through MockBigArrays. E.g:
https://github.com/elastic/elasticsearch/blob/f7a0196b454b17f7928728a26084000238c4efaa/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/AbstractFunctionTestCase.java#L1263-L1265
So I guess there was no need for this yet. But I may have missed something somewhere.

Also, I'm separating this PR from the "step 3" as integrating this (CB) in the current usages may require some refactor of external code, which may be somewhat more _dangerous_

Author: ivancea

benchmarks/src/main/java/org/elasticsearch/benchmark/tdigest/SortBench.java

@@ -21,10 +21,12 @@
 
 package org.elasticsearch.benchmark.tdigest;
 
+import org.elasticsearch.common.breaker.NoopCircuitBreaker;
+import org.elasticsearch.search.aggregations.metrics.MemoryTrackingTDigestArrays;
 import org.elasticsearch.tdigest.Sort;
+import org.elasticsearch.tdigest.arrays.TDigestArrays;
 import org.elasticsearch.tdigest.arrays.TDigestDoubleArray;
 import org.elasticsearch.tdigest.arrays.TDigestIntArray;
-import org.elasticsearch.tdigest.arrays.WrapperTDigestArrays;
 import org.openjdk.jmh.annotations.Benchmark;
 import org.openjdk.jmh.annotations.BenchmarkMode;
 import org.openjdk.jmh.annotations.Fork;
@@ -51,7 +53,8 @@
 @State(Scope.Thread)
 public class SortBench {
     private final int size = 100000;
-    private final TDigestDoubleArray values = WrapperTDigestArrays.INSTANCE.newDoubleArray(size);
+    private final TDigestArrays arrays = new MemoryTrackingTDigestArrays(new NoopCircuitBreaker("default-wrapper-tdigest-arrays"));
+    private final TDigestDoubleArray values = arrays.newDoubleArray(size);
 
     @Param({ "0", "1", "-1" })
     public int sortDirection;
@@ -72,7 +75,7 @@ public void setup() {
 
     @Benchmark
     public void stableSort() {
-        TDigestIntArray order = WrapperTDigestArrays.INSTANCE.newIntArray(size);
+        TDigestIntArray order = arrays.newIntArray(size);
         for (int i = 0; i < size; i++) {
             order.set(i, i);
         }

benchmarks/src/main/java/org/elasticsearch/benchmark/tdigest/TDigestBench.java

@@ -21,9 +21,11 @@
 
 package org.elasticsearch.benchmark.tdigest;
 
+import org.elasticsearch.common.breaker.NoopCircuitBreaker;
+import org.elasticsearch.search.aggregations.metrics.MemoryTrackingTDigestArrays;
 import org.elasticsearch.tdigest.MergingDigest;
 import org.elasticsearch.tdigest.TDigest;
-import org.elasticsearch.tdigest.arrays.WrapperTDigestArrays;
+import org.elasticsearch.tdigest.arrays.TDigestArrays;
 import org.openjdk.jmh.annotations.Benchmark;
 import org.openjdk.jmh.annotations.BenchmarkMode;
 import org.openjdk.jmh.annotations.Fork;
@@ -56,24 +58,25 @@
 @Threads(1)
 @State(Scope.Thread)
 public class TDigestBench {
+    private static final TDigestArrays arrays = new MemoryTrackingTDigestArrays(new NoopCircuitBreaker("default-wrapper-tdigest-arrays"));
 
     public enum TDigestFactory {
         MERGE {
             @Override
             TDigest create(double compression) {
-                return new MergingDigest(WrapperTDigestArrays.INSTANCE, compression, (int) (10 * compression));
+                return new MergingDigest(arrays, compression, (int) (10 * compression));
             }
         },
         AVL_TREE {
             @Override
             TDigest create(double compression) {
-                return TDigest.createAvlTreeDigest(WrapperTDigestArrays.INSTANCE, compression);
+                return TDigest.createAvlTreeDigest(arrays, compression);
             }
         },
         HYBRID {
             @Override
             TDigest create(double compression) {
-                return TDigest.createHybridDigest(WrapperTDigestArrays.INSTANCE, compression);
+                return TDigest.createHybridDigest(arrays, compression);
             }
         };
 

libs/tdigest/build.gradle

@@ -22,6 +22,8 @@ apply plugin: 'elasticsearch.build'
 apply plugin: 'elasticsearch.publish'
 
 dependencies {
+  api project(':libs:elasticsearch-core')
+
   testImplementation(project(":test:framework")) {
     exclude group: 'org.elasticsearch', module: 'elasticsearch-tdigest'
   }

libs/tdigest/src/main/java/module-info.java

@@ -18,6 +18,8 @@
  */
 
 module org.elasticsearch.tdigest {
+    requires org.elasticsearch.base;
+
     exports org.elasticsearch.tdigest;
     exports org.elasticsearch.tdigest.arrays;
 }

libs/tdigest/src/main/java/org/elasticsearch/tdigest/AVLGroupTree.java

@@ -21,6 +21,8 @@
 
 package org.elasticsearch.tdigest;
 
+import org.elasticsearch.core.Releasable;
+import org.elasticsearch.core.Releasables;
 import org.elasticsearch.tdigest.arrays.TDigestArrays;
 import org.elasticsearch.tdigest.arrays.TDigestDoubleArray;
 import org.elasticsearch.tdigest.arrays.TDigestLongArray;
@@ -31,7 +33,7 @@
 /**
  * A tree of t-digest centroids.
  */
-final class AVLGroupTree extends AbstractCollection<Centroid> {
+final class AVLGroupTree extends AbstractCollection<Centroid> implements Releasable {
     /* For insertions into the tree */
     private double centroid;
     private long count;
@@ -267,4 +269,8 @@ private void checkAggregates(int node) {
         }
     }
 
+    @Override
+    public void close() {
+        Releasables.close(centroids, counts, aggregatedCounts, tree);
+    }
 }

libs/tdigest/src/main/java/org/elasticsearch/tdigest/AVLTreeDigest.java

@@ -21,6 +21,7 @@
 
 package org.elasticsearch.tdigest;
 
+import org.elasticsearch.core.Releasables;
 import org.elasticsearch.tdigest.arrays.TDigestArrays;
 
 import java.util.Collection;
@@ -153,26 +154,27 @@ public void compress() {
         }
         needsCompression = false;
 
-        AVLGroupTree centroids = summary;
-        this.summary = new AVLGroupTree(arrays);
+        try (AVLGroupTree centroids = summary) {
+            this.summary = new AVLGroupTree(arrays);
 
-        final int[] nodes = new int[centroids.size()];
-        nodes[0] = centroids.first();
-        for (int i = 1; i < nodes.length; ++i) {
-            nodes[i] = centroids.next(nodes[i - 1]);
-            assert nodes[i] != IntAVLTree.NIL;
-        }
-        assert centroids.next(nodes[nodes.length - 1]) == IntAVLTree.NIL;
+            final int[] nodes = new int[centroids.size()];
+            nodes[0] = centroids.first();
+            for (int i = 1; i < nodes.length; ++i) {
+                nodes[i] = centroids.next(nodes[i - 1]);
+                assert nodes[i] != IntAVLTree.NIL;
+            }
+            assert centroids.next(nodes[nodes.length - 1]) == IntAVLTree.NIL;
 
-        for (int i = centroids.size() - 1; i > 0; --i) {
-            final int other = gen.nextInt(i + 1);
-            final int tmp = nodes[other];
-            nodes[other] = nodes[i];
-            nodes[i] = tmp;
-        }
+            for (int i = centroids.size() - 1; i > 0; --i) {
+                final int other = gen.nextInt(i + 1);
+                final int tmp = nodes[other];
+                nodes[other] = nodes[i];
+                nodes[i] = tmp;
+            }
 
-        for (int node : nodes) {
-            add(centroids.mean(node), centroids.count(node));
+            for (int node : nodes) {
+                add(centroids.mean(node), centroids.count(node));
+            }
         }
     }
 
@@ -356,4 +358,9 @@ public int byteSize() {
         compress();
         return 64 + summary.size() * 13;
     }
+
+    @Override
+    public void close() {
+        Releasables.close(summary);
+    }
 }

libs/tdigest/src/main/java/org/elasticsearch/tdigest/HybridDigest.java

@@ -19,6 +19,7 @@
 
 package org.elasticsearch.tdigest;
 
+import org.elasticsearch.core.Releasables;
 import org.elasticsearch.tdigest.arrays.TDigestArrays;
 
 import java.util.Collection;
@@ -110,6 +111,7 @@ public void reserve(long size) {
             }
             mergingDigest.reserve(size);
             // Release the allocated SortingDigest.
+            sortingDigest.close();
             sortingDigest = null;
         } else {
             sortingDigest.reserve(size);
@@ -196,4 +198,9 @@ public int byteSize() {
         }
         return sortingDigest.byteSize();
     }
+
+    @Override
+    public void close() {
+        Releasables.close(sortingDigest, mergingDigest);
+    }
 }

libs/tdigest/src/main/java/org/elasticsearch/tdigest/IntAVLTree.java

@@ -21,6 +21,8 @@
 
 package org.elasticsearch.tdigest;
 
+import org.elasticsearch.core.Releasable;
+import org.elasticsearch.core.Releasables;
 import org.elasticsearch.tdigest.arrays.TDigestArrays;
 import org.elasticsearch.tdigest.arrays.TDigestByteArray;
 import org.elasticsearch.tdigest.arrays.TDigestIntArray;
@@ -33,7 +35,7 @@
  * want to add data to the nodes, typically by using arrays and node
  * identifiers as indices.
  */
-abstract class IntAVLTree {
+abstract class IntAVLTree implements Releasable {
     /**
      * We use <code>0</code> instead of <code>-1</code> so that left(NIL) works without
      * condition.
@@ -586,4 +588,8 @@ int size() {
 
     }
 
+    @Override
+    public void close() {
+        Releasables.close(parent, left, right, depth);
+    }
 }

libs/tdigest/src/main/java/org/elasticsearch/tdigest/MergingDigest.java

@@ -21,6 +21,7 @@
 
 package org.elasticsearch.tdigest;
 
+import org.elasticsearch.core.Releasables;
 import org.elasticsearch.tdigest.arrays.TDigestArrays;
 import org.elasticsearch.tdigest.arrays.TDigestDoubleArray;
 import org.elasticsearch.tdigest.arrays.TDigestIntArray;
@@ -66,8 +67,6 @@
  * what the AVLTreeDigest uses and no dynamic allocation is required at all.
  */
 public class MergingDigest extends AbstractTDigest {
-    private final TDigestArrays arrays;
-
     private int mergeCount = 0;
 
     private final double publicCompression;
@@ -138,8 +137,6 @@ public MergingDigest(TDigestArrays arrays, double compression, int bufferSize) {
      * @param size        Size of main buffer
      */
     public MergingDigest(TDigestArrays arrays, double compression, int bufferSize, int size) {
-        this.arrays = arrays;
-
         // ensure compression >= 10
         // default size = 2 * ceil(compression)
         // default bufferSize = 5 * size
@@ -274,9 +271,6 @@ private void merge(
         incomingWeight.set(incomingCount, weight, 0, lastUsedCell);
         incomingCount += lastUsedCell;
 
-        if (incomingOrder == null) {
-            incomingOrder = arrays.newIntArray(incomingCount);
-        }
         Sort.stableSort(incomingOrder, incomingMean, incomingCount);
 
         totalWeight += unmergedWeight;
@@ -581,4 +575,9 @@ public String toString() {
             + "-"
             + (useTwoLevelCompression ? "twoLevel" : "oneLevel");
     }
+
+    @Override
+    public void close() {
+        Releasables.close(weight, mean, tempWeight, tempMean, order);
+    }
 }

libs/tdigest/src/main/java/org/elasticsearch/tdigest/SortingDigest.java

@@ -19,6 +19,7 @@
 
 package org.elasticsearch.tdigest;
 
+import org.elasticsearch.core.Releasables;
 import org.elasticsearch.tdigest.arrays.TDigestArrays;
 import org.elasticsearch.tdigest.arrays.TDigestDoubleArray;
 
@@ -137,4 +138,9 @@ public void reserve(long size) {
     public int byteSize() {
         return values.size() * 8;
     }
+
+    @Override
+    public void close() {
+        Releasables.close(values);
+    }
 }