Shard SampleRecorder graveyard lists across 16 cache-line-aligned shards tracked by a non_empty_shards_ bitmask to eliminate global graveyard mutex contention and cache-line bouncing during concurrent sample registration and unregistration. PiperOrigin-RevId: 1003282003
diff --git a/absl/profiling/BUILD.bazel b/absl/profiling/BUILD.bazel index 2d9ecc6..c3d13fe 100644 --- a/absl/profiling/BUILD.bazel +++ b/absl/profiling/BUILD.bazel
@@ -70,6 +70,22 @@ ], ) +cc_binary( + name = "sample_recorder_benchmark", + testonly = True, + srcs = ["internal/sample_recorder_benchmark.cc"], + copts = ABSL_TEST_COPTS, + linkopts = ABSL_DEFAULT_LINKOPTS, + tags = ["benchmark"], + visibility = ["//visibility:private"], + deps = [ + ":sample_recorder", + "//absl/base:config", + "//absl/time", + "@google_benchmark//:benchmark_main", + ], +) + cc_library( name = "exponential_biased", srcs = ["internal/exponential_biased.cc"],
diff --git a/absl/profiling/CMakeLists.txt b/absl/profiling/CMakeLists.txt index 0169182..ddf2c83 100644 --- a/absl/profiling/CMakeLists.txt +++ b/absl/profiling/CMakeLists.txt
@@ -21,7 +21,9 @@ ${ABSL_DEFAULT_COPTS} DEPS absl::base + absl::core_headers absl::synchronization + absl::time ) absl_cc_test( @@ -34,6 +36,8 @@ DEPS absl::random_random absl::sample_recorder + absl::synchronization + absl::thread_pool absl::time GTest::gmock_main )
diff --git a/absl/profiling/internal/sample_recorder.h b/absl/profiling/internal/sample_recorder.h index 8663d6c..0796ed7 100644 --- a/absl/profiling/internal/sample_recorder.h +++ b/absl/profiling/internal/sample_recorder.h
@@ -46,7 +46,7 @@ // prevents races with sampling and resurrecting an object. absl::Mutex init_mu; T* next = nullptr; - T* dead ABSL_GUARDED_BY(init_mu) = nullptr; + std::atomic<T*> dead{nullptr}; int64_t weight; // How many sampling events were required to sample this one. }; @@ -82,44 +82,46 @@ void SetMaxSamples(size_t max); private: + static constexpr size_t kNumGraveyardShards = 16; + + struct alignas(64) GraveyardShard { + absl::Mutex mu; + std::atomic<T*> head{nullptr}; + }; + + static size_t ThreadShardIndex() { + thread_local size_t shard = []() { + static std::atomic<size_t> next_shard{0}; + return next_shard.fetch_add(1, std::memory_order_relaxed); + }(); + return shard & (kNumGraveyardShards - 1); + } + void PushNew(T* sample); void PushDead(T* sample); template <typename... Targs> T* PopDead(Targs... args); + alignas(64) std::atomic<size_t> size_estimate_; + alignas(64) std::atomic<T*> all_; + alignas(64) std::atomic<size_t> max_samples_{1 << 20}; + std::atomic<DisposeCallback> dispose_; std::atomic<size_t> dropped_samples_; - std::atomic<size_t> size_estimate_; - std::atomic<size_t> max_samples_{1 << 20}; + alignas(64) std::atomic<uint32_t> non_empty_shards_{0}; // Intrusive lock free linked lists for tracking samples. // // `all_` records all samples (they are never removed from this list) and is // terminated with a `nullptr`. // - // `graveyard_.dead` is a circular linked list. When it is empty, - // `graveyard_.dead == &graveyard`. The list is circular so that - // every item on it (even the last) has a non-null dead pointer. This allows - // `Iterate` to determine if a given sample is live or dead using only - // information on the sample itself. - // - // For example, nodes [A, B, C, D, E] with [A, C, E] alive and [B, D] dead - // looks like this (G is the Graveyard): - // - // +---+ +---+ +---+ +---+ +---+ - // all -->| A |--->| B |--->| C |--->| D |--->| E | - // | | | | | | | | | | - // +---+ | | +->| |-+ | | +->| |-+ | | - // | G | +---+ | +---+ | +---+ | +---+ | +---+ - // | | | | | | - // | | --------+ +--------+ | - // +---+ | - // ^ | - // +--------------------------------------+ - // - std::atomic<T*> all_; + // `graveyard_shards_` partitions dead samples across cache-line-aligned + // circular linked lists terminated by `&graveyard_`. Every dead item has a + // non-null `dead` pointer so `Iterate` can determine if a sample is live or + // dead using only per-sample state, while per-thread sharding and the + // `non_empty_shards_` summary mask avoid both global lock contention and + // empty-shard probing. T graveyard_; - - std::atomic<DisposeCallback> dispose_; + GraveyardShard graveyard_shards_[kNumGraveyardShards]; }; template <typename T> @@ -130,9 +132,17 @@ template <typename T> SampleRecorder<T>::SampleRecorder() - : dropped_samples_(0), size_estimate_(0), all_(nullptr), dispose_(nullptr) { - absl::MutexLock l(graveyard_.init_mu); - graveyard_.dead = &graveyard_; + : size_estimate_(0), + all_(nullptr), + dispose_(nullptr), + dropped_samples_(0) { + { + absl::MutexLock l(graveyard_.init_mu); + graveyard_.dead.store(&graveyard_, std::memory_order_relaxed); + } + for (size_t i = 0; i < kNumGraveyardShards; ++i) { + graveyard_shards_[i].head.store(&graveyard_, std::memory_order_relaxed); + } } template <typename T> @@ -160,38 +170,80 @@ dispose(*sample); } - absl::MutexLock graveyard_lock(graveyard_.init_mu); - absl::MutexLock sample_lock(sample->init_mu); - sample->dead = graveyard_.dead; - graveyard_.dead = sample; + // Mark `sample` dead under `sample->init_mu` before publishing it to the + // thread's home graveyard shard. + { + absl::MutexLock sample_lock(sample->init_mu); + sample->dead.store(&graveyard_, std::memory_order_relaxed); + } + const size_t shard_idx = ThreadShardIndex(); + GraveyardShard& shard = graveyard_shards_[shard_idx]; + { + absl::MutexLock shard_lock(shard.mu); + T* prev_head = shard.head.load(std::memory_order_relaxed); + sample->dead.store(prev_head, std::memory_order_relaxed); + shard.head.store(sample, std::memory_order_relaxed); + if (prev_head == &graveyard_) { + non_empty_shards_.fetch_or(uint32_t{1} << shard_idx, + std::memory_order_relaxed); + } + } } template <typename T> template <typename... Targs> T* SampleRecorder<T>::PopDead(Targs... args) { - T* sample; - { - absl::MutexLock graveyard_lock(graveyard_.init_mu); - - // The list is circular, so eventually it collapses down to - // graveyard_.dead == &graveyard_ - // when it is empty. - sample = graveyard_.dead; - if (sample == &graveyard_) return nullptr; - - graveyard_.dead = ABSL_TS_UNCHECKED_READ(sample->dead); - // Release the global graveyard lock early, before the potentially slow - // preparation. + T* sample = nullptr; + const size_t home = ThreadShardIndex(); + GraveyardShard& home_shard = graveyard_shards_[home]; + if (home_shard.head.load(std::memory_order_relaxed) != &graveyard_) { + absl::MutexLock shard_lock(home_shard.mu); + T* head = home_shard.head.load(std::memory_order_relaxed); + if (head != &graveyard_) { + T* next_dead = head->dead.load(std::memory_order_relaxed); + home_shard.head.store(next_dead, std::memory_order_relaxed); + if (next_dead == &graveyard_) { + non_empty_shards_.fetch_and(~(uint32_t{1} << home), + std::memory_order_relaxed); + } + sample = head; + } } - // `sample` is detached from the graveyard and will not be used in PopDead. - // SampleRecorder<T>::Iterate will acquire per sample lock, so there will be - // no data race either. + if (sample == nullptr) { + uint32_t mask = non_empty_shards_.load(std::memory_order_relaxed); + while (mask != 0) { + const size_t idx = static_cast<size_t>(__builtin_ctz(mask)); + mask &= mask - 1; + if (idx == home) continue; + GraveyardShard& shard = graveyard_shards_[idx]; + if (shard.head.load(std::memory_order_relaxed) == &graveyard_) { + continue; + } + absl::MutexLock shard_lock(shard.mu); + T* head = shard.head.load(std::memory_order_relaxed); + if (head != &graveyard_) { + T* next_dead = head->dead.load(std::memory_order_relaxed); + shard.head.store(next_dead, std::memory_order_relaxed); + if (next_dead == &graveyard_) { + non_empty_shards_.fetch_and(~(uint32_t{1} << idx), + std::memory_order_relaxed); + } + sample = head; + break; + } + } + } + if (sample == nullptr) return nullptr; + + // `sample` is detached from its shard and will not be observed by PopDead. + // SampleRecorder<T>::Iterate acquires the per-sample lock before inspecting + // `sample->dead`. absl::MutexLock sample_lock(sample->init_mu); // Note: consider PrepareForSampling out of the lock. // Currently many clients have PrepareForSampling marked as // ABSL_EXCLUSIVE_LOCKS_REQUIRED. sample->PrepareForSampling(std::forward<Targs>(args)...); - sample->dead = nullptr; + sample->dead.store(nullptr, std::memory_order_relaxed); return sample; } @@ -239,7 +291,7 @@ T* s = all_.load(std::memory_order_acquire); while (s != nullptr) { absl::MutexLock l(s->init_mu); - if (s->dead == nullptr) { + if (s->dead.load(std::memory_order_relaxed) == nullptr) { f(*s); } s = s->next;
diff --git a/absl/profiling/internal/sample_recorder_test.cc b/absl/profiling/internal/sample_recorder_test.cc index 4e1437e..db1d10c 100644 --- a/absl/profiling/internal/sample_recorder_test.cc +++ b/absl/profiling/internal/sample_recorder_test.cc
@@ -187,6 +187,100 @@ sampler.Unregister(info2); } +TEST(SampleRecorderTest, ConcurrentChurnResurrectionAndIteration) { + SampleRecorder<Info> sampler; + static std::atomic<int64_t> disposed_count{0}; + disposed_count.store(0, std::memory_order_relaxed); + sampler.SetDisposeCallback([](const Info& info) { + EXPECT_GE(info.weight, 100); + disposed_count.fetch_add(1, std::memory_order_relaxed); + }); + + Notification stop; + ThreadPool pool(8); + for (int t = 0; t < 6; ++t) { + pool.Schedule([&sampler, &stop, t]() { + const int64_t w = 100 + t; + while (!stop.HasBeenNotified()) { + Info* a = sampler.Register(w); + Info* b = sampler.Register(w); + ASSERT_NE(a, nullptr); + ASSERT_NE(b, nullptr); + ASSERT_NE(a, b); + a->size.store(static_cast<size_t>(t + 1), std::memory_order_release); + b->size.store(static_cast<size_t>(t + 1), std::memory_order_release); + sampler.Unregister(a); + sampler.Unregister(b); + } + }); + } + for (int t = 0; t < 2; ++t) { + pool.Schedule([&sampler, &stop]() { + while (!stop.HasBeenNotified()) { + sampler.Iterate([](const Info& info) { + EXPECT_GE(info.weight, 100); + EXPECT_LT(info.weight, 106); + }); + } + }); + } + absl::SleepFor(absl::Milliseconds(500)); + stop.Notify(); + sampler.SetDisposeCallback(nullptr); + EXPECT_GT(disposed_count.load(std::memory_order_relaxed), 0); +} + +TEST(SampleRecorderTest, CrossThreadShardStealingAndMaxSamples) { + SampleRecorder<Info> sampler; + // Register() checks `size > max_samples_` where `size` is pre-increment, so + // max_samples = 7 allows indices 0..7 (8 concurrent live samples). + sampler.SetMaxSamples(7); + + // Unregister samples from one worker thread and resurrect them all from + // a different thread to exercise cross-shard graveyard stealing and soft + // max_samples enforcement without leaking or duplicating nodes. + std::vector<Info*> seeded; + { + ThreadPool producer_pool(1); + Notification done; + producer_pool.Schedule([&]() { + for (int i = 0; i < 8; ++i) { + Info* s = sampler.Register(200 + i); + ASSERT_NE(s, nullptr); + seeded.push_back(s); + } + for (Info* s : seeded) { + sampler.Unregister(s); + } + done.Notify(); + }); + done.WaitForNotification(); + } + + std::vector<Info*> resurrected; + for (int i = 0; i < 8; ++i) { + Info* s = sampler.Register(300 + i); + ASSERT_NE(s, nullptr); + resurrected.push_back(s); + } + // Every sample must have been stolen from the producer's graveyard shard + // rather than allocating duplicate heap entries. + std::sort(seeded.begin(), seeded.end()); + std::vector<Info*> sorted_resurrected = resurrected; + std::sort(sorted_resurrected.begin(), sorted_resurrected.end()); + EXPECT_EQ(seeded, sorted_resurrected); + + // Exceeding max_samples drops registrations cleanly and tracks dropped count. + EXPECT_EQ(sampler.Register(999), nullptr); + size_t live = 0; + EXPECT_EQ(sampler.Iterate([&](const Info&) { ++live; }), 1); + EXPECT_EQ(live, 8); + + for (Info* s : resurrected) { + sampler.Unregister(s); + } +} + } // namespace } // namespace profiling_internal ABSL_NAMESPACE_END