1 files changed, 14 insertions, 6 deletions
diff --git a/absl/container/internal/hashtablez_sampler.cc b/absl/container/internal/hashtablez_sampler.cc
index b21beef0..e17ba14d 100644
--- a/absl/container/internal/hashtablez_sampler.cc
+++ b/absl/container/internal/hashtablez_sampler.cc
@@ -18,13 +18,18 @@
 #include <atomic>
 #include <cassert>
 #include <cmath>
+#include <cstddef>
+#include <cstdint>
 #include <functional>
 #include <limits>
 
 #include "absl/base/attributes.h"
 #include "absl/base/config.h"
+#include "absl/base/internal/per_thread_tls.h"
 #include "absl/base/internal/raw_logging.h"
+#include "absl/base/macros.h"
 #include "absl/base/no_destructor.h"
+#include "absl/base/optimization.h"
 #include "absl/debugging/stacktrace.h"
 #include "absl/memory/memory.h"
 #include "absl/profiling/internal/exponential_biased.h"
@@ -73,7 +78,8 @@ HashtablezInfo::HashtablezInfo() = default;
 HashtablezInfo::~HashtablezInfo() = default;
 
 void HashtablezInfo::PrepareForSampling(int64_t stride,
-                                        size_t inline_element_size_value) {
+                                        size_t inline_element_size_value,
+                                        uint16_t soo_capacity_value) {
   capacity.store(0, std::memory_order_relaxed);
   size.store(0, std::memory_order_relaxed);
   num_erases.store(0, std::memory_order_relaxed);
@@ -93,6 +99,7 @@ void HashtablezInfo::PrepareForSampling(int64_t stride,
   depth = absl::GetStackTrace(stack, HashtablezInfo::kMaxStackDepth,
                               /* skip_count= */ 0);
   inline_element_size = inline_element_size_value;
+  soo_capacity = soo_capacity_value;
 }
 
 static bool ShouldForceSampling() {
@@ -116,12 +123,12 @@ static bool ShouldForceSampling() {
 }
 
 HashtablezInfo* SampleSlow(SamplingState& next_sample,
-                           size_t inline_element_size) {
+                           size_t inline_element_size, uint16_t soo_capacity) {
   if (ABSL_PREDICT_FALSE(ShouldForceSampling())) {
     next_sample.next_sample = 1;
     const int64_t old_stride = exchange(next_sample.sample_stride, 1);
-    HashtablezInfo* result =
-        GlobalHashtablezSampler().Register(old_stride, inline_element_size);
+    HashtablezInfo* result = GlobalHashtablezSampler().Register(
+        old_stride, inline_element_size, soo_capacity);
     return result;
   }
 
@@ -151,10 +158,11 @@ HashtablezInfo* SampleSlow(SamplingState& next_sample,
   // that case.
   if (first) {
     if (ABSL_PREDICT_TRUE(--next_sample.next_sample > 0)) return nullptr;
-    return SampleSlow(next_sample, inline_element_size);
+    return SampleSlow(next_sample, inline_element_size, soo_capacity);
   }
 
-  return GlobalHashtablezSampler().Register(old_stride, inline_element_size);
+  return GlobalHashtablezSampler().Register(old_stride, inline_element_size,
+                                            soo_capacity);
 #endif
 }