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    /// Moves every live subscription from an old runtime owner key to the
290    /// same pane under its new runtime owner without changing subscription ids.
291    pub fn rekey_pane(
292        &mut self,
293        previous: &PaneOutputSubscriptionKey,
294        current: PaneOutputSubscriptionKey,
295    ) -> usize {
296        if previous == &current {
297            return self.by_pane.get(previous).map_or(0, HashSet::len);
298        }
299        let Some(ids) = self.by_pane.remove(previous) else {
300            return 0;
301        };
302        for id in &ids {
303            if let Some(record) = self.records.get_mut(id) {
304                record.pane = current.clone();
305            }
306        }
307        let count = ids.len();
308        self.by_pane.entry(current).or_default().extend(ids);
309        count
310    }
311
312    /// Returns whether at least one live subscription targets `pane`.
313    #[must_use]
314    pub fn contains_pane(&self, pane: &PaneOutputSubscriptionKey) -> bool {
315        self.by_pane.get(pane).is_some_and(|ids| !ids.is_empty())
316    }
317
318    /// Returns live subscription ids targeting `pane`.
319    #[must_use]
320    pub fn ids_for_pane(&self, pane: &PaneOutputSubscriptionKey) -> Vec<PaneOutputSubscriptionId> {
321        self.by_pane
322            .get(pane)
323            .map(|ids| ids.iter().copied().collect())
324            .unwrap_or_default()
325    }
326
327    /// Removes subscriptions that have not been touched within the stale TTL.
328    pub fn cleanup_stale(&mut self, now: Instant) -> Vec<OutputSubscriptionRecord> {
329        let ttl = self.limits.stale_ttl;
330        let ids = self
331            .records
332            .iter()
333            .filter_map(|(id, record)| (now.duration_since(record.last_seen) >= ttl).then_some(*id))
334            .collect::<Vec<_>>();
335        self.remove_ids(ids)
336    }
337
338    /// Returns the total number of live subscriptions.
339    #[must_use]
340    pub fn len(&self) -> usize {
341        self.records.len()
342    }
343
344    /// Returns whether the registry is empty.
345    #[must_use]
346    pub fn is_empty(&self) -> bool {
347        self.records.is_empty()
348    }
349
350    fn allocate_id(&mut self) -> PaneOutputSubscriptionId {
351        let id = PaneOutputSubscriptionId::new(self.next_id);
352        self.next_id = self
353            .next_id
354            .checked_add(1)
355            .expect("pane output subscription id space exhausted");
356        id
357    }
358
359    fn remove_ids(
360        &mut self,
361        ids: impl IntoIterator<Item = PaneOutputSubscriptionId>,
362    ) -> Vec<OutputSubscriptionRecord> {
363        let mut removed = Vec::new();
364        for id in ids {
365            if let Some(record) = self.records.remove(&id) {
366                self.remove_indexes(&record);
367                removed.push(record);
368            }
369        }
370        removed
371    }
372
373    fn remove_indexes(&mut self, record: &OutputSubscriptionRecord) {
374        if let Some(ids) = self.by_connection.get_mut(&record.connection_id) {
375            ids.remove(&record.id);
376            if ids.is_empty() {
377                self.by_connection.remove(&record.connection_id);
378            }
379        }
380        if let Some(ids) = self.by_pane.get_mut(&record.pane) {
381            ids.remove(&record.id);
382            if ids.is_empty() {
383                self.by_pane.remove(&record.pane);
384            }
385        }
386    }
387}
388
389impl Default for SubscriptionRegistry {
390    fn default() -> Self {
391        Self::new(SubscriptionLimits::default())
392    }
393}
394
395#[cfg(test)]
396mod tests {
397    use super::*;
398
399    fn session() -> SessionName {
400        SessionName::new("alpha").expect("valid session")
401    }
402
403    fn pane(id: u32) -> PaneOutputSubscriptionKey {
404        PaneOutputSubscriptionKey::new(session(), PaneId::new(id))
405    }
406
407    #[test]
408    fn caps_are_released_exactly_once_across_overlapping_removals() {
409        let limits = SubscriptionLimits::new(1, 1, 64, Duration::from_secs(300));
410        let mut registry = SubscriptionRegistry::new(limits);
411        let now = Instant::now();
412        let first = registry
413            .subscribe(7, pane(1), now)
414            .expect("first subscription");
415
416        assert!(matches!(
417            registry.subscribe(7, pane(2), now),
418            Err(SubscriptionLimitError::PerConnection { limit: 1 })
419        ));
420        assert!(matches!(
421            registry.subscribe(8, pane(1), now),
422            Err(SubscriptionLimitError::PerPane { limit: 1 })
423        ));
424
425        assert_eq!(
426            registry.unsubscribe(first.id()).map(|record| record.id()),
427            Some(first.id())
428        );
429        assert!(registry.unsubscribe(first.id()).is_none());
430
431        let second = registry
432            .subscribe(8, pane(1), now)
433            .expect("cap released after unsubscribe");
434        let removed_by_pane = registry.remove_pane(second.pane());
435        assert_eq!(removed_by_pane.len(), 1);
436        assert_eq!(removed_by_pane[0].id(), second.id());
437        assert!(registry
438            .remove_connection(second.connection_id())
439            .is_empty());
440        assert!(registry.unsubscribe(second.id()).is_none());
441
442        let third = registry
443            .subscribe(9, pane(1), now)
444            .expect("cap released after pane removal");
445        let removed_by_connection = registry.remove_connection(third.connection_id());
446        assert_eq!(removed_by_connection.len(), 1);
447        assert_eq!(removed_by_connection[0].id(), third.id());
448        assert!(registry.remove_pane(third.pane()).is_empty());
449        assert!(registry.subscribe(10, pane(1), now).is_ok());
450    }
451
452    #[test]
453    fn stale_cleanup_releases_connection_and_pane_caps() {
454        let limits = SubscriptionLimits::new(1, 1, 64, Duration::from_millis(10));
455        let mut registry = SubscriptionRegistry::new(limits);
456        let now = Instant::now();
457        let first = registry
458            .subscribe(1, pane(1), now)
459            .expect("first subscription");
460
461        let removed = registry.cleanup_stale(now + Duration::from_millis(10));
462        assert_eq!(removed.len(), 1);
463        assert_eq!(removed[0].id(), first.id());
464        assert!(registry
465            .cleanup_stale(now + Duration::from_millis(20))
466            .is_empty());
467        assert!(registry.subscribe(1, pane(1), now).is_ok());
468    }
469
470    #[test]
471    fn runtime_owner_rekey_preserves_ids_and_all_registry_indexes() {
472        let mut registry = SubscriptionRegistry::default();
473        let now = Instant::now();
474        let previous = pane(7);
475        let current = PaneOutputSubscriptionKey::new(
476            SessionName::new("beta").expect("valid session"),
477            PaneId::new(7),
478        );
479        let first = registry
480            .subscribe(11, previous.clone(), now)
481            .expect("first subscription");
482        let second = registry
483            .subscribe(12, previous.clone(), now)
484            .expect("second subscription");
485
486        assert_eq!(registry.rekey_pane(&previous, current.clone()), 2);
487        assert!(!registry.contains_pane(&previous));
488        assert!(registry.contains_pane(&current));
489        assert_eq!(
490            registry.get(first.id()).map(OutputSubscriptionRecord::pane),
491            Some(&current)
492        );
493        assert_eq!(
494            registry
495                .get(second.id())
496                .map(OutputSubscriptionRecord::pane),
497            Some(&current)
498        );
499        assert_eq!(registry.remove_pane(&current).len(), 2);
500        assert!(registry.remove_connection(11).is_empty());
501        assert!(registry.remove_connection(12).is_empty());
502    }
503}
rmux_core/events/registry.rs

rmux_core/events/
registry.rs
