fixed scheduled executor bug

connections/implementation/mediums/ble/instant_on_lost_manager_test.cc

@@ -124,6 +124,15 @@ TEST_F(InstantOnLostManagerTest, ShutdownMultipleTimes) {
   EXPECT_FALSE(instant_on_lost_manager().Shutdown());
 }
 
+TEST_F(InstantOnLostManagerTest, ShutdownWhileOnLostAdvertising) {
+  instant_on_lost_manager().OnAdvertisingStarted(
+      std::string(kServiceIdA), ByteArray(kData1.data(), kData1.size()));
+  instant_on_lost_manager().OnAdvertisingStopped(std::string(kServiceIdA));
+  EXPECT_TRUE(instant_on_lost_manager().IsOnLostAdvertisingForTesting());
+
+  EXPECT_TRUE(instant_on_lost_manager().Shutdown());
+}
+
 TEST_F(InstantOnLostManagerTest, AdvertisingOnShutdownManager) {
   instant_on_lost_manager().Shutdown();
   instant_on_lost_manager().OnAdvertisingStarted(

internal/platform/implementation/linux/scheduled_executor.cc

@@ -14,8 +14,9 @@
 
 #include "internal/platform/implementation/linux/scheduled_executor.h"
 
-#include <algorithm>
+#include <chrono>
 #include <memory>
+#include <mutex>
 #include <utility>
 
 #include "absl/time/time.h"
@@ -26,40 +27,80 @@ namespace linux {
 
 ScheduledExecutor::ScheduledExecutor()
     : executor_(std::make_unique<nearby::linux::Executor>()),
+      scheduler_state_(std::make_shared<SchedulerState>()),
+      scheduler_thread_([this]() { RunScheduler(); }),
       shut_down_(false) {}
 
+ScheduledExecutor::~ScheduledExecutor() {
+  if (!shut_down_) {
+    Shutdown();
+  }
+}
+
+bool ScheduledExecutor::ScheduledTask::Cancel() {
+  State expected = State::kPending;
+  if (!state_.compare_exchange_strong(expected, State::kCancelled)) {
+    return false;
+  }
+  if (std::shared_ptr<SchedulerState> state = scheduler_state_.lock()) {
+    state->generation.fetch_add(1);
+    state->condition.notify_one();
+  }
+  return true;
+}
+
+bool ScheduledExecutor::ScheduledTask::TryDispatch() {
+  State expected = State::kPending;
+  return state_.compare_exchange_strong(expected, State::kDispatched);
+}
+
+bool ScheduledExecutor::ScheduledTask::IsCancelled() const {
+  return state_.load() == State::kCancelled;
+}
+
+void ScheduledExecutor::ScheduledTask::Run() {
+  if (task_ != nullptr) {
+    task_();
+  }
+}
+
 // Cancelable is kept both in the executor context, and in the caller context.
 // We want Cancelable to live until both caller and executor are done with it.
 // Exclusive ownership model does not work for this case;
 // using std:shared_ptr<> instead of std::unique_ptr<>.
 std::shared_ptr<api::Cancelable> ScheduledExecutor::Schedule(
-    Runnable &&runnable, absl::Duration duration) {
+    Runnable&& runnable, absl::Duration duration) {
   if (shut_down_) {
-    LOG(ERROR) << __func__
-                       << ": Attempt to Schedule on a shut down executor.";
+    LOG(ERROR) << __func__ << ": Attempt to Schedule on a shut down executor.";
 
     return nullptr;
   }
 
-  // Cleans completed tasks
-  scheduled_tasks_.erase(
-      std::remove_if(
-          scheduled_tasks_.begin(), scheduled_tasks_.end(),
-          [](std::shared_ptr<ScheduledTask> &task) { return task->IsDone(); }),
-      scheduled_tasks_.end());
+  auto deadline = std::chrono::steady_clock::now();
+  if (duration == absl::InfiniteDuration()) {
+    deadline = std::chrono::steady_clock::time_point::max();
+  } else if (duration > absl::ZeroDuration()) {
+    deadline += absl::ToChronoNanoseconds(duration);
+  }
 
-  std::shared_ptr<ScheduledTask> task =
-      std::make_shared<ScheduledTask>(std::move(runnable), duration);
-
-  scheduled_tasks_.push_back(task);
-  executor_->Execute([task]() { task->Start(); });
+  std::shared_ptr<ScheduledTask> task;
+  {
+    std::lock_guard<std::mutex> lock(scheduler_state_->mutex);
+    if (scheduler_state_->shut_down) {
+      return nullptr;
+    }
+    task = std::make_shared<ScheduledTask>(std::move(runnable), deadline,
+                                           next_sequence_++, scheduler_state_);
+    scheduled_tasks_.push(task);
+    scheduler_state_->generation.fetch_add(1);
+  }
+  scheduler_state_->condition.notify_one();
   return task;
 }
 
-void ScheduledExecutor::Execute(Runnable &&runnable) {
+void ScheduledExecutor::Execute(Runnable&& runnable) {
   if (shut_down_) {
-    LOG(ERROR) << __func__
-                       << ": Attempt to Execute on a shut down executor.";
+    LOG(ERROR) << __func__ << ": Attempt to Execute on a shut down executor.";
     return;
   }
 
@@ -67,18 +108,71 @@ void ScheduledExecutor::Execute(Runnable &&runnable) {
 }
 
 void ScheduledExecutor::Shutdown() {
-  if (!shut_down_) {
-    shut_down_ = true;
-    for (auto &task : scheduled_tasks_) {
-      task->Cancel();
+  bool expected = false;
+  if (shut_down_.compare_exchange_strong(expected, true)) {
+    {
+      std::lock_guard<std::mutex> lock(scheduler_state_->mutex);
+      scheduler_state_->shut_down = true;
+      while (!scheduled_tasks_.empty()) {
+        scheduled_tasks_.top()->Cancel();
+        scheduled_tasks_.pop();
+      }
+      scheduler_state_->generation.fetch_add(1);
+    }
+    scheduler_state_->condition.notify_all();
+    if (scheduler_thread_.joinable()) {
+      scheduler_thread_.join();
     }
-
-    scheduled_tasks_.clear();
     executor_->Shutdown();
     return;
   }
-  LOG(ERROR) << __func__
-                     << ": Attempt to Shutdown on a shut down executor.";
+  LOG(ERROR) << __func__ << ": Attempt to Shutdown on a shut down executor.";
+}
+
+void ScheduledExecutor::RunScheduler() {
+  while (true) {
+    std::shared_ptr<ScheduledTask> task;
+    {
+      std::unique_lock<std::mutex> lock(scheduler_state_->mutex);
+      while (!scheduler_state_->shut_down) {
+        if (scheduled_tasks_.empty()) {
+          scheduler_state_->condition.wait(lock);
+          continue;
+        }
+
+        task = scheduled_tasks_.top();
+        if (task->IsCancelled()) {
+          scheduled_tasks_.pop();
+          task.reset();
+          continue;
+        }
+
+        if (task->deadline() > std::chrono::steady_clock::now()) {
+          uint64_t generation = scheduler_state_->generation.load();
+          scheduler_state_->condition.wait_until(
+              lock, task->deadline(), [this, generation]() {
+                return scheduler_state_->shut_down ||
+                       scheduler_state_->generation.load() != generation;
+              });
+          task.reset();
+          continue;
+        }
+
+        scheduled_tasks_.pop();
+        if (!task->TryDispatch()) {
+          task.reset();
+          continue;
+        }
+        break;
+      }
+
+      if (scheduler_state_->shut_down) {
+        return;
+      }
+    }
+
+    executor_->Execute([task = std::move(task)]() { task->Run(); });
+  }
 }
 }  // namespace linux
 }  // namespace nearby

internal/platform/implementation/linux/scheduled_executor.h

@@ -15,11 +15,17 @@
 #ifndef PLATFORM_IMPL_LINUX_SCHEDULED_EXECUTOR_H_
 #define PLATFORM_IMPL_LINUX_SCHEDULED_EXECUTOR_H_
 
+#include <atomic>
+#include <chrono>
+#include <condition_variable>
+#include <cstdint>
 #include <memory>
+#include <mutex>
+#include <queue>
+#include <thread>
 #include <utility>
 #include <vector>
 
-#include "absl/synchronization/notification.h"
 #include "absl/time/time.h"
 #include "internal/platform/implementation/cancelable.h"
 #include "internal/platform/implementation/linux/executor.h"
@@ -38,7 +44,7 @@ class ScheduledExecutor : public api::ScheduledExecutor {
  public:
   ScheduledExecutor();
 
-  ~ScheduledExecutor() override = default;
+  ~ScheduledExecutor() override;
 
   // Cancelable is kept both in the executor context, and in the caller context.
   // We want Cancelable to live until both caller and executor are done with it.
@@ -54,43 +60,62 @@ class ScheduledExecutor : public api::ScheduledExecutor {
   void Shutdown() override;
 
  private:
-  class ScheduledTask : public api::Cancelable {
-   public:
-    explicit ScheduledTask(Runnable&& task, absl::Duration duration)
-        : task_(std::move(task)), duration_(duration) {}
-
-    bool Cancel() override {
-      if (is_executed_ || is_cancelled_) {
-        return false;
-      }
-
-      is_cancelled_ = true;
-      notification_.Notify();
-      return true;
-    };
-
-    void Start() {
-      if (is_executed_ ||
-          notification_.WaitForNotificationWithTimeout(duration_)) {
-        return;
-      }
-
-      is_executed_ = true;
-      task_();
-    }
-
-    bool IsDone() const { return is_cancelled_ || is_executed_; }
-
-   private:
-    Runnable task_;
-    absl::Duration duration_;
-    absl::Notification notification_;
-    bool is_cancelled_ = false;
-    bool is_executed_ = false;
+  struct SchedulerState {
+    std::mutex mutex;
+    std::condition_variable condition;
+    std::atomic<uint64_t> generation = 0;
+    bool shut_down = false;
   };
 
+  class ScheduledTask : public api::Cancelable {
+   public:
+    ScheduledTask(Runnable&& task,
+                  std::chrono::steady_clock::time_point deadline,
+                  uint64_t sequence,
+                  std::weak_ptr<SchedulerState> scheduler_state)
+        : task_(std::move(task)),
+          deadline_(deadline),
+          sequence_(sequence),
+          scheduler_state_(std::move(scheduler_state)) {}
+
+    bool Cancel() override;
+    bool TryDispatch();
+    bool IsCancelled() const;
+    void Run();
+
+    std::chrono::steady_clock::time_point deadline() const { return deadline_; }
+    uint64_t sequence() const { return sequence_; }
+
+   private:
+    enum class State { kPending, kDispatched, kCancelled };
+
+    Runnable task_;
+    const std::chrono::steady_clock::time_point deadline_;
+    const uint64_t sequence_;
+    std::weak_ptr<SchedulerState> scheduler_state_;
+    std::atomic<State> state_{State::kPending};
+  };
+
+  struct ScheduledTaskCompare {
+    bool operator()(const std::shared_ptr<ScheduledTask>& lhs,
+                    const std::shared_ptr<ScheduledTask>& rhs) const {
+      if (lhs->deadline() == rhs->deadline()) {
+        return lhs->sequence() > rhs->sequence();
+      }
+      return lhs->deadline() > rhs->deadline();
+    }
+  };
+
+  void RunScheduler();
+
   std::unique_ptr<nearby::linux::Executor> executor_ = nullptr;
-  std::vector<std::shared_ptr<ScheduledTask>> scheduled_tasks_;
+  std::shared_ptr<SchedulerState> scheduler_state_;
+  std::priority_queue<std::shared_ptr<ScheduledTask>,
+                      std::vector<std::shared_ptr<ScheduledTask>>,
+                      ScheduledTaskCompare>
+      scheduled_tasks_;
+  std::thread scheduler_thread_;
+  uint64_t next_sequence_ = 0;
   std::atomic_bool shut_down_ = false;
 };
 

internal/platform/implementation/linux/scheduled_executor_test.cc

@@ -12,6 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+#include <atomic>
 #include <memory>
 #include <utility>
 
@@ -98,6 +99,36 @@ TEST(ScheduledExecutorTests, ScheduleSucceeds) {
   ASSERT_EQ(output, expected);
 }
 
+TEST(ScheduledExecutorTests, DelayedTaskDoesNotBlockImmediateTask) {
+  absl::Notification notification;
+  ScheduledExecutor executor;
+
+  auto delayed = executor.Schedule([]() {}, absl::Seconds(2));
+  executor.Execute([&notification]() { notification.Notify(); });
+
+  EXPECT_TRUE(
+      notification.WaitForNotificationWithTimeout(absl::Milliseconds(200)));
+  EXPECT_TRUE(delayed->Cancel());
+  executor.Shutdown();
+}
+
+TEST(ScheduledExecutorTests, TasksRunInDeadlineOrder) {
+  absl::Notification short_task_ran;
+  std::atomic_bool long_task_ran = false;
+  ScheduledExecutor executor;
+
+  auto long_task = executor.Schedule(
+      [&long_task_ran]() { long_task_ran = true; }, absl::Seconds(2));
+  executor.Schedule([&short_task_ran]() { short_task_ran.Notify(); },
+                    absl::Milliseconds(20));
+
+  EXPECT_TRUE(
+      short_task_ran.WaitForNotificationWithTimeout(absl::Milliseconds(200)));
+  EXPECT_FALSE(long_task_ran);
+  EXPECT_TRUE(long_task->Cancel());
+  executor.Shutdown();
+}
+
 TEST(ScheduledExecutorTests, CancelSucceeds) {
   absl::Notification notification;
   // Arrange