Add mechanism to store data alongside CancellationState

Summary: Enables implementing coro::Promise with a single allocation. Not intended for direct use by folly users.

Differential Revision: D31561033

fbshipit-source-id: ae9d350694c05710dcf7abd507f11795786da09d

folly/CancellationToken-inl.h

@@ -17,6 +17,7 @@
 #include <algorithm>
 #include <cstdint>
 #include <limits>
+#include <tuple>
 #include <utility>
 
 #include <glog/logging.h>
@@ -116,6 +117,21 @@ class FixedMergingCancellationState : public CancellationState {
   std::array<CancellationCallback, N> callbacks_;
 };
 
+template <typename... Data>
+class CancellationStateWithData : public CancellationState {
+  template <typename... Args>
+  CancellationStateWithData(Args&&... data);
+
+ public:
+  template <typename... Args>
+  FOLLY_NODISCARD static std::
+      pair<CancellationStateSourcePtr, std::tuple<Data...>*>
+      create(Args&&... data);
+
+ private:
+  std::tuple<Data...> data_;
+};
+
 inline void CancellationStateTokenDeleter::operator()(
     CancellationState* state) noexcept {
   state->removeTokenReference();
@@ -370,6 +386,9 @@ inline CancellationStateTokenPtr FixedMergingCancellationState<N>::create(
       new FixedMergingCancellationState<N>(std::forward<Ts>(tokens)...)};
 }
 
+template <typename... Data>
+struct WithDataTag {};
+
 template <size_t N>
 template <typename... Ts>
 inline FixedMergingCancellationState<N>::FixedMergingCancellationState(
@@ -378,8 +397,30 @@ inline FixedMergingCancellationState<N>::FixedMergingCancellationState(
       callbacks_{
           {{std::forward<Ts>(tokens), [this] { requestCancellation(); }}...}} {}
 
+template <typename... Data>
+template <typename... Args>
+CancellationStateWithData<Data...>::CancellationStateWithData(Args&&... data)
+    : data_(std::forward<Args>(data)...) {}
+
+template <typename... Data>
+template <typename... Args>
+std::pair<CancellationStateSourcePtr, std::tuple<Data...>*>
+CancellationStateWithData<Data...>::create(Args&&... data) {
+  auto* state =
+      new CancellationStateWithData<Data...>(std::forward<Args>(data)...);
+  return {CancellationStateSourcePtr{state}, &state->data_};
+}
+
 } // namespace detail
 
+template <typename... Data, typename... Args>
+std::pair<CancellationSource, std::tuple<Data...>*> CancellationSource::create(
+    detail::WithDataTag<Data...>, Args&&... data) {
+  auto [state, dataPtr] = detail::CancellationStateWithData<Data...>::create(
+      std::forward<Args>(data)...);
+  return {CancellationSource{std::move(state)}, dataPtr};
+}
+
 template <typename... Ts>
 inline CancellationToken CancellationToken::merge(Ts&&... tokens) {
   bool canBeCancelled = (tokens.canBeCancelled() || ...);

folly/CancellationToken.h

@@ -46,6 +46,8 @@ using CancellationStateTokenPtr =
     std::unique_ptr<CancellationState, CancellationStateTokenDeleter>;
 using CancellationStateSourcePtr =
     std::unique_ptr<CancellationState, CancellationStateSourceDeleter>;
+template <typename...>
+struct WithDataTag;
 } // namespace detail
 
 // A CancellationToken is an object that can be passed into an function or
@@ -202,6 +204,10 @@ class CancellationSource {
   friend bool operator==(
       const CancellationSource& a, const CancellationSource& b) noexcept;
 
+  template <typename... Data, typename... Args>
+  static std::pair<CancellationSource, std::tuple<Data...>*> create(
+      detail::WithDataTag<Data...>, Args&&...);
+
  private:
   explicit CancellationSource(
       detail::CancellationStateSourcePtr&& state) noexcept;

folly/test/CancellationTokenTest.cpp

@@ -292,3 +292,25 @@ TEST(CancellationTokenTest, MergedToken) {
   token = CancellationToken::merge(CancellationToken());
   EXPECT_FALSE(token.canBeCancelled());
 }
+
+TEST(CancellationTokenTest, TokenWithData) {
+  struct Guard {
+    int& counter;
+    explicit Guard(int& c) : counter(c) {}
+    ~Guard() { ++counter; }
+  };
+  int counter = 0;
+
+  {
+    CancellationToken token;
+    {
+      auto [source, data] =
+          CancellationSource::create(detail::WithDataTag<Guard>{}, counter);
+      EXPECT_EQ(counter, 0);
+      token = source.getToken();
+      EXPECT_EQ(counter, 0);
+    }
+    EXPECT_EQ(counter, 0);
+  }
+  EXPECT_EQ(counter, 1);
+}