1use std::fs::{File, OpenOptions};
57use std::mem::size_of;
58use std::path::Path;
59use std::sync::atomic::{AtomicU32, AtomicU64, Ordering};
60
61use memmap2::{MmapMut, MmapOptions};
62
63pub const TOPOLOGY_MAGIC: u64 = 0x4150_544F_504F_3031;
64
65pub const DEFAULT_FAN_OUT_THRESHOLD: u32 = 3;
68pub const DEFAULT_FAN_IN_THRESHOLD: u32 = 3;
69
70#[derive(Debug, Clone, Copy, PartialEq, Eq)]
72#[repr(u32)]
73pub enum TopologyKind {
74 PointToPoint = 0,
75 BroadcastTree = 1,
76 AllToAllMesh = 2,
77}
78
79impl TopologyKind {
80 pub fn from_u32(v: u32) -> Self {
81 match v {
82 1 => Self::BroadcastTree,
83 2 => Self::AllToAllMesh,
84 _ => Self::PointToPoint,
85 }
86 }
87}
88
89#[repr(C, align(64))]
90pub struct TopologyHeader {
91 pub magic: u64,
92 pub n_nodes: u32,
93 pub fan_out_threshold: AtomicU32,
94 pub fan_in_threshold: AtomicU32,
95 _pad1: u32,
96 pub total_msgs: AtomicU64,
97 pub recommendation: AtomicU32,
98 pub recommendation_epoch: AtomicU64,
99 pub broadcast_root: AtomicU32,
100 _pad2: u32,
101 _pad3: [u8; 8],
102}
103
104const _: () = {
105 assert!(size_of::<TopologyHeader>() == 64);
106};
107
108pub const fn topology_file_size(n_nodes: usize) -> usize {
109 size_of::<TopologyHeader>() + n_nodes * n_nodes * size_of::<AtomicU64>()
110}
111
112#[derive(Debug, Clone, Copy, PartialEq, Eq)]
113pub enum TopologyError {
114 NodeIndexOutOfBounds,
115 LayoutMismatch,
116 IoError(std::io::ErrorKind),
117}
118
119impl From<std::io::Error> for TopologyError {
120 fn from(e: std::io::Error) -> Self { Self::IoError(e.kind()) }
121}
122
123#[derive(Debug, Clone, Copy, PartialEq, Eq)]
124pub struct TopologyStats {
125 pub total_msgs: u64,
126 pub max_fan_out: u32,
127 pub max_fan_in: u32,
128 pub max_fan_out_src: u32,
129 pub max_fan_in_dst: u32,
130 pub current_recommendation: TopologyKind,
131 pub recommendation_epoch: u64,
132}
133
134pub struct SharedTopologyMap {
135 _file: File,
136 mmap: MmapMut,
137 n_nodes: usize,
138 header_sidecar: subetha_core::HandshakeHeader,
139 ring_sidecar: Box<subetha_core::ObservationRing>,
140}
141
142unsafe impl Send for SharedTopologyMap {}
143unsafe impl Sync for SharedTopologyMap {}
144
145impl subetha_sidecar::AdaptiveInstance for SharedTopologyMap {
146 fn header(&self) -> &subetha_core::HandshakeHeader { &self.header_sidecar }
147 fn ring(&self) -> &subetha_core::ObservationRing { &self.ring_sidecar }
148 fn make_policy(&self) -> Box<dyn subetha_sidecar::Policy> {
149 Box::new(subetha_sidecar::NoMigrationPolicy)
150 }
151}
152
153impl SharedTopologyMap {
154 pub fn create(
155 path: impl AsRef<Path>,
156 n_nodes: usize,
157 ) -> Result<Self, TopologyError> {
158 Self::create_with_thresholds(
159 path, n_nodes,
160 DEFAULT_FAN_OUT_THRESHOLD, DEFAULT_FAN_IN_THRESHOLD,
161 )
162 }
163
164 pub fn create_with_thresholds(
165 path: impl AsRef<Path>,
166 n_nodes: usize,
167 fan_out_threshold: u32,
168 fan_in_threshold: u32,
169 ) -> Result<Self, TopologyError> {
170 assert!(n_nodes >= 1);
171 assert!(n_nodes <= u32::MAX as usize);
172 let total = topology_file_size(n_nodes);
173 let file = OpenOptions::new()
174 .read(true).write(true).create(true).truncate(true)
175 .open(path.as_ref())?;
176 file.set_len(total as u64)?;
177 let mut mmap = unsafe { MmapOptions::new().len(total).map_mut(&file)? };
178 let hdr = mmap.as_mut_ptr() as *mut TopologyHeader;
179 unsafe {
180 std::ptr::write_bytes(hdr as *mut u8, 0, size_of::<TopologyHeader>());
181 (*hdr).magic = TOPOLOGY_MAGIC;
182 (*hdr).n_nodes = n_nodes as u32;
183 (*hdr).fan_out_threshold.store(fan_out_threshold, Ordering::Release);
184 (*hdr).fan_in_threshold.store(fan_in_threshold, Ordering::Release);
185 (*hdr).recommendation.store(
186 TopologyKind::PointToPoint as u32, Ordering::Release,
187 );
188 }
189 Ok(Self {
191 _file: file, mmap, n_nodes,
192 header_sidecar: subetha_core::HandshakeHeader::new(),
193 ring_sidecar: Box::new(subetha_core::ObservationRing::new()),
194 })
195 }
196
197 pub fn open(
198 path: impl AsRef<Path>,
199 expected_n_nodes: usize,
200 ) -> Result<Self, TopologyError> {
201 let total = topology_file_size(expected_n_nodes);
202 let file = OpenOptions::new().read(true).write(true).open(path.as_ref())?;
203 if file.metadata()?.len() < total as u64 {
204 return Err(TopologyError::LayoutMismatch);
205 }
206 let mmap = unsafe { MmapOptions::new().len(total).map_mut(&file)? };
207 let hdr = unsafe { &*(mmap.as_ptr() as *const TopologyHeader) };
208 if hdr.magic != TOPOLOGY_MAGIC || hdr.n_nodes != expected_n_nodes as u32 {
209 return Err(TopologyError::LayoutMismatch);
210 }
211 Ok(Self {
212 _file: file, mmap, n_nodes: expected_n_nodes,
213 header_sidecar: subetha_core::HandshakeHeader::new(),
214 ring_sidecar: Box::new(subetha_core::ObservationRing::new()),
215 })
216 }
217
218 #[inline]
219 pub fn n_nodes(&self) -> usize { self.n_nodes }
220
221 fn header(&self) -> &TopologyHeader {
222 unsafe { &*(self.mmap.as_ptr() as *const TopologyHeader) }
223 }
224
225 fn edge(&self, src: u32, dst: u32) -> &AtomicU64 {
226 let idx = (src as usize) * self.n_nodes + (dst as usize);
227 let base = unsafe { self.mmap.as_ptr().add(size_of::<TopologyHeader>()) };
228 unsafe { &*(base.add(idx * size_of::<AtomicU64>()) as *const AtomicU64) }
229 }
230
231 pub fn record_send(&self, src: u32, dst: u32) -> Result<u64, TopologyError> {
234 if src as usize >= self.n_nodes || dst as usize >= self.n_nodes {
235 self.ring_sidecar
236 .push_op(crate::sidecar_ops::topology::OP_RECORD, 1);
237 return Err(TopologyError::NodeIndexOutOfBounds);
238 }
239 let prev = self.edge(src, dst).fetch_add(1, Ordering::AcqRel);
240 self.header().total_msgs.fetch_add(1, Ordering::AcqRel);
241 self.ring_sidecar
242 .push_op(crate::sidecar_ops::topology::OP_RECORD, 0);
243 Ok(prev + 1)
244 }
245
246 pub fn fan_out(&self, src: u32) -> u32 {
248 if src as usize >= self.n_nodes {
249 self.ring_sidecar
250 .push_op(crate::sidecar_ops::topology::OP_FAN_OUT, 1);
251 return 0;
252 }
253 let mut count = 0u32;
254 for d in 0..self.n_nodes as u32 {
255 if self.edge(src, d).load(Ordering::Acquire) > 0 {
256 count += 1;
257 }
258 }
259 self.ring_sidecar
260 .push_op(crate::sidecar_ops::topology::OP_FAN_OUT, 0);
261 count
262 }
263
264 pub fn fan_in(&self, dst: u32) -> u32 {
266 if dst as usize >= self.n_nodes {
267 self.ring_sidecar
268 .push_op(crate::sidecar_ops::topology::OP_FAN_IN, 1);
269 return 0;
270 }
271 let mut count = 0u32;
272 for s in 0..self.n_nodes as u32 {
273 if self.edge(s, dst).load(Ordering::Acquire) > 0 {
274 count += 1;
275 }
276 }
277 self.ring_sidecar
278 .push_op(crate::sidecar_ops::topology::OP_FAN_IN, 0);
279 count
280 }
281
282 pub fn max_fan_out(&self) -> (u32, u32) {
284 let mut max = 0u32;
285 let mut who = 0u32;
286 for s in 0..self.n_nodes as u32 {
287 let fo = self.fan_out(s);
288 if fo > max { max = fo; who = s; }
289 }
290 (max, who)
291 }
292
293 pub fn max_fan_in(&self) -> (u32, u32) {
295 let mut max = 0u32;
296 let mut who = 0u32;
297 for d in 0..self.n_nodes as u32 {
298 let fi = self.fan_in(d);
299 if fi > max { max = fi; who = d; }
300 }
301 (max, who)
302 }
303
304 pub fn recommend(&self) -> TopologyKind {
309 let (max_fan_out, _) = self.max_fan_out();
310 let (max_fan_in, _) = self.max_fan_in();
311 let fo_threshold = self.header().fan_out_threshold.load(Ordering::Acquire);
312 let fi_threshold = self.header().fan_in_threshold.load(Ordering::Acquire);
313 let kind = if max_fan_out >= fo_threshold && max_fan_in >= fi_threshold {
314 TopologyKind::AllToAllMesh
315 } else if max_fan_out >= fo_threshold {
316 TopologyKind::BroadcastTree
317 } else {
318 TopologyKind::PointToPoint
319 };
320 self.ring_sidecar
321 .push_op(crate::sidecar_ops::topology::OP_RECOMMEND, 0);
322 kind
323 }
324
325 pub fn publish_recommendation(&self) -> TopologyKind {
330 let kind = self.recommend();
331 let hdr = self.header();
332 hdr.recommendation.store(kind as u32, Ordering::Release);
333 hdr.recommendation_epoch.fetch_add(1, Ordering::Release);
334 if kind == TopologyKind::BroadcastTree {
337 let (_, root) = self.max_fan_out();
338 hdr.broadcast_root.store(root, Ordering::Release);
339 }
340 kind
341 }
342
343 pub fn read_recommendation(&self) -> TopologyKind {
345 TopologyKind::from_u32(
346 self.header().recommendation.load(Ordering::Acquire)
347 )
348 }
349
350 pub fn broadcast_root(&self) -> u32 {
354 self.header().broadcast_root.load(Ordering::Acquire)
355 }
356
357 pub fn recommendation_epoch(&self) -> u64 {
361 self.header().recommendation_epoch.load(Ordering::Acquire)
362 }
363
364 pub fn total_msgs(&self) -> u64 {
366 self.header().total_msgs.load(Ordering::Acquire)
367 }
368
369 pub fn stats(&self) -> TopologyStats {
371 let (max_fan_out, max_fan_out_src) = self.max_fan_out();
372 let (max_fan_in, max_fan_in_dst) = self.max_fan_in();
373 TopologyStats {
374 total_msgs: self.total_msgs(),
375 max_fan_out, max_fan_out_src,
376 max_fan_in, max_fan_in_dst,
377 current_recommendation: self.read_recommendation(),
378 recommendation_epoch: self.recommendation_epoch(),
379 }
380 }
381
382 pub fn reset_observations(&self) {
387 for s in 0..self.n_nodes as u32 {
388 for d in 0..self.n_nodes as u32 {
389 self.edge(s, d).store(0, Ordering::Release);
390 }
391 }
392 self.header().total_msgs.store(0, Ordering::Release);
393 }
394
395 pub fn set_thresholds(&self, fan_out: u32, fan_in: u32) {
398 self.header().fan_out_threshold.store(fan_out, Ordering::Release);
399 self.header().fan_in_threshold.store(fan_in, Ordering::Release);
400 }
401
402 pub fn flush(&self) -> Result<(), TopologyError> {
403 self.mmap.flush()?;
404 Ok(())
405 }
406
407 pub fn flush_async(&self) -> Result<(), TopologyError> {
411 self.mmap.flush_async()?;
412 Ok(())
413 }
414}
415
416#[cfg(test)]
417mod tests {
418 use super::*;
419 use std::sync::Arc;
420 use std::thread;
421
422 fn tmp(name: &str) -> std::path::PathBuf {
423 let mut p = std::env::temp_dir();
424 let pid = std::process::id();
425 p.push(format!("subetha-topology-{name}-{pid}.bin"));
426 p
427 }
428
429 #[test]
430 fn create_initial_state() {
431 let p = tmp("init");
432 let t = SharedTopologyMap::create(&p, 4).unwrap();
433 assert_eq!(t.n_nodes(), 4);
434 assert_eq!(t.total_msgs(), 0);
435 assert_eq!(t.max_fan_out(), (0, 0));
436 assert_eq!(t.max_fan_in(), (0, 0));
437 assert_eq!(t.recommend(), TopologyKind::PointToPoint);
438 std::fs::remove_file(&p).ok();
439 }
440
441 #[test]
442 fn record_send_increments_edge_and_total() {
443 let p = tmp("record");
444 let t = SharedTopologyMap::create(&p, 4).unwrap();
445 let new_count = t.record_send(0, 1).unwrap();
446 assert_eq!(new_count, 1);
447 assert_eq!(t.total_msgs(), 1);
448 t.record_send(0, 1).unwrap();
449 assert_eq!(t.total_msgs(), 2);
450 std::fs::remove_file(&p).ok();
451 }
452
453 #[test]
454 fn fan_out_in_count_distinct_edges() {
455 let p = tmp("fan");
456 let t = SharedTopologyMap::create(&p, 4).unwrap();
457 t.record_send(0, 1).unwrap();
459 t.record_send(0, 2).unwrap();
460 t.record_send(0, 3).unwrap();
461 assert_eq!(t.fan_out(0), 3);
462 assert_eq!(t.fan_in(0), 0);
463 t.record_send(0, 1).unwrap();
465 t.record_send(0, 1).unwrap();
466 assert_eq!(t.fan_out(0), 3);
467 assert_eq!(t.fan_in(2), 1);
469 std::fs::remove_file(&p).ok();
470 }
471
472 #[test]
473 fn out_of_bounds_record_returns_error() {
474 let p = tmp("oob");
475 let t = SharedTopologyMap::create(&p, 4).unwrap();
476 assert_eq!(t.record_send(4, 0).err(), Some(TopologyError::NodeIndexOutOfBounds));
477 assert_eq!(t.record_send(0, 99).err(), Some(TopologyError::NodeIndexOutOfBounds));
478 std::fs::remove_file(&p).ok();
479 }
480
481 #[test]
482 fn recommend_point_to_point_when_fan_low() {
483 let p = tmp("rec-p2p");
484 let t = SharedTopologyMap::create(&p, 4).unwrap();
485 for _ in 0..100 { t.record_send(0, 1).unwrap(); }
487 assert_eq!(t.recommend(), TopologyKind::PointToPoint);
488 std::fs::remove_file(&p).ok();
489 }
490
491 #[test]
492 fn recommend_broadcast_when_fan_out_high_only() {
493 let p = tmp("rec-bcast");
494 let t = SharedTopologyMap::create(&p, 5).unwrap();
495 for d in 1..5 { t.record_send(0, d).unwrap(); }
497 assert_eq!(t.fan_out(0), 4);
498 assert_eq!(t.max_fan_in(), (1, 1)); assert_eq!(t.recommend(), TopologyKind::BroadcastTree);
500 std::fs::remove_file(&p).ok();
501 }
502
503 #[test]
504 fn recommend_all_to_all_when_both_high() {
505 let p = tmp("rec-mesh");
506 let t = SharedTopologyMap::create(&p, 5).unwrap();
507 for s in 0..5u32 {
509 for d in 0..5u32 {
510 if s != d { t.record_send(s, d).unwrap(); }
511 }
512 }
513 assert_eq!(t.max_fan_out().0, 4);
514 assert_eq!(t.max_fan_in().0, 4);
515 assert_eq!(t.recommend(), TopologyKind::AllToAllMesh);
516 std::fs::remove_file(&p).ok();
517 }
518
519 #[test]
520 fn publish_recommendation_caches_for_o1_reads() {
521 let p = tmp("publish");
522 let t = SharedTopologyMap::create(&p, 5).unwrap();
523 for d in 1..5 { t.record_send(0, d).unwrap(); }
524 assert_eq!(t.recommendation_epoch(), 0);
525 let published = t.publish_recommendation();
526 assert_eq!(published, TopologyKind::BroadcastTree);
527 assert_eq!(t.read_recommendation(), TopologyKind::BroadcastTree);
528 assert_eq!(t.recommendation_epoch(), 1);
529 assert_eq!(t.broadcast_root(), 0);
531 std::fs::remove_file(&p).ok();
532 }
533
534 #[test]
535 fn cross_handle_observation_visible() {
536 let p = tmp("cross-handle");
537 let writer = SharedTopologyMap::create(&p, 4).unwrap();
538 let observer = SharedTopologyMap::open(&p, 4).unwrap();
539 writer.record_send(0, 1).unwrap();
540 writer.record_send(0, 2).unwrap();
541 writer.record_send(0, 3).unwrap();
542 assert_eq!(observer.fan_out(0), 3);
543 writer.publish_recommendation();
544 assert_eq!(observer.read_recommendation(), TopologyKind::BroadcastTree);
545 std::fs::remove_file(&p).ok();
546 }
547
548 #[test]
549 fn reset_observations_clears_all_edges() {
550 let p = tmp("reset");
551 let t = SharedTopologyMap::create(&p, 4).unwrap();
552 for s in 0..4u32 {
553 for d in 0..4u32 {
554 t.record_send(s, d).unwrap();
555 }
556 }
557 assert_eq!(t.total_msgs(), 16);
558 t.reset_observations();
559 assert_eq!(t.total_msgs(), 0);
560 assert_eq!(t.fan_out(0), 0);
561 assert_eq!(t.fan_in(0), 0);
562 std::fs::remove_file(&p).ok();
563 }
564
565 #[test]
566 fn concurrent_record_sends_count_correctly() {
567 let p = tmp("concurrent");
568 let t = Arc::new(SharedTopologyMap::create(&p, 4).unwrap());
569 let n_threads = 4;
570 let per_thread = 100;
571 let mut handles = vec![];
572 for src in 0..n_threads as u32 {
573 let t = t.clone();
574 handles.push(thread::spawn(move || {
575 for _ in 0..per_thread {
576 for dst in 0..4u32 {
577 if src != dst { t.record_send(src, dst).unwrap(); }
578 }
579 }
580 }));
581 }
582 for h in handles { h.join().unwrap(); }
583 assert_eq!(t.total_msgs(), (n_threads * 3 * per_thread) as u64);
585 for s in 0..n_threads as u32 {
587 assert_eq!(t.fan_out(s), 3, "src {s} should have fan_out 3");
588 }
589 assert_eq!(t.recommend(), TopologyKind::AllToAllMesh);
591 std::fs::remove_file(&p).ok();
592 }
593
594 #[test]
595 fn set_thresholds_adjusts_recommendation_policy() {
596 let p = tmp("thresholds");
597 let t = SharedTopologyMap::create(&p, 5).unwrap();
598 t.record_send(0, 1).unwrap();
600 t.record_send(0, 2).unwrap();
601 assert_eq!(t.recommend(), TopologyKind::PointToPoint);
602 t.set_thresholds(2, 3);
604 assert_eq!(t.recommend(), TopologyKind::BroadcastTree);
605 std::fs::remove_file(&p).ok();
606 }
607
608 #[test]
609 fn stats_snapshot_returns_full_picture() {
610 let p = tmp("stats");
611 let t = SharedTopologyMap::create(&p, 5).unwrap();
612 for d in 1..5u32 { t.record_send(0, d).unwrap(); }
613 t.publish_recommendation();
614 let s = t.stats();
615 assert_eq!(s.total_msgs, 4);
616 assert_eq!(s.max_fan_out, 4);
617 assert_eq!(s.max_fan_out_src, 0);
618 assert_eq!(s.max_fan_in, 1);
619 assert_eq!(s.current_recommendation, TopologyKind::BroadcastTree);
620 assert_eq!(s.recommendation_epoch, 1);
621 std::fs::remove_file(&p).ok();
622 }
623
624 #[test]
625 fn topology_kind_from_u32_round_trip() {
626 assert_eq!(TopologyKind::from_u32(0), TopologyKind::PointToPoint);
627 assert_eq!(TopologyKind::from_u32(1), TopologyKind::BroadcastTree);
628 assert_eq!(TopologyKind::from_u32(2), TopologyKind::AllToAllMesh);
629 assert_eq!(TopologyKind::from_u32(999), TopologyKind::PointToPoint);
631 }
632
633 #[test]
634 fn disk_persistence_survives_reopen() {
635 let p = tmp("disk");
636 {
637 let t = SharedTopologyMap::create(&p, 4).unwrap();
638 for d in 1..4u32 { t.record_send(0, d).unwrap(); }
639 t.publish_recommendation();
640 t.flush().unwrap();
641 }
642 let t2 = SharedTopologyMap::open(&p, 4).unwrap();
643 assert_eq!(t2.total_msgs(), 3);
644 assert_eq!(t2.fan_out(0), 3);
645 assert_eq!(t2.read_recommendation(), TopologyKind::BroadcastTree);
646 std::fs::remove_file(&p).ok();
647 }
648
649 #[test]
650 fn recommendation_demotes_when_observations_drop() {
651 let p = tmp("demote");
652 let t = SharedTopologyMap::create(&p, 5).unwrap();
653 for d in 1..5u32 { t.record_send(0, d).unwrap(); }
655 assert_eq!(t.recommend(), TopologyKind::BroadcastTree);
656 t.reset_observations();
658 for _ in 0..10 { t.record_send(0, 1).unwrap(); }
659 assert_eq!(t.recommend(), TopologyKind::PointToPoint);
660 std::fs::remove_file(&p).ok();
661 }
662}