Fix bin count when bins are being split

Author: vickenty

dogstatsd-http-core/src/main/java/com/datadoghq/dogstatsd/http/serializer/SketchMetric.java

@@ -12,6 +12,31 @@
 
 /** Builder for sketch timeseries. */
 public class SketchMetric extends Metric<SketchMetric> {
+    static final long maxBinCount = (1L << 32) - 1;
+    static final long maxBinBytes = ProtoUtil.varintLen(maxBinCount);
+
+    static class BinConsumer implements Sketch.BinConsumer {
+        int numBins;
+        ColumnarBuffer r;
+        DeltaEncoder dk = new DeltaEncoder();
+
+        BinConsumer(ColumnarBuffer record) {
+            r = record;
+        }
+
+        @Override
+        public void consumeBin(short key, long count) {
+            while (count > maxBinCount) {
+                r.putSint64(Column.sketchBinKeys, dk.encode(key));
+                r.putUint64(Column.sketchBinCnts, maxBinCount);
+                count -= maxBinCount;
+                numBins++;
+            }
+            r.putSint64(Column.sketchBinKeys, dk.encode(key));
+            r.putUint64(Column.sketchBinCnts, count);
+            numBins++;
+        }
+    }
 
     SketchMetric(PayloadBuilder pb, int type, String name) {
         super(pb, type, name);
@@ -30,9 +55,6 @@ protected SketchMetric self() {
      * @return This.
      */
     public SketchMetric addPoint(long timestamp, Sketch sketch) {
-        final long maxBinCount = (1L << 32) - 1;
-        final long maxBinBytes = ProtoUtil.varintLen(maxBinCount);
-
         // Skip doing the work if just the bin data would exceed payload size limit.
         if (sketch.count() / maxBinCount * maxBinBytes >= pb.maxPayloadSize) {
             throw new BufferOverflowException();
@@ -45,22 +67,10 @@ public SketchMetric addPoint(long timestamp, Sketch sketch) {
         pb.counts.put(sketch.count());
 
         final ColumnarBuffer r = pb.currentRecord();
-        final DeltaEncoder dk = new DeltaEncoder();
-
-        r.putUint64(Column.sketchNumBins, sketch.size());
-        sketch.bins(
-                new Sketch.BinConsumer() {
-                    @Override
-                    public void consumeBin(short key, long count) {
-                        while (count > maxBinCount) {
-                            r.putSint64(Column.sketchBinKeys, dk.encode(key));
-                            r.putUint64(Column.sketchBinCnts, maxBinCount);
-                            count -= maxBinCount;
-                        }
-                        r.putSint64(Column.sketchBinKeys, dk.encode(key));
-                        r.putUint64(Column.sketchBinCnts, count);
-                    }
-                });
+        final BinConsumer bc = new BinConsumer(r);
+
+        sketch.bins(bc);
+        r.putUint64(Column.sketchNumBins, bc.numBins);
 
         return this;
     }

dogstatsd-http-core/src/test/java/com/datadoghq/dogstatsd/http/serializer/PayloadBuilderTest.java

@@ -284,7 +284,7 @@ public void handle(byte[] p) {
                     1,
                     2,
                     2,
-                    2,
+                    3,
                     // sketchBinKeys
                     (21 << 3) | 2,
                     1,