fix: bound striped accumulator snapshot latency

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

Author: zeitlinger

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

@@ -2,7 +2,6 @@
 
 import java.util.Set;
 import java.util.concurrent.ConcurrentHashMap;
-import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicLong;
 
@@ -22,9 +21,6 @@
 final class ClassicOnlyAccumulator {
 
   private static final long NOT_WRITING = -1;
-  // A stalled recorder must not make a scrape wait indefinitely. The next snapshot will retry
-  // the buffer after the recorder has left its epoch.
-  private static final long SNAPSHOT_WAIT_NANOS = TimeUnit.MILLISECONDS.toNanos(1);
 
   private final int bucketCount;
   private final AtomicLong epoch = new AtomicLong();
@@ -78,14 +74,17 @@ void observe(int bucket, double value) {
     }
   }
 
-  @SuppressWarnings({"ModifyCollectionInEnhancedForLoop", "ThreadPriorityCheck"})
+  @SuppressWarnings("ModifyCollectionInEnhancedForLoop")
   synchronized Snapshot snapshot() {
+    // A snapshot is intentionally allowed to be stale for a cell whose recorder is paused. Do not
+    // wait here: this keeps collect(), getCount(), and getSum() bounded by the registered-cell and
+    // bucket counts, independent of writer stalls, while a subsequent snapshot includes the
+    // delayed observation after the recorder publishes NOT_WRITING.
     long inactiveEpoch = epoch.getAndIncrement();
     int inactiveBuffer = (int) (inactiveEpoch & 1);
-    long waitDeadline = System.nanoTime() + SNAPSHOT_WAIT_NANOS;
 
     for (Cell cell : cells) {
-      if (!awaitInactiveBuffer(cell, inactiveBuffer, waitDeadline)) {
+      if (!canDrainInactiveBuffer(cell, inactiveBuffer)) {
         // The writer may be paused indefinitely. Leave this buffer untouched; a later snapshot
         // will collect it after the writer has published NOT_WRITING.
         continue;
@@ -113,21 +112,12 @@ && isEmpty(cell.buffers[1])
     return new Snapshot(collectedBuckets.clone(), collectedCount, collectedSum);
   }
 
-  @SuppressWarnings("ThreadPriorityCheck")
-  private boolean awaitInactiveBuffer(Cell cell, int inactiveBuffer, long waitDeadline) {
-    while (true) {
-      long writingEpoch = cell.writingEpoch;
-      // An old writer can still be in the same parity after a pair of epoch flips. It is not
-      // enough to compare with inactiveEpoch: draining while that writer is active would race
-      // with its plain bucket writes.
-      if (writingEpoch == NOT_WRITING || (writingEpoch & 1) != inactiveBuffer) {
-        return true;
-      }
-      if (System.nanoTime() >= waitDeadline) {
-        return false;
-      }
-      Thread.yield();
-    }
+  private static boolean canDrainInactiveBuffer(Cell cell, int inactiveBuffer) {
+    long writingEpoch = cell.writingEpoch;
+    // An old writer can still be in the same parity after a pair of epoch flips. It is not enough
+    // to compare with the current epoch: draining while that writer is active would race with its
+    // plain bucket writes.
+    return writingEpoch == NOT_WRITING || (writingEpoch & 1) != inactiveBuffer;
   }
 
   private static boolean isEmpty(CellBuffer buffer) {

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

@@ -232,6 +232,8 @@ private DataPoint() {
     @Override
     public double getSum() {
       if (classicOnlyAccumulator != null) {
+        // A paused recorder may make this exact value temporarily stale. The next call retries
+        // its buffer rather than blocking the scrape indefinitely.
         return classicOnlyAccumulator.snapshot().sum;
       }
       return sum.sum();
@@ -240,6 +242,8 @@ public double getSum() {
     @Override
     public long getCount() {
       if (classicOnlyAccumulator != null) {
+        // A paused recorder may make this exact value temporarily stale. The next call retries
+        // its buffer rather than blocking the scrape indefinitely.
         return classicOnlyAccumulator.snapshot().count;
       }
       return count.sum();
@@ -311,6 +315,8 @@ private void doObserve(double value, boolean fromBuffer) {
     private HistogramSnapshot.HistogramDataPointSnapshot collect(Labels labels) {
       Exemplars exemplars = exemplarSampler != null ? exemplarSampler.collect() : Exemplars.EMPTY;
       if (classicOnlyAccumulator != null) {
+        // collect() is intentionally allowed to return a stale snapshot for a paused recorder;
+        // the following collection retries its buffer without an unbounded wait.
         ClassicOnlyAccumulator.Snapshot snapshot = classicOnlyAccumulator.snapshot();
         return new HistogramSnapshot.HistogramDataPointSnapshot(
             ClassicHistogramBuckets.of(classicUpperBounds, snapshot.buckets),

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

@@ -37,6 +37,66 @@ void stalledWriterDoesNotBlockSnapshotAndIsCollectedLater() throws Exception {
     assertThat(collected.sum).isEqualTo(1.0);
   }
 
+  @Test
+  void stalledCellDoesNotPreventHealthyCellsFromBeingCollected() throws Exception {
+    ClassicOnlyAccumulator accumulator = new ClassicOnlyAccumulator(1);
+    accumulator.observe(0, 1.0);
+    for (int i = 0; i < 4; i++) {
+      Thread recorder = new Thread(() -> accumulator.observe(0, 1.0));
+      recorder.start();
+      recorder.join();
+    }
+
+    Object stalledCell = onlyCell(accumulator);
+    Field writingEpoch = stalledCell.getClass().getDeclaredField("writingEpoch");
+    writingEpoch.setAccessible(true);
+    writingEpoch.setLong(stalledCell, 0);
+
+    ClassicOnlyAccumulator.Snapshot first = accumulator.snapshot();
+    // The four healthy cells are drained even though the first cell consumes its own wait budget.
+    assertThat(first.count).isEqualTo(4);
+
+    writingEpoch.setLong(stalledCell, -1);
+    accumulator.snapshot();
+    ClassicOnlyAccumulator.Snapshot finalSnapshot = accumulator.snapshot();
+    assertThat(finalSnapshot.count).isEqualTo(5);
+  }
+
+  @Test
+  void activeWriterCanResumeAfterAStaleSnapshot() throws Exception {
+    ClassicOnlyAccumulator accumulator = new ClassicOnlyAccumulator(1);
+    accumulator.observe(0, 1.0);
+    Object cell = onlyCell(accumulator);
+    Field writingEpoch = cell.getClass().getDeclaredField("writingEpoch");
+    writingEpoch.setAccessible(true);
+    CountDownLatch entered = new CountDownLatch(1);
+    CountDownLatch release = new CountDownLatch(1);
+    Thread writer =
+        new Thread(
+            () -> {
+              try {
+                writingEpoch.setLong(cell, 0);
+                entered.countDown();
+                release.await();
+                writingEpoch.setLong(cell, -1);
+              } catch (InterruptedException e) {
+                Thread.currentThread().interrupt();
+              } catch (IllegalAccessException e) {
+                throw new AssertionError(e);
+              }
+            });
+    writer.start();
+    entered.await();
+
+    ClassicOnlyAccumulator.Snapshot stale = accumulator.snapshot();
+    assertThat(stale.count).isZero();
+    release.countDown();
+    writer.join();
+    accumulator.snapshot();
+    ClassicOnlyAccumulator.Snapshot resumed = accumulator.snapshot();
+    assertThat(resumed.count).isEqualTo(1);
+  }
+
   @Test
   void emptyCellsAreReclaimedAndCanBeReused() throws Exception {
     ClassicOnlyAccumulator accumulator = new ClassicOnlyAccumulator(1);