fix: keep late observations out of subsequent collection buffers

prometheus-metrics-core/src/main/java/io/prometheus/metrics/core/metrics/Buffer.java

@@ -62,23 +62,34 @@ private static final class Generation {
   private final Condition bufferSpaceAvailable = appendLock.newCondition();
   private final long maxSpinWaitNanos;
   private final int maxBufferSize;
-  private final Runnable beforeAppendLock;
+  // These hooks are test seams only; production buffers use no-op callbacks.
+  private final Runnable beforeGenerationRead;
+  private final Runnable afterGenerationRead;
 
   Buffer() {
-    this(DEFAULT_MAX_SPIN_WAIT_NANOS, DEFAULT_MAX_BUFFER_SIZE, () -> {});
+    this(DEFAULT_MAX_SPIN_WAIT_NANOS, DEFAULT_MAX_BUFFER_SIZE, () -> {}, () -> {});
   }
 
   Buffer(long maxSpinWaitNanos) {
-    this(maxSpinWaitNanos, DEFAULT_MAX_BUFFER_SIZE, () -> {});
+    this(maxSpinWaitNanos, DEFAULT_MAX_BUFFER_SIZE, () -> {}, () -> {});
   }
 
-  Buffer(long maxSpinWaitNanos, int maxBufferSize, Runnable beforeAppendLock) {
+  Buffer(long maxSpinWaitNanos, int maxBufferSize, Runnable beforeGenerationRead) {
+    this(maxSpinWaitNanos, maxBufferSize, beforeGenerationRead, () -> {});
+  }
+
+  Buffer(
+      long maxSpinWaitNanos,
+      int maxBufferSize,
+      Runnable beforeGenerationRead,
+      Runnable afterGenerationRead) {
     if (maxBufferSize <= 0) {
       throw new IllegalArgumentException("maxBufferSize must be positive");
     }
     this.maxSpinWaitNanos = maxSpinWaitNanos;
     this.maxBufferSize = maxBufferSize;
-    this.beforeAppendLock = beforeAppendLock;
+    this.beforeGenerationRead = beforeGenerationRead;
+    this.afterGenerationRead = afterGenerationRead;
     stripedObservationCounts = new AtomicLong[Runtime.getRuntime().availableProcessors()];
     generationStartCounts = new long[stripedObservationCounts.length];
     for (int i = 0; i < stripedObservationCounts.length; i++) {
@@ -87,6 +98,7 @@ private static final class Generation {
   }
 
   boolean append(double value) {
+    // Keep the uncontended hot path small enough for the JIT to inline into observations.
     int stripe = stripeIndex(Thread.currentThread().getId(), stripedObservationCounts.length);
     AtomicLong counter = stripedObservationCounts[stripe];
     long count = counter.incrementAndGet();
@@ -97,9 +109,14 @@ boolean append(double value) {
     if ((count & BUFFER_ACTIVE_BIT) == 0) {
       return false;
     }
+    return appendToActiveGeneration(value, stripe, count);
+  }
+
+  private boolean appendToActiveGeneration(double value, int stripe, long count) {
     // Allow tests to pause between allocating an observation ticket and reading the generation.
-    beforeAppendLock.run();
+    beforeGenerationRead.run();
     Generation generation = activeGeneration;
+    afterGenerationRead.run();
     if (generation == null) {
       return false;
     }
@@ -193,7 +210,7 @@ <T extends DataPointSnapshot> T run(
         long total = 0;
         for (int i = 0; i < stripedObservationCounts.length; i++) {
           long count = stripedObservationCounts[i].getAndAdd(BUFFER_ACTIVE_BIT);
-          generationStartCounts[i] = count;
+          generationStartCounts[i] = count & ~BUFFER_ACTIVE_BIT;
           total += count;
         }
         expectedCount = total - observationCountOffset;

prometheus-metrics-core/src/test/java/io/prometheus/metrics/core/metrics/BufferTest.java

@@ -223,15 +223,21 @@ void interruptedAppenderLeavesBoundedBufferWait() throws InterruptedException {
 
   @Test
   void lateAppenderCountedByNextGenerationMustNotBeBufferedAgain() throws Exception {
-    assertLateAppenderHandoff(false);
+    assertLateAppenderHandoff(false, true);
   }
 
   @Test
   void lateAppenderHandoffUsesAbsoluteStripeCountsAfterReset() throws Exception {
-    assertLateAppenderHandoff(true);
+    assertLateAppenderHandoff(true, true);
   }
 
-  private static void assertLateAppenderHandoff(boolean reset) throws Exception {
+  @Test
+  void lateAppenderAfterGenerationReadMustNotBeBufferedAgain() throws Exception {
+    assertLateAppenderHandoff(false, false);
+  }
+
+  private static void assertLateAppenderHandoff(boolean reset, boolean pauseBeforeGenerationRead)
+      throws Exception {
     CountDownLatch firstSnapshotStarted = new CountDownLatch(1);
     CountDownLatch finishFirstSnapshot = new CountDownLatch(1);
     CountDownLatch observationCounted = new CountDownLatch(1);
@@ -240,16 +246,19 @@ private static void assertLateAppenderHandoff(boolean reset) throws Exception {
     AtomicLong completedObservations = new AtomicLong();
     AtomicLong secondExpectedCount = new AtomicLong();
     AtomicBoolean pauseFirstAppender = new AtomicBoolean(true);
+    Runnable pauseHook =
+        () -> {
+          if (pauseFirstAppender.compareAndSet(true, false)) {
+            observationCounted.countDown();
+            awaitLatch(readGeneration);
+          }
+        };
     Buffer buffer =
         new Buffer(
             TimeUnit.SECONDS.toNanos(5),
             16,
-            () -> {
-              if (pauseFirstAppender.compareAndSet(true, false)) {
-                observationCounted.countDown();
-                awaitLatch(readGeneration);
-              }
-            });
+            pauseBeforeGenerationRead ? pauseHook : () -> {},
+            pauseBeforeGenerationRead ? () -> {} : pauseHook);
     if (reset) {
       assertThat(buffer.append(1.0)).isFalse();
       buffer.observeDirect(completedObservations::incrementAndGet);
@@ -329,8 +338,8 @@ private static void assertLateAppenderHandoff(boolean reset) throws Exception {
       finishFirstSnapshot.countDown();
       readGeneration.countDown();
       executor.shutdownNow();
-      assertThat(executor.awaitTermination(10, TimeUnit.SECONDS)).isTrue();
     }
+    assertThat(executor.awaitTermination(10, TimeUnit.SECONDS)).as("executor terminated").isTrue();
   }
 
   private static void awaitLatch(CountDownLatch latch) {