From 5efc8f11d2c4f6c5c3ad9a7069bf67af19befe9d Mon Sep 17 00:00:00 2001 From: Shahar Mike Date: Sun, 6 Oct 2024 22:00:57 +0300 Subject: [PATCH] opt: Optimize `AllocationTracker` to be efficient when enabled (#3875) Today there's a cost to enabling AllocationTracker, even for rarely used memory bands. This PR slightly optimizes the "happy path" (i.e. allocations outside the tracked range), and also for the case where we track 100% of the allocations. Also, add a unit test for this class. --- src/core/CMakeLists.txt | 1 + src/core/allocation_tracker.cc | 24 +++- src/core/allocation_tracker.h | 4 + src/core/allocation_tracker_test.cc | 166 ++++++++++++++++++++++++++++ 4 files changed, 192 insertions(+), 3 deletions(-) create mode 100644 src/core/allocation_tracker_test.cc diff --git a/src/core/CMakeLists.txt b/src/core/CMakeLists.txt index dccc1065bd1b..bec89fce15fd 100644 --- a/src/core/CMakeLists.txt +++ b/src/core/CMakeLists.txt @@ -28,3 +28,4 @@ cxx_test(bptree_set_test dfly_core LABELS DFLY) cxx_test(score_map_test dfly_core LABELS DFLY) cxx_test(flatbuffers_test dfly_core TRDP::flatbuffers LABELS DFLY) cxx_test(bloom_test dfly_core LABELS DFLY) +cxx_test(allocation_tracker_test dfly_core absl::random_random LABELS DFLY) diff --git a/src/core/allocation_tracker.cc b/src/core/allocation_tracker.cc index 42b64cecc2b6..e5a7d392c04b 100644 --- a/src/core/allocation_tracker.cc +++ b/src/core/allocation_tracker.cc @@ -24,6 +24,9 @@ bool AllocationTracker::Add(const TrackingInfo& info) { } tracking_.push_back(info); + + UpdateAbsSizes(); + return true; } @@ -37,6 +40,8 @@ bool AllocationTracker::Remove(size_t lower_bound, size_t upper_bound) { }), tracking_.end()); + UpdateAbsSizes(); + return before_size != tracking_.size(); } @@ -49,7 +54,7 @@ absl::Span AllocationTracker::GetRanges() } void AllocationTracker::ProcessNew(void* ptr, size_t size) { - if (tracking_.empty()) { + if (size < abs_min_size_ || size > abs_max_size_) { return; } @@ -59,9 +64,13 @@ void AllocationTracker::ProcessNew(void* ptr, size_t size) { // Prevent endless recursion, in case logging allocates memory inside_tracker_ = true; - double random = absl::Uniform(g_bitgen, 0.0, 1.0); for (const auto& band : tracking_) { - if (random >= band.sample_odds || size > band.upper_bound || size < band.lower_bound) { + if (size > band.upper_bound || size < band.lower_bound) { + continue; + } + + // Micro optimization: in case sample_odds == 1.0 - do not draw a random number + if (band.sample_odds != 1.0 && absl::Uniform(g_bitgen, 0.0, 1.0) >= band.sample_odds) { continue; } @@ -92,4 +101,13 @@ void AllocationTracker::ProcessDelete(void* ptr) { inside_tracker_ = false; } +void AllocationTracker::UpdateAbsSizes() { + abs_min_size_ = 0; + abs_max_size_ = 0; + for (const auto& tracker : tracking_) { + abs_min_size_ = std::min(abs_min_size_, tracker.lower_bound); + abs_max_size_ = std::max(abs_max_size_, tracker.upper_bound); + } +} + } // namespace dfly diff --git a/src/core/allocation_tracker.h b/src/core/allocation_tracker.h index 58d3bd057c42..405f8f79d93c 100644 --- a/src/core/allocation_tracker.h +++ b/src/core/allocation_tracker.h @@ -46,8 +46,12 @@ class AllocationTracker { void ProcessDelete(void* ptr); private: + void UpdateAbsSizes(); + absl::InlinedVector tracking_; bool inside_tracker_ = false; + size_t abs_min_size_ = 0; + size_t abs_max_size_ = 0; }; } // namespace dfly diff --git a/src/core/allocation_tracker_test.cc b/src/core/allocation_tracker_test.cc new file mode 100644 index 000000000000..ff1e8ffefa34 --- /dev/null +++ b/src/core/allocation_tracker_test.cc @@ -0,0 +1,166 @@ +#include +#include +#include + +#include +#include + +#include "base/gtest.h" +#include "base/logging.h" + +#define INJECT_ALLOCATION_TRACKER +#include "core/allocation_tracker.h" + +namespace dfly { +namespace { +using namespace std; +using namespace testing; + +class LogSink : public google::LogSink { + public: + void send(google::LogSeverity severity, const char* full_filename, const char* base_filename, + int line, const struct tm* tm_time, const char* message, size_t message_len) override { + logs_.push_back(string(message, message_len)); + } + + const vector& GetLogs() const { + return logs_; + } + + void Clear() { + logs_.clear(); + } + + private: + vector logs_; +}; + +class AllocationTrackerTest : public Test { + protected: + AllocationTrackerTest() { + google::AddLogSink(&log_sink_); + } + + ~AllocationTrackerTest() { + google::RemoveLogSink(&log_sink_); + AllocationTracker::Get().Clear(); + } + + vector GetLogsDelta() { + auto logs = log_sink_.GetLogs(); + log_sink_.Clear(); + return logs; + } + + void Allocate(size_t s) { + CHECK(buffer_.empty()); + buffer_.resize(s); // allocate 1mb before setting up tracking + } + + void Deallocate() { + buffer_.clear(); + // Force deallocation + buffer_.shrink_to_fit(); + } + + private: + LogSink log_sink_; + string buffer_; +}; + +TEST_F(AllocationTrackerTest, UnusedTracker) { + Allocate(1'000'000); // allocate 1mb before setting up tracking + EXPECT_THAT(GetLogsDelta(), Not(Contains(HasSubstr("Allocating")))); + Deallocate(); +} + +TEST_F(AllocationTrackerTest, UsedTracker) { + AllocationTracker::Get().Add( + {.lower_bound = 1'000'000, .upper_bound = 2'000'000, .sample_odds = 1.0}); + Allocate(1'000'000); // allocate 1mb before setting up tracking + EXPECT_THAT(GetLogsDelta(), Contains(HasSubstr("Allocating"))); + EXPECT_THAT(GetLogsDelta(), Not(Contains(HasSubstr("Deallocating")))); + Deallocate(); + EXPECT_THAT(GetLogsDelta(), Contains(HasSubstr("Deallocating"))); + + // Allocate below threshold + Allocate(100'000); + EXPECT_THAT(GetLogsDelta(), Not(Contains(HasSubstr("Allocating")))); + Deallocate(); + EXPECT_THAT(GetLogsDelta(), Not(Contains(HasSubstr("Deallocating")))); + + // Allocate above threshold + Allocate(10'000'000); + EXPECT_THAT(GetLogsDelta(), Not(Contains(HasSubstr("Allocating")))); + Deallocate(); + EXPECT_THAT(GetLogsDelta(), Not(Contains(HasSubstr("Deallocating")))); + + // Remove allocator - stops logging + EXPECT_TRUE(AllocationTracker::Get().Remove(1'000'000, 2'000'000)); + Allocate(1'000'000); // allocate 1mb before setting up tracking + EXPECT_THAT(GetLogsDelta(), Not(Contains(HasSubstr("Allocating")))); + Deallocate(); + EXPECT_THAT(GetLogsDelta(), Not(Contains(HasSubstr("Deallocating")))); +} + +TEST_F(AllocationTrackerTest, MultipleRanges) { + AllocationTracker::Get().Add( + {.lower_bound = 1'000'000, .upper_bound = 2'000'000, .sample_odds = 1.0}); + AllocationTracker::Get().Add( + {.lower_bound = 100'000'000, .upper_bound = 200'000'000, .sample_odds = 1.0}); + + // Below all ranges + Allocate(100'000); + EXPECT_THAT(GetLogsDelta(), Not(Contains(HasSubstr("Allocating")))); + Deallocate(); + + // Between ranges + Allocate(10'000'000); + EXPECT_THAT(GetLogsDelta(), Not(Contains(HasSubstr("Allocating")))); + Deallocate(); + + // Above all ranges + Allocate(500'000'000); + EXPECT_THAT(GetLogsDelta(), Not(Contains(HasSubstr("Allocating")))); + Deallocate(); + + // First range + Allocate(1'000'000); + EXPECT_THAT(GetLogsDelta(), Contains(HasSubstr("Allocating"))); + Deallocate(); + + // Second range + Allocate(100'000'000); + EXPECT_THAT(GetLogsDelta(), Contains(HasSubstr("Allocating"))); + Deallocate(); +} + +TEST_F(AllocationTrackerTest, Sampling) { + // Statistically, 80% of logs should be logged + AllocationTracker::Get().Add( + {.lower_bound = 1'000'000, .upper_bound = 2'000'000, .sample_odds = 0.8}); + + const int kIterations = 10'000; + for (int i = 0; i < kIterations; ++i) { + Allocate(1'000'000); + Deallocate(); + } + + int allocations = 0; + int deallocations = 0; + for (const string& s : GetLogsDelta()) { + if (absl::StrContains(s, "Allocating")) { + ++allocations; + } + if (absl::StrContains(s, "Deallocating")) { + ++deallocations; + } + } + + EXPECT_GE(allocations, kIterations * 0.7); + EXPECT_LE(allocations, kIterations * 0.9); + EXPECT_EQ(deallocations, 0); // we only track deletions when sample_odds == 1.0 +} + +} // namespace +} // namespace dfly