Skip to main content

bark_runtime/
cancel.rs

1/// A token for signaling cancellation on native platforms.
2///
3/// Wraps `tokio_util::sync::CancellationToken` for efficient async cancellation.
4#[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/// A token for signaling cancellation on WASM platforms.
43///
44/// Uses `Arc<AtomicBool>` with polling-based waiting.
45#[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	/// Waiting on a cancelled token returns.
109	#[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	/// Cancelling while a waiter is pending wakes the waiter.
118	#[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	/// A clone reports the cancellation triggered on the original token.
133	#[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	/// Cancelling a clone also cancels the original token.
143	#[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	/// Waiting on a token that was never cancelled never returns.
153	#[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	/// A default token starts out uncancelled.
162	#[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