Fix hashtablez oversampling for SOO tables on repeated 0->1 element transitions

Currently, when an absl::flat_hash_set or flat_hash_map with SOO (Single Object Optimization) transitions from 0 to 1 element, it queries should_sample_soo() to determine if it should be sampled by hashtablez. If it is NOT sampled, it remains an SOO table. If that element is subsequently erased, the table goes back to 0 elements. The next time it goes from 0 to 1, it queries should_sample_soo() again.

Repeated insert/erase cycles on an SOO table effectively evaluate the sampling logic hundreds of times, making it extremely likely that the table eventually gets sampled and moved to the heap. This causes hashtablez to vastly overstate the memory usage of such maps in production.

This CL fixes this by tracking whether an SOO table has already been evaluated for sampling, so it only evaluates the logic once per table object instance. Since SOO tables do not hash keys, they do not need the per-table generic seed value. This CL utilizes the lowest bit of the unused seed field within HashtableSize (accessed via data_ & 1) to act as a flag soo_has_tried_sampling.

When an SOO table transitions from 0 elements to 1 element, this bit is checked. If it is false, should_sample_soo() is queried and the bit is set to true. set_full_soo() and set_empty_soo() have been modified to preserve the soo_has_tried_sampling bit across erase or clear operations.

PiperOrigin-RevId: 881570273
Change-Id: Ia3b0eff8b3e137b0d162a207220a184923e3a2ca

Author: toddlipcon

absl/container/internal/raw_hash_set.h

@@ -459,26 +459,41 @@ class PerTableSeed {
 
 // The size and also has additionally
 // 1) one bit that stores whether we have infoz.
-// 2) PerTableSeed::kBitCount bits for the seed.
+// 2) PerTableSeed::kBitCount bits for the seed. (For SOO tables, the lowest
+//    bit of the seed is repurposed to track if sampling has been tried).
 class HashtableSize {
  public:
   static constexpr size_t kSizeBitCount = 64 - PerTableSeed::kBitCount - 1;
 
   explicit HashtableSize(uninitialized_tag_t) {}
   explicit HashtableSize(no_seed_empty_tag_t) : data_(0) {}
-  explicit HashtableSize(full_soo_tag_t) : data_(kSizeOneNoMetadata) {}
+  HashtableSize(full_soo_tag_t, bool has_tried_sampling)
+      : data_(kSizeOneNoMetadata |
+              (has_tried_sampling ? kSooHasTriedSamplingMask : 0)) {}
 
   // Returns actual size of the table.
   size_t size() const { return static_cast<size_t>(data_ >> kSizeShift); }
   void increment_size() { data_ += kSizeOneNoMetadata; }
   void increment_size(size_t size) {
-    data_ += static_cast<uint64_t>(size) * kSizeOneNoMetadata;
+    data_ += static_cast<uint64_t>(size) << kSizeShift;
   }
   void decrement_size() { data_ -= kSizeOneNoMetadata; }
   // Returns true if the table is empty.
   bool empty() const { return data_ < kSizeOneNoMetadata; }
-  // Sets the size to zero, but keeps all the metadata bits.
-  void set_size_to_zero_keep_metadata() { data_ = data_ & kMetadataMask; }
+
+  // Returns true if an empty SOO table has already queried should_sample_soo().
+  bool soo_has_tried_sampling() const {
+    return (data_ & kSooHasTriedSamplingMask) != 0;
+  }
+
+  // Records that an empty SOO table has tried sampling.
+  void set_soo_has_tried_sampling() { data_ |= kSooHasTriedSamplingMask; }
+
+  // Sets the size, but keeps all the metadata bits.
+  void set_size_keep_metadata(size_t size) {
+    data_ =
+        (data_ & kMetadataMask) | (static_cast<uint64_t>(size) << kSizeShift);
+  }
 
   PerTableSeed seed() const {
     return PerTableSeed(static_cast<size_t>(data_) & kSeedMask);
@@ -509,13 +524,20 @@ class HashtableSize {
 
  private:
   void set_seed(uint16_t seed) { data_ = (data_ & ~kSeedMask) | seed; }
+  // Bit layout of `data_`:
+  // [63 ... 17] (47 bits) : size
+  // [16]        (1 bit)   : has_infoz
+  // [15 ...  0] (16 bits) : seed
   static constexpr size_t kSizeShift = 64 - kSizeBitCount;
   static constexpr uint64_t kSizeOneNoMetadata = uint64_t{1} << kSizeShift;
   static constexpr uint64_t kMetadataMask = kSizeOneNoMetadata - 1;
   static constexpr uint64_t kSeedMask =
       (uint64_t{1} << PerTableSeed::kBitCount) - 1;
   // The next bit after the seed.
   static constexpr uint64_t kHasInfozMask = kSeedMask + 1;
+  // For SOO tables, the seed is unused, and bit 0 is repurposed to track
+  // whether the table has already queried should_sample_soo().
+  static constexpr uint64_t kSooHasTriedSamplingMask = 1;
   uint64_t data_;
 };
 
@@ -907,8 +929,8 @@ class CommonFields : public CommonFieldsGenerationInfo {
  public:
   explicit CommonFields(soo_tag_t)
       : capacity_(SooCapacity()), size_(no_seed_empty_tag_t{}) {}
-  explicit CommonFields(full_soo_tag_t)
-      : capacity_(SooCapacity()), size_(full_soo_tag_t{}) {}
+  explicit CommonFields(full_soo_tag_t, bool has_tried_sampling)
+      : capacity_(SooCapacity()), size_(full_soo_tag_t{}, has_tried_sampling) {}
   explicit CommonFields(non_soo_tag_t)
       : capacity_(0), size_(no_seed_empty_tag_t{}) {}
   // For use in swapping.
@@ -963,14 +985,14 @@ class CommonFields : public CommonFieldsGenerationInfo {
   // The number of filled slots.
   size_t size() const { return size_.size(); }
   // Sets the size to zero, but keeps hashinfoz bit and seed.
-  void set_size_to_zero() { size_.set_size_to_zero_keep_metadata(); }
+  void set_size_to_zero() { size_.set_size_keep_metadata(0); }
   void set_empty_soo() {
     AssertInSooMode();
-    size_ = HashtableSize(no_seed_empty_tag_t{});
+    size_.set_size_keep_metadata(0);
   }
   void set_full_soo() {
     AssertInSooMode();
-    size_ = HashtableSize(full_soo_tag_t{});
+    size_.set_size_keep_metadata(1);
   }
   void increment_size() {
     ABSL_SWISSTABLE_ASSERT(size() < capacity());
@@ -985,6 +1007,8 @@ class CommonFields : public CommonFieldsGenerationInfo {
     size_.decrement_size();
   }
   bool empty() const { return size_.empty(); }
+  void set_soo_has_tried_sampling() { size_.set_soo_has_tried_sampling(); }
+  bool soo_has_tried_sampling() const { return size_.soo_has_tried_sampling(); }
 
   // The seed used for the hash function.
   PerTableSeed seed() const { return size_.seed(); }
@@ -2271,7 +2295,8 @@ class raw_hash_set {
          // Note: we avoid using exchange for better generated code.
         settings_(PolicyTraits::transfer_uses_memcpy() || !that.is_full_soo()
                       ? std::move(that.common())
-                      : CommonFields{full_soo_tag_t{}},
+                      : CommonFields{full_soo_tag_t{},
+                                     that.common().soo_has_tried_sampling()},
                   that.hash_ref(), that.eq_ref(), that.char_alloc_ref()) {
     if (!PolicyTraits::transfer_uses_memcpy() && that.is_full_soo()) {
       transfer(soo_slot(), that.soo_slot());
@@ -3011,12 +3036,23 @@ class raw_hash_set {
     }
   }
 
-  // Returns true if the table needs to be sampled.
+  // Returns true if the table needs to be sampled. This keeps track of whether
+  // sampling has already been evaluated and ensures that it can only return
+  // true on its first evaluation. All subsequent calls will return false.
+  //
   // This should be called on insertion into an empty SOO table and in copy
   // construction when the size can fit in SOO capacity.
-  bool should_sample_soo() const {
+  bool should_sample_soo() {
     ABSL_SWISSTABLE_ASSERT(is_soo());
     if (!ShouldSampleHashtablezInfoForAlloc<CharAlloc>()) return false;
+    if (common().soo_has_tried_sampling()) {
+      // Already evaluated sampling on this SOO table; do not re-evaluate
+      // sampling each time it transitions from empty to full SOO state.
+      return false;
+    }
+    // TODO: b/396049910 -- consider managing this flag on the 1->0 size
+    // transition of SOO tables rather than the 0->1 transition.
+    common().set_soo_has_tried_sampling();
     return ABSL_PREDICT_FALSE(ShouldSampleNextTable());
   }
 

absl/container/internal/raw_hash_set_test.cc

@@ -2932,6 +2932,70 @@ TEST(RawHashSamplerTest, SooTableInsertToEmpty) {
   }
 }
 
+// Verifies that repeated insertions and erasures on an SOO table do not cause
+// the sampling decision to be evaluated multiple times, preventing
+// oversampling.
+TEST(RawHashSamplerTest, SooTableRepeatedInsertEraseDoesNotOversample) {
+  if (SooInt32Table().capacity() != SooCapacity()) {
+    CHECK_LT(sizeof(void*), 8) << "missing SOO coverage";
+    GTEST_SKIP() << "not SOO on this platform";
+  }
+  std::vector<const HashtablezInfo*> infos =
+      SampleSooMutation([](SooInt32Table& t) {
+        for (int i = 0; i < 10; ++i) {
+          t.insert(1);
+          t.erase(1);
+        }
+      });
+
+  // SampleSooMutation checks EXPECT_NEAR(sampled/total, 1%, 0.5%).
+  // If the sampling logic is incorrectly evaluated on every 0->1 element
+  // transition, the chance of being sampled per table approaches 10% (1 -
+  // 0.99^10), which is enough to cause this test to fail. By passing, this test
+  // verifies that the sampling decision is evaluated exactly once per SOO table
+  // instance.
+}
+
+// Verifies that copy-constructing or copy-assigning an SOO table does not
+// incorrectly trigger new sampling evaluations.
+TEST(RawHashSamplerTest, SooTableCopyDoesNotOversample) {
+  if (SooInt32Table().capacity() != SooCapacity()) {
+    CHECK_LT(sizeof(void*), 8) << "missing SOO coverage";
+    GTEST_SKIP() << "not SOO on this platform";
+  }
+  std::vector<const HashtablezInfo*> infos =
+      SampleSooMutation([](SooInt32Table& t) {
+        t.insert(1);
+        t.erase(1);
+        SooInt32Table t_copy(t);
+        for (int i = 0; i < 10; ++i) {
+          t_copy.insert(1);
+          t_copy.erase(1);
+        }
+        t = std::move(t_copy);
+      });
+}
+
+// Verifies that move-constructing or move-assigning an SOO table correctly
+// transfers the sampling state and does not trigger oversampling.
+TEST(RawHashSamplerTest, SooTableMoveDoesNotOversample) {
+  if (SooInt32Table().capacity() != SooCapacity()) {
+    CHECK_LT(sizeof(void*), 8) << "missing SOO coverage";
+    GTEST_SKIP() << "not SOO on this platform";
+  }
+  std::vector<const HashtablezInfo*> infos =
+      SampleSooMutation([](SooInt32Table& t) {
+        t.insert(1);
+        t.erase(1);
+        SooInt32Table t_moved(std::move(t));
+        for (int i = 0; i < 10; ++i) {
+          t_moved.insert(1);
+          t_moved.erase(1);
+        }
+        t = std::move(t_moved);
+      });
+}
+
 TEST(RawHashSamplerTest, SooTableReserveToEmpty) {
   if (SooInt32Table().capacity() != SooCapacity()) {
     CHECK_LT(sizeof(void*), 8) << "missing SOO coverage";
@@ -3029,6 +3093,63 @@ TEST(RawHashSamplerTest, SooTableRehashShrinkWhenSizeFitsInSoo) {
     ASSERT_EQ(info->total_probe_length, 0);
   }
 }
+
+// Verifies that a moved-from table does not retain the sampled state of the
+// original table, allowing it to be used without incorrectly updating global
+// sampling stats.
+TEST(RawHashSamplerTest, MovedFromTableIsNotSampled) {
+  if (SooInt32Table().capacity() != SooCapacity()) {
+    CHECK_LT(sizeof(void*), 8) << "missing SOO coverage";
+    GTEST_SKIP() << "not SOO on this platform";
+  }
+
+  SetSamplingRateTo1Percent();
+  auto& sampler = GlobalHashtablezSampler();
+  size_t start_size = 0;
+  absl::flat_hash_set<const HashtablezInfo*> preexisting_info;
+  sampler.Iterate([&](const HashtablezInfo& info) {
+    preexisting_info.insert(&info);
+    ++start_size;
+  });
+
+  SooInt32Table t1;
+  // Loop until t1 is sampled
+  while (true) {
+    t1 = SooInt32Table();
+    t1.insert(1);
+    size_t new_size = 0;
+    sampler.Iterate([&](const HashtablezInfo&) { ++new_size; });
+    if (new_size > start_size) break;
+  }
+
+  // Move the table
+  SooInt32Table t2 = std::move(t1);
+
+  // Disable sampling to ensure any new sampling is a bug.
+  SetHashtablezEnabled(false);
+
+  // t2 is now the sampled table. t1 is moved-from.
+  // We want to verify that t2 is still sampled, and that t1 isn't sampled
+  // anymore, even if we insert a new entry into it.
+  t1.clear();  // Must clear before using a moved-from table.
+  t1.insert(2);
+  t2.insert(2);
+
+  // Verify no new sample was generated, and t2's sample size is now 2.
+  size_t final_size = 0;
+  const HashtablezInfo* latest_info = nullptr;
+  size_t dropped = sampler.Iterate([&](const HashtablezInfo& info) {
+    ++final_size;
+    if (!preexisting_info.contains(&info)) {
+      latest_info = &info;
+    }
+  });
+  EXPECT_EQ(0, dropped);
+  EXPECT_EQ(final_size, start_size + 1);
+  ASSERT_NE(latest_info, nullptr);
+  EXPECT_EQ(latest_info->size.load(std::memory_order_relaxed), 2);
+}
+
 #endif  // ABSL_INTERNAL_HASHTABLEZ_SAMPLE
 
 TEST(RawHashSamplerTest, DoNotSampleCustomAllocators) {