1 files changed, 22 insertions, 19 deletions

diff --git a/absl/container/internal/hashtablez_sampler_test.cc b/absl/container/internal/hashtablez_sampler_test.cc
index a7307a20..ac6ed9b1 100644
--- a/absl/container/internal/hashtablez_sampler_test.cc
+++ b/absl/container/internal/hashtablez_sampler_test.cc
@@ -70,7 +70,8 @@ std::vector<size_t> GetSizes(HashtablezSampler* s) {
 }
 
 HashtablezInfo* Register(HashtablezSampler* s, size_t size) {
-  auto* info = s->Register();
+  const size_t test_element_size = 17;
+  auto* info = s->Register(test_element_size);
   assert(info != nullptr);
   info->size.store(size);
   return info;
@@ -81,9 +82,8 @@ TEST(HashtablezInfoTest, PrepareForSampling) {
   const size_t test_element_size = 17;
   HashtablezInfo info;
   absl::MutexLock l(&info.init_mu);
-  info.PrepareForSampling();
+  info.PrepareForSampling(test_element_size);
 
-  info.inline_element_size.store(test_element_size, std::memory_order_relaxed);
   EXPECT_EQ(info.capacity.load(), 0);
   EXPECT_EQ(info.size.load(), 0);
   EXPECT_EQ(info.num_erases.load(), 0);
@@ -95,7 +95,7 @@ TEST(HashtablezInfoTest, PrepareForSampling) {
   EXPECT_EQ(info.hashes_bitwise_xor.load(), 0);
   EXPECT_EQ(info.max_reserve.load(), 0);
   EXPECT_GE(info.create_time, test_start);
-  EXPECT_EQ(info.inline_element_size.load(), test_element_size);
+  EXPECT_EQ(info.inline_element_size, test_element_size);
 
   info.capacity.store(1, std::memory_order_relaxed);
   info.size.store(1, std::memory_order_relaxed);
@@ -108,7 +108,7 @@ TEST(HashtablezInfoTest, PrepareForSampling) {
   info.max_reserve.store(1, std::memory_order_relaxed);
   info.create_time = test_start - absl::Hours(20);
 
-  info.PrepareForSampling();
+  info.PrepareForSampling(test_element_size);
   EXPECT_EQ(info.capacity.load(), 0);
   EXPECT_EQ(info.size.load(), 0);
   EXPECT_EQ(info.num_erases.load(), 0);
@@ -119,14 +119,15 @@ TEST(HashtablezInfoTest, PrepareForSampling) {
   EXPECT_EQ(info.hashes_bitwise_and.load(), ~size_t{});
   EXPECT_EQ(info.hashes_bitwise_xor.load(), 0);
   EXPECT_EQ(info.max_reserve.load(), 0);
-  EXPECT_EQ(info.inline_element_size.load(), test_element_size);
+  EXPECT_EQ(info.inline_element_size, test_element_size);
   EXPECT_GE(info.create_time, test_start);
 }
 
 TEST(HashtablezInfoTest, RecordStorageChanged) {
   HashtablezInfo info;
   absl::MutexLock l(&info.init_mu);
-  info.PrepareForSampling();
+  const size_t test_element_size = 19;
+  info.PrepareForSampling(test_element_size);
   RecordStorageChangedSlow(&info, 17, 47);
   EXPECT_EQ(info.size.load(), 17);
   EXPECT_EQ(info.capacity.load(), 47);
@@ -138,7 +139,8 @@ TEST(HashtablezInfoTest, RecordStorageChanged) {
 TEST(HashtablezInfoTest, RecordInsert) {
   HashtablezInfo info;
   absl::MutexLock l(&info.init_mu);
-  info.PrepareForSampling();
+  const size_t test_element_size = 23;
+  info.PrepareForSampling(test_element_size);
   EXPECT_EQ(info.max_probe_length.load(), 0);
   RecordInsertSlow(&info, 0x0000FF00, 6 * kProbeLength);
   EXPECT_EQ(info.max_probe_length.load(), 6);
@@ -161,8 +163,7 @@ TEST(HashtablezInfoTest, RecordErase) {
   const size_t test_element_size = 29;
   HashtablezInfo info;
   absl::MutexLock l(&info.init_mu);
-  info.PrepareForSampling();
-  info.inline_element_size.store(test_element_size);
+  info.PrepareForSampling(test_element_size);
   EXPECT_EQ(info.num_erases.load(), 0);
   EXPECT_EQ(info.size.load(), 0);
   RecordInsertSlow(&info, 0x0000FF00, 6 * kProbeLength);
@@ -170,15 +171,14 @@ TEST(HashtablezInfoTest, RecordErase) {
   RecordEraseSlow(&info);
   EXPECT_EQ(info.size.load(), 0);
   EXPECT_EQ(info.num_erases.load(), 1);
-  EXPECT_EQ(info.inline_element_size.load(), test_element_size);
+  EXPECT_EQ(info.inline_element_size, test_element_size);
 }
 
 TEST(HashtablezInfoTest, RecordRehash) {
   const size_t test_element_size = 31;
   HashtablezInfo info;
   absl::MutexLock l(&info.init_mu);
-  info.PrepareForSampling();
-  info.inline_element_size.store(test_element_size);
+  info.PrepareForSampling(test_element_size);
   RecordInsertSlow(&info, 0x1, 0);
   RecordInsertSlow(&info, 0x2, kProbeLength);
   RecordInsertSlow(&info, 0x4, kProbeLength);
@@ -197,13 +197,14 @@ TEST(HashtablezInfoTest, RecordRehash) {
   EXPECT_EQ(info.total_probe_length.load(), 3);
   EXPECT_EQ(info.num_erases.load(), 0);
   EXPECT_EQ(info.num_rehashes.load(), 1);
-  EXPECT_EQ(info.inline_element_size.load(), test_element_size);
+  EXPECT_EQ(info.inline_element_size, test_element_size);
 }
 
 TEST(HashtablezInfoTest, RecordReservation) {
   HashtablezInfo info;
   absl::MutexLock l(&info.init_mu);
-  info.PrepareForSampling();
+  const size_t test_element_size = 33;
+  info.PrepareForSampling(test_element_size);
   RecordReservationSlow(&info, 3);
   EXPECT_EQ(info.max_reserve.load(), 3);
 
@@ -266,7 +267,8 @@ TEST(HashtablezSamplerTest, Sample) {
 
 TEST(HashtablezSamplerTest, Handle) {
   auto& sampler = GlobalHashtablezSampler();
-  HashtablezInfoHandle h(sampler.Register());
+  const size_t test_element_size = 39;
+  HashtablezInfoHandle h(sampler.Register(test_element_size));
   auto* info = HashtablezInfoHandlePeer::GetInfo(&h);
   info->hashes_bitwise_and.store(0x12345678, std::memory_order_relaxed);
 
@@ -338,18 +340,19 @@ TEST(HashtablezSamplerTest, MultiThreaded) {
   ThreadPool pool(10);
 
   for (int i = 0; i < 10; ++i) {
-    pool.Schedule([&sampler, &stop]() {
+    const size_t elt_size = 10 + i % 2;
+    pool.Schedule([&sampler, &stop, elt_size]() {
       std::random_device rd;
       std::mt19937 gen(rd());
 
       std::vector<HashtablezInfo*> infoz;
       while (!stop.HasBeenNotified()) {
         if (infoz.empty()) {
-          infoz.push_back(sampler.Register());
+          infoz.push_back(sampler.Register(elt_size));
         }
         switch (std::uniform_int_distribution<>(0, 2)(gen)) {
           case 0: {
-            infoz.push_back(sampler.Register());
+            infoz.push_back(sampler.Register(elt_size));
             break;
           }
           case 1: {