rmux_core/events/
wait.rs

1use std::collections::{HashMap, HashSet};
2
3use rmux_proto::{SdkWaitId, SdkWaitOwnerId};
4
5/// Stable key for one daemon-backed SDK wait.
6#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
7pub struct SdkWaitKey {
8    owner_id: SdkWaitOwnerId,
9    wait_id: SdkWaitId,
10}
11
12impl SdkWaitKey {
13    /// Builds a wait key from its owner and per-owner id.
14    #[must_use]
15    pub const fn new(owner_id: SdkWaitOwnerId, wait_id: SdkWaitId) -> Self {
16        Self { owner_id, wait_id }
17    }
18
19    /// Returns the SDK owner id.
20    #[must_use]
21    pub const fn owner_id(self) -> SdkWaitOwnerId {
22        self.owner_id
23    }
24
25    /// Returns the per-owner wait id.
26    #[must_use]
27    pub const fn wait_id(self) -> SdkWaitId {
28        self.wait_id
29    }
30}
31
32/// Registered daemon-backed SDK wait metadata.
33#[derive(Debug, Clone, Copy, PartialEq, Eq)]
34pub struct SdkWaitRecord {
35    key: SdkWaitKey,
36    connection_id: u64,
37}
38
39impl SdkWaitRecord {
40    /// Returns the stable wait key.
41    #[must_use]
42    pub const fn key(self) -> SdkWaitKey {
43        self.key
44    }
45
46    /// Returns the server-private connection id that owns this wait.
47    #[must_use]
48    pub const fn connection_id(self) -> u64 {
49        self.connection_id
50    }
51}
52
53/// Registry for daemon-backed SDK wait identities and cleanup accounting.
54#[derive(Debug, Clone, Default)]
55pub struct SdkWaitRegistry {
56    next_id_by_owner: HashMap<SdkWaitOwnerId, u64>,
57    records: HashMap<SdkWaitKey, SdkWaitRecord>,
58    by_connection: HashMap<u64, HashSet<SdkWaitKey>>,
59}
60
61impl SdkWaitRegistry {
62    /// Allocates the next wait id for one SDK owner.
63    ///
64    /// IDs intentionally start at one so a zero value remains conspicuous in
65    /// logs and diagnostics. The same numeric id can exist under different
66    /// owners, which is the per-connection scoping contract.
67    pub fn allocate_id(&mut self, owner_id: SdkWaitOwnerId) -> SdkWaitId {
68        let next = self.next_id_by_owner.entry(owner_id).or_insert(1);
69        let wait_id = SdkWaitId::new(*next);
70        *next = next
71            .checked_add(1)
72            .expect("SDK wait id space exhausted for owner");
73        wait_id
74    }
75
76    /// Registers an active wait.
77    ///
78    /// Returns `false` when the `(owner_id, wait_id)` pair is already active.
79    /// Duplicate registration is rejected without disturbing the existing
80    /// record.
81    pub fn register(
82        &mut self,
83        connection_id: u64,
84        owner_id: SdkWaitOwnerId,
85        wait_id: SdkWaitId,
86    ) -> bool {
87        let key = SdkWaitKey::new(owner_id, wait_id);
88        if self.records.contains_key(&key) {
89            return false;
90        }
91
92        let record = SdkWaitRecord { key, connection_id };
93        self.records.insert(key, record);
94        self.by_connection
95            .entry(connection_id)
96            .or_default()
97            .insert(key);
98        true
99    }
100
101    /// Removes one wait by owner and id.
102    ///
103    /// Duplicate and late cancellations are idempotent: `None` means there was
104    /// no live wait to remove.
105    pub fn remove(
106        &mut self,
107        owner_id: SdkWaitOwnerId,
108        wait_id: SdkWaitId,
109    ) -> Option<SdkWaitRecord> {
110        self.remove_key(SdkWaitKey::new(owner_id, wait_id))
111    }
112
113    /// Removes every wait owned by an actual server connection.
114    pub fn remove_connection(&mut self, connection_id: u64) -> Vec<SdkWaitRecord> {
115        let keys = self
116            .by_connection
117            .remove(&connection_id)
118            .unwrap_or_default()
119            .into_iter()
120            .collect::<Vec<_>>();
121        keys.into_iter()
122            .filter_map(|key| self.remove_record_without_connection_index(key))
123            .collect()
124    }
125
126    /// Returns a live wait record.
127    #[must_use]
128    pub fn get(&self, owner_id: SdkWaitOwnerId, wait_id: SdkWaitId) -> Option<SdkWaitRecord> {
129        self.records
130            .get(&SdkWaitKey::new(owner_id, wait_id))
131            .copied()
132    }
133
134    /// Returns the number of active waits.
135    #[must_use]
136    pub fn len(&self) -> usize {
137        self.records.len()
138    }
139
140    /// Returns whether no waits are registered.
141    #[must_use]
142    pub fn is_empty(&self) -> bool {
143        self.records.is_empty()
144    }
145
146    fn remove_key(&mut self, key: SdkWaitKey) -> Option<SdkWaitRecord> {
147        let record = self.records.remove(&key)?;
148        if let Some(keys) = self.by_connection.get_mut(&record.connection_id) {
149            keys.remove(&key);
150            if keys.is_empty() {
151                self.by_connection.remove(&record.connection_id);
152            }
153        }
154        Some(record)
155    }
156
157    fn remove_record_without_connection_index(&mut self, key: SdkWaitKey) -> Option<SdkWaitRecord> {
158        self.records.remove(&key)
159    }
160}
161
162#[cfg(test)]
163mod tests {
164    use super::*;
165
166    fn owner(value: u64) -> SdkWaitOwnerId {
167        SdkWaitOwnerId::new(value)
168    }
169
170    fn wait(value: u64) -> SdkWaitId {
171        SdkWaitId::new(value)
172    }
173
174    #[test]
175    fn wait_ids_are_unique_within_owner_and_scoped_across_owners() {
176        let mut registry = SdkWaitRegistry::default();
177
178        assert_eq!(registry.allocate_id(owner(10)), wait(1));
179        assert_eq!(registry.allocate_id(owner(10)), wait(2));
180        assert_eq!(registry.allocate_id(owner(11)), wait(1));
181        assert_eq!(registry.allocate_id(owner(10)), wait(3));
182    }
183
184    #[test]
185    fn duplicate_late_and_repeated_cancel_are_idempotent() {
186        let mut registry = SdkWaitRegistry::default();
187
188        assert!(registry.register(7, owner(1), wait(1)));
189        assert!(!registry.register(8, owner(1), wait(1)));
190        assert_eq!(registry.len(), 1);
191
192        let removed = registry
193            .remove(owner(1), wait(1))
194            .expect("first cancel removes wait");
195        assert_eq!(removed.connection_id(), 7);
196        assert!(registry.remove(owner(1), wait(1)).is_none());
197        assert!(registry.remove(owner(1), wait(9)).is_none());
198        assert!(registry.is_empty());
199    }
200
201    #[test]
202    fn connection_teardown_removes_only_that_connections_waits() {
203        let mut registry = SdkWaitRegistry::default();
204        assert!(registry.register(1, owner(10), wait(1)));
205        assert!(registry.register(1, owner(10), wait(2)));
206        assert!(registry.register(2, owner(20), wait(1)));
207
208        let removed = registry.remove_connection(1);
209        let removed_keys = removed
210            .iter()
211            .map(|record| record.key())
212            .collect::<HashSet<_>>();
213        assert_eq!(
214            removed_keys,
215            HashSet::from([
216                SdkWaitKey::new(owner(10), wait(1)),
217                SdkWaitKey::new(owner(10), wait(2)),
218            ])
219        );
220        assert_eq!(registry.len(), 1);
221        assert!(registry.get(owner(20), wait(1)).is_some());
222        assert!(registry.remove_connection(1).is_empty());
223        assert_eq!(registry.remove_connection(2).len(), 1);
224        assert!(registry.is_empty());
225    }
226}
rmux_core/events/wait.rs

rmux_core/events/
wait.rs
