orbit_stream/
capacity.rs

1//! Scoped creator-capacity observation, independent of stream byte readers.
2
3use std::collections::BTreeMap;
4use std::future::Future;
5use std::pin::Pin;
6use std::sync::atomic::Ordering;
7use std::sync::{Arc, Mutex};
8use std::task::{Context, Poll, Waker};
9
10use crate::layout::Doorbell;
11use crate::table::Table;
12use crate::{Result, lock_unpoisoned};
13
14#[derive(Default)]
15pub(crate) struct Waiters(Mutex<Registered>);
16
17#[derive(Default)]
18struct Registered {
19    next: u64,
20    wakers: BTreeMap<u64, Option<Waker>>,
21}
22
23impl Waiters {
24    fn insert(
25        &self,
26        bell: &Doorbell,
27        incarnation: u64,
28    ) -> (u64, u32) {
29        let mut entries = lock_unpoisoned(&self.0);
30        let id = entries.next;
31        entries.next =
32            id.checked_add(1).expect("capacity registration exhausted");
33        let seen = bell.capacity_changes.load(Ordering::SeqCst);
34        entries.wakers.insert(id, None);
35        bell.incarnation.store(incarnation, Ordering::SeqCst);
36        // The caller retries allocation after arming. A concurrent release
37        // either observes this flag or is visible to that retry.
38        bell.capacity_waiting.store(1, Ordering::SeqCst);
39        // Pair the arm with the caller's subsequent Acquire slot scan. If a
40        // releaser saw the old flag, its SeqCst EMPTY publication precedes
41        // this fence and cannot be missed by that retry.
42        std::sync::atomic::fence(Ordering::SeqCst);
43        (id, seen)
44    }
45
46    fn register(
47        &self,
48        id: u64,
49        waker: &Waker,
50    ) {
51        let mut entries = lock_unpoisoned(&self.0);
52        if let Some(place) = entries.wakers.get_mut(&id)
53            && place.as_ref().is_none_or(|old| !old.will_wake(waker))
54        {
55            *place = Some(waker.clone());
56        }
57    }
58
59    fn remove(
60        &self,
61        id: u64,
62        bell: &Doorbell,
63    ) {
64        let mut entries = lock_unpoisoned(&self.0);
65        entries.wakers.remove(&id);
66        if entries.wakers.is_empty() {
67            bell.capacity_waiting.store(0, Ordering::SeqCst);
68        }
69    }
70
71    pub(crate) fn wake(&self) {
72        let wakers: Vec<_> = lock_unpoisoned(&self.0)
73            .wakers
74            .values_mut()
75            .filter_map(Option::take)
76            .collect();
77        for waker in wakers {
78            waker.wake();
79        }
80    }
81}
82
83/// A release in this creator's lane may have freed a slot or made a retained
84/// endpoint rearmable. This is a hint, not a reservation: retry the operation.
85///
86/// Construct **before** retrying a failed create/rearm, then await only if the
87/// retry still cannot proceed. Registration is armed at construction, and
88/// dropping the observer removes it without touching streams or their owners.
89/// Multiple callers may observe the same creator lane independently.
90#[must_use]
91pub struct CapacityChanged {
92    table: Arc<Table>,
93    id: Option<u64>,
94    seen: u32,
95    epoch: u64,
96}
97
98impl CapacityChanged {
99    pub(crate) fn new(table: Arc<Table>) -> Self {
100        let epoch = table.epoch();
101        let (id, seen) =
102            table.capacity.insert(table.own_doorbell(), table.incarnation());
103        Self { table, id: Some(id), seen, epoch }
104    }
105
106    fn changed(&self) -> bool {
107        self.epoch != self.table.epoch()
108            || self.seen
109                != self
110                    .table
111                    .own_doorbell()
112                    .capacity_changes
113                    .load(Ordering::SeqCst)
114    }
115
116    fn finish(&mut self) {
117        if let Some(id) = self.id.take() {
118            self.table.capacity.remove(id, self.table.own_doorbell());
119        }
120    }
121}
122
123impl Future for CapacityChanged {
124    type Output = Result<()>;
125
126    fn poll(
127        self: Pin<&mut Self>,
128        cx: &mut Context<'_>,
129    ) -> Poll<Self::Output> {
130        let this = self.get_mut();
131        let Some(id) = this.id else { return Poll::Ready(Ok(())) };
132        if this.changed() {
133            this.finish();
134            return Poll::Ready(Ok(()));
135        }
136        this.table.capacity.register(id, cx.waker());
137        if let Err(error) = this.table.ensure_driver() {
138            this.finish();
139            return Poll::Ready(Err(error));
140        }
141        if this.changed() {
142            this.finish();
143            Poll::Ready(Ok(()))
144        } else {
145            Poll::Pending
146        }
147    }
148}
149
150impl Drop for CapacityChanged {
151    fn drop(&mut self) {
152        self.finish();
153    }
154}
155
156#[cfg(test)]
157#[path = "../tests/unit/capacity.rs"]
158mod tests;
orbit_stream/capacity.rs

orbit_stream/
capacity.rs
