Merge branch 'google3' up to cl/358328147.

cpp/platform/base/feature_flags.h

@@ -39,6 +39,9 @@ class FeatureFlags {
     // Let endpoint_manager erase deleted endpoint from endpoints_ inside
     // function RemoveEndpoint.
     bool endpoint_manager_ensure_workers_terminated_inside_remove = true;
+    // If a scheduled runnable is already running, Cancel() will synchronously
+    // wait for the task to complete.
+    bool cancel_waits_for_running_tasks = true;
   };
 
   static const FeatureFlags& GetInstance() {

cpp/platform/public/BUILD

@@ -22,6 +22,7 @@ cc_library(
         "atomic_reference.h",
         "cancelable.h",
         "cancelable_alarm.h",
+        "cancellable_task.h",
         "condition_variable.h",
         "count_down_latch.h",
         "crypto.h",
@@ -107,6 +108,7 @@ cc_library(
 cc_test(
     name = "public_test",
     size = "small",
+    timeout = "moderate",
     srcs = [
         "atomic_boolean_test.cc",
         "atomic_reference_test.cc",

cpp/platform/public/cancelable.h

@@ -19,6 +19,7 @@
 #include <utility>
 
 #include "platform/api/cancelable.h"
+#include "platform/public/cancellable_task.h"
 
 namespace location {
 namespace nearby {
@@ -35,14 +36,21 @@ class Cancelable final {
 
   // This constructor is used internally only,
   // by other classes in "//platform/public/".
-  explicit Cancelable(std::shared_ptr<api::Cancelable> impl)
-      : impl_(std::move(impl)) {}
+  explicit Cancelable(std::shared_ptr<CancellableTask> task,
+                      std::shared_ptr<api::Cancelable> impl)
+      : task_{task}, impl_(std::move(impl)) {}
 
-  bool Cancel() { return impl_ ? impl_->Cancel() : false; }
+  bool Cancel() {
+    if (!impl_) return false;
+    bool result = impl_->Cancel();
+    task_->CancelAndWaitIfStarted();
+    return result;
+  }
 
   bool IsValid() { return impl_ != nullptr; }
 
  private:
+  std::shared_ptr<CancellableTask> task_;
   std::shared_ptr<api::Cancelable> impl_;
 };
 

cpp/platform/public/cancellable_task.h

@@ -0,0 +1,56 @@
+#ifndef PLATFORM_PUBLIC_CANCELLABLE_TASK_H_
+#define PLATFORM_PUBLIC_CANCELLABLE_TASK_H_
+
+#include <utility>
+
+#include "platform/base/feature_flags.h"
+#include "platform/base/runnable.h"
+#include "platform/public/atomic_boolean.h"
+#include "platform/public/future.h"
+
+namespace location {
+namespace nearby {
+
+/**
+ * Runnable wrapper that allows one to wait for the task
+ * to complete if it is already running.
+ */
+class CancellableTask {
+ public:
+  explicit CancellableTask(Runnable&& runnable)
+      : runnable_{std::move(runnable)} {}
+
+  /**
+   * Try to cancel the task and wait until completion if the task is already
+   * running.
+   */
+  void CancelAndWaitIfStarted() {
+    if (started_or_cancelled_.Set(true)) {
+      if (FeatureFlags::GetInstance()
+              .GetFlags()
+              .cancel_waits_for_running_tasks) {
+        // task could still be running, wait until finish
+        finished_.Get();
+      }
+    } else {
+      // mark as finished to support multiple calls to this method
+      finished_.Set(true);
+    }
+  }
+
+  void operator()() {
+    if (started_or_cancelled_.Set(true)) return;
+    runnable_();
+    finished_.Set(true);
+  }
+
+ private:
+  AtomicBoolean started_or_cancelled_;
+  Future<bool> finished_;
+  Runnable runnable_;
+};
+
+}  // namespace nearby
+}  // namespace location
+
+#endif  // PLATFORM_PUBLIC_CANCELLABLE_TASK_H_

cpp/platform/public/scheduled_executor.h

@@ -23,6 +23,7 @@
 #include "platform/api/scheduled_executor.h"
 #include "platform/base/runnable.h"
 #include "platform/public/cancelable.h"
+#include "platform/public/cancellable_task.h"
 #include "platform/public/mutex.h"
 #include "platform/public/mutex_lock.h"
 #include "absl/time/time.h"
@@ -73,8 +74,13 @@ class ScheduledExecutor final {
   Cancelable Schedule(Runnable&& runnable, absl::Duration duration)
       ABSL_LOCKS_EXCLUDED(mutex_) {
     MutexLock lock(&mutex_);
-    return impl_ ? Cancelable(impl_->Schedule(std::move(runnable), duration))
-                 : Cancelable();
+    if (impl_) {
+      auto task = std::make_shared<CancellableTask>(std::move(runnable));
+      return Cancelable(task,
+                        impl_->Schedule([task]() { (*task)(); }, duration));
+    } else {
+      return Cancelable();
+    }
   }
 
  private:

cpp/platform/public/scheduled_executor_test.cc

@@ -18,6 +18,7 @@
 #include <functional>
 
 #include "platform/base/exception.h"
+#include "platform/public/count_down_latch.h"
 #include "gtest/gtest.h"
 #include "absl/synchronization/mutex.h"
 #include "absl/time/clock.h"
@@ -95,6 +96,20 @@ TEST(ScheduledExecutorTest, CanCancel) {
   EXPECT_EQ(value, 0);
 }
 
+TEST(ScheduledExecutorTest, CanCancelTwice) {
+  ScheduledExecutor executor;
+  std::atomic_int value = 0;
+  Cancelable cancelable =
+      executor.Schedule([&value]() { value += 1; }, kShortDelay);
+  EXPECT_EQ(value, 0);
+
+  cancelable.Cancel();
+  cancelable.Cancel();
+
+  absl::SleepFor(kLongDelay);
+  EXPECT_EQ(value, 0);
+}
+
 TEST(ScheduledExecutorTest, FailToCancel) {
   absl::Mutex mutex;
   absl::CondVar cond;
@@ -118,5 +133,47 @@ TEST(ScheduledExecutorTest, FailToCancel) {
   EXPECT_EQ(value, 1);
 }
 
+TEST(ScheduledExecutorTest,
+     CancelWhileRunning_TaskCompletesBeforeCancelReturns) {
+  CountDownLatch start_latch(1);
+  ScheduledExecutor executor;
+  std::atomic_int value = 0;
+  // A task that takes a little bit of time to complete
+  Cancelable cancelable = executor.Schedule(
+      [&start_latch, &value]() {
+        start_latch.CountDown();
+        absl::SleepFor(kLongDelay);
+        value += 1;
+      },
+      absl::ZeroDuration());
+
+  start_latch.Await();
+  cancelable.Cancel();
+
+  EXPECT_EQ(value, 1);
+}
+
+TEST(ScheduledExecutorTest,
+     CancelTwiceWhileRunning_TaskCompletesBeforeCancelReturns) {
+  CountDownLatch start_latch(1);
+  ScheduledExecutor executor;
+  std::atomic_int value = 0;
+  // A task that takes a little bit of time to complete
+  Cancelable cancelable = executor.Schedule(
+      [&start_latch, &value]() {
+        start_latch.CountDown();
+        absl::SleepFor(kLongDelay);
+        value += 1;
+      },
+      absl::ZeroDuration());
+
+  start_latch.Await();
+
+  cancelable.Cancel();
+  cancelable.Cancel();
+
+  EXPECT_EQ(value, 1);
+}
+
 }  // namespace nearby
 }  // namespace location