fix: harden Buffer collection and stripe indexing

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

@@ -2,6 +2,7 @@
 
 import io.prometheus.metrics.model.snapshots.DataPointSnapshot;
 import java.util.Arrays;
+import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicLong;
 import java.util.concurrent.locks.Condition;
 import java.util.concurrent.locks.ReentrantLock;
@@ -18,6 +19,7 @@
 class Buffer {
 
   private static final long bufferActiveBit = 1L << 63;
+  private static final long DEFAULT_MAX_SPIN_WAIT_NANOS = TimeUnit.SECONDS.toNanos(1);
   // Tracking observation counts requires an AtomicLong for coordination between recording and
   // collecting. AtomicLong does much worse under contention than the LongAdder instances used
   // elsewhere to hold aggregated state. To improve, we stripe the AtomicLong into N instances,
@@ -34,16 +36,22 @@ class Buffer {
   ReentrantLock appendLock = new ReentrantLock();
   ReentrantLock runLock = new ReentrantLock();
   Condition bufferFilled = appendLock.newCondition();
+  private final long maxSpinWaitNanos;
 
   Buffer() {
+    this(DEFAULT_MAX_SPIN_WAIT_NANOS);
+  }
+
+  Buffer(long maxSpinWaitNanos) {
+    this.maxSpinWaitNanos = maxSpinWaitNanos;
     stripedObservationCounts = new AtomicLong[Runtime.getRuntime().availableProcessors()];
     for (int i = 0; i < stripedObservationCounts.length; i++) {
       stripedObservationCounts[i] = new AtomicLong(0);
     }
   }
 
   boolean append(double value) {
-    int index = Math.abs((int) Thread.currentThread().getId()) % stripedObservationCounts.length;
+    int index = stripeIndex(Thread.currentThread().getId(), stripedObservationCounts.length);
     AtomicLong observationCountForThread = stripedObservationCounts[index];
     long count = observationCountForThread.incrementAndGet();
     if ((count & bufferActiveBit) == 0) {
@@ -54,6 +62,10 @@ boolean append(double value) {
     }
   }
 
+  static int stripeIndex(long threadId, int stripeCount) {
+    return (int) Math.floorMod(threadId, stripeCount);
+  }
+
   private void doAppend(double amount) {
     appendLock.lock();
     try {
@@ -74,14 +86,15 @@ void reset() {
     reset = true;
   }
 
-  @SuppressWarnings("ThreadPriorityCheck")
+  @SuppressWarnings({"NullAway", "ThreadPriorityCheck"})
   <T extends DataPointSnapshot> T run(
       Function<Long, Boolean> complete,
       Supplier<T> createResult,
       Consumer<Double> observeFunction) {
     double[] buffer;
     int bufferSize;
     T result;
+    boolean timedOut = false;
 
     runLock.lock();
     try {
@@ -91,14 +104,19 @@ <T extends DataPointSnapshot> T run(
         expectedCount += observationCount.getAndAdd(bufferActiveBit);
       }
 
+      long deadline = System.nanoTime() + maxSpinWaitNanos;
       while (!complete.apply(expectedCount)) {
         // Wait until all in-flight threads have added their observations to the histogram /
         // summary.
         // we can't use a condition here, because the other thread doesn't have a lock as it's on
         // the fast path.
+        if (System.nanoTime() - deadline >= 0) {
+          timedOut = true;
+          break;
+        }
         Thread.yield();
       }
-      result = createResult.get();
+      result = timedOut ? null : createResult.get();
 
       // Signal that the buffer is inactive.
       long expectedBufferSize = 0;
@@ -137,6 +155,9 @@ <T extends DataPointSnapshot> T run(
     for (int i = 0; i < bufferSize; i++) {
       observeFunction.accept(buffer[i]);
     }
+    if (timedOut) {
+      throw new IllegalStateException("Timed out while waiting for in-flight observations.");
+    }
     return result;
   }
 }

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

@@ -0,0 +1,73 @@
+package io.prometheus.metrics.core.metrics;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatExceptionOfType;
+
+import io.prometheus.metrics.model.snapshots.CounterSnapshot;
+import io.prometheus.metrics.model.snapshots.Labels;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
+import org.junit.jupiter.api.Test;
+
+class BufferTest {
+
+  @Test
+  void stripeIndexDoesNotOverflowWhenThreadIdNarrowsToIntegerMinValue() {
+    assertThat(Buffer.stripeIndex(2_147_483_648L, 3)).isEqualTo(2);
+    assertThat(Buffer.stripeIndex(2_147_483_648L, 6)).isEqualTo(2);
+    assertThat(Buffer.stripeIndex(2_147_483_648L, 12)).isEqualTo(8);
+  }
+
+  @Test
+  void timeoutDeactivatesBufferAndReplaysBufferedObservations() throws InterruptedException {
+    Buffer buffer = new Buffer(TimeUnit.SECONDS.toNanos(1));
+    CountDownLatch spinWaitStarted = new CountDownLatch(1);
+    AtomicBoolean keepWaiting = new AtomicBoolean(true);
+    List<Double> replayedObservations = new ArrayList<>();
+    AtomicBoolean timedOut = new AtomicBoolean(false);
+
+    Thread runner =
+        new Thread(
+            () -> {
+              try {
+                buffer.run(
+                    expectedCount -> {
+                      spinWaitStarted.countDown();
+                      return !keepWaiting.get();
+                    },
+                    () -> new CounterSnapshot.CounterDataPointSnapshot(0, Labels.EMPTY, null, 0),
+                    replayedObservations::add);
+              } catch (IllegalStateException expected) {
+                timedOut.set(true);
+              }
+            });
+    runner.start();
+
+    assertThat(spinWaitStarted.await(5, TimeUnit.SECONDS)).isTrue();
+    assertThat(buffer.append(1.0)).isTrue();
+    runner.join(5_000);
+
+    assertThat(timedOut).isTrue();
+    assertThat(replayedObservations).containsExactly(1.0);
+    assertThat(buffer.append(2.0)).isFalse();
+  }
+
+  @Test
+  void timeoutDoesNotCreateSnapshot() {
+    Buffer buffer = new Buffer(TimeUnit.MILLISECONDS.toNanos(1));
+
+    assertThatExceptionOfType(IllegalStateException.class)
+        .isThrownBy(
+            () ->
+                buffer.run(
+                    expectedCount -> false,
+                    () -> {
+                      throw new AssertionError("snapshot should not be created");
+                    },
+                    ignored -> {}))
+        .withMessage("Timed out while waiting for in-flight observations.");
+  }
+}