Fix TDigestState.read CB leaks (elastic#114303)

Closes elastic#114194

Fixes `TDigestState.read()` CB leak on error on `.reserve()`.

Author: ivancea

docs/changelog/114303.yaml

@@ -0,0 +1,6 @@
+pr: 114303
+summary: Fix TDigestState.read CB leaks
+area: ES|QL
+type: bug
+issues:
+ - 114194

muted-tests.yml

@@ -328,8 +328,6 @@ tests:
 - class: org.elasticsearch.action.bulk.IncrementalBulkIT
   method: testIncrementalBulkLowWatermarkBackOff
   issue: https://github.com/elastic/elasticsearch/issues/114182
-- class: org.elasticsearch.xpack.esql.action.EsqlActionBreakerIT
-  issue: https://github.com/elastic/elasticsearch/issues/114194
 - class: org.elasticsearch.xpack.ilm.ExplainLifecycleIT
   method: testStepInfoPreservedOnAutoRetry
   issue: https://github.com/elastic/elasticsearch/issues/114220

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

@@ -84,7 +84,7 @@ public static TDigestState create(CircuitBreaker breaker, double compression) {
         }
     }
 
-    static TDigestState create(CircuitBreaker breaker, Type type, double compression) {
+    static TDigestState createOfType(CircuitBreaker breaker, Type type, double compression) {
         breaker.addEstimateBytesAndMaybeBreak(SHALLOW_SIZE, "tdigest-state-create-with-type");
         try {
             return new TDigestState(breaker, type, compression);
@@ -196,28 +196,38 @@ public static TDigestState read(StreamInput in) throws IOException {
 
     public static TDigestState read(CircuitBreaker breaker, StreamInput in) throws IOException {
         double compression = in.readDouble();
-        TDigestState state;
+        TDigestState state = null;
         long size = 0;
-        breaker.addEstimateBytesAndMaybeBreak(SHALLOW_SIZE, "tdigest-state-read");
+        boolean success = false;
         try {
-            if (in.getTransportVersion().onOrAfter(TransportVersions.V_8_9_X)) {
-                state = new TDigestState(breaker, Type.valueOf(in.readString()), compression);
-                size = in.readVLong();
-            } else {
-                state = new TDigestState(breaker, Type.valueForHighAccuracy(), compression);
+            breaker.addEstimateBytesAndMaybeBreak(SHALLOW_SIZE, "tdigest-state-read");
+            try {
+                if (in.getTransportVersion().onOrAfter(TransportVersions.V_8_9_X)) {
+                    state = new TDigestState(breaker, Type.valueOf(in.readString()), compression);
+                    size = in.readVLong();
+                } else {
+                    state = new TDigestState(breaker, Type.valueForHighAccuracy(), compression);
+                }
+            } finally {
+                if (state == null) {
+                    breaker.addWithoutBreaking(-SHALLOW_SIZE);
+                }
+            }
+
+            int n = in.readVInt();
+            if (size > 0) {
+                state.tdigest.reserve(size);
+            }
+            for (int i = 0; i < n; i++) {
+                state.add(in.readDouble(), in.readVLong());
+            }
+            success = true;
+            return state;
+        } finally {
+            if (success == false) {
+                Releasables.close(state);
             }
-        } catch (Exception e) {
-            breaker.addWithoutBreaking(-SHALLOW_SIZE);
-            throw e;
-        }
-        int n = in.readVInt();
-        if (size > 0) {
-            state.tdigest.reserve(size);
-        }
-        for (int i = 0; i < n; i++) {
-            state.add(in.readDouble(), in.readVLong());
         }
-        return state;
     }
 
     @Override

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

@@ -13,9 +13,14 @@
 
 import org.elasticsearch.common.breaker.CircuitBreaker;
 import org.elasticsearch.common.breaker.CircuitBreakingException;
+import org.elasticsearch.common.io.stream.BytesStreamOutput;
+import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.unit.ByteSizeValue;
+import org.elasticsearch.core.CheckedFunction;
+import org.elasticsearch.indices.CrankyCircuitBreakerService;
 import org.elasticsearch.test.ESTestCase;
 
+import java.io.IOException;
 import java.util.Arrays;
 
 import static org.hamcrest.Matchers.equalTo;
@@ -32,28 +37,84 @@ public TDigestStateReleasingTests(TDigestState.Type digestType) {
         this.digestType = digestType;
     }
 
+    public void testCreateOfType() {
+        testCircuitBreakerTrip(circuitBreaker -> TDigestState.createOfType(circuitBreaker, digestType, 100));
+    }
+
+    public void testCreateUsingParamsFrom() {
+        testCircuitBreakerTrip(circuitBreaker -> {
+            try (TDigestState example = TDigestState.createOfType(newLimitedBreaker(ByteSizeValue.ofMb(100)), digestType, 100)) {
+                return TDigestState.createUsingParamsFrom(example);
+            }
+        });
+    }
+
+    /**
+     * This test doesn't use the {@code digestType} param.
+     */
+    public void testCreate() {
+        testCircuitBreakerTrip(circuitBreaker -> TDigestState.create(circuitBreaker, 100));
+    }
+
+    /**
+     * This test doesn't use the {@code digestType} param.
+     */
+    public void testCreateOptimizedForAccuracy() {
+        testCircuitBreakerTrip(circuitBreaker -> TDigestState.createOptimizedForAccuracy(circuitBreaker, 100));
+    }
+
+    public void testRead() throws IOException {
+        try (
+            TDigestState state = TDigestState.createOfType(newLimitedBreaker(ByteSizeValue.ofMb(100)), digestType, 100);
+            BytesStreamOutput output = new BytesStreamOutput()
+        ) {
+            TDigestState.write(state, output);
+
+            testCircuitBreakerTrip(circuitBreaker -> {
+                try (StreamInput input = output.bytes().streamInput()) {
+                    return TDigestState.read(circuitBreaker, input);
+                }
+            });
+        }
+    }
+
+    public void testReadWithData() throws IOException {
+        try (
+            TDigestState state = TDigestState.createOfType(newLimitedBreaker(ByteSizeValue.ofMb(100)), digestType, 100);
+            BytesStreamOutput output = new BytesStreamOutput()
+        ) {
+            for (int i = 0; i < 1000; i++) {
+                state.add(randomDoubleBetween(-Double.MAX_VALUE, Double.MAX_VALUE, true));
+            }
+
+            TDigestState.write(state, output);
+
+            testCircuitBreakerTrip(circuitBreaker -> {
+                try (StreamInput input = output.bytes().streamInput()) {
+                    return TDigestState.read(circuitBreaker, input);
+                }
+            });
+        }
+    }
+
     /**
      * Tests that a circuit breaker trip leaves no unreleased memory.
      */
-    public void testCircuitBreakerTrip() {
-        for (int bytes = randomIntBetween(0, 16); bytes < 50_000; bytes += 17) {
-            CircuitBreaker breaker = newLimitedBreaker(ByteSizeValue.ofBytes(bytes));
+    public <E extends Exception> void testCircuitBreakerTrip(CheckedFunction<CircuitBreaker, TDigestState, E> tDigestStateFactory)
+        throws E {
+        try (CrankyCircuitBreakerService circuitBreakerService = new CrankyCircuitBreakerService()) {
+            CircuitBreaker breaker = circuitBreakerService.getBreaker("test");
 
-            try (TDigestState state = TDigestState.create(breaker, digestType, 100)) {
+            try (TDigestState state = tDigestStateFactory.apply(breaker)) {
                 // Add some data to make it trip. It won't work in all digest types
-                for (int i = 0; i < 100; i++) {
+                for (int i = 0; i < 10; i++) {
                     state.add(randomDoubleBetween(-Double.MAX_VALUE, Double.MAX_VALUE, true));
                 }
-
-                // Testing with more memory shouldn't change anything, we finished the test
-                return;
             } catch (CircuitBreakingException e) {
                 // Expected
             } finally {
-                assertThat("unreleased bytes with a " + bytes + " bytes limit", breaker.getUsed(), equalTo(0L));
+                assertThat("unreleased bytes", breaker.getUsed(), equalTo(0L));
             }
         }
-
-        fail("Test case didn't reach a non-tripping breaker limit");
     }
 }