net_mux/util/
id.rs

1//! Stream identifier types and the per-session allocator.
2
3use std::sync::atomic::{AtomicU32, Ordering};
4
5/// Unique identifier of a logical stream within a session.
6pub type StreamId = u32;
7
8/// Endpoint role of a session.
9#[derive(Debug, Clone, Copy, PartialEq, Eq)]
10pub(crate) enum Role {
11    /// Initiator side: allocates odd stream ids starting at 1.
12    Client,
13    /// Acceptor side: allocates even stream ids starting at 2.
14    Server,
15}
16
17impl Role {
18    pub(crate) fn first_id(self) -> StreamId {
19        match self {
20            Role::Client => 1,
21            Role::Server => 2,
22        }
23    }
24
25    /// Whether the given id was allocated by *this* role.
26    pub(crate) fn owns(self, id: StreamId) -> bool {
27        match self {
28            Role::Client => !id.is_multiple_of(2),
29            Role::Server => id != 0 && id.is_multiple_of(2),
30        }
31    }
32}
33
34/// Lock-free, parity-respecting stream id allocator.
35///
36/// Client-allocated ids are odd (1, 3, 5, …) and server-allocated ids are
37/// even (2, 4, 6, …); id `0` is reserved for session-scoped frames such as
38/// `Ping` and `GoAway`.
39#[derive(Debug)]
40pub(crate) struct StreamIdAllocator {
41    next: AtomicU32,
42}
43
44impl StreamIdAllocator {
45    pub(crate) fn new(role: Role) -> Self {
46        Self {
47            next: AtomicU32::new(role.first_id()),
48        }
49    }
50
51    /// Allocate the next stream id, advancing by 2. Returns `None` once the
52    /// 32-bit id space is exhausted.
53    pub(crate) fn allocate(&self) -> Option<StreamId> {
54        loop {
55            let cur = self.next.load(Ordering::Relaxed);
56            let next = cur.checked_add(2)?;
57            if self
58                .next
59                .compare_exchange_weak(cur, next, Ordering::Relaxed, Ordering::Relaxed)
60                .is_ok()
61            {
62                return Some(cur);
63            }
64        }
65    }
66}
67
68#[cfg(test)]
69mod tests {
70    use super::*;
71
72    #[test]
73    fn allocates_with_correct_parity() {
74        let client = StreamIdAllocator::new(Role::Client);
75        assert_eq!(client.allocate(), Some(1));
76        assert_eq!(client.allocate(), Some(3));
77        assert_eq!(client.allocate(), Some(5));
78
79        let server = StreamIdAllocator::new(Role::Server);
80        assert_eq!(server.allocate(), Some(2));
81        assert_eq!(server.allocate(), Some(4));
82    }
83
84    #[test]
85    fn role_owns_check() {
86        assert!(Role::Client.owns(1));
87        assert!(!Role::Client.owns(2));
88        assert!(Role::Server.owns(2));
89        assert!(!Role::Server.owns(0));
90        assert!(!Role::Server.owns(1));
91    }
92}
net_mux/util/id.rs

net_mux/util/
id.rs
