[native] Add ability to cancel Tasks with stuck Drivers.

Author: spershin

presto-native-execution/presto_cpp/main/TaskManager.cpp

@@ -1098,16 +1098,47 @@ bool TaskManager::getStuckOpCalls(
     std::vector<velox::exec::Task::OpCallInfo>& stuckOpCalls) const {
   const auto thresholdDurationMs =
       SystemConfig::instance()->driverStuckOperatorThresholdMs();
+  const auto thresholdCancelMs =
+      SystemConfig::instance()
+          ->driverCancelTasksWithStuckOperatorsThresholdMs();
+  stuckOpCalls.clear();
+
   const std::chrono::milliseconds lockTimeoutMs(thresholdDurationMs);
   auto taskMap = taskMap_.rlock(lockTimeoutMs);
   if (!taskMap) {
     return false;
   }
-  for (const auto& [_, prestoTask] : *taskMap) {
-    if (prestoTask->task != nullptr &&
-        !prestoTask->task->getLongRunningOpCalls(
-            lockTimeoutMs, thresholdDurationMs, stuckOpCalls)) {
-      deadlockTasks.push_back(prestoTask->task->taskId());
+
+  for (const auto& [id, prestoTask] : *taskMap) {
+    if (prestoTask->task != nullptr) {
+      const auto numPrevStuckOps = stuckOpCalls.size();
+      if (!prestoTask->task->getLongRunningOpCalls(
+              lockTimeoutMs, thresholdDurationMs, stuckOpCalls)) {
+        deadlockTasks.push_back(id);
+        continue;
+      }
+      // See if we need to cancel the Task - it should be running, the cancel
+      // threshold should be valid and it should have at least one stuck driver
+      // that was stuck for enough time.
+      if (numPrevStuckOps < stuckOpCalls.size() && thresholdCancelMs != 0 &&
+          prestoTask->task->isRunning()) {
+        for (auto it = stuckOpCalls.begin() + numPrevStuckOps;
+             it != stuckOpCalls.end();
+             ++it) {
+          if (it->durationMs >= thresholdCancelMs) {
+            std::stringstream ss;
+            ss << "Task " << id
+               << " cancelled due to stuck operator: tid=" << it->tid
+               << " opCall=" << it->opCall
+               << " duration= " << velox::succinctMillis(it->durationMs);
+            const std::string msg = ss.str();
+            LOG(ERROR) << msg;
+            prestoTask->task->setError(msg);
+            RECORD_METRIC_VALUE(kCounterNumCancelledTasksByStuckDriver, 1);
+            break;
+          }
+        }
+      }
     }
   }
   return true;

presto-native-execution/presto_cpp/main/common/Configs.cpp

@@ -157,6 +157,8 @@ SystemConfig::SystemConfig() {
           NUM_PROP(kDriverNumCpuThreadsHwMultiplier, 4.0),
           BOOL_PROP(kDriverThreadsBatchSchedulingEnabled, false),
           NUM_PROP(kDriverStuckOperatorThresholdMs, 30 * 60 * 1000),
+          NUM_PROP(
+              kDriverCancelTasksWithStuckOperatorsThresholdMs, 40 * 60 * 1000),
           NUM_PROP(kSpillerNumCpuThreadsHwMultiplier, 1.0),
           STR_PROP(kSpillerFileCreateConfig, ""),
           NONE_PROP(kSpillerSpillPath),
@@ -377,6 +379,12 @@ size_t SystemConfig::driverStuckOperatorThresholdMs() const {
   return optionalProperty<size_t>(kDriverStuckOperatorThresholdMs).value();
 }
 
+size_t SystemConfig::driverCancelTasksWithStuckOperatorsThresholdMs() const {
+  return optionalProperty<size_t>(
+             kDriverCancelTasksWithStuckOperatorsThresholdMs)
+      .value();
+}
+
 double SystemConfig::spillerNumCpuThreadsHwMultiplier() const {
   return optionalProperty<double>(kSpillerNumCpuThreadsHwMultiplier).value();
 }

presto-native-execution/presto_cpp/main/common/Configs.h

@@ -235,6 +235,13 @@ class SystemConfig : public ConfigBase {
   static constexpr std::string_view kDriverStuckOperatorThresholdMs{
       "driver.stuck-operator-threshold-ms"};
 
+  /// Immediately cancels any Task when it is detected that it has at least one
+  /// stuck Operator for at least the time specified by this threshold.
+  /// Use zero to disable canceling.
+  static constexpr std::string_view
+      kDriverCancelTasksWithStuckOperatorsThresholdMs{
+          "driver.cancel-tasks-with-stuck-operators-threshold-ms"};
+
   /// Floating point number used in calculating how many threads we would use
   /// for Spiller CPU executor: hw_concurrency x multiplier.
   /// If 0.0 then spilling is disabled.
@@ -705,6 +712,8 @@ class SystemConfig : public ConfigBase {
 
   size_t driverStuckOperatorThresholdMs() const;
 
+  size_t driverCancelTasksWithStuckOperatorsThresholdMs() const;
+
   double spillerNumCpuThreadsHwMultiplier() const;
 
   std::string spillerFileCreateConfig() const;

presto-native-execution/presto_cpp/main/common/Counters.cpp

@@ -44,6 +44,8 @@ void registerPrestoMetrics() {
   DEFINE_METRIC(kCounterNumZombiePrestoTasks, facebook::velox::StatType::AVG);
   DEFINE_METRIC(
       kCounterNumTasksWithStuckOperator, facebook::velox::StatType::AVG);
+  DEFINE_METRIC(
+      kCounterNumCancelledTasksByStuckDriver, facebook::velox::StatType::COUNT);
   DEFINE_METRIC(kCounterNumTasksDeadlock, facebook::velox::StatType::AVG);
   DEFINE_METRIC(
       kCounterNumTaskManagerLockTimeOut, facebook::velox::StatType::AVG);

presto-native-execution/presto_cpp/main/common/Counters.h

@@ -67,6 +67,8 @@ constexpr folly::StringPiece kCounterNumZombiePrestoTasks{
     "presto_cpp.num_zombie_presto_tasks"};
 constexpr folly::StringPiece kCounterNumTasksWithStuckOperator{
     "presto_cpp.num_tasks_with_stuck_operator"};
+constexpr folly::StringPiece kCounterNumCancelledTasksByStuckDriver{
+    "presto_cpp.num_cancelled_tasks_by_stuck_driver"};
 constexpr folly::StringPiece kCounterNumTasksDeadlock{
     "presto_cpp.num_tasks_deadlock"};
 constexpr folly::StringPiece kCounterNumTaskManagerLockTimeOut{