optional max distance in shortest paths

configs/shortest_paths/coordinator.toml

@@ -1,6 +1,7 @@
 source = "https://www.cdc.gov/healthywater/swimming/"
 host = "0.0.0.0:5000"
 output_path = "data/shortest_paths"
+max_distance = 4
 
 [gossip]
 addr = "0.0.0.0:5001"

crates/core/src/config/mod.rs

@@ -686,6 +686,7 @@ pub struct ShortestPathCoordinatorConfig {
     pub gossip: GossipConfig,
     pub host: SocketAddr,
     pub output_path: String,
+    pub max_distance: Option<u64>,
 }
 
 #[derive(Clone, Debug, serde::Deserialize, serde::Serialize)]

crates/core/src/entrypoint/ampc/shortest_path/coordinator.rs

@@ -78,9 +78,12 @@ impl Setup for ShortestPathSetup {
     }
 
     fn setup_round(&self, dht: &Self::DhtTables) {
+        let meta = dht.meta.get(()).unwrap();
+
         dht.meta.set(
             (),
             Meta {
+                round: meta.round + 1,
                 round_had_changes: false,
             },
         );
@@ -92,18 +95,28 @@ impl Setup for ShortestPathSetup {
             (),
             Meta {
                 round_had_changes: true,
+                round: 0,
             },
         );
     }
 }
 
-pub struct ShortestPathFinish;
+pub struct ShortestPathFinish {
+    max_distance: Option<u64>,
+}
 
 impl Finisher for ShortestPathFinish {
     type Job = ShortestPathJob;
 
     fn is_finished(&self, dht: &ShortestPathTables) -> bool {
-        !dht.meta.get(()).unwrap().round_had_changes
+        let meta = dht.meta.get(()).unwrap();
+        if let Some(max_distance) = self.max_distance {
+            if meta.round >= max_distance {
+                return true;
+            }
+        }
+
+        !meta.round_had_changes
     }
 }
 
@@ -202,7 +215,12 @@ pub fn run(config: ShortestPathCoordinatorConfig) -> Result<()> {
     tracing::info!("starting {} jobs", jobs.len());
 
     let coordinator = build(&cluster.dht, cluster.workers.clone(), source);
-    let res = coordinator.run(jobs, ShortestPathFinish)?;
+    let res = coordinator.run(
+        jobs,
+        ShortestPathFinish {
+            max_distance: config.max_distance,
+        },
+    )?;
 
     let output_path = Path::new(&config.output_path);
 

crates/core/src/entrypoint/ampc/shortest_path/mapper.rs

@@ -21,7 +21,7 @@ use rustc_hash::FxHashMap;
 use super::{
     updated_nodes::{UpdatedNodes, UpdatedNodesKind},
     worker::ShortestPathWorker,
-    DhtTable as _, Mapper, Meta, ShortestPathJob, ShortestPathTables,
+    DhtTable as _, Mapper, ShortestPathJob, ShortestPathTables,
 };
 use crate::{
     ampc::{
@@ -237,13 +237,12 @@ impl Mapper for ShortestPathMapper {
                 dht.next()
                     .changed_nodes
                     .set(worker.shard(), new_changed_nodes.lock().unwrap().clone());
-                dht.next().meta.set(
-                    (),
-                    Meta {
-                        round_had_changes: round_had_changes
-                            .load(std::sync::atomic::Ordering::Relaxed),
-                    },
-                );
+
+                let mut meta = dht.next().meta.get(()).unwrap();
+                meta.round_had_changes =
+                    round_had_changes.load(std::sync::atomic::Ordering::Relaxed);
+
+                dht.next().meta.set((), meta);
             }
             ShortestPathMapper::UpdateChangedNodes => {
                 let all_changed_nodes: Vec<_> =

crates/core/src/entrypoint/ampc/shortest_path/mod.rs

@@ -44,6 +44,7 @@ use self::worker::{RemoteShortestPathWorker, ShortestPathWorker};
 )]
 pub struct Meta {
     round_had_changes: bool,
+    round: u64,
 }
 
 #[derive(bincode::Encode, bincode::Decode, Debug, Clone)]