flush old addrs

Cargo.toml

@@ -1,6 +1,11 @@
 [workspace]
 resolver = "3"
-members = ["rust/networking", "rust/exo_pyo3_bindings", "rust/util"]
+members = [
+  "rust/networking",
+  "rust/exo_pyo3_bindings",
+  "rust/util",
+  "rust/babblerd",
+]
 
 [workspace.package]
 version = "0.0.1"
@@ -43,7 +48,6 @@ log = "0.4"
 
 # networking
 libp2p = "0.56"
-libp2p-tcp = "0.44"
 
 [workspace.lints.rust]
 static_mut_refs = "warn"      # Or use "warn" instead of deny

bench/test_mlx_bandwidth.py

@@ -0,0 +1,377 @@
+# type: ignore
+import argparse
+import json
+import os
+import statistics
+import sys
+import tempfile
+import time
+
+import mlx.core as mx
+
+DTYPE_MAP = {
+    "float32": (mx.float32, 4),
+    "float16": (mx.float16, 2),
+    "bfloat16": (mx.bfloat16, 2),
+}
+
+SIZES = [
+    1 * 1024,
+    4 * 1024,
+    16 * 1024,
+    64 * 1024,
+    256 * 1024,
+    1 * 1024 * 1024,
+    4 * 1024 * 1024,
+    16 * 1024 * 1024,
+    64 * 1024 * 1024,
+    256 * 1024 * 1024,
+    1 * 1024 * 1024 * 1024,
+    2 * 1024 * 1024 * 1024,
+    4 * 1024 * 1024 * 1024,
+    8 * 1024 * 1024 * 1024,
+]
+
+
+def format_bytes(n: int) -> str:
+    if n >= 1024 * 1024 * 1024:
+        return f"{n / (1024 * 1024 * 1024):.0f} GB"
+    if n >= 1024 * 1024:
+        return f"{n / (1024 * 1024):.0f} MB"
+    if n >= 1024:
+        return f"{n / 1024:.0f} KB"
+    return f"{n} B"
+
+
+def format_time(seconds: float) -> str:
+    if seconds >= 1.0:
+        return f"{seconds:.3f} s"
+    if seconds >= 0.001:
+        return f"{seconds * 1000:.2f} ms"
+    return f"{seconds * 1_000_000:.1f} us"
+
+
+def format_bandwidth(bytes_per_sec: float) -> str:
+    if bytes_per_sec >= 1024 * 1024 * 1024:
+        return f"{bytes_per_sec / (1024 * 1024 * 1024):.2f} GB/s"
+    if bytes_per_sec >= 1024 * 1024:
+        return f"{bytes_per_sec / (1024 * 1024):.1f} MB/s"
+    return f"{bytes_per_sec / 1024:.1f} KB/s"
+
+
+def barrier(group: mx.distributed.Group) -> None:
+    mx.eval(mx.distributed.all_sum(mx.array(1.0), group=group))
+
+
+def init_ring(
+    rank: int, self_ip: str, peer_ip: str, port: int, tmpdir: str
+) -> mx.distributed.Group:
+    if rank == 0:
+        hosts = [f"{self_ip}:{port}", f"{peer_ip}:{port}"]
+    else:
+        hosts = [f"{peer_ip}:{port}", f"{self_ip}:{port}"]
+
+    hostfile = os.path.join(tmpdir, "hosts.json")
+    with open(hostfile, "w") as f:
+        json.dump(hosts, f)
+
+    for var in ("MLX_HOSTFILE", "MLX_RANK", "MLX_IBV_DEVICES", "MLX_JACCL_COORDINATOR"):
+        os.environ.pop(var, None)
+
+    os.environ["MLX_HOSTFILE"] = hostfile
+    os.environ["MLX_RANK"] = str(rank)
+    return mx.distributed.init(backend="ring", strict=True)
+
+
+def init_jaccl(
+    rank: int, interface: str, coordinator: str, port: int, tmpdir: str
+) -> mx.distributed.Group:
+    devices = [[None, interface], [interface, None]]
+    devfile = os.path.join(tmpdir, "devices.json")
+    with open(devfile, "w") as f:
+        json.dump(devices, f)
+
+    for var in ("MLX_HOSTFILE", "MLX_RANK", "MLX_IBV_DEVICES", "MLX_JACCL_COORDINATOR"):
+        os.environ.pop(var, None)
+
+    os.environ["MLX_IBV_DEVICES"] = devfile
+    os.environ["MLX_RANK"] = str(rank)
+    if rank == 0:
+        os.environ["MLX_JACCL_COORDINATOR"] = f"0.0.0.0:{port}"
+    else:
+        os.environ["MLX_JACCL_COORDINATOR"] = coordinator
+
+    return mx.distributed.init(backend="jaccl", strict=True)
+
+
+def bench_unidirectional(
+    group: mx.distributed.Group,
+    rank: int,
+    size_bytes: int,
+    dtype: mx.Dtype,
+    element_size: int,
+    warmup: int,
+    iterations: int,
+) -> list[float]:
+    n_elements = size_bytes // element_size
+    tensor = mx.random.normal(shape=(n_elements,)).astype(dtype)
+    mx.eval(tensor)
+
+    for _ in range(warmup):
+        if rank == 0:
+            sent = mx.distributed.send(tensor, dst=1, group=group)
+            mx.eval(sent)
+        else:
+            received = mx.distributed.recv_like(tensor, src=0, group=group)
+            mx.eval(received)
+        barrier(group)
+
+    times: list[float] = []
+    for _ in range(iterations):
+        barrier(group)
+        t0 = time.perf_counter()
+        if rank == 0:
+            sent = mx.distributed.send(tensor, dst=1, group=group)
+            mx.eval(sent)
+        else:
+            received = mx.distributed.recv_like(tensor, src=0, group=group)
+            mx.eval(received)
+        barrier(group)
+        t1 = time.perf_counter()
+        times.append(t1 - t0)
+
+    return times
+
+
+def bench_rtt(
+    group: mx.distributed.Group,
+    rank: int,
+    size_bytes: int,
+    dtype: mx.Dtype,
+    element_size: int,
+    warmup: int,
+    iterations: int,
+) -> list[float]:
+    n_elements = size_bytes // element_size
+    tensor = mx.random.normal(shape=(n_elements,)).astype(dtype)
+    mx.eval(tensor)
+
+    for _ in range(warmup):
+        if rank == 0:
+            sent = mx.distributed.send(tensor, dst=1, group=group)
+            mx.eval(sent)
+            received = mx.distributed.recv_like(tensor, src=1, group=group)
+            mx.eval(received)
+        else:
+            received = mx.distributed.recv_like(tensor, src=0, group=group)
+            mx.eval(received)
+            sent = mx.distributed.send(received, dst=0, group=group)
+            mx.eval(sent)
+        barrier(group)
+
+    times: list[float] = []
+    for _ in range(iterations):
+        barrier(group)
+        t0 = time.perf_counter()
+        if rank == 0:
+            sent = mx.distributed.send(tensor, dst=1, group=group)
+            mx.eval(sent)
+            received = mx.distributed.recv_like(tensor, src=1, group=group)
+            mx.eval(received)
+        else:
+            received = mx.distributed.recv_like(tensor, src=0, group=group)
+            mx.eval(received)
+            sent = mx.distributed.send(received, dst=0, group=group)
+            mx.eval(sent)
+        barrier(group)
+        t1 = time.perf_counter()
+        times.append(t1 - t0)
+
+    return times
+
+
+def bench_all_gather(
+    group: mx.distributed.Group,
+    rank: int,
+    size_bytes: int,
+    dtype: mx.Dtype,
+    element_size: int,
+    warmup: int,
+    iterations: int,
+) -> list[float]:
+    n_elements = (size_bytes // 2) // element_size
+    tensor = mx.random.normal(shape=(n_elements,)).astype(dtype)
+    mx.eval(tensor)
+
+    for _ in range(warmup):
+        gathered = mx.distributed.all_gather(tensor, group=group)
+        mx.eval(gathered)
+        barrier(group)
+
+    times: list[float] = []
+    for _ in range(iterations):
+        barrier(group)
+        t0 = time.perf_counter()
+        gathered = mx.distributed.all_gather(tensor, group=group)
+        mx.eval(gathered)
+        t1 = time.perf_counter()
+        times.append(t1 - t0)
+
+    return times
+
+
+def print_table(title: str, rows: list[dict[str, str]]) -> None:
+    print(f"\n=== {title} ===")
+    headers = ["Size", "Median", "Min", "Max", "Bandwidth"]
+    widths = [
+        max(len(h), max((len(r[h]) for r in rows), default=0)) + 2 for h in headers
+    ]
+    header_line = "".join(h.ljust(w) for h, w in zip(headers, widths, strict=True))
+    print(header_line)
+    print("-" * len(header_line))
+    for row in rows:
+        print("".join(row[h].ljust(w) for h, w in zip(headers, widths, strict=True)))
+
+
+def run_bench(
+    name: str,
+    bench_fn,
+    group: mx.distributed.Group,
+    rank: int,
+    dtype: mx.Dtype,
+    element_size: int,
+    warmup: int,
+    iterations: int,
+    bw_multiplier: int = 1,
+) -> None:
+    rows: list[dict[str, str]] = []
+    for size in SIZES:
+        if rank == 0:
+            print(f"  {name}: {format_bytes(size)}...", end="", flush=True)
+        times = bench_fn(group, rank, size, dtype, element_size, warmup, iterations)
+        if rank == 0:
+            med = statistics.median(times)
+            mn = min(times)
+            mx_ = max(times)
+            bw = (size * bw_multiplier) / med
+            rows.append(
+                {
+                    "Size": format_bytes(size),
+                    "Median": format_time(med),
+                    "Min": format_time(mn),
+                    "Max": format_time(mx_),
+                    "Bandwidth": format_bandwidth(bw),
+                }
+            )
+            print(f" {format_bandwidth(bw)}")
+    if rank == 0:
+        print_table(name, rows)
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        description="MLX Distributed Communication Benchmark"
+    )
+    subparsers = parser.add_subparsers(dest="backend", required=True)
+
+    ring_parser = subparsers.add_parser("ring")
+    ring_parser.add_argument("--rank", type=int, required=True, choices=[0, 1])
+    ring_parser.add_argument("--self-ip", required=True)
+    ring_parser.add_argument("--peer-ip", required=True)
+    ring_parser.add_argument("--port", type=int, default=5555)
+
+    jaccl_parser = subparsers.add_parser("jaccl")
+    jaccl_parser.add_argument("--rank", type=int, required=True, choices=[0, 1])
+    jaccl_parser.add_argument("--interface", required=True)
+    jaccl_parser.add_argument(
+        "--coordinator",
+        type=str,
+        default=None,
+        help="IP:PORT of rank 0 (required for rank 1)",
+    )
+    jaccl_parser.add_argument(
+        "--port", type=int, default=9999, help="Coordinator port (rank 0 only)"
+    )
+
+    for p in [ring_parser, jaccl_parser]:
+        p.add_argument("--warmup", type=int, default=3)
+        p.add_argument("--iterations", type=int, default=10)
+        p.add_argument("--dtype", choices=list(DTYPE_MAP.keys()), default="float32")
+
+    args = parser.parse_args()
+
+    if args.backend == "jaccl" and args.rank == 1 and args.coordinator is None:
+        jaccl_parser.error("--coordinator is required for rank 1")
+
+    return args
+
+
+def main() -> int:
+    args = parse_args()
+    dtype, element_size = DTYPE_MAP[args.dtype]
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        if args.backend == "ring":
+            print(f"Initializing ring backend (rank {args.rank})...")
+            group = init_ring(args.rank, args.self_ip, args.peer_ip, args.port, tmpdir)
+        else:
+            print(f"Initializing jaccl backend (rank {args.rank})...")
+            group = init_jaccl(
+                args.rank, args.interface, args.coordinator or "", args.port, tmpdir
+            )
+
+        print(f"Rank {group.rank()} of {group.size()} initialized")
+        barrier(group)
+
+        if args.rank == 0:
+            print("\nMLX Distributed Communication Benchmark")
+            print(
+                f"Backend: {args.backend} | Dtype: {args.dtype} | Warmup: {args.warmup} | Iterations: {args.iterations}"
+            )
+
+        run_bench(
+            "Unidirectional (rank 0 -> rank 1)",
+            bench_unidirectional,
+            group,
+            args.rank,
+            dtype,
+            element_size,
+            args.warmup,
+            args.iterations,
+        )
+        run_bench(
+            "Round-Trip (ping-pong)",
+            bench_rtt,
+            group,
+            args.rank,
+            dtype,
+            element_size,
+            args.warmup,
+            args.iterations,
+            bw_multiplier=2,
+        )
+        run_bench(
+            "All-Gather",
+            bench_all_gather,
+            group,
+            args.rank,
+            dtype,
+            element_size,
+            args.warmup,
+            args.iterations,
+        )
+
+        if args.rank == 0:
+            print("\nDone.")
+        else:
+            print("Rank 1 complete.")
+
+    return 0
+
+
+if __name__ == "__main__":
+    try:
+        sys.exit(main())
+    except KeyboardInterrupt:
+        print("\nInterrupted.")
+        sys.exit(1)

flake.nix

@@ -112,7 +112,9 @@
             };
           };
 
-          packages = lib.optionalAttrs pkgs.stdenv.hostPlatform.isDarwin (
+          packages = {
+            babeld = pkgs.callPackage ./nix/babeld.nix { };
+          } // lib.optionalAttrs pkgs.stdenv.hostPlatform.isDarwin (
             let
               uvLock = builtins.fromTOML (builtins.readFile ./uv.lock);
               mlxPackage = builtins.head (builtins.filter (p: p.name == "mlx" && p.source ? git) uvLock.package);
@@ -157,9 +159,6 @@
                 just
                 jq
               ]
-              ++ lib.optionals stdenv.isLinux [
-                unixtools.ifconfig
-              ]
               ++ lib.optionals stdenv.isDarwin [
                 macmon
               ];

nix/babeld.nix

@@ -0,0 +1,25 @@
+{ stdenv
+, lib
+, fetchurl
+}:
+stdenv.mkDerivation rec {
+  pname = "babeld";
+  version = "1.13.1";
+
+  src = fetchurl {
+    url = "https://www.irif.fr/~jch/software/files/${pname}-${version}.tar.gz";
+    hash = "sha256-FfJNJtoMz8Bzq83vAwnygeRoTyqnESb4JlcsTIRejdk=";
+  };
+
+  outputs = [
+    "out"
+    "man"
+  ];
+
+  makeFlags = [
+    "PREFIX=${placeholder "out"}"
+    "ETCDIR=${placeholder "out"}/etc"
+  ]
+  ++ lib.optional stdenv.isDarwin "LDLIBS=''";
+}
+

python/parts.nix

@@ -185,6 +185,7 @@
         exo-eval = mkBenchScript "exo-eval" (inputs.self + /bench/exo_eval.py);
         exo-eval-tool-calls = mkBenchScript "exo-eval-tool-calls" (inputs.self + /bench/eval_tool_calls.py);
         exo-get-all-models-on-cluster = mkSimplePythonScript "exo-get-all-models-on-cluster" (inputs.self + /tests/get_all_models_on_cluster.py);
+        exo-mlx-bandwidth-test = mkPythonScript "exo-mlx-bandwidth-test" (inputs.self + /bench/test_mlx_bandwidth.py);
       };
 
       checks = {

result

@@ -0,0 +1 @@
+/nix/store/41rd4g2p69a2106qscl3xvkzmh6j7nry-babblerd

rust/babblerd/Cargo.toml

@@ -0,0 +1,19 @@
+[package]
+name = "babblerd"
+version.workspace = true
+edition.workspace = true
+
+[dependencies]
+color-eyre = "0.6.5"
+futures-lite.workspace = true
+ipnet = "2.12.0"
+libc = "0.2.183"
+n0-watcher = "0.6.1"
+netwatch = "0.14.0"
+rand = "0.10.0"
+tokio = { workspace = true, features = ["full"] }
+tracing = "0.1.44"
+tracing-subscriber = { version = "0.3.22", features = ["env-filter"] }
+
+[lints]
+workspace = true

rust/babblerd/src/lib.rs

@@ -0,0 +1,548 @@
+pub use babel::{babel, handle_listener};
+pub use error::{BabbleError, Result};
+pub use if_watcher::{PREFIX, watch};
+pub mod error {
+    use std::io;
+
+    pub type Result<T> = core::result::Result<T, BabbleError>;
+    #[derive(Debug)]
+    pub enum BabbleError {
+        Io(io::Error),
+        Unspecified,
+        BabeldCrashed(Option<i32>),
+        FailedToSetIp,
+        Other(String),
+    }
+    impl std::fmt::Display for BabbleError {
+        // use the debug display for now
+        #[inline]
+        fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+            write!(f, "{self:?}")
+        }
+    }
+    impl std::error::Error for BabbleError {}
+    impl From<io::Error> for BabbleError {
+        #[inline]
+        fn from(value: io::Error) -> Self {
+            Self::Io(value)
+        }
+    }
+}
+pub mod if_watcher {
+    #[cfg(target_os = "linux")]
+    use std::path::PathBuf;
+    use std::{
+        collections::HashSet,
+        net::{IpAddr, Ipv6Addr},
+    };
+
+    use futures_lite::StreamExt;
+    use ipnet::Ipv6Net;
+    use n0_watcher::Watcher;
+    use netwatch::interfaces::{Interface, IpNet};
+    use tokio::sync::mpsc;
+
+    use crate::{
+        BabbleError, Result,
+        babel::Babble,
+        ip_manager::{add_ip, remove_ip},
+    };
+
+    pub const PREFIX: Ipv6Net = Ipv6Net::new_assert(
+        Ipv6Addr::new(0xfde0, 0x20c6, 0x1fa7, 0xffff, 0, 0, 0, 0),
+        48,
+    );
+
+    trait IfaceExt {
+        fn has_link_local_v6(&self) -> bool;
+        fn is_real_interface(&self) -> bool;
+        fn will_babel(&self) -> bool;
+        fn get_v6_in(&self, prefix: Ipv6Net) -> Option<Ipv6Net>;
+    }
+    impl IfaceExt for Interface {
+        fn will_babel(&self) -> bool {
+            self.has_link_local_v6() && self.is_real_interface() && self.is_up()
+        }
+
+        fn has_link_local_v6(&self) -> bool {
+            let mut has = false;
+            for addr in self.addrs() {
+                let IpAddr::V6(a) = addr.addr() else {
+                    continue;
+                };
+                if a.is_unicast_link_local() {
+                    has = true;
+                    break;
+                }
+            }
+            has
+        }
+        fn is_real_interface(&self) -> bool {
+            // macos is weird. en0 & en1 are ethernet & wifi (varies which is which by device). en3+ is thunderbolt, but at some point becomes usb ethernet.
+            if self.name().strip_prefix("en").is_none()
+            //.and_then(|s| s.parse::<u8>().ok())
+            //.is_none_or(|_n| false)
+            {
+                return false;
+            }
+            #[cfg(target_os = "linux")]
+            {
+                if !PathBuf::from(format!("/sys/class/net/{}/device", self.name())).exists() {
+                    tracing::debug!(
+                        "skipping interface {} as it doesn't correspond to a physical link",
+                        self.name()
+                    );
+                    return false;
+                }
+                let dev_type_path = PathBuf::from(format!("/sys/class/net/{}/type", self.name()));
+                if !dev_type_path.exists() {
+                    tracing::debug!(
+                        "skipping interface {} with no type file at {:?}",
+                        self.name(),
+                        dev_type_path.to_str()
+                    );
+                    return false;
+                }
+                let Ok(dev_type) = std::fs::read_to_string(dev_type_path) else {
+                    return false;
+                };
+                if dev_type.trim() != "1" {
+                    tracing::debug!(
+                        "skipping interface {} with type {:?}",
+                        self.name(),
+                        dev_type
+                    );
+                    return false;
+                }
+            }
+            true
+        }
+        fn get_v6_in(&self, prefix: Ipv6Net) -> Option<Ipv6Net> {
+            for addr in self.addrs() {
+                if let IpNet::V6(v6) = addr
+                    && prefix.contains(&v6.addr())
+                {
+                    return Some(v6);
+                }
+            }
+            None
+        }
+    }
+
+    #[tracing::instrument(skip(send))]
+    pub async fn watch(my_range: Ipv6Net, send: mpsc::Sender<Babble>) -> Result<()> {
+        assert!(PREFIX.contains(&my_range));
+        let mut ready_ifaces = HashSet::new();
+        // 0 is reserved for the first loopback address
+        let mut iface_num: u16 = 1;
+
+        tracing::info!("starting interface monitor");
+        let mon = netwatch::netmon::Monitor::new()
+            .await
+            .map_err(|_| BabbleError::Unspecified)?;
+        let mut mon_stream = mon.interface_state().stream();
+
+        while let Some(s) = mon_stream.next().await {
+            for iface in s.interfaces.values() {
+                if !iface.is_real_interface() {
+                    continue;
+                }
+                for addr in iface.addrs() {
+                    if let IpNet::V6(v6) = addr
+                        && PREFIX.contains(&v6.addr())
+                        && !my_range.contains(&v6.addr())
+                    {
+                        tracing::info!("removing stale ip {v6} from {}", iface.name());
+                        // don't really care if this fails
+                        if let Err(e) = remove_ip(v6, iface).await {
+                            tracing::warn!(%e, "failed to remove ip");
+                        }
+                    }
+                }
+                if !iface.will_babel() {
+                    continue;
+                }
+                if iface.get_v6_in(my_range).is_none() {
+                    let addr = Ipv6Net::new_assert(
+                        Ipv6Addr::from_bits(my_range.addr().to_bits() | u128::from(iface_num)),
+                        128,
+                    );
+                    iface_num += 1;
+                    assert!(iface_num < u16::MAX, "Really? u16::MAX interfaces?");
+                    tracing::info!("adding new ip {addr} to {}", iface.name());
+                    add_ip(addr, iface).await?;
+                }
+
+                if ready_ifaces.insert(iface.name().to_owned()) {
+                    tracing::info!("telling babeld to watch {}", iface.name());
+                    let Ok(()) = send.send(Babble::AddIface(iface.name().to_owned())).await else {
+                        return Ok(());
+                    };
+                }
+            }
+        }
+        tracing::info!("stopping interface monitor");
+
+        Ok(())
+    }
+}
+pub mod babel {
+    #[tracing::instrument(skip_all)]
+    pub async fn handle_listener(sock: UnixStream, mut receiver: broadcast::Receiver<String>) {
+        tracing::info!("new socket conn");
+        let (reader, mut write) = sock.into_split();
+        let mut reader = BufReader::new(reader).lines();
+        loop {
+            tokio::select! {
+                read = reader.next_line() => {
+                    let Ok(Some(_)) = read else { break; };
+                }
+                recv = receiver.recv() => {
+                    let res = match recv {
+                        Err(broadcast::error::RecvError::Lagged(_)) => { tracing::warn!("receiver lagged"); continue; },
+                        Err(broadcast::error::RecvError::Closed) => { tracing::debug!("receiver closed, dropping connection"); break; },
+                        Ok(s) => write.write_all(format!("{s}\n").as_bytes()).await,
+                    };
+                    if let Err(e) = res { tracing::warn!(error=%e, "failed to write to socket"); break; }
+                }
+            };
+        }
+        tracing::info!("closing socket conn");
+        _ = write.shutdown().await;
+    }
+
+    #[cfg(target_os = "macos")]
+    const PRIVATE_SOCK_PATH: &str = "/var/run/babbler/private/babeld.sock";
+    #[cfg(target_os = "linux")]
+    const PRIVATE_SOCK_PATH: &str = "/run/babbler/private/babeld.sock";
+    #[cfg(target_os = "macos")]
+    const PRIVATE_DIR: &str = "/var/run/babbler/private";
+    #[cfg(target_os = "linux")]
+    const PRIVATE_DIR: &str = "/run/babbler/private";
+
+    use ipnet::Ipv6Net;
+    use std::fs::Permissions;
+    use std::io;
+    use std::os::unix::fs::PermissionsExt;
+    use tokio::time::Duration;
+
+    use futures_lite::FutureExt;
+    use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader, Lines};
+    use tokio::net::UnixStream;
+    use tokio::net::unix::{OwnedReadHalf, OwnedWriteHalf};
+    use tokio::process::Command;
+    use tokio::sync::{broadcast, mpsc};
+
+    use crate::{BabbleError, Result};
+
+    #[derive(Debug)]
+    pub enum Babble {
+        AddIface(String),
+    }
+
+    pub struct BabeldProcess {
+        proc: tokio::process::Child,
+    }
+    impl Drop for BabeldProcess {
+        #[inline]
+        fn drop(&mut self) {
+            // emergency sigkill babeld process to prevent leakage
+            match self.proc.try_wait() {
+                Ok(None) => {}
+                Ok(Some(sc)) => {
+                    if !sc.success() {
+                        _ = self.proc.start_kill();
+                    }
+                }
+                _ => {
+                    _ = self.proc.start_kill();
+                }
+            }
+        }
+    }
+    impl BabeldProcess {
+        #[tracing::instrument]
+        async fn spawn(my_range: Ipv6Net, iface: String) -> Result<Self> {
+            tokio::fs::create_dir_all(PRIVATE_DIR).await?;
+            tokio::fs::set_permissions(PRIVATE_DIR, Permissions::from_mode(0o0700)).await?;
+            tracing::info!("spawning babeld socket in {PRIVATE_SOCK_PATH}");
+            let res = match Command::new("babeld")
+                .arg("-G")
+                .arg(PRIVATE_SOCK_PATH)
+                .arg("-I")
+                .arg(format!("{PRIVATE_DIR}/babeld.pid"))
+                .arg("-C")
+                .arg(format!("redistribute local ip {my_range}"))
+                .arg("-C")
+                .arg("redistribute local deny")
+                .arg(iface)
+                .spawn()
+            {
+                Ok(proc) => Ok(Self { proc }),
+                Err(e) => {
+                    tracing::warn!(error=%e, "failed to spawn babeld");
+                    Err(e.into())
+                }
+            };
+            tokio::time::sleep(Duration::from_millis(10)).await;
+            while !matches!(tokio::fs::try_exists(PRIVATE_SOCK_PATH).await, Ok(true)) {
+                tracing::info!("where is the sock");
+                tokio::time::sleep(Duration::from_secs(1)).await;
+            }
+            res
+        }
+
+        #[tracing::instrument(skip(read, write, send))]
+        async fn query(
+            read: &mut Lines<BufReader<OwnedReadHalf>>,
+            write: &mut OwnedWriteHalf,
+            send: &broadcast::Sender<String>,
+            cmd: &str,
+        ) -> io::Result<Option<bool>> {
+            write.write_all(cmd.as_bytes()).await?;
+            loop {
+                let Some(line) = read.next_line().await? else {
+                    tracing::warn!("babeld closed unexpectedly");
+                    return Ok(None);
+                };
+                tracing::info!("[babel] {:?}", line);
+                let ret = match line.as_str() {
+                    "ok" => Ok(Some(true)),
+                    "bad" => {
+                        tracing::warn!("malformed message sent to babeld");
+                        Ok(Some(false))
+                    }
+                    _ if line.starts_with("no") => {
+                        tracing::warn!("message rejected");
+                        Ok(Some(false))
+                    }
+                    _ => Ok(None),
+                };
+                let Ok(_) = send.send(line) else {
+                    return Ok(None);
+                };
+                if !matches!(ret, Ok(None)) {
+                    return ret;
+                }
+            }
+        }
+
+        #[tracing::instrument(skip_all)]
+        async fn supervise(
+            &self,
+            mut recv: mpsc::Receiver<Babble>,
+            send: broadcast::Sender<String>,
+        ) -> Result<()> {
+            std::fs::set_permissions(PRIVATE_SOCK_PATH, Permissions::from_mode(0o0600))?;
+            tracing::debug!("connecting to babeld.sock");
+            let (reader, mut writer) = UnixStream::connect(PRIVATE_SOCK_PATH).await?.into_split();
+            let mut babel_lines = BufReader::new(reader).lines();
+            while let Some(s) = babel_lines.next_line().await? {
+                tracing::debug!("[babeld] {}", s);
+                if s == "ok" {
+                    break;
+                }
+            }
+            tracing::info!("babeld ok");
+            /* TODO(evan): push rather than pull
+            if Self::query(&mut babel_lines, &mut writer, "monitor\n")
+                .await?
+                .is_none()
+            {
+                return Ok(());
+            };
+            */
+
+            let mut interval = tokio::time::interval(Duration::from_secs(5));
+            interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
+            loop {
+                tokio::select! {
+                    _ = interval.tick() => {
+                        if Self::query(&mut babel_lines, &mut writer, &send, "dump\n")
+                            .await?
+                            .is_none()
+                        {
+                            return Ok(());
+                        }
+                    },
+                    babble = recv.recv() => {
+                        tracing::debug!("[babble] {:?}", babble);
+                        let Some(babble) = babble else {
+                            break;
+                        };
+                        match babble {
+                            Babble::AddIface(iface) => {
+                                Self::query(&mut babel_lines, &mut writer, &send, format!("interface {iface}\n").as_ref()).await?;
+                            }
+                        }
+                    },
+                    line = babel_lines.next_line() => {
+                        let Ok(Some(line)) = line else {
+                            break;
+                        };
+                        tracing::debug!("[babeld] {}", line);
+                        let Ok(_) = send.send(line) else {
+                            break;
+                        };
+                    },
+                }
+            }
+            Ok(())
+        }
+
+        async fn shutdown(mut self) -> Result<()> {
+            let kill_res = if let Some(pid) = self.proc.id() {
+                let pid: libc::pid_t = pid.try_into().expect("pid overflow");
+                // SAFETY: pid >= 0, freshly checked.
+                let rc = unsafe { libc::kill(pid, libc::SIGINT) };
+                let rc_err = if rc != 0 && rc != libc::ESRCH {
+                    Err(io::Error::last_os_error().into())
+                } else {
+                    Ok(())
+                };
+                let exit_code = async { Some(self.proc.wait().await) }
+                    .or(async {
+                        tokio::time::sleep(Duration::from_secs(5)).await;
+                        None
+                    })
+                    .await;
+                match exit_code {
+                    Some(Ok(code)) => {
+                        if code.success() {
+                            rc_err
+                        } else {
+                            rc_err.and_then(|()| Err(BabbleError::BabeldCrashed(code.code())))
+                        }
+                    }
+                    Some(Err(e)) => Err(e.into()),
+                    None => {
+                        self.proc.kill().await?;
+                        rc_err.and(Err(BabbleError::BabeldCrashed(None)))
+                    }
+                }
+            } else {
+                Ok(())
+            };
+            let rem_res = match std::fs::remove_file(PRIVATE_SOCK_PATH) {
+                Ok(()) => Ok(()),
+                Err(e) if e.kind() == io::ErrorKind::NotFound => Ok(()),
+                Err(e) => Err(e.into()),
+            };
+            kill_res.and(rem_res)
+        }
+    }
+
+    #[tracing::instrument(skip(send, recv))]
+    pub async fn babel(
+        my_range: Ipv6Net,
+        mut recv: mpsc::Receiver<Babble>,
+        send: broadcast::Sender<String>,
+    ) -> Result<()> {
+        let iface = loop {
+            match recv.recv().await {
+                Some(Babble::AddIface(iface)) => {
+                    break iface;
+                }
+                None => return Ok(()),
+            }
+        };
+
+        let babel = BabeldProcess::spawn(my_range, iface).await?;
+        let res1 = babel.supervise(recv, send).await;
+        let res2 = babel.shutdown().await;
+        res1.and(res2)
+    }
+}
+pub(crate) mod ip_manager {
+    pub use sys::add_ip;
+    pub use sys::remove_ip;
+
+    #[cfg(target_os = "linux")]
+    mod sys {
+        use ipnet::Ipv6Net;
+        use netwatch::interfaces::Interface;
+
+        use crate::{BabbleError, Result};
+        use tokio::process::Command;
+
+        #[tracing::instrument]
+        pub async fn add_ip(subnet: Ipv6Net, iface: &Interface) -> Result<()> {
+            let out = Command::new("ip")
+                .arg("addr")
+                .arg("add")
+                .arg(format!("{subnet}"))
+                .arg("dev")
+                .arg(iface.name())
+                .output()
+                .await?;
+            if out.status.success() {
+                Ok(())
+            } else {
+                Err(BabbleError::FailedToSetIp)
+            }
+        }
+
+        #[tracing::instrument]
+        pub async fn remove_ip(v6: Ipv6Net, iface: &Interface) -> Result<()> {
+            let out = Command::new("ip")
+                .arg("addr")
+                .arg("del")
+                .arg(format!("{v6}"))
+                .arg("dev")
+                .arg(iface.name())
+                .output()
+                .await?;
+            if out.status.success() {
+                Ok(())
+            } else {
+                let std_err = String::from_utf8_lossy(&out.stdout);
+                tracing::debug!(%std_err);
+                Err(BabbleError::FailedToSetIp)
+            }
+        }
+    }
+
+    #[cfg(target_os = "macos")]
+    mod sys {
+        use ipnet::Ipv6Net;
+        use netwatch::interfaces::Interface;
+
+        use crate::BabbleError;
+        use crate::Result;
+        use tokio::process::Command;
+
+        #[tracing::instrument]
+        pub async fn add_ip(subnet: Ipv6Net, iface: &Interface) -> Result<()> {
+            let out = Command::new("ifconfig")
+                .arg(iface.name())
+                .arg("inet6")
+                .arg(format!("{subnet}"))
+                .arg("add")
+                .output()
+                .await?;
+            if out.status.success() {
+                Ok(())
+            } else {
+                Err(BabbleError::FailedToSetIp)
+            }
+        }
+
+        #[tracing::instrument]
+        pub async fn remove_ip(v6: Ipv6Net, iface: &Interface) -> Result<()> {
+            let out = Command::new("ifconfig")
+                .arg(iface.name())
+                .arg("inet6")
+                .arg(format!("{v6}"))
+                .arg("delete")
+                .output()
+                .await?;
+            if out.status.success() {
+                Ok(())
+            } else {
+                let std_err = String::from_utf8_lossy(&out.stdout);
+                tracing::debug!(%std_err);
+                Err(BabbleError::FailedToSetIp)
+            }
+        }
+    }
+}

rust/babblerd/src/main.rs

@@ -0,0 +1,167 @@
+use std::{fs::Permissions, io, net::Ipv6Addr, os::unix::fs::PermissionsExt};
+
+use babblerd::babel::handle_listener;
+use color_eyre::eyre::{WrapErr, eyre};
+use ipnet::Ipv6Net;
+use tokio::{
+    net::UnixListener,
+    signal,
+    sync::{broadcast, mpsc},
+    task::{JoinHandle, JoinSet},
+};
+
+#[cfg(target_os = "macos")]
+const PUBLIC_DIR: &str = "/var/run/babbler";
+#[cfg(target_os = "linux")]
+const PUBLIC_DIR: &str = "/run/babbler";
+#[cfg(target_os = "macos")]
+const PUBLIC_SOCK_PATH: &str = "/var/run/babbler/babblerd.sock";
+#[cfg(target_os = "linux")]
+const PUBLIC_SOCK_PATH: &str = "/run/babbler/babblerd.sock";
+
+enum State {
+    Idle,
+    Active {
+        recv: broadcast::Receiver<String>,
+        babel: JoinHandle<babblerd::Result<()>>,
+        watcher: JoinHandle<babblerd::Result<()>>,
+        listeners: JoinSet<()>,
+    },
+}
+
+#[tokio::main]
+async fn main() -> color_eyre::Result<()> {
+    color_eyre::install()?;
+    tracing_subscriber::fmt()
+        .with_env_filter(tracing_subscriber::EnvFilter::from_default_env())
+        .init();
+    // cleanup old data
+    match std::fs::remove_file(PUBLIC_SOCK_PATH) {
+        Err(e) if e.kind() != io::ErrorKind::NotFound => {
+            return Err(e.into());
+        }
+        Ok(()) => {
+            tracing::info!("cleaned up old file at {PUBLIC_SOCK_PATH}");
+        }
+        _ => {}
+    }
+    std::fs::create_dir_all(PUBLIC_DIR)?;
+    // make our directory world readable
+    if let Err(e) = std::fs::set_permissions(PUBLIC_DIR, Permissions::from_mode(0o0755)) {
+        if e.kind() == io::ErrorKind::PermissionDenied {
+            return Err(eyre!(
+                "Insufficient permissions to run daemon -- did you forget sudo?"
+            ));
+        }
+        return Err(e.into());
+    }
+
+    let res = inner_main().await;
+    _ = std::fs::remove_file(PUBLIC_SOCK_PATH);
+    res
+}
+
+async fn inner_main() -> color_eyre::Result<()> {
+    tracing::info!("creating socket at {PUBLIC_SOCK_PATH}");
+    let public_socket = UnixListener::bind(PUBLIC_SOCK_PATH)?;
+    // make our socket world accessible
+    std::fs::set_permissions(PUBLIC_SOCK_PATH, Permissions::from_mode(0o0666))?;
+    let mut babbler: State = State::Idle;
+
+    loop {
+        babbler = match babbler {
+            State::Idle => {
+                tracing::info!("waiting for connection");
+                tokio::select! {
+                    sig = signal::ctrl_c() => {
+                        sig?;
+                        break;
+                    }
+                    sock = public_socket.accept() => {
+                        let sock = sock?.0;
+                        tracing::info!("starting babeld");
+                        let (br_send, br_recv) = broadcast::channel(1024);
+                        let (mp_send, mp_recv) = mpsc::channel(32);
+                        // node id is a PREFIX/64, NODE_ID/48 and an INTERFACE_ID/16
+                        let ip_node_id = u128::from(rand::random::<u64>() & (u64::MAX>>16)) << 16;
+                        let my_range = Ipv6Net::new_assert(
+                            Ipv6Addr::from_bits(babblerd::PREFIX.addr().to_bits() | ip_node_id),
+                            112,
+                        );
+                        let babel = tokio::spawn(babblerd::babel(my_range, mp_recv, br_send));
+                        let watcher = tokio::spawn(babblerd::watch(my_range, mp_send));
+                        let mut listeners = JoinSet::new();
+                        listeners.spawn(handle_listener(sock, br_recv.resubscribe()));
+                        State::Active { recv: br_recv, babel, watcher, listeners }
+                    }
+                }
+            }
+            State::Active {
+                recv,
+                mut babel,
+                mut watcher,
+                mut listeners,
+            } => {
+                tokio::select! {
+                    sig = signal::ctrl_c() => {
+                        sig?;
+                        drop(recv);
+                        watcher.abort();
+                        if let Ok(e) = watcher.await {
+                            e.wrap_err("while ctrl-c")?;
+                        }
+                        babel.await?.wrap_err("while ctrl-c")?;
+                        while let Some(res) = listeners.join_next().await {
+                            res.wrap_err("while ctrl-c")?;
+                        }
+                        break;
+                    }
+                    next_join_result = listeners.join_next(), if !listeners.is_empty() => {
+                        next_join_result.expect("checked")?;
+                        tracing::info!("dropped a listener");
+                        if listeners.is_empty() {
+                            tracing::info!("stopping babeld");
+                            drop(recv);
+
+                            watcher.abort();
+                            if let Ok(e) = watcher.await {
+                                e.wrap_err("while closing listeners")?;
+                            }
+                            babel.await?.wrap_err("while closing listeners")?;
+
+                            State::Idle
+                        } else {
+                            State::Active { recv, babel, watcher, listeners }
+                        }
+                    }
+                    sock = public_socket.accept() => {
+                        listeners.spawn(handle_listener(sock?.0, recv.resubscribe()));
+                        State::Active { recv, babel, watcher, listeners }
+                    }
+                    res = &mut watcher => {
+                        res??;
+                        drop(recv);
+                        babel.await?.wrap_err("while closing watcher")?;
+                        while let Some(res2) = listeners.join_next().await {
+                            res2.wrap_err("while closing watcher")?;
+                        }
+                        State::Idle
+                    }
+                    res = &mut babel => {
+                        res??;
+                        drop(recv);
+                        watcher.abort();
+                        if let Ok(e) = watcher.await {
+                            e.wrap_err("while closing babeld")?;
+                        }
+                        while let Some(res2) = listeners.join_next().await {
+                            res2.wrap_err("while closing babeld")?;
+                        }
+                        State::Idle
+                    }
+                }
+            }
+        };
+    }
+    Ok(())
+}

rust/parts.nix

@@ -1,7 +1,7 @@
 { inputs, ... }:
 {
   perSystem =
-    { inputs', pkgs, lib, ... }:
+    { inputs', self', pkgs, lib, ... }:
     let
       # Fenix nightly toolchain with all components
       rustToolchain = inputs'.fenix.packages.stable.withComponents [
@@ -79,7 +79,7 @@
       };
 
       config = {
-        packages = {
+        packages = rec {
           # Python bindings wheel via maturin
           exo_pyo3_bindings = craneLib.buildPackage (
             commonArgs
@@ -110,6 +110,19 @@
               '';
             }
           );
+          babblerd-unwrapped = craneLib.buildPackage (
+            commonArgs // {
+              inherit cargoArtifacts;
+              pname = "babblerd-unwrapped";
+            }
+          );
+          babblerd = pkgs.writeShellApplication {
+            name = "babblerd";
+            runtimeInputs = [ self'.packages.babeld ];
+            text = ''
+              exec ${babblerd-unwrapped}/bin/babblerd "$@"
+            '';
+          };
         };
 
         checks = {