Add ContinuousProfilerOrchestrator::Flush() to harvest unconsumed tail events on demand before draining chunks.

tsl/profiler/lib/continuous_profiler_orchestrator.h

@@ -28,6 +28,7 @@ limitations under the License.
 #include "absl/status/status.h"
 #include "absl/status/statusor.h"
 #include "absl/synchronization/mutex.h"
+#include "absl/time/clock.h"
 #include "absl/time/time.h"
 #include "xla/tsl/platform/env.h"
 #include "xla/tsl/platform/errors.h"  // IWYU pragma: keep
@@ -53,12 +54,19 @@ class ContinuousProfilerOrchestrator : public ProfilerInterface {
   static constexpr absl::Duration kDefaultPollingInterval = absl::Seconds(2);
   static constexpr absl::Duration kMinPollingInterval = absl::Seconds(2);
   static constexpr absl::Duration kMaxPollingInterval = absl::Seconds(5);
+  // Upper bound on how long Flush() will block waiting for the ingestion
+  // thread. Generous relative to kMaxPollingInterval (which bounds how long
+  // the ingestion thread sleeps) but small enough that a wedged Consume()
+  // cannot hold an RPC handler hostage.
+  static constexpr absl::Duration kFlushTimeout = absl::Seconds(30);
 
   explicit ContinuousProfilerOrchestrator(
       std::unique_ptr<ProfilerType> profiler,
-      size_t max_buffer_bytes = kDefaultMaxBufferBytes)
+      size_t max_buffer_bytes = kDefaultMaxBufferBytes,
+      absl::Duration flush_timeout = kFlushTimeout)
       : profiler_(std::move(profiler)),
         max_buffer_bytes_(max_buffer_bytes),
+        flush_timeout_(flush_timeout),
         is_running_(false),
         polling_interval_(kDefaultPollingInterval) {}
 
@@ -132,6 +140,41 @@ class ContinuousProfilerOrchestrator : public ProfilerInterface {
     return spaces;
   }
 
+  // Asks the ingestion thread to run a Consume() cycle now and waits for it to
+  // finish, so that everything recorded before this call is in the buffer.
+  // Callers that need the most recent data should call Flush() before
+  // SerializeChunks(); otherwise SerializeChunks() only returns what the
+  // ingestion loop has already harvested, which can lag by up to one polling
+  // interval.
+  //
+  // No-op if the orchestrator is not running. The wait is bounded by the
+  // `flush_timeout` passed to the constructor: a Consume() implementation that
+  // wedges must not be able to block the caller (typically an RPC handler)
+  // indefinitely. On timeout this simply returns, and a following
+  // SerializeChunks() sees whatever was already buffered.
+  //
+  // Note this can cost up to two Consume() cycles in the worst case: a
+  // Consume() that was already in flight when the flush was requested cannot
+  // be credited against it, because it may have started before the data the
+  // caller cares about was recorded.
+  void Flush() {
+    absl::MutexLock lock(mutex_);
+    if (!is_running_) return;
+    uint64_t target_seq = ++flush_requested_seq_;
+    cv_.SignalAll();
+    const absl::Time deadline = absl::Now() + flush_timeout_;
+    while (is_running_ && flush_completed_seq_ < target_seq) {
+      if (cv_.WaitWithDeadline(&mutex_, deadline)) {
+        if (flush_completed_seq_ >= target_seq) break;
+        LOG_EVERY_N_SEC(WARNING, 30)
+            << "ContinuousProfilerOrchestrator::Flush timed out after "
+            << flush_timeout_
+            << " waiting for the ingestion thread; returning partial data.";
+        return;
+      }
+    }
+  }
+
   // Returns the current polling interval (primarily for testing).
   absl::Duration polling_interval() const {
     absl::MutexLock lock(mutex_);
@@ -225,9 +268,11 @@ class ContinuousProfilerOrchestrator : public ProfilerInterface {
   void IngestionLoop() {
     LOG(INFO) << "ContinuousProfilerOrchestrator::IngestionLoop started";
     while (true) {
+      uint64_t seq_before_consume = 0;
       {
         absl::MutexLock lock(mutex_);
         if (!is_running_) break;
+        seq_before_consume = flush_requested_seq_;
       }
       absl::StatusOr<ConsumeResult> result = profiler_->Consume();
 
@@ -238,9 +283,14 @@ class ContinuousProfilerOrchestrator : public ProfilerInterface {
         AdjustIntervalLocked(chunk_size);
       }
 
+      if (seq_before_consume > flush_completed_seq_) {
+        flush_completed_seq_ = seq_before_consume;
+        cv_.SignalAll();
+      }
       if (!is_running_) break;
-
-      // Wait using absl::CondVar on absl::Mutex
+      if (flush_requested_seq_ > flush_completed_seq_) {
+        continue;  // Immediate consume requested while Consume() was in flight.
+      }
       cv_.WaitWithTimeout(&mutex_, polling_interval_);
       if (!is_running_) break;
     }
@@ -272,6 +322,7 @@ class ContinuousProfilerOrchestrator : public ProfilerInterface {
 
   std::unique_ptr<ProfilerType> profiler_;
   const size_t max_buffer_bytes_;
+  const absl::Duration flush_timeout_;
 
   mutable absl::Mutex mutex_;
   absl::CondVar cv_;
@@ -284,6 +335,8 @@ class ContinuousProfilerOrchestrator : public ProfilerInterface {
   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;
+  uint64_t flush_requested_seq_ ABSL_GUARDED_BY(mutex_) = 0;
+  uint64_t flush_completed_seq_ ABSL_GUARDED_BY(mutex_) = 0;
 };
 
 }  // namespace profiler

tsl/profiler/lib/continuous_profiler_orchestrator_test.cc

@@ -17,12 +17,14 @@ limitations under the License.
 #include <any>
 #include <atomic>
 #include <memory>
+#include <thread>  // NOLINT(build/c++11)
 #include <utility>
 #include <vector>
 
 #include "absl/status/status.h"
 #include "absl/status/statusor.h"
 #include "absl/synchronization/notification.h"
+#include "absl/time/clock.h"
 #include "absl/time/time.h"
 #include "xla/tsl/platform/test.h"
 #include "tsl/profiler/lib/profiler_interface.h"
@@ -494,6 +496,168 @@ TEST(ContinuousProfilerOrchestratorTest, StopDrainAccounting) {
   EXPECT_EQ(std::any_cast<int>(chunks[1]), 2);
 }
 
+TEST(ContinuousProfilerOrchestratorTest, FlushForcesAnImmediateConsume) {
+  // The polling interval is >= 2s. Flush() must not wait for it: it asks the
+  // ingestion thread for a Consume() now.
+  auto mock_profiler = std::make_unique<MockProfiler>();
+  MockProfiler* mock = mock_profiler.get();
+
+  EXPECT_CALL(*mock, Start()).WillOnce(Return(absl::OkStatus()));
+  EXPECT_CALL(*mock, Stop()).WillOnce(Return(absl::OkStatus()));
+
+  std::atomic<int> consume_count(0);
+  EXPECT_CALL(*mock, Consume())
+      .WillRepeatedly([&]() -> absl::StatusOr<ConsumeResult> {
+        return ConsumeResult{.data = std::any(++consume_count),
+                             .estimated_size_bytes = 10 * 1024 * 1024};
+      });
+  EXPECT_CALL(*mock, SerializeMock(::testing::_, ::testing::_))
+      .WillRepeatedly(Return(absl::OkStatus()));
+
+  ContinuousProfilerOrchestrator<ProfilerInterface> orchestrator(
+      std::move(mock_profiler));
+
+  ASSERT_OK(orchestrator.Start());
+
+  absl::Time start = absl::Now();
+  orchestrator.Flush();
+  absl::Duration elapsed = absl::Now() - start;
+  std::vector<tensorflow::profiler::XSpace> spaces =
+      orchestrator.SerializeChunks();
+
+  EXPECT_FALSE(spaces.empty());
+  // Two Consume() cycles is the documented worst case (an in-flight Consume()
+  // cannot be credited against the flush), so allow for it but not for a full
+  // polling interval of idling on top.
+  EXPECT_LT(elapsed, absl::Seconds(5))
+      << "Flush() should not have waited out the polling interval";
+  EXPECT_GT(consume_count.load(), 0);
+
+  ASSERT_OK(orchestrator.Stop());
+}
+
+TEST(ContinuousProfilerOrchestratorTest, SerializeChunksDoesNotFlush) {
+  // SerializeChunks() only drains what is already buffered. With a wedged
+  // Consume() and the default (30s) flush timeout, it must still return
+  // right away.
+  auto mock_profiler = std::make_unique<MockProfiler>();
+  MockProfiler* mock = mock_profiler.get();
+
+  EXPECT_CALL(*mock, Start()).WillOnce(Return(absl::OkStatus()));
+  EXPECT_CALL(*mock, Stop()).WillOnce(Return(absl::OkStatus()));
+
+  absl::Notification release_consume;
+  absl::Notification consume_entered;
+  EXPECT_CALL(*mock, Consume())
+      .WillRepeatedly([&]() -> absl::StatusOr<ConsumeResult> {
+        if (!consume_entered.HasBeenNotified()) consume_entered.Notify();
+        release_consume.WaitForNotification();
+        return absl::OutOfRangeError("End of stream");
+      });
+
+  ContinuousProfilerOrchestrator<ProfilerInterface> orchestrator(
+      std::move(mock_profiler));
+  ASSERT_OK(orchestrator.Start());
+  consume_entered.WaitForNotification();
+
+  absl::Time start = absl::Now();
+  std::vector<tensorflow::profiler::XSpace> spaces =
+      orchestrator.SerializeChunks();
+  absl::Duration elapsed = absl::Now() - start;
+
+  EXPECT_TRUE(spaces.empty());
+  EXPECT_LT(elapsed, absl::Seconds(10));
+
+  release_consume.Notify();
+  ASSERT_OK(orchestrator.Stop());
+}
+
+TEST(ContinuousProfilerOrchestratorTest, FlushIsBoundedWhenConsumeWedges) {
+  // A Consume() that never returns must not hold Flush() (and hence the
+  // calling RPC handler) forever.
+  auto mock_profiler = std::make_unique<MockProfiler>();
+  MockProfiler* mock = mock_profiler.get();
+
+  EXPECT_CALL(*mock, Start()).WillOnce(Return(absl::OkStatus()));
+  EXPECT_CALL(*mock, Stop()).WillOnce(Return(absl::OkStatus()));
+
+  absl::Notification release_consume;
+  absl::Notification consume_entered;
+  EXPECT_CALL(*mock, Consume())
+      .WillRepeatedly([&]() -> absl::StatusOr<ConsumeResult> {
+        if (!consume_entered.HasBeenNotified()) consume_entered.Notify();
+        release_consume.WaitForNotification();
+        return absl::OutOfRangeError("End of stream");
+      });
+
+  constexpr absl::Duration kTestFlushTimeout = absl::Milliseconds(200);
+  ContinuousProfilerOrchestrator<ProfilerInterface> orchestrator(
+      std::move(mock_profiler), kDefaultMaxBufferBytes, kTestFlushTimeout);
+  ASSERT_OK(orchestrator.Start());
+  consume_entered.WaitForNotification();
+
+  absl::Time start = absl::Now();
+  orchestrator.Flush();
+  absl::Duration elapsed = absl::Now() - start;
+
+  EXPECT_GE(elapsed, kTestFlushTimeout - absl::Milliseconds(50));
+  EXPECT_LT(elapsed, kTestFlushTimeout + absl::Seconds(5));
+
+  release_consume.Notify();
+  ASSERT_OK(orchestrator.Stop());
+}
+
+TEST(ContinuousProfilerOrchestratorTest, FlushOnAStoppedOrchestratorReturns) {
+  auto mock_profiler = std::make_unique<MockProfiler>();
+  MockProfiler* mock = mock_profiler.get();
+  EXPECT_CALL(*mock, Consume())
+      .WillRepeatedly(Return(absl::OutOfRangeError("End of stream")));
+
+  ContinuousProfilerOrchestrator<ProfilerInterface> orchestrator(
+      std::move(mock_profiler));
+
+  // Never started: Flush() must be a no-op rather than waiting for an ingestion
+  // thread that does not exist.
+  absl::Time start = absl::Now();
+  orchestrator.Flush();
+  EXPECT_LT(absl::Now() - start, absl::Seconds(5));
+  EXPECT_TRUE(orchestrator.SerializeChunks().empty());
+}
+
+TEST(ContinuousProfilerOrchestratorTest, ConcurrentFlushesAllComplete) {
+  auto mock_profiler = std::make_unique<MockProfiler>();
+  MockProfiler* mock = mock_profiler.get();
+
+  EXPECT_CALL(*mock, Start()).WillOnce(Return(absl::OkStatus()));
+  EXPECT_CALL(*mock, Stop()).WillOnce(Return(absl::OkStatus()));
+
+  std::atomic<int> consume_count(0);
+  EXPECT_CALL(*mock, Consume())
+      .WillRepeatedly([&]() -> absl::StatusOr<ConsumeResult> {
+        return ConsumeResult{.data = std::any(++consume_count),
+                             .estimated_size_bytes = 1024};
+      });
+
+  ContinuousProfilerOrchestrator<ProfilerInterface> orchestrator(
+      std::move(mock_profiler));
+  ASSERT_OK(orchestrator.Start());
+
+  // The flush sequence counters are shared; several waiters must all be woken.
+  std::atomic<int> completed(0);
+  std::vector<std::thread> threads;
+  threads.reserve(4);
+  for (int i = 0; i < 4; ++i) {
+    threads.emplace_back([&] {
+      orchestrator.Flush();
+      completed.fetch_add(1);
+    });
+  }
+  for (auto& t : threads) t.join();
+  EXPECT_EQ(completed.load(), 4);
+
+  ASSERT_OK(orchestrator.Stop());
+}
+
 }  // namespace
 }  // namespace profiler
 }  // namespace tsl