priority_semaphore/
waiter.rs

1//! Acquire future and direct-handoff state.
2
3use crate::{
4    error::AcquireError,
5    permit::Permit,
6    queue::WaitKey,
7    semaphore::{Priority, PrioritySemaphore, RegisterResult},
8};
9use alloc::sync::Arc;
10use core::{
11    future::Future,
12    pin::Pin,
13    sync::atomic::{AtomicU8, Ordering},
14    task::{Context, Poll},
15};
16
17const WAITING: u8 = 0;
18const ASSIGNED: u8 = 1;
19const CLOSED: u8 = 2;
20
21/// State shared between a queued future and the thread returning a permit.
22#[derive(Debug)]
23pub(crate) struct Waiter(AtomicU8);
24
25impl Waiter {
26    pub(crate) const fn new() -> Self {
27        Self(AtomicU8::new(WAITING))
28    }
29
30    pub(crate) fn assign(&self) {
31        self.0.store(ASSIGNED, Ordering::Release);
32    }
33
34    pub(crate) fn close(&self) {
35        self.0.store(CLOSED, Ordering::Release);
36    }
37
38    pub(crate) fn is_waiting(&self) -> bool {
39        self.0.load(Ordering::Acquire) == WAITING
40    }
41
42    pub(crate) fn is_assigned(&self) -> bool {
43        self.0.load(Ordering::Acquire) == ASSIGNED
44    }
45
46    fn status(&self) -> u8 {
47        self.0.load(Ordering::Acquire)
48    }
49}
50
51#[derive(Debug)]
52enum Phase {
53    Initial,
54    Waiting { key: WaitKey, waiter: Arc<Waiter> },
55    Complete,
56}
57
58/// Future returned by [`PrioritySemaphore::acquire`](crate::PrioritySemaphore::acquire).
59///
60/// Dropping this future is cancellation-safe in every state, including after
61/// a permit has been assigned but before the executor polls it again.
62#[derive(Debug)]
63#[must_use = "futures do nothing unless polled or awaited"]
64pub struct AcquireFuture {
65    // Option lets a completed future move its existing Arc directly into the
66    // permit instead of paying for an increment/decrement pair per acquire.
67    root: Option<Arc<PrioritySemaphore>>,
68    priority: Priority,
69    phase: Phase,
70}
71
72impl AcquireFuture {
73    pub(crate) fn new(root: Arc<PrioritySemaphore>, priority: Priority) -> Self {
74        Self {
75            root: Some(root),
76            priority,
77            phase: Phase::Initial,
78        }
79    }
80}
81
82impl Future for AcquireFuture {
83    type Output = Result<Permit, AcquireError>;
84
85    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
86        let this = self.get_mut();
87        match &this.phase {
88            Phase::Initial => match this.root.as_ref().unwrap().try_take() {
89                Ok(()) => {
90                    let root = this.root.take().unwrap();
91                    this.phase = Phase::Complete;
92                    Poll::Ready(Ok(Permit::new(root)))
93                }
94                Err(crate::TryAcquireError::Closed) => {
95                    this.root = None;
96                    this.phase = Phase::Complete;
97                    Poll::Ready(Err(AcquireError::Closed))
98                }
99                Err(crate::TryAcquireError::NoPermits) => {
100                    match this
101                        .root
102                        .as_ref()
103                        .unwrap()
104                        .register(this.priority, cx.waker())
105                    {
106                        RegisterResult::Acquired => {
107                            let root = this.root.take().unwrap();
108                            this.phase = Phase::Complete;
109                            Poll::Ready(Ok(Permit::new(root)))
110                        }
111                        RegisterResult::Closed => {
112                            this.root = None;
113                            this.phase = Phase::Complete;
114                            Poll::Ready(Err(AcquireError::Closed))
115                        }
116                        RegisterResult::Queued { key, waiter } => {
117                            this.phase = Phase::Waiting { key, waiter };
118                            Poll::Pending
119                        }
120                    }
121                }
122            },
123            Phase::Waiting { key, waiter } => match waiter.status() {
124                ASSIGNED => {
125                    let root = this.root.take().unwrap();
126                    this.phase = Phase::Complete;
127                    Poll::Ready(Ok(Permit::new(root)))
128                }
129                CLOSED => {
130                    this.root = None;
131                    this.phase = Phase::Complete;
132                    Poll::Ready(Err(AcquireError::Closed))
133                }
134                WAITING => {
135                    this.root
136                        .as_ref()
137                        .unwrap()
138                        .refresh_waker(*key, waiter, cx.waker());
139                    // The status may have changed before refresh_waker took
140                    // the queue lock. In that case the corresponding wake is
141                    // already guaranteed, so Pending remains correct.
142                    Poll::Pending
143                }
144                _ => unreachable!("invalid waiter state"),
145            },
146            Phase::Complete => panic!("AcquireFuture polled after completion"),
147        }
148    }
149}
150
151impl Drop for AcquireFuture {
152    fn drop(&mut self) {
153        if let (Some(root), Phase::Waiting { key, waiter }) = (&self.root, &self.phase) {
154            root.cancel_waiter(*key, waiter);
155        }
156    }
157}
priority_semaphore/waiter.rs

priority_semaphore/
waiter.rs
