diff --git a/src/core/CMakeLists.txt b/src/core/CMakeLists.txt index dccc1065b..bec89fce1 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 42b64cecc..e5a7d392c 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 58d3bd057..405f8f79d 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 000000000..ff1e8ffef --- /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