forked from tokio-rs/tokio
-
Notifications
You must be signed in to change notification settings - Fork 0
/
task_local.rs
64 lines (50 loc) · 1.33 KB
/
task_local.rs
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
#![cfg(feature = "full")]
tokio::task_local! {
static REQ_ID: u32;
pub static FOO: bool;
}
#[tokio::test(flavor = "multi_thread")]
async fn local() {
let j1 = tokio::spawn(REQ_ID.scope(1, async move {
assert_eq!(REQ_ID.get(), 1);
assert_eq!(REQ_ID.get(), 1);
}));
let j2 = tokio::spawn(REQ_ID.scope(2, async move {
REQ_ID.with(|v| {
assert_eq!(REQ_ID.get(), 2);
assert_eq!(*v, 2);
});
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
assert_eq!(REQ_ID.get(), 2);
}));
let j3 = tokio::spawn(FOO.scope(true, async move {
assert!(FOO.get());
}));
j1.await.unwrap();
j2.await.unwrap();
j3.await.unwrap();
}
tokio::task_local! {
static KEY: u32;
}
struct Guard(u32);
impl Drop for Guard {
fn drop(&mut self) {
assert_eq!(KEY.try_with(|x| *x), Ok(self.0));
}
}
#[tokio::test(flavor = "multi_thread")]
async fn task_local_available_on_drop() {
let (tx, rx) = tokio::sync::oneshot::channel();
let h = tokio::spawn(KEY.scope(42, async move {
let _g = Guard(42);
let _ = tx.send(());
std::future::pending::<()>().await;
}));
rx.await.unwrap();
h.abort();
let err = h.await.unwrap_err();
if err.is_panic() {
panic!("{}", err);
}
}