1use std::collections::{BTreeMap, HashMap};
4use std::hash::Hash;
5
6use crate::command::{CommandEnvelope, CommandRejectReason};
7use crate::ids::{ClientId, StationId, Tick};
8
9const HASHED_GATEWAY_SESSION_MIN_ENTRIES: usize = 1_024;
10
11#[derive(Clone, Debug)]
12enum AdaptiveSessionMap<K, V> {
13 Ordered(BTreeMap<K, V>),
14 Hashed(HashMap<K, V>),
15}
16
17impl<K: Copy + Eq + Hash + Ord, V> AdaptiveSessionMap<K, V> {
18 fn new() -> Self {
19 Self::Ordered(BTreeMap::new())
20 }
21
22 fn len(&self) -> usize {
23 match self {
24 Self::Ordered(entries) => entries.len(),
25 Self::Hashed(entries) => entries.len(),
26 }
27 }
28
29 fn is_empty(&self) -> bool {
30 match self {
31 Self::Ordered(entries) => entries.is_empty(),
32 Self::Hashed(entries) => entries.is_empty(),
33 }
34 }
35
36 fn get(&self, key: &K) -> Option<&V> {
37 match self {
38 Self::Ordered(entries) => entries.get(key),
39 Self::Hashed(entries) => entries.get(key),
40 }
41 }
42
43 fn get_mut(&mut self, key: &K) -> Option<&mut V> {
44 match self {
45 Self::Ordered(entries) => entries.get_mut(key),
46 Self::Hashed(entries) => entries.get_mut(key),
47 }
48 }
49
50 fn insert(&mut self, key: K, value: V) -> Option<V> {
51 let promote = match self {
52 Self::Ordered(entries) => {
53 entries.len() >= HASHED_GATEWAY_SESSION_MIN_ENTRIES.saturating_sub(1)
54 && !entries.contains_key(&key)
55 }
56 Self::Hashed(_) => false,
57 };
58 if promote {
59 let Self::Ordered(ordered) = std::mem::replace(self, Self::Hashed(HashMap::new()))
60 else {
61 unreachable!("promotion starts from ordered session storage");
62 };
63 let mut hashed = HashMap::with_capacity(ordered.len().saturating_add(1));
64 hashed.extend(ordered);
65 *self = Self::Hashed(hashed);
66 }
67 match self {
68 Self::Ordered(entries) => entries.insert(key, value),
69 Self::Hashed(entries) => entries.insert(key, value),
70 }
71 }
72
73 fn retain<F>(&mut self, mut keep: F)
74 where
75 F: FnMut(&K, &mut V) -> bool,
76 {
77 match self {
78 Self::Ordered(entries) => entries.retain(|key, value| keep(key, value)),
79 Self::Hashed(entries) => entries.retain(|key, value| keep(key, value)),
80 }
81 }
82
83 #[cfg(test)]
84 fn is_hashed(&self) -> bool {
85 matches!(self, Self::Hashed(_))
86 }
87}
88
89#[derive(Clone, Copy, Debug, PartialEq, Eq)]
91pub struct GatewayConfig {
92 pub max_sessions: usize,
94 pub reconnect_grace_ticks: u64,
97 pub max_commands_per_tick: usize,
99}
100
101impl Default for GatewayConfig {
102 fn default() -> Self {
103 Self {
104 max_sessions: 65_536,
105 reconnect_grace_ticks: 20 * 60,
106 max_commands_per_tick: 64,
107 }
108 }
109}
110
111#[derive(Clone, Copy, Debug, PartialEq, Eq)]
113pub struct GatewayRoute {
114 pub client_id: ClientId,
116 pub station_id: StationId,
118 pub generation: u64,
120 pub route_epoch: u64,
122}
123
124#[derive(Clone, Copy, Debug, PartialEq, Eq)]
126pub enum GatewaySessionState {
127 Connected,
129 Disconnected {
131 since: Tick,
133 },
134}
135
136#[derive(Clone, Debug, PartialEq, Eq)]
138pub struct GatewaySession {
139 pub client_id: ClientId,
141 pub station_id: StationId,
143 pub connected_at: Tick,
145 pub last_seen: Tick,
147 pub generation: u64,
150 pub route_epoch: u64,
152 pub state: GatewaySessionState,
154 last_sequence: Option<u64>,
155 command_tick: Tick,
156 commands_this_tick: usize,
157}
158
159impl GatewaySession {
160 pub const fn route(&self) -> GatewayRoute {
162 GatewayRoute {
163 client_id: self.client_id,
164 station_id: self.station_id,
165 generation: self.generation,
166 route_epoch: self.route_epoch,
167 }
168 }
169
170 pub const fn last_sequence(&self) -> Option<u64> {
172 self.last_sequence
173 }
174
175 pub const fn commands_this_tick(&self) -> usize {
177 self.commands_this_tick
178 }
179
180 pub const fn is_connected(&self) -> bool {
182 matches!(self.state, GatewaySessionState::Connected)
183 }
184}
185
186#[derive(Clone, Copy, Debug, PartialEq, Eq)]
188pub enum GatewayConnectOutcome {
189 Created,
191 AlreadyConnected,
193 Reconnected {
195 disconnected_for: u64,
197 },
198 ReplacedExpired {
200 disconnected_for: u64,
202 },
203}
204
205#[derive(Clone, Copy, Debug, PartialEq, Eq)]
207pub struct GatewayConnectReport {
208 pub outcome: GatewayConnectOutcome,
210 pub route: GatewayRoute,
212}
213
214#[derive(Clone, Copy, Debug, PartialEq, Eq)]
216pub struct GatewayCommandAdmission {
217 pub route: GatewayRoute,
219 pub sequence: u64,
221 pub commands_this_tick: usize,
223}
224
225#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
227pub struct GatewayStats {
228 pub sessions_created: usize,
230 pub sessions_reconnected: usize,
232 pub sessions_expired: usize,
234 pub routes_changed: usize,
236 pub commands_admitted: usize,
238 pub commands_rejected_replay: usize,
240 pub commands_rejected_rate_limit: usize,
242}
243
244#[derive(Clone, Copy, Debug, PartialEq, Eq)]
246pub enum GatewayError {
247 CapacityFull {
249 capacity: usize,
251 },
252 MissingSession(ClientId),
254 SessionDisconnected {
256 client_id: ClientId,
258 since: Tick,
260 },
261 BadGeneration {
263 expected: u64,
265 actual: u64,
267 },
268 ReplayOrStale {
270 last_sequence: Option<u64>,
272 sequence: u64,
274 },
275 RateLimited {
277 limit: usize,
279 attempted: usize,
281 },
282}
283
284impl GatewayError {
285 pub const fn command_reject_reason(&self) -> Option<CommandRejectReason> {
287 match self {
288 Self::ReplayOrStale { .. } => Some(CommandRejectReason::ReplayOrStale),
289 Self::RateLimited { .. } => Some(CommandRejectReason::RateLimited),
290 Self::MissingSession(_)
291 | Self::SessionDisconnected { .. }
292 | Self::BadGeneration { .. }
293 | Self::CapacityFull { .. } => None,
294 }
295 }
296}
297
298impl core::fmt::Display for GatewayError {
299 fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
300 match self {
301 Self::CapacityFull { capacity } => {
302 write!(f, "gateway session table is full at capacity {capacity}")
303 }
304 Self::MissingSession(client_id) => {
305 write!(
306 f,
307 "gateway session for client {} is missing",
308 client_id.get()
309 )
310 }
311 Self::SessionDisconnected { client_id, since } => write!(
312 f,
313 "gateway session for client {} disconnected at tick {}",
314 client_id.get(),
315 since.get()
316 ),
317 Self::BadGeneration { expected, actual } => write!(
318 f,
319 "gateway reconnect generation mismatch: expected {expected}, actual {actual}"
320 ),
321 Self::ReplayOrStale {
322 last_sequence,
323 sequence,
324 } => write!(
325 f,
326 "gateway command sequence {sequence} is not newer than {last_sequence:?}"
327 ),
328 Self::RateLimited { limit, attempted } => write!(
329 f,
330 "gateway command rate limited: limit {limit}, attempted {attempted}"
331 ),
332 }
333 }
334}
335
336impl std::error::Error for GatewayError {}
337
338#[derive(Clone, Debug)]
340pub struct GatewaySessionTable {
341 config: GatewayConfig,
342 sessions: AdaptiveSessionMap<ClientId, GatewaySession>,
343 stats: GatewayStats,
344}
345
346impl GatewaySessionTable {
347 pub fn new(config: GatewayConfig) -> Self {
349 Self {
350 config,
351 sessions: AdaptiveSessionMap::new(),
352 stats: GatewayStats::default(),
353 }
354 }
355
356 pub const fn config(&self) -> GatewayConfig {
358 self.config
359 }
360
361 pub const fn stats(&self) -> GatewayStats {
363 self.stats
364 }
365
366 pub fn len(&self) -> usize {
368 self.sessions.len()
369 }
370
371 pub fn is_empty(&self) -> bool {
373 self.sessions.is_empty()
374 }
375
376 pub fn session(&self, client_id: ClientId) -> Option<&GatewaySession> {
378 self.sessions.get(&client_id)
379 }
380
381 pub fn route(&self, client_id: ClientId) -> Result<GatewayRoute, GatewayError> {
383 let session = self
384 .sessions
385 .get(&client_id)
386 .ok_or(GatewayError::MissingSession(client_id))?;
387 match session.state {
388 GatewaySessionState::Connected => Ok(session.route()),
389 GatewaySessionState::Disconnected { since } => {
390 Err(GatewayError::SessionDisconnected { client_id, since })
391 }
392 }
393 }
394
395 pub fn connect(
397 &mut self,
398 client_id: ClientId,
399 station_id: StationId,
400 now: Tick,
401 ) -> Result<GatewayConnectReport, GatewayError> {
402 let reconnect_grace_ticks = self.config.reconnect_grace_ticks;
403 if let Some(session) = self.sessions.get_mut(&client_id) {
404 return Ok(Self::connect_existing(
405 session,
406 station_id,
407 now,
408 reconnect_grace_ticks,
409 &mut self.stats,
410 ));
411 }
412
413 if self.sessions.len() >= self.config.max_sessions {
414 return Err(GatewayError::CapacityFull {
415 capacity: self.config.max_sessions,
416 });
417 }
418
419 let session = GatewaySession {
420 client_id,
421 station_id,
422 connected_at: now,
423 last_seen: now,
424 generation: 1,
425 route_epoch: 1,
426 state: GatewaySessionState::Connected,
427 last_sequence: None,
428 command_tick: now,
429 commands_this_tick: 0,
430 };
431 let route = session.route();
432 self.sessions.insert(client_id, session);
433 self.stats.sessions_created = self.stats.sessions_created.saturating_add(1);
434 Ok(GatewayConnectReport {
435 outcome: GatewayConnectOutcome::Created,
436 route,
437 })
438 }
439
440 pub fn reconnect(
442 &mut self,
443 client_id: ClientId,
444 generation: u64,
445 now: Tick,
446 ) -> Result<GatewayConnectReport, GatewayError> {
447 let session = self
448 .sessions
449 .get_mut(&client_id)
450 .ok_or(GatewayError::MissingSession(client_id))?;
451 if session.generation != generation {
452 return Err(GatewayError::BadGeneration {
453 expected: session.generation,
454 actual: generation,
455 });
456 }
457
458 match session.state {
459 GatewaySessionState::Connected => {
460 session.last_seen = now;
461 Ok(GatewayConnectReport {
462 outcome: GatewayConnectOutcome::AlreadyConnected,
463 route: session.route(),
464 })
465 }
466 GatewaySessionState::Disconnected { since } => {
467 let disconnected_for = now.get().saturating_sub(since.get());
468 if disconnected_for > self.config.reconnect_grace_ticks {
469 session.connected_at = now;
470 session.last_seen = now;
471 session.generation = session.generation.saturating_add(1);
472 session.route_epoch = session.route_epoch.saturating_add(1);
473 session.state = GatewaySessionState::Connected;
474 session.last_sequence = None;
475 session.command_tick = now;
476 session.commands_this_tick = 0;
477 self.stats.sessions_expired = self.stats.sessions_expired.saturating_add(1);
478 Ok(GatewayConnectReport {
479 outcome: GatewayConnectOutcome::ReplacedExpired { disconnected_for },
480 route: session.route(),
481 })
482 } else {
483 session.last_seen = now;
484 session.state = GatewaySessionState::Connected;
485 self.stats.sessions_reconnected =
486 self.stats.sessions_reconnected.saturating_add(1);
487 Ok(GatewayConnectReport {
488 outcome: GatewayConnectOutcome::Reconnected { disconnected_for },
489 route: session.route(),
490 })
491 }
492 }
493 }
494 }
495
496 pub fn disconnect(&mut self, client_id: ClientId, now: Tick) -> Result<(), GatewayError> {
498 let session = self
499 .sessions
500 .get_mut(&client_id)
501 .ok_or(GatewayError::MissingSession(client_id))?;
502 session.last_seen = now;
503 session.state = GatewaySessionState::Disconnected { since: now };
504 Ok(())
505 }
506
507 pub fn expire_disconnected(&mut self, now: Tick) -> usize {
509 let grace = self.config.reconnect_grace_ticks;
510 let before = self.sessions.len();
511 self.sessions.retain(|_, session| match session.state {
512 GatewaySessionState::Connected => true,
513 GatewaySessionState::Disconnected { since } => {
514 now.get().saturating_sub(since.get()) <= grace
515 }
516 });
517 let expired = before - self.sessions.len();
518 self.stats.sessions_expired = self.stats.sessions_expired.saturating_add(expired);
519 expired
520 }
521
522 pub fn reroute(
524 &mut self,
525 client_id: ClientId,
526 station_id: StationId,
527 now: Tick,
528 ) -> Result<GatewayRoute, GatewayError> {
529 let session = self
530 .sessions
531 .get_mut(&client_id)
532 .ok_or(GatewayError::MissingSession(client_id))?;
533 match session.state {
534 GatewaySessionState::Connected => {
535 session.last_seen = now;
536 if session.station_id != station_id {
537 session.station_id = station_id;
538 session.route_epoch = session.route_epoch.saturating_add(1);
539 self.stats.routes_changed = self.stats.routes_changed.saturating_add(1);
540 }
541 Ok(session.route())
542 }
543 GatewaySessionState::Disconnected { since } => {
544 Err(GatewayError::SessionDisconnected { client_id, since })
545 }
546 }
547 }
548
549 pub fn admit_command(
552 &mut self,
553 command: &CommandEnvelope,
554 ) -> Result<GatewayCommandAdmission, GatewayError> {
555 self.admit_sequence(command.client_id, command.sequence, command.received_at)
556 }
557
558 pub fn admit_sequence(
560 &mut self,
561 client_id: ClientId,
562 sequence: u64,
563 now: Tick,
564 ) -> Result<GatewayCommandAdmission, GatewayError> {
565 let session = self
566 .sessions
567 .get_mut(&client_id)
568 .ok_or(GatewayError::MissingSession(client_id))?;
569 match session.state {
570 GatewaySessionState::Connected => {}
571 GatewaySessionState::Disconnected { since } => {
572 return Err(GatewayError::SessionDisconnected { client_id, since });
573 }
574 }
575
576 if session
577 .last_sequence
578 .is_some_and(|last_sequence| sequence <= last_sequence)
579 {
580 self.stats.commands_rejected_replay =
581 self.stats.commands_rejected_replay.saturating_add(1);
582 return Err(GatewayError::ReplayOrStale {
583 last_sequence: session.last_sequence,
584 sequence,
585 });
586 }
587
588 if session.command_tick != now {
589 session.command_tick = now;
590 session.commands_this_tick = 0;
591 }
592 let attempted = session.commands_this_tick.saturating_add(1);
593 if attempted > self.config.max_commands_per_tick {
594 self.stats.commands_rejected_rate_limit =
595 self.stats.commands_rejected_rate_limit.saturating_add(1);
596 return Err(GatewayError::RateLimited {
597 limit: self.config.max_commands_per_tick,
598 attempted,
599 });
600 }
601
602 session.commands_this_tick = attempted;
603 session.last_sequence = Some(sequence);
604 session.last_seen = now;
605 self.stats.commands_admitted = self.stats.commands_admitted.saturating_add(1);
606 Ok(GatewayCommandAdmission {
607 route: session.route(),
608 sequence,
609 commands_this_tick: attempted,
610 })
611 }
612
613 fn connect_existing(
614 session: &mut GatewaySession,
615 station_id: StationId,
616 now: Tick,
617 reconnect_grace_ticks: u64,
618 stats: &mut GatewayStats,
619 ) -> GatewayConnectReport {
620 match session.state {
621 GatewaySessionState::Connected => {
622 session.last_seen = now;
623 if session.station_id != station_id {
624 session.station_id = station_id;
625 session.route_epoch = session.route_epoch.saturating_add(1);
626 stats.routes_changed = stats.routes_changed.saturating_add(1);
627 }
628 GatewayConnectReport {
629 outcome: GatewayConnectOutcome::AlreadyConnected,
630 route: session.route(),
631 }
632 }
633 GatewaySessionState::Disconnected { since } => {
634 let disconnected_for = now.get().saturating_sub(since.get());
635 if disconnected_for > reconnect_grace_ticks {
636 session.station_id = station_id;
637 session.connected_at = now;
638 session.last_seen = now;
639 session.generation = session.generation.saturating_add(1);
640 session.route_epoch = session.route_epoch.saturating_add(1);
641 session.state = GatewaySessionState::Connected;
642 session.last_sequence = None;
643 session.command_tick = now;
644 session.commands_this_tick = 0;
645 stats.sessions_expired = stats.sessions_expired.saturating_add(1);
646 GatewayConnectReport {
647 outcome: GatewayConnectOutcome::ReplacedExpired { disconnected_for },
648 route: session.route(),
649 }
650 } else {
651 session.station_id = station_id;
652 session.last_seen = now;
653 session.state = GatewaySessionState::Connected;
654 stats.sessions_reconnected = stats.sessions_reconnected.saturating_add(1);
655 GatewayConnectReport {
656 outcome: GatewayConnectOutcome::Reconnected { disconnected_for },
657 route: session.route(),
658 }
659 }
660 }
661 }
662 }
663}
664
665impl Default for GatewaySessionTable {
666 fn default() -> Self {
667 Self::new(GatewayConfig::default())
668 }
669}
670
671#[cfg(test)]
672mod tests {
673 use super::*;
674 use crate::command::{CommandEnvelope, CommandPriority};
675 use crate::ids::{CommandId, EntityId};
676
677 fn command(client_id: ClientId, sequence: u64, tick: u64) -> CommandEnvelope {
678 CommandEnvelope {
679 id: CommandId::new(sequence),
680 client_id,
681 entity_id: EntityId::new(10),
682 sequence,
683 received_at: Tick::new(tick),
684 kind: 1,
685 priority: CommandPriority::Normal,
686 payload: Vec::new(),
687 }
688 }
689
690 #[test]
691 fn connects_routes_and_reroutes_sessions() {
692 let client_id = ClientId::new(7);
693 let mut table = GatewaySessionTable::new(GatewayConfig {
694 max_sessions: 4,
695 reconnect_grace_ticks: 3,
696 max_commands_per_tick: 4,
697 });
698
699 let connected = table
700 .connect(client_id, StationId::new(1), Tick::new(10))
701 .expect("connect should work");
702 assert_eq!(connected.outcome, GatewayConnectOutcome::Created);
703 assert_eq!(connected.route.station_id, StationId::new(1));
704 assert_eq!(connected.route.generation, 1);
705 assert_eq!(connected.route.route_epoch, 1);
706
707 let route = table
708 .reroute(client_id, StationId::new(2), Tick::new(11))
709 .expect("reroute should work");
710 assert_eq!(route.station_id, StationId::new(2));
711 assert_eq!(route.route_epoch, 2);
712 assert_eq!(table.stats().routes_changed, 1);
713 assert_eq!(
714 table
715 .route(client_id)
716 .expect("route should exist")
717 .station_id,
718 StationId::new(2)
719 );
720 }
721
722 #[test]
723 fn gateway_sessions_promote_without_losing_route_admission_or_expiry_state() {
724 let mut table = GatewaySessionTable::new(GatewayConfig {
725 max_sessions: HASHED_GATEWAY_SESSION_MIN_ENTRIES + 1,
726 reconnect_grace_ticks: 2,
727 max_commands_per_tick: 4,
728 });
729 for index in 0..HASHED_GATEWAY_SESSION_MIN_ENTRIES - 1 {
730 table
731 .connect(
732 ClientId::new(u64::try_from(index).expect("test client id fits u64")),
733 StationId::new(u32::try_from(index % 4).expect("test station id fits u32")),
734 Tick::new(0),
735 )
736 .expect("session should connect");
737 }
738 assert!(!table.sessions.is_hashed());
739 table
740 .connect(ClientId::new(0), StationId::new(3), Tick::new(1))
741 .expect("existing session should refresh without promotion");
742 assert!(!table.sessions.is_hashed());
743 table
744 .admit_sequence(ClientId::new(0), 1, Tick::new(1))
745 .expect("command should admit before promotion");
746 table
747 .disconnect(ClientId::new(1), Tick::new(1))
748 .expect("session should disconnect before promotion");
749
750 let final_client = ClientId::new(
751 u64::try_from(HASHED_GATEWAY_SESSION_MIN_ENTRIES - 1).expect("test client id fits u64"),
752 );
753 table
754 .connect(final_client, StationId::new(2), Tick::new(2))
755 .expect("threshold session should connect");
756 assert!(table.sessions.is_hashed());
757 assert_eq!(
758 table
759 .route(ClientId::new(0))
760 .expect("route should survive")
761 .station_id,
762 StationId::new(3)
763 );
764 assert_eq!(
765 table
766 .admit_sequence(ClientId::new(0), 2, Tick::new(2))
767 .expect("admission should survive")
768 .sequence,
769 2
770 );
771 assert_eq!(table.expire_disconnected(Tick::new(4)), 1);
772 assert!(table.sessions.is_hashed());
773 assert_eq!(table.len(), HASHED_GATEWAY_SESSION_MIN_ENTRIES - 1);
774 assert_eq!(
775 table.stats.sessions_created,
776 HASHED_GATEWAY_SESSION_MIN_ENTRIES
777 );
778 assert_eq!(table.stats.sessions_expired, 1);
779 assert_eq!(table.stats.commands_admitted, 2);
780 }
781
782 #[test]
783 fn reconnects_with_generation_and_expires_disconnected_sessions() {
784 let client_id = ClientId::new(7);
785 let mut table = GatewaySessionTable::new(GatewayConfig {
786 max_sessions: 4,
787 reconnect_grace_ticks: 3,
788 max_commands_per_tick: 4,
789 });
790 let connected = table
791 .connect(client_id, StationId::new(1), Tick::new(10))
792 .expect("connect should work");
793 table
794 .disconnect(client_id, Tick::new(12))
795 .expect("disconnect should work");
796 assert!(matches!(
797 table.route(client_id),
798 Err(GatewayError::SessionDisconnected { .. })
799 ));
800
801 let bad = table
802 .reconnect(client_id, connected.route.generation + 1, Tick::new(13))
803 .expect_err("bad generation should fail");
804 assert_eq!(
805 bad,
806 GatewayError::BadGeneration {
807 expected: 1,
808 actual: 2
809 }
810 );
811
812 let reconnected = table
813 .reconnect(client_id, connected.route.generation, Tick::new(14))
814 .expect("reconnect should work");
815 assert_eq!(
816 reconnected.outcome,
817 GatewayConnectOutcome::Reconnected {
818 disconnected_for: 2
819 }
820 );
821 assert_eq!(reconnected.route.generation, 1);
822
823 table
824 .disconnect(client_id, Tick::new(15))
825 .expect("disconnect should work");
826 assert_eq!(table.expire_disconnected(Tick::new(19)), 1);
827 assert_eq!(table.len(), 0);
828 assert_eq!(table.stats().sessions_expired, 1);
829 }
830
831 #[test]
832 fn expiry_retains_connected_and_grace_boundary_sessions() {
833 let mut table = GatewaySessionTable::new(GatewayConfig {
834 max_sessions: 4,
835 reconnect_grace_ticks: 3,
836 max_commands_per_tick: 4,
837 });
838 for client in 1..=3 {
839 table
840 .connect(ClientId::new(client), StationId::new(1), Tick::new(10))
841 .expect("session should connect");
842 }
843 table
844 .disconnect(ClientId::new(2), Tick::new(12))
845 .expect("boundary session disconnects");
846 table
847 .disconnect(ClientId::new(3), Tick::new(11))
848 .expect("expired session disconnects");
849
850 assert_eq!(table.expire_disconnected(Tick::new(15)), 1);
851 assert!(table.session(ClientId::new(1)).is_some());
852 assert!(table.session(ClientId::new(2)).is_some());
853 assert!(table.session(ClientId::new(3)).is_none());
854 assert_eq!(table.stats().sessions_expired, 1);
855 assert_eq!(table.expire_disconnected(Tick::new(16)), 1);
856 assert_eq!(table.len(), 1);
857 assert_eq!(table.stats().sessions_expired, 2);
858 }
859
860 #[test]
861 fn connect_replaces_stale_disconnected_session() {
862 let client_id = ClientId::new(7);
863 let mut table = GatewaySessionTable::new(GatewayConfig {
864 max_sessions: 4,
865 reconnect_grace_ticks: 3,
866 max_commands_per_tick: 4,
867 });
868 let connected = table
869 .connect(client_id, StationId::new(1), Tick::new(10))
870 .expect("connect should work");
871 table
872 .admit_sequence(client_id, 10, Tick::new(11))
873 .expect("first command should admit");
874 table
875 .disconnect(client_id, Tick::new(12))
876 .expect("disconnect should work");
877
878 let replaced = table
879 .connect(client_id, StationId::new(2), Tick::new(20))
880 .expect("stale reconnect should replace generation");
881 assert_eq!(
882 replaced.outcome,
883 GatewayConnectOutcome::ReplacedExpired {
884 disconnected_for: 8
885 }
886 );
887 assert_eq!(replaced.route.generation, connected.route.generation + 1);
888 assert_eq!(replaced.route.station_id, StationId::new(2));
889 assert_eq!(
890 table
891 .admit_sequence(client_id, 1, Tick::new(21))
892 .expect("new generation should reset sequence")
893 .sequence,
894 1
895 );
896 }
897
898 #[test]
899 fn command_admission_rejects_replay_and_rate_limit() {
900 let client_id = ClientId::new(7);
901 let mut table = GatewaySessionTable::new(GatewayConfig {
902 max_sessions: 4,
903 reconnect_grace_ticks: 3,
904 max_commands_per_tick: 2,
905 });
906 table
907 .connect(client_id, StationId::new(1), Tick::new(10))
908 .expect("connect should work");
909
910 let first = table
911 .admit_command(&command(client_id, 1, 10))
912 .expect("first command should admit");
913 assert_eq!(first.commands_this_tick, 1);
914 assert_eq!(first.route.station_id, StationId::new(1));
915
916 let replay = table
917 .admit_command(&command(client_id, 1, 10))
918 .expect_err("same sequence should reject");
919 assert_eq!(
920 replay,
921 GatewayError::ReplayOrStale {
922 last_sequence: Some(1),
923 sequence: 1
924 }
925 );
926 assert_eq!(
927 replay.command_reject_reason(),
928 Some(CommandRejectReason::ReplayOrStale)
929 );
930
931 table
932 .admit_command(&command(client_id, 2, 10))
933 .expect("second command should admit");
934 let limited = table
935 .admit_command(&command(client_id, 3, 10))
936 .expect_err("third same-tick command should rate limit");
937 assert_eq!(
938 limited,
939 GatewayError::RateLimited {
940 limit: 2,
941 attempted: 3
942 }
943 );
944 assert_eq!(
945 limited.command_reject_reason(),
946 Some(CommandRejectReason::RateLimited)
947 );
948
949 let next_tick = table
950 .admit_command(&command(client_id, 3, 11))
951 .expect("next tick should reset rate count");
952 assert_eq!(next_tick.commands_this_tick, 1);
953 assert_eq!(table.stats().commands_admitted, 3);
954 assert_eq!(table.stats().commands_rejected_replay, 1);
955 assert_eq!(table.stats().commands_rejected_rate_limit, 1);
956 }
957
958 #[test]
959 fn capacity_limit_is_enforced() {
960 let mut table = GatewaySessionTable::new(GatewayConfig {
961 max_sessions: 1,
962 reconnect_grace_ticks: 3,
963 max_commands_per_tick: 4,
964 });
965 table
966 .connect(ClientId::new(1), StationId::new(1), Tick::new(0))
967 .expect("first session should fit");
968 let error = table
969 .connect(ClientId::new(2), StationId::new(1), Tick::new(0))
970 .expect_err("second session should exceed capacity");
971 assert_eq!(error, GatewayError::CapacityFull { capacity: 1 });
972 }
973}