Skip to main content

sectorsync_core/
gateway.rs

1//! Low-level gateway session and client routing primitives.
2
3use std::cmp::Reverse;
4use std::collections::{BTreeMap, BinaryHeap, HashMap};
5use std::hash::Hash;
6
7use crate::command::{CommandEnvelope, CommandRejectReason};
8use crate::ids::{ClientId, StationId, Tick};
9
10const HASHED_GATEWAY_SESSION_MIN_ENTRIES: usize = 1_024;
11
12#[derive(Clone, Debug)]
13enum AdaptiveSessionMap<K, V> {
14    Ordered(BTreeMap<K, V>),
15    Hashed(HashMap<K, V>),
16}
17
18impl<K: Copy + Eq + Hash + Ord, V> AdaptiveSessionMap<K, V> {
19    fn new() -> Self {
20        Self::Ordered(BTreeMap::new())
21    }
22
23    fn len(&self) -> usize {
24        match self {
25            Self::Ordered(entries) => entries.len(),
26            Self::Hashed(entries) => entries.len(),
27        }
28    }
29
30    fn is_empty(&self) -> bool {
31        match self {
32            Self::Ordered(entries) => entries.is_empty(),
33            Self::Hashed(entries) => entries.is_empty(),
34        }
35    }
36
37    fn get(&self, key: &K) -> Option<&V> {
38        match self {
39            Self::Ordered(entries) => entries.get(key),
40            Self::Hashed(entries) => entries.get(key),
41        }
42    }
43
44    fn get_mut(&mut self, key: &K) -> Option<&mut V> {
45        match self {
46            Self::Ordered(entries) => entries.get_mut(key),
47            Self::Hashed(entries) => entries.get_mut(key),
48        }
49    }
50
51    fn insert(&mut self, key: K, value: V) -> Option<V> {
52        let promote = match self {
53            Self::Ordered(entries) => {
54                entries.len() >= HASHED_GATEWAY_SESSION_MIN_ENTRIES.saturating_sub(1)
55                    && !entries.contains_key(&key)
56            }
57            Self::Hashed(_) => false,
58        };
59        if promote {
60            let Self::Ordered(ordered) = std::mem::replace(self, Self::Hashed(HashMap::new()))
61            else {
62                unreachable!("promotion starts from ordered session storage");
63            };
64            let mut hashed = HashMap::with_capacity(ordered.len().saturating_add(1));
65            hashed.extend(ordered);
66            *self = Self::Hashed(hashed);
67        }
68        match self {
69            Self::Ordered(entries) => entries.insert(key, value),
70            Self::Hashed(entries) => entries.insert(key, value),
71        }
72    }
73
74    fn remove(&mut self, key: &K) -> Option<V> {
75        match self {
76            Self::Ordered(entries) => entries.remove(key),
77            Self::Hashed(entries) => entries.remove(key),
78        }
79    }
80
81    fn iter(&self) -> impl Iterator<Item = (&K, &V)> {
82        match self {
83            Self::Ordered(entries) => SessionMapIter::Ordered(entries.iter()),
84            Self::Hashed(entries) => SessionMapIter::Hashed(entries.iter()),
85        }
86    }
87
88    #[cfg(test)]
89    fn is_hashed(&self) -> bool {
90        matches!(self, Self::Hashed(_))
91    }
92}
93
94enum SessionMapIter<'a, K, V> {
95    Ordered(std::collections::btree_map::Iter<'a, K, V>),
96    Hashed(std::collections::hash_map::Iter<'a, K, V>),
97}
98
99impl<'a, K, V> Iterator for SessionMapIter<'a, K, V> {
100    type Item = (&'a K, &'a V);
101
102    fn next(&mut self) -> Option<Self::Item> {
103        match self {
104            Self::Ordered(entries) => entries.next(),
105            Self::Hashed(entries) => entries.next(),
106        }
107    }
108}
109
110/// Gateway/session table configuration.
111#[derive(Clone, Copy, Debug, PartialEq, Eq)]
112pub struct GatewayConfig {
113    /// Maximum tracked client sessions.
114    pub max_sessions: usize,
115    /// Number of ticks a disconnected session may reconnect with its current
116    /// generation.
117    pub reconnect_grace_ticks: u64,
118    /// Maximum admitted commands per client per tick.
119    pub max_commands_per_tick: usize,
120}
121
122impl Default for GatewayConfig {
123    fn default() -> Self {
124        Self {
125            max_sessions: 65_536,
126            reconnect_grace_ticks: 20 * 60,
127            max_commands_per_tick: 64,
128        }
129    }
130}
131
132/// Gateway route snapshot for a client.
133#[derive(Clone, Copy, Debug, PartialEq, Eq)]
134pub struct GatewayRoute {
135    /// Routed client.
136    pub client_id: ClientId,
137    /// Current target station for client commands.
138    pub station_id: StationId,
139    /// Session generation used by reconnect handshakes.
140    pub generation: u64,
141    /// Route epoch incremented every time a client changes station route.
142    pub route_epoch: u64,
143}
144
145/// Session connection state.
146#[derive(Clone, Copy, Debug, PartialEq, Eq)]
147pub enum GatewaySessionState {
148    /// Client is connected and may submit commands.
149    Connected,
150    /// Client disconnected and may reconnect until the grace window expires.
151    Disconnected {
152        /// Tick at which the disconnect was observed.
153        since: Tick,
154    },
155}
156
157/// Tracked client gateway session.
158#[derive(Clone, Debug, PartialEq, Eq)]
159pub struct GatewaySession {
160    /// Client id.
161    pub client_id: ClientId,
162    /// Current target station.
163    pub station_id: StationId,
164    /// Tick at which this generation was first connected.
165    pub connected_at: Tick,
166    /// Last tick at which the client was observed.
167    pub last_seen: Tick,
168    /// Reconnect generation. External gateways may expose this as an opaque
169    /// token after adding their own authentication.
170    pub generation: u64,
171    /// Route epoch incremented on station route changes.
172    pub route_epoch: u64,
173    /// Connection state.
174    pub state: GatewaySessionState,
175    last_sequence: Option<u64>,
176    command_tick: Tick,
177    commands_this_tick: usize,
178}
179
180impl GatewaySession {
181    /// Returns a route snapshot.
182    pub const fn route(&self) -> GatewayRoute {
183        GatewayRoute {
184            client_id: self.client_id,
185            station_id: self.station_id,
186            generation: self.generation,
187            route_epoch: self.route_epoch,
188        }
189    }
190
191    /// Returns the latest accepted command sequence.
192    pub const fn last_sequence(&self) -> Option<u64> {
193        self.last_sequence
194    }
195
196    /// Returns command count admitted in the currently tracked tick.
197    pub const fn commands_this_tick(&self) -> usize {
198        self.commands_this_tick
199    }
200
201    /// Returns whether the session is connected.
202    pub const fn is_connected(&self) -> bool {
203        matches!(self.state, GatewaySessionState::Connected)
204    }
205}
206
207/// Result of connecting or reconnecting a client.
208#[derive(Clone, Copy, Debug, PartialEq, Eq)]
209pub enum GatewayConnectOutcome {
210    /// A new session entry was created.
211    Created,
212    /// An already connected session was refreshed.
213    AlreadyConnected,
214    /// A disconnected session reconnected within its grace window.
215    Reconnected {
216        /// Ticks between disconnect and reconnect.
217        disconnected_for: u64,
218    },
219    /// A stale disconnected session was replaced with a new generation.
220    ReplacedExpired {
221        /// Ticks between disconnect and replacement.
222        disconnected_for: u64,
223    },
224}
225
226/// Connection report.
227#[derive(Clone, Copy, Debug, PartialEq, Eq)]
228pub struct GatewayConnectReport {
229    /// Outcome.
230    pub outcome: GatewayConnectOutcome,
231    /// Current route.
232    pub route: GatewayRoute,
233}
234
235/// Result of admitting a command through gateway metadata checks.
236#[derive(Clone, Copy, Debug, PartialEq, Eq)]
237pub struct GatewayCommandAdmission {
238    /// Route that should receive the command.
239    pub route: GatewayRoute,
240    /// Accepted client-side command sequence.
241    pub sequence: u64,
242    /// Commands accepted for this client in the current tick.
243    pub commands_this_tick: usize,
244}
245
246/// Gateway session table statistics.
247#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
248pub struct GatewayStats {
249    /// Sessions created.
250    pub sessions_created: usize,
251    /// Successful reconnects.
252    pub sessions_reconnected: usize,
253    /// Disconnected sessions expired or replaced after grace.
254    pub sessions_expired: usize,
255    /// Route changes.
256    pub routes_changed: usize,
257    /// Commands admitted.
258    pub commands_admitted: usize,
259    /// Commands rejected as replay/stale.
260    pub commands_rejected_replay: usize,
261    /// Commands rejected by per-client rate limit.
262    pub commands_rejected_rate_limit: usize,
263    /// Due deadline entries examined by expiry maintenance.
264    pub expiry_deadlines_popped: usize,
265    /// Due entries ignored because the session reconnected or disconnected again.
266    pub stale_expiry_deadlines: usize,
267    /// Stale deadline heap compactions performed after bounded growth.
268    pub expiry_deadline_compactions: usize,
269}
270
271/// Gateway/session error.
272#[derive(Clone, Copy, Debug, PartialEq, Eq)]
273pub enum GatewayError {
274    /// Session table is full.
275    CapacityFull {
276        /// Configured maximum session count.
277        capacity: usize,
278    },
279    /// Session does not exist.
280    MissingSession(ClientId),
281    /// Session is disconnected.
282    SessionDisconnected {
283        /// Client id.
284        client_id: ClientId,
285        /// Disconnect tick.
286        since: Tick,
287    },
288    /// Session reconnect generation does not match.
289    BadGeneration {
290        /// Expected generation.
291        expected: u64,
292        /// Provided generation.
293        actual: u64,
294    },
295    /// Command sequence was stale or replayed.
296    ReplayOrStale {
297        /// Latest accepted sequence, if any.
298        last_sequence: Option<u64>,
299        /// Submitted sequence.
300        sequence: u64,
301    },
302    /// Client exceeded the per-tick command admission limit.
303    RateLimited {
304        /// Configured limit.
305        limit: usize,
306        /// Attempted count in the tick.
307        attempted: usize,
308    },
309}
310
311impl GatewayError {
312    /// Maps gateway metadata errors into generic command reject reasons.
313    pub const fn command_reject_reason(&self) -> Option<CommandRejectReason> {
314        match self {
315            Self::ReplayOrStale { .. } => Some(CommandRejectReason::ReplayOrStale),
316            Self::RateLimited { .. } => Some(CommandRejectReason::RateLimited),
317            Self::MissingSession(_)
318            | Self::SessionDisconnected { .. }
319            | Self::BadGeneration { .. }
320            | Self::CapacityFull { .. } => None,
321        }
322    }
323}
324
325impl core::fmt::Display for GatewayError {
326    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
327        match self {
328            Self::CapacityFull { capacity } => {
329                write!(f, "gateway session table is full at capacity {capacity}")
330            }
331            Self::MissingSession(client_id) => {
332                write!(
333                    f,
334                    "gateway session for client {} is missing",
335                    client_id.get()
336                )
337            }
338            Self::SessionDisconnected { client_id, since } => write!(
339                f,
340                "gateway session for client {} disconnected at tick {}",
341                client_id.get(),
342                since.get()
343            ),
344            Self::BadGeneration { expected, actual } => write!(
345                f,
346                "gateway reconnect generation mismatch: expected {expected}, actual {actual}"
347            ),
348            Self::ReplayOrStale {
349                last_sequence,
350                sequence,
351            } => write!(
352                f,
353                "gateway command sequence {sequence} is not newer than {last_sequence:?}"
354            ),
355            Self::RateLimited { limit, attempted } => write!(
356                f,
357                "gateway command rate limited: limit {limit}, attempted {attempted}"
358            ),
359        }
360    }
361}
362
363impl std::error::Error for GatewayError {}
364
365/// Bounded in-memory gateway session and route table.
366#[derive(Clone, Debug)]
367pub struct GatewaySessionTable {
368    config: GatewayConfig,
369    sessions: AdaptiveSessionMap<ClientId, GatewaySession>,
370    expiry_deadlines: BinaryHeap<Reverse<(u64, ClientId, u64)>>,
371    stats: GatewayStats,
372}
373
374impl GatewaySessionTable {
375    /// Creates an empty gateway session table.
376    pub fn new(config: GatewayConfig) -> Self {
377        Self {
378            config,
379            sessions: AdaptiveSessionMap::new(),
380            expiry_deadlines: BinaryHeap::new(),
381            stats: GatewayStats::default(),
382        }
383    }
384
385    /// Returns configuration.
386    pub const fn config(&self) -> GatewayConfig {
387        self.config
388    }
389
390    /// Returns statistics.
391    pub const fn stats(&self) -> GatewayStats {
392        self.stats
393    }
394
395    /// Returns tracked session count.
396    pub fn len(&self) -> usize {
397        self.sessions.len()
398    }
399
400    /// Returns whether no sessions are tracked.
401    pub fn is_empty(&self) -> bool {
402        self.sessions.is_empty()
403    }
404
405    /// Returns a session by client id.
406    pub fn session(&self, client_id: ClientId) -> Option<&GatewaySession> {
407        self.sessions.get(&client_id)
408    }
409
410    /// Returns a connected route by client id.
411    pub fn route(&self, client_id: ClientId) -> Result<GatewayRoute, GatewayError> {
412        let session = self
413            .sessions
414            .get(&client_id)
415            .ok_or(GatewayError::MissingSession(client_id))?;
416        match session.state {
417            GatewaySessionState::Connected => Ok(session.route()),
418            GatewaySessionState::Disconnected { since } => {
419                Err(GatewayError::SessionDisconnected { client_id, since })
420            }
421        }
422    }
423
424    /// Connects a client, reconnecting within grace or replacing stale sessions.
425    pub fn connect(
426        &mut self,
427        client_id: ClientId,
428        station_id: StationId,
429        now: Tick,
430    ) -> Result<GatewayConnectReport, GatewayError> {
431        let reconnect_grace_ticks = self.config.reconnect_grace_ticks;
432        if let Some(session) = self.sessions.get_mut(&client_id) {
433            return Ok(Self::connect_existing(
434                session,
435                station_id,
436                now,
437                reconnect_grace_ticks,
438                &mut self.stats,
439            ));
440        }
441
442        if self.sessions.len() >= self.config.max_sessions {
443            return Err(GatewayError::CapacityFull {
444                capacity: self.config.max_sessions,
445            });
446        }
447
448        let session = GatewaySession {
449            client_id,
450            station_id,
451            connected_at: now,
452            last_seen: now,
453            generation: 1,
454            route_epoch: 1,
455            state: GatewaySessionState::Connected,
456            last_sequence: None,
457            command_tick: now,
458            commands_this_tick: 0,
459        };
460        let route = session.route();
461        self.sessions.insert(client_id, session);
462        self.stats.sessions_created = self.stats.sessions_created.saturating_add(1);
463        Ok(GatewayConnectReport {
464            outcome: GatewayConnectOutcome::Created,
465            route,
466        })
467    }
468
469    /// Reconnects a disconnected client by generation.
470    pub fn reconnect(
471        &mut self,
472        client_id: ClientId,
473        generation: u64,
474        now: Tick,
475    ) -> Result<GatewayConnectReport, GatewayError> {
476        let session = self
477            .sessions
478            .get_mut(&client_id)
479            .ok_or(GatewayError::MissingSession(client_id))?;
480        if session.generation != generation {
481            return Err(GatewayError::BadGeneration {
482                expected: session.generation,
483                actual: generation,
484            });
485        }
486
487        match session.state {
488            GatewaySessionState::Connected => {
489                session.last_seen = now;
490                Ok(GatewayConnectReport {
491                    outcome: GatewayConnectOutcome::AlreadyConnected,
492                    route: session.route(),
493                })
494            }
495            GatewaySessionState::Disconnected { since } => {
496                let disconnected_for = now.get().saturating_sub(since.get());
497                if disconnected_for > self.config.reconnect_grace_ticks {
498                    session.connected_at = now;
499                    session.last_seen = now;
500                    session.generation = session.generation.saturating_add(1);
501                    session.route_epoch = session.route_epoch.saturating_add(1);
502                    session.state = GatewaySessionState::Connected;
503                    session.last_sequence = None;
504                    session.command_tick = now;
505                    session.commands_this_tick = 0;
506                    self.stats.sessions_expired = self.stats.sessions_expired.saturating_add(1);
507                    Ok(GatewayConnectReport {
508                        outcome: GatewayConnectOutcome::ReplacedExpired { disconnected_for },
509                        route: session.route(),
510                    })
511                } else {
512                    session.last_seen = now;
513                    session.state = GatewaySessionState::Connected;
514                    self.stats.sessions_reconnected =
515                        self.stats.sessions_reconnected.saturating_add(1);
516                    Ok(GatewayConnectReport {
517                        outcome: GatewayConnectOutcome::Reconnected { disconnected_for },
518                        route: session.route(),
519                    })
520                }
521            }
522        }
523    }
524
525    /// Marks a connected session disconnected.
526    pub fn disconnect(&mut self, client_id: ClientId, now: Tick) -> Result<(), GatewayError> {
527        let session = self
528            .sessions
529            .get_mut(&client_id)
530            .ok_or(GatewayError::MissingSession(client_id))?;
531        session.last_seen = now;
532        session.state = GatewaySessionState::Disconnected { since: now };
533        self.expiry_deadlines.push(Reverse((
534            expiry_deadline(now, self.config.reconnect_grace_ticks),
535            client_id,
536            now.get(),
537        )));
538        self.compact_expiry_deadlines_if_needed();
539        Ok(())
540    }
541
542    /// Removes disconnected sessions whose grace window has expired.
543    pub fn expire_disconnected(&mut self, now: Tick) -> usize {
544        let mut expired = 0_usize;
545        while self
546            .expiry_deadlines
547            .peek()
548            .is_some_and(|Reverse((deadline, _, _))| *deadline <= now.get())
549        {
550            let Reverse((_, client_id, expected_since)) = self
551                .expiry_deadlines
552                .pop()
553                .expect("peeked deadline remains available");
554            self.stats.expiry_deadlines_popped =
555                self.stats.expiry_deadlines_popped.saturating_add(1);
556            let current_matches = self.sessions.get(&client_id).is_some_and(|session| {
557                matches!(
558                    session.state,
559                    GatewaySessionState::Disconnected { since }
560                        if since.get() == expected_since
561                            && now.get().saturating_sub(since.get())
562                                > self.config.reconnect_grace_ticks
563                )
564            });
565            if current_matches {
566                self.sessions.remove(&client_id);
567                expired = expired.saturating_add(1);
568            } else {
569                self.stats.stale_expiry_deadlines =
570                    self.stats.stale_expiry_deadlines.saturating_add(1);
571            }
572        }
573        self.stats.sessions_expired = self.stats.sessions_expired.saturating_add(expired);
574        expired
575    }
576
577    /// Deadline entries retained for caller-driven expiry maintenance.
578    pub fn expiry_deadline_len(&self) -> usize {
579        self.expiry_deadlines.len()
580    }
581
582    /// Allocated deadline-entry capacity retained by the table.
583    pub fn expiry_deadline_capacity(&self) -> usize {
584        self.expiry_deadlines.capacity()
585    }
586
587    /// Changes the target station for a connected client.
588    pub fn reroute(
589        &mut self,
590        client_id: ClientId,
591        station_id: StationId,
592        now: Tick,
593    ) -> Result<GatewayRoute, GatewayError> {
594        let session = self
595            .sessions
596            .get_mut(&client_id)
597            .ok_or(GatewayError::MissingSession(client_id))?;
598        match session.state {
599            GatewaySessionState::Connected => {
600                session.last_seen = now;
601                if session.station_id != station_id {
602                    session.station_id = station_id;
603                    session.route_epoch = session.route_epoch.saturating_add(1);
604                    self.stats.routes_changed = self.stats.routes_changed.saturating_add(1);
605                }
606                Ok(session.route())
607            }
608            GatewaySessionState::Disconnected { since } => {
609                Err(GatewayError::SessionDisconnected { client_id, since })
610            }
611        }
612    }
613
614    /// Applies session metadata checks to a command and returns the route that
615    /// should receive it.
616    pub fn admit_command(
617        &mut self,
618        command: &CommandEnvelope,
619    ) -> Result<GatewayCommandAdmission, GatewayError> {
620        self.admit_sequence(command.client_id, command.sequence, command.received_at)
621    }
622
623    /// Applies session metadata checks to a client command sequence.
624    pub fn admit_sequence(
625        &mut self,
626        client_id: ClientId,
627        sequence: u64,
628        now: Tick,
629    ) -> Result<GatewayCommandAdmission, GatewayError> {
630        let session = self
631            .sessions
632            .get_mut(&client_id)
633            .ok_or(GatewayError::MissingSession(client_id))?;
634        match session.state {
635            GatewaySessionState::Connected => {}
636            GatewaySessionState::Disconnected { since } => {
637                return Err(GatewayError::SessionDisconnected { client_id, since });
638            }
639        }
640
641        if session
642            .last_sequence
643            .is_some_and(|last_sequence| sequence <= last_sequence)
644        {
645            self.stats.commands_rejected_replay =
646                self.stats.commands_rejected_replay.saturating_add(1);
647            return Err(GatewayError::ReplayOrStale {
648                last_sequence: session.last_sequence,
649                sequence,
650            });
651        }
652
653        if session.command_tick != now {
654            session.command_tick = now;
655            session.commands_this_tick = 0;
656        }
657        let attempted = session.commands_this_tick.saturating_add(1);
658        if attempted > self.config.max_commands_per_tick {
659            self.stats.commands_rejected_rate_limit =
660                self.stats.commands_rejected_rate_limit.saturating_add(1);
661            return Err(GatewayError::RateLimited {
662                limit: self.config.max_commands_per_tick,
663                attempted,
664            });
665        }
666
667        session.commands_this_tick = attempted;
668        session.last_sequence = Some(sequence);
669        session.last_seen = now;
670        self.stats.commands_admitted = self.stats.commands_admitted.saturating_add(1);
671        Ok(GatewayCommandAdmission {
672            route: session.route(),
673            sequence,
674            commands_this_tick: attempted,
675        })
676    }
677
678    fn connect_existing(
679        session: &mut GatewaySession,
680        station_id: StationId,
681        now: Tick,
682        reconnect_grace_ticks: u64,
683        stats: &mut GatewayStats,
684    ) -> GatewayConnectReport {
685        match session.state {
686            GatewaySessionState::Connected => {
687                session.last_seen = now;
688                if session.station_id != station_id {
689                    session.station_id = station_id;
690                    session.route_epoch = session.route_epoch.saturating_add(1);
691                    stats.routes_changed = stats.routes_changed.saturating_add(1);
692                }
693                GatewayConnectReport {
694                    outcome: GatewayConnectOutcome::AlreadyConnected,
695                    route: session.route(),
696                }
697            }
698            GatewaySessionState::Disconnected { since } => {
699                let disconnected_for = now.get().saturating_sub(since.get());
700                if disconnected_for > reconnect_grace_ticks {
701                    session.station_id = station_id;
702                    session.connected_at = now;
703                    session.last_seen = now;
704                    session.generation = session.generation.saturating_add(1);
705                    session.route_epoch = session.route_epoch.saturating_add(1);
706                    session.state = GatewaySessionState::Connected;
707                    session.last_sequence = None;
708                    session.command_tick = now;
709                    session.commands_this_tick = 0;
710                    stats.sessions_expired = stats.sessions_expired.saturating_add(1);
711                    GatewayConnectReport {
712                        outcome: GatewayConnectOutcome::ReplacedExpired { disconnected_for },
713                        route: session.route(),
714                    }
715                } else {
716                    session.station_id = station_id;
717                    session.last_seen = now;
718                    session.state = GatewaySessionState::Connected;
719                    stats.sessions_reconnected = stats.sessions_reconnected.saturating_add(1);
720                    GatewayConnectReport {
721                        outcome: GatewayConnectOutcome::Reconnected { disconnected_for },
722                        route: session.route(),
723                    }
724                }
725            }
726        }
727    }
728
729    fn compact_expiry_deadlines_if_needed(&mut self) {
730        let limit = self.config.max_sessions.max(1).saturating_mul(2);
731        if self.expiry_deadlines.len() <= limit {
732            return;
733        }
734        let mut deadlines = BinaryHeap::with_capacity(self.sessions.len());
735        for (client_id, session) in self.sessions.iter() {
736            if let GatewaySessionState::Disconnected { since } = session.state {
737                deadlines.push(Reverse((
738                    expiry_deadline(since, self.config.reconnect_grace_ticks),
739                    *client_id,
740                    since.get(),
741                )));
742            }
743        }
744        self.expiry_deadlines = deadlines;
745        self.stats.expiry_deadline_compactions =
746            self.stats.expiry_deadline_compactions.saturating_add(1);
747    }
748}
749
750const fn expiry_deadline(since: Tick, grace: u64) -> u64 {
751    since.get().saturating_add(grace).saturating_add(1)
752}
753
754impl Default for GatewaySessionTable {
755    fn default() -> Self {
756        Self::new(GatewayConfig::default())
757    }
758}
759
760#[cfg(test)]
761mod tests {
762    use super::*;
763    use crate::command::{CommandEnvelope, CommandPriority};
764    use crate::ids::{CommandId, EntityId};
765
766    fn command(client_id: ClientId, sequence: u64, tick: u64) -> CommandEnvelope {
767        CommandEnvelope {
768            id: CommandId::new(sequence),
769            client_id,
770            entity_id: EntityId::new(10),
771            sequence,
772            received_at: Tick::new(tick),
773            kind: 1,
774            priority: CommandPriority::Normal,
775            payload: Vec::new(),
776        }
777    }
778
779    #[test]
780    fn connects_routes_and_reroutes_sessions() {
781        let client_id = ClientId::new(7);
782        let mut table = GatewaySessionTable::new(GatewayConfig {
783            max_sessions: 4,
784            reconnect_grace_ticks: 3,
785            max_commands_per_tick: 4,
786        });
787
788        let connected = table
789            .connect(client_id, StationId::new(1), Tick::new(10))
790            .expect("connect should work");
791        assert_eq!(connected.outcome, GatewayConnectOutcome::Created);
792        assert_eq!(connected.route.station_id, StationId::new(1));
793        assert_eq!(connected.route.generation, 1);
794        assert_eq!(connected.route.route_epoch, 1);
795
796        let route = table
797            .reroute(client_id, StationId::new(2), Tick::new(11))
798            .expect("reroute should work");
799        assert_eq!(route.station_id, StationId::new(2));
800        assert_eq!(route.route_epoch, 2);
801        assert_eq!(table.stats().routes_changed, 1);
802        assert_eq!(
803            table
804                .route(client_id)
805                .expect("route should exist")
806                .station_id,
807            StationId::new(2)
808        );
809    }
810
811    #[test]
812    fn gateway_sessions_promote_without_losing_route_admission_or_expiry_state() {
813        let mut table = GatewaySessionTable::new(GatewayConfig {
814            max_sessions: HASHED_GATEWAY_SESSION_MIN_ENTRIES + 1,
815            reconnect_grace_ticks: 2,
816            max_commands_per_tick: 4,
817        });
818        for index in 0..HASHED_GATEWAY_SESSION_MIN_ENTRIES - 1 {
819            table
820                .connect(
821                    ClientId::new(u64::try_from(index).expect("test client id fits u64")),
822                    StationId::new(u32::try_from(index % 4).expect("test station id fits u32")),
823                    Tick::new(0),
824                )
825                .expect("session should connect");
826        }
827        assert!(!table.sessions.is_hashed());
828        table
829            .connect(ClientId::new(0), StationId::new(3), Tick::new(1))
830            .expect("existing session should refresh without promotion");
831        assert!(!table.sessions.is_hashed());
832        table
833            .admit_sequence(ClientId::new(0), 1, Tick::new(1))
834            .expect("command should admit before promotion");
835        table
836            .disconnect(ClientId::new(1), Tick::new(1))
837            .expect("session should disconnect before promotion");
838
839        let final_client = ClientId::new(
840            u64::try_from(HASHED_GATEWAY_SESSION_MIN_ENTRIES - 1).expect("test client id fits u64"),
841        );
842        table
843            .connect(final_client, StationId::new(2), Tick::new(2))
844            .expect("threshold session should connect");
845        assert!(table.sessions.is_hashed());
846        assert_eq!(
847            table
848                .route(ClientId::new(0))
849                .expect("route should survive")
850                .station_id,
851            StationId::new(3)
852        );
853        assert_eq!(
854            table
855                .admit_sequence(ClientId::new(0), 2, Tick::new(2))
856                .expect("admission should survive")
857                .sequence,
858            2
859        );
860        assert_eq!(table.expire_disconnected(Tick::new(4)), 1);
861        assert!(table.sessions.is_hashed());
862        assert_eq!(table.len(), HASHED_GATEWAY_SESSION_MIN_ENTRIES - 1);
863        assert_eq!(
864            table.stats.sessions_created,
865            HASHED_GATEWAY_SESSION_MIN_ENTRIES
866        );
867        assert_eq!(table.stats.sessions_expired, 1);
868        assert_eq!(table.stats.commands_admitted, 2);
869    }
870
871    #[test]
872    fn reconnects_with_generation_and_expires_disconnected_sessions() {
873        let client_id = ClientId::new(7);
874        let mut table = GatewaySessionTable::new(GatewayConfig {
875            max_sessions: 4,
876            reconnect_grace_ticks: 3,
877            max_commands_per_tick: 4,
878        });
879        let connected = table
880            .connect(client_id, StationId::new(1), Tick::new(10))
881            .expect("connect should work");
882        table
883            .disconnect(client_id, Tick::new(12))
884            .expect("disconnect should work");
885        assert!(matches!(
886            table.route(client_id),
887            Err(GatewayError::SessionDisconnected { .. })
888        ));
889
890        let bad = table
891            .reconnect(client_id, connected.route.generation + 1, Tick::new(13))
892            .expect_err("bad generation should fail");
893        assert_eq!(
894            bad,
895            GatewayError::BadGeneration {
896                expected: 1,
897                actual: 2
898            }
899        );
900
901        let reconnected = table
902            .reconnect(client_id, connected.route.generation, Tick::new(14))
903            .expect("reconnect should work");
904        assert_eq!(
905            reconnected.outcome,
906            GatewayConnectOutcome::Reconnected {
907                disconnected_for: 2
908            }
909        );
910        assert_eq!(reconnected.route.generation, 1);
911
912        table
913            .disconnect(client_id, Tick::new(15))
914            .expect("disconnect should work");
915        assert_eq!(table.expire_disconnected(Tick::new(19)), 1);
916        assert_eq!(table.len(), 0);
917        assert_eq!(table.stats().sessions_expired, 1);
918    }
919
920    #[test]
921    fn stale_expiry_deadline_cannot_remove_reconnected_session() {
922        let client_id = ClientId::new(7);
923        let mut table = GatewaySessionTable::new(GatewayConfig {
924            max_sessions: 1,
925            reconnect_grace_ticks: 3,
926            max_commands_per_tick: 4,
927        });
928        let connected = table
929            .connect(client_id, StationId::new(1), Tick::new(0))
930            .expect("connect");
931        table
932            .disconnect(client_id, Tick::new(1))
933            .expect("disconnect");
934        table
935            .reconnect(client_id, connected.route.generation, Tick::new(2))
936            .expect("reconnect");
937
938        assert_eq!(table.expire_disconnected(Tick::new(10)), 0);
939        assert!(
940            table
941                .session(client_id)
942                .is_some_and(GatewaySession::is_connected)
943        );
944        assert_eq!(table.stats().expiry_deadlines_popped, 1);
945        assert_eq!(table.stats().stale_expiry_deadlines, 1);
946    }
947
948    #[test]
949    fn repeated_disconnect_deadlines_compact_at_bounded_growth() {
950        let client_id = ClientId::new(7);
951        let mut table = GatewaySessionTable::new(GatewayConfig {
952            max_sessions: 1,
953            reconnect_grace_ticks: 10,
954            max_commands_per_tick: 4,
955        });
956        let generation = table
957            .connect(client_id, StationId::new(1), Tick::new(0))
958            .expect("connect")
959            .route
960            .generation;
961        for tick in 1..=3 {
962            table
963                .disconnect(client_id, Tick::new(tick * 2))
964                .expect("disconnect");
965            table
966                .reconnect(client_id, generation, Tick::new(tick * 2 + 1))
967                .expect("reconnect");
968        }
969
970        assert_eq!(table.stats().expiry_deadline_compactions, 1);
971        assert!(table.expiry_deadline_len() <= table.config().max_sessions);
972        assert!(
973            table
974                .session(client_id)
975                .is_some_and(GatewaySession::is_connected)
976        );
977    }
978
979    #[test]
980    fn expiry_retains_connected_and_grace_boundary_sessions() {
981        let mut table = GatewaySessionTable::new(GatewayConfig {
982            max_sessions: 4,
983            reconnect_grace_ticks: 3,
984            max_commands_per_tick: 4,
985        });
986        for client in 1..=3 {
987            table
988                .connect(ClientId::new(client), StationId::new(1), Tick::new(10))
989                .expect("session should connect");
990        }
991        table
992            .disconnect(ClientId::new(2), Tick::new(12))
993            .expect("boundary session disconnects");
994        table
995            .disconnect(ClientId::new(3), Tick::new(11))
996            .expect("expired session disconnects");
997
998        assert_eq!(table.expire_disconnected(Tick::new(15)), 1);
999        assert!(table.session(ClientId::new(1)).is_some());
1000        assert!(table.session(ClientId::new(2)).is_some());
1001        assert!(table.session(ClientId::new(3)).is_none());
1002        assert_eq!(table.stats().sessions_expired, 1);
1003        assert_eq!(table.expire_disconnected(Tick::new(16)), 1);
1004        assert_eq!(table.len(), 1);
1005        assert_eq!(table.stats().sessions_expired, 2);
1006    }
1007
1008    #[test]
1009    fn connect_replaces_stale_disconnected_session() {
1010        let client_id = ClientId::new(7);
1011        let mut table = GatewaySessionTable::new(GatewayConfig {
1012            max_sessions: 4,
1013            reconnect_grace_ticks: 3,
1014            max_commands_per_tick: 4,
1015        });
1016        let connected = table
1017            .connect(client_id, StationId::new(1), Tick::new(10))
1018            .expect("connect should work");
1019        table
1020            .admit_sequence(client_id, 10, Tick::new(11))
1021            .expect("first command should admit");
1022        table
1023            .disconnect(client_id, Tick::new(12))
1024            .expect("disconnect should work");
1025
1026        let replaced = table
1027            .connect(client_id, StationId::new(2), Tick::new(20))
1028            .expect("stale reconnect should replace generation");
1029        assert_eq!(
1030            replaced.outcome,
1031            GatewayConnectOutcome::ReplacedExpired {
1032                disconnected_for: 8
1033            }
1034        );
1035        assert_eq!(replaced.route.generation, connected.route.generation + 1);
1036        assert_eq!(replaced.route.station_id, StationId::new(2));
1037        assert_eq!(
1038            table
1039                .admit_sequence(client_id, 1, Tick::new(21))
1040                .expect("new generation should reset sequence")
1041                .sequence,
1042            1
1043        );
1044    }
1045
1046    #[test]
1047    fn command_admission_rejects_replay_and_rate_limit() {
1048        let client_id = ClientId::new(7);
1049        let mut table = GatewaySessionTable::new(GatewayConfig {
1050            max_sessions: 4,
1051            reconnect_grace_ticks: 3,
1052            max_commands_per_tick: 2,
1053        });
1054        table
1055            .connect(client_id, StationId::new(1), Tick::new(10))
1056            .expect("connect should work");
1057
1058        let first = table
1059            .admit_command(&command(client_id, 1, 10))
1060            .expect("first command should admit");
1061        assert_eq!(first.commands_this_tick, 1);
1062        assert_eq!(first.route.station_id, StationId::new(1));
1063
1064        let replay = table
1065            .admit_command(&command(client_id, 1, 10))
1066            .expect_err("same sequence should reject");
1067        assert_eq!(
1068            replay,
1069            GatewayError::ReplayOrStale {
1070                last_sequence: Some(1),
1071                sequence: 1
1072            }
1073        );
1074        assert_eq!(
1075            replay.command_reject_reason(),
1076            Some(CommandRejectReason::ReplayOrStale)
1077        );
1078
1079        table
1080            .admit_command(&command(client_id, 2, 10))
1081            .expect("second command should admit");
1082        let limited = table
1083            .admit_command(&command(client_id, 3, 10))
1084            .expect_err("third same-tick command should rate limit");
1085        assert_eq!(
1086            limited,
1087            GatewayError::RateLimited {
1088                limit: 2,
1089                attempted: 3
1090            }
1091        );
1092        assert_eq!(
1093            limited.command_reject_reason(),
1094            Some(CommandRejectReason::RateLimited)
1095        );
1096
1097        let next_tick = table
1098            .admit_command(&command(client_id, 3, 11))
1099            .expect("next tick should reset rate count");
1100        assert_eq!(next_tick.commands_this_tick, 1);
1101        assert_eq!(table.stats().commands_admitted, 3);
1102        assert_eq!(table.stats().commands_rejected_replay, 1);
1103        assert_eq!(table.stats().commands_rejected_rate_limit, 1);
1104    }
1105
1106    #[test]
1107    fn capacity_limit_is_enforced() {
1108        let mut table = GatewaySessionTable::new(GatewayConfig {
1109            max_sessions: 1,
1110            reconnect_grace_ticks: 3,
1111            max_commands_per_tick: 4,
1112        });
1113        table
1114            .connect(ClientId::new(1), StationId::new(1), Tick::new(0))
1115            .expect("first session should fit");
1116        let error = table
1117            .connect(ClientId::new(2), StationId::new(1), Tick::new(0))
1118            .expect_err("second session should exceed capacity");
1119        assert_eq!(error, GatewayError::CapacityFull { capacity: 1 });
1120    }
1121}