1 mod bindings {
2     wit_bindgen::generate!({
3         path: "../misc/component-async-tests/wit",
4         world: "cancel-callee",
5     });
6 }
7 
8 use test_programs::async_::{
9     CALLBACK_CODE_EXIT, CALLBACK_CODE_WAIT, EVENT_CANCELLED, EVENT_NONE, EVENT_SUBTASK,
10     STATUS_RETURN_CANCELLED, STATUS_RETURNED, STATUS_STARTED, context_get, context_set,
11     subtask_cancel, subtask_drop, task_cancel, waitable_join, waitable_set_drop, waitable_set_new,
12 };
13 
14 #[cfg(target_arch = "wasm32")]
15 #[link(wasm_import_module = "[export]local:local/yield-with-options")]
16 unsafe extern "C" {
17     #[link_name = "[task-return]yield-times"]
task_return_yield_times()18     fn task_return_yield_times();
19 }
20 #[cfg(not(target_arch = "wasm32"))]
task_return_yield_times()21 unsafe extern "C" fn task_return_yield_times() {
22     unreachable!()
23 }
24 
25 #[cfg(target_arch = "wasm32")]
26 #[link(wasm_import_module = "local:local/yield")]
27 unsafe extern "C" {
28     #[link_name = "yield-times"]
yield_times(_: u64)29     fn yield_times(_: u64);
30 }
31 #[cfg(not(target_arch = "wasm32"))]
yield_times(_: u64)32 unsafe fn yield_times(_: u64) {
33     unreachable!()
34 }
35 
36 #[cfg(target_arch = "wasm32")]
37 #[link(wasm_import_module = "local:local/yield")]
38 unsafe extern "C" {
39     #[link_name = "[async-lower]yield-times"]
yield_times_async(ms: u64) -> u3240     fn yield_times_async(ms: u64) -> u32;
41 }
42 #[cfg(not(target_arch = "wasm32"))]
yield_times_async(_ms: u64) -> u3243 unsafe fn yield_times_async(_ms: u64) -> u32 {
44     unreachable!()
45 }
46 
47 const ON_CANCEL_TASK_RETURN: u8 = 0;
48 const ON_CANCEL_TASK_CANCEL: u8 = 1;
49 
50 const _MODE_NORMAL: u8 = 0;
51 const _MODE_TRAP_CANCEL_GUEST_AFTER_START_CANCELLED: u8 = 1;
52 const _MODE_TRAP_CANCEL_GUEST_AFTER_RETURN_CANCELLED: u8 = 2;
53 const _MODE_TRAP_CANCEL_GUEST_AFTER_RETURN: u8 = 3;
54 const MODE_TRAP_CANCEL_HOST_AFTER_RETURN_CANCELLED: u8 = 4;
55 const MODE_TRAP_CANCEL_HOST_AFTER_RETURN: u8 = 5;
56 const MODE_LEAK_TASK_AFTER_CANCEL: u8 = 6;
57 
58 #[derive(Clone, Copy)]
59 struct YieldParams {
60     times: u64,
61     on_cancel: u8,
62     on_cancel_delay_times: u64,
63     synchronous_delay: bool,
64     mode: u8,
65 }
66 
67 enum State {
68     S0(YieldParams),
69     S1 {
70         set: u32,
71         waitable: u32,
72         params: YieldParams,
73     },
74     S2 {
75         set: u32,
76         waitable: u32,
77         params: YieldParams,
78     },
79 }
80 
81 #[unsafe(export_name = "local:local/backpressure#set-backpressure")]
export_set_backpressure(enabled: bool)82 unsafe extern "C" fn export_set_backpressure(enabled: bool) {
83     if enabled {
84         wit_bindgen::backpressure_inc();
85     } else {
86         wit_bindgen::backpressure_dec();
87     }
88 }
89 
90 #[unsafe(export_name = "local:local/backpressure#inc-backpressure")]
export_inc_backpressure()91 unsafe extern "C" fn export_inc_backpressure() {
92     wit_bindgen::backpressure_inc();
93 }
94 
95 #[unsafe(export_name = "local:local/backpressure#dec-backpressure")]
export_dec_backpressure()96 unsafe extern "C" fn export_dec_backpressure() {
97     wit_bindgen::backpressure_dec();
98 }
99 
100 #[unsafe(export_name = "local:local/yield#yield-times")]
export_yield_yield_times(times: u64)101 unsafe extern "C" fn export_yield_yield_times(times: u64) {
102     unsafe {
103         yield_times(times);
104     }
105 }
106 
107 #[unsafe(export_name = "[async-lift]local:local/yield-with-options#yield-times")]
export_yield_with_options_yield_times( times: u64, on_cancel: u8, on_cancel_delay_times: u64, synchronous_delay: bool, mode: u8, ) -> u32108 unsafe extern "C" fn export_yield_with_options_yield_times(
109     times: u64,
110     on_cancel: u8,
111     on_cancel_delay_times: u64,
112     synchronous_delay: bool,
113     mode: u8,
114 ) -> u32 {
115     unsafe {
116         context_set(
117             u32::try_from(Box::into_raw(Box::new(State::S0(YieldParams {
118                 times,
119                 on_cancel,
120                 on_cancel_delay_times,
121                 synchronous_delay,
122                 mode,
123             }))) as usize)
124             .unwrap(),
125         );
126         callback_yield_with_options_yield_times(EVENT_NONE, 0, 0)
127     }
128 }
129 
130 #[unsafe(export_name = "[callback][async-lift]local:local/yield-with-options#yield-times")]
callback_yield_with_options_yield_times( event0: u32, event1: u32, event2: u32, ) -> u32131 unsafe extern "C" fn callback_yield_with_options_yield_times(
132     event0: u32,
133     event1: u32,
134     event2: u32,
135 ) -> u32 {
136     unsafe {
137         let state = &mut *(usize::try_from(context_get()).unwrap() as *mut State);
138         match state {
139             State::S0(params) => {
140                 assert_eq!(event0, EVENT_NONE);
141 
142                 let status = yield_times_async(params.times);
143 
144                 let waitable = status >> 4;
145                 let status = status & 0xF;
146 
147                 assert_eq!(status, STATUS_STARTED);
148 
149                 let set = waitable_set_new();
150                 waitable_join(waitable, set);
151 
152                 *state = State::S1 {
153                     set,
154                     waitable,
155                     params: *params,
156                 };
157 
158                 CALLBACK_CODE_WAIT | (set << 4)
159             }
160 
161             State::S1 {
162                 set,
163                 waitable,
164                 params,
165             } => {
166                 assert_eq!(event0, EVENT_CANCELLED);
167 
168                 let result = subtask_cancel(*waitable);
169 
170                 assert_eq!(result, STATUS_RETURN_CANCELLED);
171 
172                 if params.mode == MODE_TRAP_CANCEL_HOST_AFTER_RETURN_CANCELLED {
173                     // This should trap, since `waitable` has already been
174                     // cancelled:
175                     subtask_cancel(*waitable);
176                     unreachable!()
177                 }
178 
179                 waitable_join(*waitable, 0);
180 
181                 if params.mode != MODE_LEAK_TASK_AFTER_CANCEL {
182                     subtask_drop(*waitable);
183                 }
184 
185                 if params.on_cancel_delay_times == 0 {
186                     match params.on_cancel {
187                         ON_CANCEL_TASK_RETURN => task_return_yield_times(),
188                         ON_CANCEL_TASK_CANCEL => task_cancel(),
189                         _ => unreachable!(),
190                     }
191 
192                     CALLBACK_CODE_EXIT
193                 } else if params.synchronous_delay {
194                     yield_times(params.on_cancel_delay_times);
195 
196                     match params.on_cancel {
197                         ON_CANCEL_TASK_RETURN => task_return_yield_times(),
198                         ON_CANCEL_TASK_CANCEL => task_cancel(),
199                         _ => unreachable!(),
200                     }
201 
202                     CALLBACK_CODE_EXIT
203                 } else {
204                     let status = yield_times_async(params.on_cancel_delay_times);
205 
206                     let waitable = status >> 4;
207                     let status = status & 0xF;
208 
209                     assert_eq!(status, STATUS_STARTED);
210 
211                     waitable_join(waitable, *set);
212 
213                     let set = *set;
214 
215                     *state = State::S2 {
216                         set,
217                         waitable,
218                         params: *params,
219                     };
220 
221                     CALLBACK_CODE_WAIT | (set << 4)
222                 }
223             }
224 
225             State::S2 {
226                 set,
227                 waitable,
228                 params,
229             } => {
230                 assert_eq!(event0, EVENT_SUBTASK);
231                 assert_eq!(event1, *waitable);
232                 assert_eq!(event2, STATUS_RETURNED);
233 
234                 if params.mode == MODE_TRAP_CANCEL_HOST_AFTER_RETURN {
235                     // This should trap, since `waitable` has already returned:
236                     subtask_cancel(*waitable);
237                     unreachable!()
238                 }
239 
240                 waitable_join(*waitable, 0);
241                 subtask_drop(*waitable);
242                 waitable_set_drop(*set);
243 
244                 match params.on_cancel {
245                     ON_CANCEL_TASK_RETURN => task_return_yield_times(),
246                     ON_CANCEL_TASK_CANCEL => task_cancel(),
247                     _ => unreachable!(),
248                 }
249 
250                 CALLBACK_CODE_EXIT
251             }
252         }
253     }
254 }
255 
256 // Unused function; required since this file is built as a `bin`:
main()257 fn main() {}
258