1 files changed, 21 insertions, 6 deletions

diff --git a/absl/profiling/internal/sample_recorder.h b/absl/profiling/internal/sample_recorder.h
index 5e04a9cd..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
@@ -59,7 +60,8 @@ class SampleRecorder {
   ~SampleRecorder();
 
   // Registers for sampling.  Returns an opaque registration info.
-  T* Register();
+  template <typename... Targs>
+  T* Register(Targs&&... args);
 
   // Unregisters the sample.
   void Unregister(T* sample);
@@ -75,12 +77,14 @@ class SampleRecorder {
   // samples that have been dropped.
   int64_t Iterate(const std::function<void(const T& stack)>& f);
 
+  int32_t GetMaxSamples() const;
   void SetMaxSamples(int32_t max);
 
  private:
   void PushNew(T* sample);
   void PushDead(T* sample);
-  T* PopDead();
+  template <typename... Targs>
+  T* PopDead(Targs... args);
 
   std::atomic<size_t> dropped_samples_;
   std::atomic<size_t> size_estimate_;
@@ -162,7 +166,8 @@ void SampleRecorder<T>::PushDead(T* sample) {
 }
 
 template <typename T>
-T* SampleRecorder<T>::PopDead() {
+template <typename... Targs>
+T* SampleRecorder<T>::PopDead(Targs... args) {
   absl::MutexLock graveyard_lock(&graveyard_.init_mu);
 
   // The list is circular, so eventually it collapses down to
@@ -174,12 +179,13 @@ T* SampleRecorder<T>::PopDead() {
   absl::MutexLock sample_lock(&sample->init_mu);
   graveyard_.dead = sample->dead;
   sample->dead = nullptr;
-  sample->PrepareForSampling();
+  sample->PrepareForSampling(std::forward<Targs>(args)...);
   return sample;
 }
 
 template <typename T>
-T* SampleRecorder<T>::Register() {
+template <typename... Targs>
+T* SampleRecorder<T>::Register(Targs&&... args) {
   int64_t size = size_estimate_.fetch_add(1, std::memory_order_relaxed);
   if (size > max_samples_.load(std::memory_order_relaxed)) {
     size_estimate_.fetch_sub(1, std::memory_order_relaxed);
@@ -187,10 +193,14 @@ T* SampleRecorder<T>::Register() {
     return nullptr;
   }
 
-  T* sample = PopDead();
+  T* sample = PopDead(args...);
   if (sample == nullptr) {
     // Resurrection failed.  Hire a new warlock.
     sample = new T();
+    {
+      absl::MutexLock sample_lock(&sample->init_mu);
+      sample->PrepareForSampling(std::forward<Targs>(args)...);
+    }
     PushNew(sample);
   }
 
@@ -223,6 +233,11 @@ void SampleRecorder<T>::SetMaxSamples(int32_t max) {
   max_samples_.store(max, std::memory_order_release);
 }
 
+template <typename T>
+int32_t SampleRecorder<T>::GetMaxSamples() const {
+  return max_samples_.load(std::memory_order_acquire);
+}
+
 }  // namespace profiling_internal
 ABSL_NAMESPACE_END
 }  // namespace absl