Make more tests pass TSAN checks

The db_cloud_tests currently do not pass TSAN but the other cloud-specific ones should now pass.

Author: mrambacher

cloud/cloud_scheduler.cc

@@ -46,6 +46,7 @@ class CloudSchedulerImpl : public CloudScheduler {
                             std::function<void(void*)> callback,
                             void* arg) override;
   bool CancelJob(long handle) override;
+  bool IsScheduled(long handle) override;
 
  private:
   void DoWork();
@@ -70,8 +71,14 @@ class LocalCloudScheduler : public CloudScheduler {
  public:
   LocalCloudScheduler(const std::shared_ptr<CloudScheduler>& scheduler,
                       long local_id)
-      : scheduler_(scheduler), next_local_id_(local_id) {}
+      : scheduler_(scheduler),
+        next_local_id_(local_id),
+        shutting_down_(false) {}
   ~LocalCloudScheduler() override {
+    {
+      std::lock_guard<std::mutex> lk(job_mutex_);
+      shutting_down_ = true;
+    }
     for (const auto& job : jobs_) {
       scheduler_->CancelJob(job.second);
     }
@@ -80,6 +87,9 @@ class LocalCloudScheduler : public CloudScheduler {
 
   long ScheduleJob(std::chrono::microseconds when,
                    std::function<void(void*)> callback, void* arg) override {
+    if (shutting_down_) {
+      return -1;
+    }
     std::lock_guard<std::mutex> lk(job_mutex_);
     long local_id = next_local_id_++;
     auto job = [this, local_id, callback](void* a) {
@@ -95,12 +105,37 @@ class LocalCloudScheduler : public CloudScheduler {
                             std::chrono::microseconds frequency,
                             std::function<void(void*)> callback,
                             void* arg) override {
+    if (shutting_down_) {
+      return -1;
+    }
     auto job = scheduler_->ScheduleRecurringJob(when, frequency, callback, arg);
     std::lock_guard<std::mutex> lk(job_mutex_);
     long local_id = next_local_id_++;
     jobs_[local_id] = job;
     return local_id;
   }
+
+  bool IsScheduled(long handle) override {
+    if (shutting_down_) {
+      return false;
+    } else {
+      std::lock_guard<std::mutex> lk(job_mutex_);
+      const auto& it = jobs_.find(handle);
+      if (it == jobs_.end()) {
+        // We do not have the job in our queue.  Return false
+        return false;
+      } else if (scheduler_->IsScheduled(it->second)) {
+        // The job is still scheduled.  Return false
+        return true;
+      } else {
+        // We have the job in our queue but it has already
+        // completed.  Erase from our queue and return false
+        jobs_.erase(it);
+        return false;
+      }
+    }
+  }
+
   // Cancels the job referred to by handle if it is active and associated with
   // this scheduler
   bool CancelJob(long handle) override {
@@ -123,6 +158,7 @@ class LocalCloudScheduler : public CloudScheduler {
   std::mutex job_mutex_;
   std::shared_ptr<CloudScheduler> scheduler_;
   long next_local_id_;
+  bool shutting_down_;
   std::unordered_map<long, long> jobs_;
 };
 
@@ -187,6 +223,22 @@ long CloudSchedulerImpl::ScheduleRecurringJob(
   return id;
 }
 
+bool CloudSchedulerImpl::IsScheduled(long id) {
+  if (id < 0) {
+    return false;
+  } else {
+    std::unique_lock<std::mutex> lk(mutex_);
+    if (!scheduled_jobs_.empty()) {
+      for (const auto& job : scheduled_jobs_) {
+        if (job.id == id) {
+          return true;
+        }
+      }
+    }
+    return false;
+  }
+}
+
 bool CloudSchedulerImpl::CancelJob(long id) {
   if (id < 0) {
     return false;

cloud/cloud_scheduler.h

@@ -37,6 +37,9 @@ class CloudScheduler {
   // was canceled, false otherwise.
   virtual bool CancelJob(long handle) = 0;
 
+  // Returns true if the job represented by handle has not completed,
+  // meaning that it is either currently schedule or actively running
+  virtual bool IsScheduled(long handle) = 0;
   // Returns a new instance of a cloud scheduler.  The caller is responsible
   // for freeing the scheduler when it is no longer required.
   static std::shared_ptr<CloudScheduler> Get();

cloud/cloud_scheduler_test.cc

@@ -23,6 +23,21 @@ class CloudSchedulerTest : public testing::Test {
   ~CloudSchedulerTest() {}
 
   std::shared_ptr<CloudScheduler> scheduler_;
+  void WaitForJobs(const std::vector<long> &jobs, uint32_t delay) {
+    bool running = true;
+    while (running) {
+      running = false;
+      for (const auto &job : jobs) {
+        if (scheduler_->IsScheduled(job)) {
+          running = true;
+          break;
+        }
+      }
+      if (running) {
+        usleep(delay);
+      }
+    }
+  }
 };
 
 // This test tests basic scheduling function. There are 2 jobs, job2 is
@@ -34,10 +49,13 @@ TEST_F(CloudSchedulerTest, TestSchedule) {
   auto job1 = [&p1](void *) { p1 = std::chrono::steady_clock::now(); };
   auto job2 = [&p2](void *) { p2 = std::chrono::steady_clock::now(); };
 
-  scheduler_->ScheduleJob(std::chrono::milliseconds(300), job1, nullptr);
-  scheduler_->ScheduleJob(std::chrono::milliseconds(100), job2, nullptr);
-
-  std::this_thread::sleep_for(std::chrono::milliseconds(400));
+  auto h1 =
+      scheduler_->ScheduleJob(std::chrono::milliseconds(300), job1, nullptr);
+  auto h2 =
+      scheduler_->ScheduleJob(std::chrono::milliseconds(100), job2, nullptr);
+  while (scheduler_->IsScheduled(h1) && scheduler_->IsScheduled(h2)) {
+    std::this_thread::sleep_for(std::chrono::milliseconds(400));
+  }
   ASSERT_LT(p2, p1);
 }
 
@@ -51,16 +69,16 @@ TEST_F(CloudSchedulerTest, TestCancel) {
   auto handle2 =
       scheduler_->ScheduleJob(std::chrono::microseconds(200), doJob, &job2);
   ASSERT_TRUE(scheduler_->CancelJob(handle2));
-  usleep(300);
+  WaitForJobs({handle1, handle2}, 300);
   ASSERT_EQ(job1, 2);
   ASSERT_EQ(job2, 0);
   ASSERT_FALSE(scheduler_->CancelJob(handle1));
   ASSERT_FALSE(scheduler_->CancelJob(handle2));
 }
 
 TEST_F(CloudSchedulerTest, TestRecurring) {
-  int job1 = 1;
-  int job2 = 1;
+  std::atomic<int> job1(1);
+  std::atomic<int> job2(1);
   auto doJob1 = [&job1](void *) { job1++; };
   auto doJob2 = [&job2](void *) { job2++; };
 
@@ -70,22 +88,24 @@ TEST_F(CloudSchedulerTest, TestRecurring) {
   scheduler_->ScheduleRecurringJob(std::chrono::microseconds(120),
                                    std::chrono::microseconds(100), doJob2,
                                    nullptr);
-  usleep(700);
-  ASSERT_GE(job2, 4);
-  ASSERT_GT(job1, job2);
+  while (job2 <= 4) {
+    usleep(100);
+  }
+  ASSERT_GE(job2.load(), 4);
+  ASSERT_GT(job1.load(), job2);
   ASSERT_TRUE(scheduler_->CancelJob(handle1));
-  auto old1 = job1;
-  auto old2 = job2;
+  auto old1 = job1.load();
+  auto old2 = job2.load();
   usleep(200);
-  ASSERT_EQ(job1, old1);
-  ASSERT_GT(job2, old2);
+  ASSERT_EQ(job1.load(), old1);
+  ASSERT_GT(job2.load(), old2);
 }
 
 TEST_F(CloudSchedulerTest, TestMultipleSchedulers) {
   auto scheduler2 = CloudScheduler::Get();
 
-  int job1 = 1;
-  int job2 = 1;
+  std::atomic<int> job1(1);
+  std::atomic<int> job2(1);
   auto doJob1 = [&job1](void *) { job1++; };
   auto doJob2 = [&job2](void *) { job2++; };
 
@@ -108,8 +128,8 @@ TEST_F(CloudSchedulerTest, TestMultipleSchedulers) {
                                    std::chrono::microseconds(20), doJob2,
                                    nullptr);
   scheduler2.reset();
-  auto old1 = job1;
-  auto old2 = job2;
+  auto old1 = job1.load();
+  auto old2 = job2.load();
   usleep(200);
   ASSERT_EQ(job2, old2);
   ASSERT_GT(job1, old1);

cloud/db_cloud_test.cc

@@ -64,8 +64,9 @@ class CloudTest : public testing::Test {
 
     CloudEnv* aenv;
     // create a dummy aws env
-    CloudEnv::NewAwsEnv(base_env_, cloud_env_options_, options_.info_log,
-                        &aenv);
+    ASSERT_OK(CloudEnv::NewAwsEnv(base_env_, cloud_env_options_,
+                                  options_.info_log, &aenv));
+    ASSERT_NE(aenv, nullptr);
     aenv_.reset(aenv);
     // delete all pre-existing contents from the bucket
     Status st = aenv_->GetCloudEnvOptions().storage_provider->EmptyBucket(