Skip to content
rachel/workerd
File

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

8.2 KB
1#include "deferred-proxy.h"
2 
3#include <kj/test.h>
4 
5namespace workerd::api {
6namespace {
7 
8KJ_TEST("kj::Promise<DeferredProxy<T>>: early co_return implicitly fulfills outer promise") {
9  kj::EventLoop loop;
10  kj::WaitScope waitScope(loop);
11 
12  {
13    // Implicit void co_return.
14    auto coro = []() -> kj::Promise<DeferredProxy<void>> {
15      co_await kj::Promise<void>(kj::READY_NOW);
16    };
17    auto promise = coro();
18    KJ_EXPECT(promise.poll(waitScope));
19    auto proxyTask = promise.wait(waitScope).proxyTask;
20    KJ_EXPECT(proxyTask.poll(waitScope));
21    proxyTask.wait(waitScope);
22  }
23  {
24    // Explicit void co_return.
25    auto coro = []() -> kj::Promise<DeferredProxy<void>> { co_return; };
26    auto promise = coro();
27    KJ_EXPECT(promise.poll(waitScope));
28    auto proxyTask = promise.wait(waitScope).proxyTask;
29    KJ_EXPECT(proxyTask.poll(waitScope));
30    proxyTask.wait(waitScope);
31  }
32  {
33    // Valueful co_return.
34    auto coro = []() -> kj::Promise<DeferredProxy<int>> { co_return 123; };
35    auto promise = coro();
36    KJ_EXPECT(promise.poll(waitScope));
37    auto proxyTask = promise.wait(waitScope).proxyTask;
38    KJ_EXPECT(proxyTask.poll(waitScope));
39    KJ_EXPECT(proxyTask.wait(waitScope) == 123);
40  }
41}
42 
43KJ_TEST("kj::Promise<DeferredProxy<T>>: `KJ_CO_MAGIC BEGIN_DEFERRED_PROXYING` fulfills outer "
44        "promise") {
45  kj::EventLoop loop;
46  kj::WaitScope waitScope(loop);
47 
48  auto paf1 = kj::newPromiseAndFulfiller<void>();
49  auto paf2 = kj::newPromiseAndFulfiller<int>();
50 
51  auto coro = [&]() -> kj::Promise<DeferredProxy<int>> {
52    co_await paf1.promise;
53    KJ_CO_MAGIC BEGIN_DEFERRED_PROXYING;
54    co_return co_await paf2.promise;
55  };
56 
57  auto promise = coro();
58 
59  // paf1 unfulfilled, so we don't have a DeferredProxy<T> yet.
60  KJ_EXPECT(!promise.poll(waitScope));
61 
62  paf1.fulfiller->fulfill();
63 
64  KJ_EXPECT(promise.poll(waitScope));
65  auto proxyTask = promise.wait(waitScope).proxyTask;
66 
67  // paf2 unfulfilled, so we don't have a T yet.
68  KJ_EXPECT(!proxyTask.poll(waitScope));
69 
70  paf2.fulfiller->fulfill(123);
71 
72  KJ_EXPECT(proxyTask.poll(waitScope));
73  KJ_EXPECT(proxyTask.wait(waitScope) == 123);
74}
75 
76KJ_TEST("kj::Promise<DeferredProxy<T>>: unhandled exception before "
77        "`KJ_CO_MAGIC BEGIN_DEFERRED_PROXYING`") {
78  kj::EventLoop loop;
79  kj::WaitScope waitScope(loop);
80 
81  auto paf = kj::newPromiseAndFulfiller<void>();
82 
83  auto coro = [&]() -> kj::Promise<DeferredProxy<int>> {
84    co_await paf.promise;
85    KJ_FAIL_ASSERT("promise should have been rejected");
86  };
87 
88  auto promise = coro();
89 
90  // paf unfulfilled, so we don't have a DeferredProxy<T> yet.
91  KJ_EXPECT(!promise.poll(waitScope));
92 
93  paf.fulfiller->reject(KJ_EXCEPTION(FAILED, "test error"));
94 
95  KJ_EXPECT(promise.poll(waitScope));
96  KJ_EXPECT_THROW_MESSAGE("test error", promise.wait(waitScope));
97}
98 
99KJ_TEST("kj::Promise<DeferredProxy<T>>: unhandled exception after "
100        "`KJ_CO_MAGIC BEGIN_DEFERRED_PROXYING`") {
101  kj::EventLoop loop;
102  kj::WaitScope waitScope(loop);
103 
104  auto paf1 = kj::newPromiseAndFulfiller<void>();
105  auto paf2 = kj::newPromiseAndFulfiller<int>();
106 
107  auto coro = [&]() -> kj::Promise<DeferredProxy<int>> {
108    co_await paf1.promise;
109    KJ_CO_MAGIC BEGIN_DEFERRED_PROXYING;
110    co_return co_await paf2.promise;
111  };
112 
113  auto promise = coro();
114 
115  // paf1 unfulfilled, so we don't have a DeferredProxy<T> yet.
116  KJ_EXPECT(!promise.poll(waitScope));
117 
118  paf1.fulfiller->fulfill();
119 
120  KJ_EXPECT(promise.poll(waitScope));
121  auto proxyTask = promise.wait(waitScope).proxyTask;
122 
123  // paf2 unfulfilled, so we don't have a T yet.
124  KJ_EXPECT(!proxyTask.poll(waitScope));
125 
126  paf2.fulfiller->reject(KJ_EXCEPTION(FAILED, "test error"));
127 
128  KJ_EXPECT(proxyTask.poll(waitScope));
129  KJ_EXPECT_THROW_MESSAGE("test error", proxyTask.wait(waitScope));
130}
131 
132KJ_TEST("kj::Promise<DeferredProxy<T>>: can be `co_await`ed from another coroutine") {
133  kj::EventLoop loop;
134  kj::WaitScope waitScope(loop);
135 
136  auto paf1 = kj::newPromiseAndFulfiller<void>();
137  auto paf2 = kj::newPromiseAndFulfiller<int>();
138 
139  auto nestedCoro = [&]() -> kj::Promise<DeferredProxy<int>> {
140    co_await paf1.promise;
141    KJ_CO_MAGIC BEGIN_DEFERRED_PROXYING;
142    co_return co_await paf2.promise;
143  };
144 
145  auto coro = [&]() -> kj::Promise<DeferredProxy<int>> {
146    auto deferred = co_await nestedCoro();
147    KJ_CO_MAGIC BEGIN_DEFERRED_PROXYING;
148    co_return co_await deferred.proxyTask;
149  };
150 
151  auto promise = coro();
152 
153  // paf1 unfulfilled, so we don't have a DeferredProxy<T> yet.
154  KJ_EXPECT(!promise.poll(waitScope));
155 
156  paf1.fulfiller->fulfill();
157 
158  KJ_EXPECT(promise.poll(waitScope));
159  auto proxyTask = promise.wait(waitScope).proxyTask;
160 
161  // paf2 unfulfilled, so we don't have a T yet.
162  KJ_EXPECT(!proxyTask.poll(waitScope));
163 
164  paf2.fulfiller->fulfill(123);
165 
166  KJ_EXPECT(proxyTask.poll(waitScope));
167  KJ_EXPECT(proxyTask.wait(waitScope) == 123);
168}
169 
170struct Counter {
171  size_t& wind;
172  size_t& unwind;
173  Counter(size_t& wind, size_t& unwind): wind(wind), unwind(unwind) {
174    ++wind;
175  }
176  ~Counter() {
177    ++unwind;
178  }
179  KJ_DISALLOW_COPY_AND_MOVE(Counter);
180};
181 
182kj::Promise<DeferredProxy<void>> cancellationTester(kj::Promise<void> preDeferredProxying,
183    kj::Promise<void> postDeferredProxying,
184    size_t& wind,
185    size_t& unwind) {
186  Counter preCounter(wind, unwind);
187  co_await preDeferredProxying;
188  KJ_CO_MAGIC BEGIN_DEFERRED_PROXYING;
189  Counter postCounter(wind, unwind);
190  co_await postDeferredProxying;
191};
192 
193KJ_TEST("kj::Promise<DeferredProxy<T>>: can be canceled while suspended before deferred proxying") {
194  kj::EventLoop loop;
195  kj::WaitScope waitScope(loop);
196 
197  size_t wind = 0, unwind = 0;
198 
199  {
200    auto neverDone1 = kj::Promise<void>(kj::NEVER_DONE);
201    auto neverDone2 = kj::Promise<void>(kj::NEVER_DONE);
202    neverDone1 = neverDone1.attach(kj::heap<Counter>(wind, unwind));
203    neverDone2 = neverDone2.attach(kj::heap<Counter>(wind, unwind));
204    auto promise = cancellationTester(kj::mv(neverDone1), kj::mv(neverDone2), wind, unwind);
205    KJ_EXPECT(!promise.poll(waitScope));
206  }
207 
208  KJ_EXPECT(wind == 3);
209  KJ_EXPECT(unwind == 3);
210}
211 
212KJ_TEST("kj::Promise<DeferredProxy<T>>: can be canceled while suspended after deferred proxying") {
213  kj::EventLoop loop;
214  kj::WaitScope waitScope(loop);
215 
216  size_t wind = 0, unwind = 0;
217 
218  {
219    auto readyNow = kj::Promise<void>(kj::READY_NOW);
220    auto neverDone = kj::Promise<void>(kj::NEVER_DONE);
221    readyNow = readyNow.attach(kj::heap<Counter>(wind, unwind));
222    neverDone = neverDone.attach(kj::heap<Counter>(wind, unwind));
223    auto promise = cancellationTester(kj::mv(readyNow), kj::mv(neverDone), wind, unwind);
224    auto proxyTask = promise.wait(waitScope).proxyTask;
225    KJ_EXPECT(!proxyTask.poll(waitScope));
226  }
227 
228  KJ_EXPECT(wind == 4);
229  KJ_EXPECT(unwind == 4);
230}
231 
232KJ_TEST("kj::Promise<DeferredProxy<T>>: destroying inner PromiseNode before outer does not "
233        "segfault") {
234  // Destroy the inner promise before the outer promise to test our safeguard against incorrect
235  // destruction order causing segfaults.
236 
237  kj::EventLoop loop;
238  kj::WaitScope waitScope(loop);
239 
240  auto coro = []() -> kj::Promise<DeferredProxy<void>> {
241    KJ_CO_MAGIC BEGIN_DEFERRED_PROXYING;
242    co_await kj::Promise<void>(kj::NEVER_DONE);
243  };
244 
245  auto outer = coro();
246 
247  // We could call `get()` on the outer node immediately, even before it reports it is ready, but
248  // we call `poll()` for good measure, in case the DeferredProxyCoroutine implementation ever
249  // changes to disallow `get()`-before-ready. We cannot use `wait()` for this purpose, because
250  // `wait()` would avoid the segfault by (correctly) destroying the outer PromiseNode before
251  // returning the result to us.
252  KJ_EXPECT(outer.poll(waitScope));
253 
254  auto outerNode = kj::_::PromiseNode::from(kj::mv(outer));
255 
256  // `poll()`, unlike `wait()`, does not call `setSelfPointer()` on the outer PromiseNode, which
257  // would cause an assertion failure inside the outer PromiseNode's `get()` implementation, so we
258  // have to do it ourselves.
259  outerNode->setSelfPointer(&outerNode);
260 
261  kj::_::ExceptionOr<DeferredProxy<void>> result;
262  outerNode->get(result);
263 
264  {
265    // Destroy the inner promise.
266    auto inner = kj::mv(KJ_ASSERT_NONNULL(result.value).proxyTask);
267  }
268 
269  // Destroy the outer promise. At one time, this caused a segfault ... or at least it produced
270  // invalid accesses under Valgrind. :/
271  outerNode = nullptr;
272}
273 
274}  // namespace
275}  // namespace workerd::api