1use std::collections::{BTreeMap, BTreeSet};
4
5use sectorsync_core::prelude::{
6 ClientId, GatewayError, GatewayRoute, GatewaySessionTable, NodeId, StationId, Tick,
7};
8
9#[derive(Clone, Copy, Debug, PartialEq, Eq)]
11pub struct DeploymentConfig {
12 pub max_nodes: usize,
14 pub max_stations_per_node: usize,
17 pub stale_after_ticks: u64,
19}
20
21impl Default for DeploymentConfig {
22 fn default() -> Self {
23 Self {
24 max_nodes: 1024,
25 max_stations_per_node: 1024,
26 stale_after_ticks: 20 * 10,
27 }
28 }
29}
30
31#[derive(Clone, Copy, Debug, PartialEq, Eq)]
33pub enum DeploymentNodeState {
34 Online,
36 Draining,
38 Offline,
40}
41
42#[derive(Clone, Copy, Debug, PartialEq, Eq)]
44pub struct DeploymentNodeRoute {
45 pub node_id: NodeId,
47 pub state: DeploymentNodeState,
49 pub last_heartbeat: Tick,
51 pub station_capacity: usize,
53 pub assigned_stations: usize,
55 pub route_epoch: u64,
57}
58
59#[derive(Clone, Copy, Debug, PartialEq, Eq)]
61pub struct DeploymentStationRoute {
62 pub station_id: StationId,
64 pub node_id: NodeId,
66 pub route_epoch: u64,
68 pub assigned_at: Tick,
70}
71
72#[derive(Clone, Copy, Debug, PartialEq, Eq)]
74pub struct DeploymentStationMove {
75 pub previous: DeploymentStationRoute,
77 pub current: DeploymentStationRoute,
79}
80
81#[derive(Clone, Copy, Debug, PartialEq, Eq)]
83pub struct GatewayDeliveryRoute {
84 pub client_id: ClientId,
86 pub station_id: StationId,
88 pub node_id: NodeId,
90 pub gateway_generation: u64,
92 pub gateway_route_epoch: u64,
94 pub station_route_epoch: u64,
96 pub node_route_epoch: u64,
98 pub node_state: DeploymentNodeState,
100 pub assigned_at: Tick,
102}
103
104#[derive(Clone, Copy, Debug, PartialEq, Eq)]
106pub enum GatewayDeliveryError {
107 Gateway(GatewayError),
109 Deployment(DeploymentError),
111}
112
113impl core::fmt::Display for GatewayDeliveryError {
114 fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
115 match self {
116 Self::Gateway(error) => write!(f, "{error}"),
117 Self::Deployment(error) => write!(f, "{error}"),
118 }
119 }
120}
121
122impl std::error::Error for GatewayDeliveryError {
123 fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
124 match self {
125 Self::Gateway(error) => Some(error),
126 Self::Deployment(error) => Some(error),
127 }
128 }
129}
130
131#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
133pub struct DeploymentStats {
134 pub nodes_registered: usize,
136 pub nodes_draining: usize,
138 pub nodes_offline: usize,
140 pub stations_assigned: usize,
142 pub stations_moved: usize,
144 pub stations_unassigned: usize,
146 pub placements_rejected_capacity: usize,
148 pub stale_nodes_detected: usize,
150}
151
152#[derive(Clone, Copy, Debug, PartialEq, Eq)]
154pub enum DeploymentError {
155 NodeCapacityFull {
157 capacity: usize,
159 },
160 MissingNode(NodeId),
162 MissingStation(StationId),
164 NodeUnavailable {
166 node_id: NodeId,
168 state: DeploymentNodeState,
170 },
171 NodeStationCapacity {
173 node_id: NodeId,
175 capacity: usize,
177 attempted: usize,
179 },
180}
181
182impl core::fmt::Display for DeploymentError {
183 fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
184 match self {
185 Self::NodeCapacityFull { capacity } => {
186 write!(f, "deployment node table is full at capacity {capacity}")
187 }
188 Self::MissingNode(node_id) => {
189 write!(f, "deployment node {} is missing", node_id.get())
190 }
191 Self::MissingStation(station_id) => {
192 write!(f, "deployment station {} is missing", station_id.get())
193 }
194 Self::NodeUnavailable { node_id, state } => write!(
195 f,
196 "deployment node {} cannot accept placements in state {state:?}",
197 node_id.get()
198 ),
199 Self::NodeStationCapacity {
200 node_id,
201 capacity,
202 attempted,
203 } => write!(
204 f,
205 "deployment node {} station capacity exceeded: capacity {capacity}, attempted {attempted}",
206 node_id.get()
207 ),
208 }
209 }
210}
211
212impl std::error::Error for DeploymentError {}
213
214#[derive(Clone, Debug)]
215struct DeploymentNodeRecord {
216 route: DeploymentNodeRoute,
217 stations: BTreeSet<StationId>,
218}
219
220impl DeploymentNodeRecord {
221 fn new(node_id: NodeId, station_capacity: usize, now: Tick) -> Self {
222 Self {
223 route: DeploymentNodeRoute {
224 node_id,
225 state: DeploymentNodeState::Online,
226 last_heartbeat: now,
227 station_capacity,
228 assigned_stations: 0,
229 route_epoch: 1,
230 },
231 stations: BTreeSet::new(),
232 }
233 }
234
235 fn refresh_assigned_count(&mut self) {
236 self.route.assigned_stations = self.stations.len();
237 }
238}
239
240#[derive(Clone, Debug)]
242pub struct DeploymentRouteTable {
243 config: DeploymentConfig,
244 nodes: BTreeMap<NodeId, DeploymentNodeRecord>,
245 stations: BTreeMap<StationId, DeploymentStationRoute>,
246 stats: DeploymentStats,
247}
248
249impl DeploymentRouteTable {
250 pub fn new(config: DeploymentConfig) -> Self {
252 Self {
253 config,
254 nodes: BTreeMap::new(),
255 stations: BTreeMap::new(),
256 stats: DeploymentStats::default(),
257 }
258 }
259
260 pub const fn config(&self) -> DeploymentConfig {
262 self.config
263 }
264
265 pub const fn stats(&self) -> DeploymentStats {
267 self.stats
268 }
269
270 pub fn node_len(&self) -> usize {
272 self.nodes.len()
273 }
274
275 pub fn station_len(&self) -> usize {
277 self.stations.len()
278 }
279
280 pub fn register_node(
282 &mut self,
283 node_id: NodeId,
284 station_capacity: usize,
285 now: Tick,
286 ) -> Result<DeploymentNodeRoute, DeploymentError> {
287 let station_capacity = station_capacity.min(self.config.max_stations_per_node);
288 if let Some(node) = self.nodes.get_mut(&node_id) {
289 node.route.last_heartbeat = now;
290 node.route.state = DeploymentNodeState::Online;
291 if node.route.station_capacity != station_capacity {
292 node.route.station_capacity = station_capacity;
293 node.route.route_epoch = node.route.route_epoch.saturating_add(1);
294 }
295 node.refresh_assigned_count();
296 return Ok(node.route);
297 }
298
299 if self.nodes.len() >= self.config.max_nodes {
300 return Err(DeploymentError::NodeCapacityFull {
301 capacity: self.config.max_nodes,
302 });
303 }
304
305 let node = DeploymentNodeRecord::new(node_id, station_capacity, now);
306 let route = node.route;
307 self.nodes.insert(node_id, node);
308 self.stats.nodes_registered = self.stats.nodes_registered.saturating_add(1);
309 Ok(route)
310 }
311
312 pub fn heartbeat(
314 &mut self,
315 node_id: NodeId,
316 now: Tick,
317 ) -> Result<DeploymentNodeRoute, DeploymentError> {
318 let node = self
319 .nodes
320 .get_mut(&node_id)
321 .ok_or(DeploymentError::MissingNode(node_id))?;
322 node.route.last_heartbeat = now;
323 Ok(node.route)
324 }
325
326 pub fn mark_draining(
328 &mut self,
329 node_id: NodeId,
330 ) -> Result<DeploymentNodeRoute, DeploymentError> {
331 let node = self
332 .nodes
333 .get_mut(&node_id)
334 .ok_or(DeploymentError::MissingNode(node_id))?;
335 if node.route.state != DeploymentNodeState::Draining {
336 node.route.state = DeploymentNodeState::Draining;
337 node.route.route_epoch = node.route.route_epoch.saturating_add(1);
338 self.stats.nodes_draining = self.stats.nodes_draining.saturating_add(1);
339 }
340 Ok(node.route)
341 }
342
343 pub fn mark_offline(
346 &mut self,
347 node_id: NodeId,
348 ) -> Result<DeploymentNodeRoute, DeploymentError> {
349 let node = self
350 .nodes
351 .get_mut(&node_id)
352 .ok_or(DeploymentError::MissingNode(node_id))?;
353 if node.route.state != DeploymentNodeState::Offline {
354 node.route.state = DeploymentNodeState::Offline;
355 node.route.route_epoch = node.route.route_epoch.saturating_add(1);
356 self.stats.nodes_offline = self.stats.nodes_offline.saturating_add(1);
357 }
358 Ok(node.route)
359 }
360
361 pub fn assign_station(
363 &mut self,
364 station_id: StationId,
365 node_id: NodeId,
366 now: Tick,
367 ) -> Result<DeploymentStationRoute, DeploymentError> {
368 self.ensure_node_can_accept(node_id)?;
369
370 if let Some(existing) = self.stations.get(&station_id).copied() {
371 if existing.node_id == node_id {
372 return Ok(existing);
373 }
374 return self
375 .move_station(station_id, node_id, now)
376 .map(|move_report| move_report.current);
377 }
378
379 let node = self
380 .nodes
381 .get_mut(&node_id)
382 .ok_or(DeploymentError::MissingNode(node_id))?;
383 let attempted = node.stations.len().saturating_add(1);
384 if attempted > node.route.station_capacity {
385 self.stats.placements_rejected_capacity =
386 self.stats.placements_rejected_capacity.saturating_add(1);
387 return Err(DeploymentError::NodeStationCapacity {
388 node_id,
389 capacity: node.route.station_capacity,
390 attempted,
391 });
392 }
393
394 node.stations.insert(station_id);
395 node.refresh_assigned_count();
396 let route = DeploymentStationRoute {
397 station_id,
398 node_id,
399 route_epoch: 1,
400 assigned_at: now,
401 };
402 self.stations.insert(station_id, route);
403 self.stats.stations_assigned = self.stats.stations_assigned.saturating_add(1);
404 Ok(route)
405 }
406
407 pub fn move_station(
409 &mut self,
410 station_id: StationId,
411 target_node: NodeId,
412 now: Tick,
413 ) -> Result<DeploymentStationMove, DeploymentError> {
414 self.ensure_node_can_accept(target_node)?;
415 let previous = self
416 .stations
417 .get(&station_id)
418 .copied()
419 .ok_or(DeploymentError::MissingStation(station_id))?;
420 if previous.node_id == target_node {
421 return Ok(DeploymentStationMove {
422 previous,
423 current: previous,
424 });
425 }
426
427 {
428 let target = self
429 .nodes
430 .get(&target_node)
431 .ok_or(DeploymentError::MissingNode(target_node))?;
432 let attempted = target.stations.len().saturating_add(1);
433 if attempted > target.route.station_capacity {
434 self.stats.placements_rejected_capacity =
435 self.stats.placements_rejected_capacity.saturating_add(1);
436 return Err(DeploymentError::NodeStationCapacity {
437 node_id: target_node,
438 capacity: target.route.station_capacity,
439 attempted,
440 });
441 }
442 }
443
444 if let Some(source) = self.nodes.get_mut(&previous.node_id) {
445 source.stations.remove(&station_id);
446 source.refresh_assigned_count();
447 }
448 let target = self
449 .nodes
450 .get_mut(&target_node)
451 .ok_or(DeploymentError::MissingNode(target_node))?;
452 target.stations.insert(station_id);
453 target.refresh_assigned_count();
454
455 let current = DeploymentStationRoute {
456 station_id,
457 node_id: target_node,
458 route_epoch: previous.route_epoch.saturating_add(1),
459 assigned_at: now,
460 };
461 self.stations.insert(station_id, current);
462 self.stats.stations_moved = self.stats.stations_moved.saturating_add(1);
463 Ok(DeploymentStationMove { previous, current })
464 }
465
466 pub fn unassign_station(
468 &mut self,
469 station_id: StationId,
470 ) -> Result<DeploymentStationRoute, DeploymentError> {
471 let route = self
472 .stations
473 .remove(&station_id)
474 .ok_or(DeploymentError::MissingStation(station_id))?;
475 if let Some(node) = self.nodes.get_mut(&route.node_id) {
476 node.stations.remove(&station_id);
477 node.refresh_assigned_count();
478 }
479 self.stats.stations_unassigned = self.stats.stations_unassigned.saturating_add(1);
480 Ok(route)
481 }
482
483 pub fn node_route(&self, node_id: NodeId) -> Result<DeploymentNodeRoute, DeploymentError> {
485 self.nodes
486 .get(&node_id)
487 .map(|node| node.route)
488 .ok_or(DeploymentError::MissingNode(node_id))
489 }
490
491 pub fn station_route(
493 &self,
494 station_id: StationId,
495 ) -> Result<DeploymentStationRoute, DeploymentError> {
496 self.stations
497 .get(&station_id)
498 .copied()
499 .ok_or(DeploymentError::MissingStation(station_id))
500 }
501
502 pub fn stations_on_node(&self, node_id: NodeId) -> Result<Vec<StationId>, DeploymentError> {
504 let node = self
505 .nodes
506 .get(&node_id)
507 .ok_or(DeploymentError::MissingNode(node_id))?;
508 Ok(node.stations.iter().copied().collect())
509 }
510
511 pub fn resolve_gateway_route(
516 &self,
517 route: GatewayRoute,
518 ) -> Result<GatewayDeliveryRoute, DeploymentError> {
519 let station_route = self.station_route(route.station_id)?;
520 let node_route = self.node_route(station_route.node_id)?;
521 if node_route.state == DeploymentNodeState::Offline {
522 return Err(DeploymentError::NodeUnavailable {
523 node_id: node_route.node_id,
524 state: node_route.state,
525 });
526 }
527
528 Ok(GatewayDeliveryRoute {
529 client_id: route.client_id,
530 station_id: route.station_id,
531 node_id: node_route.node_id,
532 gateway_generation: route.generation,
533 gateway_route_epoch: route.route_epoch,
534 station_route_epoch: station_route.route_epoch,
535 node_route_epoch: node_route.route_epoch,
536 node_state: node_route.state,
537 assigned_at: station_route.assigned_at,
538 })
539 }
540
541 pub fn resolve_gateway_client(
543 &self,
544 gateway: &GatewaySessionTable,
545 client_id: ClientId,
546 ) -> Result<GatewayDeliveryRoute, GatewayDeliveryError> {
547 let route = gateway
548 .route(client_id)
549 .map_err(GatewayDeliveryError::Gateway)?;
550 self.resolve_gateway_route(route)
551 .map_err(GatewayDeliveryError::Deployment)
552 }
553
554 pub fn stale_nodes(&self, now: Tick) -> Vec<NodeId> {
556 self.nodes
557 .iter()
558 .filter_map(|(node_id, node)| {
559 (node.route.state != DeploymentNodeState::Offline
560 && now.get().saturating_sub(node.route.last_heartbeat.get())
561 > self.config.stale_after_ticks)
562 .then_some(*node_id)
563 })
564 .collect::<Vec<_>>()
565 }
566
567 pub fn mark_stale_offline(&mut self, now: Tick) -> usize {
569 let stale = self.stale_nodes(now);
570 self.stats.stale_nodes_detected =
571 self.stats.stale_nodes_detected.saturating_add(stale.len());
572 for node_id in &stale {
573 let _ = self.mark_offline(*node_id);
574 }
575 stale.len()
576 }
577
578 fn ensure_node_can_accept(&self, node_id: NodeId) -> Result<(), DeploymentError> {
579 let node = self
580 .nodes
581 .get(&node_id)
582 .ok_or(DeploymentError::MissingNode(node_id))?;
583 match node.route.state {
584 DeploymentNodeState::Online => Ok(()),
585 DeploymentNodeState::Draining | DeploymentNodeState::Offline => {
586 Err(DeploymentError::NodeUnavailable {
587 node_id,
588 state: node.route.state,
589 })
590 }
591 }
592 }
593}
594
595impl Default for DeploymentRouteTable {
596 fn default() -> Self {
597 Self::new(DeploymentConfig::default())
598 }
599}
600
601#[cfg(test)]
602mod tests {
603 use super::*;
604
605 fn config() -> DeploymentConfig {
606 DeploymentConfig {
607 max_nodes: 2,
608 max_stations_per_node: 2,
609 stale_after_ticks: 3,
610 }
611 }
612
613 #[test]
614 fn registers_nodes_and_assigns_station_routes() {
615 let mut table = DeploymentRouteTable::new(config());
616 let node = table
617 .register_node(NodeId::new(1), 2, Tick::new(10))
618 .expect("node should register");
619 assert_eq!(node.route_epoch, 1);
620 assert_eq!(node.state, DeploymentNodeState::Online);
621
622 let route = table
623 .assign_station(StationId::new(11), NodeId::new(1), Tick::new(10))
624 .expect("station should assign");
625 assert_eq!(route.node_id, NodeId::new(1));
626 assert_eq!(route.route_epoch, 1);
627 assert_eq!(
628 table.stations_on_node(NodeId::new(1)).expect("node exists"),
629 vec![StationId::new(11)]
630 );
631 assert_eq!(
632 table
633 .node_route(NodeId::new(1))
634 .expect("node route should exist")
635 .assigned_stations,
636 1
637 );
638 }
639
640 #[test]
641 fn enforces_node_and_station_capacity() {
642 let mut table = DeploymentRouteTable::new(DeploymentConfig {
643 max_nodes: 1,
644 max_stations_per_node: 1,
645 stale_after_ticks: 3,
646 });
647 table
648 .register_node(NodeId::new(1), 2, Tick::new(0))
649 .expect("first node should register");
650 assert_eq!(
651 table
652 .register_node(NodeId::new(2), 1, Tick::new(0))
653 .expect_err("second node should exceed capacity"),
654 DeploymentError::NodeCapacityFull { capacity: 1 }
655 );
656 table
657 .assign_station(StationId::new(10), NodeId::new(1), Tick::new(0))
658 .expect("first station should fit");
659 assert_eq!(
660 table
661 .assign_station(StationId::new(11), NodeId::new(1), Tick::new(0))
662 .expect_err("second station should exceed node capacity"),
663 DeploymentError::NodeStationCapacity {
664 node_id: NodeId::new(1),
665 capacity: 1,
666 attempted: 2
667 }
668 );
669 assert_eq!(table.stats().placements_rejected_capacity, 1);
670 }
671
672 #[test]
673 fn draining_nodes_keep_routes_but_reject_new_placements() {
674 let mut table = DeploymentRouteTable::new(config());
675 table
676 .register_node(NodeId::new(1), 2, Tick::new(0))
677 .expect("node should register");
678 table
679 .assign_station(StationId::new(10), NodeId::new(1), Tick::new(0))
680 .expect("station should assign");
681 let draining = table
682 .mark_draining(NodeId::new(1))
683 .expect("node should drain");
684 assert_eq!(draining.state, DeploymentNodeState::Draining);
685 assert_eq!(
686 table
687 .station_route(StationId::new(10))
688 .expect("existing route remains")
689 .node_id,
690 NodeId::new(1)
691 );
692 assert_eq!(
693 table
694 .assign_station(StationId::new(11), NodeId::new(1), Tick::new(1))
695 .expect_err("draining node should reject new placement"),
696 DeploymentError::NodeUnavailable {
697 node_id: NodeId::new(1),
698 state: DeploymentNodeState::Draining
699 }
700 );
701 }
702
703 #[test]
704 fn moves_station_routes_between_nodes() {
705 let mut table = DeploymentRouteTable::new(config());
706 table
707 .register_node(NodeId::new(1), 2, Tick::new(0))
708 .expect("source node should register");
709 table
710 .register_node(NodeId::new(2), 2, Tick::new(0))
711 .expect("target node should register");
712 table
713 .assign_station(StationId::new(10), NodeId::new(1), Tick::new(0))
714 .expect("station should assign");
715
716 let moved = table
717 .move_station(StationId::new(10), NodeId::new(2), Tick::new(5))
718 .expect("station should move");
719 assert_eq!(moved.previous.node_id, NodeId::new(1));
720 assert_eq!(moved.current.node_id, NodeId::new(2));
721 assert_eq!(moved.current.route_epoch, 2);
722 assert!(
723 table
724 .stations_on_node(NodeId::new(1))
725 .expect("source exists")
726 .is_empty()
727 );
728 assert_eq!(
729 table
730 .stations_on_node(NodeId::new(2))
731 .expect("target exists"),
732 vec![StationId::new(10)]
733 );
734 }
735
736 #[test]
737 fn resolves_gateway_clients_to_deployment_delivery_routes() {
738 let client_id = ClientId::new(7);
739 let station_id = StationId::new(10);
740 let node_id = NodeId::new(1);
741 let mut gateway = GatewaySessionTable::default();
742 let connected = gateway
743 .connect(client_id, station_id, Tick::new(10))
744 .expect("client should connect");
745 let mut table = DeploymentRouteTable::new(config());
746 table
747 .register_node(node_id, 2, Tick::new(10))
748 .expect("node should register");
749 table
750 .assign_station(station_id, node_id, Tick::new(10))
751 .expect("station should assign");
752
753 let route = table
754 .resolve_gateway_client(&gateway, client_id)
755 .expect("route should resolve");
756 assert_eq!(route.client_id, client_id);
757 assert_eq!(route.station_id, station_id);
758 assert_eq!(route.node_id, node_id);
759 assert_eq!(route.gateway_generation, connected.route.generation);
760 assert_eq!(route.gateway_route_epoch, connected.route.route_epoch);
761 assert_eq!(route.station_route_epoch, 1);
762 assert_eq!(route.node_route_epoch, 1);
763 assert_eq!(route.node_state, DeploymentNodeState::Online);
764
765 table
766 .mark_draining(node_id)
767 .expect("draining node should remain routable");
768 assert_eq!(
769 table
770 .resolve_gateway_client(&gateway, client_id)
771 .expect("draining node should still route")
772 .node_state,
773 DeploymentNodeState::Draining
774 );
775
776 table
777 .mark_offline(node_id)
778 .expect("offline node should mark");
779 assert_eq!(
780 table
781 .resolve_gateway_client(&gateway, client_id)
782 .expect_err("offline node should reject delivery"),
783 GatewayDeliveryError::Deployment(DeploymentError::NodeUnavailable {
784 node_id,
785 state: DeploymentNodeState::Offline
786 })
787 );
788 }
789
790 #[test]
791 fn detects_and_marks_stale_nodes_offline() {
792 let mut table = DeploymentRouteTable::new(config());
793 table
794 .register_node(NodeId::new(1), 2, Tick::new(10))
795 .expect("node should register");
796 assert!(table.stale_nodes(Tick::new(13)).is_empty());
797 assert_eq!(table.stale_nodes(Tick::new(14)), vec![NodeId::new(1)]);
798 assert_eq!(table.mark_stale_offline(Tick::new(14)), 1);
799 assert_eq!(
800 table
801 .node_route(NodeId::new(1))
802 .expect("node route should exist")
803 .state,
804 DeploymentNodeState::Offline
805 );
806 }
807}