Replace DeferredProxyPromise<T> with kj::Promise<DeferredProxy<T>>

Author: harrishancock

src/workerd/api/blob.c++

@@ -183,8 +183,8 @@ public:
     }
 
     // We can't defer the write to the proxy stage since it depends on `blob` which lives in the
-    // isolate.
-    co_return newNoopDeferredProxy();
+    // isolate, so we don't `KJ_CO_MAGIC BEGIN_DEFERRED_PROXYING`.
+    co_return;
   }
 
 private:

src/workerd/api/cache.c++

@@ -332,9 +332,10 @@ jsg::Promise<void> Cache::put(jsg::Lock& js, Request::Info requestOrUrl,
           //   destruction timing in the event of an error; see below.
 
       // The next step is a bit complicated as it occurs in two separate async flows.
-      // First, we await the serialization promise, then construct a DeferredProxy.
-      // That DeferredProxy contains the second async flow that actually handles the
-      // request and response.
+      // First, we await the serialization promise, then enter "deferred proxying" by issuing
+      // `KJ_CO_MAGIC BEGIN_DEFERRED_PROXYING` from our coroutine. Everything after that
+      // `KJ_CO_MAGIC` constitutes the second async flow that actually handles the request and
+      // response.
       //
       // Weird: It's important that these objects be torn down in the right order and that the
       // DeferredProxy promise is handled separately from the inner promise.
@@ -385,62 +386,9 @@ jsg::Promise<void> Cache::put(jsg::Lock& js, Request::Info requestOrUrl,
       // will properly treat it as pending I/O, but only *after* the outer promise completes. This
       // gets us everything we want.
       //
-      // Hence, what you see here: we first await the serializePromise then add all our
-      // additional work onto the inner "deferred proxy" promise. Then we `awaitDeferredProxy()`
-      // the whole thing.
-
-      // We use the same exception handling logic for both the promise for the DeferredProxy
-      // and the inner promise handling the request/response.
-      static auto constexpr handleException = [](kj::Exception&& exception) {
-        if (exception.getType() != kj::Exception::Type::DISCONNECTED) {
-          kj::throwFatalException(kj::mv(exception));
-        }
-        // If the origin or the cache disconnected, we don't treat this as an error, as put()
-        // doesn't guarantee that it stores anything anyway.
-        //
-        // TODO(someday): I (Kenton) don't undestand why we'd explicitly want to hide this
-        //   error, even though hiding it is technically not a violation of the contract. To me
-        //   this seems undesirable, especially when it was the origin that failed. The caller
-        //   can always choose to ignore errors if they want (and many do, by passing to
-        //   waitUntil()). However, there is at least one test which depends on this behavior,
-        //   and probably production Workers in the wild, so I'm not changing it for now.
-      };
-
-      // Here we handle the inner promise -- the one held by the DeferredProxy that is
-      // returned below by the handleSerialize() method.
-      static auto constexpr handleResponse =
-          [](DeferredProxy<void>& deferred,
-             kj::Promise<void> writePayloadHeadersPromise,
-             kj::Promise<void> pumpRequestBodyPromise,
-             kj::Own<kj::AsyncOutputStream> bodyStream,
-             kj::Promise<kj::HttpClient::Response> responsePromise,
-             kj::Own<kj::HttpClient> client,
-             kj::Own<kj::AsyncInputStream> payloadStream) -> kj::Promise<void> {
-        try {
-          co_await deferred.proxyTask;
-          // Make sure headers get written even if the body was empty -- see comments earlier.
-          co_await writePayloadHeadersPromise;
-          // Make sure the request body is done being pumped and had no errors. If serialization
-          // completed successfully, then this should also complete immediately thereafter.
-          co_await pumpRequestBodyPromise;
-          // It is important to destroy the bodyStream before actually waiting on the
-          // responsePromise to ensure that the terminal chunk is written since the bodyStream
-          // may only write the terminal chunk in the streams destructor.
-          bodyStream = nullptr;
-          payloadStream = nullptr;
-          auto response = co_await responsePromise;
-          // We expect to see either 204 (success) or 413 (failure). Any other status code is a
-          // violation of the contract between us and the cache, and is an internal
-          // error, which we log. However, there's no need to throw, since the Cache API is an
-          // ephemeral K/V store, and we never guaranteed the script we'd actually cache anything.
-          if (response.statusCode != 204 && response.statusCode != 413) {
-            LOG_CACHE_ERROR_ONCE(
-                "Response to Cache API PUT was neither 204 nor 413: ", response);
-          }
-        } catch (...) {
-          handleException(kj::getCaughtExceptionAsKj());
-        }
-      };
+      // Hence, what you see here: we first await the serializePromise, then enter deferred proxying
+      // with our magic `KJ_CO_MAGIC`, then perform all our additional work. Then we
+      // `awaitDeferredProxy()` the whole thing.
 
       // Here we handle the promise for the DeferredProxy itself.
       static auto constexpr handleSerialize = [](
@@ -477,18 +425,45 @@ jsg::Promise<void> Cache::put(jsg::Lock& js, Request::Info requestOrUrl,
         });
         try {
           auto deferred = co_await serialize;
-          co_return DeferredProxy<void> {
-            handleResponse(deferred,
-                           kj::mv(writePayloadHeadersPromise),
-                           kj::mv(pumpRequestBodyPromise),
-                           kj::mv(bodyStream),
-                           kj::mv(responsePromise),
-                           kj::mv(httpClient),
-                           kj::mv(payloadStream))
-          };
+
+          // With our `serialize` promise having resolved to a DeferredProxy, we can now enter
+          // deferred proxying ourselves.
+          KJ_CO_MAGIC BEGIN_DEFERRED_PROXYING;
+
+          co_await deferred.proxyTask;
+          // Make sure headers get written even if the body was empty -- see comments earlier.
+          co_await writePayloadHeadersPromise;
+          // Make sure the request body is done being pumped and had no errors. If serialization
+          // completed successfully, then this should also complete immediately thereafter.
+          co_await pumpRequestBodyPromise;
+          // It is important to destroy the bodyStream before actually waiting on the
+          // responsePromise to ensure that the terminal chunk is written since the bodyStream
+          // may only write the terminal chunk in the streams destructor.
+          bodyStream = nullptr;
+          payloadStream = nullptr;
+          auto response = co_await responsePromise;
+          // We expect to see either 204 (success) or 413 (failure). Any other status code is a
+          // violation of the contract between us and the cache, and is an internal
+          // error, which we log. However, there's no need to throw, since the Cache API is an
+          // ephemeral K/V store, and we never guaranteed the script we'd actually cache anything.
+          if (response.statusCode != 204 && response.statusCode != 413) {
+            LOG_CACHE_ERROR_ONCE(
+                "Response to Cache API PUT was neither 204 nor 413: ", response);
+          }
         } catch (...) {
-          handleException(kj::getCaughtExceptionAsKj());
-          co_return newNoopDeferredProxy();
+          auto exception = kj::getCaughtExceptionAsKj();
+          if (exception.getType() != kj::Exception::Type::DISCONNECTED) {
+            kj::throwFatalException(kj::mv(exception));
+          }
+          // If the origin or the cache disconnected, we don't treat this as an error, as put()
+          // doesn't guarantee that it stores anything anyway.
+          //
+          // TODO(someday): I (Kenton) don't undestand why we'd explicitly want to hide this
+          //   error, even though hiding it is technically not a violation of the contract. To me
+          //   this seems undesirable, especially when it was the origin that failed. The caller
+          //   can always choose to ignore errors if they want (and many do, by passing to
+          //   waitUntil()). However, there is at least one test which depends on this behavior,
+          //   and probably production Workers in the wild, so I'm not changing it for now.
         }
       };
 

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

@@ -4,13 +4,13 @@
 namespace workerd::api {
 namespace {
 
-KJ_TEST("DeferredProxyPromise<T>: early co_return implicitly fulfills outer promise") {
+KJ_TEST("kj::Promise<DeferredProxy<T>>: early co_return implicitly fulfills outer promise") {
   kj::EventLoop loop;
   kj::WaitScope waitScope(loop);
 
   {
     // Implicit void co_return.
-    auto coro = []() -> DeferredProxyPromise<void> {
+    auto coro = []() -> kj::Promise<DeferredProxy<void>> {
       co_await kj::Promise<void>(kj::READY_NOW);
     };
     auto promise = coro();
@@ -21,7 +21,7 @@ KJ_TEST("DeferredProxyPromise<T>: early co_return implicitly fulfills outer prom
   }
   {
     // Explicit void co_return.
-    auto coro = []() -> DeferredProxyPromise<void> {
+    auto coro = []() -> kj::Promise<DeferredProxy<void>> {
       co_return;
     };
     auto promise = coro();
@@ -32,7 +32,7 @@ KJ_TEST("DeferredProxyPromise<T>: early co_return implicitly fulfills outer prom
   }
   {
     // Valueful co_return.
-    auto coro = []() -> DeferredProxyPromise<int> {
+    auto coro = []() -> kj::Promise<DeferredProxy<int>> {
       co_return 123;
     };
     auto promise = coro();
@@ -43,14 +43,15 @@ KJ_TEST("DeferredProxyPromise<T>: early co_return implicitly fulfills outer prom
   }
 }
 
-KJ_TEST("DeferredProxyPromise<T>: `co_yield BEGIN_DEFERRED_PROXYING` fulfills outer promise") {
+KJ_TEST("kj::Promise<DeferredProxy<T>>: `KJ_CO_MAGIC 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> {
+  auto coro = [&]() -> kj::Promise<DeferredProxy<int>> {
     co_await paf1.promise;
     KJ_CO_MAGIC BEGIN_DEFERRED_PROXYING;
     co_return co_await paf2.promise;
@@ -75,13 +76,14 @@ KJ_TEST("DeferredProxyPromise<T>: `co_yield BEGIN_DEFERRED_PROXYING` fulfills ou
   KJ_EXPECT(proxyTask.wait(waitScope) == 123);
 }
 
-KJ_TEST("DeferredProxyPromise<T>: unhandled exception before `co_yield BEGIN_DEFERRED_PROXYING`") {
+KJ_TEST("kj::Promise<DeferredProxy<T>>: unhandled exception before "
+    "`KJ_CO_MAGIC BEGIN_DEFERRED_PROXYING`") {
   kj::EventLoop loop;
   kj::WaitScope waitScope(loop);
 
   auto paf = kj::newPromiseAndFulfiller<void>();
 
-  auto coro = [&]() -> DeferredProxyPromise<int> {
+  auto coro = [&]() -> kj::Promise<DeferredProxy<int>> {
     co_await paf.promise;
     KJ_FAIL_ASSERT("promise should have been rejected");
   };
@@ -97,14 +99,15 @@ KJ_TEST("DeferredProxyPromise<T>: unhandled exception before `co_yield BEGIN_DEF
   KJ_EXPECT_THROW_MESSAGE("test error", promise.wait(waitScope));
 }
 
-KJ_TEST("DeferredProxyPromise<T>: unhandled exception after `co_yield BEGIN_DEFERRED_PROXYING`") {
+KJ_TEST("kj::Promise<DeferredProxy<T>>: unhandled exception after "
+    "`KJ_CO_MAGIC BEGIN_DEFERRED_PROXYING`") {
   kj::EventLoop loop;
   kj::WaitScope waitScope(loop);
 
   auto paf1 = kj::newPromiseAndFulfiller<void>();
   auto paf2 = kj::newPromiseAndFulfiller<int>();
 
-  auto coro = [&]() -> DeferredProxyPromise<int> {
+  auto coro = [&]() -> kj::Promise<DeferredProxy<int>> {
     co_await paf1.promise;
     KJ_CO_MAGIC BEGIN_DEFERRED_PROXYING;
     co_return co_await paf2.promise;
@@ -129,20 +132,20 @@ KJ_TEST("DeferredProxyPromise<T>: unhandled exception after `co_yield BEGIN_DEFE
   KJ_EXPECT_THROW_MESSAGE("test error", proxyTask.wait(waitScope));
 }
 
-KJ_TEST("DeferredProxyPromise<T>: can be `co_await`ed from another coroutine") {
+KJ_TEST("kj::Promise<DeferredProxy<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> {
+  auto nestedCoro = [&]() -> kj::Promise<DeferredProxy<int>> {
     co_await paf1.promise;
     KJ_CO_MAGIC BEGIN_DEFERRED_PROXYING;
     co_return co_await paf2.promise;
   };
 
-  auto coro = [&]() -> DeferredProxyPromise<int> {
+  auto coro = [&]() -> kj::Promise<DeferredProxy<int>> {
     auto deferred = co_await nestedCoro();
     KJ_CO_MAGIC BEGIN_DEFERRED_PROXYING;
     co_return co_await deferred.proxyTask;

src/workerd/api/deferred-proxy.h

@@ -61,69 +61,84 @@ inline kj::Promise<DeferredProxy<void>> addNoopDeferredProxy(kj::Promise<void> p
 // ---------------------------------------------------------
 // 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.
+// If a coroutine returns a kj::Promise<DeferredProxy<T>>, the coroutine implementation gains the
+// following features:
+//
+// - `KJ_CO_MAGIC 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 `KJ_CO_MAGIC` 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 `KJ_CO_MAGIC` described above), then reject the inner `proxyTask`.
+//
+// It is not possible to write a "regular" coroutine which returns kj::Promise<DeferredProxy<T>>;
+// that is, `co_return DeferredProxy<T> { ... }` is a compile error. You must initiate deferred
+// proxying using `KJ_CO_MAGIC BEGIN_DEFERRED_PROXYING`.
+
+template <typename T, typename... Args>
+class DeferredProxyCoroutine;
+// The coroutine adapter class, required for the compiler to know how to create coroutines
+// returning kj::Promise<DeferredProxy<T>>. We declare it here so we can name it in our
+// `coroutine_traits` specialization below.
 
-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`.
+}  // namespace workerd::api
 
-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.
+namespace KJ_COROUTINE_STD_NAMESPACE {
+// Enter the `std` or `std::experimental` namespace, depending on whether we're using C++20
+// coroutines or the Coroutines TS.
+
+template <class T, class... Args>
+struct coroutine_traits<kj::Promise<workerd::api::DeferredProxy<T>>, Args...> {
+  using promise_type = workerd::api::DeferredProxyCoroutine<T, Args...>;
 };
 
+}  // namespace KJ_COROUTINE_STD_NAMESPACE
+
+namespace workerd::api {
+
+class BeginDeferredProxyingConstant final {};
+constexpr BeginDeferredProxyingConstant BEGIN_DEFERRED_PROXYING {};
+// A magic constant which a DeferredProxyPromise<T> coroutine can `KJ_CO_MAGIC` to indicate that the
+// deferred proxying phase of its operation has begun.
+
 template <typename T>
-class DeferredProxyPromise<T>::Coroutine:
-    public kj::_::CoroutineMixin<DeferredProxyPromise<T>::Coroutine, T> {
+class DeferredProxyCoroutine: public kj::_::CoroutineMixin<DeferredProxyCoroutine<T>, T> {
   // The coroutine adapter type for DeferredProxyPromise<T>. Most of the work is forwarded to the
   // regular kj::Promise<T> coroutine adapter.
 
   using InnerCoroutineAdapter =
       typename kj::_::stdcoro::coroutine_traits<kj::Promise<T>, Args...>::promise_type;
 
 public:
-  using Handle = kj::_::stdcoro::coroutine_handle<Coroutine>;
+  using Handle = kj::_::stdcoro::coroutine_handle<DeferredProxyCoroutine>;
 
-  Coroutine(kj::SourceLocation location = {}): inner(Handle::from_promise(*this), location) {}
+  DeferredProxyCoroutine(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.
+    // Promise owns `this`. So, returning a Promise with `inner.get_return_object()` in a `.then()`
+    // gives transitive ownership of the coroutine to the caller.
+    //
+    // Note that we must use a `.then()` instead of a coroutine to implement this, because trying to
+    // write a coroutine returning `kj::Promise<DeferredProxy<T>>` will just recurse into this
+    // coroutine adapter class.
+    //
+    // TODO(perf): We could avoid the `newPromiseAndFulfiller()` and `.then()` overhead by having
+    //   DeferredProxyCoroutine<T> implement PromiseNode. `this` could own `proxyTask` temporarily,
+    //   forwarding `PromiseNode::destroy()` calls to `proxyTask`'s PromiseNode (if it still
+    //   exists), which takes care of the ownership reference cycle. `yield_value()`, `fulfill()`,
+    //   and `unhandled_exception()`could arm our OnReadyEvent, and our `get()` implementation would
+    //   return `proxyTask`, or a stored exception.
     auto proxyTask = inner.get_return_object();
-    co_await beginDeferredProxying.promise;
-    co_return DeferredProxy<T> { kj::mv(proxyTask) };
+    return beginDeferredProxying.promise.then([proxyTask=kj::mv(proxyTask)]() mutable {
+      return DeferredProxy<T> { kj::mv(proxyTask) };
+    });
   }
 
   auto initial_suspend() { return inner.initial_suspend(); }

src/workerd/api/http.c++

@@ -458,7 +458,7 @@ public:
       if (end) co_await output.end();
     }
 
-    co_return newNoopDeferredProxy();
+    co_return;
   }
 
 private: