Commit 25374a6a authored by Lewis Baker's avatar Lewis Baker Committed by Facebook Github Bot

Add folly::CancellationToken

Summary:
Adds a general-purpose CancellationToken abstraction that can be used to build APIs that allow the caller to pass in a CancellationToken that the caller can later use to communicate a request to cancel the operation.

An operation can either poll for cancellation by calling the isCancellationRequested() method or can register for notification of a cancellation request by attaching a callback to the CancellationToken using the CancellationCallback class.

The caller first constructs a CancellationSource, which allows them to request cancellation, and uses the CancellationSource to obtain CancellationToken objects which it can then pass into cancellable functions.

This implementation is based on the reference implementation for the  interrupt_token/stop_token abstraction proposed for C++20.

```
void polling_operation(folly::CancellationToken ct)
{
  while (!ct.isCancellationRequested())
  {
    do_work();
  }
}

void blocking_operation(folly::CancellationToken ct)
{
  folly::Baton baton;

  // Register a callback.
  folly::CancellationCallback cb{ct, [&] { baton.post(); }};

  // Blocks until cancelled.
  baton.wait();
}

void caller()
{
  CancellationSource src;
  std::thread t1{ [&] {
    polling_operation(src.getToken());
    } };
  std::thread t2{ [&] {
    blocking_operation(src.getToken());
    } };
  std::this_thread::sleep_for(1s);
  src.requestCancellation();
  t1.join();
  t2.join();
}
```

Reviewed By: andriigrynenko

Differential Revision: D10522066

fbshipit-source-id: 11ad3c104eda6650d11081485509981c9b1ea110
parent 51635047
This diff is collapsed.
/*
* Copyright 2018-present Facebook, Inc.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
#include <folly/CancellationToken.h>
#include <folly/Optional.h>
#include <folly/synchronization/detail/Sleeper.h>
#include <glog/logging.h>
#include <algorithm>
#include <new>
#include <thread>
#include <tuple>
namespace folly {
namespace detail {
CancellationState::~CancellationState() {
DCHECK(head_ == nullptr);
DCHECK(!isLocked(state_.load(std::memory_order_relaxed)));
DCHECK(
state_.load(std::memory_order_relaxed) < kTokenReferenceCountIncrement);
}
bool CancellationState::tryAddCallback(
CancellationCallback* callback,
bool incrementRefCountIfSuccessful) noexcept {
// Try to acquire the lock, but abandon trying to acquire the lock if
// cancellation has already been requested (we can just immediately invoke
// the callback) or if cancellation can never be requested (we can just
// skip registration).
if (!tryLock([callback](std::uint64_t oldState) noexcept {
if (isCancellationRequested(oldState)) {
callback->invokeCallback();
return false;
}
return canBeCancelled(oldState);
})) {
return false;
}
// We've acquired the lock and cancellation has not yet been requested.
// Push this callback onto the head of the list.
if (head_ != nullptr) {
head_->prevNext_ = &callback->next_;
}
callback->next_ = head_;
callback->prevNext_ = &head_;
head_ = callback;
if (incrementRefCountIfSuccessful) {
// Combine multiple atomic operations into a single atomic operation.
unlockAndIncrementTokenCount();
} else {
unlock();
}
// Successfully added the callback.
return true;
}
void CancellationState::removeCallback(
CancellationCallback* callback) noexcept {
DCHECK(callback != nullptr);
lock();
if (callback->prevNext_ != nullptr) {
// Still registered in the list => not yet executed.
// Just remove it from the list.
*callback->prevNext_ = callback->next_;
if (callback->next_ != nullptr) {
callback->next_->prevNext_ = callback->prevNext_;
}
unlockAndDecrementTokenCount();
return;
}
unlock();
// Callback has either already executed or is executing concurrently on
// another thread.
if (signallingThreadId_ == std::this_thread::get_id()) {
// Callback executed on this thread or is still currently executing
// and is deregistering itself from within the callback.
if (callback->destructorHasRunInsideCallback_ != nullptr) {
// Currently inside the callback, let the requestCancellation() method
// know the object is about to be destructed and that it should
// not try to access the object when the callback returns.
*callback->destructorHasRunInsideCallback_ = true;
}
} else {
// Callback is currently executing on another thread, block until it
// finishes executing.
folly::detail::Sleeper sleeper;
while (!callback->callbackCompleted_.load(std::memory_order_acquire)) {
sleeper.wait();
}
}
removeTokenReference();
}
bool CancellationState::requestCancellation() noexcept {
if (!tryLockAndCancelUnlessCancelled()) {
// Was already marked as cancelled
return true;
}
// This thread marked as cancelled and acquired the lock
signallingThreadId_ = std::this_thread::get_id();
while (head_ != nullptr) {
// Dequeue the first item on the queue.
CancellationCallback* callback = head_;
head_ = callback->next_;
const bool anyMore = head_ != nullptr;
if (anyMore) {
head_->prevNext_ = &head_;
}
// Mark this item as removed from the list.
callback->prevNext_ = nullptr;
// Don't hold the lock while executing the callback
// as we don't want to block other threads from
// deregistering callbacks.
unlock();
// TRICKY: Need to store a flag on the stack here that the callback
// can use to signal that the destructor was executed inline
// during the call.
// If the destructor was executed inline then it's not safe to
// dereference 'callback' after 'invokeCallback()' returns.
// If the destructor runs on some other thread then the other
// thread will block waiting for this thread to signal that the
// callback has finished executing.
bool destructorHasRunInsideCallback = false;
callback->destructorHasRunInsideCallback_ = &destructorHasRunInsideCallback;
callback->invokeCallback();
if (!destructorHasRunInsideCallback) {
callback->destructorHasRunInsideCallback_ = nullptr;
callback->callbackCompleted_.store(true, std::memory_order_release);
}
if (!anyMore) {
// This was the last item in the queue when we dequeued it.
// No more items should be added to the queue after we have
// marked the state as cancelled, only removed from the queue.
// Avoid acquring/releasing the lock in this case.
return false;
}
lock();
}
unlock();
return false;
}
void CancellationState::lock() noexcept {
folly::detail::Sleeper sleeper;
std::uint64_t oldState = state_.load(std::memory_order_relaxed);
do {
while (isLocked(oldState)) {
sleeper.wait();
oldState = state_.load(std::memory_order_relaxed);
}
} while (!state_.compare_exchange_weak(
oldState,
oldState | kLockedFlag,
std::memory_order_acquire,
std::memory_order_relaxed));
}
void CancellationState::unlock() noexcept {
state_.fetch_sub(kLockedFlag, std::memory_order_release);
}
void CancellationState::unlockAndIncrementTokenCount() noexcept {
state_.fetch_sub(
kLockedFlag - kTokenReferenceCountIncrement, std::memory_order_release);
}
void CancellationState::unlockAndDecrementTokenCount() noexcept {
auto oldState = state_.fetch_sub(
kLockedFlag + kTokenReferenceCountIncrement, std::memory_order_acq_rel);
if (oldState < (kLockedFlag + 2 * kTokenReferenceCountIncrement)) {
delete this;
}
}
bool CancellationState::tryLockAndCancelUnlessCancelled() noexcept {
folly::detail::Sleeper sleeper;
std::uint64_t oldState = state_.load(std::memory_order_acquire);
while (true) {
if (isCancellationRequested(oldState)) {
return false;
} else if (isLocked(oldState)) {
sleeper.wait();
oldState = state_.load(std::memory_order_acquire);
} else if (state_.compare_exchange_weak(
oldState,
oldState | kLockedFlag | kCancellationRequestedFlag,
std::memory_order_acq_rel,
std::memory_order_acquire)) {
return true;
}
}
}
template <typename Predicate>
bool CancellationState::tryLock(Predicate predicate) noexcept {
folly::detail::Sleeper sleeper;
std::uint64_t oldState = state_.load(std::memory_order_acquire);
while (true) {
if (!predicate(oldState)) {
return false;
} else if (isLocked(oldState)) {
sleeper.wait();
oldState = state_.load(std::memory_order_acquire);
} else if (state_.compare_exchange_weak(
oldState,
oldState | kLockedFlag,
std::memory_order_acquire)) {
return true;
}
}
}
} // namespace detail
} // namespace folly
This diff is collapsed.
/*
* Copyright 2018-present Facebook, Inc.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
#include <folly/CancellationToken.h>
#include <folly/Optional.h>
#include <folly/portability/GTest.h>
#include <folly/synchronization/Baton.h>
#include <chrono>
#include <thread>
using namespace folly;
using namespace std::literals::chrono_literals;
TEST(CancellationTokenTest, DefaultCancellationTokenIsNotCancellable) {
CancellationToken t;
EXPECT_FALSE(t.isCancellationRequested());
EXPECT_FALSE(t.canBeCancelled());
CancellationToken tCopy = t;
EXPECT_FALSE(tCopy.isCancellationRequested());
EXPECT_FALSE(tCopy.canBeCancelled());
CancellationToken tMoved = std::move(t);
EXPECT_FALSE(tMoved.isCancellationRequested());
EXPECT_FALSE(tMoved.canBeCancelled());
}
TEST(CancellationTokenTest, Polling) {
CancellationSource src;
EXPECT_FALSE(src.isCancellationRequested());
EXPECT_TRUE(src.canBeCancelled());
CancellationToken token = src.getToken();
EXPECT_FALSE(token.isCancellationRequested());
EXPECT_TRUE(token.canBeCancelled());
CancellationToken tokenCopy = token;
EXPECT_FALSE(tokenCopy.isCancellationRequested());
EXPECT_TRUE(tokenCopy.canBeCancelled());
src.requestCancellation();
EXPECT_TRUE(token.isCancellationRequested());
EXPECT_TRUE(tokenCopy.isCancellationRequested());
}
TEST(CancellationTokenTest, MultiThreadedPolling) {
CancellationSource src;
std::thread t1{[t = src.getToken()] {
while (!t.isCancellationRequested()) {
std::this_thread::yield();
}
}};
src.requestCancellation();
t1.join();
}
TEST(CancellationTokenTest, TokenIsNotCancellableOnceLastSourceIsDestroyed) {
CancellationToken token;
{
CancellationSource src;
token = src.getToken();
{
CancellationSource srcCopy1;
CancellationSource srcCopy2;
EXPECT_TRUE(token.canBeCancelled());
}
EXPECT_TRUE(token.canBeCancelled());
}
EXPECT_FALSE(token.canBeCancelled());
}
TEST(
CancellationTokenTest,
TokenRemainsCancellableEvenOnceLastSourceIsDestroyed) {
CancellationToken token;
{
CancellationSource src;
token = src.getToken();
{
CancellationSource srcCopy1;
CancellationSource srcCopy2;
EXPECT_TRUE(token.canBeCancelled());
}
EXPECT_TRUE(token.canBeCancelled());
src.requestCancellation();
}
EXPECT_TRUE(token.canBeCancelled());
EXPECT_TRUE(token.isCancellationRequested());
}
TEST(CancellationTokenTest, CallbackRegistration) {
CancellationSource src;
bool callbackExecuted = false;
CancellationCallback cb{src.getToken(), [&] { callbackExecuted = true; }};
EXPECT_FALSE(callbackExecuted);
src.requestCancellation();
EXPECT_TRUE(callbackExecuted);
}
TEST(CancellationTokenTest, CallbackExecutesImmediatelyIfAlreadyCancelled) {
CancellationSource src;
src.requestCancellation();
bool callbackExecuted = false;
CancellationCallback cb{src.getToken(), [&] { callbackExecuted = true; }};
EXPECT_TRUE(callbackExecuted);
}
TEST(CancellationTokenTest, CallbackShouldNotBeExecutedMultipleTimes) {
CancellationSource src;
int callbackExecutionCount = 0;
CancellationCallback cb{src.getToken(), [&] { ++callbackExecutionCount; }};
src.requestCancellation();
EXPECT_EQ(1, callbackExecutionCount);
src.requestCancellation();
EXPECT_EQ(1, callbackExecutionCount);
}
TEST(CancellationTokenTest, RegisterMultipleCallbacks) {
CancellationSource src;
bool executed1 = false;
CancellationCallback cb1{src.getToken(), [&] { executed1 = true; }};
bool executed2 = false;
CancellationCallback cb2{src.getToken(), [&] { executed2 = true; }};
EXPECT_FALSE(executed1);
EXPECT_FALSE(executed2);
src.requestCancellation();
EXPECT_TRUE(executed1);
EXPECT_TRUE(executed2);
}
TEST(CancellationTokenTest, DeregisteredCallbacksDontExecute) {
CancellationSource src;
bool executed1 = false;
bool executed2 = false;
CancellationCallback cb1{src.getToken(), [&] { executed1 = true; }};
{
CancellationCallback cb2{src.getToken(), [&] { executed2 = true; }};
}
src.requestCancellation();
EXPECT_TRUE(executed1);
EXPECT_FALSE(executed2);
}
TEST(CancellationTokenTest, CallbackThatDeregistersItself) {
CancellationSource src;
// Check that this doesn't deadlock when a callback tries to deregister
// itself from within the callback.
folly::Optional<CancellationCallback> callback;
callback.emplace(src.getToken(), [&] { callback.clear(); });
src.requestCancellation();
}
TEST(CancellationTokenTest, ManyCallbacks) {
// This test checks that the CancellationSource internal state is able to
// grow to accommodate a large number of callbacks and that there are no
// memory leaks when it's all eventually destroyed.
CancellationSource src;
auto addLotsOfCallbacksAndWait = [t = src.getToken()] {
int counter = 0;
std::vector<std::unique_ptr<CancellationCallback>> callbacks;
for (int i = 0; i < 100; ++i) {
callbacks.push_back(
std::make_unique<CancellationCallback>(t, [&] { ++counter; }));
}
Baton<> baton;
CancellationCallback cb{t, [&] { baton.post(); }};
baton.wait();
};
std::thread t1{addLotsOfCallbacksAndWait};
std::thread t2{addLotsOfCallbacksAndWait};
std::thread t3{addLotsOfCallbacksAndWait};
std::thread t4{addLotsOfCallbacksAndWait};
src.requestCancellation();
t1.join();
t2.join();
t3.join();
t4.join();
}
TEST(CancellationTokenTest, ManyConcurrentCallbackAddRemove) {
auto runTest = [](CancellationToken ct) {
auto cb = [] { std::this_thread::sleep_for(1ms); };
while (!ct.isCancellationRequested()) {
CancellationCallback cb1{ct, cb};
CancellationCallback cb2{ct, cb};
CancellationCallback cb3{ct, cb};
CancellationCallback cb5{ct, cb};
CancellationCallback cb6{ct, cb};
CancellationCallback cb7{ct, cb};
CancellationCallback cb8{ct, cb};
}
};
CancellationSource src;
std::vector<std::thread> threads;
for (int i = 0; i < 10; ++i) {
threads.emplace_back([&, t = src.getToken()] { runTest(t); });
}
std::this_thread::sleep_for(1s);
src.requestCancellation();
for (auto& t : threads) {
t.join();
}
}
Markdown is supported
0%
or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment