diff --git a/cpp/platform/base/feature_flags.h b/cpp/platform/base/feature_flags.h index db58c6fa..ee4a89cd 100644 --- a/cpp/platform/base/feature_flags.h +++ b/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() { diff --git a/cpp/platform/public/BUILD b/cpp/platform/public/BUILD index 6b394eec..24b4878e 100644 --- a/cpp/platform/public/BUILD +++ b/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", diff --git a/cpp/platform/public/cancelable.h b/cpp/platform/public/cancelable.h index ef9fd6ae..89b957b1 100644 --- a/cpp/platform/public/cancelable.h +++ b/cpp/platform/public/cancelable.h @@ -19,6 +19,7 @@ #include #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 impl) - : impl_(std::move(impl)) {} + explicit Cancelable(std::shared_ptr task, + std::shared_ptr 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 task_; std::shared_ptr impl_; }; diff --git a/cpp/platform/public/cancellable_task.h b/cpp/platform/public/cancellable_task.h new file mode 100644 index 00000000..70c6e013 --- /dev/null +++ b/cpp/platform/public/cancellable_task.h @@ -0,0 +1,56 @@ +#ifndef PLATFORM_PUBLIC_CANCELLABLE_TASK_H_ +#define PLATFORM_PUBLIC_CANCELLABLE_TASK_H_ + +#include + +#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 finished_; + Runnable runnable_; +}; + +} // namespace nearby +} // namespace location + +#endif // PLATFORM_PUBLIC_CANCELLABLE_TASK_H_ diff --git a/cpp/platform/public/scheduled_executor.h b/cpp/platform/public/scheduled_executor.h index 4dba913f..4ce64867 100644 --- a/cpp/platform/public/scheduled_executor.h +++ b/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(std::move(runnable)); + return Cancelable(task, + impl_->Schedule([task]() { (*task)(); }, duration)); + } else { + return Cancelable(); + } } private: diff --git a/cpp/platform/public/scheduled_executor_test.cc b/cpp/platform/public/scheduled_executor_test.cc index d7e371ec..f2469373 100644 --- a/cpp/platform/public/scheduled_executor_test.cc +++ b/cpp/platform/public/scheduled_executor_test.cc @@ -18,6 +18,7 @@ #include #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