fix: keep late observations out of subsequent collection buffers

Signed-off-by: Gregor Zeitlinger <gregor.zeitlinger@grafana.com>

Author: zeitlinger

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

@@ -49,6 +49,10 @@ private static final class Generation {
   // available processors. This is simpler than the striping used by LongAdder, so hot spots remain
   // possible when several recording threads resolve to the same stripe.
   private final AtomicLong[] stripedObservationCounts;
+  // Protected by appendLock. These are absolute per-stripe observation counts at activation, not
+  // the reset-adjusted count used by complete. Reused across generations to avoid scrape
+  // allocations.
+  private final long[] generationStartCounts;
   private final ReentrantLock observationLock = new ReentrantLock();
   private boolean reset;
   private long observationCountOffset;
@@ -76,33 +80,41 @@ private static final class Generation {
     this.maxBufferSize = maxBufferSize;
     this.beforeAppendLock = beforeAppendLock;
     stripedObservationCounts = new AtomicLong[Runtime.getRuntime().availableProcessors()];
+    generationStartCounts = new long[stripedObservationCounts.length];
     for (int i = 0; i < stripedObservationCounts.length; i++) {
       stripedObservationCounts[i] = new AtomicLong();
     }
   }
 
   boolean append(double value) {
-    AtomicLong counter =
-        stripedObservationCounts[
-            stripeIndex(Thread.currentThread().getId(), stripedObservationCounts.length)];
+    int stripe = stripeIndex(Thread.currentThread().getId(), stripedObservationCounts.length);
+    AtomicLong counter = stripedObservationCounts[stripe];
     long count = counter.incrementAndGet();
     // The active bit is the exact handoff decision. An observation either increments its stripe
     // before the collector's getAndAdd(BUFFER_ACTIVE_BIT) and takes the direct path, or sees the
-    // active bit and is buffered in the current generation.
+    // active bit and may be buffered. The stripe ticket below also checks that it was not counted
+    // by a later collection that started before this thread read activeGeneration.
     if ((count & BUFFER_ACTIVE_BIT) == 0) {
       return false;
     }
+    // Allow tests to pause between allocating an observation ticket and reading the generation.
+    beforeAppendLock.run();
     Generation generation = activeGeneration;
     if (generation == null) {
       return false;
     }
-    beforeAppendLock.run();
     appendLock.lock();
     try {
       Generation current = activeGeneration;
       if (current != generation || !generation.active) {
         return false;
       }
+      if ((count & ~BUFFER_ACTIVE_BIT) <= generationStartCounts[stripe]) {
+        // This observation incremented its stripe in an earlier generation. The current collector
+        // already includes it in expectedCount, so buffering it here would make the collector wait
+        // for an observation that is only replayed after that same wait finishes.
+        return false;
+      }
       while (generation.size >= maxBufferSize && generation.active) {
         try {
           bufferSpaceAvailable.await();
@@ -179,8 +191,10 @@ <T extends DataPointSnapshot> T run(
       try {
         activeGeneration = generation;
         long total = 0;
-        for (AtomicLong counter : stripedObservationCounts) {
-          total += counter.getAndAdd(BUFFER_ACTIVE_BIT);
+        for (int i = 0; i < stripedObservationCounts.length; i++) {
+          long count = stripedObservationCounts[i].getAndAdd(BUFFER_ACTIVE_BIT);
+          generationStartCounts[i] = count;
+          total += count;
         }
         expectedCount = total - observationCountOffset;
       } finally {

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

@@ -8,6 +8,9 @@
 import java.util.ArrayList;
 import java.util.List;
 import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicLong;
@@ -180,83 +183,123 @@ void interruptedAppenderLeavesBoundedBufferWait() throws InterruptedException {
   }
 
   @Test
-  void lateAppenderCannotBeAddedToTheNextGeneration() throws InterruptedException {
-    CountDownLatch firstRunStarted = new CountDownLatch(1);
-    CountDownLatch firstRunMayFinish = new CountDownLatch(1);
-    CountDownLatch stalled = new CountDownLatch(1);
-    CountDownLatch release = new CountDownLatch(1);
+  void lateAppenderCountedByNextGenerationMustNotBeBufferedAgain() throws Exception {
+    assertLateAppenderHandoff(false);
+  }
+
+  @Test
+  void lateAppenderHandoffUsesAbsoluteStripeCountsAfterReset() throws Exception {
+    assertLateAppenderHandoff(true);
+  }
+
+  private static void assertLateAppenderHandoff(boolean reset) throws Exception {
+    CountDownLatch firstSnapshotStarted = new CountDownLatch(1);
+    CountDownLatch finishFirstSnapshot = new CountDownLatch(1);
+    CountDownLatch observationCounted = new CountDownLatch(1);
+    CountDownLatch readGeneration = new CountDownLatch(1);
     CountDownLatch secondRunStarted = new CountDownLatch(1);
-    AtomicBoolean appended = new AtomicBoolean();
     AtomicLong completedObservations = new AtomicLong();
+    AtomicLong secondExpectedCount = new AtomicLong();
+    AtomicBoolean pauseFirstAppender = new AtomicBoolean(true);
     Buffer buffer =
         new Buffer(
-            TimeUnit.SECONDS.toNanos(1),
+            TimeUnit.SECONDS.toNanos(5),
             16,
             () -> {
-              stalled.countDown();
-              try {
-                release.await();
-              } catch (InterruptedException e) {
-                Thread.currentThread().interrupt();
+              if (pauseFirstAppender.compareAndSet(true, false)) {
+                observationCounted.countDown();
+                awaitLatch(readGeneration);
               }
             });
-    Thread firstRun =
-        new Thread(
-            () ->
-                buffer.run(
-                    ignored -> {
-                      firstRunStarted.countDown();
-                      return firstRunMayFinish.getCount() == 0;
-                    },
-                    () -> new CounterSnapshot.CounterDataPointSnapshot(0, Labels.EMPTY, null, 0),
-                    ignored -> {}),
-            "buffer-first-runner");
-    firstRun.setDaemon(true);
-    firstRun.start();
-    assertThat(firstRunStarted.await(5, TimeUnit.SECONDS)).isTrue();
+    if (reset) {
+      assertThat(buffer.append(1.0)).isFalse();
+      buffer.observeDirect(completedObservations::incrementAndGet);
+    }
+    ExecutorService executor = Executors.newFixedThreadPool(2);
+    try {
+      Future<CounterSnapshot.CounterDataPointSnapshot> firstRun =
+          executor.submit(
+              () ->
+                  buffer.run(
+                      expectedCount -> completedObservations.get() == expectedCount,
+                      () -> {
+                        firstSnapshotStarted.countDown();
+                        awaitLatch(finishFirstSnapshot);
+                        CounterSnapshot.CounterDataPointSnapshot snapshot =
+                            new CounterSnapshot.CounterDataPointSnapshot(
+                                completedObservations.get(), Labels.EMPTY, null, 0);
+                        if (reset) {
+                          completedObservations.set(0);
+                          buffer.reset();
+                        }
+                        return snapshot;
+                      },
+                      ignored -> completedObservations.incrementAndGet()));
+      awaitLatch(firstSnapshotStarted);
 
-    Thread appender =
-        new Thread(
-            () -> {
-              appended.set(buffer.append(1.0));
-              if (!appended.get()) {
-                buffer.observeDirect(
-                    () -> {
-                      completedObservations.incrementAndGet();
-                      return null;
-                    });
-              }
-            },
-            "buffer-late-appender");
-    appender.setDaemon(true);
-    appender.start();
-    assertThat(stalled.await(5, TimeUnit.SECONDS)).isTrue();
+      // Increment while generation A is active, but do not read activeGeneration yet.
+      Future<Boolean> appender =
+          executor.submit(
+              () -> {
+                boolean appended = buffer.append(1.0);
+                if (!appended) {
+                  buffer.observeDirect(completedObservations::incrementAndGet);
+                }
+                return appended;
+              });
+      awaitLatch(observationCounted);
+      finishFirstSnapshot.countDown();
+      assertThat(firstRun.get(10, TimeUnit.SECONDS).getValue()).isEqualTo(reset ? 1 : 0);
 
-    firstRunMayFinish.countDown();
-    firstRun.join(5_000);
-    assertThat(firstRun.isAlive()).isFalse();
+      Future<CounterSnapshot.CounterDataPointSnapshot> secondRun =
+          executor.submit(
+              () ->
+                  buffer.run(
+                      expectedCount -> {
+                        secondExpectedCount.set(expectedCount);
+                        secondRunStarted.countDown();
+                        return completedObservations.get() == expectedCount;
+                      },
+                      () ->
+                          new CounterSnapshot.CounterDataPointSnapshot(
+                              completedObservations.get(), Labels.EMPTY, null, 0),
+                      ignored -> completedObservations.incrementAndGet()));
+      awaitLatch(secondRunStarted);
+      assertThat(secondExpectedCount).hasValue(1);
+      // An observation arriving after B's activation still belongs in B's buffer. It must not
+      // appear in B's snapshot and must be replayed exactly once before the following collection.
+      assertThat(buffer.append(1.0)).isTrue();
 
-    Thread secondRun =
-        new Thread(
-            () ->
-                buffer.run(
-                    expectedCount -> {
-                      secondRunStarted.countDown();
-                      return completedObservations.get() == expectedCount;
-                    },
-                    () -> new CounterSnapshot.CounterDataPointSnapshot(0, Labels.EMPTY, null, 0),
-                    ignored -> {}),
-            "buffer-second-runner");
-    secondRun.setDaemon(true);
-    secondRun.start();
-    assertThat(secondRunStarted.await(5, TimeUnit.SECONDS)).isTrue();
-    release.countDown();
-    appender.join(5_000);
-    secondRun.join(5_000);
+      // B includes the paused observation in expectedCount. Buffering it in B would make B wait
+      // until its own timeout/replay; it must instead complete via the direct observation path.
+      readGeneration.countDown();
+      assertThat(secondRun.get(10, TimeUnit.SECONDS).getValue()).isEqualTo(1);
+      assertThat(appender.get(10, TimeUnit.SECONDS)).isFalse();
+      assertThat(completedObservations).hasValue(2);
+      assertThat(
+              buffer
+                  .run(
+                      expectedCount -> completedObservations.get() == expectedCount,
+                      () ->
+                          new CounterSnapshot.CounterDataPointSnapshot(
+                              completedObservations.get(), Labels.EMPTY, null, 0),
+                      ignored -> completedObservations.incrementAndGet())
+                  .getValue())
+          .isEqualTo(2);
+    } finally {
+      finishFirstSnapshot.countDown();
+      readGeneration.countDown();
+      executor.shutdownNow();
+      assertThat(executor.awaitTermination(10, TimeUnit.SECONDS)).isTrue();
+    }
+  }
 
-    assertThat(appender.isAlive()).isFalse();
-    assertThat(secondRun.isAlive()).isFalse();
-    assertThat(appended).isFalse();
-    assertThat(completedObservations).hasValue(1);
+  private static void awaitLatch(CountDownLatch latch) {
+    try {
+      assertThat(latch.await(10, TimeUnit.SECONDS)).isTrue();
+    } catch (InterruptedException e) {
+      Thread.currentThread().interrupt();
+      throw new RuntimeException(e);
+    }
   }
 }