Fix duplicate and biased thread picks in the wall-clock reservoir sampler

ddprof-lib/src/main/cpp/reservoirSampler.h

@@ -1,5 +1,5 @@
 /*
- * Copyright 2024 Datadog
+ * Copyright 2024, 2026 Datadog
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -18,7 +18,6 @@
 #define RESERVOIR_SAMPLER_H
 
 #include "xorshift.h"
-#include <cassert>
 #include <math.h>
 #include <vector>
 
@@ -69,16 +68,35 @@ class ReservoirSampler {
         for (int i = 0; i < _size && i < (int)input.size(); i++) {
             _reservoir.push_back(input[i]);
         }
+        // Algorithm L (Li, 1994). `target` is the 0-based index of the next
+        // input item to place, so the first candidate is input[_size] plus a
+        // skip, and every later one is one past the item just placed plus a
+        // skip. Without that +1 a zero skip places the same item again,
+        // duplicating it in the sample.
+        const int n = (int)input.size();
         double weight = exp(log(nextUniform()) / _size);
-        int target = _size + (int) (log(nextUniform()) / log(1 - weight));
-        assert(target >= 0);
-        while (target < (int)input.size()) {
+        int target = advance(_size, n, weight);
+        while (target < n) {
             _reservoir[xorshift::boundedIndex(xorshift::next(_rng), (u32)_size)] = input[target];
             weight *= exp(log(nextUniform()) / _size);
-            target += (int) (log(nextUniform()) / log(1 - weight));
+            target = advance(target + 1, n, weight);
         }
         return _reservoir;
     }
+
+private:
+    // Returns `from` plus a geometric skip, or `n` when the skip runs past the
+    // input. The skip is compared as a double before the int conversion: as
+    // `weight` shrinks it can exceed INT_MAX, and once `1 - weight` rounds to
+    // 1.0 the quotient is -inf (log(u) < 0 over log(1.0) == +0.0). Either way
+    // the stream is exhausted, and converting such a value would be undefined.
+    int advance(int from, int n, double weight) {
+        double skip = floor(log(nextUniform()) / log(1 - weight));
+        if (!(skip >= 0.0 && skip < (double)(n - from))) {
+            return n;
+        }
+        return from + (int)skip;
+    }
 };
 
 #endif //RESERVOIR_SAMPLER_H

ddprof-lib/src/test/cpp/xorshift_ut.cpp

@@ -407,3 +407,55 @@ TEST(ReservoirSamplerTest, ReuseAcrossCallsDoesNotAccumulate) {
         ASSERT_EQ(4u, out.size()) << "call " << call;
     }
 }
+
+// Algorithm L steps one past the item it just placed before adding the skip.
+// Without that step a zero skip places the same input item a second time, so a
+// sample of distinct threads can name one thread twice and miss another.
+TEST(ReservoirSamplerTest, SampleOfDistinctInputHasNoDuplicates) {
+    const int cases[][2] = {{4, 1}, {10, 3}, {17, 16}, {100, 16}, {2048, 16}, {16, 16}};
+    for (const auto& c : cases) {
+        const int n = c[0];
+        const int k = c[1];
+        std::vector<int> input;
+        for (int i = 0; i < n; i++) {
+            input.push_back(i);
+        }
+        for (int run = 0; run < 2000; run++) {
+            ReservoirSampler<int> sampler(k, (u64)run);
+            std::vector<int>& out = sampler.sample(input);
+            ASSERT_EQ((size_t)k, out.size()) << "n=" << n << " k=" << k;
+            std::set<int> distinct(out.begin(), out.end());
+            ASSERT_EQ((size_t)k, distinct.size())
+                << "duplicate selection for n=" << n << " k=" << k << " run=" << run;
+        }
+    }
+}
+
+// Each input item must land in the sample with probability k/n. The bound is
+// +/-5% of the expectation, which is over ten standard deviations wide for
+// these draw counts; the missing step put item k at roughly 1.25x (n=4, k=1)
+// and 1.8x (n=10, k=3) its fair share.
+TEST(ReservoirSamplerTest, InclusionFrequencyIsUniform) {
+    const int cases[][2] = {{4, 1}, {10, 3}, {40, 16}};
+    const int draws = 100000;
+    for (const auto& c : cases) {
+        const int n = c[0];
+        const int k = c[1];
+        std::vector<int> input;
+        for (int i = 0; i < n; i++) {
+            input.push_back(i);
+        }
+        std::vector<long> hits(n, 0);
+        for (int run = 0; run < draws; run++) {
+            ReservoirSampler<int> sampler(k, (u64)run);
+            for (int value : sampler.sample(input)) {
+                hits[value]++;
+            }
+        }
+        const double expected = (double)draws * k / n;
+        for (int i = 0; i < n; i++) {
+            EXPECT_NEAR((double)hits[i], expected, expected * 0.05)
+                << "item " << i << " for n=" << n << " k=" << k;
+        }
+    }
+}