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

Summary:
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.

Differential Revision: D96055408

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_++;
 }
 
@@ -994,6 +995,9 @@ void Driver::close() {
   if (!isOnThread() && !isTerminated()) {
     LOG(FATAL) << "Driver::close is only allowed from the Driver's thread";
   }
+  // A driver that completes normally should not have a blocking future.
+  // If it does, it indicates a bug in the driver lifecycle.
+  VELOX_CHECK(!state_.hasBlockingFuture);
   closeOperators();
   updateStats();
   closed_ = true;
@@ -1003,6 +1007,15 @@ void Driver::close() {
 void Driver::closeByTask() {
   VELOX_CHECK(isOnThread());
   VELOX_CHECK(isTerminated());
+  // 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 && state_.blockingStartUs != 0 &&
+      blockedOperatorId_ < operators_.size()) {
+    TestValue::adjust(
+        "facebook::velox::exec::Driver::closeByTask::recordBlockingTime", this);
+    operators_[blockedOperatorId_]->recordBlockingTime(
+        state_.blockingStartUs, blockingReason_);
+  }
   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;

velox/exec/tests/TaskTest.cpp

@@ -3559,4 +3559,87 @@ 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;
+
+  // Track if recordBlockingTime is called in closeByTask via TestValue.
+  std::atomic_bool recordBlockingTimeCalled{false};
+  SCOPED_TESTVALUE_SET(
+      "facebook::velox::exec::Driver::closeByTask::recordBlockingTime",
+      std::function<void(Driver*)>(
+          [&](Driver* /*driver*/) { recordBlockingTimeCalled = true; }));
+
+  // 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 recordBlockingTime was called in closeByTask.
+  // This is the key assertion that validates the fix.
+  EXPECT_TRUE(recordBlockingTimeCalled)
+      << "recordBlockingTime should be called in closeByTask when "
+      << "driver is terminated while blocked";
+
+  // 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