1;;! component_model_async = true
2;;! component_model_async_stackful = true
3;;! component_model_async_builtins = true
4;;! component_model_threading = true
5;;! reference_types = true
6
7;; Tests that cancellation works with the async threading intrinsics.
8;; Consists of two components, C and D. C implements functions that mix cancellable and uncancellable yields and suspensions.
9;; D calls these functions and cancels the resulting subtasks, ensuring that cancellation is only seen when expected.
10
11;; -- Component C --
12
13;; `run-yield`: Yields twice, first with an uncancellable yield, then with a cancellable yield.
14;;      The caller cancels the subtask during the first yield, and ensures that the cancellation only takes effect
15;;      on the second yield.
16
17;; `run-yield-to`: Yields twice to a spawned thread, first with an uncancellable yield, then with a cancellable yield.
18;;      A complication is that we can't guarantee that if the spawned thread yields, the supertask will be scheduled to
19;;      cancel the subtask before the subtask's implicit thread is rescheduled. To handle this, the subtask's implicit
20;;      thread first waits on a future to be written by the supertask, then yields to the spawned thread.
21
22;; `run-suspend`: More complex, because executing an uncancellable suspension requires another
23;;      thread in the same subtask to explicitly wake it up. This is done by the subtask spawning a new thread that
24;;      waits on a future to be written by the supertask, and then resumes the main thread once that happens.
25;;      After setting up this thread, `run-suspend` performs an uncancellable suspend, then a cancellable suspend.
26;;      The caller cancels the subtask during the first suspend, writes to the future to make the spawned thread
27;;      resume the implicit thread, and ensures that the cancellation only takes effect on the second suspend.
28
29;; `run-switch-to`: Similar to `run-suspend`, but uses `thread.switch-to` instead of `thread.suspend`.
30
31;; -- Component D --
32
33;; `run-test`: Calls one of the functions in C based on a test id, cancels the resulting subtask, and ensures that
34;;      cancellation is only seen when expected.
35
36;; `run`: Calls `run-test` for each of the functions in C.
37
38(component
39    (component $C
40        (type $FT (future))
41        (core module $Memory (memory (export "mem") 1))
42        (core instance $memory (instantiate $Memory))
43        ;; Defines the table for the thread start functions, of which there are two
44        (core module $libc
45            (table (export "__indirect_function_table") 2 funcref))
46        (core module $CM
47            (import "" "mem" (memory 1))
48            (import "" "task.cancel" (func $task-cancel))
49            (import "" "thread.new-indirect" (func $thread-new-indirect (param i32 i32) (result i32)))
50            (import "" "thread.suspend" (func $thread-suspend (result i32)))
51            (import "" "thread.suspend-cancellable" (func $thread-suspend-cancellable (result i32)))
52            (import "" "thread.yield-to" (func $thread-yield-to (param i32) (result i32)))
53            (import "" "thread.yield-to-cancellable" (func $thread-yield-to-cancellable (param i32) (result i32)))
54            (import "" "thread.switch-to" (func $thread-switch-to (param i32) (result i32)))
55            (import "" "thread.switch-to-cancellable" (func $thread-switch-to-cancellable (param i32) (result i32)))
56            (import "" "thread.yield" (func $thread-yield (result i32)))
57            (import "" "thread.yield-cancellable" (func $thread-yield-cancellable (result i32)))
58            (import "" "thread.index" (func $thread-index (result i32)))
59            (import "" "thread.resume-later" (func $thread-resume-later (param i32)))
60            (import "" "future.read" (func $future.read (param i32 i32) (result i32)))
61            (import "" "waitable.join" (func $waitable.join (param i32 i32)))
62            (import "" "waitable-set.new" (func $waitable-set.new (result i32)))
63            (import "" "waitable-set.wait" (func $waitable-set.wait (param i32 i32) (result i32)))
64            (import "libc" "__indirect_function_table" (table $indirect-function-table 2 funcref))
65
66            ;; Indices into the function table for the thread start functions
67            (global $wake-from-suspend-ftbl-idx i32 (i32.const 0))
68            (global $just-yield-ftbl-idx i32 (i32.const 1))
69
70            (func (export "run-yield")
71                ;; Yield back to the caller, who will attempt to cancel us, but we won't see it
72                ;; because we're using an uncancellable yield
73                (if (i32.ne (call $thread-yield) (i32.const 0)) (then unreachable))
74                ;; Yield back to the caller again. This time, we should receive the cancellation immediately.
75                (if (i32.ne (call $thread-yield-cancellable) (i32.const 1)) (then unreachable))
76                (call $task-cancel)
77            )
78
79            (func $wait-for-future-write (param i32)
80                (local $ret i32)
81                ;; Perform a future.read, which will block, waiting for the supertask to write
82                (local.set $ret (call $future.read (local.get 0) (i32.const 0xba5eba11)))
83                (if (i32.ne (i32.const 0 (; COMPLETED ;)) (local.get $ret))
84                    (then unreachable))
85            )
86
87            (func $wake-from-suspend (param i32)
88                ;; Extract the thread index and future to wait on from the argument structure
89                (local $thread-index i32) (local $future i32)
90                (local.set $thread-index (i32.load offset=0 (local.get 0)))
91                (local.set $future (i32.load offset=4 (local.get 0)))
92
93                ;; Wait for the supertask to signal us to wake up suspended thread.
94                (call $wait-for-future-write (local.get $future))
95                ;; Resume the main thread, which is suspended in an uncancellable suspend
96                (call $thread-resume-later (local.get $thread-index))
97            )
98
99            (func $just-yield (param $explicit-thread-idx i32)
100                ;; Yield nondeterministically, either back to the supertask, who will then wait on cancellation to be acknowledged,
101                ;; or to the implicit thread, who will acknowledge the cancellation.
102                (if (i32.ne (call $thread-yield) (i32.const 0)) (then unreachable))
103            )
104
105            ;; Initialize the function table that will be used by thread.new-indirect
106            (elem (table $indirect-function-table) (i32.const 0 (; wake-from-suspend-ftbl-idx ;)) func $wake-from-suspend)
107            (elem (table $indirect-function-table) (i32.const 1 (; just-yield-ftbl-idx ;)) func $just-yield)
108
109            (func (export "run-yield-to") (param $futr i32)
110                (local $thread-index i32)
111                ;; Spawn a new thread that will wake us up from our uncancellable suspend; we'll switch to it next
112                (local.set $thread-index
113                    (call $thread-new-indirect (global.get $just-yield-ftbl-idx) (call $thread-index)))
114
115                ;; We can't guarantee that the supertask will be scheduled to cancel us before we're rescheduled, so we first
116                ;; wait on the future to be written, then yield to the spawned thread. This means that cancellation will be
117                ;; sent while we're waiting on the future rather than at the yield point, but the cancel will still be pending
118                ;; when we reach the yield point, so it should still be ignored by the uncancellable yield and only take effect
119                ;; when we reach the second, cancellable yield.
120                (call $wait-for-future-write (local.get $futr))
121
122                ;; Yield to the spawned thread uncancellably. We should eventually be rescheduled without being notified
123                ;; of the pending cancellation.
124                (if (i32.ne (call $thread-yield-to (local.get $thread-index)) (i32.const 0)) (then unreachable))
125                ;; Yield to the spawned thread again. This time we should see the cancellation immediately.
126                (if (i32.ne (call $thread-yield-to-cancellable (local.get $thread-index)) (i32.const 1)) (then unreachable))
127                (call $task-cancel)
128            )
129
130            (func (export "run-suspend") (param $futr i32)
131                ;; Set up the arguments for the wake-for-suspend thread start function.
132                ;; It expects a pointer to a structure containing the thread index to resume
133                ;; and the future to wait on before resuming it.
134                (local $wake-from-suspend-argp i32)
135                (local.set $wake-from-suspend-argp (i32.const 4))
136                (i32.store offset=0 (local.get $wake-from-suspend-argp) (call $thread-index))
137                (i32.store offset=4 (local.get $wake-from-suspend-argp) (local.get $futr))
138
139                ;; Spawn a new thread that will wake us up from our uncancellable suspend and schedule
140                ;; it to resume after we suspend.
141                (call $thread-resume-later
142                    (call $thread-new-indirect (global.get $wake-from-suspend-ftbl-idx) (local.get $wake-from-suspend-argp)))
143
144                ;; Request suspension. We will not be woken up by cancellation, because this is an uncancellable
145                ;; suspend. We will be woken up by the other thread we spawned above, which will be resumed after
146                ;; the supertask cancels our subtask.
147                (if (i32.ne (call $thread-suspend) (i32.const 0)) (then unreachable))
148                ;; Request suspension again. This time we should see the cancellation immediately.
149                (if (i32.ne (call $thread-suspend-cancellable) (i32.const 1)) (then unreachable))
150                (call $task-cancel)
151            )
152
153            (func (export "run-switch-to") (param $futr i32)
154                (local $thread-index i32)
155                ;; Set up the arguments for the wake-for-suspend thread start function.
156                ;; It expects a pointer to a structure containing the thread index to resume
157                ;; and the future to wait on before resuming it.
158                (local $wake-from-suspend-argp i32)
159                (local.set $wake-from-suspend-argp (i32.const 4))
160                (i32.store offset=0 (local.get $wake-from-suspend-argp) (call $thread-index))
161                (i32.store offset=4 (local.get $wake-from-suspend-argp) (local.get $futr))
162
163                ;; Spawn a new thread that will wake us up from our uncancellable suspend; we'll switch to it next
164                (local.set $thread-index
165                    (call $thread-new-indirect (global.get $wake-from-suspend-ftbl-idx) (local.get $wake-from-suspend-argp)))
166
167                ;; Request suspension by switching to the spawned thread.
168                ;; We will not be woken up by cancellation, because this is an uncancellable suspend.
169                ;; We will be woken up by the other thread we spawned above, which will be resumed after
170                ;; the supertask cancels our subtask.
171                (if (i32.ne (call $thread-switch-to (local.get $thread-index)) (i32.const 0)) (then unreachable))
172                ;; Request suspension again. This time we should see the cancellation immediately.
173                (if (i32.ne (call $thread-switch-to-cancellable (local.get $thread-index)) (i32.const 1)) (then unreachable))
174                (call $task-cancel)
175            )
176        )
177
178        ;; Instantiate the libc module to get the table
179        (core instance $libc (instantiate $libc))
180        ;; Get access to `thread.new-indirect` that uses the table from libc
181        (core type $start-func-ty (func (param i32)))
182        (alias core export $libc "__indirect_function_table" (core table $indirect-function-table))
183
184        (core func $task-cancel (canon task.cancel))
185        (core func $thread-new-indirect
186            (canon thread.new-indirect $start-func-ty (table $indirect-function-table)))
187        (core func $thread-yield (canon thread.yield))
188        (core func $thread-yield-cancellable (canon thread.yield cancellable))
189        (core func $thread-index (canon thread.index))
190        (core func $thread-yield-to (canon thread.yield-to))
191        (core func $thread-yield-to-cancellable (canon thread.yield-to cancellable))
192        (core func $thread-resume-later (canon thread.resume-later))
193        (core func $thread-switch-to (canon thread.switch-to))
194        (core func $thread-switch-to-cancellable (canon thread.switch-to cancellable))
195        (core func $thread-suspend (canon thread.suspend))
196        (core func $thread-suspend-cancellable (canon thread.suspend cancellable))
197        (core func $future.read (canon future.read $FT (memory $memory "mem")))
198        (core func $waitable-set.new (canon waitable-set.new))
199        (core func $waitable.join (canon waitable.join))
200        (core func $waitable-set.wait (canon waitable-set.wait (memory $memory "mem")))
201
202        ;; Instantiate the main module
203        (core instance $cm (
204            instantiate $CM
205                (with "" (instance
206                    (export "mem" (memory $memory "mem"))
207                    (export "task.cancel" (func $task-cancel))
208                    (export "thread.new-indirect" (func $thread-new-indirect))
209                    (export "thread.index" (func $thread-index))
210                    (export "thread.yield-to" (func $thread-yield-to))
211                    (export "thread.yield-to-cancellable" (func $thread-yield-to-cancellable))
212                    (export "thread.yield" (func $thread-yield))
213                    (export "thread.yield-cancellable" (func $thread-yield-cancellable))
214                    (export "thread.switch-to" (func $thread-switch-to))
215                    (export "thread.switch-to-cancellable" (func $thread-switch-to-cancellable))
216                    (export "thread.suspend" (func $thread-suspend))
217                    (export "thread.suspend-cancellable" (func $thread-suspend-cancellable))
218                    (export "thread.resume-later" (func $thread-resume-later))
219                    (export "future.read" (func $future.read))
220                    (export "waitable.join" (func $waitable.join))
221                    (export "waitable-set.wait" (func $waitable-set.wait))
222                    (export "waitable-set.new" (func $waitable-set.new))))
223                (with "libc" (instance $libc))))
224
225        (func (export "run-yield") async (result u32) (canon lift (core func $cm "run-yield") async))
226        (func (export "run-yield-to") async (param "fut" $FT) (result u32) (canon lift (core func $cm "run-yield-to") async))
227        (func (export "run-suspend") async (param "fut" $FT) (result u32) (canon lift (core func $cm "run-suspend") async))
228        (func (export "run-switch-to") async (param "fut" $FT) (result u32) (canon lift (core func $cm "run-switch-to") async))
229    )
230
231    (component $D
232        (type $FT (future))
233        (import "run-yield" (func $run-yield async (result u32)))
234        (import "run-yield-to" (func $run-yield-to async (param "fut" $FT) (result u32)))
235        (import "run-suspend" (func $run-suspend async (param "fut" $FT) (result u32)))
236        (import "run-switch-to" (func $run-switch-to async (param "fut" $FT) (result u32)))
237
238        (core module $Memory (memory (export "mem") 1))
239        (core instance $memory (instantiate $Memory))
240        (core module $DM
241            (import "" "mem" (memory 1))
242            (import "" "subtask.cancel" (func $subtask.cancel (param i32) (result i32)))
243            (import "" "run-yield" (func $run-yield (param i32) (result i32)))
244            (import "" "run-yield-to" (func $run-yield-to (param i32 i32) (result i32)))
245            (import "" "run-suspend" (func $run-suspend (param i32 i32) (result i32)))
246            (import "" "run-switch-to" (func $run-switch-to (param i32 i32) (result i32)))
247            (import "" "waitable.join" (func $waitable.join (param i32 i32)))
248            (import "" "waitable-set.new" (func $waitable-set.new (result i32)))
249            (import "" "waitable-set.wait" (func $waitable-set.wait (param i32 i32) (result i32)))
250            (import "" "future.new" (func $future.new (result i64)))
251            (import "" "future.write" (func $future.write (param i32 i32) (result i32)))
252            (import "" "thread.yield" (func $thread-yield (result i32)))
253
254            (func $run-test (param $test-id i32) (result i32)
255                (local $ret i32) (local $subtask i32)
256                (local $ws i32) (local $event_code i32)
257                (local $run-retp i32) (local $wait-retp i32)
258                (local $ret64 i64) (local $futr i32) (local $futw i32)
259
260                ;; Set up return value storage for run-suspend/switch-to and waitable-set.wait
261                (local.set $run-retp (i32.const 4))
262                (local.set $wait-retp (i32.const 8))
263                (i32.store (local.get $run-retp) (i32.const 0xbad0bad0))
264                (i32.store (local.get $wait-retp) (i32.const 0xbad0bad0))
265
266                ;; Create a future that the subtask may wait on
267                (local.set $ret64 (call $future.new))
268                (local.set $futr (i32.wrap_i64 (local.get $ret64)))
269                (local.set $futw (i32.wrap_i64 (i64.shr_u (local.get $ret64) (i64.const 32))))
270
271                ;; Calling run-suspend/switch-to will start the thread, which will suspend.
272                ;; This is basically a switch statement:
273                ;; 0: run-yield
274                ;; 1: run-yield-to
275                ;; 2: run-suspend
276                ;; 3: run-switch-to
277                (if (i32.eq (local.get $test-id) (i32.const 0))
278                    (then (local.set $ret (call $run-yield (local.get $run-retp))))
279                    (else (if (i32.eq (local.get $test-id) (i32.const 1))
280                        (then (local.set $ret (call $run-yield-to (local.get $futr) (local.get $run-retp))))
281                        (else (if (i32.eq (local.get $test-id) (i32.const 2))
282                            (then (local.set $ret (call $run-suspend (local.get $futr) (local.get $run-retp))))
283                            (else (if (i32.eq (local.get $test-id) (i32.const 3))
284                                (then (local.set $ret (call $run-switch-to (local.get $futr) (local.get $run-retp))))
285                                (else unreachable))))))))
286
287                ;; Ensure that the thread started
288                (if (i32.ne (i32.and (local.get $ret) (i32.const 0xF)) (i32.const 1 (; STARTED ;)))
289                 (then unreachable))
290                ;; Extract the subtask index
291                (local.set $subtask (i32.shr_u (local.get $ret) (i32.const 4)))
292                ;; Cancel the subtask, which should block, because the initial suspend/yield is uncancellable
293                (local.set $ret (call $subtask.cancel (local.get $subtask)))
294                ;; Ensure the cancellation blocked
295                (if (i32.ne (local.get $ret) (i32.const -1 (; BLOCKED ;)))
296                 (then unreachable))
297
298                ;; If we're not testing run-yield, the subtask is expecting a write to our future, so write to it
299                (if (i32.ne (local.get $test-id) (i32.const 0))
300                    (then
301                        (local.set $ret (call $future.write (local.get $futw) (i32.const 0xdeadbeef)))
302                        ;; The write should succeed
303                        (if (i32.ne (i32.const 0 (; COMPLETED ;)) (local.get $ret))
304                            (then unreachable))))
305
306                ;; Wait on the subtask, which will eventually progress to a cancellable yield/suspend and acknowledge the cancellation
307                (local.set $ws (call $waitable-set.new))
308                (call $waitable.join (local.get $subtask) (local.get $ws))
309                (local.set $event_code (call $waitable-set.wait (local.get $ws) (local.get $wait-retp)))
310                ;; Ensure we got the subtask event
311                (if (i32.ne (local.get $event_code) (i32.const 1 (; SUBTASK ;)))
312                 (then unreachable))
313                ;; Ensure the subtask index matches
314                (if (i32.ne (local.get $subtask) (i32.load (local.get $wait-retp)))
315                  (then unreachable))
316                ;; Ensure the subtask was cancelled before it returned
317                (if (i32.ne (i32.const 4 (; CANCELLED_BEFORE_RETURNED=4 | (0<<4) ;))
318                            (i32.load offset=4 (local.get $wait-retp)))
319                  (then unreachable))
320
321                ;; Return success
322                (i32.const 42)
323            )
324
325            (func $run (export "run") (result i32)
326                ;; test-id 0: run-yield
327                (if (i32.ne (call $run-test (i32.const 0)) (i32.const 42))
328                    (then unreachable))
329
330                ;; test-id 1: run-yield-to
331                (if (i32.ne (call $run-test (i32.const 1)) (i32.const 42))
332                    (then unreachable))
333
334                ;; test-id 2: run-suspend
335                (if (i32.ne (call $run-test (i32.const 2)) (i32.const 42))
336                    (then unreachable))
337
338                ;; test-id 3: run-switch-to
339                (if (i32.ne (call $run-test (i32.const 3)) (i32.const 42))
340                    (then unreachable))
341
342                ;; Return success
343                (i32.const 42)
344            )
345        )
346
347        (core func $waitable-set.new (canon waitable-set.new))
348        (core func $waitable-set.wait (canon waitable-set.wait (memory $memory "mem")))
349        (core func $waitable.join (canon waitable.join))
350        (core func $subtask.cancel (canon subtask.cancel async))
351        (core func $future.new (canon future.new $FT))
352        (core func $future.write (canon future.write $FT (memory $memory "mem")))
353        (core func $thread.yield (canon thread.yield))
354        (canon lower (func $run-yield) async (memory $memory "mem") (core func $run-yield'))
355        (canon lower (func $run-suspend) async (memory $memory "mem") (core func $run-suspend'))
356        (canon lower (func $run-switch-to) async (memory $memory "mem") (core func $run-switch-to'))
357        (canon lower (func $run-yield-to) async (memory $memory "mem") (core func $run-yield-to'))
358        (core instance $dm (instantiate $DM (with "" (instance
359            (export "mem" (memory $memory "mem"))
360            (export "run-yield" (func $run-yield'))
361            (export "run-suspend" (func $run-suspend'))
362            (export "run-switch-to" (func $run-switch-to'))
363            (export "run-yield-to" (func $run-yield-to'))
364            (export "waitable.join" (func $waitable.join))
365            (export "waitable-set.new" (func $waitable-set.new))
366            (export "waitable-set.wait" (func $waitable-set.wait))
367            (export "subtask.cancel" (func $subtask.cancel))
368            (export "future.new" (func $future.new))
369            (export "future.write" (func $future.write))
370            (export "thread.yield" (func $thread.yield))
371        ))))
372        (func (export "run") async (result u32) (canon lift (core func $dm "run")))
373    )
374
375    (instance $c (instantiate $C))
376    (instance $d (instantiate $D
377        (with "run-yield" (func $c "run-yield"))
378        (with "run-yield-to" (func $c "run-yield-to"))
379        (with "run-suspend" (func $c "run-suspend"))
380        (with "run-switch-to" (func $c "run-switch-to"))
381    ))
382  (func (export "run") (alias export $d "run"))
383)
384
385(assert_return (invoke "run") (u32.const 42))
386