2 files changed, 27 insertions, 13 deletions

diff --git a/absl/profiling/internal/sample_recorder.h b/absl/profiling/internal/sample_recorder.h
index fc269bdd..5f65983b 100644
--- a/absl/profiling/internal/sample_recorder.h
+++ b/absl/profiling/internal/sample_recorder.h
@@ -46,6 +46,7 @@ struct Sample {
   absl::Mutex init_mu;
   T* next = nullptr;
   T* dead ABSL_GUARDED_BY(init_mu) = nullptr;
+  int64_t weight;  // How many sampling events were required to sample this one.
 };
 
 // Holds samples and their associated stack traces with a soft limit of
diff --git a/absl/profiling/internal/sample_recorder_test.cc b/absl/profiling/internal/sample_recorder_test.cc
index ec6e0fa2..3373329a 100644
--- a/absl/profiling/internal/sample_recorder_test.cc
+++ b/absl/profiling/internal/sample_recorder_test.cc
@@ -36,7 +36,7 @@ using ::testing::UnorderedElementsAre;
 
 struct Info : public Sample<Info> {
  public:
-  void PrepareForSampling() {}
+  void PrepareForSampling(int64_t w) { weight = w; }
   std::atomic<size_t> size;
   absl::Time create_time;
 };
@@ -49,8 +49,14 @@ std::vector<size_t> GetSizes(SampleRecorder<Info>* s) {
   return res;
 }
 
-Info* Register(SampleRecorder<Info>* s, size_t size) {
-  auto* info = s->Register();
+std::vector<int64_t> GetWeights(SampleRecorder<Info>* s) {
+  std::vector<int64_t> res;
+  s->Iterate([&](const Info& info) { res.push_back(info.weight); });
+  return res;
+}
+
+Info* Register(SampleRecorder<Info>* s, int64_t weight, size_t size) {
+  auto* info = s->Register(weight);
   assert(info != nullptr);
   info->size.store(size);
   return info;
@@ -58,13 +64,15 @@ Info* Register(SampleRecorder<Info>* s, size_t size) {
 
 TEST(SampleRecorderTest, Registration) {
   SampleRecorder<Info> sampler;
-  auto* info1 = Register(&sampler, 1);
+  auto* info1 = Register(&sampler, 31, 1);
   EXPECT_THAT(GetSizes(&sampler), UnorderedElementsAre(1));
+  EXPECT_THAT(GetWeights(&sampler), UnorderedElementsAre(31));
 
-  auto* info2 = Register(&sampler, 2);
+  auto* info2 = Register(&sampler, 32, 2);
   EXPECT_THAT(GetSizes(&sampler), UnorderedElementsAre(1, 2));
   info1->size.store(3);
   EXPECT_THAT(GetSizes(&sampler), UnorderedElementsAre(3, 2));
+  EXPECT_THAT(GetWeights(&sampler), UnorderedElementsAre(31, 32));
 
   sampler.Unregister(info1);
   sampler.Unregister(info2);
@@ -74,18 +82,22 @@ TEST(SampleRecorderTest, Unregistration) {
   SampleRecorder<Info> sampler;
   std::vector<Info*> infos;
   for (size_t i = 0; i < 3; ++i) {
-    infos.push_back(Register(&sampler, i));
+    infos.push_back(Register(&sampler, 33 + i, i));
   }
   EXPECT_THAT(GetSizes(&sampler), UnorderedElementsAre(0, 1, 2));
+  EXPECT_THAT(GetWeights(&sampler), UnorderedElementsAre(33, 34, 35));
 
   sampler.Unregister(infos[1]);
   EXPECT_THAT(GetSizes(&sampler), UnorderedElementsAre(0, 2));
+  EXPECT_THAT(GetWeights(&sampler), UnorderedElementsAre(33, 35));
 
-  infos.push_back(Register(&sampler, 3));
-  infos.push_back(Register(&sampler, 4));
+  infos.push_back(Register(&sampler, 36, 3));
+  infos.push_back(Register(&sampler, 37, 4));
   EXPECT_THAT(GetSizes(&sampler), UnorderedElementsAre(0, 2, 3, 4));
+  EXPECT_THAT(GetWeights(&sampler), UnorderedElementsAre(33, 35, 36, 37));
   sampler.Unregister(infos[3]);
   EXPECT_THAT(GetSizes(&sampler), UnorderedElementsAre(0, 2, 4));
+  EXPECT_THAT(GetWeights(&sampler), UnorderedElementsAre(33, 35, 37));
 
   sampler.Unregister(infos[0]);
   sampler.Unregister(infos[2]);
@@ -99,18 +111,18 @@ TEST(SampleRecorderTest, MultiThreaded) {
   ThreadPool pool(10);
 
   for (int i = 0; i < 10; ++i) {
-    pool.Schedule([&sampler, &stop]() {
+    pool.Schedule([&sampler, &stop, i]() {
       std::random_device rd;
       std::mt19937 gen(rd());
 
       std::vector<Info*> infoz;
       while (!stop.HasBeenNotified()) {
         if (infoz.empty()) {
-          infoz.push_back(sampler.Register());
+          infoz.push_back(sampler.Register(i));
         }
         switch (std::uniform_int_distribution<>(0, 2)(gen)) {
           case 0: {
-            infoz.push_back(sampler.Register());
+            infoz.push_back(sampler.Register(i));
             break;
           }
           case 1: {
@@ -119,6 +131,7 @@ TEST(SampleRecorderTest, MultiThreaded) {
             Info* info = infoz[p];
             infoz[p] = infoz.back();
             infoz.pop_back();
+            EXPECT_EQ(info->weight, i);
             sampler.Unregister(info);
             break;
           }
@@ -143,8 +156,8 @@ TEST(SampleRecorderTest, MultiThreaded) {
 TEST(SampleRecorderTest, Callback) {
   SampleRecorder<Info> sampler;
 
-  auto* info1 = Register(&sampler, 1);
-  auto* info2 = Register(&sampler, 2);
+  auto* info1 = Register(&sampler, 39, 1);
+  auto* info2 = Register(&sampler, 40, 2);
 
   static const Info* expected;