From 018de0cafdaed248963dd998db411a7b19cc9a51 Mon Sep 17 00:00:00 2001 From: Vitaly Goldshteyn Date: Thu, 6 Feb 2025 11:55:42 -0800 Subject: [PATCH] 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 | 1 + absl/container/CMakeLists.txt | 1 + absl/container/internal/hashtablez_sampler.cc | 6 +++ absl/container/internal/hashtablez_sampler.h | 5 ++ absl/container/internal/raw_hash_set_test.cc | 46 +++++++++++++------ absl/container/sample_element_size_test.cc | 11 ++--- 6 files changed, 49 insertions(+), 21 deletions(-) diff --git a/absl/container/BUILD.bazel b/absl/container/BUILD.bazel index 19dc91de..cd0438ee 100644 --- a/absl/container/BUILD.bazel +++ b/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", diff --git a/absl/container/CMakeLists.txt b/absl/container/CMakeLists.txt index 85a4ce21..ce65359c 100644 --- a/absl/container/CMakeLists.txt +++ b/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 ) diff --git a/absl/container/internal/hashtablez_sampler.cc b/absl/container/internal/hashtablez_sampler.cc index 7b9a463f..ad5a493c 100644 --- a/absl/container/internal/hashtablez_sampler.cc +++ b/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, diff --git a/absl/container/internal/hashtablez_sampler.h b/absl/container/internal/hashtablez_sampler.h index 0c30a7e8..305dc855 100644 --- a/absl/container/internal/hashtablez_sampler.h +++ b/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) { diff --git a/absl/container/internal/raw_hash_set_test.cc b/absl/container/internal/raw_hash_set_test.cc index 36749ff6..2d332908 100644 --- a/absl/container/internal/raw_hash_set_test.cc +++ b/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(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::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 preexisting_info; + + // Reserve these utility tables, so that if they sampled, they'll be + // preexisting. + absl::flat_hash_set preexisting_info(10); + absl::flat_hash_map observed_checksums(10); + absl::flat_hash_map 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 observed_checksums; - absl::flat_hash_map 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 SampleSooMutation( absl::FunctionRef 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 preexisting_info; + // Reserve the table, so that if it sampled, it'll be preexisting. + absl::flat_hash_set 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); diff --git a/absl/container/sample_element_size_test.cc b/absl/container/sample_element_size_test.cc index 22470b49..ccc616f0 100644 --- a/absl/container/sample_element_size_test.cc +++ b/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& tables, const std::vector& 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_map_tables; @@ -92,14 +95,6 @@ TEST(FlatHashMap, SampleElementSize) { std::vector> 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.