Fix flaky tests due to sampling by introducing utility to refresh sampling counters for the current thread.

Non refreshed counters may cause unknown (implementation detail) delay for sampling after reenabling.

Disabling sampling doesn't have effect on `ShouldSampleNextTable`, so SOO tables can be resized to size 3 even after disabling. Refreshing counters will guarantee that not happening during the test.

PiperOrigin-RevId: 724020036
Change-Id: I814c90cbab18b66e66283e7049151e6819302559

absl/container/BUILD.bazel

@@ -1008,6 +1008,7 @@ cc_test(
     deps = [
         ":flat_hash_map",
         ":flat_hash_set",
+        ":hashtablez_sampler",
         ":node_hash_map",
         ":node_hash_set",
         "@googletest//:gtest",

absl/container/CMakeLists.txt

@@ -1028,5 +1028,6 @@ absl_cc_test(
     absl::flat_hash_set
     absl::node_hash_map
     absl::node_hash_set
+    absl::hashtablez_sampler
     GTest::gmock_main
 )

absl/container/internal/hashtablez_sampler.cc

@@ -134,10 +134,16 @@ HashtablezInfoHandle ForcedTrySample(size_t inline_element_size,
                                          inline_element_size, key_size,
                                          value_size, soo_capacity));
 }
+void TestOnlyRefreshSamplingStateForCurrentThread() {
+  global_next_sample.next_sample =
+      g_hashtablez_sample_parameter.load(std::memory_order_relaxed);
+  global_next_sample.sample_stride = global_next_sample.next_sample;
+}
 #else
 HashtablezInfoHandle ForcedTrySample(size_t, size_t, size_t, uint16_t) {
   return HashtablezInfoHandle{nullptr};
 }
+void TestOnlyRefreshSamplingStateForCurrentThread() {}
 #endif  // ABSL_INTERNAL_HASHTABLEZ_SAMPLE
 
 HashtablezInfo* SampleSlow(SamplingState& next_sample,

absl/container/internal/hashtablez_sampler.h

@@ -241,6 +241,11 @@ HashtablezInfoHandle ForcedTrySample(size_t inline_element_size,
                                      size_t key_size, size_t value_size,
                                      uint16_t soo_capacity);
 
+// In case sampling needs to be disabled and re-enabled in tests, this function
+// can be used to reset the sampling state for the current thread.
+// It is useful to avoid sampling attempts and sampling delays in tests.
+void TestOnlyRefreshSamplingStateForCurrentThread();
+
 // Returns a sampling handle.
 inline HashtablezInfoHandle Sample(size_t inline_element_size, size_t key_size,
                                    size_t value_size, uint16_t soo_capacity) {

absl/container/internal/raw_hash_set_test.cc

@@ -99,6 +99,23 @@ using ::testing::UnorderedElementsAre;
 // Convenience function to static cast to ctrl_t.
 ctrl_t CtrlT(int i) { return static_cast<ctrl_t>(i); }
 
+// Enables sampling with 1 percent sampling rate and
+// resets the rate counter for the current thread.
+void SetSamplingRateTo1Percent() {
+  SetHashtablezEnabled(true);
+  SetHashtablezSampleParameter(100);  // Sample ~1% of tables.
+  // Reset rate counter for the current thread.
+  TestOnlyRefreshSamplingStateForCurrentThread();
+}
+
+// Disables sampling and resets the rate counter for the current thread.
+void DisableSampling() {
+  SetHashtablezEnabled(false);
+  SetHashtablezSampleParameter(1 << 16);
+  // Reset rate counter for the current thread.
+  TestOnlyRefreshSamplingStateForCurrentThread();
+}
+
 TEST(GrowthInfoTest, GetGrowthLeft) {
   GrowthInfo gi;
   gi.InitGrowthLeftNoDeleted(5);
@@ -1061,6 +1078,7 @@ TYPED_TEST(SmallTableResizeTest, ResizeGrowSmallTables) {
 }
 
 TYPED_TEST(SmallTableResizeTest, ResizeReduceSmallTables) {
+  DisableSampling();
   for (size_t source_size = 0; source_size < 32; ++source_size) {
     for (size_t target_size = 0; target_size <= source_size; ++target_size) {
       TypeParam t;
@@ -2233,7 +2251,7 @@ TYPED_TEST(SooTest, FindFullDeletedRegression) {
 
 TYPED_TEST(SooTest, ReplacingDeletedSlotDoesNotRehash) {
   // We need to disable hashtablez to avoid issues related to SOO and sampling.
-  SetHashtablezEnabled(false);
+  DisableSampling();
 
   size_t n;
   {
@@ -2601,7 +2619,7 @@ TYPED_TEST(SooTest, IterationOrderChangesOnRehash) {
 // This prevents dependency on pointer stability on small tables.
 TYPED_TEST(SooTest, UnstablePointers) {
   // We need to disable hashtablez to avoid issues related to SOO and sampling.
-  SetHashtablezEnabled(false);
+  DisableSampling();
 
   TypeParam table;
 
@@ -2739,12 +2757,17 @@ TYPED_TEST_SUITE(RawHashSamplerTest, RawHashSamplerTestTypes);
 TYPED_TEST(RawHashSamplerTest, Sample) {
   constexpr bool soo_enabled = std::is_same<SooIntTable, TypeParam>::value;
   // Enable the feature even if the prod default is off.
-  SetHashtablezEnabled(true);
-  SetHashtablezSampleParameter(100);  // Sample ~1% of tables.
+  SetSamplingRateTo1Percent();
 
   auto& sampler = GlobalHashtablezSampler();
   size_t start_size = 0;
-  absl::flat_hash_set<const HashtablezInfo*> preexisting_info;
+
+  // Reserve these utility tables, so that if they sampled, they'll be
+  // preexisting.
+  absl::flat_hash_set<const HashtablezInfo*> preexisting_info(10);
+  absl::flat_hash_map<size_t, int> observed_checksums(10);
+  absl::flat_hash_map<ssize_t, int> reservations(10);
+
   start_size += sampler.Iterate([&](const HashtablezInfo& info) {
     preexisting_info.insert(&info);
     ++start_size;
@@ -2771,8 +2794,6 @@ TYPED_TEST(RawHashSamplerTest, Sample) {
     }
   }
   size_t end_size = 0;
-  absl::flat_hash_map<size_t, int> observed_checksums;
-  absl::flat_hash_map<ssize_t, int> reservations;
   end_size += sampler.Iterate([&](const HashtablezInfo& info) {
     ++end_size;
     if (preexisting_info.contains(&info)) return;
@@ -2811,12 +2832,12 @@ TYPED_TEST(RawHashSamplerTest, Sample) {
 std::vector<const HashtablezInfo*> SampleSooMutation(
     absl::FunctionRef<void(SooIntTable&)> mutate_table) {
   // Enable the feature even if the prod default is off.
-  SetHashtablezEnabled(true);
-  SetHashtablezSampleParameter(100);  // Sample ~1% of tables.
+  SetSamplingRateTo1Percent();
 
   auto& sampler = GlobalHashtablezSampler();
   size_t start_size = 0;
-  absl::flat_hash_set<const HashtablezInfo*> preexisting_info;
+  // Reserve the table, so that if it sampled, it'll be preexisting.
+  absl::flat_hash_set<const HashtablezInfo*> preexisting_info(10);
   start_size += sampler.Iterate([&](const HashtablezInfo& info) {
     preexisting_info.insert(&info);
     ++start_size;
@@ -2940,8 +2961,7 @@ TEST(RawHashSamplerTest, SooTableRehashShrinkWhenSizeFitsInSoo) {
 
 TEST(RawHashSamplerTest, DoNotSampleCustomAllocators) {
   // Enable the feature even if the prod default is off.
-  SetHashtablezEnabled(true);
-  SetHashtablezSampleParameter(100);  // Sample ~1% of tables.
+  SetSamplingRateTo1Percent();
 
   auto& sampler = GlobalHashtablezSampler();
   size_t start_size = 0;
@@ -3630,7 +3650,7 @@ TEST(Table, RehashToSooUnsampled) {
 
   // We disable hashtablez sampling for this test to ensure that the table isn't
   // sampled. When the table is sampled, it won't rehash down to SOO.
-  SetHashtablezEnabled(false);
+  DisableSampling();
 
   t.reserve(100);
   t.insert(0);

absl/container/sample_element_size_test.cc

@@ -21,6 +21,7 @@
 #include "gtest/gtest.h"
 #include "absl/container/flat_hash_map.h"
 #include "absl/container/flat_hash_set.h"
+#include "absl/container/internal/hashtablez_sampler.h"
 #include "absl/container/node_hash_map.h"
 #include "absl/container/node_hash_set.h"
 
@@ -46,6 +47,7 @@ void TestInlineElementSize(
     std::vector<Table>& tables,
     const std::vector<typename Table::value_type>& values,
     size_t expected_element_size) {
+  EXPECT_GT(values.size(), 0);
   for (int i = 0; i < 10; ++i) {
     // We create a new table and must store it somewhere so that when we store
     // a pointer to the resulting `HashtablezInfo` into `preexisting_info`
@@ -82,6 +84,7 @@ TEST(FlatHashMap, SampleElementSize) {
   // Enable sampling even if the prod default is off.
   SetHashtablezEnabled(true);
   SetHashtablezSampleParameter(1);
+  TestOnlyRefreshSamplingStateForCurrentThread();
 
   auto& sampler = GlobalHashtablezSampler();
   std::vector<flat_hash_map<int, bigstruct>> flat_map_tables;
@@ -92,14 +95,6 @@ TEST(FlatHashMap, SampleElementSize) {
   std::vector<std::pair<const int, bigstruct>> map_values = {{0, bigstruct{}},
                                                              {1, bigstruct{}}};
 
-  // It takes thousands of new tables after changing the sampling parameters
-  // before you actually get some instrumentation.  And if you must actually
-  // put something into those tables.
-  for (int i = 0; i < 10000; ++i) {
-    flat_map_tables.emplace_back();
-    flat_map_tables.back()[i] = bigstruct{};
-  }
-
   // clang-tidy gives a false positive on this declaration.  This unordered set
   // cannot be a flat_hash_set, however, since that would introduce a mutex
   // deadlock.