From d8f4d4963bbd9a469ae75abceb2851c53a923332 Mon Sep 17 00:00:00 2001 From: hai007 Date: Fri, 19 Feb 2021 11:32:44 -0800 Subject: [PATCH] 358328147 Synchronous scheduled task cancel --- cpp/platform/base/feature_flags.h | 3 + cpp/platform/public/BUILD | 2 + cpp/platform/public/cancelable.h | 14 ++++- cpp/platform/public/cancellable_task.h | 56 ++++++++++++++++++ cpp/platform/public/scheduled_executor.h | 10 +++- .../public/scheduled_executor_test.cc | 57 +++++++++++++++++++ 6 files changed, 137 insertions(+), 5 deletions(-) create mode 100644 cpp/platform/public/cancellable_task.h diff --git a/cpp/platform/base/feature_flags.h b/cpp/platform/base/feature_flags.h index 6880c1b7..e85342e2 100644 --- a/cpp/platform/base/feature_flags.h +++ b/cpp/platform/base/feature_flags.h @@ -25,6 +25,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 3c1514bb..400afc55 100644 --- a/cpp/platform/public/BUILD +++ b/cpp/platform/public/BUILD @@ -8,6 +8,7 @@ cc_library( "atomic_reference.h", "cancelable.h", "cancelable_alarm.h", + "cancellable_task.h", "condition_variable.h", "count_down_latch.h", "crypto.h", @@ -93,6 +94,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 a565e669..ec0273a1 100644 --- a/cpp/platform/public/cancelable.h +++ b/cpp/platform/public/cancelable.h @@ -5,6 +5,7 @@ #include #include "platform/api/cancelable.h" +#include "platform/public/cancellable_task.h" namespace location { namespace nearby { @@ -21,14 +22,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 6dea9e55..e250b44f 100644 --- a/cpp/platform/public/scheduled_executor.h +++ b/cpp/platform/public/scheduled_executor.h @@ -9,6 +9,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" @@ -59,8 +60,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 95720137..49155c0d 100644 --- a/cpp/platform/public/scheduled_executor_test.cc +++ b/cpp/platform/public/scheduled_executor_test.cc @@ -4,6 +4,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" @@ -81,6 +82,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; @@ -104,5 +119,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