Make Task::removeDriver thread safe (#177)

Summary: Pull Request resolved: #177 Reviewed By: zzhao0 Differential Revision: D30822400 fbshipit-source-id: 841867a5793dfeacf9e32a3daef11c7242a10719
facebookincubator · Sep 9, 2021 · ef29d02 · ef29d02
1 parent 703507c
commit ef29d02
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 7 deletions.
diff --git a/velox/exec/Task.cpp b/velox/exec/Task.cpp
@@ -187,14 +187,15 @@ void Task::resume(std::shared_ptr<Task> self) {
 
 // static
 void Task::removeDriver(std::shared_ptr<Task> self, Driver* driver) {
+  std::lock_guard<std::mutex> cancelPoolLock(*self->cancelPool()->mutex());
   for (auto& driverPtr : self->drivers_) {
     if (driverPtr.get() == driver) {
       driverPtr = nullptr;
-      self->driverClosed(driver);
+      self->driverClosed();
       return;
     }
   }
-  VELOX_CHECK(false, "Trying to delete a Driver twice from its Task");
+  VELOX_FAIL("Trying to delete a Driver twice from its Task");
 }
 
 void Task::setMaxSplitSequenceId(
@@ -387,8 +388,7 @@ void Task::setAllOutputConsumed() {
   }
 }
 
-void Task::driverClosed(Driver* /* unused */) {
-  std::lock_guard<std::mutex> cancelPoolLock(*cancelPool()->mutex());
+void Task::driverClosed() {
   --numDrivers_;
   if ((numDrivers_ == 0) && (state_ == kRunning)) {
     std::lock_guard<std::mutex> l(mutex_);

diff --git a/velox/exec/Task.h b/velox/exec/Task.h
@@ -228,8 +228,6 @@ class Task {
     return message;
   }
 
-  void driverClosed(Driver* FOLLY_NONNULL instance);
-
   std::shared_ptr<core::QueryCtx> queryCtx() const {
     return queryCtx_;
   }
@@ -369,7 +367,8 @@ class Task {
     SplitsState& operator=(SplitsState const&) = delete;
   };
 
- private:
+  void driverClosed();
+
   std::shared_ptr<ExchangeClient> addExchangeClient();
 
   void stateChangedLocked();