fix: Add blockedWaitFor runtime metrics in Driver::closeByTask() (#16711)

Summary:
Pull Request resolved: #16711

When a query fails and some drivers were blocked we fail to submit the
blockedWaitFor runtime metrics.
This deprives us from possibly valuable debugging information, such in cases
with stuck queries.

Reviewed By: xiaoxmeng, vandreykiv, amitkdutta

Differential Revision: D96055408

fbshipit-source-id: 1eb06adfede26857a580f22f400b5bca9dc6c6fa

Author: Sergey Pershin

velox/exec/Driver.cpp

@@ -202,6 +202,7 @@ BlockingState::BlockingState(
               .count()) {
   // Set before leaving the thread.
   driver_->state().hasBlockingFuture = true;
+  driver_->state().blockingStartUs = sinceUs_;
   numBlockedDrivers_++;
 }
 
@@ -897,6 +898,15 @@ void Driver::updateStats() {
   task()->addDriverStats(ctx_->pipelineId, std::move(stats));
 }
 
+void Driver::updateOperatorBlockingStats() {
+  // Record blocked time if the driver was blocked when terminated.
+  // This ensures we don't lose blocked time metrics when a query is aborted.
+  if (state_.hasBlockingFuture && blockedOperatorId_ < operators_.size()) {
+    operators_[blockedOperatorId_]->recordBlockingTime(
+        state_.blockingStartUs, blockingReason_);
+  }
+}
+
 void Driver::startBarrier() {
   VELOX_CHECK(ctx_->task->underBarrier());
   VELOX_CHECK(
@@ -1003,6 +1013,7 @@ void Driver::close() {
 void Driver::closeByTask() {
   VELOX_CHECK(isOnThread());
   VELOX_CHECK(isTerminated());
+  updateOperatorBlockingStats();
   closeOperators();
   updateStats();
   closed_ = true;

velox/exec/Driver.h

@@ -105,6 +105,9 @@ struct ThreadState {
   /// True if there is a future outstanding that will schedule this on an
   /// executor thread when some promise is realized.
   bool hasBlockingFuture{false};
+  /// Timestamp in microseconds when the driver became blocked. Used to record
+  /// blocked time when the driver is terminated while still blocked.
+  uint64_t blockingStartUs{0};
   /// The number of suspension requests on a on-thread driver. If > 0, this
   /// driver thread is in a (recursive) section waiting for RPC or memory
   /// strategy decision. The thread is not supposed to access its memory, which
@@ -170,6 +173,7 @@ struct ThreadState {
     obj["isTerminated"] = isTerminated.load();
     obj["isEnqueued"] = isEnqueued.load();
     obj["hasBlockingFuture"] = hasBlockingFuture;
+    obj["blockingStartUs"] = blockingStartUs;
     obj["isSuspended"] = suspended();
     obj["startExecTime"] = startExecTimeMs;
     return obj;
@@ -595,6 +599,10 @@ class Driver : public std::enable_shared_from_this<Driver> {
 
   void updateStats();
 
+  /// Records operator blocked time ins case the driver was off the thrread and
+  /// blocked when terminated.
+  void updateOperatorBlockingStats();
+
   // Defines the driver barrier processing state.
   struct BarrierState {
     // True if the driver is under barrier processing. This is set by

velox/exec/tests/TaskTest.cpp

@@ -3559,4 +3559,74 @@ DEBUG_ONLY_TEST_F(TaskTest, operatorShouldYieldMethod) {
   }
 }
 
+// Verifies that blocked wait time is recorded even when an operator is
+// terminated while blocked.
+DEBUG_ONLY_TEST_F(TaskTest, blockedWaitTimeOnAbort) {
+  // Test that blocked time is recorded even when a task is aborted while
+  // an operator is blocked.
+  //
+  // We use a simple table scan that blocks waiting for splits. By starting
+  // the task without adding splits, the scan operator will block on
+  // kWaitForSplit. We then abort the task while blocked and verify that
+  // the blocked time is recorded via closeByTask().
+  constexpr int kBlockTimeMs = 100;
+
+  // Build a simple plan with a table scan.
+  core::PlanNodeId scanNodeId;
+  auto plan = PlanBuilder()
+                  .tableScan(ROW({"c0"}, {BIGINT()}))
+                  .capturePlanNodeId(scanNodeId)
+                  .planFragment();
+
+  auto task = Task::create(
+      "blockedWaitTimeOnAbort",
+      plan,
+      0,
+      core::QueryCtx::create(driverExecutor_.get()),
+      Task::ExecutionMode::kParallel);
+
+  task->start(1, 1);
+
+  // Wait for the driver to become blocked waiting for splits.
+  auto startTime = std::chrono::steady_clock::now();
+  while (BlockingState::numBlockedDrivers() == 0) {
+    std::this_thread::sleep_for(std::chrono::milliseconds(10));
+    auto elapsed = std::chrono::duration_cast<std::chrono::seconds>(
+                       std::chrono::steady_clock::now() - startTime)
+                       .count();
+    if (elapsed > 5) {
+      task->requestAbort().wait();
+      GTEST_SKIP() << "Operator did not block in time";
+    }
+  }
+
+  // Let some time pass while blocked to have measurable blocked time.
+  std::this_thread::sleep_for(std::chrono::milliseconds(kBlockTimeMs));
+
+  // Abort the task while the operator is still blocked waiting for splits.
+  task->requestAbort().wait();
+
+  // Wait for the task to be fully aborted.
+  ASSERT_TRUE(waitForTaskAborted(task.get()));
+
+  // Verify that blocked wait time was recorded despite the abort.
+  const auto stats = task->taskStats().pipelineStats;
+  ASSERT_FALSE(stats.empty());
+  ASSERT_FALSE(stats[0].operatorStats.empty());
+
+  // Find operator stats with blocked time recorded.
+  bool foundBlockedTime = false;
+  for (const auto& opStats : stats[0].operatorStats) {
+    if (opStats.blockedWallNanos > 0) {
+      foundBlockedTime = true;
+      // Verify the blocked time is at least what we waited (with tolerance).
+      EXPECT_GE(opStats.blockedWallNanos, (kBlockTimeMs - 30) * 1'000'000)
+          << "Blocked time should be at least " << (kBlockTimeMs - 30) << "ms";
+      break;
+    }
+  }
+  EXPECT_TRUE(foundBlockedTime)
+      << "Operator should have recorded blocked time despite task abort";
+}
+
 } // namespace facebook::velox::exec::test