Skip to main content

rmux_core/events/
registry.rs

1use std::collections::{HashMap, HashSet};
2use std::time::{Duration, Instant};
3
4use rmux_proto::{PaneId, PaneOutputSubscriptionId, SessionName};
5
6/// Default maximum live pane-output subscriptions owned by one connection.
7pub const DEFAULT_MAX_SUBSCRIPTIONS_PER_CONNECTION: usize = 16;
8/// Default maximum live pane-output subscriptions attached to one pane.
9pub const DEFAULT_MAX_SUBSCRIPTIONS_PER_PANE: usize = 64;
10/// Default maximum output events returned by one cursor poll.
11pub const DEFAULT_SUBSCRIPTION_BATCH_EVENTS: usize = 64;
12/// Default idle TTL before a subscription can be removed as stale.
13pub const DEFAULT_SUBSCRIPTION_STALE_TTL: Duration = Duration::from_secs(300);
14
15/// Runtime identity of a pane-output stream used for subscription accounting.
16#[derive(Debug, Clone, PartialEq, Eq, Hash)]
17pub struct PaneOutputSubscriptionKey {
18    runtime_session_name: SessionName,
19    pane_id: PaneId,
20}
21
22impl PaneOutputSubscriptionKey {
23    /// Builds a subscription key from the runtime session owner and pane id.
24    #[must_use]
25    pub fn new(runtime_session_name: SessionName, pane_id: PaneId) -> Self {
26        Self {
27            runtime_session_name,
28            pane_id,
29        }
30    }
31
32    /// Returns the runtime session name that owns the pane output stream.
33    #[must_use]
34    pub fn runtime_session_name(&self) -> &SessionName {
35        &self.runtime_session_name
36    }
37
38    /// Returns the stable pane id.
39    #[must_use]
40    pub const fn pane_id(&self) -> PaneId {
41        self.pane_id
42    }
43}
44
45/// Numeric limits enforced by the pane-output subscription registry.
46#[derive(Debug, Clone, Copy, PartialEq, Eq)]
47pub struct SubscriptionLimits {
48    max_per_connection: usize,
49    max_per_pane: usize,
50    batch_events: usize,
51    stale_ttl: Duration,
52}
53
54impl SubscriptionLimits {
55    /// Builds explicit subscription limits.
56    #[must_use]
57    pub const fn new(
58        max_per_connection: usize,
59        max_per_pane: usize,
60        batch_events: usize,
61        stale_ttl: Duration,
62    ) -> Self {
63        Self {
64            max_per_connection,
65            max_per_pane,
66            batch_events,
67            stale_ttl,
68        }
69    }
70
71    /// Returns the maximum subscriptions one connection may own.
72    #[must_use]
73    pub const fn max_per_connection(self) -> usize {
74        self.max_per_connection
75    }
76
77    /// Returns the maximum subscriptions one pane may have.
78    #[must_use]
79    pub const fn max_per_pane(self) -> usize {
80        self.max_per_pane
81    }
82
83    /// Returns the maximum output events returned by one cursor response.
84    #[must_use]
85    pub const fn batch_events(self) -> usize {
86        self.batch_events
87    }
88
89    /// Returns the idle TTL used by stale cleanup.
90    #[must_use]
91    pub const fn stale_ttl(self) -> Duration {
92        self.stale_ttl
93    }
94}
95
96impl Default for SubscriptionLimits {
97    fn default() -> Self {
98        Self::new(
99            DEFAULT_MAX_SUBSCRIPTIONS_PER_CONNECTION,
100            DEFAULT_MAX_SUBSCRIPTIONS_PER_PANE,
101            DEFAULT_SUBSCRIPTION_BATCH_EVENTS,
102            DEFAULT_SUBSCRIPTION_STALE_TTL,
103        )
104    }
105}
106
107/// Rejection raised when a subscription cap would be exceeded.
108#[derive(Debug, Clone, Copy, PartialEq, Eq)]
109pub enum SubscriptionLimitError {
110    /// The connection already owns the configured maximum subscription count.
111    PerConnection {
112        /// Configured per-connection limit.
113        limit: usize,
114    },
115    /// The pane already has the configured maximum subscription count.
116    PerPane {
117        /// Configured per-pane limit.
118        limit: usize,
119    },
120}
121
122/// Registered subscription metadata.
123#[derive(Debug, Clone, PartialEq, Eq)]
124pub struct OutputSubscriptionRecord {
125    id: PaneOutputSubscriptionId,
126    connection_id: u64,
127    pane: PaneOutputSubscriptionKey,
128    created_at: Instant,
129    last_seen: Instant,
130}
131
132impl OutputSubscriptionRecord {
133    /// Returns the subscription id.
134    #[must_use]
135    pub const fn id(&self) -> PaneOutputSubscriptionId {
136        self.id
137    }
138
139    /// Returns the owning connection id.
140    #[must_use]
141    pub const fn connection_id(&self) -> u64 {
142        self.connection_id
143    }
144
145    /// Returns the subscribed pane key.
146    #[must_use]
147    pub fn pane(&self) -> &PaneOutputSubscriptionKey {
148        &self.pane
149    }
150
151    /// Returns when the subscription was created.
152    #[must_use]
153    pub fn created_at(&self) -> Instant {
154        self.created_at
155    }
156
157    /// Returns when the subscription was last used.
158    #[must_use]
159    pub fn last_seen(&self) -> Instant {
160        self.last_seen
161    }
162}
163
164/// Pane-output subscription registry with connection and pane cap accounting.
165#[derive(Debug, Clone)]
166pub struct SubscriptionRegistry {
167    limits: SubscriptionLimits,
168    next_id: u64,
169    records: HashMap<PaneOutputSubscriptionId, OutputSubscriptionRecord>,
170    by_connection: HashMap<u64, HashSet<PaneOutputSubscriptionId>>,
171    by_pane: HashMap<PaneOutputSubscriptionKey, HashSet<PaneOutputSubscriptionId>>,
172}
173
174impl SubscriptionRegistry {
175    /// Builds an empty registry using explicit limits.
176    #[must_use]
177    pub fn new(limits: SubscriptionLimits) -> Self {
178        Self {
179            limits,
180            next_id: 1,
181            records: HashMap::new(),
182            by_connection: HashMap::new(),
183            by_pane: HashMap::new(),
184        }
185    }
186
187    /// Returns the configured limits.
188    #[must_use]
189    pub const fn limits(&self) -> SubscriptionLimits {
190        self.limits
191    }
192
193    /// Registers a new subscription if both caps allow it.
194    pub fn subscribe(
195        &mut self,
196        connection_id: u64,
197        pane: PaneOutputSubscriptionKey,
198        now: Instant,
199    ) -> Result<OutputSubscriptionRecord, SubscriptionLimitError> {
200        let _ = self.cleanup_stale(now);
201
202        let connection_count = self
203            .by_connection
204            .get(&connection_id)
205            .map_or(0, HashSet::len);
206        if connection_count >= self.limits.max_per_connection {
207            return Err(SubscriptionLimitError::PerConnection {
208                limit: self.limits.max_per_connection,
209            });
210        }
211
212        let pane_count = self.by_pane.get(&pane).map_or(0, HashSet::len);
213        if pane_count >= self.limits.max_per_pane {
214            return Err(SubscriptionLimitError::PerPane {
215                limit: self.limits.max_per_pane,
216            });
217        }
218
219        let id = self.allocate_id();
220        let record = OutputSubscriptionRecord {
221            id,
222            connection_id,
223            pane: pane.clone(),
224            created_at: now,
225            last_seen: now,
226        };
227        self.records.insert(id, record.clone());
228        self.by_connection
229            .entry(connection_id)
230            .or_default()
231            .insert(id);
232        self.by_pane.entry(pane).or_default().insert(id);
233        Ok(record)
234    }
235
236    /// Returns a registered subscription.
237    #[must_use]
238    pub fn get(&self, id: PaneOutputSubscriptionId) -> Option<&OutputSubscriptionRecord> {
239        self.records.get(&id)
240    }
241
242    /// Marks a subscription as recently used and returns its updated record.
243    pub fn touch(
244        &mut self,
245        id: PaneOutputSubscriptionId,
246        now: Instant,
247    ) -> Option<OutputSubscriptionRecord> {
248        let _ = self.cleanup_stale(now);
249        let record = self.records.get_mut(&id)?;
250        record.last_seen = now;
251        Some(record.clone())
252    }
253
254    /// Removes one subscription exactly once.
255    pub fn unsubscribe(
256        &mut self,
257        id: PaneOutputSubscriptionId,
258    ) -> Option<OutputSubscriptionRecord> {
259        let record = self.records.remove(&id)?;
260        self.remove_indexes(&record);
261        Some(record)
262    }
263
264    /// Removes every subscription owned by a connection.
265    pub fn remove_connection(&mut self, connection_id: u64) -> Vec<OutputSubscriptionRecord> {
266        let ids = self
267            .by_connection
268            .remove(&connection_id)
269            .unwrap_or_default()
270            .into_iter()
271            .collect::<Vec<_>>();
272        self.remove_ids(ids)
273    }
274
275    /// Removes every subscription attached to a pane.
276    pub fn remove_pane(
277        &mut self,
278        pane: &PaneOutputSubscriptionKey,
279    ) -> Vec<OutputSubscriptionRecord> {
280        let ids = self
281            .by_pane
282            .remove(pane)
283            .unwrap_or_default()
284            .into_iter()
285            .collect::<Vec<_>>();
286        self.remove_ids(ids)
287    }
288
289    /// Returns whether at least one live subscription targets `pane`.
290    #[must_use]
291    pub fn contains_pane(&self, pane: &PaneOutputSubscriptionKey) -> bool {
292        self.by_pane.get(pane).is_some_and(|ids| !ids.is_empty())
293    }
294
295    /// Returns live subscription ids targeting `pane`.
296    #[must_use]
297    pub fn ids_for_pane(&self, pane: &PaneOutputSubscriptionKey) -> Vec<PaneOutputSubscriptionId> {
298        self.by_pane
299            .get(pane)
300            .map(|ids| ids.iter().copied().collect())
301            .unwrap_or_default()
302    }
303
304    /// Removes subscriptions that have not been touched within the stale TTL.
305    pub fn cleanup_stale(&mut self, now: Instant) -> Vec<OutputSubscriptionRecord> {
306        let ttl = self.limits.stale_ttl;
307        let ids = self
308            .records
309            .iter()
310            .filter_map(|(id, record)| (now.duration_since(record.last_seen) >= ttl).then_some(*id))
311            .collect::<Vec<_>>();
312        self.remove_ids(ids)
313    }
314
315    /// Returns the total number of live subscriptions.
316    #[must_use]
317    pub fn len(&self) -> usize {
318        self.records.len()
319    }
320
321    /// Returns whether the registry is empty.
322    #[must_use]
323    pub fn is_empty(&self) -> bool {
324        self.records.is_empty()
325    }
326
327    fn allocate_id(&mut self) -> PaneOutputSubscriptionId {
328        let id = PaneOutputSubscriptionId::new(self.next_id);
329        self.next_id = self
330            .next_id
331            .checked_add(1)
332            .expect("pane output subscription id space exhausted");
333        id
334    }
335
336    fn remove_ids(
337        &mut self,
338        ids: impl IntoIterator<Item = PaneOutputSubscriptionId>,
339    ) -> Vec<OutputSubscriptionRecord> {
340        let mut removed = Vec::new();
341        for id in ids {
342            if let Some(record) = self.records.remove(&id) {
343                self.remove_indexes(&record);
344                removed.push(record);
345            }
346        }
347        removed
348    }
349
350    fn remove_indexes(&mut self, record: &OutputSubscriptionRecord) {
351        if let Some(ids) = self.by_connection.get_mut(&record.connection_id) {
352            ids.remove(&record.id);
353            if ids.is_empty() {
354                self.by_connection.remove(&record.connection_id);
355            }
356        }
357        if let Some(ids) = self.by_pane.get_mut(&record.pane) {
358            ids.remove(&record.id);
359            if ids.is_empty() {
360                self.by_pane.remove(&record.pane);
361            }
362        }
363    }
364}
365
366impl Default for SubscriptionRegistry {
367    fn default() -> Self {
368        Self::new(SubscriptionLimits::default())
369    }
370}
371
372#[cfg(test)]
373mod tests {
374    use super::*;
375
376    fn session() -> SessionName {
377        SessionName::new("alpha").expect("valid session")
378    }
379
380    fn pane(id: u32) -> PaneOutputSubscriptionKey {
381        PaneOutputSubscriptionKey::new(session(), PaneId::new(id))
382    }
383
384    #[test]
385    fn caps_are_released_exactly_once_across_overlapping_removals() {
386        let limits = SubscriptionLimits::new(1, 1, 64, Duration::from_secs(300));
387        let mut registry = SubscriptionRegistry::new(limits);
388        let now = Instant::now();
389        let first = registry
390            .subscribe(7, pane(1), now)
391            .expect("first subscription");
392
393        assert!(matches!(
394            registry.subscribe(7, pane(2), now),
395            Err(SubscriptionLimitError::PerConnection { limit: 1 })
396        ));
397        assert!(matches!(
398            registry.subscribe(8, pane(1), now),
399            Err(SubscriptionLimitError::PerPane { limit: 1 })
400        ));
401
402        assert_eq!(
403            registry.unsubscribe(first.id()).map(|record| record.id()),
404            Some(first.id())
405        );
406        assert!(registry.unsubscribe(first.id()).is_none());
407
408        let second = registry
409            .subscribe(8, pane(1), now)
410            .expect("cap released after unsubscribe");
411        let removed_by_pane = registry.remove_pane(second.pane());
412        assert_eq!(removed_by_pane.len(), 1);
413        assert_eq!(removed_by_pane[0].id(), second.id());
414        assert!(registry
415            .remove_connection(second.connection_id())
416            .is_empty());
417        assert!(registry.unsubscribe(second.id()).is_none());
418
419        let third = registry
420            .subscribe(9, pane(1), now)
421            .expect("cap released after pane removal");
422        let removed_by_connection = registry.remove_connection(third.connection_id());
423        assert_eq!(removed_by_connection.len(), 1);
424        assert_eq!(removed_by_connection[0].id(), third.id());
425        assert!(registry.remove_pane(third.pane()).is_empty());
426        assert!(registry.subscribe(10, pane(1), now).is_ok());
427    }
428
429    #[test]
430    fn stale_cleanup_releases_connection_and_pane_caps() {
431        let limits = SubscriptionLimits::new(1, 1, 64, Duration::from_millis(10));
432        let mut registry = SubscriptionRegistry::new(limits);
433        let now = Instant::now();
434        let first = registry
435            .subscribe(1, pane(1), now)
436            .expect("first subscription");
437
438        let removed = registry.cleanup_stale(now + Duration::from_millis(10));
439        assert_eq!(removed.len(), 1);
440        assert_eq!(removed[0].id(), first.id());
441        assert!(registry
442            .cleanup_stale(now + Duration::from_millis(20))
443            .is_empty());
444        assert!(registry.subscribe(1, pane(1), now).is_ok());
445    }
446}