1
// Copyright (c) 2017-2022 Cloudflare, Inc.
2
// Licensed under the Apache 2.0 license found in the LICENSE file or at:
3
//     https://opensource.org/licenses/Apache-2.0
4
 
5
#pragma once
6
 
7
#include <kj/async.h>
8
#include <kj/debug.h>
9
 
10
namespace workerd::api {
11
 
12
// =======================================================================================
13
 
14
// Some API methods return Promise<DeferredProxy<T>> when the task can be separated into two
15
// parts: some work that must be done with the IoContext still live, and some part that
16
// can occur after the IoContext completes, but which should still be performed before
17
// the overall task is "done".
18
//
19
// In particular, when an HTTP event ends up proxying the response body stream (or WebSocket
20
// stream) directly to/from origin, then that streaming can take place without pinning the
21
// isolate in memory, and without holding the IoContext open. So,
22
// `ServiceWorkerGlobalScope::request()` returns `Promise<DeferredProxy<void>>`. The outer
23
// Promise waits for the JavaScript work to be done, and the inner DeferredProxy<void> represents
24
// the proxying step.
25
//
26
// Note that if you're performing a task that resolves to DeferredProxy but JavaScript is
27
// actually waiting for the result of the task, then it's your responsibility to call
28
// IoContext::current().registerPendingEvent() and attach it to `proxyTask`, otherwise
29
// the request might be canceled as the proxy task won't be recognized as something that the
30
// request is waiting on.
31
template <typename T>
32
struct DeferredProxy {
33
  // TODO(cleanup): Now that we have jsg::Promise, it might make sense for deferred proxying to
34
  //    be represented as `jsg::Promise<api::DeferredProxy<T>>`, since the outer promise is
35
  //    intended to represent activity that happens in JavaScript while the inner one represents
36
  //    pure I/O. This will require some refactoring, though.
37
 
38
  kj::Promise<T> proxyTask;
39
};
40
 
41
inline DeferredProxy<void> newNoopDeferredProxy() {
42
  return DeferredProxy<void>{kj::READY_NOW};
43
}
44
 
45
template <typename T>
46
inline DeferredProxy<T> newNoopDeferredProxy(T&& value) {
47
  return DeferredProxy<T>{kj::mv(value)};
48
}
49
 
50
// Helper method to use when you need to return `Promise<DeferredProxy<T>>` but no part of the
51
// operation you are returning is eligible to be deferred past the IoContext lifetime.
52
template <typename T>
53
inline kj::Promise<DeferredProxy<T>> addNoopDeferredProxy(kj::Promise<T> promise) {
54
  co_return newNoopDeferredProxy(co_await promise);
55
}
56
inline kj::Promise<DeferredProxy<void>> addNoopDeferredProxy(kj::Promise<void> promise) {
57
  co_await promise;
58
  co_return newNoopDeferredProxy();
59
}
60
 
61
// ---------------------------------------------------------
62
// Deferred proxy coroutine integration
63
 
64
// If a coroutine returns a kj::Promise<DeferredProxy<T>>, the coroutine implementation gains the
65
// following features:
66
//
67
// - `KJ_CO_MAGIC BEGIN_DEFERRED_PROXYING` fulfills the outer kj::Promise<DeferredProxy<T>>. The
68
//   resulting DeferredProxy<T> object contains a `proxyTask` Promise which owns the coroutine.
69
//
70
// - `co_return` implicitly fulfills the outer Promise for the DeferredProxy<T> (if it has not
71
//   already been fulfilled by the magic `KJ_CO_MAGIC` described above), then fulfills the inner
72
//   `proxyTask`.
73
//
74
// - Unhandled exceptions reject the outer kj::Promise<DeferredProxy<T>> (if it has not already
75
//   been fulfilled by the magic `KJ_CO_MAGIC` described above), then reject the inner `proxyTask`.
76
//
77
// It is not possible to write a "regular" coroutine which returns kj::Promise<DeferredProxy<T>>;
78
// that is, `co_return DeferredProxy<T> { ... }` is a compile error. You must initiate deferred
79
// proxying using `KJ_CO_MAGIC BEGIN_DEFERRED_PROXYING`.
80
 
81
// The coroutine adapter class, required for the compiler to know how to create coroutines
82
// returning kj::Promise<DeferredProxy<T>>. We declare it here so we can name it in our
83
// `coroutine_traits` specialization below.
84
template <typename T, typename... Args>
85
class DeferredProxyCoroutine;
86
}  // namespace workerd::api
87
 
88
namespace KJ_COROUTINE_STD_NAMESPACE {
89
// Enter the `std` or `std::experimental` namespace, depending on whether we're using C++20
90
// coroutines or the Coroutines TS.
91
 
92
template <class T, class... Args>
93
struct coroutine_traits<kj::Promise<workerd::api::DeferredProxy<T>>, Args...> {
94
  using promise_type = workerd::api::DeferredProxyCoroutine<T, Args...>;
95
};
96
 
97
}  // namespace KJ_COROUTINE_STD_NAMESPACE
98
 
99
namespace workerd::api {
100
 
101
class BeginDeferredProxyingConstant final {};
102
// A magic constant which a DeferredProxyPromise<T> coroutine can `KJ_CO_MAGIC` to indicate that the
103
// deferred proxying phase of its operation has begun.
104
constexpr BeginDeferredProxyingConstant BEGIN_DEFERRED_PROXYING{};
105
 
106
// A concept which is true if C is a coroutine adapter which supports the `co_yield` operator for
107
// type T. We could also check that the expression results in an awaitable, but that is already a
108
// compile error in other ways.
109
template <typename T, typename C>
110
concept CoroutineYieldValue = requires(T&& v, C coroutineAdapter) {
111
  { coroutineAdapter.yield_value(kj::fwd<T>(v)) };
112
};
113
 
114
// The coroutine adapter type for DeferredProxyPromise<T>. Most of the work is forwarded to the
115
// regular kj::Promise<T> coroutine adapter.
116
template <typename T, typename... Args>
117
class DeferredProxyCoroutine: public kj::_::PromiseNode,
118
                              public kj::_::CoroutineMixin<DeferredProxyCoroutine<T, Args...>, T> {
119
  using InnerCoroutineAdapter =
120
      kj::_::stdcoro::coroutine_traits<kj::Promise<T>, Args...>::promise_type;
121
 
122
 public:
123
  using Handle = kj::_::stdcoro::coroutine_handle<DeferredProxyCoroutine>;
124
 
125
  DeferredProxyCoroutine(kj::SourceLocation location = {})
126
      : inner(Handle::from_promise(*this), location) {}
127
 
128
  kj::Promise<DeferredProxy<T>> get_return_object() {
129
    // We need to return a RAII object which will destroy this (as in, `this`) coroutine adapter.
130
    // The logic which calls `coroutine_handle<>::destroy()` is tucked away in our inner coroutine
131
    // adapter, however, leading to the weird situation where the `inner.get_return_object()`
132
    // Promise owns `this`. And `this` owns `inner.get_return_object()` transitively via `result`!
133
    //
134
    // Fortunately, DeferredProxyCoroutine implements the PromiseNode interface, meaning when our
135
    // returned Promise is eventually dropped, our `PromiseNode::destroy()` implementation will be
136
    // called. This gives us the opportunity (that is, in `destroy()`) to destroy our
137
    // `inner.get_return_object()` Promise, breaking the ownership cycle and destroying `this`.
138
 
139
    result.value = DeferredProxy<T>{inner.get_return_object()};
140
    return kj::_::PromiseNode::to<kj::Promise<DeferredProxy<T>>>(kj::_::OwnPromiseNode(this));
141
  }
142
 
143
  auto initial_suspend() {
144
    return inner.initial_suspend();
145
  }
146
  auto final_suspend() noexcept {
147
    return inner.final_suspend();
148
  }
149
  // Just trivially forward these.
150
 
151
  void unhandled_exception() {
152
    // Reject our outer promise if it hasn't yet been fulfilled, or forward to the inner
153
    // implementation.
154
 
155
    if (!deferredProxyingHasBegun) {
156
      result.addException(kj::getCaughtExceptionAsKj());
157
      onReadyEvent.arm();
158
      deferredProxyingHasBegun = true;
159
    } else {
160
      inner.unhandled_exception();
161
    }
162
  }
163
 
164
  kj::_::stdcoro::suspend_never yield_value(decltype(BEGIN_DEFERRED_PROXYING)) {
165
    // This allows us to write `KJ_CO_MAGIC` within a DeferredProxyPromise<T> coroutine to fulfill
166
    // the coroutine's outer promise with a DeferredProxy<T>.
167
    //
168
    // This could alternatively be implemented as an await_transform() with a magic parameter type.
169
 
170
    fulfillOuterPromise();
171
    return {};
172
  }
173
 
174
  template <CoroutineYieldValue<InnerCoroutineAdapter> U>
175
  auto yield_value(U&& value) {
176
    // Forward all other `co_yield`s to the inner coroutine, if it has a `yield_value()`
177
    // implementation -- it might implement some magic, too.
178
    return inner.yield_value(kj::fwd<U>(value));
179
  }
180
 
181
  void fulfill(kj::_::FixVoid<T>&& value) {
182
    // Required by CoroutineMixin implementation to implement `co_return`.
183
 
184
    fulfillOuterPromise();
185
    inner.fulfill(kj::mv(value));
186
  }
187
 
188
  template <typename U>
189
  decltype(auto) await_transform(U&& awaitable) {
190
    // Trivially forward everything, so we can await anything a kj::Promise<T> can.
191
    return inner.await_transform(kj::fwd<U>(awaitable));
192
  }
193
 
194
  operator kj::_::CoroutineBase&() {
195
    return inner;
196
  }
197
  // Required by Awaiter<T>::await_suspend() to support awaiting Promises.
198
 
199
 private:
200
  void fulfillOuterPromise() {
201
    // Fulfill the outer promise if it hasn't already settled.
202
 
203
    if (!deferredProxyingHasBegun) {
204
      // Our `result` is put in place already by `get_return_object()`, so all we have to do is arm
205
      // the event.
206
      onReadyEvent.arm();
207
      deferredProxyingHasBegun = true;
208
    }
209
  }
210
 
211
  // PromiseNode implementation
212
 
213
  void setSelfPointer(kj::_::OwnPromiseNode* selfPtr) noexcept override {
214
    this->selfPtr = selfPtr;
215
  }
216
 
217
  void destroy() override {
218
    // The promise returned by `inner.get_return_object()` is what actually owns this coroutine
219
    // frame. We temporarily store that in `result` until our outer promise is fulfilled. So, to
220
    // destroy ourselves, we must manually drop `result`.
221
    //
222
    // On the other hand, if our outer promise has already been fulfilled, then `result` has already
223
    // been delivered to wherever it is going, and someone else directly owns the coroutine now, not
224
    // us. In this case, this `destroy()` override will have already been called (and it will have
225
    // been a no-op), because our own OwnPromiseNode will have already been dropped in `get()`.
226
 
227
    auto drop = kj::mv(result);
228
  }
229
 
230
  void onReady(kj::_::Event* event) noexcept override {
231
    onReadyEvent.init(event);
232
  }
233
 
234
  void get(kj::_::ExceptionOrValue& output) noexcept override {
235
    // Make sure that the outer PromiseNode (`this` one) is destroyed before the inner PromiseNode.
236
    // kj-async should already provide us this guarantee, but since incorrect destruction order
237
    // would cause invalid memory access, we provide a stronger guarantee. Also see the comment for
238
    // the `result` data member.
239
    KJ_ASSERT(selfPtr != nullptr);
240
    KJ_DEFER(*selfPtr = nullptr);
241
 
242
    static_cast<decltype(result)&>(output) = kj::mv(result);
243
  }
244
 
245
  void tracePromise(kj::_::TraceBuilder& builder, bool stopAtNextEvent) override {
246
    // The PromiseNode we're waiting on is whatever the coroutine is waiting on.
247
    static_cast<kj::_::PromiseNode&>(inner).tracePromise(builder, stopAtNextEvent);
248
 
249
    // Maybe returning the address of get() will give us a function name with meaningful type
250
    // information.
251
    builder.add(getMethodStartAddress(implicitCast<PromiseNode&>(*this), &PromiseNode::get));
252
  }
253
 
254
  // We defer the majority of the implementation to the regular kj::Promise<T> coroutine adapter.
255
  InnerCoroutineAdapter inner;
256
 
257
  // Helper to arm the event which fires when the outer promise (that is, `this` PromiseNode) for
258
  // the DeferredProxy<T> is ready.
259
  OnReadyEvent onReadyEvent;
260
 
261
  // Stores the result for the outer promise.
262
  //
263
  // WARNING: This object owns `this` PromiseNode! If `result` is ever moved away, as is done in
264
  // `get()`, we must arrange to make sure that no one ever tries to use `this` PromiseNode again.
265
  // Stated another way, we must guarantee that the outer PromiseNode (for `DeferredProxy<T>`) is
266
  // always destroyed before the inner PromiseNode (for `T`). kj-async always does this anyway, but
267
  // we implement an additional safeguard by immediately destroying our own `OwnPromiseNode` (which
268
  // we have access to via `setSelfPointer()`) when we move `result` away in `get()`.
269
  kj::_::ExceptionOr<DeferredProxy<T>> result;
270
 
271
  // Used to drop ourselves in `get()` -- see comment for `result`.
272
  kj::_::OwnPromiseNode* selfPtr = nullptr;
273
 
274
  // Set to true when deferred proxying has begun -- that is, when the outer DeferredProxy<T>
275
  // promise is fulfilled by calling `onReadyEvent.arm()`.
276
  bool deferredProxyingHasBegun = false;
277
};
278
 
279
}  // namespace workerd::api