Implement KJ_CO_MAGIC BEGIN_DEFERRED_PROXYING

This commit adds a new coroutine return type, DeferredProxyPromise<T>, which is derived from kj::Promise<DeferredProxy<T>>, and allows us to `KJ_CO_MAGIC BEGIN_DEFERRED_PROXYING` from inside the coroutine to fulfill the outer promise early.

Author: harrishancock

src/workerd/api/deferred-proxy-test.c++

@@ -0,0 +1,171 @@
+#include "util.h"
+#include <kj/test.h>
+
+namespace workerd::api {
+namespace {
+
+KJ_TEST("DeferredProxyPromise<T>: early co_return implicitly fulfills outer promise") {
+  kj::EventLoop loop;
+  kj::WaitScope waitScope(loop);
+
+  {
+    // Implicit void co_return.
+    auto coro = []() -> DeferredProxyPromise<void> {
+      co_await kj::Promise<void>(kj::READY_NOW);
+    };
+    auto promise = coro();
+    KJ_EXPECT(promise.poll(waitScope));
+    auto proxyTask = promise.wait(waitScope).proxyTask;
+    KJ_EXPECT(proxyTask.poll(waitScope));
+    proxyTask.wait(waitScope);
+  }
+  {
+    // Explicit void co_return.
+    auto coro = []() -> DeferredProxyPromise<void> {
+      co_return;
+    };
+    auto promise = coro();
+    KJ_EXPECT(promise.poll(waitScope));
+    auto proxyTask = promise.wait(waitScope).proxyTask;
+    KJ_EXPECT(proxyTask.poll(waitScope));
+    proxyTask.wait(waitScope);
+  }
+  {
+    // Valueful co_return.
+    auto coro = []() -> DeferredProxyPromise<int> {
+      co_return 123;
+    };
+    auto promise = coro();
+    KJ_EXPECT(promise.poll(waitScope));
+    auto proxyTask = promise.wait(waitScope).proxyTask;
+    KJ_EXPECT(proxyTask.poll(waitScope));
+    KJ_EXPECT(proxyTask.wait(waitScope) == 123);
+  }
+}
+
+KJ_TEST("DeferredProxyPromise<T>: `co_yield BEGIN_DEFERRED_PROXYING` fulfills outer promise") {
+  kj::EventLoop loop;
+  kj::WaitScope waitScope(loop);
+
+  auto paf1 = kj::newPromiseAndFulfiller<void>();
+  auto paf2 = kj::newPromiseAndFulfiller<int>();
+
+  auto coro = [&]() -> DeferredProxyPromise<int> {
+    co_await paf1.promise;
+    KJ_CO_MAGIC BEGIN_DEFERRED_PROXYING;
+    co_return co_await paf2.promise;
+  };
+
+  auto promise = coro();
+
+  // paf1 unfulfilled, so we don't have a DeferredProxy<T> yet.
+  KJ_EXPECT(!promise.poll(waitScope));
+
+  paf1.fulfiller->fulfill();
+
+  KJ_EXPECT(promise.poll(waitScope));
+  auto proxyTask = promise.wait(waitScope).proxyTask;
+
+  // paf2 unfulfilled, so we don't have a T yet.
+  KJ_EXPECT(!proxyTask.poll(waitScope));
+
+  paf2.fulfiller->fulfill(123);
+
+  KJ_EXPECT(proxyTask.poll(waitScope));
+  KJ_EXPECT(proxyTask.wait(waitScope) == 123);
+}
+
+KJ_TEST("DeferredProxyPromise<T>: unhandled exception before `co_yield BEGIN_DEFERRED_PROXYING`") {
+  kj::EventLoop loop;
+  kj::WaitScope waitScope(loop);
+
+  auto paf = kj::newPromiseAndFulfiller<void>();
+
+  auto coro = [&]() -> DeferredProxyPromise<int> {
+    co_await paf.promise;
+    KJ_FAIL_ASSERT("promise should have been rejected");
+  };
+
+  auto promise = coro();
+
+  // paf unfulfilled, so we don't have a DeferredProxy<T> yet.
+  KJ_EXPECT(!promise.poll(waitScope));
+
+  paf.fulfiller->reject(KJ_EXCEPTION(FAILED, "test error"));
+
+  KJ_EXPECT(promise.poll(waitScope));
+  KJ_EXPECT_THROW_MESSAGE("test error", promise.wait(waitScope));
+}
+
+KJ_TEST("DeferredProxyPromise<T>: unhandled exception after `co_yield BEGIN_DEFERRED_PROXYING`") {
+  kj::EventLoop loop;
+  kj::WaitScope waitScope(loop);
+
+  auto paf1 = kj::newPromiseAndFulfiller<void>();
+  auto paf2 = kj::newPromiseAndFulfiller<int>();
+
+  auto coro = [&]() -> DeferredProxyPromise<int> {
+    co_await paf1.promise;
+    KJ_CO_MAGIC BEGIN_DEFERRED_PROXYING;
+    co_return co_await paf2.promise;
+  };
+
+  auto promise = coro();
+
+  // paf1 unfulfilled, so we don't have a DeferredProxy<T> yet.
+  KJ_EXPECT(!promise.poll(waitScope));
+
+  paf1.fulfiller->fulfill();
+
+  KJ_EXPECT(promise.poll(waitScope));
+  auto proxyTask = promise.wait(waitScope).proxyTask;
+
+  // paf2 unfulfilled, so we don't have a T yet.
+  KJ_EXPECT(!proxyTask.poll(waitScope));
+
+  paf2.fulfiller->reject(KJ_EXCEPTION(FAILED, "test error"));
+
+  KJ_EXPECT(proxyTask.poll(waitScope));
+  KJ_EXPECT_THROW_MESSAGE("test error", proxyTask.wait(waitScope));
+}
+
+KJ_TEST("DeferredProxyPromise<T>: can be `co_await`ed from another coroutine") {
+  kj::EventLoop loop;
+  kj::WaitScope waitScope(loop);
+
+  auto paf1 = kj::newPromiseAndFulfiller<void>();
+  auto paf2 = kj::newPromiseAndFulfiller<int>();
+
+  auto nestedCoro = [&]() -> DeferredProxyPromise<int> {
+    co_await paf1.promise;
+    KJ_CO_MAGIC BEGIN_DEFERRED_PROXYING;
+    co_return co_await paf2.promise;
+  };
+
+  auto coro = [&]() -> DeferredProxyPromise<int> {
+    auto deferred = co_await nestedCoro();
+    KJ_CO_MAGIC BEGIN_DEFERRED_PROXYING;
+    co_return co_await deferred.proxyTask;
+  };
+
+  auto promise = coro();
+
+  // paf1 unfulfilled, so we don't have a DeferredProxy<T> yet.
+  KJ_EXPECT(!promise.poll(waitScope));
+
+  paf1.fulfiller->fulfill();
+
+  KJ_EXPECT(promise.poll(waitScope));
+  auto proxyTask = promise.wait(waitScope).proxyTask;
+
+  // paf2 unfulfilled, so we don't have a T yet.
+  KJ_EXPECT(!proxyTask.poll(waitScope));
+
+  paf2.fulfiller->fulfill(123);
+
+  KJ_EXPECT(proxyTask.poll(waitScope));
+  KJ_EXPECT(proxyTask.wait(waitScope) == 123);
+}
+
+}  // namespace
+}  // namespace workerd::api

src/workerd/api/util.h

@@ -152,6 +152,127 @@ inline kj::Promise<DeferredProxy<void>> addNoopDeferredProxy(kj::Promise<void> p
   co_return newNoopDeferredProxy();
 }
 
+// ---------------------------------------------------------
+// Deferred proxy coroutine integration
+
+class BeginDeferredProxyingConstant final {};
+constexpr BeginDeferredProxyingConstant BEGIN_DEFERRED_PROXYING {};
+// A magic constant which a DeferredProxyPromise<T> coroutine can `co_yield` to indicate that the
+// deferred proxying phase of its operation has begun.
+
+template <typename T>
+class DeferredProxyPromise: public kj::Promise<DeferredProxy<T>> {
+  // A "strong typedef" for a kj::Promise<DeferredProxy<T>>. DeferredProxyPromise<T> is intended to
+  // be used as the return type for coroutines, in which case the coroutine implementation gains the
+  // following features:
+  //
+  // - `co_yield BEGIN_DEFERRED_PROXYING` fulfills the outer kj::Promise<DeferredProxy<T>>. The
+  //   resulting DeferredProxy<T> object contains a `proxyTask` Promise which owns the coroutine.
+  //
+  // - `co_return` implicitly fulfills the outer Promise for the DeferredProxy<T> (if it has not
+  //   already been fulfilled by the magic `co_yield` described above), then fulfills the inner
+  //   `proxyTask`.
+  //
+  // - Unhandled exceptions reject the outer kj::Promise<DeferredProxy<T>> (if it has not already
+  //   been fulfilled by the magic `co_yield` described above), then reject the inner `proxyTask`.
+
+public:
+  DeferredProxyPromise(kj::Promise<DeferredProxy<T>> promise)
+      : kj::Promise<DeferredProxy<T>>(kj::mv(promise)) {}
+  // Allow conversion from a regular Promise. This allows our `promise_type::get_return_object()`
+  // implementation to be implemented as a regular Promise-returning coroutine.
+
+  class Coroutine;
+  using promise_type = Coroutine;
+  // The coroutine adapter class, required for the compiler to know how to create coroutines
+  // returning DeferredProxyPromise<T>. Since the whole point of DeferredProxyPromise<T> is to serve
+  // as a coroutine return type, there's not really any point hiding promise_type inside of a
+  // coroutine_traits specialization, like kj::Promise<T> does.
+};
+
+template <typename T>
+class DeferredProxyPromise<T>::Coroutine:
+    public kj::_::CoroutineMixin<DeferredProxyPromise<T>::Coroutine, T> {
+  // The coroutine adapter type for DeferredProxyPromise<T>. Most of the work is forwarded to the
+  // regular kj::Promise<T> coroutine adapter.
+
+public:
+  using Handle = kj::_::stdcoro::coroutine_handle<Coroutine>;
+
+  Coroutine(kj::SourceLocation location = {}): inner(Handle::from_promise(*this), location) {}
+
+  kj::Promise<DeferredProxy<T>> get_return_object() {
+    // We need to return a RAII object which will destroy this (as in, `this`) coroutine adapter.
+    // The logic which calls `coroutine_handle<>::destroy()` is tucked away in our inner coroutine
+    // adapter, however, leading to the weird situation where the `inner.get_return_object()`
+    // Promise owns `this`. Thus, we cannot store the inner promise in our own coroutine adapter
+    // class, because that would cause a reference cycle. Fortunately, we can implement our own
+    // `get_return_object()` as a regular Promise-returning coroutine and keep the inner Promise in
+    // our coroutine frame, giving the caller transitive ownership of the DeferredProxyPromise<T>
+    // coroutine by way of the kj::Promise<DeferredProxy<T>> coroutine. Later on, when the outer
+    // Promise is fulfilled, the caller will gain direct ownership of the DeferredProxyPromise<T>
+    // coroutine via the `proxyTask` promise.
+    auto proxyTask = inner.get_return_object();
+    co_await beginDeferredProxying.promise;
+    co_return DeferredProxy<T> { kj::mv(proxyTask) };
+  }
+
+  auto initial_suspend() { return inner.initial_suspend(); }
+  auto final_suspend() noexcept { return inner.final_suspend(); }
+  // Just trivially forward these.
+
+  void unhandled_exception() {
+    // If the outer promise hasn't yet been fulfilled, it needs to be rejected now.
+    if (beginDeferredProxying.fulfiller->isWaiting()) {
+      beginDeferredProxying.fulfiller->reject(kj::getCaughtExceptionAsKj());
+    }
+
+    inner.unhandled_exception();
+  }
+
+  kj::_::stdcoro::suspend_never yield_value(decltype(BEGIN_DEFERRED_PROXYING)) {
+    // This allows us to write `co_yield` within a DeferredProxyPromise<T> coroutine to fulfill
+    // the coroutine's outer promise with a DeferredProxy<T>.
+    // This could alternatively be an await_transform() with a magic parameter type.
+    if (beginDeferredProxying.fulfiller->isWaiting()) {
+      beginDeferredProxying.fulfiller->fulfill();
+    }
+
+    return {};
+  }
+
+  void fulfill(kj::_::FixVoid<T>&& value) {
+    // Required by CoroutineMixin implementation to implement `co_return`.
+
+    // Fulfill the outer promise if it hasn't already been fulfilled.
+    if (beginDeferredProxying.fulfiller->isWaiting()) {
+      beginDeferredProxying.fulfiller->fulfill();
+    }
+
+    inner.fulfill(kj::mv(value));
+  }
+
+  template <typename U>
+  auto await_transform(U&& awaitable) {
+    // Trivially forward everything, so we can await anything a kj::Promise<T> can.
+    return inner.await_transform(kj::fwd<U>(awaitable));
+  }
+
+  operator kj::_::CoroutineBase&() { return inner; }
+  // Required by Awaiter<T>::await_suspend() to support awaiting Promises.
+
+private:
+  typename kj::_::stdcoro::coroutine_traits<kj::Promise<T>>::promise_type inner;
+  // We defer the majority of the implementation to the regular kj::Promise<T> coroutine adapter.
+
+  kj::PromiseFulfillerPair<void> beginDeferredProxying = kj::newPromiseAndFulfiller<void>();
+  // Our `get_return_object()` function returns a kj::Promise<DeferredProxy<T>>, waits on this
+  // `beginDeferredProxying.promise`, then fulfills its Promise with the result of
+  // `inner.get_return_object()`.
+};
+
+// =======================================================================================
+
 kj::Maybe<jsg::V8Ref<v8::Object>> cloneRequestCf(
     jsg::Lock& js, kj::Maybe<jsg::V8Ref<v8::Object>> maybeCf);