1#[cfg(not(target_arch = "wasm32"))]
5pub struct NativeCancellationToken {
6 inner: tokio_util::sync::CancellationToken,
7}
8
9#[cfg(not(target_arch = "wasm32"))]
10impl Clone for NativeCancellationToken {
11 fn clone(&self) -> Self {
12 Self {
13 inner: self.inner.clone(),
14 }
15 }
16}
17
18#[cfg(not(target_arch = "wasm32"))]
19impl NativeCancellationToken {
20 pub fn new() -> Self {
21 Self {
22 inner: tokio_util::sync::CancellationToken::new(),
23 }
24 }
25
26 pub fn cancel(&self) {
27 self.inner.cancel();
28 }
29
30 pub async fn cancelled(&self) {
31 self.inner.cancelled().await;
32 }
33}
34
35#[cfg(not(target_arch = "wasm32"))]
36impl Default for NativeCancellationToken {
37 fn default() -> Self {
38 Self::new()
39 }
40}
41
42#[cfg(target_arch = "wasm32")]
46pub struct WasmCancellationToken {
47 inner: std::sync::Arc<std::sync::atomic::AtomicBool>,
48}
49
50#[cfg(target_arch = "wasm32")]
51impl Clone for WasmCancellationToken {
52 fn clone(&self) -> Self {
53 Self {
54 inner: self.inner.clone(),
55 }
56 }
57}
58
59#[cfg(target_arch = "wasm32")]
60impl WasmCancellationToken {
61 pub fn new() -> Self {
62 Self {
63 inner: std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false)),
64 }
65 }
66
67 pub fn cancel(&self) {
68 self.inner
69 .store(true, std::sync::atomic::Ordering::SeqCst);
70 }
71
72 pub async fn cancelled(&self) {
73 loop {
74 if self.inner.load(std::sync::atomic::Ordering::SeqCst) {
75 break;
76 }
77 crate::sleep(std::time::Duration::from_millis(10)).await;
78 }
79 }
80}
81
82#[cfg(target_arch = "wasm32")]
83impl Default for WasmCancellationToken {
84 fn default() -> Self {
85 Self::new()
86 }
87}
88
89#[cfg(not(target_arch = "wasm32"))]
90pub use NativeCancellationToken as CancellationToken;
91
92#[cfg(target_arch = "wasm32")]
93pub use WasmCancellationToken as CancellationToken;
94
95#[cfg(test)]
96mod test {
97 use std::time::Duration;
98
99 use crate::timeout;
100
101 use super::*;
102
103 #[cfg(target_arch = "wasm32")]
104 use wasm_bindgen_test::wasm_bindgen_test;
105 #[cfg(target_arch = "wasm32")]
106 wasm_bindgen_test::wasm_bindgen_test_configure!(run_in_browser);
107
108 #[cfg_attr(not(target_arch = "wasm32"), tokio::test)]
110 #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test)]
111 async fn cancelled_returns_after_cancel() {
112 let token = CancellationToken::new();
113 token.cancel();
114 token.cancelled().await;
115 }
116
117 #[cfg_attr(not(target_arch = "wasm32"), tokio::test)]
119 #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test)]
120 async fn cancel_wakes_a_pending_waiter() {
121 let token = CancellationToken::new();
122 let waiter = token.clone();
123 crate::spawn(async move {
124 crate::sleep(Duration::from_millis(10)).await;
125 token.cancel();
126 });
127
128 timeout(Duration::from_secs(5), waiter.cancelled()).await
129 .expect("cancellation should have woken the waiter");
130 }
131
132 #[cfg_attr(not(target_arch = "wasm32"), tokio::test)]
134 #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test)]
135 async fn clone_shares_the_cancellation_state() {
136 let token = CancellationToken::new();
137 let clone = token.clone();
138 token.cancel();
139 clone.cancelled().await;
140 }
141
142 #[cfg_attr(not(target_arch = "wasm32"), tokio::test)]
144 #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test)]
145 async fn cancelling_a_clone_cancels_the_original() {
146 let token = CancellationToken::new();
147 let clone = token.clone();
148 clone.cancel();
149 token.cancelled().await;
150 }
151
152 #[cfg_attr(not(target_arch = "wasm32"), tokio::test)]
154 #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test)]
155 async fn cancelled_stays_pending_without_cancel() {
156 let token = CancellationToken::new();
157 timeout(Duration::from_millis(50), token.cancelled()).await
158 .expect_err("an uncancelled token should keep waiting");
159 }
160
161 #[cfg_attr(not(target_arch = "wasm32"), tokio::test)]
163 #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test)]
164 async fn default_token_is_not_cancelled() {
165 let token = CancellationToken::default();
166 timeout(Duration::from_millis(50), token.cancelled()).await
167 .expect_err("a fresh token should keep waiting");
168 token.cancel();
169 token.cancelled().await;
170 }
171}
172