src: improve thread safety of TaskQueue

Author: codebytere

src/node_platform.cc

@@ -40,9 +40,9 @@ static void PlatformWorkerThread(void* data) {
     worker_data->platform_workers_ready->Signal(lock);
   }
 
-  while (std::unique_ptr<Task> task = pending_worker_tasks->BlockingPop()) {
+  while (std::unique_ptr<Task> task = pending_worker_tasks->Lock().BlockingPop()) {
     task->Run();
-    pending_worker_tasks->NotifyOfCompletion();
+    pending_worker_tasks->Lock().NotifyOfCompletion();
   }
 }
 
@@ -73,13 +73,15 @@ class WorkerThreadsTaskRunner::DelayedTaskScheduler {
   }
 
   void PostDelayedTask(std::unique_ptr<Task> task, double delay_in_seconds) {
-    tasks_.Push(std::make_unique<ScheduleTask>(this, std::move(task),
-                                               delay_in_seconds));
+    auto locked = tasks_.Lock();
+    locked.Push(std::make_unique<ScheduleTask>(this, std::move(task),
+                                           delay_in_seconds));
     uv_async_send(&flush_tasks_);
   }
 
   void Stop() {
-    tasks_.Push(std::make_unique<StopTask>(this));
+    auto locked = tasks_.Lock();
+    locked.Push(std::make_unique<StopTask>(this));
     uv_async_send(&flush_tasks_);
   }
 
@@ -100,7 +102,8 @@ class WorkerThreadsTaskRunner::DelayedTaskScheduler {
   static void FlushTasks(uv_async_t* flush_tasks) {
     DelayedTaskScheduler* scheduler =
         ContainerOf(&DelayedTaskScheduler::loop_, flush_tasks->loop);
-    while (std::unique_ptr<Task> task = scheduler->tasks_.Pop())
+    auto locked = scheduler->tasks_.Lock();
+    while (std::unique_ptr<Task> task = locked.Pop())
       task->Run();
   }
 
@@ -149,7 +152,7 @@ class WorkerThreadsTaskRunner::DelayedTaskScheduler {
   static void RunTask(uv_timer_t* timer) {
     DelayedTaskScheduler* scheduler =
         ContainerOf(&DelayedTaskScheduler::loop_, timer->loop);
-    scheduler->pending_worker_tasks_->Push(scheduler->TakeTimerTask(timer));
+    scheduler->pending_worker_tasks_->Lock().Push(scheduler->TakeTimerTask(timer));
   }
 
   std::unique_ptr<Task> TakeTimerTask(uv_timer_t* timer) {
@@ -203,7 +206,7 @@ WorkerThreadsTaskRunner::WorkerThreadsTaskRunner(int thread_pool_size) {
 }
 
 void WorkerThreadsTaskRunner::PostTask(std::unique_ptr<Task> task) {
-  pending_worker_tasks_.Push(std::move(task));
+  pending_worker_tasks_.Lock().Push(std::move(task));
 }
 
 void WorkerThreadsTaskRunner::PostDelayedTask(std::unique_ptr<Task> task,
@@ -212,11 +215,11 @@ void WorkerThreadsTaskRunner::PostDelayedTask(std::unique_ptr<Task> task,
 }
 
 void WorkerThreadsTaskRunner::BlockingDrain() {
-  pending_worker_tasks_.BlockingDrain();
+  pending_worker_tasks_.Lock().BlockingDrain();
 }
 
 void WorkerThreadsTaskRunner::Shutdown() {
-  pending_worker_tasks_.Stop();
+  pending_worker_tasks_.Lock().Stop();
   delayed_task_scheduler_->Stop();
   for (size_t i = 0; i < threads_.size(); i++) {
     CHECK_EQ(0, uv_thread_join(threads_[i].get()));
@@ -253,20 +256,22 @@ void PerIsolatePlatformData::PostIdleTaskImpl(
 
 void PerIsolatePlatformData::PostTaskImpl(std::unique_ptr<Task> task,
                                           const v8::SourceLocation& location) {
-  if (flush_tasks_ == nullptr) {
+  // Check if flush_tasks is valid before using it
+  if (!flush_tasks_valid_.load()) {
     // V8 may post tasks during Isolate disposal. In that case, the only
     // sensible path forward is to discard the task.
     return;
   }
-  foreground_tasks_.Push(std::move(task));
+  foreground_tasks_.Lock().Push(std::move(task));
   uv_async_send(flush_tasks_);
 }
 
 void PerIsolatePlatformData::PostDelayedTaskImpl(
     std::unique_ptr<Task> task,
     double delay_in_seconds,
     const v8::SourceLocation& location) {
-  if (flush_tasks_ == nullptr) {
+  // Check if flush_tasks is valid before using it
+  if (!flush_tasks_valid_.load()) {
     // V8 may post tasks during Isolate disposal. In that case, the only
     // sensible path forward is to discard the task.
     return;
@@ -275,7 +280,7 @@ void PerIsolatePlatformData::PostDelayedTaskImpl(
   delayed->task = std::move(task);
   delayed->platform_data = shared_from_this();
   delayed->timeout = delay_in_seconds;
-  foreground_delayed_tasks_.Push(std::move(delayed));
+  foreground_delayed_tasks_.Lock().Push(std::move(delayed));
   uv_async_send(flush_tasks_);
 }
 
@@ -301,15 +306,19 @@ void PerIsolatePlatformData::AddShutdownCallback(void (*callback)(void*),
 }
 
 void PerIsolatePlatformData::Shutdown() {
-  if (flush_tasks_ == nullptr)
+  // First mark the flush_tasks as invalid
+  bool expected = true;
+  if (!flush_tasks_valid_.compare_exchange_strong(expected, false)) {
+    // Already marked invalid
     return;
+  }
 
   // While there should be no V8 tasks in the queues at this point, it is
   // possible that Node.js-internal tasks from e.g. the inspector are still
   // lying around. We clear these queues and ignore the return value,
   // effectively deleting the tasks instead of running them.
-  foreground_delayed_tasks_.PopAll();
-  foreground_tasks_.PopAll();
+  foreground_delayed_tasks_.Lock().PopAll();
+  foreground_tasks_.Lock().PopAll();
   scheduled_delayed_tasks_.clear();
 
   // Both destroying the scheduled_delayed_tasks_ lists and closing
@@ -472,33 +481,36 @@ void NodePlatform::DrainTasks(Isolate* isolate) {
 bool PerIsolatePlatformData::FlushForegroundTasksInternal() {
   bool did_work = false;
 
-  while (std::unique_ptr<DelayedTask> delayed =
-      foreground_delayed_tasks_.Pop()) {
-    did_work = true;
-    uint64_t delay_millis = llround(delayed->timeout * 1000);
-
-    delayed->timer.data = static_cast<void*>(delayed.get());
-    uv_timer_init(loop_, &delayed->timer);
-    // Timers may not guarantee queue ordering of events with the same delay if
-    // the delay is non-zero. This should not be a problem in practice.
-    uv_timer_start(&delayed->timer, RunForegroundTask, delay_millis, 0);
-    uv_unref(reinterpret_cast<uv_handle_t*>(&delayed->timer));
-    uv_handle_count_++;
-
-    scheduled_delayed_tasks_.emplace_back(delayed.release(),
+  {
+    auto locked_tasks = foreground_delayed_tasks_.Lock();
+    while (std::unique_ptr<DelayedTask> delayed = locked_tasks.Pop()) {
+      did_work = true;
+      uint64_t delay_millis = llround(delayed->timeout * 1000);
+
+      delayed->timer.data = static_cast<void*>(delayed.get());
+      uv_timer_init(loop_, &delayed->timer);
+      // Timers may not guarantee queue ordering of events with the same delay if
+      // the delay is non-zero. This should not be a problem in practice.
+      uv_timer_start(&delayed->timer, RunForegroundTask, delay_millis, 0);
+      uv_unref(reinterpret_cast<uv_handle_t*>(&delayed->timer));
+      uv_handle_count_++;
+
+      scheduled_delayed_tasks_.emplace_back(delayed.release(),
                                           [](DelayedTask* delayed) {
-      uv_close(reinterpret_cast<uv_handle_t*>(&delayed->timer),
+        uv_close(reinterpret_cast<uv_handle_t*>(&delayed->timer),
                [](uv_handle_t* handle) {
-        std::unique_ptr<DelayedTask> task {
-            static_cast<DelayedTask*>(handle->data) };
-        task->platform_data->DecreaseHandleCount();
+          std::unique_ptr<DelayedTask> task {
+              static_cast<DelayedTask*>(handle->data) };
+          task->platform_data->DecreaseHandleCount();
+        });
       });
-    });
+    }
   }
+  
   // Move all foreground tasks into a separate queue and flush that queue.
   // This way tasks that are posted while flushing the queue will be run on the
   // next call of FlushForegroundTasksInternal.
-  std::queue<std::unique_ptr<Task>> tasks = foreground_tasks_.PopAll();
+  std::queue<std::unique_ptr<Task>> tasks = foreground_tasks_.Lock().PopAll();
   while (!tasks.empty()) {
     std::unique_ptr<Task> task = std::move(tasks.front());
     tasks.pop();
@@ -593,68 +605,4 @@ TaskQueue<T>::TaskQueue()
     : lock_(), tasks_available_(), tasks_drained_(),
       outstanding_tasks_(0), stopped_(false), task_queue_() { }
 
-template <class T>
-void TaskQueue<T>::Push(std::unique_ptr<T> task) {
-  Mutex::ScopedLock scoped_lock(lock_);
-  outstanding_tasks_++;
-  task_queue_.push(std::move(task));
-  tasks_available_.Signal(scoped_lock);
-}
-
-template <class T>
-std::unique_ptr<T> TaskQueue<T>::Pop() {
-  Mutex::ScopedLock scoped_lock(lock_);
-  if (task_queue_.empty()) {
-    return std::unique_ptr<T>(nullptr);
-  }
-  std::unique_ptr<T> result = std::move(task_queue_.front());
-  task_queue_.pop();
-  return result;
-}
-
-template <class T>
-std::unique_ptr<T> TaskQueue<T>::BlockingPop() {
-  Mutex::ScopedLock scoped_lock(lock_);
-  while (task_queue_.empty() && !stopped_) {
-    tasks_available_.Wait(scoped_lock);
-  }
-  if (stopped_) {
-    return std::unique_ptr<T>(nullptr);
-  }
-  std::unique_ptr<T> result = std::move(task_queue_.front());
-  task_queue_.pop();
-  return result;
-}
-
-template <class T>
-void TaskQueue<T>::NotifyOfCompletion() {
-  Mutex::ScopedLock scoped_lock(lock_);
-  if (--outstanding_tasks_ == 0) {
-    tasks_drained_.Broadcast(scoped_lock);
-  }
-}
-
-template <class T>
-void TaskQueue<T>::BlockingDrain() {
-  Mutex::ScopedLock scoped_lock(lock_);
-  while (outstanding_tasks_ > 0) {
-    tasks_drained_.Wait(scoped_lock);
-  }
-}
-
-template <class T>
-void TaskQueue<T>::Stop() {
-  Mutex::ScopedLock scoped_lock(lock_);
-  stopped_ = true;
-  tasks_available_.Broadcast(scoped_lock);
-}
-
-template <class T>
-std::queue<std::unique_ptr<T>> TaskQueue<T>::PopAll() {
-  Mutex::ScopedLock scoped_lock(lock_);
-  std::queue<std::unique_ptr<T>> result;
-  result.swap(task_queue_);
-  return result;
-}
-
 }  // namespace node

src/node_platform.h

@@ -22,16 +22,72 @@ class PerIsolatePlatformData;
 template <class T>
 class TaskQueue {
  public:
+  class Locked {
+   public:
+    explicit Locked(TaskQueue* queue)
+        : queue_(queue), lock_(queue->lock_) {}
+
+    void Push(std::unique_ptr<T> task) {
+      queue_->outstanding_tasks_++;
+      queue_->task_queue_.push(std::move(task));
+      queue_->tasks_available_.Signal(lock_);
+    }
+
+    std::unique_ptr<T> Pop() {
+      if (queue_->task_queue_.empty()) {
+        return std::unique_ptr<T>(nullptr);
+      }
+      std::unique_ptr<T> result = std::move(queue_->task_queue_.front());
+      queue_->task_queue_.pop();
+      return result;
+    }
+
+    std::unique_ptr<T> BlockingPop() {
+      while (queue_->task_queue_.empty() && !queue_->stopped_) {
+        queue_->tasks_available_.Wait(lock_);
+      }
+      if (queue_->stopped_) {
+        return std::unique_ptr<T>(nullptr);
+      }
+      std::unique_ptr<T> result = std::move(queue_->task_queue_.front());
+      queue_->task_queue_.pop();
+      return result;
+    }
+
+    void NotifyOfCompletion() {
+      if (--queue_->outstanding_tasks_ == 0) {
+        queue_->tasks_drained_.Broadcast(lock_);
+      }
+    }
+
+    void BlockingDrain() {
+      while (queue_->outstanding_tasks_ > 0) {
+        queue_->tasks_drained_.Wait(lock_);
+      }
+    }
+
+    void Stop() {
+      queue_->stopped_ = true;
+      queue_->tasks_available_.Broadcast(lock_);
+    }
+
+    std::queue<std::unique_ptr<T>> PopAll() {
+      std::queue<std::unique_ptr<T>> result;
+      result.swap(queue_->task_queue_);
+      return result;
+    }
+
+   private:
+    TaskQueue* queue_;
+    Mutex::ScopedLock lock_;
+  };
+
   TaskQueue();
   ~TaskQueue() = default;
 
-  void Push(std::unique_ptr<T> task);
-  std::unique_ptr<T> Pop();
-  std::unique_ptr<T> BlockingPop();
-  std::queue<std::unique_ptr<T>> PopAll();
-  void NotifyOfCompletion();
-  void BlockingDrain();
-  void Stop();
+  Locked Lock() {
+    return Locked(this);
+  }
 
  private:
   Mutex lock_;
@@ -98,6 +154,9 @@ class PerIsolatePlatformData
   void RunForegroundTask(std::unique_ptr<v8::Task> task);
   static void RunForegroundTask(uv_timer_t* timer);
 
+  uv_async_t* flush_tasks_ = nullptr;
+  std::atomic<bool> flush_tasks_valid_{true};
+
   struct ShutdownCallback {
     void (*cb)(void*);
     void* data;
@@ -110,7 +169,6 @@ class PerIsolatePlatformData
 
   v8::Isolate* const isolate_;
   uv_loop_t* const loop_;
-  uv_async_t* flush_tasks_ = nullptr;
   TaskQueue<v8::Task> foreground_tasks_;
   TaskQueue<DelayedTask> foreground_delayed_tasks_;