When CancelJob, wait for it to finish if it's running

Summary: Right now, when we CancelJob, we only remove it from the queue. However, in some cases, the job already started, and it might refer to objects that were already destroyed. This diff forces CancelJob to wait for the job to finish if it has started.

Test Plan: Add a unit test.

Reviewers: dhruba, igor, mdcallag

Reviewed By: igor

Subscribers: ben, ngbronson

Differential Revision: https://rockset.phacility.com/D9853

Author: Duc Hieu Pham

cloud/cloud_scheduler.cc

@@ -2,6 +2,7 @@
 #ifndef ROCKSDB_LITE
 #include "cloud/cloud_scheduler.h"
 
+#include <chrono>
 #include <condition_variable>
 #include <mutex>
 #include <thread>
@@ -23,6 +24,8 @@ struct ScheduledJob {
   std::chrono::steady_clock::time_point when;
   std::chrono::microseconds frequency;
   std::function<void(void*)> callback;
+
+  // Caller is responsible for the lifetime of arg.
   void* arg;
 };
 
@@ -49,9 +52,13 @@ class CloudSchedulerImpl : public CloudScheduler {
   long next_id_;
 
   std::mutex mutex_;
-  // Notified when the earliest job to be scheduled has changed.
+  // Notified when the earliest job to be scheduled has changed or
+  // currently_running_ has changed.
   std::condition_variable jobs_changed_cv_;
   std::multiset<ScheduledJob, Comp> scheduled_jobs_;
+  // Whether the job in the front of scheduled_jobs_ is running
+  bool currently_running_;
+
   bool shutting_down_{false};
 
   std::unique_ptr<std::thread> thread_;
@@ -97,14 +104,19 @@ class LocalCloudScheduler : public CloudScheduler {
   // Cancels the job referred to by handle if it is active and associated with
   // this scheduler
   bool CancelJob(long handle) override {
-    std::lock_guard<std::mutex> lk(job_mutex_);
-    const auto& it = jobs_.find(handle);
-    if (it != jobs_.end()) {
-      jobs_.erase(it);
-      return scheduler_->CancelJob(it->second);
-    } else {
-      return false;
+    long internal_job_id = -1;
+    {
+      std::lock_guard<std::mutex> lk(job_mutex_);
+      const auto& it = jobs_.find(handle);
+      if (it != jobs_.end()) {
+        internal_job_id = it->second;
+        jobs_.erase(it);
+      } else {
+        return false;
+      }
     }
+
+    return scheduler_->CancelJob(internal_job_id);
   }
 
  private:
@@ -127,6 +139,7 @@ std::shared_ptr<CloudScheduler> CloudScheduler::Get() {
 
 CloudSchedulerImpl::CloudSchedulerImpl() {
   next_id_ = 1;
+  currently_running_ = false;
   auto lambda = [this]() { DoWork(); };
   thread_.reset(new std::thread(lambda));
 }
@@ -175,7 +188,28 @@ long CloudSchedulerImpl::ScheduleRecurringJob(
 }
 
 bool CloudSchedulerImpl::CancelJob(long id) {
-  std::lock_guard<std::mutex> lk(mutex_);
+  if (id < 0) {
+    return false;
+  }
+
+  std::unique_lock<std::mutex> lk(mutex_);
+  jobs_changed_cv_.wait(lk, [this, id]() {
+    // No job is scheduled. Good to continue.
+    if (scheduled_jobs_.empty()) {
+      return true;
+    }
+
+    // Front of the queue is not this job. That means this job is not running.
+    // We're clear to cancel this job.
+    if (scheduled_jobs_.begin()->id != id) {
+      return true;
+    }
+
+    // Front of the queue is the current job. That means only continue if this
+    // job is not running. If it's still running, wait.
+    return !currently_running_;
+  });
+
   for (auto it = scheduled_jobs_.begin(); it != scheduled_jobs_.end(); ++it) {
     if (it->id == id) {
       bool is_first = (it == scheduled_jobs_.begin());
@@ -205,18 +239,28 @@ void CloudSchedulerImpl::DoWork() {
       jobs_changed_cv_.wait_until(lk, earliest_job_time);
       continue;
     }
-    // invoke the function
+
+    currently_running_ = true;
     lk.unlock();
 
+    // invoke the function
     earliest_job->callback(earliest_job->arg);
 
     lk.lock();
+    // Finished running the job. Remove it from the queue.
+    currently_running_ = false;
+
+    // If this is a recurring job, add back to the queue.
     if (earliest_job->frequency.count() > 0) {
       ScheduledJob new_job = *earliest_job;
       new_job.when = std::chrono::steady_clock::now() + new_job.frequency;
       scheduled_jobs_.emplace(new_job);
     }
     scheduled_jobs_.erase(earliest_job);
+
+    // We might be waiting for the change in currently_running_job_id_ when
+    // cancelling a job.
+    jobs_changed_cv_.notify_all();
   }
 }
 }  // namespace ROCKSDB_NAMESPACE

cloud/cloud_scheduler_test.cc

@@ -2,7 +2,14 @@
 
 #include "cloud/cloud_scheduler.h"
 
+#include <gtest/gtest.h>
+
+#include <atomic>
 #include <chrono>
+#include <condition_variable>
+#include <iostream>
+#include <mutex>
+#include <thread>
 #include <unordered_set>
 
 #include "test_util/testharness.h"
@@ -17,29 +24,20 @@ class CloudSchedulerTest : public testing::Test {
   std::shared_ptr<CloudScheduler> scheduler_;
 };
 
+// This test tests basic scheduling function. There are 2 jobs, job2 is
+// scheduled before job1. We check that job2 is actually run before job1.
 TEST_F(CloudSchedulerTest, TestSchedule) {
-  static int job1 = 1;
-  static int job2 = 1;
-  auto doJob = [](void *arg) { (*(reinterpret_cast<int *>(arg)))++; };
+  std::chrono::steady_clock::time_point p1;
+  std::chrono::steady_clock::time_point p2;
 
-  scheduler_->ScheduleJob(std::chrono::microseconds(100), doJob, &job1);
-  scheduler_->ScheduleJob(std::chrono::microseconds(300), doJob, &job2);
+  auto job1 = [&p1](void *) { p1 = std::chrono::steady_clock::now(); };
+  auto job2 = [&p2](void *) { p2 = std::chrono::steady_clock::now(); };
 
-  usleep(200);
-  ASSERT_EQ(job2, 1);
-  ASSERT_EQ(job1, 2);
-  usleep(200);
-  ASSERT_EQ(job2, 2);
-  ASSERT_EQ(job1, 2);
+  scheduler_->ScheduleJob(std::chrono::milliseconds(300), job1, nullptr);
+  scheduler_->ScheduleJob(std::chrono::milliseconds(100), job2, nullptr);
 
-  scheduler_->ScheduleJob(std::chrono::microseconds(300), doJob, &job1);
-  scheduler_->ScheduleJob(std::chrono::microseconds(100), doJob, &job2);
-  usleep(200);
-  ASSERT_EQ(job1, 2);
-  ASSERT_EQ(job2, 3);
-  usleep(200);
-  ASSERT_EQ(job1, 3);
-  ASSERT_EQ(job2, 3);
+  std::this_thread::sleep_for(std::chrono::milliseconds(400));
+  ASSERT_LT(p2, p1);
 }
 
 TEST_F(CloudSchedulerTest, TestCancel) {
@@ -116,6 +114,45 @@ TEST_F(CloudSchedulerTest, TestMultipleSchedulers) {
   ASSERT_GT(job1, old1);
 }
 
+// This test tests the scenario where a job is started before we attempt to
+// cancel it. It should wait for the job to finish.
+TEST_F(CloudSchedulerTest, TestLongRunningJobCancel) {
+  enum class JobStatus {
+    NOT_STARTED = 0,
+    STARTED = 1,
+    FINISHED = 2,
+  };
+
+  std::mutex statusLock;
+  std::atomic<JobStatus> status{JobStatus::NOT_STARTED};
+  std::condition_variable statusVar;
+
+  auto doJob = [&](void *) {
+    {
+      std::unique_lock<std::mutex> lk(statusLock);
+      status.store(JobStatus::STARTED);
+      statusVar.notify_all();
+    }
+
+    std::this_thread::sleep_for(std::chrono::seconds(1));
+    {
+      std::unique_lock<std::mutex> lk(statusLock);
+      status.store(JobStatus::FINISHED);
+    }
+  };
+
+  auto handle =
+      scheduler_->ScheduleJob(std::chrono::microseconds(0), doJob, nullptr);
+
+  {
+    std::unique_lock<std::mutex> lg(statusLock);
+    statusVar.wait(lg, [&]() { return status == JobStatus::STARTED; });
+  }
+
+  scheduler_->CancelJob(handle);
+  ASSERT_EQ(status.load(), JobStatus::FINISHED);
+}
+
 }  //  namespace ROCKSDB_NAMESPACE
 
 int main(int argc, char **argv) {