Fix thread-safety of FixedCapacityExponentialHistogram.valueCount (elastic#137357)

Author: JonasKunz

libs/exponential-histogram/src/main/java/org/elasticsearch/exponentialhistogram/FixedCapacityExponentialHistogram.java

@@ -30,6 +30,10 @@
  * <br>
  * Consumers must ensure that if the histogram is mutated, all previously acquired {@link BucketIterator}
  * instances are no longer used.
+ * <br>
+ * This implementation is thread-safe for all operations provided via {@link ReleasableExponentialHistogram} and its superclasses,
+ * as long as it is not mutated concurrently using any of the methods declared in addition in this class
+ * (e.g. {@link #tryAddBucket(long, long, boolean)}).
  */
 final class FixedCapacityExponentialHistogram extends AbstractExponentialHistogram implements ReleasableExponentialHistogram {
 
@@ -243,8 +247,10 @@ private class Buckets implements ExponentialHistogram.Buckets {
 
         private final boolean isPositive;
         private int numBuckets;
-        private int cachedValueSumForNumBuckets;
-        private long cachedValueSum;
+
+        private record CachedCountsSum(int numBuckets, long countsSum) {}
+
+        private CachedCountsSum cachedCountsSum;
 
         /**
          * @param isPositive true, if this object should represent the positive bucket range, false for the negative range
@@ -263,8 +269,7 @@ int startSlot() {
 
         final void reset() {
             numBuckets = 0;
-            cachedValueSumForNumBuckets = 0;
-            cachedValueSum = 0;
+            cachedCountsSum = null;
         }
 
         boolean tryAddBucket(long index, long count) {
@@ -297,12 +302,26 @@ public OptionalLong maxBucketIndex() {
 
         @Override
         public long valueCount() {
+            // copy a reference to the field to avoid problems with concurrent updates
+            CachedCountsSum cachedVal = cachedCountsSum;
+            if (cachedVal != null && cachedVal.numBuckets == numBuckets) {
+                return cachedVal.countsSum;
+            }
+
+            long countsSum = 0;
+            int position = 0;
+            if (cachedVal != null) {
+                countsSum = cachedVal.countsSum;
+                position = cachedVal.numBuckets;
+            }
+
             int startSlot = startSlot();
-            while (cachedValueSumForNumBuckets < numBuckets) {
-                cachedValueSum += bucketCounts[startSlot + cachedValueSumForNumBuckets];
-                cachedValueSumForNumBuckets++;
+            while (position < numBuckets) {
+                countsSum += bucketCounts[startSlot + position];
+                position++;
             }
-            return cachedValueSum;
+            this.cachedCountsSum = new CachedCountsSum(position, countsSum);
+            return countsSum;
         }
 
         @Override

libs/exponential-histogram/src/test/java/org/elasticsearch/exponentialhistogram/FixedCapacityExponentialHistogramTests.java

@@ -24,11 +24,39 @@
 import org.elasticsearch.common.breaker.CircuitBreaker;
 import org.elasticsearch.common.unit.ByteSizeValue;
 
+import java.util.ArrayList;
+import java.util.List;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
+import java.util.stream.IntStream;
+
 import static org.hamcrest.Matchers.equalTo;
 import static org.hamcrest.Matchers.greaterThan;
 
 public class FixedCapacityExponentialHistogramTests extends ExponentialHistogramTestCase {
 
+    public void testConcurrentHashCode() throws ExecutionException, InterruptedException {
+        List<ExponentialHistogram> originalHistograms = IntStream.range(0, 1000)
+            .mapToObj(i -> ExponentialHistogramTestUtils.randomHistogram())
+            .toList();
+
+        List<? extends ExponentialHistogram> copies = originalHistograms.stream()
+            .map(histo -> ExponentialHistogram.builder(histo, ExponentialHistogramCircuitBreaker.noop()).build())
+            .toList();
+
+        // Compute potentially lazy data correctly on the originals
+        originalHistograms.forEach(Object::hashCode);
+        concurrentTest(() -> {
+            for (int i = 0; i < originalHistograms.size(); i++) {
+                ExponentialHistogram original = originalHistograms.get(i);
+                ExponentialHistogram copy = copies.get(i);
+                assertThat(copy.hashCode(), equalTo(original.hashCode()));
+            }
+        });
+    }
+
     public void testValueCountUpdatedCorrectly() {
 
         FixedCapacityExponentialHistogram histogram = FixedCapacityExponentialHistogram.create(100, breaker());
@@ -69,4 +97,21 @@ public void testMemoryAccounting() {
         }
         assertThat(esBreaker.getUsed(), equalTo(0L));
     }
+
+    protected void concurrentTest(Runnable r) throws InterruptedException, ExecutionException {
+        int threads = 5;
+        int tasks = threads * 2;
+        ExecutorService exec = Executors.newFixedThreadPool(threads);
+        try {
+            List<Future<?>> results = new ArrayList<>();
+            for (int t = 0; t < tasks; t++) {
+                results.add(exec.submit(r));
+            }
+            for (Future<?> f : results) {
+                f.get();
+            }
+        } finally {
+            exec.shutdown();
+        }
+    }
 }

x-pack/plugin/esql/qa/testFixtures/src/main/java/org/elasticsearch/xpack/esql/EsqlTestUtils.java

@@ -1060,13 +1060,6 @@ private static ExponentialHistogram randomExponentialHistogram() {
             ExponentialHistogramCircuitBreaker.noop(),
             rawValues
         );
-        /*
-         * hashcode mutates the histogram by doing some lazy work.
-         * That can change the value of the hashCode if you call it concurrently....
-         * This works around that by running it once up front. Jonas will have a
-         * look at this one soon.
-         */
-        histo.hashCode();
         return histo;
     }