From 7e50ddbcc1664216266224adfeef504200119115 Mon Sep 17 00:00:00 2001 From: Sannidhya Chauhan Date: Thu, 27 Aug 2026 02:43:15 -0700 Subject: [PATCH] Implement byte-based memory budgeting and telemetry in ContinuousProfilerOrchestrator PiperOrigin-RevId: 971832824 --- .../lib/continuous_profiler_orchestrator.h | 115 ++++++- .../continuous_profiler_orchestrator_test.cc | 310 +++++++++++++++++- 2 files changed, 407 insertions(+), 18 deletions(-) diff --git a/tsl/profiler/lib/continuous_profiler_orchestrator.h b/tsl/profiler/lib/continuous_profiler_orchestrator.h index d12a8be24..58954be37 100644 --- a/tsl/profiler/lib/continuous_profiler_orchestrator.h +++ b/tsl/profiler/lib/continuous_profiler_orchestrator.h @@ -18,6 +18,7 @@ limitations under the License. #include #include #include +#include #include #include #include @@ -37,6 +38,15 @@ limitations under the License. namespace tsl { namespace profiler { +inline constexpr size_t kDefaultMaxBufferBytes = + 4ULL * 1024 * 1024 * 1024; // 4GB + +struct DrainedBuffer { + std::vector chunks; + uint64_t cumulative_dropped_chunks = 0; + uint64_t cumulative_dropped_bytes = 0; +}; + template class ContinuousProfilerOrchestrator : public ProfilerInterface { public: @@ -45,8 +55,10 @@ class ContinuousProfilerOrchestrator : public ProfilerInterface { static constexpr absl::Duration kMaxPollingInterval = absl::Seconds(5); explicit ContinuousProfilerOrchestrator( - std::unique_ptr profiler) + std::unique_ptr profiler, + size_t max_buffer_bytes = kDefaultMaxBufferBytes) : profiler_(std::move(profiler)), + max_buffer_bytes_(max_buffer_bytes), is_running_(false), polling_interval_(kDefaultPollingInterval) {} @@ -77,10 +89,10 @@ class ContinuousProfilerOrchestrator : public ProfilerInterface { // Stops background thread and profiling. absl::Status Stop() override { absl::Status status = StopInternal(); - auto result = profiler_->Consume(); + absl::StatusOr result = profiler_->Consume(); if (result.ok()) { absl::MutexLock lock(mutex_); - circular_buffer_.push_back(std::move(result->data)); + PushChunkLocked(std::move(*result)); } else if (!absl::IsUnimplemented(result.status())) { LOG(WARNING) << "Final Consume failed during Stop: " << result.status(); } @@ -129,18 +141,87 @@ class ContinuousProfilerOrchestrator : public ProfilerInterface { ProfilerType* profiler() { return profiler_.get(); } const ProfilerType* profiler() const { return profiler_.get(); } - std::vector PopBuffer() { + size_t max_buffer_bytes() const { return max_buffer_bytes_; } + + size_t total_buffered_bytes() const { absl::MutexLock lock(mutex_); + return total_buffered_bytes_; + } + + uint64_t dropped_chunks_count() const { + absl::MutexLock lock(mutex_); + return dropped_chunks_count_; + } + + uint64_t dropped_bytes_count() const { + absl::MutexLock lock(mutex_); + return dropped_bytes_count_; + } + + DrainedBuffer PopBufferWithTelemetry() { + std::deque local_buffer; + uint64_t dropped_chunks = 0; + uint64_t dropped_bytes = 0; + { + absl::MutexLock lock(mutex_); + local_buffer.swap(circular_buffer_); + chunk_sizes_.clear(); + total_buffered_bytes_ = 0; + dropped_chunks = dropped_chunks_count_; + dropped_bytes = dropped_bytes_count_; + } + std::vector chunks; - chunks.reserve(circular_buffer_.size()); - for (auto& item : circular_buffer_) { - chunks.push_back(std::move(item)); + chunks.reserve(local_buffer.size()); + for (auto& item : local_buffer) { + if (item.has_value()) { + chunks.push_back(std::move(item)); + } } - circular_buffer_.clear(); - return chunks; + return DrainedBuffer{ + .chunks = std::move(chunks), + .cumulative_dropped_chunks = dropped_chunks, + .cumulative_dropped_bytes = dropped_bytes, + }; } + std::vector PopBuffer() { return PopBufferWithTelemetry().chunks; } + private: + void PushChunkLocked(ConsumeResult chunk) + ABSL_EXCLUSIVE_LOCKS_REQUIRED(mutex_) { + if (!chunk.data.has_value()) { + return; + } + + if (chunk.estimated_size_bytes > max_buffer_bytes_) { + LOG_EVERY_N_SEC(WARNING, 30) + << "ContinuousProfilerOrchestrator rejected chunk of size " + << chunk.estimated_size_bytes << " bytes exceeding max buffer limit " + << max_buffer_bytes_ << " bytes."; + dropped_chunks_count_ += 1; + dropped_bytes_count_ += chunk.estimated_size_bytes; + return; + } + + while (!circular_buffer_.empty() && + (total_buffered_bytes_ + chunk.estimated_size_bytes > + max_buffer_bytes_)) { + size_t front_size = chunk_sizes_.front(); + total_buffered_bytes_ = (total_buffered_bytes_ > front_size) + ? total_buffered_bytes_ - front_size + : 0; + dropped_chunks_count_ += 1; + dropped_bytes_count_ += front_size; + circular_buffer_.pop_front(); + chunk_sizes_.pop_front(); + } + + total_buffered_bytes_ += chunk.estimated_size_bytes; + circular_buffer_.push_back(std::move(chunk.data)); + chunk_sizes_.push_back(chunk.estimated_size_bytes); + } + void IngestionLoop() { LOG(INFO) << "ContinuousProfilerOrchestrator::IngestionLoop started"; while (true) { @@ -152,14 +233,9 @@ class ContinuousProfilerOrchestrator : public ProfilerInterface { absl::MutexLock lock(mutex_); if (result.ok()) { - circular_buffer_.push_back(std::move(result->data)); - - // Cap circular buffer to prevent infinite memory growth. - if (circular_buffer_.size() > 100) { - circular_buffer_.pop_front(); - } - - AdjustIntervalLocked(result->estimated_size_bytes); + const size_t chunk_size = result->estimated_size_bytes; + PushChunkLocked(std::move(*result)); + AdjustIntervalLocked(chunk_size); } if (!is_running_) break; @@ -195,6 +271,7 @@ class ContinuousProfilerOrchestrator : public ProfilerInterface { } std::unique_ptr profiler_; + const size_t max_buffer_bytes_; mutable absl::Mutex mutex_; absl::CondVar cv_; @@ -203,6 +280,10 @@ class ContinuousProfilerOrchestrator : public ProfilerInterface { absl::Duration polling_interval_ ABSL_GUARDED_BY(mutex_); std::deque circular_buffer_ ABSL_GUARDED_BY(mutex_); + std::deque chunk_sizes_ ABSL_GUARDED_BY(mutex_); + size_t total_buffered_bytes_ ABSL_GUARDED_BY(mutex_) = 0; + uint64_t dropped_chunks_count_ ABSL_GUARDED_BY(mutex_) = 0; + uint64_t dropped_bytes_count_ ABSL_GUARDED_BY(mutex_) = 0; }; } // namespace profiler diff --git a/tsl/profiler/lib/continuous_profiler_orchestrator_test.cc b/tsl/profiler/lib/continuous_profiler_orchestrator_test.cc index a1c81fde6..4341e5eeb 100644 --- a/tsl/profiler/lib/continuous_profiler_orchestrator_test.cc +++ b/tsl/profiler/lib/continuous_profiler_orchestrator_test.cc @@ -77,7 +77,7 @@ TEST(ContinuousProfilerOrchestratorTest, } return ConsumeResult{ .data = std::any(count), - .estimated_size_bytes = 1000 * 1024 * 1024 // 1000MB (>512MB) + .estimated_size_bytes = 600 * 1024 * 1024 // 600MB (>512MB) }; }); @@ -186,6 +186,314 @@ TEST(ContinuousProfilerOrchestratorTest, SerializeChunks) { EXPECT_EQ(spaces.size(), 1); } +TEST(ContinuousProfilerOrchestratorTest, CustomMemoryBudget) { + auto mock_profiler = std::make_unique(); + ContinuousProfilerOrchestrator default_orchestrator( + std::move(mock_profiler)); + EXPECT_EQ(default_orchestrator.max_buffer_bytes(), kDefaultMaxBufferBytes); + + auto mock_profiler2 = std::make_unique(); + ContinuousProfilerOrchestrator custom_orchestrator( + std::move(mock_profiler2), 500 * 1024 * 1024); + EXPECT_EQ(custom_orchestrator.max_buffer_bytes(), 500 * 1024 * 1024); +} + +TEST(ContinuousProfilerOrchestratorTest, ByteBudgetAccounting) { + auto mock_profiler = std::make_unique(); + MockProfiler* mock = mock_profiler.get(); + + EXPECT_CALL(*mock, Start()).WillOnce(Return(absl::OkStatus())); + EXPECT_CALL(*mock, Stop()).WillOnce(Return(absl::OkStatus())); + + absl::Notification consumed_2; + std::atomic consume_count(0); + EXPECT_CALL(*mock, Consume()) + .WillRepeatedly([&]() -> absl::StatusOr { + int count = ++consume_count; + if (count <= 2) { + if (count == 2 && !consumed_2.HasBeenNotified()) { + consumed_2.Notify(); + } + return ConsumeResult{ + .data = std::any(count), + .estimated_size_bytes = 25 * 1024 * 1024, // 25MB + }; + } + return absl::OutOfRangeError("End of stream"); + }); + + ContinuousProfilerOrchestrator orchestrator( + std::move(mock_profiler), 100 * 1024 * 1024); + EXPECT_EQ(orchestrator.total_buffered_bytes(), 0); + EXPECT_EQ(orchestrator.dropped_chunks_count(), 0); + EXPECT_EQ(orchestrator.dropped_bytes_count(), 0); + + ASSERT_OK(orchestrator.Start()); + consumed_2.WaitForNotification(); + ASSERT_OK(orchestrator.Stop()); + + EXPECT_EQ(orchestrator.total_buffered_bytes(), 50 * 1024 * 1024); + EXPECT_EQ(orchestrator.dropped_chunks_count(), 0); + EXPECT_EQ(orchestrator.dropped_bytes_count(), 0); + + std::vector chunks = orchestrator.PopBuffer(); + EXPECT_EQ(chunks.size(), 2); + EXPECT_EQ(orchestrator.total_buffered_bytes(), 0); +} + +TEST(ContinuousProfilerOrchestratorTest, FIFOEvictionOnMemoryCap) { + auto mock_profiler = std::make_unique(); + MockProfiler* mock = mock_profiler.get(); + + EXPECT_CALL(*mock, Start()).WillOnce(Return(absl::OkStatus())); + EXPECT_CALL(*mock, Stop()).WillOnce(Return(absl::OkStatus())); + + absl::Notification consumed_3; + std::atomic consume_count(0); + EXPECT_CALL(*mock, Consume()) + .WillRepeatedly([&]() -> absl::StatusOr { + int count = ++consume_count; + if (count <= 3) { + if (count == 3 && !consumed_3.HasBeenNotified()) { + consumed_3.Notify(); + } + return ConsumeResult{ + .data = std::any(count), + .estimated_size_bytes = 40 * 1024 * 1024, // 40MB + }; + } + return absl::OutOfRangeError("End of stream"); + }); + + ContinuousProfilerOrchestrator orchestrator( + std::move(mock_profiler), 100 * 1024 * 1024); // 100MB limit + + ASSERT_OK(orchestrator.Start()); + consumed_3.WaitForNotification(); + ASSERT_OK(orchestrator.Stop()); + + // Chunks: 1 (40MB), 2 (40MB), 3 (40MB). Total = 120MB > 100MB limit. + // Chunk 1 is evicted. Retained: chunks 2 & 3 (80MB). + EXPECT_EQ(orchestrator.total_buffered_bytes(), 80 * 1024 * 1024); + EXPECT_EQ(orchestrator.dropped_chunks_count(), 1); + EXPECT_EQ(orchestrator.dropped_bytes_count(), 40 * 1024 * 1024); + + DrainedBuffer drained = orchestrator.PopBufferWithTelemetry(); + ASSERT_EQ(drained.chunks.size(), 2); + EXPECT_EQ(std::any_cast(drained.chunks[0]), 2); + EXPECT_EQ(std::any_cast(drained.chunks[1]), 3); + EXPECT_EQ(drained.cumulative_dropped_chunks, 1); + EXPECT_EQ(drained.cumulative_dropped_bytes, 40 * 1024 * 1024); + + // Buffer is empty after drain, but cumulative drop counters persist + EXPECT_EQ(orchestrator.total_buffered_bytes(), 0); + EXPECT_EQ(orchestrator.dropped_chunks_count(), 1); + EXPECT_EQ(orchestrator.dropped_bytes_count(), 40 * 1024 * 1024); +} + +TEST(ContinuousProfilerOrchestratorTest, OversizedChunkRejection) { + auto mock_profiler = std::make_unique(); + MockProfiler* mock = mock_profiler.get(); + + EXPECT_CALL(*mock, Start()).WillOnce(Return(absl::OkStatus())); + EXPECT_CALL(*mock, Stop()).WillOnce(Return(absl::OkStatus())); + + absl::Notification consumed_1; + std::atomic consume_count(0); + EXPECT_CALL(*mock, Consume()) + .WillRepeatedly([&]() -> absl::StatusOr { + int count = ++consume_count; + if (count == 1) { + consumed_1.Notify(); + return ConsumeResult{ + .data = std::any(count), + .estimated_size_bytes = 60 * 1024 * 1024, // 60MB > 50MB limit + }; + } + return absl::OutOfRangeError("End of stream"); + }); + + ContinuousProfilerOrchestrator orchestrator( + std::move(mock_profiler), 50 * 1024 * 1024); + + ASSERT_OK(orchestrator.Start()); + consumed_1.WaitForNotification(); + ASSERT_OK(orchestrator.Stop()); + + EXPECT_EQ(orchestrator.total_buffered_bytes(), 0); + EXPECT_EQ(orchestrator.dropped_chunks_count(), 1); + EXPECT_EQ(orchestrator.dropped_bytes_count(), 60 * 1024 * 1024); + EXPECT_TRUE(orchestrator.PopBuffer().empty()); +} + +TEST(ContinuousProfilerOrchestratorTest, IdleEmptyDataChunkNoOp) { + auto mock_profiler = std::make_unique(); + MockProfiler* mock = mock_profiler.get(); + + EXPECT_CALL(*mock, Start()).WillOnce(Return(absl::OkStatus())); + EXPECT_CALL(*mock, Stop()).WillOnce(Return(absl::OkStatus())); + + absl::Notification consumed_1; + std::atomic consume_count(0); + EXPECT_CALL(*mock, Consume()) + .WillRepeatedly([&]() -> absl::StatusOr { + int count = ++consume_count; + if (count == 1) { + consumed_1.Notify(); + return ConsumeResult{ + .data = std::any(), // empty data (no value) + .estimated_size_bytes = 0, + }; + } + return absl::OutOfRangeError("End of stream"); + }); + + ContinuousProfilerOrchestrator orchestrator( + std::move(mock_profiler)); + + ASSERT_OK(orchestrator.Start()); + consumed_1.WaitForNotification(); + ASSERT_OK(orchestrator.Stop()); + + EXPECT_EQ(orchestrator.total_buffered_bytes(), 0); + EXPECT_EQ(orchestrator.dropped_chunks_count(), 0); + EXPECT_EQ(orchestrator.dropped_bytes_count(), 0); + EXPECT_TRUE(orchestrator.PopBuffer().empty()); +} + +TEST(ContinuousProfilerOrchestratorTest, ZeroByteChunkWithDataIsBuffered) { + auto mock_profiler = std::make_unique(); + MockProfiler* mock = mock_profiler.get(); + + EXPECT_CALL(*mock, Start()).WillOnce(Return(absl::OkStatus())); + EXPECT_CALL(*mock, Stop()).WillOnce(Return(absl::OkStatus())); + + absl::Notification consumed_1; + std::atomic consume_count(0); + EXPECT_CALL(*mock, Consume()) + .WillRepeatedly([&]() -> absl::StatusOr { + int count = ++consume_count; + if (count == 1) { + consumed_1.Notify(); + return ConsumeResult{ + .data = std::any(42), // valid payload + .estimated_size_bytes = 0, // 0 estimated bytes + }; + } + return absl::OutOfRangeError("End of stream"); + }); + + ContinuousProfilerOrchestrator orchestrator( + std::move(mock_profiler)); + + ASSERT_OK(orchestrator.Start()); + consumed_1.WaitForNotification(); + ASSERT_OK(orchestrator.Stop()); + + EXPECT_EQ(orchestrator.total_buffered_bytes(), 0); + EXPECT_EQ(orchestrator.dropped_chunks_count(), 0); + EXPECT_EQ(orchestrator.dropped_bytes_count(), 0); + std::vector chunks = orchestrator.PopBuffer(); + ASSERT_EQ(chunks.size(), 1); + EXPECT_EQ(std::any_cast(chunks[0]), 42); +} + +TEST(ContinuousProfilerOrchestratorTest, DropCountersAccuracy) { + auto mock_profiler = std::make_unique(); + MockProfiler* mock = mock_profiler.get(); + + EXPECT_CALL(*mock, Start()).WillOnce(Return(absl::OkStatus())); + EXPECT_CALL(*mock, Stop()).WillOnce(Return(absl::OkStatus())); + + absl::Notification consumed_3; + std::atomic consume_count(0); + EXPECT_CALL(*mock, Consume()) + .WillRepeatedly([&]() -> absl::StatusOr { + int count = ++consume_count; + if (count == 1) { + return ConsumeResult{ + .data = std::any(1), + .estimated_size_bytes = 60 * 1024 * 1024, // 60MB + }; + } + if (count == 2) { + return ConsumeResult{ + .data = std::any(2), + .estimated_size_bytes = 120 * 1024 * 1024, // 120MB (rejected) + }; + } + if (count == 3) { + consumed_3.Notify(); + return ConsumeResult{ + .data = std::any(3), + .estimated_size_bytes = + 60 * 1024 * 1024, // 60MB (evicts chunk 1) + }; + } + return absl::OutOfRangeError("End of stream"); + }); + + ContinuousProfilerOrchestrator orchestrator( + std::move(mock_profiler), 100 * 1024 * 1024); + + ASSERT_OK(orchestrator.Start()); + consumed_3.WaitForNotification(); + ASSERT_OK(orchestrator.Stop()); + + EXPECT_EQ(orchestrator.total_buffered_bytes(), 60 * 1024 * 1024); + EXPECT_EQ(orchestrator.dropped_chunks_count(), 2); + EXPECT_EQ(orchestrator.dropped_bytes_count(), (120 + 60) * 1024 * 1024); + + DrainedBuffer drained = orchestrator.PopBufferWithTelemetry(); + ASSERT_EQ(drained.chunks.size(), 1); + EXPECT_EQ(std::any_cast(drained.chunks[0]), 3); + EXPECT_EQ(drained.cumulative_dropped_chunks, 2); + EXPECT_EQ(drained.cumulative_dropped_bytes, 180 * 1024 * 1024); +} + +TEST(ContinuousProfilerOrchestratorTest, StopDrainAccounting) { + auto mock_profiler = std::make_unique(); + MockProfiler* mock = mock_profiler.get(); + + EXPECT_CALL(*mock, Start()).WillOnce(Return(absl::OkStatus())); + EXPECT_CALL(*mock, Stop()).WillOnce(Return(absl::OkStatus())); + + absl::Notification consumed_1; + std::atomic consume_count(0); + EXPECT_CALL(*mock, Consume()) + .WillRepeatedly([&]() -> absl::StatusOr { + int count = ++consume_count; + if (count == 1) { + consumed_1.Notify(); + return ConsumeResult{ + .data = std::any(1), + .estimated_size_bytes = 20 * 1024 * 1024, // 20MB + }; + } + if (count == 2) { + // Returned during Stop() + return ConsumeResult{ + .data = std::any(2), + .estimated_size_bytes = 30 * 1024 * 1024, // 30MB + }; + } + return absl::OutOfRangeError("End of stream"); + }); + + ContinuousProfilerOrchestrator orchestrator( + std::move(mock_profiler), 100 * 1024 * 1024); + + ASSERT_OK(orchestrator.Start()); + consumed_1.WaitForNotification(); + ASSERT_OK(orchestrator.Stop()); + + EXPECT_EQ(orchestrator.total_buffered_bytes(), 50 * 1024 * 1024); + std::vector chunks = orchestrator.PopBuffer(); + ASSERT_EQ(chunks.size(), 2); + EXPECT_EQ(std::any_cast(chunks[0]), 1); + EXPECT_EQ(std::any_cast(chunks[1]), 2); +} + } // namespace } // namespace profiler } // namespace tsl