CoroutineHandle

Summary: Extended coroutine handle interface for attaching new features

Reviewed By: capickett

Differential Revision: D31450126

fbshipit-source-id: 2aa6b8a4fb72e5007d08f38fcde5c7abd8eabd67

folly/experimental/coro/Coroutine.h

@@ -166,6 +166,68 @@ class variant_awaitable : private std::variant<A...> {
 
 #endif // __has_include(<variant>)
 
+// Extended promise interface folly::coro types are expected to implement
+class ExtendedCoroutinePromise {
+ public:
+  virtual coroutine_handle<> getHandle() = 0;
+
+ protected:
+  ~ExtendedCoroutinePromise() = default;
+};
+
+template <typename Promise>
+class ExtendedCoroutinePromiseImpl : public ExtendedCoroutinePromise {
+ public:
+  coroutine_handle<> getHandle() final {
+    return coroutine_handle<Promise>::from_promise(
+        *static_cast<Promise*>(this));
+  }
+
+ protected:
+  ~ExtendedCoroutinePromiseImpl() = default;
+};
+
+// Extended version of coroutine_handle<void>
+// Assumes (and enforces) assumption that coroutine_handle is a pointer
+class ExtendedCoroutineHandle {
+ public:
+  template <typename Promise>
+  /*implicit*/ ExtendedCoroutineHandle(
+      coroutine_handle<Promise> handle) noexcept
+      : basic_(handle), extended_(fromBasic(handle)) {}
+
+  /*implicit*/ ExtendedCoroutineHandle(coroutine_handle<> handle) noexcept
+      : basic_(handle) {}
+
+  /*implicit*/ ExtendedCoroutineHandle(ExtendedCoroutinePromise* ptr) noexcept
+      : basic_(ptr->getHandle()), extended_(ptr) {}
+
+  ExtendedCoroutineHandle() noexcept = default;
+
+  void resume() { basic_.resume(); }
+
+  void destroy() { basic_.destroy(); }
+
+  coroutine_handle<> getHandle() const noexcept { return basic_; }
+
+  ExtendedCoroutinePromise* getPromise() const noexcept { return extended_; }
+
+  explicit operator bool() const noexcept { return !!extended_; }
+
+ private:
+  template <typename Promise>
+  static auto fromBasic(coroutine_handle<Promise> handle) noexcept {
+    if constexpr (std::is_convertible_v<Promise*, ExtendedCoroutinePromise*>) {
+      return static_cast<ExtendedCoroutinePromise*>(&handle.promise());
+    } else {
+      return nullptr;
+    }
+  }
+
+  coroutine_handle<> basic_;
+  ExtendedCoroutinePromise* extended_{nullptr};
+};
+
 } // namespace folly::coro
 
 #endif

folly/experimental/coro/Task.h

@@ -76,7 +76,7 @@ class TaskPromiseBase {
       if (promise.ownsAsyncFrame_) {
         folly::popAsyncStackFrameCallee(promise.asyncFrame_);
       }
-      return promise.continuation_;
+      return promise.continuation_.getHandle();
     }
 
     [[noreturn]] void await_resume() noexcept { folly::assume_unreachable(); }
@@ -151,11 +151,11 @@ class TaskPromiseBase {
       coroutine_handle<> continuation) noexcept {
     return {
         std::exchange(p.ownsAsyncFrame_, false),
-        std::exchange(p.continuation_, continuation),
+        std::exchange(p.continuation_, {continuation}).getHandle(),
     };
   }
 
-  coroutine_handle<> continuation_;
+  ExtendedCoroutineHandle continuation_;
   folly::AsyncStackFrame asyncFrame_;
   folly::Executor::KeepAlive<> executor_;
   folly::CancellationToken cancelToken_;
@@ -164,7 +164,8 @@ class TaskPromiseBase {
 };
 
 template <typename T>
-class TaskPromise : public TaskPromiseBase {
+class TaskPromise final : public TaskPromiseBase,
+                          public ExtendedCoroutinePromiseImpl<TaskPromise<T>> {
  public:
   static_assert(
       !std::is_rvalue_reference_v<T>,
@@ -224,7 +225,9 @@ class TaskPromise : public TaskPromiseBase {
 };
 
 template <>
-class TaskPromise<void> : public TaskPromiseBase {
+class TaskPromise<void> final
+    : public TaskPromiseBase,
+      public ExtendedCoroutinePromiseImpl<TaskPromise<void>> {
  public:
   friend class TaskPromiseBase;