Coverage Report

Created: 2025-10-29 07:05

next uncovered line (L), next uncovered region (R), next uncovered branch (B)
/rust/registry/src/index.crates.io-1949cf8c6b5b557f/tokio-1.44.2/src/runtime/park.rs
Line
Count
Source
1
#![cfg_attr(not(feature = "full"), allow(dead_code))]
2
3
use crate::loom::sync::atomic::AtomicUsize;
4
use crate::loom::sync::{Arc, Condvar, Mutex};
5
6
use std::sync::atomic::Ordering::SeqCst;
7
use std::time::Duration;
8
9
#[derive(Debug)]
10
pub(crate) struct ParkThread {
11
    inner: Arc<Inner>,
12
}
13
14
/// Unblocks a thread that was blocked by `ParkThread`.
15
#[derive(Clone, Debug)]
16
pub(crate) struct UnparkThread {
17
    inner: Arc<Inner>,
18
}
19
20
#[derive(Debug)]
21
struct Inner {
22
    state: AtomicUsize,
23
    mutex: Mutex<()>,
24
    condvar: Condvar,
25
}
26
27
const EMPTY: usize = 0;
28
const PARKED: usize = 1;
29
const NOTIFIED: usize = 2;
30
31
tokio_thread_local! {
32
    static CURRENT_PARKER: ParkThread = ParkThread::new();
33
}
34
35
// Bit of a hack, but it is only for loom
36
#[cfg(loom)]
37
tokio_thread_local! {
38
    pub(crate) static CURRENT_THREAD_PARK_COUNT: AtomicUsize = AtomicUsize::new(0);
39
}
40
41
// ==== impl ParkThread ====
42
43
impl ParkThread {
44
0
    pub(crate) fn new() -> Self {
45
0
        Self {
46
0
            inner: Arc::new(Inner {
47
0
                state: AtomicUsize::new(EMPTY),
48
0
                mutex: Mutex::new(()),
49
0
                condvar: Condvar::new(),
50
0
            }),
51
0
        }
52
0
    }
53
54
0
    pub(crate) fn unpark(&self) -> UnparkThread {
55
0
        let inner = self.inner.clone();
56
0
        UnparkThread { inner }
57
0
    }
58
59
0
    pub(crate) fn park(&mut self) {
60
        #[cfg(loom)]
61
        CURRENT_THREAD_PARK_COUNT.with(|count| count.fetch_add(1, SeqCst));
62
0
        self.inner.park();
63
0
    }
64
65
0
    pub(crate) fn park_timeout(&mut self, duration: Duration) {
66
        #[cfg(loom)]
67
        CURRENT_THREAD_PARK_COUNT.with(|count| count.fetch_add(1, SeqCst));
68
0
        self.inner.park_timeout(duration);
69
0
    }
70
71
0
    pub(crate) fn shutdown(&mut self) {
72
0
        self.inner.shutdown();
73
0
    }
74
}
75
76
// ==== impl Inner ====
77
78
impl Inner {
79
0
    fn park(&self) {
80
        // If we were previously notified then we consume this notification and
81
        // return quickly.
82
0
        if self
83
0
            .state
84
0
            .compare_exchange(NOTIFIED, EMPTY, SeqCst, SeqCst)
85
0
            .is_ok()
86
        {
87
0
            return;
88
0
        }
89
90
        // Otherwise we need to coordinate going to sleep
91
0
        let mut m = self.mutex.lock();
92
93
0
        match self.state.compare_exchange(EMPTY, PARKED, SeqCst, SeqCst) {
94
0
            Ok(_) => {}
95
            Err(NOTIFIED) => {
96
                // We must read here, even though we know it will be `NOTIFIED`.
97
                // This is because `unpark` may have been called again since we read
98
                // `NOTIFIED` in the `compare_exchange` above. We must perform an
99
                // acquire operation that synchronizes with that `unpark` to observe
100
                // any writes it made before the call to unpark. To do that we must
101
                // read from the write it made to `state`.
102
0
                let old = self.state.swap(EMPTY, SeqCst);
103
0
                debug_assert_eq!(old, NOTIFIED, "park state changed unexpectedly");
104
105
0
                return;
106
            }
107
0
            Err(actual) => panic!("inconsistent park state; actual = {actual}"),
108
        }
109
110
        loop {
111
0
            m = self.condvar.wait(m).unwrap();
112
113
0
            if self
114
0
                .state
115
0
                .compare_exchange(NOTIFIED, EMPTY, SeqCst, SeqCst)
116
0
                .is_ok()
117
            {
118
                // got a notification
119
0
                return;
120
0
            }
121
122
            // spurious wakeup, go back to sleep
123
        }
124
0
    }
125
126
    /// Parks the current thread for at most `dur`.
127
0
    fn park_timeout(&self, dur: Duration) {
128
        // Like `park` above we have a fast path for an already-notified thread,
129
        // and afterwards we start coordinating for a sleep. Return quickly.
130
0
        if self
131
0
            .state
132
0
            .compare_exchange(NOTIFIED, EMPTY, SeqCst, SeqCst)
133
0
            .is_ok()
134
        {
135
0
            return;
136
0
        }
137
138
0
        if dur == Duration::from_millis(0) {
139
0
            return;
140
0
        }
141
142
0
        let m = self.mutex.lock();
143
144
0
        match self.state.compare_exchange(EMPTY, PARKED, SeqCst, SeqCst) {
145
0
            Ok(_) => {}
146
            Err(NOTIFIED) => {
147
                // We must read again here, see `park`.
148
0
                let old = self.state.swap(EMPTY, SeqCst);
149
0
                debug_assert_eq!(old, NOTIFIED, "park state changed unexpectedly");
150
151
0
                return;
152
            }
153
0
            Err(actual) => panic!("inconsistent park_timeout state; actual = {actual}"),
154
        }
155
156
        #[cfg(not(all(target_family = "wasm", not(target_feature = "atomics"))))]
157
        // Wait with a timeout, and if we spuriously wake up or otherwise wake up
158
        // from a notification, we just want to unconditionally set the state back to
159
        // empty, either consuming a notification or un-flagging ourselves as
160
        // parked.
161
0
        let (_m, _result) = self.condvar.wait_timeout(m, dur).unwrap();
162
163
        #[cfg(all(target_family = "wasm", not(target_feature = "atomics")))]
164
        // Wasm without atomics doesn't have threads, so just sleep.
165
        {
166
            let _m = m;
167
            std::thread::sleep(dur);
168
        }
169
170
0
        match self.state.swap(EMPTY, SeqCst) {
171
0
            NOTIFIED => {} // got a notification, hurray!
172
0
            PARKED => {}   // no notification, alas
173
0
            n => panic!("inconsistent park_timeout state: {n}"),
174
        }
175
0
    }
176
177
0
    fn unpark(&self) {
178
        // To ensure the unparked thread will observe any writes we made before
179
        // this call, we must perform a release operation that `park` can
180
        // synchronize with. To do that we must write `NOTIFIED` even if `state`
181
        // is already `NOTIFIED`. That is why this must be a swap rather than a
182
        // compare-and-swap that returns if it reads `NOTIFIED` on failure.
183
0
        match self.state.swap(NOTIFIED, SeqCst) {
184
0
            EMPTY => return,    // no one was waiting
185
0
            NOTIFIED => return, // already unparked
186
0
            PARKED => {}        // gotta go wake someone up
187
0
            _ => panic!("inconsistent state in unpark"),
188
        }
189
190
        // There is a period between when the parked thread sets `state` to
191
        // `PARKED` (or last checked `state` in the case of a spurious wake
192
        // up) and when it actually waits on `cvar`. If we were to notify
193
        // during this period it would be ignored and then when the parked
194
        // thread went to sleep it would never wake up. Fortunately, it has
195
        // `lock` locked at this stage so we can acquire `lock` to wait until
196
        // it is ready to receive the notification.
197
        //
198
        // Releasing `lock` before the call to `notify_one` means that when the
199
        // parked thread wakes it doesn't get woken only to have to wait for us
200
        // to release `lock`.
201
0
        drop(self.mutex.lock());
202
203
0
        self.condvar.notify_one();
204
0
    }
205
206
0
    fn shutdown(&self) {
207
0
        self.condvar.notify_all();
208
0
    }
209
}
210
211
impl Default for ParkThread {
212
0
    fn default() -> Self {
213
0
        Self::new()
214
0
    }
215
}
216
217
// ===== impl UnparkThread =====
218
219
impl UnparkThread {
220
0
    pub(crate) fn unpark(&self) {
221
0
        self.inner.unpark();
222
0
    }
223
}
224
225
use crate::loom::thread::AccessError;
226
use std::future::Future;
227
use std::marker::PhantomData;
228
use std::rc::Rc;
229
use std::task::{RawWaker, RawWakerVTable, Waker};
230
231
/// Blocks the current thread using a condition variable.
232
#[derive(Debug)]
233
pub(crate) struct CachedParkThread {
234
    _anchor: PhantomData<Rc<()>>,
235
}
236
237
impl CachedParkThread {
238
    /// Creates a new `ParkThread` handle for the current thread.
239
    ///
240
    /// This type cannot be moved to other threads, so it should be created on
241
    /// the thread that the caller intends to park.
242
0
    pub(crate) fn new() -> CachedParkThread {
243
0
        CachedParkThread {
244
0
            _anchor: PhantomData,
245
0
        }
246
0
    }
247
248
0
    pub(crate) fn waker(&self) -> Result<Waker, AccessError> {
249
0
        self.unpark().map(UnparkThread::into_waker)
250
0
    }
251
252
0
    fn unpark(&self) -> Result<UnparkThread, AccessError> {
253
0
        self.with_current(ParkThread::unpark)
254
0
    }
255
256
0
    pub(crate) fn park(&mut self) {
257
0
        self.with_current(|park_thread| park_thread.inner.park())
258
0
            .unwrap();
259
0
    }
260
261
0
    pub(crate) fn park_timeout(&mut self, duration: Duration) {
262
0
        self.with_current(|park_thread| park_thread.inner.park_timeout(duration))
263
0
            .unwrap();
264
0
    }
265
266
    /// Gets a reference to the `ParkThread` handle for this thread.
267
0
    fn with_current<F, R>(&self, f: F) -> Result<R, AccessError>
268
0
    where
269
0
        F: FnOnce(&ParkThread) -> R,
270
    {
271
0
        CURRENT_PARKER.try_with(|inner| f(inner))
Unexecuted instantiation: <tokio::runtime::park::CachedParkThread>::with_current::<<tokio::runtime::park::CachedParkThread>::park_timeout::{closure#0}, ()>::{closure#0}
Unexecuted instantiation: <tokio::runtime::park::CachedParkThread>::with_current::<<tokio::runtime::park::CachedParkThread>::park::{closure#0}, ()>::{closure#0}
Unexecuted instantiation: <tokio::runtime::park::CachedParkThread>::with_current::<<tokio::runtime::park::ParkThread>::unpark, tokio::runtime::park::UnparkThread>::{closure#0}
272
0
    }
Unexecuted instantiation: <tokio::runtime::park::CachedParkThread>::with_current::<<tokio::runtime::park::CachedParkThread>::park_timeout::{closure#0}, ()>
Unexecuted instantiation: <tokio::runtime::park::CachedParkThread>::with_current::<<tokio::runtime::park::CachedParkThread>::park::{closure#0}, ()>
Unexecuted instantiation: <tokio::runtime::park::CachedParkThread>::with_current::<<tokio::runtime::park::ParkThread>::unpark, tokio::runtime::park::UnparkThread>
273
274
0
    pub(crate) fn block_on<F: Future>(&mut self, f: F) -> Result<F::Output, AccessError> {
275
        use std::task::Context;
276
        use std::task::Poll::Ready;
277
278
0
        let waker = self.waker()?;
279
0
        let mut cx = Context::from_waker(&waker);
280
281
0
        pin!(f);
282
283
        loop {
284
0
            if let Ready(v) = crate::task::coop::budget(|| f.as_mut().poll(&mut cx)) {
Unexecuted instantiation: <tokio::runtime::park::CachedParkThread>::block_on::<core::pin::Pin<alloc::boxed::Box<tracing::instrument::Instrumented<ztunnel::app::new_data_plane_pool::{closure#0}::{closure#1}>>>>::{closure#0}
Unexecuted instantiation: <tokio::runtime::park::CachedParkThread>::block_on::<tracing::instrument::Instrumented<ztunnel::app::new_data_plane_pool::{closure#0}::{closure#1}>>::{closure#0}
Unexecuted instantiation: <tokio::runtime::park::CachedParkThread>::block_on::<core::future::poll_fn::PollFn<<tokio::runtime::scheduler::current_thread::CurrentThread>::block_on<core::pin::Pin<alloc::boxed::Box<tracing::instrument::Instrumented<ztunnel::app::new_data_plane_pool::{closure#0}::{closure#1}>>>>::{closure#0}::{closure#0}>>::{closure#0}
Unexecuted instantiation: <tokio::runtime::park::CachedParkThread>::block_on::<core::future::poll_fn::PollFn<<tokio::runtime::scheduler::current_thread::CurrentThread>::block_on<tracing::instrument::Instrumented<ztunnel::app::new_data_plane_pool::{closure#0}::{closure#1}>>::{closure#0}::{closure#0}>>::{closure#0}
Unexecuted instantiation: <tokio::runtime::park::CachedParkThread>::block_on::<&mut tokio::sync::oneshot::Receiver<()>>::{closure#0}
285
0
                return Ok(v);
286
0
            }
287
288
0
            self.park();
289
        }
290
0
    }
Unexecuted instantiation: <tokio::runtime::park::CachedParkThread>::block_on::<core::pin::Pin<alloc::boxed::Box<tracing::instrument::Instrumented<ztunnel::app::new_data_plane_pool::{closure#0}::{closure#1}>>>>
Unexecuted instantiation: <tokio::runtime::park::CachedParkThread>::block_on::<tracing::instrument::Instrumented<ztunnel::app::new_data_plane_pool::{closure#0}::{closure#1}>>
Unexecuted instantiation: <tokio::runtime::park::CachedParkThread>::block_on::<core::future::poll_fn::PollFn<<tokio::runtime::scheduler::current_thread::CurrentThread>::block_on<core::pin::Pin<alloc::boxed::Box<tracing::instrument::Instrumented<ztunnel::app::new_data_plane_pool::{closure#0}::{closure#1}>>>>::{closure#0}::{closure#0}>>
Unexecuted instantiation: <tokio::runtime::park::CachedParkThread>::block_on::<core::future::poll_fn::PollFn<<tokio::runtime::scheduler::current_thread::CurrentThread>::block_on<tracing::instrument::Instrumented<ztunnel::app::new_data_plane_pool::{closure#0}::{closure#1}>>::{closure#0}::{closure#0}>>
Unexecuted instantiation: <tokio::runtime::park::CachedParkThread>::block_on::<&mut tokio::sync::oneshot::Receiver<()>>
291
}
292
293
impl UnparkThread {
294
0
    pub(crate) fn into_waker(self) -> Waker {
295
        unsafe {
296
0
            let raw = unparker_to_raw_waker(self.inner);
297
0
            Waker::from_raw(raw)
298
        }
299
0
    }
300
}
301
302
impl Inner {
303
    #[allow(clippy::wrong_self_convention)]
304
0
    fn into_raw(this: Arc<Inner>) -> *const () {
305
0
        Arc::into_raw(this) as *const ()
306
0
    }
307
308
0
    unsafe fn from_raw(ptr: *const ()) -> Arc<Inner> {
309
0
        Arc::from_raw(ptr as *const Inner)
310
0
    }
311
}
312
313
0
unsafe fn unparker_to_raw_waker(unparker: Arc<Inner>) -> RawWaker {
314
0
    RawWaker::new(
315
0
        Inner::into_raw(unparker),
316
0
        &RawWakerVTable::new(clone, wake, wake_by_ref, drop_waker),
317
    )
318
0
}
319
320
0
unsafe fn clone(raw: *const ()) -> RawWaker {
321
0
    Arc::increment_strong_count(raw as *const Inner);
322
0
    unparker_to_raw_waker(Inner::from_raw(raw))
323
0
}
324
325
0
unsafe fn drop_waker(raw: *const ()) {
326
0
    drop(Inner::from_raw(raw));
327
0
}
328
329
0
unsafe fn wake(raw: *const ()) {
330
0
    let unparker = Inner::from_raw(raw);
331
0
    unparker.unpark();
332
0
}
333
334
0
unsafe fn wake_by_ref(raw: *const ()) {
335
0
    let raw = raw as *const Inner;
336
0
    (*raw).unpark();
337
0
}
338
339
#[cfg(loom)]
340
pub(crate) fn current_thread_park_count() -> usize {
341
    CURRENT_THREAD_PARK_COUNT.with(|count| count.load(SeqCst))
342
}