ExponentialHistograms: refactor tests to use builder, add copy method

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

@@ -216,6 +216,16 @@ static ExponentialHistogramBuilder builder(int scale, ExponentialHistogramCircui
         return new ExponentialHistogramBuilder(scale, breaker);
     }
 
+    /**
+     * Create a builder for an exponential histogram, which is initialized to copy the given histogram.
+     * @param toCopy the histogram to copy
+     * @param breaker the circuit breaker to use
+     * @return a new builder
+     */
+    static ExponentialHistogramBuilder builder(ExponentialHistogram toCopy, ExponentialHistogramCircuitBreaker breaker) {
+        return new ExponentialHistogramBuilder(toCopy, breaker);
+    }
+
     /**
      * Creates a histogram representing the distribution of the given values with at most the given number of buckets.
      * If the given {@code maxBucketCount} is greater than or equal to the number of values, the resulting histogram will have a

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

@@ -27,13 +27,13 @@
 
 /**
  * A builder for building a {@link ReleasableExponentialHistogram} directly from buckets.
- * Note that this class is not optimized regarding memory allocations, so it is not intended for high-throughput usage.
+ * Note that this class is not optimized regarding memory allocations or performance, so it is not intended for high-throughput usage.
  */
 public class ExponentialHistogramBuilder {
 
     private final ExponentialHistogramCircuitBreaker breaker;
 
-    private final int scale;
+    private int scale;
     private ZeroBucket zeroBucket = ZeroBucket.minimalEmpty();
     private Double sum;
     private Double min;
@@ -47,6 +47,29 @@ public class ExponentialHistogramBuilder {
         this.scale = scale;
     }
 
+    ExponentialHistogramBuilder(ExponentialHistogram toCopy, ExponentialHistogramCircuitBreaker breaker) {
+        this(toCopy.scale(), breaker);
+        zeroBucket(toCopy.zeroBucket());
+        sum(toCopy.sum());
+        min(toCopy.min());
+        max(toCopy.max());
+        BucketIterator negBuckets = toCopy.negativeBuckets().iterator();
+        while (negBuckets.hasNext()) {
+            setNegativeBucket(negBuckets.peekIndex(), negBuckets.peekCount());
+            negBuckets.advance();
+        }
+        BucketIterator posBuckets = toCopy.positiveBuckets().iterator();
+        while (posBuckets.hasNext()) {
+            setPositiveBucket(posBuckets.peekIndex(), posBuckets.peekCount());
+            posBuckets.advance();
+        }
+    }
+
+    public ExponentialHistogramBuilder scale(int scale) {
+        this.scale = scale;
+        return this;
+    }
+
     public ExponentialHistogramBuilder zeroBucket(ZeroBucket zeroBucket) {
         this.zeroBucket = zeroBucket;
         return this;
@@ -83,39 +106,33 @@ public ExponentialHistogramBuilder max(double max) {
     }
 
     /**
-     * Adds the given bucket to the positive buckets.
-     * Buckets may be added in arbitrary order, but each bucket can only be added once.
+     * Sets the given bucket of the positive buckets.
+     * Buckets may be set in arbitrary order. If the bucket already exists, it will be overridden.
      *
      * @param index the index of the bucket
      * @param count the count of the bucket, must be at least 1
      * @return the builder
      */
-    public ExponentialHistogramBuilder addPositiveBucket(long index, long count) {
+    public ExponentialHistogramBuilder setPositiveBucket(long index, long count) {
         if (count < 1) {
             throw new IllegalArgumentException("Bucket count must be at least 1");
         }
-        if (positiveBuckets.containsKey(index)) {
-            throw new IllegalArgumentException("Positive bucket already exists: " + index);
-        }
         positiveBuckets.put(index, count);
         return this;
     }
 
     /**
-     * Adds the given bucket to the negative buckets.
-     * Buckets may be added in arbitrary order, but each bucket can only be added once.
+     * Sets the given bucket of the negative buckets.
+     * Buckets may be set in arbitrary order. If the bucket already exists, it will be overridden.
      *
      * @param index the index of the bucket
      * @param count the count of the bucket, must be at least 1
      * @return the builder
      */
-    public ExponentialHistogramBuilder addNegativeBucket(long index, long count) {
+    public ExponentialHistogramBuilder setNegativeBucket(long index, long count) {
         if (count < 1) {
             throw new IllegalArgumentException("Bucket count must be at least 1");
         }
-        if (negativeBuckets.containsKey(index)) {
-            throw new IllegalArgumentException("Negative bucket already exists: " + index);
-        }
         negativeBuckets.put(index, count);
         return this;
     }
@@ -152,8 +169,6 @@ public ReleasableExponentialHistogram build() {
                 Releasables.close(result);
             }
         }
-
-        // Create histogram
         return result;
     }
 }

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

@@ -27,18 +27,18 @@
 
 public class ExponentialHistogramBuilderTests extends ExponentialHistogramTestCase {
 
-    public void testBuildWithAllFieldsSet() {
+    public void testBuildAndCopyWithAllFieldsSet() {
         ZeroBucket zeroBucket = ZeroBucket.create(1, 2);
         ExponentialHistogramBuilder builder = ExponentialHistogram.builder(3, breaker())
             .zeroBucket(zeroBucket)
             .sum(100.0)
             .min(1.0)
             .max(50.0)
-            .addPositiveBucket(2, 10)
-            .addPositiveBucket(0, 1)
-            .addPositiveBucket(5, 2)
-            .addNegativeBucket(-2, 5)
-            .addNegativeBucket(1, 2);
+            .setPositiveBucket(2, 10)
+            .setPositiveBucket(0, 1)
+            .setPositiveBucket(5, 2)
+            .setNegativeBucket(-2, 5)
+            .setNegativeBucket(1, 2);
 
         try (ReleasableExponentialHistogram histogram = builder.build()) {
             assertThat(histogram.scale(), equalTo(3));
@@ -48,14 +48,19 @@ public void testBuildWithAllFieldsSet() {
             assertThat(histogram.max(), equalTo(50.0));
             assertBuckets(histogram.positiveBuckets(), List.of(0L, 2L, 5L), List.of(1L, 10L, 2L));
             assertBuckets(histogram.negativeBuckets(), List.of(-2L, 1L), List.of(5L, 2L));
+
+            try (ReleasableExponentialHistogram copy = ExponentialHistogram.builder(histogram, breaker()).build()) {
+                assertThat(copy, equalTo(histogram));
+            }
         }
+
     }
 
     public void testBuildWithEstimation() {
         ExponentialHistogramBuilder builder = ExponentialHistogram.builder(0, breaker())
-            .addPositiveBucket(0, 1)
-            .addPositiveBucket(1, 1)
-            .addNegativeBucket(0, 4);
+            .setPositiveBucket(0, 1)
+            .setPositiveBucket(1, 1)
+            .setNegativeBucket(0, 4);
 
         try (ReleasableExponentialHistogram histogram = builder.build()) {
             assertThat(histogram.scale(), equalTo(0));
@@ -68,18 +73,6 @@ public void testBuildWithEstimation() {
         }
     }
 
-    public void testAddDuplicatePositiveBucketThrows() {
-        ExponentialHistogramBuilder builder = ExponentialHistogram.builder(0, breaker());
-        builder.addPositiveBucket(1, 10);
-        expectThrows(IllegalArgumentException.class, () -> builder.addPositiveBucket(1, 5));
-    }
-
-    public void testAddDuplicateNegativeBucketThrows() {
-        ExponentialHistogramBuilder builder = ExponentialHistogram.builder(0, breaker());
-        builder.addNegativeBucket(-1, 10);
-        expectThrows(IllegalArgumentException.class, () -> builder.addNegativeBucket(-1, 5));
-    }
-
     private static void assertBuckets(ExponentialHistogram.Buckets buckets, List<Long> indices, List<Long> counts) {
         List<Long> actualIndices = new java.util.ArrayList<>();
         List<Long> actualCounts = new java.util.ArrayList<>();

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

@@ -25,7 +25,6 @@
 
 import java.util.ArrayList;
 import java.util.Arrays;
-import java.util.Collections;
 import java.util.List;
 
 import static org.elasticsearch.exponentialhistogram.ExponentialHistogram.MAX_INDEX;
@@ -86,46 +85,32 @@ private ExponentialHistogram randomHistogram() {
     }
 
     private ExponentialHistogram copyWithModification(ExponentialHistogram toCopy, Modification modification) {
-        int totalBucketCount = getBucketCount(toCopy.positiveBuckets(), toCopy.negativeBuckets());
-        FixedCapacityExponentialHistogram copy = createAutoReleasedHistogram(totalBucketCount + 2);
-        if (modification == Modification.SCALE) {
-            copy.resetBuckets((int) createRandomLongBetweenOtherThan(MIN_SCALE, MAX_SCALE, toCopy.scale()));
-        } else {
-            copy.resetBuckets(toCopy.scale());
-        }
-        if (modification == Modification.SUM) {
-            copy.setSum(randomDouble());
-        } else {
-            copy.setSum(toCopy.sum());
-        }
-        if (modification == Modification.MIN) {
-            copy.setMin(randomDouble());
-        } else {
-            copy.setMin(toCopy.min());
-        }
-        if (modification == Modification.MAX) {
-            copy.setMax(randomDouble());
-        } else {
-            copy.setMax(toCopy.max());
-        }
-        long zeroCount = toCopy.zeroBucket().count();
-        double zeroThreshold = toCopy.zeroBucket().zeroThreshold();
-        if (modification == Modification.ZERO_COUNT) {
-            zeroCount = createRandomLongBetweenOtherThan(0, Long.MAX_VALUE, zeroCount);
-        } else if (modification == Modification.ZERO_THRESHOLD) {
-            zeroThreshold = randomDouble();
-        }
-        copy.setZeroBucket(ZeroBucket.create(zeroThreshold, zeroCount));
-        copyBuckets(copy, toCopy.negativeBuckets(), modification == Modification.NEGATIVE_BUCKETS, false);
-        copyBuckets(copy, toCopy.positiveBuckets(), modification == Modification.POSITIVE_BUCKETS, true);
-
-        return copy;
+        ExponentialHistogramBuilder copyBuilder = ExponentialHistogram.builder(toCopy, breaker());
+        copyBuilder = switch (modification) {
+            case null -> copyBuilder;
+            case SCALE -> copyBuilder.scale((int) createRandomLongBetweenOtherThan(MIN_SCALE, MAX_SCALE, toCopy.scale()));
+            case SUM -> copyBuilder.sum(randomDouble());
+            case MIN -> copyBuilder.min(randomDouble());
+            case MAX -> copyBuilder.max(randomDouble());
+            case ZERO_THRESHOLD -> copyBuilder.zeroBucket(ZeroBucket.create(randomDouble(), toCopy.zeroBucket().count()));
+            case ZERO_COUNT -> copyBuilder.zeroBucket(
+                ZeroBucket.create(
+                    toCopy.zeroBucket().zeroThreshold(),
+                    createRandomLongBetweenOtherThan(0, Long.MAX_VALUE, toCopy.zeroBucket().count())
+                )
+            );
+            case POSITIVE_BUCKETS -> modifyBuckets(copyBuilder, toCopy.positiveBuckets(), true);
+            case NEGATIVE_BUCKETS -> modifyBuckets(copyBuilder, toCopy.negativeBuckets(), false);
+        };
+
+        ReleasableExponentialHistogram result = copyBuilder.build();
+        autoReleaseOnTestEnd(result);
+        return result;
     }
 
-    private void copyBuckets(
-        FixedCapacityExponentialHistogram into,
+    private ExponentialHistogramBuilder modifyBuckets(
+        ExponentialHistogramBuilder builder,
         ExponentialHistogram.Buckets buckets,
-        boolean modify,
         boolean isPositive
     ) {
         List<Long> indices = new ArrayList<>();
@@ -136,29 +121,28 @@ private void copyBuckets(
             counts.add(it.peekCount());
             it.advance();
         }
-        if (modify) {
-            if (counts.isEmpty() == false && randomBoolean()) {
-                int toModify = randomIntBetween(0, indices.size() - 1);
-                counts.set(toModify, createRandomLongBetweenOtherThan(1, Long.MAX_VALUE, counts.get(toModify)));
-            } else {
-                insertRandomBucket(indices, counts);
-            }
+        long toModifyIndex;
+        long toModifyCount;
+        if (counts.isEmpty() == false && randomBoolean()) {
+            // Modify existing bucket
+            int position = randomIntBetween(0, indices.size() - 1);
+            toModifyIndex = indices.get(position);
+            toModifyCount = createRandomLongBetweenOtherThan(1, Long.MAX_VALUE, counts.get(position));
+        } else {
+            // Add new bucket
+            long minIndex = indices.stream().mapToLong(Long::longValue).min().orElse(MIN_INDEX);
+            long maxIndex = indices.stream().mapToLong(Long::longValue).min().orElse(MAX_INDEX);
+            do {
+                toModifyIndex = randomLongBetween(Math.max(MIN_INDEX, minIndex - 10), Math.min(MAX_INDEX, maxIndex + 10));
+            } while (indices.contains(toModifyIndex));
+            toModifyCount = randomLongBetween(1, Long.MAX_VALUE);
         }
-        for (int i = 0; i < indices.size(); i++) {
-            into.tryAddBucket(indices.get(i), counts.get(i), isPositive);
+        if (isPositive) {
+            builder.setPositiveBucket(toModifyIndex, toModifyCount);
+        } else {
+            builder.setNegativeBucket(toModifyIndex, toModifyCount);
         }
-    }
-
-    private void insertRandomBucket(List<Long> indices, List<Long> counts) {
-        long minIndex = indices.stream().mapToLong(Long::longValue).min().orElse(MIN_INDEX);
-        long maxIndex = indices.stream().mapToLong(Long::longValue).min().orElse(MAX_INDEX);
-        long newIndex;
-        do {
-            newIndex = randomLongBetween(Math.max(MIN_INDEX, minIndex - 10), Math.min(MAX_INDEX, maxIndex + 10));
-        } while (indices.contains(newIndex));
-        int position = -(Collections.binarySearch(indices, newIndex) + 1);
-        indices.add(position, newIndex);
-        counts.add(position, randomLongBetween(1, Long.MAX_VALUE));
+        return builder;
     }
 
     private static long createRandomLongBetweenOtherThan(long min, long max, long notAllowedValue) {
@@ -169,16 +153,4 @@ private static long createRandomLongBetweenOtherThan(long min, long max, long no
         } while (result == notAllowedValue);
         return result;
     }
-
-    private static int getBucketCount(ExponentialHistogram.Buckets... buckets) {
-        int count = 0;
-        for (ExponentialHistogram.Buckets val : buckets) {
-            BucketIterator it = val.iterator();
-            while (it.hasNext()) {
-                count++;
-                it.advance();
-            }
-        }
-        return count;
-    }
 }