Allow repeated task in scheduled executor.

PiperOrigin-RevId: 818758380

internal/platform/cancellable_task.h

@@ -17,10 +17,10 @@
 
 #include <utility>
 
-#include "internal/platform/feature_flags.h"
-#include "internal/platform/runnable.h"
 #include "internal/platform/atomic_boolean.h"
+#include "internal/platform/feature_flags.h"
 #include "internal/platform/future.h"
+#include "internal/platform/runnable.h"
 
 namespace nearby {
 
@@ -31,7 +31,10 @@ namespace nearby {
 class CancellableTask {
  public:
   explicit CancellableTask(Runnable&& runnable)
-      : runnable_{std::move(runnable)} {}
+      : CancellableTask(std::move(runnable), /*is_repeated=*/false) {}
+
+  explicit CancellableTask(Runnable&& runnable, bool is_repeated_)
+      : is_repeated_{is_repeated_}, runnable_{std::move(runnable)} {}
 
   /**
    * Try to cancel the task and wait until completion if the task is already
@@ -53,12 +56,17 @@ class CancellableTask {
 
   void operator()() {
     if (started_or_cancelled_.Set(true)) return;
+    finished_ = Future<bool>();
     runnable_();
     finished_.Set(true);
+    if (is_repeated_) {
+      started_or_cancelled_.Set(false);
+    }
   }
 
  private:
-  AtomicBoolean started_or_cancelled_;
+  const bool is_repeated_;
+  AtomicBoolean started_or_cancelled_{false};
   Future<bool> finished_;
   Runnable runnable_;
 };

internal/platform/scheduled_executor.h

@@ -21,8 +21,10 @@
 
 #include "absl/base/thread_annotations.h"
 #include "absl/time/time.h"
+#include "internal/platform/atomic_boolean.h"
 #include "internal/platform/cancelable.h"
 #include "internal/platform/cancellable_task.h"
+#include "internal/platform/implementation/cancelable.h"
 #include "internal/platform/implementation/platform.h"
 #include "internal/platform/implementation/scheduled_executor.h"
 #include "internal/platform/lockable.h"
@@ -82,7 +84,97 @@ class ABSL_LOCKABLE ScheduledExecutor final : public Lockable {
     }
   }
 
+  Cancelable ScheduleRepeatedly(Runnable&& runnable, absl::Duration duration)
+      ABSL_LOCKS_EXCLUDED(mutex_) {
+    {
+      MutexLock lock(&mutex_);
+      if (!impl_) {
+        return Cancelable();
+      }
+    }
+
+    auto cancellable_task = std::make_shared<CancellableTask>(
+        ThreadCheckRunnable(this, std::move(runnable)), /*is_repeated=*/true);
+
+    auto repeated_task_handler =
+        std::make_shared<RepeatedTask>(this, cancellable_task, duration);
+
+    // Start the first execution.
+    repeated_task_handler->Start();
+    return Cancelable(cancellable_task, repeated_task_handler);
+  }
+
  private:
+  // This class encapsulates the state and re-scheduling logic for a single
+  // repeated task. An instance is created for each call to
+  // ScheduleRepeatedly.
+  class RepeatedTask : public api::Cancelable,
+                       public std::enable_shared_from_this<RepeatedTask> {
+   public:
+    RepeatedTask(ScheduledExecutor* executor,
+                 std::shared_ptr<CancellableTask> cancellable_task,
+                 absl::Duration delay)
+        : executor_(executor),
+          cancellable_task_(std::move(cancellable_task)),
+          delay_(delay) {}
+
+    // Starts the first execution of the task.
+    void Start() {
+      MutexLock lock(&executor_->mutex_);
+      ScheduleNextUnderLock();
+    }
+
+    // Implementation of api::Cancelable. This cancels future executions.
+    bool Cancel() override {
+      if (cancelled_.Set(true)) {
+        return true;  // Already cancelled
+      }
+
+      // Cancel the pending scheduled task, if any.
+      MutexLock lock(&executor_->mutex_);
+      if (pending_future_) {
+        pending_future_->Cancel();
+      }
+      return true;
+    }
+
+   private:
+    void Run() {
+      if (cancelled_.Get()) {
+        return;
+      }
+
+      (*cancellable_task_)();
+
+      if (cancelled_.Get()) {
+        return;
+      }
+
+      // Re-schedule the next execution.
+      MutexLock lock(&executor_->mutex_);
+      ScheduleNextUnderLock();
+    }
+
+    void ScheduleNextUnderLock()
+        ABSL_EXCLUSIVE_LOCKS_REQUIRED(executor_->mutex_) {
+      if (cancelled_.Get() || !executor_->impl_) {
+        return;
+      }
+      // Schedule the next run, capturing a shared_ptr to ourself to keep
+      // this object alive.
+      pending_future_ = executor_->impl_->Schedule(
+          [self = shared_from_this()]() { self->Run(); }, delay_);
+    }
+
+    ScheduledExecutor* const executor_;
+    const std::shared_ptr<CancellableTask> cancellable_task_;
+    const absl::Duration delay_;
+
+    AtomicBoolean cancelled_{false};
+    std::shared_ptr<api::Cancelable> pending_future_
+        ABSL_GUARDED_BY(executor_->mutex_);
+  };
+
   void DoShutdown() ABSL_LOCKS_EXCLUDED(mutex_) {
     std::unique_ptr<api::ScheduledExecutor> executor = ReleaseExecutor();
     if (executor) {

internal/platform/scheduled_executor_test.cc

@@ -301,4 +301,152 @@ TEST(ScheduledExecutorTest, ThreadCheck_Schedule) {
   EXPECT_TRUE(notification.WaitForNotificationWithTimeout(absl::Seconds(2)));
 }
 
+TEST(ScheduledExecutorTest, CanScheduleRepeatedly) {
+  constexpr int kNumIterations = 3;
+  ScheduledExecutor executor;
+  std::atomic_int value = 0;
+  CountDownLatch latch(kNumIterations);
+
+  Cancelable cancelable = executor.ScheduleRepeatedly(
+      [&]() {
+        value++;
+        latch.CountDown();
+      },
+      kShortDelay);
+
+  latch.Await();
+  EXPECT_GE(value, kNumIterations);
+  cancelable.Cancel();
+}
+
+TEST(ScheduledExecutorTest, CanCancelRepeatedly) {
+  ScheduledExecutor executor;
+  std::atomic_int value = 0;
+  CountDownLatch latch(1);
+
+  Cancelable cancelable = executor.ScheduleRepeatedly(
+      [&]() {
+        value++;
+        latch.CountDown();
+      },
+      kLongDelay);
+
+  // Wait for the first execution.
+  latch.Await();
+  EXPECT_EQ(value, 1);
+  EXPECT_TRUE(cancelable.Cancel());
+
+  // Wait for a bit to see if it runs again.
+  absl::SleepFor(kLongDelay);
+  EXPECT_EQ(value, 1);
+}
+
+TEST(ScheduledExecutorTest, ShutdownDoesNotRescheduleRepeatedTask) {
+  ScheduledExecutor executor;
+  std::atomic_int value = 0;
+  CountDownLatch latch(1);
+  executor.ScheduleRepeatedly(
+      [&]() {
+        value++;
+        latch.CountDown();
+      },
+      kShortDelay);
+
+  // Wait for first execution to complete.
+  latch.Await();
+  EXPECT_EQ(value, 1);
+
+  executor.Shutdown();
+
+  // After shutdown, the task should not run again.
+  absl::SleepFor(kLongDelay);
+  EXPECT_EQ(value, 1);
+}
+
+TEST(ScheduledExecutorTest, CanCancelOneOfTwoRepeatedTasks) {
+  ScheduledExecutor executor;
+  std::atomic_int valueA = 0;
+  std::atomic_int valueB = 0;
+  CountDownLatch latchA(1);
+  CountDownLatch latchB(1);
+
+  Cancelable cancelableA = executor.ScheduleRepeatedly(
+      [&]() {
+        valueA++;
+        latchA.CountDown();
+      },
+      kShortDelay);
+
+  Cancelable cancelableB = executor.ScheduleRepeatedly(
+      [&]() {
+        valueB++;
+        latchB.CountDown();
+      },
+      kShortDelay);
+
+  // Wait for both to execute once.
+  latchA.Await();
+  latchB.Await();
+  EXPECT_EQ(valueA, 1);
+  EXPECT_EQ(valueB, 1);
+
+  // Cancel the first task.
+  cancelableA.Cancel();
+
+  // Wait for a while and check that only the second task continues to run.
+  absl::SleepFor(kShortDelay * 3);
+  EXPECT_EQ(valueA, 1);
+  EXPECT_GE(valueB, 2);
+
+  cancelableB.Cancel();
+}
+
+TEST(ScheduledExecutorTest, SimulatedClockCanScheduleRepeatedly) {
+  MediumEnvironment::Instance().Start({.use_simulated_clock = true});
+  FakeClock* fake_clock =
+      MediumEnvironment::Instance().GetSimulatedClock().value();
+  ScheduledExecutor executor;
+  std::atomic_int value = 0;
+  std::atomic_int i = 0;
+  CountDownLatch latch[] = {CountDownLatch(1), CountDownLatch(1)};
+
+  Cancelable cancelable = executor.ScheduleRepeatedly(
+      [&]() {
+        value++;
+        latch[i.fetch_add(1)].CountDown();
+      },
+      kShortDelay);
+
+  EXPECT_EQ(value, 0);
+  // Advance to just before the first execution.
+  fake_clock->FastForward(kShortDelay - absl::Milliseconds(1));
+  EXPECT_EQ(value, 0);
+
+  // Advance past the first execution.
+  fake_clock->FastForward(absl::Milliseconds(1));
+  latch[0].Await(absl::Seconds(1));
+  EXPECT_EQ(value, 1);
+
+  // Wait for the second execution to schedule.
+  absl::SleepFor(kShortDelay);
+
+  // Advance to just before the second execution.
+  fake_clock->FastForward(kShortDelay - absl::Milliseconds(1));
+  EXPECT_EQ(value, 1);
+
+  // Advance past the second execution.
+  fake_clock->FastForward(absl::Milliseconds(1));
+  latch[1].Await(absl::Seconds(1));
+  EXPECT_EQ(value, 2);
+
+  // Cancel the task.
+  cancelable.Cancel();
+
+  // Advance a long time and make sure it doesn't run again.
+  fake_clock->FastForward(kLongDelay * 5);
+  EXPECT_EQ(value, 2);
+
+  MediumEnvironment::Instance().Stop();
+}
+
 }  // namespace nearby