1use std::collections::{HashMap, HashSet};
2use std::time::{Duration, Instant};
3
4use rmux_proto::{PaneId, PaneOutputSubscriptionId, SessionName};
5
6pub const DEFAULT_MAX_SUBSCRIPTIONS_PER_CONNECTION: usize = 16;
8pub const DEFAULT_MAX_SUBSCRIPTIONS_PER_PANE: usize = 64;
10pub const DEFAULT_SUBSCRIPTION_BATCH_EVENTS: usize = 64;
12pub const DEFAULT_SUBSCRIPTION_STALE_TTL: Duration = Duration::from_secs(300);
14
15#[derive(Debug, Clone, PartialEq, Eq, Hash)]
17pub struct PaneOutputSubscriptionKey {
18 runtime_session_name: SessionName,
19 pane_id: PaneId,
20}
21
22impl PaneOutputSubscriptionKey {
23 #[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 #[must_use]
34 pub fn runtime_session_name(&self) -> &SessionName {
35 &self.runtime_session_name
36 }
37
38 #[must_use]
40 pub const fn pane_id(&self) -> PaneId {
41 self.pane_id
42 }
43}
44
45#[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 #[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 #[must_use]
73 pub const fn max_per_connection(self) -> usize {
74 self.max_per_connection
75 }
76
77 #[must_use]
79 pub const fn max_per_pane(self) -> usize {
80 self.max_per_pane
81 }
82
83 #[must_use]
85 pub const fn batch_events(self) -> usize {
86 self.batch_events
87 }
88
89 #[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#[derive(Debug, Clone, Copy, PartialEq, Eq)]
109pub enum SubscriptionLimitError {
110 PerConnection {
112 limit: usize,
114 },
115 PerPane {
117 limit: usize,
119 },
120}
121
122#[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 #[must_use]
135 pub const fn id(&self) -> PaneOutputSubscriptionId {
136 self.id
137 }
138
139 #[must_use]
141 pub const fn connection_id(&self) -> u64 {
142 self.connection_id
143 }
144
145 #[must_use]
147 pub fn pane(&self) -> &PaneOutputSubscriptionKey {
148 &self.pane
149 }
150
151 #[must_use]
153 pub fn created_at(&self) -> Instant {
154 self.created_at
155 }
156
157 #[must_use]
159 pub fn last_seen(&self) -> Instant {
160 self.last_seen
161 }
162}
163
164#[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 #[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 #[must_use]
189 pub const fn limits(&self) -> SubscriptionLimits {
190 self.limits
191 }
192
193 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 #[must_use]
238 pub fn get(&self, id: PaneOutputSubscriptionId) -> Option<&OutputSubscriptionRecord> {
239 self.records.get(&id)
240 }
241
242 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 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 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 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 #[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 #[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 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 #[must_use]
317 pub fn len(&self) -> usize {
318 self.records.len()
319 }
320
321 #[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}