Compare commits

..
Author SHA1 Message Date
ciaranbor 91a9d0e10e Use sliding window 2026-06-02 12:09:28 -07:00
ciaranbor a2dfc57d50 Always take the most recent snapshot 2026-06-01 11:08:05 -07:00
Andrei Cravtov 629c55d6ba Rename exo_pyo3_bindings to exo_rs (#2131)
## Motivation

(I think it) Makes Evan's massive PR easier to merge later on

## Changes

- Renamed exo_pyo3_bindings to exo_rs
- Upgraded versions of pyo3-based dependencies
- Renamed PyFromSwarm to just FromSwarm, and PyNetworkingHandle to just
NetworkingHandle
2026-05-31 19:23:41 +01:00
Andrei Cravtov f9f8cbb3c3 fix: make app builds work again (#2127)
## Motivation

They didn't

## Changes

They now do

## Why It Works

I changed an env flag, and added a keyword
2026-05-29 18:37:47 +01:00
ciaranbor 051a64e3b4 Capture energy in prefill and ageneration separately (#2124)
## Motivation

Energy was reported as a single aggregate. Split into prefill vs.
generation so each phase can be analysed independently.

## Changes

- `PowerSampler`: `mark_prefill_done()` + `trapezoidal_energy_range()`
helper; `result()` now emits per-phase splits.
- `PowerUsage` / `NodePowerStats`: optional `prefill_*` / `generation_*`
fields (back-compat: `None` if unmarked).
- API marks the boundary on the first non-`PrefillProgressChunk`.
- `bench/exo_bench.py` surfaces the split in the log line and persists
`power_usage` to JSON.
- METHODOLOGY: one sentence + one bullet.

## Why It Works

First non-prefill chunk *is* the boundary. Anchoring a sample there and
interpolating power at the boundary makes phase energies sum exactly to
the unsplit total.

## Test Plan

### Manual Testing

`eco`-reserved nodes:
- M3 Ultra, Qwen3-VL-4B, pp=8192/tg=1024: server 1940 J vs client 1931 J
(+0.5 %)
- M4 Pro, Qwen3.6-27B, pp=16384/tg=2048: server 20,292 J vs client
20,221 J (+0.35 %)

### Automated Testing

5 new tests in `test_power_sampler.py` (range integrator,
splits-sum-to-total, `None`-when-unmarked, idempotency). 14/14 pass.
2026-05-28 14:42:36 -07:00
Andrei Cravtov a8602ea6d5 fix(bug): no longer repeated _trigger_notify_user_to_download_model (#2114)
## Motivation

Partially fixes [this](https://github.com/exo-explore/exo/issues/2098)
issue. Removed erroneous logic for telling user to download when they
already downloaded.

Could not figure out about the "spontaneous crashes" in that issue,
author should consolidate more logs and open a new issue dedicated to
that. I believe
[this](https://github.com/exo-explore/exo/commit/74e9fe15e62fe189dc7e019db86e75c83eca2721)
commit solved some EventRouter-related crashes, which was mentioned in
[this](https://github.com/exo-explore/exo/issues/2098) issue, so it may
have already been solved. If not, should be re-submitted as a new issue.

## Changes

- Consolidated _resolve_and_validate_text_model and
_validate_image_model into one function: _validate_model_has_instance;
- + They already had virtually identical logic, it being different seems
to be an artifact of history
- + Added logic to ensure that _trigger_notify_user_to_download_model is
only called when no such model is downloaded, not just if there is no
instance of it
- Added a new `/instance/await` SSE streaming endpoint to wait for when
a model has an instance available. Complements instance-placement API,
so we can wait till that is done without client-side polling.
- Updated docs and a /tmp script to reflect some of the changes
- Updated dashboard `getModelForRequest` to only return model ID if an
instance exists for it, and updated bits to use `handleChatSend` instead
of `sendMessage` because that checks for if a model instance exists
first.

## Why It Works

The problem was that there was erroneous logging for model not
downloaded. I fixed that logic. The rest is extra.
2026-05-26 14:42:39 +01:00
Andrei Cravtov a1a22b5f38 feat: added background/daemon support (#2106)
## Motivation

Addresses [this](https://github.com/exo-explore/exo/issues/1931) issue.

## Changes

You can now launch Exo as a legacy SysV-style daemin (in the background)
with `--legacy-daemon` flag.
NOTE: don't use it if you're managing Exo with systemd or launchd

SIDE FIX: the macmon process not found trace is no longer displayed on
process shutdown via ctrl+c, that error is supressed.

## Why It Works

Because I used a daemonization library and tweaked it not to break
multiprocessing.

## Test Plan

I ran it in daemon mode, non daemon mode, etc., and pid locking +
inference + everything else works just fine.

Also ran it `ssh user@host -t 'cd exo && nohup nix run .#exo --
--legacy-daemon'` on a 4-node TB mac-mini cluster and the mDNS didn't
die
2026-05-25 20:42:47 +01:00
Andrei Cravtov 74e9fe15e6 fix(bug): EventRouter lifetime-handling fixed, no more process crashes (#2102)
## Motivation

Trying to (partially) fix
[this](https://github.com/exo-explore/exo/issues/2101) issue.

## Changes

Changed channels (in channels.py) to support exception overriding.

Made EventRouter channels throw a subclass of the resource closed/broken
errors.

The current lifetime logic of EventRouter in event loop no longer blows
up because components that use channels from EventRouter now catch the
subclass exceptions in the run method: Worker, Master,
DownloadCoordinator, RunnerSupervisor.

Added logic to throw when API server exits without being asked to shut
down - this kill the sleep-forever in the task-group.
2026-05-22 14:20:04 +01:00
96 changed files with 2203 additions and 14955 deletions

No files matched your search

-7
View File
@@ -1,8 +1 @@
use flake
# creates .venv if doesn't exist and loads its environment
export VIRTUAL_ENV=".venv"
if ! [ -d "./$VIRTUAL_ENV" ]; then
uv venv
fi
layout python
+2
View File
@@ -38,6 +38,8 @@ bench/**/*.json
# tmp
tmp/models
/build/exo
/.agents
/.claude/skills
/.claude
/.codex
skills-lock.json
Generated
+156 -1072
View File
File diff suppressed because it is too large. Load diff
+2 -6
View File
@@ -1,11 +1,6 @@
[workspace]
resolver = "3"
members = [
"rust/networking",
"rust/exo_pyo3_bindings",
"rust/util",
"rust/babblerd",
]
members = ["rust/networking", "rust/exo_rs", "rust/util"]
[workspace.package]
version = "0.0.1"
@@ -33,6 +28,7 @@ delegate = "0.13"
# Utility dependencies
keccak-const = "0.2"
nix = "0.31"
# Async dependencies
async-stream = "0.3"
+18
View File
@@ -201,6 +201,12 @@ This starts the exo dashboard and API at http://localhost:52415/
uv run exo --no-worker
```
- `--legacy-daemon`: Run exo as a legacy SysV-style background daemon using double-fork daemonization. This is intended for legacy init scripts; systemd and launchd should run exo in the foreground without this flag.
```bash
uv run exo --legacy-daemon
```
**File Locations (Linux):**
exo follows the [XDG Base Directory Specification](https://specifications.freedesktop.org/basedir-spec/basedir-spec-latest.html) on Linux:
@@ -395,6 +401,18 @@ Sample response:
}
```
This command is asynchronous. Before sending inference requests, wait until the
API sees the new instance for this model:
```bash
curl -N "http://localhost:52415/instance/await?model_id=mlx-community/Llama-3.2-1B-Instruct-4bit"
```
The endpoint returns an SSE stream. A successful wait emits a message with
`"type": "ready"` and the matching instance; a timeout emits `"type": "timeout"`.
By default it waits indefinitely. Set `timeout_seconds` to a positive value to
bound the wait.
---
**3. Send a chat completion**
@@ -37,7 +37,7 @@ final class ClusterStateService: ObservableObject {
/// gain nothing from being cached on disk. Use an ephemeral session
/// with `urlCache = nil` so neither response bodies nor metadata
/// touch disk.
private static func makeNonCachingSession() -> URLSession {
nonisolated private static func makeNonCachingSession() -> URLSession {
let config = URLSessionConfiguration.ephemeral
config.urlCache = nil
config.requestCachePolicy = .reloadIgnoringLocalCacheData
+2 -1
View File
@@ -125,7 +125,7 @@ A background thread polls each node at 1 Hz, collecting:
- System power draw (W)
- CPU cluster usage (performance and efficiency cores)
**Energy** is computed via trapezoidal integration of the power samples over each inference window (the wall-clock span of each benchmark request or concurrent batch). Average power is `total_joules / total_inference_seconds`.
**Energy** is computed via trapezoidal integration of the power samples over each inference window (the wall-clock span of each benchmark request or concurrent batch). Average power is `total_joules / total_inference_seconds`. The server additionally returns a `power_usage` block in each non-stream `/bench/chat/completions` response that splits energy into prefill and generation phases, with the boundary anchored to the first non-`PrefillProgressChunk` from the runner.
---
@@ -136,6 +136,7 @@ Results are written as JSON with three top-level keys:
- **`runs`**: Array of per-request result objects, each containing:
- `elapsed_s`, `output_text_preview` (first 200 chars)
- `stats`: `{ prompt_tps, generation_tps, prompt_tokens, generation_tokens, peak_memory_usage }`
- `power_usage`: server-side total + prefill/generation split, per-node breakdown (non-stream requests only)
- Placement metadata: `model_id`, `placement_sharding`, `placement_instance_meta`, `placement_nodes`
- Run metadata: `pp_tokens`, `tg`, `repeat_index`, `concurrency`, `concurrent_index`
- `download_duration_s` (if model was freshly downloaded)
+26
View File
@@ -295,6 +295,7 @@ def run_one_completion(
elapsed = time.perf_counter() - t0
stats = out.get("generation_stats")
power_usage = out.get("power_usage")
choices = out.get("choices") or [{}]
message = choices[0].get("message", {}) if choices else {}
content = message.get("content") or ""
@@ -330,6 +331,7 @@ def run_one_completion(
elapsed = time.perf_counter() - t0
preview = "".join(text_parts)[:200]
power_usage = None
if not stats:
ttft = (first_token_time - t0) if first_token_time else elapsed
@@ -348,6 +350,7 @@ def run_one_completion(
"elapsed_s": elapsed,
"output_text_preview": preview,
"stats": stats,
"power_usage": power_usage,
}, pp_tokens
@@ -764,6 +767,7 @@ def main() -> int:
out = c.post_bench_chat_completions(_payload)
elapsed = time.perf_counter() - t0
stats = out.get("generation_stats")
power_usage = out.get("power_usage")
choices = out.get("choices") or [{}]
message = (
choices[0].get("message", {}) if choices else {}
@@ -773,6 +777,7 @@ def main() -> int:
"elapsed_s": elapsed,
"output_text_preview": text[:200],
"stats": stats,
"power_usage": power_usage,
}, _actual_pp
inf_t0 = time.monotonic()
@@ -868,6 +873,27 @@ def main() -> int:
inf_seconds = sum(t1 - t0 for t0, t1 in inference_windows)
avg_watts = joules / inf_seconds if inf_seconds > 0 else 0
summary += f" energy={joules:.1f}J ({avg_watts:.1f}W avg over {inf_seconds:.1f}s inference)"
# mean() not sum() across concurrent runs: each
# request's PowerSampler observes the same shared
# cluster state, so they all report the same figure.
prefill_energies = [
(x.get("power_usage") or {}).get("prefill_energy_joules")
for x in runs
]
gen_energies = [
(x.get("power_usage") or {}).get("generation_energy_joules")
for x in runs
]
prefill_vals = [e for e in prefill_energies if e is not None]
gen_vals = [e for e in gen_energies if e is not None]
if prefill_vals and gen_vals:
avg_pref = mean(prefill_vals)
avg_gen = mean(gen_vals)
summary += (
f" prefill_energy={avg_pref:.1f}J "
f"gen_energy={avg_gen:.1f}J"
)
logger.info(f"{summary}\n")
time.sleep(2)
finally:
+11 -5
View File
@@ -2253,10 +2253,9 @@ class AppStore {
* @returns The model ID to use, or null if none available
*/
private getModelForRequest(modelId?: string): string | null {
if (modelId) return modelId;
if (this.selectedChatModel) return this.selectedChatModel;
const requestedModelId = modelId || this.selectedChatModel;
// Try to get model from first running instance
// Only models with a placed instance can receive requests; disk downloads alone are not enough.
for (const [, instanceWrapper] of Object.entries(this.instances)) {
if (instanceWrapper && typeof instanceWrapper === "object") {
const keys = Object.keys(instanceWrapper as Record<string, unknown>);
@@ -2264,8 +2263,15 @@ class AppStore {
const instance = (instanceWrapper as Record<string, unknown>)[
keys[0]
] as { shardAssignments?: { modelId?: string } };
if (instance?.shardAssignments?.modelId) {
return instance.shardAssignments.modelId;
const instanceModelId = instance?.shardAssignments?.modelId;
// ensure to only return requestedModelId that matches an instance
// or fall back to first instance
if (
instanceModelId &&
(!requestedModelId || requestedModelId === instanceModelId)
) {
return instanceModelId;
}
}
}
+18 -8
View File
@@ -1461,6 +1461,9 @@
addToast({ type: "info", message: `Launching model...` });
// Always auto-select the newly launched model so the user chats to what they just launched
setSelectedChatModel(modelId);
userForcedIdle = false;
pendingChatModelId = modelId;
chatLaunchState = "launching";
// Record the launch in recent models history
recordRecentLaunch(modelId);
@@ -2547,12 +2550,10 @@
];
// ── Seamless chat: launch models from chat view ──
type ChatLaunchState =
| "idle"
| "launching"
| "downloading"
| "loading"
| "ready";
type InFlightChatLaunchState = "launching" | "downloading" | "loading";
type ReadyLikeChatLaunchState = "idle" | "ready";
type ChatLaunchState = InFlightChatLaunchState | ReadyLikeChatLaunchState;
let chatLaunchState = $state<ChatLaunchState>("idle");
let pendingChatModelId = $state<string | null>(null);
let selectedChatCategory = $state<string | null>(null);
@@ -3129,6 +3130,15 @@
if (model) {
pendingAutoMessage = { content, files };
userForcedIdle = false;
// The selected model is already being placed or loaded; keep the queued
// message and let the existing launch state effects send it once ready.
if (
pendingChatModelId === model &&
chatLaunchState !== "idle" &&
chatLaunchState !== "ready"
) {
return;
}
launchModelForChat(model, "picker", messages().length > 0);
return;
}
@@ -4603,7 +4613,7 @@
type="button"
onclick={() => {
completeOnboarding();
sendMessage(chip, undefined, thinkingEnabled());
handleChatSend(chip);
}}
class="px-4 py-2 rounded-full border border-white/10 bg-white/5 text-sm text-white/60 hover:bg-white/10 hover:text-white/80 hover:border-white/20 transition-all duration-200 cursor-pointer"
>
@@ -6100,7 +6110,7 @@
onclick={() => {
chatLaunchState = "idle";
selectedChatCategory = null;
sendMessage(prompt, undefined, thinkingEnabled());
handleChatSend(prompt);
}}
class="text-left px-3 py-2.5 text-xs text-exo-light-gray hover:text-white font-mono rounded-lg border border-exo-medium-gray/30 hover:border-exo-yellow/30 bg-exo-dark-gray/30 hover:bg-exo-dark-gray/60 transition-all duration-200 cursor-pointer"
>
+35 -6
View File
@@ -66,7 +66,9 @@ Creates a new model instance in the cluster.
```
**Response:**
JSON description of the created instance.
Command acknowledgement. Instance creation is asynchronous; clients should wait
for the model to appear through `/instance/await` before sending inference
requests for that model.
### Delete Instance
@@ -94,6 +96,31 @@ Returns details of a specific instance.
**Response:**
JSON description of the instance.
### Await Instance
**GET** `/instance/await?model_id=...&timeout_seconds=0`
Waits until API state contains an instance for the requested model. The response
is an SSE stream so clients receive keep-alive comments while waiting.
**Query parameters:**
* `model_id`: string, required
* `timeout_seconds`: float, optional, default `0`. `0` waits indefinitely;
positive values time out after that many seconds. Maximum positive value:
`300`.
**Stream messages:**
```text
data: {"type": "ready", "instance": {...}}
data: {"type": "timeout", "message": "No instance found for model ..."}
```
The HTTP status is `200` for both messages because the stream starts before the
final result is known. The `type` field disambiguates the terminal message.
### Preview Placements
**GET** `/instance/previews?model_id=...`
@@ -123,17 +150,18 @@ Computes a placement for a potential instance without creating it.
**Response:**
JSON object describing the proposed placement / instance configuration.
### Place Instance (Dry Operation)
### Place Instance
**POST** `/place_instance`
Performs a placement operation for an instance (planning step), without necessarily creating it.
Places an instance for a model using the server's placement logic.
**Request body:**
JSON describing the instance to be placed.
**Response:**
Placement result.
Command acknowledgement. The instance may not be ready immediately; wait for it
to appear through `/instance/await` before sending inference requests.
## 3. Models
@@ -639,10 +667,11 @@ GET /events
# Instance Management
POST /instance
GET /instance/{instance_id}
DELETE /instance/{instance_id}
GET /instance/await
GET /instance/previews
GET /instance/placement
GET /instance/{instance_id}
DELETE /instance/{instance_id}
POST /place_instance
# Models
Generated
+24 -24
View File
@@ -2,11 +2,11 @@
"nodes": {
"crane": {
"locked": {
"lastModified": 1779130139,
"narHash": "sha256-BLrtr42azquO7MdGFU5a7KiMl3YpFlTeIXqy1fT5GlQ=",
"lastModified": 1775790182,
"narHash": "sha256-pG2RWVQY0Pe+rmmXJx+Jpyi+JcgjWzS18m7fcD1B64Q=",
"owner": "ipetkov",
"repo": "crane",
"rev": "edb38893982a3338972bb4a2ec7ce7c29ba10fd9",
"rev": "534982f1c41834b101e381b07b1121a4f065a374",
"type": "github"
},
"original": {
@@ -47,11 +47,11 @@
"rust-analyzer-src": "rust-analyzer-src"
},
"locked": {
"lastModified": 1779185128,
"narHash": "sha256-Kl2bkmwZJD3n2KWDxuIlturZ7emqRK+anpD1LmDwpmY=",
"lastModified": 1775807984,
"narHash": "sha256-Redoe3D9zGN5I9QPHWL9vfMVQBehY1fKsMiRXQ83X3w=",
"owner": "nix-community",
"repo": "fenix",
"rev": "b7bd9323fe26a3b4f4bddbb2c2a1dacabced2f88",
"rev": "fcf90c0c4d368b2ca917a7afa6d08e98a397e5fd",
"type": "github"
},
"original": {
@@ -83,11 +83,11 @@
]
},
"locked": {
"lastModified": 1778716662,
"narHash": "sha256-m1Yf0wZ8j1OHjTc2UwHwyQRSnNeSgLJOd7q5Y45hzi4=",
"lastModified": 1775087534,
"narHash": "sha256-91qqW8lhL7TLwgQWijoGBbiD4t7/q75KTi8NxjVmSmA=",
"owner": "hercules-ci",
"repo": "flake-parts",
"rev": "f7c1a2d347e4c52d5fb8d10cb4d94b5884e546fb",
"rev": "3107b77cd68437b9a76194f0f7f9c55f2329ca5b",
"type": "github"
},
"original": {
@@ -118,11 +118,11 @@
},
"nixpkgs": {
"locked": {
"lastModified": 1779102034,
"narHash": "sha256-vZJZjLo513IeI8hjzHFc6TDezUd4uCE2Eq4SNO3DNNg=",
"lastModified": 1775595990,
"narHash": "sha256-OEf7YqhF9IjJFYZJyuhAypgU+VsRB5lD4DuiMws5Ltc=",
"owner": "NixOS",
"repo": "nixpkgs",
"rev": "687f05a9184cad4eaf905c48b63649e3a86f5433",
"rev": "4e92bbcdb030f3b4782be4751dc08e6b6cb6ccf2",
"type": "github"
},
"original": {
@@ -168,11 +168,11 @@
]
},
"locked": {
"lastModified": 1776659114,
"narHash": "sha256-qapCOQmR++yZSY43dzrp3wCrkOTLpod+ONtJWBk6iKU=",
"lastModified": 1773870109,
"narHash": "sha256-ZoTdqZP03DcdoyxvpFHCAek4bkPUTUPUF3oCCgc3dP4=",
"owner": "pyproject-nix",
"repo": "build-system-pkgs",
"rev": "ffaa2161dd5d63e0e94591f86b54fc239660fb2e",
"rev": "b6e74f433b02fa4b8a7965ee24680f4867e2926f",
"type": "github"
},
"original": {
@@ -188,11 +188,11 @@
]
},
"locked": {
"lastModified": 1778901413,
"narHash": "sha256-GSKXTAnFqRAMlZkJrIPcQMYf+lpMr66K3i60mB9STvc=",
"lastModified": 1775439158,
"narHash": "sha256-NHY9SJNU019n+8NCabBDtmuzRFeE2gZlYKHowp9bV24=",
"owner": "pyproject-nix",
"repo": "pyproject.nix",
"rev": "a228447c3e179d477c1b6246ef3efa8cfe3c469a",
"rev": "fb6b728260f3f32761367e9fd1e1a25b4245bcd0",
"type": "github"
},
"original": {
@@ -218,11 +218,11 @@
"rust-analyzer-src": {
"flake": false,
"locked": {
"lastModified": 1779074864,
"narHash": "sha256-0M3WqsWmtXmv9Ev/vnFfCHosWvISDwiuuhQ104UO3CI=",
"lastModified": 1775745684,
"narHash": "sha256-8MbfLwd60FNa8dRFkjE+G3TT/x21G3Rsplm1bMBQUtU=",
"owner": "rust-lang",
"repo": "rust-analyzer",
"rev": "cdfe408d4b436e806ff525cb3e67588a6a009ed1",
"rev": "64ddb549bc9a70d011328746fa46a8883f937b6b",
"type": "github"
},
"original": {
@@ -284,11 +284,11 @@
]
},
"locked": {
"lastModified": 1778664018,
"narHash": "sha256-ogNyNANNLo0SMFevIeUpbTMOL9uUDu/hXvp7JlOYbwQ=",
"lastModified": 1775706324,
"narHash": "sha256-BTb4sydzX2B5/oNbvCdQFeSbk97xEnbb8bk84CiKCOs=",
"owner": "pyproject-nix",
"repo": "uv2nix",
"rev": "b48abe99ef639cd100c224898529370e5d935294",
"rev": "5707df99097375896a3dda811d492a2fabe63500",
"type": "github"
},
"original": {
+2 -12
View File
@@ -82,7 +82,7 @@
overlays = [
inputs.nixglhost.overlays.default
(import ./nix/apple-sdk-overlay.nix)
(final: prev: {
(final: _: {
macmon = final.rustPlatform.buildRustPackage {
pname = "macmon";
version = "git";
@@ -94,15 +94,6 @@
};
cargoHash = "sha256-Epj3L+db1flGNK5y6yfSig8piEiXTz15lPo/FNkqlkA=";
};
babeld = final.callPackage ./nix/babeld.nix { };
iperf3 = prev.iperf3.overrideAttrs (_old: {
version = "3.21+local";
src = final.fetchgit {
url = "https://github.com/AndreiCravtov/iperf.git";
fetchSubmodules = true;
sha256 = "sha256-2laL7DrEVZxC7sVieaRXBACzMri7YOpu/yO7sC+t3aI=";
};
});
})
];
};
@@ -119,7 +110,7 @@
nixpkgs-fmt.enable = true;
ruff-format = {
enable = true;
excludes = [ "rust/exo_pyo3_bindings/exo_pyo3_bindings.pyi" ];
excludes = [ "rust/exo_rs/exo_rs.pyi" ];
};
rustfmt = {
enable = true;
@@ -141,7 +132,6 @@
packages = {
default = self'.packages.exo;
iperf3 = pkgs.iperf3;
} //
lib.optionalAttrs pkgs.stdenv.hostPlatform.isDarwin {
metal-toolchain = pkgs.callPackage ./nix/metal-toolchain.nix { };
+2 -2
View File
@@ -23,7 +23,7 @@ sync-clean:
rust-rebuild:
PYO3_PYTHON="$(uv run python -c 'import sys; print(sys.executable)')" cargo run --bin stub_gen
uv sync --reinstall-package exo_pyo3_bindings
uv sync --reinstall-package exo_rs
build-dashboard:
#!/usr/bin/env bash
@@ -37,7 +37,7 @@ package: build-dashboard
rm -rf build
build-app: rust-rebuild sync-clean package
xcodebuild build -project app/EXO/EXO.xcodeproj -scheme EXO -configuration Debug -derivedDataPath app/EXO/build
env -u LD xcodebuild build -project app/EXO/EXO.xcodeproj -scheme EXO -configuration Debug -derivedDataPath app/EXO/build
@echo "\nBuild complete. Run with:\n open {{justfile_directory()}}/app/EXO/build/Build/Products/Debug/EXO.app"
clean:
-26
View File
@@ -1,26 +0,0 @@
{ stdenv
, lib
, fetchgit
}:
stdenv.mkDerivation {
pname = "babeld";
version = "1.13.1+local";
# TODO: pin to specific version/revision, or better yet, use a patch file
src = fetchgit {
url = "https://github.com/AndreiCravtov/babeld.git";
fetchSubmodules = true;
sha256 = "sha256-Z4fZNh9ZdWRaTrUxgbXZnDCqvG6m4F/CND3ApyavbLw=";
};
outputs = [
"out"
"man"
];
makeFlags = [
"PREFIX=${placeholder "out"}"
"ETCDIR=${placeholder "out"}/etc"
];
}
+6
View File
@@ -0,0 +1,6 @@
{
"name": "exo",
"lockfileVersion": 3,
"requires": true,
"packages": {}
}
+5 -5
View File
@@ -15,7 +15,7 @@ dependencies = [
"huggingface-hub>=1.8.0",
"psutil>=7.0.0",
"loguru>=0.7.3",
"exo-pyo3-bindings", # rust bindings
"exo-rs", # rust bindings
"anyio==4.11.0",
"tiktoken>=0.12.0", # required for kimi k2 tokenizer
"hypercorn>=0.18.0",
@@ -26,6 +26,7 @@ dependencies = [
"msgspec>=0.19.0",
"zstandard>=0.23.0",
"transformers>=5.6.2",
"python-daemon>=3.1.2",
]
[project.scripts]
@@ -75,15 +76,14 @@ mlx-cuda13 = [
###
[tool.uv.workspace]
members = ["rust/exo_pyo3_bindings", "bench", "tools"]
members = ["rust/exo_rs", "bench", "tools"]
[tool.uv.sources]
exo-pyo3-bindings = { workspace = true }
exo-rs = { workspace = true }
mlx = [
{ git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git", branch = "address-rdma-gpu-locks", marker = "sys_platform == 'darwin'" },
{ url = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv/releases/download/mlx_cuda/mlx-0.32.0-cp313-cp313-manylinux_2_35_aarch64.whl", marker = "sys_platform == 'linux' and platform_machine == 'aarch64'" },
{ url = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv/releases/download/mlx_cuda/mlx-0.32.0-cp313-cp313-manylinux_2_35_x86_64.whl", marker = "sys_platform == 'linux' and platform_machine != 'aarch64'" },
]
mlx-lm = { git = "https://github.com/rltakashige/mlx-lm", branch = "leo/deepseek-v4" }
mflux = { git = "https://github.com/evanev7/mflux", branch = "exo2" }
@@ -240,7 +240,7 @@ torchaudio = ["torch"]
###
[tool.ruff]
extend-exclude = [".typings/**", "rust/exo_pyo3_bindings/**", "bench/vendor/**"]
extend-exclude = [".typings/**", "rust/exo_rs/**", "bench/vendor/**"]
[tool.ruff.lint]
extend-select = ["I", "N", "B", "A", "PIE", "SIM"]
+9 -8
View File
@@ -44,20 +44,21 @@ let
paths = builtins.concatMap (p: [ (lib.getBin p) (lib.getLib p) (lib.getDev p) ]) (cudaLibs ++ [ cudaPackages.cuda_nvcc cuda_cccl_compat ]);
};
exoOverlay = final: prev: {
# Replace workspace exo_pyo3_bindings with Nix-built wheel.
# Replace workspace exo_rs with Nix-built wheel.
# Preserve passthru so mkVirtualEnv can resolve dependency groups.
# Copy .pyi stub + py.typed marker so basedpyright can find the types.
exo-pyo3-bindings = pkgs.stdenv.mkDerivation {
pname = "exo-pyo3-bindings";
exo-rs = pkgs.stdenv.mkDerivation {
pname = "exo-rs";
version = "0.1.0";
src = self'.packages.exo_pyo3_bindings;
src = self'.packages.exo-rs;
# Install from pre-built wheel
nativeBuildInputs = [ final.pyprojectWheelHook ];
dontStrip = true;
passthru = prev.exo-pyo3-bindings.passthru or { };
passthru = prev.exo-rs.passthru or { };
postInstall = ''
local siteDir=$out/${final.python.sitePackages}/exo_pyo3_bindings
cp ${inputs.self}/rust/exo_pyo3_bindings/exo_pyo3_bindings.pyi $siteDir/
local siteDir=$out/${final.python.sitePackages}/exo_rs
cp ${inputs.self}/rust/exo_rs/exo_rs.pyi $siteDir/
touch $siteDir/py.typed
'';
};
@@ -223,7 +224,7 @@ let
++ lib.optionals isDarwin [ pkgs.macmon ];
passthru = {
venv = venv name;
evenv = ((pythonSet.overrideScope editableOverlay).mkVirtualEnv "${name}-evenv" (members // { exo = (members.exo or [ ]) ++ [ "dev" ]; })).overrideAttrs (_: {
evenv = ((pythonSet.overrideScope editableOverlay).mkVirtualEnv "${name}-evenv" (members // { exo = (members.exo or [ ]) ++ [ "dev" ]; exo-rs = [ ]; })).overrideAttrs (_: {
venvSkip = [ "lib/python${python.pythonVersion}/site-packages/mlx/share/cmake/*" "lib/python${python.pythonVersion}/site-packages/build_backend.py" ];
});
} // lib.optionalAttrs cudaSupport {
-1
View File
@@ -1 +0,0 @@
/nix/store/41rd4g2p69a2106qscl3xvkzmh6j7nry-babblerd
-1
View File
@@ -1 +0,0 @@
/pbprobe/*
-48
View File
@@ -1,48 +0,0 @@
[package]
name = "babblerd"
version.workspace = true
edition.workspace = true
[dependencies]
color-eyre = "0.6.5"
clap = { version = "4.5.53", features = ["derive"] }
futures-lite.workspace = true
ipnet = "2.12.0"
nix = { version = "0.31", features = [
"fs",
"signal",
"process",
"user",
"net",
"uio",
] }
netdev = "0.42"
ahash = "0.8.12"
arrayvec = "0.7.6"
crossbeam-channel = "0.5.15"
hashbrown = "0.16.0"
iroh-quinn-udp = { version = "0.8.0", default-features = false, features = [
"fast-apple-datapath",
] }
libc = "0.2"
mio = { version = "1.1.0", features = ["net", "os-ext", "os-poll"] }
n0-watcher = "0.6"
netwatch = "0.16"
rand = "0.10"
route_manager = "0.2.11"
slab = "0.4.11"
socket2 = "0.6.1"
tokio = { workspace = true, features = ["full"] }
tracing = "0.1.44"
tracing-subscriber = { version = "0.3.22", features = ["env-filter"] }
tun-rs = "2.8.1" # if you update, it causes transitive dependeny clashes that need patch.crates-io fixes or whatnot, too long to do now :)
# parsing
memchr = "2.8"
winnow = "1.0"
thiserror = "2.0"
macaddr = "1.0"
zerocopy = { version = "0.8.31", features = ["derive"] }
[lints]
workspace = true
-4
View File
@@ -1,4 +0,0 @@
fn main() -> color_eyre::Result<()> {
color_eyre::install()?;
babblerd::profiling::standalone::run_from_env()
}
-4
View File
@@ -1,4 +0,0 @@
fn main() -> color_eyre::Result<()> {
color_eyre::install()?;
babblerd::profiling::pbprobe::standalone::run_from_env()
}
@@ -1,586 +0,0 @@
# `babblerd` Future Architectural Directions
This file is not a debt list.
Use [shortcuts.md](./shortcuts.md) for concrete shortcuts, footguns, and
implementation compromises that should be cleaned up later.
This file is for directional reasoning:
- what the current architecture is trying to become,
- which major steps are worth doing next,
- and why those steps are ordered the way they are.
It should evolve as the architecture evolves.
## Current Position
`babblerd` is no longer just a thin wrapper around `babeld`.
It now has the beginnings of a real daemon architecture:
- a resident daemon process,
- a resident TUN interface,
- a keepalive-driven daemon core,
- a heavy routing stack that can turn on and off,
- a typed Babel control/runtime layer,
- a derived FIB layer,
- a dedicated dataplane thread wired into the routing stack,
- a persisted node identity,
- and a central config module.
That is enough structure to stop treating the whole project as “just Babel
plumbing”.
For bring-up, the current tree also contains a temporary internal self-client
that connects to the public socket and periodically sends keepalive commands.
That is only a testing scaffold so the routing stack stays on without a real
frontend process yet. It should be removed once a real controller exists.
It is also enough structure to improve the actual dataplane without needing to
perfect every IPC and control-plane detail first.
Stable forwarding is now proven on the four-Mac Thunderbolt lab:
- adjacent overlay traffic works,
- single-hop forwarded UDP works at modest rates,
- steady-state ICMPv6 reachability is green across the full ring,
- small generic TCP application payloads work after convergence,
- and the temporary `enN` link-cost policy successfully keeps steady-state
routes away from `en0`/`en1` and toward the intended lower-numbered direct
Thunderbolt-style links.
The current remaining gap is no longer basic dataplane correctness or forcing
Babel onto the intended route. The route-selection heuristic is now good enough
to expose the next bottleneck: raw dataplane throughput and overload behavior.
The latest lab numbers put the userspace overlay far below the direct physical
UDP baseline:
- direct UDP without the software router: about `11 Gbit/s` observed outside
the overlay,
- direct overlay UDP, `e4 -> e16`, `iperf3 -6 -u -b 0 -t 10`: about
`1.46 Gbit/s` received with negligible loss,
- direct overlay TCP, `e4 -> e16`, `iperf3 -6 -b 0 -t 10`: about
`1.24 Gbit/s` received,
- single-hop overlay UDP, `e2 -> e16`, `iperf3 -6 -u -b 0 -t 10`: server
intervals around `1.11-1.16 Gbit/s` with `12-14%` loss in one run; a later
run sent about `1.32 Gbit/s` and dataplane counters showed about `1.16M`
packets delivered, but the `iperf3` control connection broke before a valid
receiver summary was produced and the overlay path needed a `babblerd`
restart to recover,
- single-hop overlay TCP, `e2 -> e16`, `iperf3 -6 -b 0 -t 10`: about
`1.07 Gbit/s` received.
So the architecture is good enough for continued correctness bring-up, but
serious performance work should now treat packet processing cost, syscalls,
copies, batching, and backpressure as the main suspects.
## The Most Important Architectural Decision
The current codebase is already good enough to serve as the shell around the
first real dataplane.
That means the next major effort should **not** automatically be:
- replacing the line protocol with `zbus`,
- perfecting lease ownership,
- or fully polishing lifecycle semantics.
Those are still desirable, but they are not the blocking step for getting to a
working end-to-end system.
The next big milestone should be:
- make the existing UDP dataplane observable enough to explain overload,
- make the forwarding path robust when the sender exceeds what the router can
currently drain,
- and then reduce per-packet cost enough to move beyond the current
`1-1.5 Gbit/s` overlay ceiling.
In other words: the project has moved from “build the router” to “make the
router fast and predictable under load”.
## Near-Term Goal
The near-term target is now:
> a measurable end-to-end router where:
>
> - the daemon has a stable node identity,
> - the daemon can be kept alive by the frontend,
> - the daemon maintains Babel-derived routing state,
> - the daemon forwards through the UDP overlay between nodes,
> - the route heuristic selects the intended fast links after convergence,
> - and overload is visible through counters rather than guessed from `iperf3`
> alone.
This still does **not** require the final IPC architecture first.
## Recommended Next Phase
### 1. Harden and measure the UDP dataplane
This is the current major feature.
The basic pieces are now in place:
- `fib.rs` derives immutable forwarding snapshots from `BabelState`,
- those snapshots now carry the admitted interface set as well as routes,
- `dataplane.rs` provides a dedicated-thread hot-path module using `mio`,
`socket2`, `crossbeam-channel`, `hashbrown`, `ahash`, `slab`, and
`arrayvec`.
- `routing_stack.rs` now starts the dataplane and publishes coalesced
`FibSnapshot` updates into it.
- dataplane socket ownership is now driven by interfaces that currently have
live Babel neighbours rather than only by currently selected routes, and
retained sockets are refreshed when an `ifname` resolves to a new ifindex.
- socket reconcile is now best-effort under interface churn: transient
resolution/open failures are logged and retried without killing the
dataplane during reconcile.
- unchanged FIB snapshots are now deduplicated in the control plane, so the
dataplane also carries its own lightweight timer-driven reconcile retry for
admitted interfaces that still do not have usable sockets.
- the stable-link packet path is now working on the lab ring for adjacent
one-hop traffic.
- the dataplane now drains ready TUN and UDP fds up to fairness budgets instead
of handling only one packet per readiness event.
- the dataplane logs useful packet/drop counters every few seconds when active:
TUN RX/TX, UDP RX/TX, TUN-to-UDP, forwarded, local-delivered, no-route,
invalid, hop-limit, and UDP/TUN `WouldBlock` drops.
- the UDP receive path now mutates the stack buffer slice directly instead of
allocating a `Vec` per received packet.
- the dataplane compiles each `FibSnapshot` into a local fast route table that
stores direct socket slots, avoiding the old per-packet `FibEntry` clone and
route-ifname-to-socket lookup.
So the next step is no longer “invent or wire the modules”.
It is:
- explain the current `1-1.5 Gbit/s` ceiling against packet counters and host
CPU/syscall behavior,
- make full-blast UDP overload recover cleanly instead of destabilizing the
path,
- compare direct, single-hop, and multi-hop runs with the same counter set,
- and then fill the first obvious protocol gaps such as ICMPv6 error handling.
One caveat that is now proven on the lab Macs: the current macOS receive path
cannot treat "which UDP socket got the packet" as trustworthy interface
attribution. In live tests, packets sent directly over one Thunderbolt
interface are still being received by a different reuseport socket while the
peer scope-id reflects the real physical ingress interface. That means the
current multi-socket receive model is acceptable for basic forwarding bring-up,
but it is not yet a reliable source of receive-side interface truth on macOS.
The likely long-term fix is to move receive-side interface attribution onto
ancillary packet metadata (`IPV6_PKTINFO` / receive-interface data) rather than
inferring it from which socket woke up.
The forked `babeld` used by the Nix build can now start with no managed
interfaces as long as `babblerd` gives it a read-write local control socket.
That means `babblerd` no longer waits for a first interface before spawning
`babeld`; it starts `babeld` immediately and sends `interface <ifname>` commands
later as the watcher discovers eligible links.
For the current four-Mac Thunderbolt lab, the broad macOS `en*` watcher
heuristic has proven too permissive in practice. The dataplane now corrects
that somewhat by only owning sockets on interfaces that Babel has actually
formed neighbour adjacencies on, but the watcher/bootstrap side is still broad.
The env allowlist should remain an escape hatch, not the default topology
description.
The current broad-admission behavior is acceptable for v1 as long as point to
point and multihop forwarding remain reliable, but it does mean that multiple
wired interfaces can become equally admissible at once.
The desired longer-term policy is:
- admit any interface that Babel can actually form a live neighbour adjacency
on, regardless of naming convention,
- keep that broad admissibility for reachability,
- but rank competing links by measured quality rather than treating all wired
links as equivalent.
The forked `babeld` now has the primitive needed for this: the read-write local
socket accepts `neighbour-cost` commands, and neighbour monitor lines report the
active `external-bias-256` and `external-coef-256` fields. Those values are
fixed-point controls in units of `1/256`: the bias is additive, and the
coefficient multiplies the native base cost before the RTT penalty is added.
The measured scoring system is still post-MVP work. It likely requires:
- computing local link metrics such as latency, loss, and possibly sustainable
throughput without generating excessive probe traffic,
- feeding those metrics into `neighbour-cost`,
- and then validating that Babel path selection consistently prefers the better
direct link when multiple usable adjacencies exist.
For MVP, `babblerd` now applies a deliberately simpler policy: for each live
neighbour on an `enN` interface, it sends `neighbour-cost` with
`coef-256 0`. Most `enN` links get `bias-256 N * 100 * 256`, making Babel
treat the native base cost as an absolute synthetic interface-index cost of
roughly `N * 100`, so lower-numbered links such as `en2`, `en3`, and `en4`
win over high-numbered links such as `en18`. `en0` and `en1` are temporarily
deprioritized with the maximum finite `bias-256` value so shared low-index
networks do not dominate Thunderbolt-style links during throughput smoke tests.
This is not a robust scoring model; it is a temporary selection heuristic so
raw throughput work can proceed on the intended fast links.
The current sustained-throughput investigation is therefore focused on the
dataplane hot path and overload behavior:
- direct overlay UDP tops out around `1.46 Gbit/s` in the latest test, far
below the `11 Gbit/s` direct physical UDP baseline,
- single-hop overlay UDP can receive around `1.1 Gbit/s` during full-blast
`-b 0` tests; a later run sent about `1.32 Gbit/s` and delivered roughly
`1.16M` packets at the receiver according to dataplane counters, but the
`iperf3` control connection broke before a receiver summary was produced and
the overlay path needed a restart to recover,
- full-blast UDP should be treated as a stress/failure test until the
backpressure story is better,
- route selection is still important during convergence, but after the mesh
settles the `enN` policy is no longer the main explanation for the throughput
gap.
Be precise about "control traffic" during these tests. Babel's own protocol
packets should remain link-local traffic on the direct physical `en*`
interfaces that `babblerd` explicitly gives to `babeld`; the TUN/overlay
interface is not a Babel interface. The overlay does carry traffic addressed to
node ULAs, including `iperf3` payload, the `iperf3` TCP control/session
connection, and `ping6` to peer ULAs. Saturating the overlay can still disturb
Babel indirectly through shared physical NIC queues, socket buffers, and CPU
scheduling, but not because Babel packets are routed through the userspace
overlay.
That distinction matters for the next diagnosis step. Protecting or separating
control traffic may make tests less fragile and may avoid `iperf3` control
connection failures, but it does not by itself close the `7-8x` dataplane
throughput gap. When a full-load run wedges the path, capture raw Babel state,
the derived `BabelState`, the dataplane FIB/socket map, dataplane counters, and
host route state before concluding whether the failure is route churn, overlay
queue exhaustion, or application-control failure.
At small inner MTUs, approaching `11 Gbit/s` is a packet-rate problem. Ignoring
Ethernet/IP/UDP overhead, the per-packet budget is:
```text
packet_budget_seconds = dataplane_packet_bytes * 8 / target_bits_per_second
```
For `11 Gbit/s`:
- `1452` byte packets, the UDP default TUN MTU: about `947 kpps`, or
`1.06 us/packet`.
- `65535` byte packets, the forced-TCP default TUN MTU: about `21.0 kpps`, or
`47.7 us/packet`.
- `1500` byte packets: about `917 kpps`, or `1.09 us/packet`.
- `1200` byte packets: about `1.15 Mpps`, or `873 ns/packet`.
- `9000` byte jumbo packets: about `153 kpps`, or `6.55 us/packet`.
- `64` byte minimum-size packets: about `21.5 Mpps`, or `46.5 ns/packet`.
So for MTU-sized `iperf3` traffic this is not a "few nanoseconds per packet"
target, but it is roughly a one-microsecond total budget per packet. That
budget has to cover all user/kernel crossings, copies, route lookup, hop-limit
mutation on forwarded packets, UDP send/receive, TUN read/write, and scheduler
overhead. The latest direct overlay result of `1.46 Gbit/s` at `1452` bytes
corresponds to roughly `126 kpps`, or about `8 us/packet`, so getting near
`11 Gbit/s` means shrinking per-packet cost by around `7-8x` or reducing packet
rate with larger packets/aggregation.
Likely optimization directions, in priority order:
- keep improving counters and expose them over the public state surface, so
tests can distinguish no-route, UDP send backpressure, TUN reinjection
backpressure, invalid packets, forwarding, and local delivery without log
scraping;
- add recovery/backpressure policy for overload rather than just
drop-on-`WouldBlock`;
- keep expanding OS packet batching where the target platform allows it. The
dataplane now receives through `iroh-quinn-udp`, which maps to
`recvmsg_x`/`sendmsg_x` on Apple fast builds and `recvmmsg` on Linux-like
Unix. Transmit still sends one packet at a time from the forwarding loop, so
real output batching remains future work;
- consider overlay aggregation, where one outer UDP datagram carries several
inner packets, to amortize syscall and UDP/IP overhead;
- explore jumbo MTUs on the Thunderbolt links, because `9000` byte packets
reduce the `11 Gbit/s` packet rate from about `947 kpps` to about `153 kpps`;
- consider multi-core dataplane sharding once single-thread costs are
measured, because one dedicated thread is a likely ceiling for this design;
- treat kernel-bypass or moving more forwarding into the kernel as a separate
architecture track if `11 Gbit/s` at standard MTU is a hard requirement on
macOS.
Before jumbo frames or overlay aggregation, the current per-packet cost audit
points at these near-term bottlenecks:
- A transit packet still implies one UDP receive syscall and one UDP send
syscall in the forwarding process. At `11 Gbit/s` and `1452` byte packets,
that is roughly `947 kpps`, or nearly `1.9M` UDP syscalls/sec on the transit
node before counting TUN work on endpoints. That alone makes a full `7-8x`
improvement unlikely from ordinary Rust-level cleanup.
- The UDP ingress path was still cloning `socket.ifname` for every received
packet just to support logging/error context. Because `Box<str>::clone()`
allocates, that is an avoidable heap allocation per overlay packet.
- UDP ingress previously used `recv_from` even though the peer address is not
needed for forwarding. Decoding the source address is useful for debugging
but should not be mandatory hot-path work.
- TUN and UDP packet buffers were stack-created as zeroed arrays for each
packet. Reusing worker-owned buffers avoids repeated stack initialization and
keeps the packet loop closer to "syscall, parse, lookup, syscall".
- The send path still builds a `SocketAddrV6` and emits one `iroh-quinn-udp`
transmit per packet. A future connected per-neighbour output-socket model
could remove that address construction and let the kernel cache more route
state.
- The dataplane now uses `iroh-quinn-udp` with the Apple fast datapath, which
exposes `sendmsg_x`/`recvmsg_x` batching. This is not a QUIC routing change;
the useful part is Quinn's UDP socket layer. The current patch batches
receive calls and routes transmit through the same abstraction, but still
emits one transmit call per forwarded packet. The next useful version should
group same-peer/same-size packets into a single `Transmit` with
`segment_size` set.
- Any output batching must be opportunistic, not latency-gating. A single ready
packet must still be sent immediately; batching should flush at the end of a
poll/drain slice or when the next packet targets a different peer/size, never
wait for a full batch.
- The tree now also has an opt-in TCP neighbour transport for Mac Thunderbolt
experiments: `BABBLER_ROUTER_TRANSPORT=tcp`, `--router-transport tcp`, or
`--force-tcp`. UDP remains the default. TCP mode opens scoped link-local TCP
streams to next-hop neighbours, frames inner IPv6 packets with a `u16`
big-endian length, batches framed packets in bounded per-peer write buffers,
and flushes partial batches at drain/poll boundaries. This is intended to
test whether macOS Thunderbolt TCP can expose the higher native TCP path while
avoiding one syscall per inner packet.
- macOS TCP mode keeps one wildcard IPv6 listener per daemon, not one
`IPV6_BOUND_IF` listener per admitted interface. The lab showed
per-interface-bound TCP listeners can stall Thunderbolt handshakes in
`SYN_RCVD`; outbound streams are still scoped to the Babel-selected
interface. Inbound streams are accepted only from link-local peers that match
a live Babel neighbour on the accepted interface/scope.
- TCP mode now defaults the TUN MTU to `65535` to cut packet rate through
macOS `utun` on the Thunderbolt fast path; `BABBLER_TUN_MTU=<mtu>` or
`--tun-mtu <mtu>` can be used for smaller/larger sweeps. All forced-TCP peers
in one test mesh must use the same value, because receivers reject TCP frames
larger than their local TUN MTU. UDP mode still defaults to the
physical-MTU-derived `1452`.
- TCP mode uses `256 KiB` TCP read buffers and `256 KiB` opportunistic write
batch targets by default. Sweep write targets dynamically with
`BABBLER_TCP_BATCH_TARGET_BYTES=<bytes>`; useful first values are `512 KiB`,
`1 MiB`, and `2 MiB`.
- TCP socket send/receive buffers default to `4 MiB`. Sweep them with
`BABBLER_TCP_SOCKET_BUFFER_BYTES=<bytes>`; useful first values are `8 MiB`,
`16 MiB`, and `32 MiB`.
- TCP receive reads directly into the frame decoder buffer, avoiding the old
`tcp_read_buf` copy. Stream readiness is reregistered only when write interest
changes, so a busy stream should not do one `kevent` update per queued packet.
- TCP mode changes overload behavior. Instead of UDP drops on send backpressure,
it can accumulate bounded per-stream pending bytes and then drop once that
bound is reached. Its counters must be watched separately:
`tcp_tx_batches`, `tcp_tx_bytes`, `tcp_queued_packets`,
`tcp_written_frames`, `tcp_rx_batches`, `tcp_rx_bytes`, `tcp_rx_frames`,
`tcp_reregisters`, `tcp_blocked_writes`, `tcp_queue_drops`,
`tcp_frame_errors`, `tcp_stream_errors`, and TUN packet-size buckets.
The first low-risk cleanup pass has now landed: UDP ingress no longer clones
`ifname`, packet buffers are worker-owned instead of stack-created for every
packet, and UDP I/O goes through `iroh-quinn-udp` so the OS-specific fast path
is selected by the crate. These are worth doing, but they should be expected to
remove avoidable overhead rather than close the full `7-8x` gap by themselves.
The current tree now owns the kernel route that steers overlay traffic into the
resident TUN interface:
- the local node `/128` address is installed on the TUN device,
- `EXO_ULA_PREFIX -> tunX` is added when the routing stack turns on,
- that prefix route is removed when the routing stack turns off,
- and `babeld` kernel installs remain disabled.
That means local application traffic can now be steered into the overlay once
the UDP dataplane is active.
The current MVP can stay simple:
- one UDP datagram carries exactly one inner IPv6 packet,
- no custom framing,
- no batching,
- no crypto,
- no relays,
- no multiplexed control/data protocol.
The dataplane currently does this basic loop:
- reads packets from TUN,
- classifies local-delivery vs forwarding,
- looks up next-hop information from a derived forwarding view,
- sends encapsulated packets to direct neighbors over UDP,
- receives UDP packets from neighbors,
- decapsulates them,
- either injects them locally into TUN or forwards them onward.
This gives the project a real “V” and “M” to go with the current daemon/control
shell.
The current tree now hardcodes:
- physical link MTU assumption: `1500`
- outer overhead assumption: `40 bytes IPv6 + 8 bytes UDP`
- UDP TUN MTU default: `1452`
- forced-TCP TUN MTU default: `65535`, with `BABBLER_TUN_MTU`/`--tun-mtu`
override
That is acceptable for bring-up, but it is still only a temporary model.
The future direction should be:
- route-aware MTU derivation,
- PMTUD-aware behavior,
- and better support for environments where hop-to-hop links can use jumbo
frames without exposing that complexity to user traffic.
### 2. Keep the derived forwarding table separate from `BabelState`
`BabelState` should remain a mirror of what `babeld` says.
The current code now reflects that direction:
- `BabelState` is still the protocol mirror,
- `FibSnapshot` is the dataplane view.
The derived forwarding layer should keep moving toward a table that:
- is keyed by destination prefix or node address,
- only keeps the routes the dataplane should actually use,
- captures next hop / outgoing interface / any other forwarding metadata,
- and is cheap for the dataplane to consult.
This avoids mixing:
- “what Babel currently knows”
- with
- “what the UDP router should do with packets”.
### 3. Add a stronger public state/readiness model
The current `ServiceState` is useful, but it is only lifecycle state:
- `Off`
- `Starting`
- `On`
- `Stopping`
That is not the same thing as routing readiness.
Now that the dataplane exists, a separate readiness/status view should exist too.
For example, the frontend may want to distinguish:
- daemon is idle,
- daemon is starting,
- Babel is running but no eligible interfaces exist,
- interfaces exist but no neighbors are usable,
- forwarding is nominal,
- forwarding is degraded.
That should be modeled separately from `ServiceState`, not by making
`ServiceState::On` carry too much meaning.
## What Can Wait Until After Throughput Bring-Up
These are still desirable, but they do not need to block the current
throughput/backpressure work:
### `zbus` / D-Bus-style IPC
This is still the likely long-term direction.
But the current line protocol is good enough for:
- `keepalive <ttl_ms>`
- `get-state`
while the dataplane is being measured and tuned.
So `zbus` should remain a planned improvement, not the immediate blocker.
### Per-client leases
The daemon should eventually track leases per client/connection rather than via
a single global keepalive deadline.
That is a real architectural improvement, but it is control-plane polish rather
than dataplane unblocker.
It can happen after the current router path is faster and better characterized.
### Structured diagnostics/debug output
Right now diagnostics are tracing-only.
That is acceptable for development while the dataplane is being characterized.
A configurable debug stream or structured diagnostics feed should be added
later, preferably once the public IPC shape is stabilized.
## The MVP Dataplane Should Stay Intentionally Small
The MVP version should avoid solving every future overlay concern.
It should **not** attempt to solve:
- encryption,
- authentication,
- path quality metrics beyond what Babel already provides,
- relay protocols,
- or multi-transport negotiation.
Batching and packet aggregation are now valid throughput experiments, but they
should be evaluated as dataplane optimizations rather than bundled with
unrelated control-plane redesign.
The current version has proven the simplest useful thing:
- stable node addresses,
- UDP transport between neighbors,
- Babel-driven next-hop selection,
- TUN injection/extraction,
- packet forwarding that actually works end-to-end.
The rest can be improved incrementally.
## Architectural Path After the MVP Dataplane Works
Now that the basic dataplane exists and works, the likely next path is:
1. Expose dataplane counters and route/FIB state through a better public status
surface.
2. Make overload/backpressure behavior recoverable and measurable.
3. Benchmark batching, aggregation, jumbo MTUs, and eventually multi-core
dataplane options.
4. Replace the temporary `enN` link policy with measured link-quality scoring.
5. Replace the ad-hoc control socket with `zbus`.
6. Replace the single keepalive deadline with per-client leases.
7. Tighten interface admission beyond the current broad heuristic.
8. Pin and explicitly invoke the exact forked `babeld`.
9. Harden node-id file mode checks and other local security edges.
10. Revisit diagnostics streaming.
11. Revisit platform abstractions around TUN / transport / forwarding.
That ordering is intentional:
- keep the router measurable while improving throughput,
- then harden and refine the daemon architecture around it.
## Guiding Principle
The project should prefer:
- a coherent working router with a few acknowledged shortcuts
over:
- a beautifully abstract control plane that still does not move packets.
That does **not** mean ignoring architecture.
It means using the current architecture as a platform for the next real
capability, rather than polishing the control shell ahead of the dataplane's
current throughput and overload problems.
-590
View File
@@ -1,590 +0,0 @@
# babblerd Handoff
This is the current handoff for a new session picking up `babblerd` work.
## Repo / Branch / State
- Repo: `/home/royalguard/Desktop/exo-all/networking-related/exo-babbler`
- Branch: `babbler`
- Current handoff tracks the forced-TCP tuning work on branch `babbler`.
- Recent relevant commits:
- `0b7a3ad3` wires dataplane UDP receive/send through `iroh-quinn-udp`
while keeping `mio` readiness; receive-side batching is in, true transmit
batching is not
- `3cbb5758` removes several avoidable hot-path costs and adds dataplane
counter coverage
- `2f038588` deprioritizes `en0` and `en1` with maximum finite neighbour cost
- `b02cf2cb` adds temporary `enN -> N * 100` link scoring
- `5a158a51` adds Babel neighbour-cost parsing/command support
- `5bf8f62f` added iperf3
- `af0b6e17` remove first interface requirement
- `82adc5d9` it builds
- `59747529` no longer need optional build flags
- earlier dataplane bring-up commits remain relevant, but the local repo has
since moved to `networking-related/exo-babbler`
- Do not trust this file for working-tree cleanliness; run `git status --short`.
## Handoff To Next Agent
The committed baseline before the forced-TCP work was `0b7a3ad3`. Do not infer
working-tree cleanliness from this file; run `git status --short`.
What is implemented:
- `src/dataplane.rs` keeps `mio::net::UdpSocket` for readiness polling.
- Each dataplane interface socket also owns an `iroh_quinn_udp::UdpSocketState`.
- UDP receive now uses `UdpSocketState::recv`, so Apple fast builds can use
`recvmsg_x` and Linux-like Unix can use `recvmmsg` through the crate.
- UDP send now goes through `UdpSocketState::try_send`, but still one transmit
call per forwarded packet.
- Receive batching handles `RecvMeta::stride`, so GRO-style buffers containing
multiple datagrams are split back into inner packets.
- Regression test:
`dataplane::tests::udp_batch_recv_returns_single_datagram_without_full_batch`
proves a single datagram returns immediately as a one-packet batch.
- The current working tree adds an opt-in TCP neighbour transport:
`--force-tcp`, `--router-transport tcp`, or `BABBLER_ROUTER_TRANSPORT=tcp`.
UDP remains the default transport.
- TCP mode opens scoped link-local TCP streams to next-hop neighbours, frames
inner IPv6 packets with a `u16` big-endian length, batches framed packets into
bounded per-peer write buffers, and flushes partial batches at drain/poll
boundaries or when the batch reaches the target size.
- On macOS, TCP mode keeps one wildcard IPv6 listener per daemon and separate
no-fd slots for admitted interfaces. Lab testing found per-interface-bound
TCP listeners could leave e4/e16 Thunderbolt handshakes stuck in `SYN_RCVD`;
outbound TCP streams are still scoped to the Babel-selected interface.
Accepted streams are admitted only when the peer is a link-local Babel
neighbour on the accepted scope/interface.
- TCP mode is intended as an experimental Mac Thunderbolt fast path to reduce
one-syscall-per-packet overhead. It is not the default mesh transport.
- TCP mode uses a jumbo `65535` byte TUN MTU by default to reduce userspace TUN
packet rate on the Mac Thunderbolt fast path. Override with `--tun-mtu <mtu>`
or `BABBLER_TUN_MTU=<mtu>` when sweeping Mac `utun` limits. All forced-TCP
peers in a test mesh must use the same TUN MTU; a receiver rejects TCP frames
larger than its local MTU. UDP mode keeps the old `1452` byte default.
- TCP mode uses `256 KiB` TCP read buffers and `256 KiB` opportunistic write
batch targets by default. Override the write target with
`BABBLER_TCP_BATCH_TARGET_BYTES=<bytes>` when sweeping `512 KiB`, `1 MiB`, or
`2 MiB` batches. Partial batches still flush at drain/poll boundaries.
- TCP socket send/receive buffers default to `4 MiB`; override with
`BABBLER_TCP_SOCKET_BUFFER_BYTES=<bytes>` for matrix runs such as `8 MiB`,
`16 MiB`, or `32 MiB`.
- TCP receive now reads directly into the frame decoder buffer, and stream
readiness is reregistered only on `READABLE` / `READABLE|WRITABLE` changes.
What is not implemented:
- No true UDP output batching yet.
- No UDP output queue, aggregation, or waiting-to-fill behavior.
- No connected per-neighbour output sockets yet.
- Full-load remote tests have not yet been recorded for TCP mode in this
handoff.
If implementing actual UDP transmit batching next:
1. Add a worker-owned `TxBatch` scratch buffer.
2. Append only packets with the same output socket, next-hop peer, and packet
length.
3. Flush on peer change, packet-size change, full batch, end of TUN/UDP drain
slice, poll-loop boundary, snapshot/reconcile, stop, or send error.
4. Flush via one `Transmit { contents: batch_bytes, segment_size: Some(packet_len), ... }`.
5. Never wait for a full batch. Batching is syscall amortization inside an
already-ready drain turn, not a latency queue.
6. Add tests for single-packet flush, peer-change flush, size-change flush, and
full-batch flush.
Docs are part of the fix. Any future code change should update this handoff and
the relevant architecture/lab notes in the same patch, especially when it
changes what is implemented versus future work.
## Core Conclusion
The project is past the “is the overlay architecture wrong?” phase.
The current architecture is the right one:
- `babeld` is control plane only
- Babel kernel installs are disabled
- local mesh traffic is steered into a resident TUN
- userspace dataplane forwards one inner IPv6 packet per UDP datagram hop-by-hop over neighbour link-locals
So the main remaining work is now:
- dataplane throughput and overload behavior
- making full-blast UDP recover cleanly
- restart/convergence robustness around transient route choices
- eventually replacing the temporary `enN` link scoring with measured link
quality
## Why The Old Approach Failed
The original macOS idea was effectively:
- let `babeld` install routes
- try to make kernel source selection behave
That did not work cleanly for this use case:
- no usable IPv6 pref-src route install path on macOS/BSD for this design
- no native source-specific IPv6 routing model that solves the app behavior wanted here
- putting ULAs on `lo0` or `utun` did not reliably fix source selection
- app-aware binding alone was not enough in practice
That is why the design pivoted to the userspace overlay dataplane.
## Files To Read First
- `src/daemon.rs`
- `src/routing_stack.rs`
- `src/dataplane.rs`
- `src/fib.rs`
- `src/tun.rs`
- `src/babel/runtime.rs`
- `src/route_ctl.rs`
- `lab_topology_reference.md`
- `shortcuts.md`
- `future_architectural_directions.md`
## Current Intended Architecture
Model:
- one stable node `/128` on TUN
- `EXO_ULA_PREFIX -> tunX` installed by `babblerd`
- `babeld` kernel installs are disabled / ignored
- `BabelState` mirrors `babeld`
- `FibSnapshot` is a reduced immutable dataplane view
- control plane stays on Tokio
- dataplane is a dedicated thread
- one UDP datagram carries exactly one inner IPv6 packet
- no custom framing yet
- outer IPv6 destination is neighbour link-local
- outer UDP port is `router_udp_port`
This is the v1 forwarding model:
- exact-match `/128` host routes only
- interface identity in FIB is `ifname`
- dataplane owns sockets from admitted interface set
- admitted dataplane interfaces come from live Babel neighbours
## Important Design Decisions Already Landed
- typed Babel parsing/state model, not raw string handling
- monitor-driven Babel runtime, not periodic dump polling
- persistent node identity across restarts
- explicit daemon lifecycle: `Off | Starting | On | Stopping`
- resident TUN lifetime, separate heavy routing stack
- dataplane thread + immutable FIB snapshot swaps
- socket ownership from admitted interface set, not only current route set
- same-name/new-ifindex socket refresh handled
- timer-driven socket reconcile retry in dataplane, so deduped unchanged FIB snapshots do not suppress retries forever
- dataplane exit supervision back into routing stack / daemon
- macOS dataplane uses `tun-rs` packet I/O (`SyncDevice::recv/send`), not raw fd reads/writes
## Important Forked `babeld` Changes
The Nix build now uses a forked `babeld` from
`/home/royalguard/Desktop/exo-all/networking-related/babeld`, packaged as
`1.13.1+local`.
Recent fork behavior that matters to `babblerd`:
- `babeld` can start with no managed interfaces when a read-write local control
socket exists.
- `babblerd` now spawns `babeld` immediately and adds interfaces later with
local-socket `interface <ifname>` commands from the watcher.
- `kernel-install false` is still used so `babeld` performs route selection and
reports installed routes without touching kernel routes.
- `neighbour-cost <ifname> <link-local-neighbour> bias-256 <bias> coef-256 <coef>`
is available for external link-cost steering.
- `bias-256` is a signed fixed-point additive value in units of `1/256`.
`256` adds one Babel cost unit and `-256` subtracts one.
- `coef-256` is an unsigned fixed-point multiplier in units of `1/256`.
`256` is neutral, `128` halves the native base cost, and `0` ignores the
native base cost while still adding the RTT penalty.
- `dump`/`monitor` neighbour lines now include `external-bias-256` and
`external-coef-256` before `cost`.
Automatic measured link scoring is not part of the MVP. The current temporary
policy is a simple Mac heuristic: for most `enN` links, set an absolute
synthetic base cost around `N * 100` with `coef-256 0` and
`bias-256 N * 100 * 256`, so lower-numbered Thunderbolt-style interfaces are
preferred over high-numbered interfaces such as `en18`. `en0` and `en1` are
temporarily assigned the maximum finite `bias-256` value (`16776704`, yielding
cost `65534` with `coef-256 0`) so shared low-index networks do not dominate
the direct-link smoke tests. This is intentionally a temporary selection aid so
raw throughput work can assume the good direct links are chosen after
convergence. Immediately after a restart, Babel may still transiently install a
bad high-cost route until better neighbour state arrives.
## Very Important Fix After Earlier Handovers
The previously-deployed node addresses were wrong.
There was a real bug in `EXO_ULA_PREFIX` construction:
- intended prefix: `fde0:20c6:1fa7:ffff::/64`
- broken runtime prefix had become: `20c6:1fa7:ffff:0::/64`
Cause:
- `config.rs` used a `u128` left-shift construction that dropped the high `fde0` bits
Fix:
- commit `b0f508ac` changed the prefix constant to explicit hextets and added a regression test
Live verification after redeploy:
- `e4 utun5`: `fde0:20c6:1fa7:ffff:cc78:aec2:d64e:f125/128`
- `e2 utun5`: `fde0:20c6:1fa7:ffff:aeb:e53a:cb17:aa42/128`
- `e11 utun5`: `fde0:20c6:1fa7:ffff:34a:26dd:46ff:1a3f/128`
- `e16 utun5`: `fde0:20c6:1fa7:ffff:7c5d:5e2d:54df:e665/128`
So any older notes mentioning the truncated non-ULA prefix are stale.
## Current Dataplane Behavior
In `src/dataplane.rs`:
- TUN ingress:
- read inner IPv6 packet
- parse destination
- drop self-directed
- FIB lookup
- send raw inner packet as UDP payload to neighbour
- UDP ingress:
- receive UDP payload
- payload is raw inner IPv6 packet
- if destination local, inject into TUN
- else decrement inner hop limit and forward
Fast-path traits:
- dedicated OS thread
- `mio` polling
- `socket2` UDP sockets
- immutable FIB snapshot swaps over `crossbeam-channel`
- no lock on packet lookup path
- ready TUN/UDP fds are drained up to fairness budgets
- UDP receive uses a stack buffer slice directly, not `to_vec()`
- each `FibSnapshot` is compiled into dataplane-local fast routes that include
direct socket slots, so packets no longer clone `FibEntry` or do an ifname
lookup to find the output socket
- dataplane counters are logged periodically when active: TUN RX/TX, UDP RX/TX,
TUN-to-UDP, forwarded, local-delivered, no-route, invalid, hop-limit,
UDP/TUN `WouldBlock` drops, TCP batch/write/read counters, TCP reregisters,
and TUN packet-size buckets
## What Works
These things are now real:
- one-hop two-node `ping6`
- adjacent dataplane path
- small low-rate UDP matrix
- single-hop forwarded UDP at `100M` with no loss in the latest smoke test
- encapsulation / decapsulation itself
- basic generic TCP correctness after convergence
- steady-state route choice avoiding `en0`/`en1` after the temporary cost policy
has converged
- full-bandwidth direct overlay tests that show a repeatable `1-1.5 Gbit/s`
dataplane ceiling rather than a basic correctness failure
Latest route examples after convergence:
- `e2 -> e11`: direct on `en3`, metric `300`
- `e2 -> e16`: single-hop via `e4` on `en2`, metric `400`
- `e4 -> e16`: direct on `en2`, metric `200`
- `e16 -> e2`: single-hop via `e11` on `en2`, metric `400`
## What Is Still Broken
The main remaining live problem is dataplane throughput and overload behavior.
Observed latest performance shape:
- direct physical UDP without the software router is about `11 Gbit/s`
according to the latest external baseline,
- direct overlay UDP is about `1.46 Gbit/s` received,
- direct overlay TCP is about `1.24 Gbit/s` received,
- single-hop overlay TCP is about `1.07-1.08 Gbit/s` received,
- single-hop overlay UDP at `-b 0` receives around `1.11-1.16 Gbit/s` during
one run and loses `12-14%`; a later run sent about `1.32 Gbit/s` and
delivered about `1.16M` packets according to dataplane counters, but the
`iperf3` control connection broke before a valid receiver summary was
produced and the overlay path needed a `babblerd` restart to recover.
So the current blocker is:
- not “UDP overlay cannot carry packets”,
- not primarily “Babel selected the wrong steady-state route”,
- but packet processing cost, syscall/copy overhead, and drop/recovery behavior
when the dataplane is overdriven.
Restart/convergence route quality is still worth watching. Immediately after a
restart, Babel can transiently install high-cost `en0`/`en1` routes before the
better neighbours converge. But after convergence, the current `enN` policy is
good enough for throughput work.
At `11 Gbit/s` with `1452` byte inner packets, the budget is about `947 kpps`,
or `1.06 us/packet`. The latest direct overlay result at `1.46 Gbit/s` is about
`126 kpps`, or `8 us/packet`. Closing the gap means cutting per-packet cost by
roughly `7-8x`, increasing effective packet size with jumbo/aggregation, or
both.
Near-term cost audit before jumbo/aggregation:
- A transit node needs one UDP receive syscall and one UDP send syscall per
forwarded packet, so standard-MTU `11 Gbit/s` implies nearly `1.9M` UDP
syscalls/sec on that node.
- Rust-level cleanup alone is unlikely to recover a full `7-8x`, but avoidable
hot-path work should still be removed before blaming the architecture.
- First targets now landed: remove the per-packet `socket.ifname` clone on UDP
ingress, avoid source-address decoding when the peer address is not needed,
and reuse packet buffers instead of stack-zeroing a new array per packet.
- Initial `iroh-quinn-udp` wiring has landed. The dataplane still keeps `mio`
sockets for readiness, but each socket also has a `UdpSocketState`; receives
can batch through the crate's Apple `recvmsg_x` path or Linux `recvmmsg`
path, and sends go through the same abstraction. This is not QUIC.
- Next candidates are connected per-neighbour output sockets and true output
batching: collect same-peer/same-size packets and send them as one
`Transmit` with `segment_size` set, instead of one transmit call per forwarded
packet.
- Batching invariant: never wait for a full batch. Receive batching must return
whatever is already queued on the nonblocking fd, and future output batching
must flush partial batches at drain boundaries or peer/size changes.
Control traffic terminology:
- Babel protocol packets should stay on the direct link-local `en*`
interfaces that `babblerd` explicitly adds to `babeld`; the TUN/overlay
interface is not added to Babel.
- `iperf3` data, the `iperf3` TCP control/session connection, and `ping6` to a
peer ULA do traverse the overlay because they are addressed to node ULAs.
- Full overlay load can still perturb Babel indirectly through shared physical
NIC queues, kernel buffers, and CPU scheduling, but Babel packets are not
being encapsulated by the software router in the normal design.
- Therefore "protect control traffic" means making measurements and recovery
less fragile; it is not a direct explanation for the order-of-magnitude
throughput gap.
## Very Important macOS Receive-Side Finding
On macOS, receive-side socket attribution is not trustworthy in the current one-socket-per-interface model.
Observed live behavior:
- traffic sent directly over one Thunderbolt link can be delivered to a different UDP socket than expected
- the peer scope-id still reflects the real ingress interface
Implication:
- do not trust “which socket woke up” as authoritative ingress truth on macOS
- if receive-side interface attribution matters, use peer scope-id and likely ancillary packet metadata later
This is a real quirk, but it is not the primary explanation for the current
order-of-magnitude throughput gap.
## Key Local FIB Caveat
Do not assume route-choice issues are only Babels fault.
In `src/fib.rs`, `FibBuilder` collapses multiple installed host routes by choosing the lowest:
- `metric`
- then `refmetric`
- then `handle`
So if restart churn leaves multiple `installed=yes` candidates, babblerds
derived `FibSnapshot` can still be part of why traffic transiently goes via
`en0`/`en1`.
If route-choice anomalies reappear, compare all three:
1. raw Babel route events / dump
2. current `BabelState`
3. derived `FibSnapshot`
Not Babel in isolation.
## Lab Topology / Operations
Source of truth file:
- `lab_topology_reference.md`
Key facts:
- four Mac minis
- hostnames:
- `e4@e4`
- `e2@e2`
- `e11@e11`
- `e16@e16`
- ring topology:
- `e4 -> e2 -> e11 -> e16 -> e4`
- remote repo path:
- `~/babeld-exo`
- each remote must `git pull` before running
- current start command:
- `cd ~/babeld-exo && git pull && RUST_LOG=info sudo -E nix run .#babblerd --impure`
- temporary internal keepalive client exists, so no external `nc -U ...` client is needed just to keep daemon alive
- `iperf3` is provided by the flake:
- `nix run .#iperf3 -- -s`
- `nix run .#iperf3 -- -c <addr>`
- Force the experimental TCP dataplane transport with either:
- `BABBLER_ROUTER_TRANSPORT=tcp RUST_LOG=info sudo -E nix run .#babblerd --impure`
- `RUST_LOG=info sudo -E nix run .#babblerd --impure -- --force-tcp`
- In forced TCP mode on macOS, each daemon has one wildcard listener shared by
all admitted interfaces; per-peer outbound streams remain interface-scoped.
Accepted TCP streams are rejected unless the peer is link-local and matches a
live Babel neighbour on the accepted scope/interface.
- Forced TCP defaults the TUN MTU to `65535`. Use `--tun-mtu <mtu>` or
`BABBLER_TUN_MTU=<mtu>` to test smaller values such as `16384` or `32768`. Keep the
value identical on every forced-TCP peer in a run.
- Tune forced-TCP write batching and TCP socket buffers with:
- `BABBLER_TCP_BATCH_TARGET_BYTES=<bytes>`; default `262144`
- `BABBLER_TCP_SOCKET_BUFFER_BYTES=<bytes>`; default `4194304`
- The current `iperf3` source is the fork at
`/home/royalguard/Desktop/exo-all/networking-related/iperf3`.
Commit `962e05b` adds `%scopeID` rendering for link-local IPv6 output.
## Current Docs Status
Read:
- `shortcuts.md`
- `future_architectural_directions.md`
- `lab_topology_reference.md`
They correctly capture:
- broad admissibility is acceptable for v1 reachability
- flat wired costs are not enough for good best-path choice, hence the
temporary `enN` policy
- forked `babeld` now has the `neighbour-cost` primitive needed for temporary
external cost steering
- the route-selection heuristic is now active and good enough after convergence
to expose dataplane throughput limits
- the latest direct/single-hop `iperf3` results and the `11 Gbit/s` direct UDP
baseline
- the remaining debt around backpressure, batching/aggregation, jumbo MTU,
interface identity, macOS receive attribution, IPC/authz, and incomplete
ICMP/PMTUD behavior
## Important Remaining Technical Debt
Still unresolved:
- public IPC socket is too open
- `ServiceState::On` is not the same as “fully ready/routable”
- broad interface admission is still heuristic
- route ownership of `EXO_ULA_PREFIX` is aggressive
- no ICMPv6 Time Exceeded
- no Packet Too Big handling
- no real backpressure/queueing; `WouldBlock` is still drop-on-backpressure
- counters are logged but not yet exposed as a structured public status surface
- direct overlay throughput is still about `1.46 Gbit/s`, far below the
`11 Gbit/s` direct physical UDP baseline
- full-blast single-hop UDP can destabilize the overlay path after the run
- macOS receive-side interface attribution needs a better long-term path
- multi-link path selection still uses a temporary `enN -> N * 100`
absolute-cost heuristic, not measured scoring
- automatic measured link-scoring policy is not implemented yet
## What Not To Revisit Right Now
These are settled enough for now:
- overlay architecture itself
- TUN + userspace UDP forwarding model
- one-packet-per-datagram framing as the MVP correctness model; batching or
aggregation can now be evaluated as a performance extension
- control plane on Tokio, dataplane on dedicated thread
- exact-match `/128` FIB for v1
- `tun-rs` packet I/O on macOS instead of raw fd reads/writes
- disabling Babel kernel installs and owning `EXO_ULA_PREFIX -> tunX` locally
## Best Next Performance Step
Use the current route heuristic and focus on dataplane cost.
The most concrete current experiment is forced TCP transport on the Mac
Thunderbolt lab. It should be compared against the UDP default with the same
routes, same iperf pairs, same dataplane counter deltas, and same recovery
checks. TCP mode batches framed inner packets before kernel writes; UDP mode
still emits one send operation per forwarded packet.
Capture each run with:
1. `iperf3` sender/receiver summaries
2. `babblerd` dataplane counter deltas
3. CPU usage on sender, transit node, and receiver
4. route/FIB snapshots before and after the run
5. whether bidirectional `ping6` still works after the run
6. for TCP mode, `tcp_tx_batches`, `tcp_tx_bytes`, `tcp_queued_packets`,
`tcp_written_frames`, `tcp_rx_batches`, `tcp_rx_bytes`, `tcp_rx_frames`,
`tcp_reregisters`, `tcp_rejected_peers`, `tcp_queue_drops`,
`tcp_frame_errors`, `tcp_stream_errors`, and TUN packet-size bucket deltas
Goal:
- separate CPU/syscall ceiling from UDP/TUN backpressure,
- explain the single-hop UDP wedge and distinguish overlay application-control
failure from Babel route churn,
- and measure whether receive-side batching changed direct and single-hop
throughput before implementing true transmit batching, aggregation, jumbo MTU
support, or multi-core sharding.
The temporary `neighbour-cost` policy now lives in `babel/link_policy.rs`.
For each live neighbour on an `enN` interface, `babblerd` sets `coef-256 0`.
Most `enN` links use `bias-256 N * 100 * 256`, so Babel sees lower-numbered
interfaces as cheaper while keeping the distributed Babel view and dataplane
view aligned. `en0` and `en1` are the temporary exceptions: they get the
largest finite `bias-256` value so they lose to the explicit direct-link
interfaces during smoke tests.
If route-choice bugs reappear, then inspect all three together for the
problematic `/128` pair: raw Babel route events over time, current `BabelState`,
and derived `FibSnapshot`.
## Best Next Live Tests
1. full directed `ping6` matrix on node `/128`s
2. small directed UDP matrix
3. direct overlay TCP/UDP `-b 0` with counter capture
4. single-hop overlay TCP/UDP `-b 0` with counter capture
5. short soak tests on adjacent and two-hop pairs
6. restart/convergence tests
7. physical churn tests
8. for failures, always capture:
- symptom
- raw Babel route state / dump
- current `BabelState`
- derived FIB state if relevant
- dataplane logs
- relevant `ifconfig`
## Short Version
The project is now in the:
- throughput / backpressure robustness
- overload recovery
- restart convergence sanity-checking
phase.
The dataplane is basically real.
The current main question is not “can the overlay forward packets at all?”
It is:
- why the userspace router tops out around `1-1.5 Gbit/s` when direct physical
UDP can reach about `11 Gbit/s`
- and how much of that gap comes from one-packet-per-datagram syscalls/copies,
single-thread processing, TUN/UDP backpressure, or recoverability bugs under
overload
-191
View File
@@ -1,191 +0,0 @@
# Lab Topology Reference
This file records the current, still-relevant lab topology and bring-up
context for `babblerd`.
## Hosts
- The lab consists of four Mac minis.
- SSH targets:
- `e4@e4`
- `e2@e2`
- `e11@e11`
- `e16@e16`
- You can SSH into these machines directly to inspect or run commands.
## Physical Topology
- The machines are connected in a Thunderbolt ring:
- `e4 -> e2 -> e11 -> e16 -> e4`
- The Thunderbolt-facing interface names are not fixed to `en2` and `en3`.
macOS can expose additional Thunderbolt links as other `en*` interfaces such
as `en5`, `en6`, or host-specific names after reconfiguration.
- Treat `en2,en3` as an old bring-up heuristic only. For normal lab testing,
run without `BABBLER_INTERFACE_ALLOWLIST` and let `babblerd`/Babel discover
the live interfaces.
## Repository Location On The Macs
- Each machine has a checkout of the Exo repository at:
- `~/babeld-exo`
- That checkout is expected to already be on the correct branch for this work.
## Current Node Addresses
- `e2`: `fde0:20c6:1fa7:ffff:aeb:e53a:cb17:aa42`
- `e11`: `fde0:20c6:1fa7:ffff:34a:26dd:46ff:1a3f`
- `e16`: `fde0:20c6:1fa7:ffff:7c5d:5e2d:54df:e665`
- `e4`: `fde0:20c6:1fa7:ffff:cc78:aec2:d64e:f125`
## Current Route Policy
`babblerd` now uses the forked `babeld` `neighbour-cost` command to bias route
selection after neighbours appear:
- `en0` and `en1` get maximum finite cost with `bias-256 16776704` and
`coef-256 0`, yielding cost `65534`.
- Most other `enN` interfaces get `bias-256 N * 100 * 256` and `coef-256 0`,
so `en2` costs `200`, `en3` costs `300`, `en5` costs `500`, and so on.
- This is a temporary Mac lab heuristic, not measured link scoring.
- After convergence it keeps steady-state routes away from `en0`/`en1` and
toward lower-numbered direct Thunderbolt-style links. Immediately after
restart, Babel can still transiently install worse high-cost routes until
better neighbour state arrives.
## Running `babblerd`
From `~/babeld-exo`, pull first so the machine is not testing stale commits,
then start `babblerd`:
```sh
cd ~/babeld-exo
git pull
RUST_LOG=info sudo -E nix run .#babblerd --impure
```
The default dataplane transport is UDP. To force the experimental TCP
neighbour transport for Mac Thunderbolt throughput tests, start every node with
one of:
```sh
BABBLER_ROUTER_TRANSPORT=tcp RUST_LOG=info sudo -E nix run .#babblerd --impure
RUST_LOG=info sudo -E nix run .#babblerd --impure -- --force-tcp
```
TCP mode still uses Babel for route selection and still sends Babel packets on
the direct link-local `en*` interfaces. Only node-ULA overlay traffic uses the
TCP streams.
On macOS, forced-TCP listener sockets are wildcard listeners rather than
per-interface-bound listeners. Per-interface `IPV6_BOUND_IF` on TCP listeners
left e4/e16 handshakes stuck in `SYN_RCVD` during lab testing, while a plain
link-local TCP listener on the same cable completed. Outbound TCP streams remain
scoped to the Babel-selected interface. Accepted TCP streams are rejected unless
the peer is link-local and matches a live Babel neighbour on the accepted
scope/interface.
Forced TCP defaults the TUN MTU to `65535` to reduce per-packet TUN syscalls on
the Mac Thunderbolt fast path. Override it during sweeps with either command
form, but keep the value identical on every forced-TCP peer in the run:
```sh
BABBLER_TUN_MTU=16384 BABBLER_ROUTER_TRANSPORT=tcp RUST_LOG=info sudo -E nix run .#babblerd --impure
RUST_LOG=info sudo -E nix run .#babblerd --impure -- --force-tcp --tun-mtu 32768
```
TCP mode uses `256 KiB` TCP read buffers and `256 KiB` opportunistic write batch
targets by default. Sweep write targets and TCP socket buffers dynamically with
the same values on every node:
```sh
BABBLER_TCP_BATCH_TARGET_BYTES=1048576 BABBLER_TCP_SOCKET_BUFFER_BYTES=16777216 BABBLER_ROUTER_TRANSPORT=tcp RUST_LOG=info sudo -E nix run .#babblerd --impure
```
Good first matrix values are `262144`, `524288`, `1048576`, and `2097152` for
`BABBLER_TCP_BATCH_TARGET_BYTES`, plus `4194304`, `8388608`, `16777216`, and
`33554432` for `BABBLER_TCP_SOCKET_BUFFER_BYTES`. TCP receive reads directly
into the frame decoder buffer, and stream readiness is reregistered only when
write interest changes. UDP mode still defaults to `1452`, derived from a
`1500` byte physical MTU minus outer IPv6 and UDP headers.
If broad interface discovery causes unrelated links to interfere with a
specific debug run, `BABBLER_INTERFACE_ALLOWLIST` is still available as a
temporary escape hatch. Do not use it as the default lab topology description.
`babblerd` no longer needs to wait for an initial interface before starting
`babeld`. The forked `babeld` can start with no managed interfaces, and
`babblerd` adds interfaces later through the read-write local control socket.
## `iperf3`
Use the flake-provided forked `iperf3` when testing this branch:
```sh
nix run .#iperf3 -- -s
nix run .#iperf3 -- -c <addr>
```
The current fork lives at
`/home/royalguard/Desktop/exo-all/networking-related/iperf3`; commit `962e05b`
adds `%scopeID` rendering for link-local IPv6 output.
Latest useful test commands:
```sh
nix run .#iperf3 -- -s -1
nix run .#iperf3 -- -6 -b 100M -t 5 -c <node-ula>
nix run .#iperf3 -- -6 -u -b 100M -t 5 -c <node-ula>
nix run .#iperf3 -- -6 -u -b 0 -t 10 -c <node-ula>
nix run .#iperf3 -- -6 -b 0 -t 10 -c <node-ula>
```
For the forced-TCP dataplane experiment, capture both correctness and
performance:
1. Start all four nodes with TCP mode enabled.
2. Wait for convergence and record route/FIB state.
3. Run a directed `ping6` matrix over node ULAs.
4. Run adjacent TCP and UDP iperf smoke tests at `100M`.
5. Run direct overlay TCP/UDP `-b 0` on an adjacent pair such as `e4 -> e16`.
6. Run single-hop overlay TCP/UDP `-b 0` such as `e2 -> e16`.
7. Capture dataplane counter deltas, especially TCP batches/frames/errors, and
verify bidirectional `ping6` still works after each full-rate run.
Latest findings:
- Direct physical UDP without the software router is about `11 Gbit/s`.
- Direct overlay `e4 -> e16`, UDP `-b 0`: about `1.46 Gbit/s` received with
negligible loss.
- Direct overlay `e4 -> e16`, TCP `-b 0`: about `1.24 Gbit/s` received.
- Single-hop overlay `e2 -> e16`, UDP `-b 100M`: `100 Mbit/s` with no loss.
- Single-hop overlay `e2 -> e16`, UDP `-b 0`: earlier server intervals were
around `1.11-1.16 Gbit/s` with `12-14%` loss. A later run sent about
`1.32 Gbit/s`; dataplane counters showed roughly `1.16M` packets delivered
at `e16`, but the `iperf3` control connection broke before a receiver
summary was produced. Treat that as an overload/control-path failure, not as
a zero-throughput receiver result. The overlay path then needed a
`babblerd` restart to recover.
- Single-hop overlay `e2 -> e16`, TCP `-b 0`: about `1.07-1.08 Gbit/s`
received, with route/path recovery sometimes lagging briefly after the run.
For `1452` byte packets, `11 Gbit/s` is roughly a one-microsecond packet
budget: about `947 kpps`, or `1.06 us/packet`. The current direct overlay
result is roughly `126 kpps`, or `8 us/packet`. Forced TCP now defaults the TUN
MTU to `65535`, so compare packet counters before attributing any result to the
outer TCP socket alone.
Babel protocol packets should not traverse the software router. `babblerd`
starts `babeld` without startup interfaces and later adds only eligible
physical `en*` interfaces through the local control socket; the TUN/overlay
interface is not added to Babel. What does traverse the overlay is traffic
addressed to node ULAs, including `iperf3` data, the `iperf3` TCP control
connection, and `ping6` to a peer ULA. Full overlay load can still indirectly
perturb Babel by consuming shared NIC queues, kernel buffers, and CPU time on
the same physical interfaces, but it is not because Babel's link-local packets
are being encapsulated by `babblerd`.
## Important Current Note
- With the current codebase, `babblerd` has an internal dummy keepalive client.
- That means you do **not** need to connect an external client socket just to
make the daemon stay active during testing.
@@ -1,51 +0,0 @@
# PBProbe Implementation Plan
This file tracks the staged implementation and validation of a paper-faithful
PBProbe profiler for link-local lab links.
## Stage 1: Local Implementation
- Add `src/profiling/pbprobe/` as a separate module from the simple packet
train profiler.
- Implement the paper protocol:
- START initiates one direction.
- RTS requests each sample.
- the sender replies with a packet bulk of length `k`, meaning `k + 1`
packets.
- the receiver measures first and last packet arrival time, delay sum, and
dispersion.
- END reports the selected sample and estimate.
- Implement Algorithm 1:
- start with `k = 1`.
- if measured minimum dispersion is below `D_thresh`, multiply `k` by 10 and
restart.
- otherwise pace samples with `G = 2D / U`.
- stop after fixed `n` accepted samples.
- Keep the C implementation as a reference, but use the paper's units for `G`.
## Stage 2: Local Verification
- Unit-test packet encoding/decoding.
- Unit-test estimator selection by minimum delay sum.
- Unit-test bulk-length adaptation and pacing calculations.
- Compile the standalone example.
## Stage 3: Lab Validation
- Discover the current link-local addresses and interface names on the Mac mini
ring via SSH.
- Build or run the PBProbe example on the relevant remotes.
- Run the flake-provided forked `iperf3` over the same link-local scoped
addresses as the baseline. The fork at
`/home/royalguard/Desktop/exo-all/networking-related/iperf3` includes commit
`962e05b`, which renders `%scopeID` for link-local IPv6 output.
- Compare PBProbe estimates against `iperf3` with a reasonable tolerance.
- If estimates are outside tolerance, adjust only algorithm parameters or
implementation bugs, not the scoring target.
## Current Notes
- The repo license is Apache-2.0. The dropped PBProbe source has a permissive
MIT-like license header with notice retention and academic citation language.
- The C code appears to implement the core estimator, but its `G` sleep units
look inconsistent with the paper. This implementation should follow the paper.
-381
View File
@@ -1,381 +0,0 @@
# `babblerd` Shortcuts
This file tracks architectural and implementation shortcuts that were taken
deliberately during the refactors. They are acceptable for now, but they are
not meant to be the final design.
This is not a dump of every `TODO` comment in the crate. It is the curated list
of shortcuts that should be revisited later.
## Architecture / IPC
- The public control socket still uses an ad-hoc line protocol instead of the
intended `zbus`/D-Bus-style IPC surface.
Files:
- `src/daemon.rs`
- `src/main.rs`
Follow-up:
- Replace `keepalive <ttl_ms>` / `get-state` string commands with a typed IPC
API.
- The daemon core currently tracks a single global keepalive deadline, not
per-client leases.
Files:
- `src/daemon.rs`
- `src/main.rs`
Why this is a shortcut:
- It does not model multiple clients independently.
- It cannot distinguish which client is keeping the service alive.
- The current tree also includes a temporary internal self-client in
`main.rs` that periodically issues keepalive commands just to keep the
daemon/routing stack alive during bring-up.
Follow-up:
- Introduce real lease ownership/tracking in the daemon core.
- Remove the temporary internal keepalive client once a real frontend or test
harness is driving the daemon.
- Raw Babel debug output currently only goes to tracing logs.
Files:
- `src/babel/runtime.rs`
- `src/daemon.rs`
Why this is a shortcut:
- There is no configurable or structured diagnostics stream anymore.
- That is fine for now, but eventually debugging should not require tailing
daemon logs.
Follow-up:
- Add configurable debug output or a separate structured diagnostics stream
once the real IPC surface exists.
- The daemon core exposes state only through `get-state` polling and inline
command responses.
Files:
- `src/daemon.rs`
Follow-up:
- Add real state publication/signals once the IPC surface is upgraded.
## Service Lifecycle
- The daemon now has explicit `Off/Starting/On/Stopping`, but the control model
is still minimal.
Files:
- `src/daemon.rs`
Why this is a shortcut:
- There is no richer lifecycle API yet.
- There is no explicit enable/disable policy beyond keepalive-driven on/off.
Follow-up:
- Revisit the final lifecycle API once IPC is made real.
- `ServiceState::On` currently means “the routing tasks were started”, not a
stronger readiness guarantee such as “babeld is healthy, has admitted
interfaces, and is actually usable for mesh forwarding”.
Files:
- `src/daemon.rs`
- `src/routing_stack.rs`
- `src/babel/runtime.rs`
Why this is a shortcut:
- The frontend may eventually want to distinguish process/task liveness from
actual routing readiness.
Follow-up:
- Add a separate readiness field or richer public state model instead of
overloading `ServiceState::On`.
- The resident TUN vs heavy routing-stack split is now in place, but the
naming and abstractions are still transitional.
Files:
- `src/daemon.rs`
- `src/routing_stack.rs`
- `src/tun.rs`
Follow-up:
- Revisit names and boundaries after the daemon core / IPC architecture settles.
- `RoutingStack::stop` still uses abort-driven shutdown for the interface
watcher and logger task.
Files:
- `src/routing_stack.rs`
Why this is a shortcut:
- It is pragmatic, but not a carefully coordinated shutdown protocol.
Follow-up:
- Replace task abortion with explicit shutdown signaling where it matters.
## Babel Integration
- `babeld` runtime startup config is still assembled partly as raw strings.
Files:
- `src/babel/runtime.rs`
- `src/babel/command.rs`
Why this is a shortcut:
- The local-socket command side is typed, but spawn-time `-C` config is not.
Follow-up:
- Add a typed Babel config/config-statement layer.
- The runtime still depends on fork-specific `babeld` behavior
while spawning `"babeld"` from `PATH`.
Files:
- `src/babel/runtime.rs`
- `../nix/babeld.nix`
Why this is a shortcut:
- The current fork supplies `kernel-install false`, no-interface startup, and
the `neighbour-cost` local-socket command.
- It assumes the right binary is on `PATH`.
- The Nix packaging is still not pinned to a specific revision.
Follow-up:
- Pin the fork revision and make the runtime use that exact binary.
## Networking / Interface Admission
- Interface admission is still heuristic and too broad on macOS.
Files:
- `src/lib.rs` (`if_watcher`)
- `src/config.rs`
- `src/fib.rs`
Why this is a shortcut:
- Any `en*` interface with link-local IPv6 and `is_up()` can still get pulled
into Babel during bootstrap.
- This can include unrelated WiFi, built-in Ethernet, USB Ethernet, etc.
- The dataplane now narrows that broad bootstrap set back down to interfaces
that actually have live Babel neighbours, which is much closer to the real
transport set.
- But the watcher/bootstrap side is still using the coarse `en*` heuristic,
and the env allowlist is still just a bring-up escape hatch rather than the
long-term admission policy.
- When multiple admissible wired links exist, Babel's native wired scoring
still treats them essentially flatly.
- `babblerd` now applies a temporary `enN -> N * 100` absolute-cost policy
with forked `babeld`'s `neighbour-cost` command, except that `en0` and
`en1` are assigned the maximum finite cost so shared low-index networks do
not win simply because their interface indexes are small.
Follow-up:
- Replace the watcher-side bootstrap heuristic with a stronger admission
policy (neighbor proof, richer metadata, or both), so Babel does not need
broad speculative interface admission just to discover the right links.
- Keep validating route choice during restart/convergence, but treat the
temporary macOS `en0`/`en1` deprioritization plus `enN -> N * 100`
neighbour-cost policy as good enough for steady-state throughput work.
- Later, replace that heuristic with measured link-quality scoring so broadly
admissible direct links can be ranked by actual observed quality.
- The dataplane now derives immutable FIB snapshots and runs on a dedicated
thread, but it still assumes interface names are the stable long-lived
identity for socket ownership.
Files:
- `src/fib.rs`
- `src/dataplane.rs`
Why this is a shortcut:
- The dataplane now owns sockets from the admitted interface set rather than
inferring them only from current routes, and it refreshes retained sockets
when a name resolves to a new ifindex.
- The dataplane now also has a lightweight timer-driven reconcile retry for
admitted interfaces whose socket setup was skipped or failed, so unchanged
FIB snapshots no longer suppress retries completely.
- That fixes the earlier route-derived and stale-ifindex bugs, but the design
still assumes interface names are stable enough to be the long-lived
control-plane identity.
Follow-up:
- Revisit whether the long-term identity should be richer than `ifname`,
especially if interface renames/hotplug churn become common during runtime.
- On macOS in particular, the current "one socket per interface" receive
model is not trustworthy enough to identify the real ingress interface:
live testing shows packets sent directly over one Thunderbolt link can be
received on a different `MioUdpSocket` while the peer scope-id still
reflects the actual physical ingress interface.
- Revisit receive-side interface attribution on macOS, likely using ancillary
packet-info / receive-interface metadata instead of assuming the receiving
socket tells the truth.
- The current dataplane is intentionally minimal and still drops several packet
classes silently.
Files:
- `src/dataplane.rs`
- `src/fib.rs`
- `src/routing_stack.rs`
Why this is a shortcut:
- The current lab state has reliable ICMPv6 reachability, basic generic TCP
correctness after convergence, a clean `100M` single-hop UDP smoke test,
and direct/single-hop full-bandwidth `iperf3` measurements.
- Earlier `iperf3` results were confounded by route selection. That is no
longer the main explanation after the temporary `neighbour-cost` policy:
direct overlay UDP still tops out around `1.46 Gbit/s`, while direct
physical UDP without the software router is about `11 Gbit/s`.
- The first round of hot-path cleanup is in: readiness drains have fairness
budgets, UDP receive no longer allocates a `Vec`, `FibSnapshot`s are
compiled into dataplane-local fast routes with socket slots, and dataplane
counters are logged periodically.
- The current code also still treats `WouldBlock` on UDP send and TUN
reinjection as drop-on-backpressure behavior. That is now visible in logs,
but it is not yet a proper queued/backpressured forwarding model.
- Single-hop UDP at `-b 0` can receive around `1.1 Gbit/s` during the run but
with heavy loss and a post-test overlay wedge until `babblerd` is
restarted. A later `-b 0` run sent about `1.32 Gbit/s` and dataplane
counters showed about `1.16M` packets delivered, but the `iperf3` control
connection failed before a valid receiver summary and the path needed a
restart to recover. That points at overload/recovery and test-control
fragility, not basic packet decoding.
- There is no ICMPv6 Time Exceeded generation yet.
- There is no Packet Too Big handling yet.
- No-route and invalid-packet cases are mostly tracing-and-drop behavior.
Follow-up:
- Use dataplane counter deltas, CPU measurements, and route/FIB snapshots
around each `iperf3` run to separate syscall/CPU ceiling from UDP/TUN
backpressure.
- The first avoidable per-packet costs have been removed: no heap allocation
for UDP receive logging, no peer-address decoding when it is not needed,
and no per-packet zeroed buffer construction.
- Initial `iroh-quinn-udp` wiring is in: `mio` still owns readiness, but UDP
receive/send calls now go through Quinn's UDP socket layer, which selects
Apple `recvmsg_x`/`sendmsg_x` or Linux `recvmmsg` internally where
available. True output batching is still future work, and it must flush
partial batches rather than wait for a full batch.
- An opt-in TCP transport now exists for Mac Thunderbolt experiments. It is
selected with `BABBLER_ROUTER_TRANSPORT=tcp`, `--router-transport tcp`, or
`--force-tcp`. It frames inner IPv6 packets onto scoped link-local TCP
streams and batches writes through bounded per-peer buffers. This is a
deliberate experimental shortcut, not the default transport policy.
- On macOS, TCP mode uses one wildcard IPv6 listener per daemon instead of
one per-interface-bound listener per admitted interface; outbound streams
still use the Babel-selected interface. Accepted streams must be link-local
Babel neighbours on the accepted interface/scope. Bound listeners caused
e4/e16 lab handshakes to stick in `SYN_RCVD`.
- TCP mode defaults the TUN MTU to `65535`; use `--tun-mtu <mtu>` or
`BABBLER_TUN_MTU=<mtu>` for smaller/larger lab sweeps. Use the same value on
every forced-TCP peer in a run. UDP keeps the old `1452` default.
- TCP mode uses `256 KiB` TCP read buffers and `256 KiB` opportunistic write
batch targets by default. Sweep write batches with
`BABBLER_TCP_BATCH_TARGET_BYTES=<bytes>` and socket buffers with
`BABBLER_TCP_SOCKET_BUFFER_BYTES=<bytes>`; start with `512 KiB`, `1 MiB`,
`2 MiB` batches and `8 MiB`, `16 MiB`, `32 MiB` socket buffers.
- TCP receive reads directly into the frame decoder buffer, and stream
readiness is reregistered only when write interest changes.
- TCP mode relies on a lab assumption: direct cabled Mac Thunderbolt links are
low-loss enough that TCP-over-TCP pathologies should be limited during
throughput tests. It should not be treated as a general unreliable-mesh
replacement for UDP without more overload and loss testing.
- Benchmark output batching, aggregation, jumbo MTU support, and eventually
multi-core dataplane sharding. At `11 Gbit/s` with `1452` byte packets, the
budget is about `947 kpps`, or `1.06 us/packet`; the current direct overlay
result is roughly `126 kpps`, or `8 us/packet`.
- Do not describe Babel packets as going through the software router in the
normal design. Babel should be direct-interface link-local traffic on the
`en*` interfaces that `babblerd` gives to `babeld`; overlay load can still
disturb it indirectly through shared NIC/kernel/CPU resources.
- Later, replace the suffix heuristic with measured scoring and investigate
how that should interact with restart/convergence behavior.
- Add proper ICMPv6 error generation and tighter packet-validation behavior.
- `TunDevice` is still a thin platform-specific wrapper with some rough edges.
Files:
- `src/tun.rs`
- `src/dataplane.rs`
Why this is a shortcut:
- It still stores the address as `Ipv6Net` even though usage is `/128`-only.
- On macOS, the actual kernel interface name is still `utunN`; the daemon's
cross-platform naming has been cleaned up, but the OS-level interface name
is not under our control there.
- It still has hard-coded MTU and other tun-rs builder assumptions.
- The dataplane still relies on `mio::unix::SourceFd` and `AsRawFd` to poll
the TUN fd on Unix.
- On macOS, packet I/O must still go through `tun-rs`'s `SyncDevice::recv`
and `SyncDevice::send`; bypassing those with raw fd `read`/`write` breaks
utun packet-information handling even if `mio` polling itself is correct.
- That low-level fd borrowing is smaller and safer than the old
`unsafe`/owned-fd handoff, but it still keeps raw-fd details in the
dataplane hot path.
Follow-up:
- Tighten the type and revisit the platform-specific tuning once the dataplane
is implemented.
- Revisit whether a future dataplane/eventing design can remove the direct
`mio`/raw-fd dependency entirely.
- The current MTU model is still intentionally crude:
- assume physical links must support 1500-byte packets,
- use `1452` as the UDP TUN MTU default,
- use `65535` as the forced-TCP TUN MTU default,
- allow `BABBLER_TUN_MTU` or `--tun-mtu` for explicit experiments,
- reject candidate physical interfaces below 1500 MTU.
Files:
- `src/config.rs`
- `src/lib.rs`
- `src/tun.rs`
Why this is a shortcut:
- It does not handle PMTUD, VLAN overhead, per-route MTU variation, or
jumbo-frame opportunities.
Follow-up:
- Replace the current coarse MTU model with route-aware MTU derivation once
the dataplane is better characterized.
- The overlay route controller currently claims the whole overlay prefix
aggressively.
Files:
- `src/route_ctl.rs`
Why this is a shortcut:
- It removes any existing route matching `EXO_ULA_PREFIX` before adding the
daemon's own interface route, and removes all matching routes again on
shutdown.
- That is acceptable only if babblerd is the sole owner of the overlay
prefix.
Follow-up:
- Narrow route deletion so it only removes routes that this daemon installed,
or otherwise encode route ownership more precisely.
## Identity / Security / Filesystem
- The node-id file is created with `0600`, but existing files are only
owner-checked, not mode-checked.
Files:
- `src/identity.rs`
Why this is a shortcut:
- A root-owned but group/world-writable file would still be accepted.
Follow-up:
- Enforce safe permissions on reload, not just on initial creation.
- The public IPC socket is intentionally world-accessible for now.
Files:
- `src/main.rs`
Why this is a shortcut:
- Any local user can connect, issue keepalives, and drive the daemon's public
control surface.
Follow-up:
- Revisit permissions/authz once the IPC surface is finalized.
## Error Modeling
- Several orchestration-layer errors are flattened to `String`/`Arc<str>` too
early.
Files:
- `src/daemon.rs`
- `src/babel/runtime.rs`
- `src/lib.rs` (`BabbleError::Other(String)`)
Why this is a shortcut:
- It loses structure and source-chain information.
Follow-up:
- Prefer typed errors or `eyre::Report` internally, and stringify only at the
IPC/UI boundary.
## Constants / Magic Values
- A few important constants are still effectively magic values:
- EXO ULA prefix details
- default router UDP port
- various timeout/sleep durations in the Babel runtime
- tun MTU
Files:
- `src/config.rs`
- `src/babel/runtime.rs`
- `src/tun.rs`
Follow-up:
- Either justify them clearly as real protocol/runtime constants or move them
into better configuration/abstraction layers.
## Testing
- The typed Babel parser/state layers are tested, but the newer daemon-core and
routing-stack lifecycle behavior is still lightly tested.
Files:
- `src/daemon.rs`
- `src/routing_stack.rs`
- `src/main.rs`
Follow-up:
- Add focused tests for:
- keepalive-driven transitions,
- stack start/stop behavior,
- public socket command behavior,
- failure propagation from the routing stack.
-244
View File
@@ -1,244 +0,0 @@
//! Typed representation of commands sent to `babeld`'s local socket.
//!
//! This is the outbound counterpart to [`crate::babel::line`]:
//!
//! - [`crate::babel::line`] models what `babeld` emits
//! - this module models the runtime control lines that `babblerd` sends
//!
//! The scope here is intentionally narrow: this module only models the local-socket
//! commands that `babblerd` currently issues at runtime.
//!
//! NOTE: spawn-time `-C` configuration strings are still assembled in the runtime layer for now.
//! If you want to push the protocol model further, the next obvious extraction is a typed
//! configuration/config-statement layer rather than more runtime socket commands.
use std::fmt;
use std::net::Ipv6Addr;
pub const BABEL_INFINITY: i32 = 65_535;
pub const NEIGHBOUR_COST_BIAS_256_MIN: i32 = -((BABEL_INFINITY - 1) * 256);
pub const NEIGHBOUR_COST_BIAS_256_MAX: i32 = (BABEL_INFINITY - 1) * 256;
pub const NEIGHBOUR_COST_COEF_256_MIN: u32 = 0;
pub const NEIGHBOUR_COST_COEF_256_MAX: u32 = 65_535;
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum BabelCommand {
Dump,
Monitor,
Unmonitor,
Quit,
Interface(Box<str>),
NeighbourCost(NeighbourCostCommand),
}
impl BabelCommand {
/// Encode this command for the local `babeld` socket, including line framing.
#[must_use]
pub fn encode(&self) -> String {
format!("{self}\n")
}
}
/// Signed fixed-point additive neighbour-cost bias in units of 1/256.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct NeighbourCostBias256(i32);
impl NeighbourCostBias256 {
#[must_use]
pub fn new(value: i32) -> Option<Self> {
(NEIGHBOUR_COST_BIAS_256_MIN..=NEIGHBOUR_COST_BIAS_256_MAX)
.contains(&value)
.then_some(Self(value))
}
#[must_use]
pub const fn neutral() -> Self {
Self(0)
}
#[must_use]
pub const fn raw(self) -> i32 {
self.0
}
}
/// Unsigned fixed-point neighbour-cost multiplier in units of 1/256.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct NeighbourCostCoef256(u32);
impl NeighbourCostCoef256 {
#[must_use]
pub fn new(value: u32) -> Option<Self> {
(NEIGHBOUR_COST_COEF_256_MIN..=NEIGHBOUR_COST_COEF_256_MAX)
.contains(&value)
.then_some(Self(value))
}
#[must_use]
pub const fn neutral() -> Self {
Self(256)
}
#[must_use]
pub const fn raw(self) -> u32 {
self.0
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct NeighbourCostCommand {
ifname: Box<str>,
link_local_neighbour: Ipv6Addr,
bias_256: NeighbourCostBias256,
coef_256: NeighbourCostCoef256,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum NeighbourCostCommandError {
NonLinkLocalNeighbour,
}
impl NeighbourCostCommand {
pub fn new(
ifname: impl Into<Box<str>>,
link_local_neighbour: Ipv6Addr,
bias_256: NeighbourCostBias256,
coef_256: NeighbourCostCoef256,
) -> Result<Self, NeighbourCostCommandError> {
if !link_local_neighbour.is_unicast_link_local() {
return Err(NeighbourCostCommandError::NonLinkLocalNeighbour);
}
Ok(Self {
ifname: ifname.into(),
link_local_neighbour,
bias_256,
coef_256,
})
}
pub fn neutral(
ifname: impl Into<Box<str>>,
link_local_neighbour: Ipv6Addr,
) -> Result<Self, NeighbourCostCommandError> {
Self::new(
ifname,
link_local_neighbour,
NeighbourCostBias256::neutral(),
NeighbourCostCoef256::neutral(),
)
}
}
impl fmt::Display for BabelCommand {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Self::Dump => f.write_str("dump"),
Self::Monitor => f.write_str("monitor"),
Self::Unmonitor => f.write_str("unmonitor"),
Self::Quit => f.write_str("quit"),
Self::Interface(ifname) => write!(f, "interface {ifname}"),
Self::NeighbourCost(cmd) => write!(f, "{cmd}"),
}
}
}
impl fmt::Display for NeighbourCostCommand {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(
f,
"neighbour-cost {} {} bias-256 {} coef-256 {}",
self.ifname,
self.link_local_neighbour,
self.bias_256.raw(),
self.coef_256.raw()
)
}
}
#[cfg(test)]
mod tests {
use std::net::Ipv6Addr;
use super::{
BabelCommand, NeighbourCostBias256, NeighbourCostCoef256, NeighbourCostCommand,
NeighbourCostCommandError,
};
#[test]
fn renders_commands() {
assert_eq!(BabelCommand::Dump.to_string(), "dump");
assert_eq!(BabelCommand::Monitor.to_string(), "monitor");
assert_eq!(BabelCommand::Unmonitor.to_string(), "unmonitor");
assert_eq!(BabelCommand::Quit.to_string(), "quit");
assert_eq!(
BabelCommand::Interface("en2".into()).to_string(),
"interface en2"
);
assert_eq!(
BabelCommand::NeighbourCost(
NeighbourCostCommand::new(
"en18",
"fe80::42".parse().unwrap(),
NeighbourCostBias256::new(4_096).unwrap(),
NeighbourCostCoef256::new(128).unwrap(),
)
.unwrap()
)
.to_string(),
"neighbour-cost en18 fe80::42 bias-256 4096 coef-256 128"
);
}
#[test]
fn encodes_commands() {
assert_eq!(BabelCommand::Dump.encode(), "dump\n");
assert_eq!(BabelCommand::Monitor.encode(), "monitor\n");
assert_eq!(
BabelCommand::Interface("en2".into()).encode(),
"interface en2\n"
);
assert_eq!(
BabelCommand::NeighbourCost(
NeighbourCostCommand::neutral("en2", "fe80::1".parse().unwrap()).unwrap()
)
.encode(),
"neighbour-cost en2 fe80::1 bias-256 0 coef-256 256\n"
);
}
#[test]
fn validates_neighbour_cost_fixed_point_ranges() {
assert_eq!(
NeighbourCostBias256::new(super::NEIGHBOUR_COST_BIAS_256_MIN)
.unwrap()
.raw(),
-16_776_704
);
assert_eq!(
NeighbourCostBias256::new(super::NEIGHBOUR_COST_BIAS_256_MAX)
.unwrap()
.raw(),
16_776_704
);
assert!(NeighbourCostBias256::new(super::NEIGHBOUR_COST_BIAS_256_MIN - 1).is_none());
assert!(NeighbourCostBias256::new(super::NEIGHBOUR_COST_BIAS_256_MAX + 1).is_none());
assert_eq!(NeighbourCostCoef256::new(0).unwrap().raw(), 0);
assert_eq!(
NeighbourCostCoef256::new(super::NEIGHBOUR_COST_COEF_256_MAX)
.unwrap()
.raw(),
65_535
);
assert!(NeighbourCostCoef256::new(super::NEIGHBOUR_COST_COEF_256_MAX + 1).is_none());
}
#[test]
fn rejects_non_link_local_neighbour_cost_address() {
assert_eq!(
NeighbourCostCommand::neutral("en2", Ipv6Addr::LOCALHOST),
Err(NeighbourCostCommandError::NonLinkLocalNeighbour)
);
}
}
-785
View File
@@ -1,785 +0,0 @@
//! Typed representation of lines emitted by `babeld`'s local socket.
//!
//! This module models the inbound side of the Babel local control protocol:
//!
//! - [`BabelLine`] is one parsed wire line.
//! - [`HeaderLine`] covers the connection prelude.
//! - [`Status`] covers command completion lines such as `ok`, `bad`, and `no ...`.
//! - [`Event`] and its associated structs cover the asynchronous routing/interface updates
//! emitted by `dump` and `monitor`.
//!
//! The sibling parser lives in [`parse`]. Its job is to turn raw socket lines into these domain
//! types. Higher layers such as the Babel runtime/state code should depend on this module's
//! types, and keep raw strings only at the actual socket boundary.
//!
//! More concretely:
//!
//! - use [`parse::parse_line`] when reading from `babeld`
//! - reduce [`Event`] values into [`crate::babel::state::BabelState`]
//! - treat [`Status`] as command acknowledgements
//! - keep outbound socket/config commands in a separate module rather than mixing them into
//! this inbound line model
use crate::babel::Eui64;
use ipnet::IpNet;
use std::net::{IpAddr, Ipv4Addr};
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum BabelLine {
Header(HeaderLine),
Status(Status),
Event(Event),
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum HeaderLine {
Banner { major: u8, minor: u8 },
Version(Box<str>),
Host(Box<str>),
MyId(Eui64),
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum Status {
Ok,
Bad,
No(Option<Box<str>>),
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum Event {
Interface(InterfaceEvent),
Neighbour(NeighbourEvent),
XRoute(XRouteEvent),
Route(RouteEvent),
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum EventKind {
Add,
Change,
Flush,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct InterfaceEvent {
pub kind: EventKind,
pub ifname: Box<str>,
pub up: bool,
pub ipv6: Option<IpAddr>,
pub ipv4: Option<Ipv4Addr>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct NeighbourEvent {
pub kind: EventKind,
pub handle: u64,
pub address: IpAddr,
pub ifname: Box<str>,
pub reach: u16,
pub ureach: u16,
pub rxcost: u32,
pub txcost: u32,
pub rtt_millis: Option<u32>,
pub rttcost: Option<u32>,
pub external_bias_256: i32,
pub external_coef_256: u32,
pub cost: u32,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct XRouteEvent {
pub kind: EventKind,
pub prefix: IpNet,
pub from: IpNet,
pub metric: u32,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct RouteEvent {
pub kind: EventKind,
pub handle: u64,
pub prefix: IpNet,
pub from: IpNet,
pub installed: bool,
pub id: Eui64,
pub metric: u32,
pub refmetric: u32,
pub via: IpAddr,
pub ifname: Box<str>,
}
/// Parser for `babeld`'s local socket output.
///
/// This submodule is the wire-format counterpart to the parent [`crate::babel::line`] domain
/// types. It turns raw socket text into [`BabelLine`] values.
///
/// The local socket protocol implemented in `networking-related/babeld/local.c` is line-oriented
/// ASCII. The parser is split into two layers:
///
/// - [`RawLines`] does zero-copy line framing over buffered bytes with [`memchr`].
/// - [`parse_line`] parses one complete line with [`winnow`].
/// - [`ParsedLines`] is a convenience adapter for buffered transcripts such as `dump` output.
///
/// `monitor` mode uses the exact same line grammar as `dump`; it simply keeps emitting event lines
/// after the initial snapshot.
///
/// The accepted grammar is:
///
/// ```text
/// stream ::= (line "\n")* line?
/// line ::= header | status | event
///
/// header ::= banner | version | host | my-id
/// banner ::= "BABEL " uint "." uint
/// version ::= "version " text
/// host ::= "host " text
/// my-id ::= "my-id " eui64
///
/// status ::= "ok" | "bad" | ("no" (" " text)?)
///
/// event ::= kind " " (interface | neighbour | xroute | route)
/// kind ::= "add" | "change" | "flush"
///
/// interface ::= "interface " ifname " up " bool
/// (" ipv6 " ip)?
/// (" ipv4 " ipv4)?
///
/// neighbour ::= "neighbour " hex " address " ip " if " ifname
/// " reach " hex " ureach " hex
/// " rxcost " uint " txcost " uint
/// (" rtt " millis " rttcost " uint)?
/// (" external-bias-256 " int " external-coef-256 " uint)?
/// " cost " uint
///
/// xroute ::= "xroute " prefix "-" prefix
/// " prefix " prefix " from " prefix " metric " uint
///
/// route ::= "route " hex
/// " prefix " prefix " from " prefix
/// " installed " yesno
/// " id " eui64
/// " metric " uint " refmetric " uint
/// " via " ip " if " ifname
/// ```
///
/// The accepted grammar is written in a regex/BNF-ish notation:
///
/// - `e1 e2` means concatenation
/// - `e1 | e2` means choice
/// - `e*` means zero or more
/// - `e+` means one or more
/// - `e?` means optional
/// - `(e)` groups expressions
///
/// # Notes
///
/// - The `xroute` summary `prefix-from` token is parsed only to consume the wire format;
/// the later `prefix` and `from` fields are treated as the authoritative values.
/// - The parser is intentionally strict about the documented token set. Internal defensive
/// fallbacks in `babeld` such as `???` are not treated as part of the formal grammar.
pub mod parse {
use crate::babel::Eui64;
use crate::babel::command::{
NEIGHBOUR_COST_BIAS_256_MAX, NEIGHBOUR_COST_BIAS_256_MIN, NEIGHBOUR_COST_COEF_256_MAX,
NEIGHBOUR_COST_COEF_256_MIN,
};
use crate::babel::line::{
BabelLine, Event, EventKind, HeaderLine, InterfaceEvent, NeighbourEvent, RouteEvent,
Status, XRouteEvent,
};
use ipnet::IpNet;
use memchr::memchr;
use std::{
net::{IpAddr, Ipv4Addr},
str::FromStr,
};
use thiserror::Error;
use winnow::{
ascii::{dec_int, dec_uint, hex_uint, space1},
combinator::{alt, eof, opt, preceded, terminated},
error::ContextError,
prelude::*,
token::{rest, take_till},
};
#[derive(Error, Debug)]
pub enum ParseError {
#[error("invalid utf8 in babeld output: {0}")]
InvalidUtf8(#[from] std::str::Utf8Error),
#[error("failed to parse babeld line {line:?}: {error}")]
Syntax { line: String, error: String },
}
/// Zero-copy line framing for already-buffered socket output.
///
/// This is the `stream = { line }` part of the grammar: framing happens first,
/// then each line is parsed independently by `parse_line`.
#[derive(Debug, Clone)]
pub struct RawLines<'a> {
remaining: &'a [u8],
}
impl<'a> RawLines<'a> {
pub fn new(bytes: &'a [u8]) -> Self {
Self { remaining: bytes }
}
}
impl<'a> Iterator for RawLines<'a> {
type Item = Result<&'a str, ParseError>;
fn next(&mut self) -> Option<Self::Item> {
if self.remaining.is_empty() {
return None;
}
let split = memchr(b'\n', self.remaining);
let (line, rest_bytes) = match split {
Some(idx) => (&self.remaining[..idx], &self.remaining[idx + 1..]),
None => (self.remaining, &[][..]),
};
self.remaining = rest_bytes;
let line = if let Some(stripped) = line.strip_suffix(b"\r") {
stripped
} else {
line
};
Some(std::str::from_utf8(line).map_err(ParseError::InvalidUtf8))
}
}
/// Convenience adapter for parsing a fully buffered transcript, e.g. a dump.
#[derive(Debug, Clone)]
pub struct ParsedLines<'a> {
raw: RawLines<'a>,
}
impl<'a> ParsedLines<'a> {
pub fn new(bytes: &'a [u8]) -> Self {
Self {
raw: RawLines::new(bytes),
}
}
}
impl<'a> Iterator for ParsedLines<'a> {
type Item = Result<BabelLine, ParseError>;
fn next(&mut self) -> Option<Self::Item> {
self.raw.next().map(|line| line.and_then(parse_line))
}
}
pub fn parse_line(line: &str) -> Result<BabelLine, ParseError> {
terminated(parse_babel_line, eof)
.parse(line)
.map_err(|err| ParseError::Syntax {
line: line.to_owned(),
error: err.to_string(),
})
}
fn parse_babel_line(input: &mut &str) -> ModalResult<BabelLine> {
alt((
parse_banner,
parse_version,
parse_host,
parse_my_id,
parse_ok,
parse_bad,
parse_no,
parse_event,
))
.parse_next(input)
}
fn parse_banner(input: &mut &str) -> ModalResult<BabelLine> {
let _ = "BABEL ".parse_next(input)?;
let major = dec_uint::<_, u8, _>.parse_next(input)?;
let _ = '.'.parse_next(input)?;
let minor = dec_uint::<_, u8, _>.parse_next(input)?;
Ok(BabelLine::Header(HeaderLine::Banner { major, minor }))
}
fn parse_version(input: &mut &str) -> ModalResult<BabelLine> {
let _ = "version ".parse_next(input)?;
let version = Box::<str>::from(rest.parse_next(input)?);
Ok(BabelLine::Header(HeaderLine::Version(version)))
}
fn parse_host(input: &mut &str) -> ModalResult<BabelLine> {
let _ = "host ".parse_next(input)?;
let host = Box::<str>::from(rest.parse_next(input)?);
Ok(BabelLine::Header(HeaderLine::Host(host)))
}
fn parse_my_id(input: &mut &str) -> ModalResult<BabelLine> {
let _ = "my-id ".parse_next(input)?;
let id = parse_eui64.parse_next(input)?;
Ok(BabelLine::Header(HeaderLine::MyId(id)))
}
fn parse_ok(input: &mut &str) -> ModalResult<BabelLine> {
let _ = "ok".parse_next(input)?;
Ok(BabelLine::Status(Status::Ok))
}
fn parse_bad(input: &mut &str) -> ModalResult<BabelLine> {
let _ = "bad".parse_next(input)?;
Ok(BabelLine::Status(Status::Bad))
}
fn parse_no(input: &mut &str) -> ModalResult<BabelLine> {
let _ = "no".parse_next(input)?;
let message = opt(preceded(space1, rest)).parse_next(input)?;
let message = message.filter(|msg| !msg.is_empty()).map(Into::into);
Ok(BabelLine::Status(Status::No(message)))
}
fn parse_event(input: &mut &str) -> ModalResult<BabelLine> {
let kind = parse_kind.parse_next(input)?;
let _ = ' '.parse_next(input)?;
let entity = parse_word.parse_next(input)?;
match entity {
"interface" => parse_interface_event(kind, input).map(Event::Interface),
"neighbour" => parse_neighbour_event(kind, input).map(Event::Neighbour),
"xroute" => parse_xroute_event(kind, input).map(Event::XRoute),
"route" => parse_route_event(kind, input).map(Event::Route),
_ => Err(winnow::error::ErrMode::Backtrack(ContextError::new())),
}
.map(BabelLine::Event)
}
fn parse_interface_event(kind: EventKind, input: &mut &str) -> ModalResult<InterfaceEvent> {
let _ = ' '.parse_next(input)?;
let ifname = parse_word.parse_next(input)?;
let _ = " up ".parse_next(input)?;
let up = parse_bool.parse_next(input)?;
let ipv6 = opt(preceded(" ipv6 ", parse_ip_addr)).parse_next(input)?;
let ipv4 = opt(preceded(" ipv4 ", parse_ipv4_addr)).parse_next(input)?;
Ok(InterfaceEvent {
kind,
ifname: ifname.into(),
up,
ipv6,
ipv4,
})
}
fn parse_neighbour_event(kind: EventKind, input: &mut &str) -> ModalResult<NeighbourEvent> {
let _ = ' '.parse_next(input)?;
let handle = parse_hex_u64.parse_next(input)?;
let _ = " address ".parse_next(input)?;
let address = parse_ip_addr.parse_next(input)?;
let _ = " if ".parse_next(input)?;
let ifname = parse_word.parse_next(input)?;
let _ = " reach ".parse_next(input)?;
let reach = parse_hex_u16.parse_next(input)?;
let _ = " ureach ".parse_next(input)?;
let ureach = parse_hex_u16.parse_next(input)?;
let _ = " rxcost ".parse_next(input)?;
let rxcost = dec_uint::<_, u32, _>.parse_next(input)?;
let _ = " txcost ".parse_next(input)?;
let txcost = dec_uint::<_, u32, _>.parse_next(input)?;
let rtt = opt(parse_rtt_clause).parse_next(input)?;
let external_cost = opt(parse_external_cost_clause).parse_next(input)?;
let (external_bias_256, external_coef_256) = external_cost.unwrap_or((0, 256));
let _ = " cost ".parse_next(input)?;
let cost = dec_uint::<_, u32, _>.parse_next(input)?;
Ok(NeighbourEvent {
kind,
handle,
address,
ifname: ifname.into(),
reach,
ureach,
rxcost,
txcost,
rtt_millis: rtt.map(|(millis, _)| millis),
rttcost: rtt.map(|(_, cost)| cost),
external_bias_256,
external_coef_256,
cost,
})
}
fn parse_xroute_event(kind: EventKind, input: &mut &str) -> ModalResult<XRouteEvent> {
let _ = ' '.parse_next(input)?;
let _summary_prefix = parse_prefix_until('-').parse_next(input)?;
let _ = '-'.parse_next(input)?;
let _summary_from = parse_prefix.parse_next(input)?;
let _ = " prefix ".parse_next(input)?;
let prefix = parse_prefix.parse_next(input)?;
let _ = " from ".parse_next(input)?;
let from = parse_prefix.parse_next(input)?;
let _ = " metric ".parse_next(input)?;
let metric = dec_uint::<_, u32, _>.parse_next(input)?;
Ok(XRouteEvent {
kind,
prefix,
from,
metric,
})
}
fn parse_route_event<'a>(kind: EventKind, input: &mut &'a str) -> ModalResult<RouteEvent> {
let _ = ' '.parse_next(input)?;
let handle = parse_hex_u64.parse_next(input)?;
let _ = " prefix ".parse_next(input)?;
let prefix = parse_prefix.parse_next(input)?;
let _ = " from ".parse_next(input)?;
let from = parse_prefix.parse_next(input)?;
let _ = " installed ".parse_next(input)?;
let installed = parse_yes_no.parse_next(input)?;
let _ = " id ".parse_next(input)?;
let id = parse_eui64.parse_next(input)?;
let _ = " metric ".parse_next(input)?;
let metric = dec_uint::<_, u32, _>.parse_next(input)?;
let _ = " refmetric ".parse_next(input)?;
let refmetric = dec_uint::<_, u32, _>.parse_next(input)?;
let _ = " via ".parse_next(input)?;
let via = parse_ip_addr.parse_next(input)?;
let _ = " if ".parse_next(input)?;
let ifname = parse_word.parse_next(input)?;
Ok(RouteEvent {
kind,
handle,
prefix,
from,
installed,
id,
metric,
refmetric,
via,
ifname: ifname.into(),
})
}
fn parse_rtt_clause(input: &mut &str) -> ModalResult<(u32, u32)> {
let _ = " rtt ".parse_next(input)?;
let millis = parse_millis.parse_next(input)?;
let _ = " rttcost ".parse_next(input)?;
let rttcost = dec_uint::<_, u32, _>.parse_next(input)?;
Ok((millis, rttcost))
}
fn parse_external_cost_clause(input: &mut &str) -> ModalResult<(i32, u32)> {
let _ = " external-bias-256 ".parse_next(input)?;
let bias_256 = parse_external_bias_256.parse_next(input)?;
let _ = " external-coef-256 ".parse_next(input)?;
let coef_256 = parse_external_coef_256.parse_next(input)?;
Ok((bias_256, coef_256))
}
fn parse_external_bias_256(input: &mut &str) -> ModalResult<i32> {
let value = dec_int::<_, i32, _>.parse_next(input)?;
if (NEIGHBOUR_COST_BIAS_256_MIN..=NEIGHBOUR_COST_BIAS_256_MAX).contains(&value) {
Ok(value)
} else {
Err(winnow::error::ErrMode::Backtrack(ContextError::new()))
}
}
fn parse_external_coef_256(input: &mut &str) -> ModalResult<u32> {
let value = dec_uint::<_, u32, _>.parse_next(input)?;
if (NEIGHBOUR_COST_COEF_256_MIN..=NEIGHBOUR_COST_COEF_256_MAX).contains(&value) {
Ok(value)
} else {
Err(winnow::error::ErrMode::Backtrack(ContextError::new()))
}
}
fn parse_kind(input: &mut &str) -> ModalResult<EventKind> {
alt((
"add".value(EventKind::Add),
"change".value(EventKind::Change),
"flush".value(EventKind::Flush),
))
.parse_next(input)
}
fn parse_bool(input: &mut &str) -> ModalResult<bool> {
alt(("true".value(true), "false".value(false))).parse_next(input)
}
fn parse_yes_no(input: &mut &str) -> ModalResult<bool> {
alt(("yes".value(true), "no".value(false))).parse_next(input)
}
fn parse_ip_addr(input: &mut &str) -> ModalResult<IpAddr> {
parse_word.try_map(IpAddr::from_str).parse_next(input)
}
fn parse_ipv4_addr(input: &mut &str) -> ModalResult<Ipv4Addr> {
parse_word.try_map(Ipv4Addr::from_str).parse_next(input)
}
fn parse_prefix(input: &mut &str) -> ModalResult<IpNet> {
parse_word.try_map(IpNet::from_str).parse_next(input)
}
fn parse_prefix_until(separator: char) -> impl FnMut(&mut &str) -> ModalResult<IpNet> {
move |input: &mut &str| {
let token = take_till(1.., |c: char| c == separator).parse_next(input)?;
IpNet::from_str(token)
.map_err(|_| winnow::error::ErrMode::Backtrack(ContextError::new()))
}
}
fn parse_eui64(input: &mut &str) -> ModalResult<Eui64> {
parse_word.try_map(Eui64::from_str).parse_next(input)
}
fn parse_hex_u64(input: &mut &str) -> ModalResult<u64> {
hex_uint.parse_next(input)
}
fn parse_hex_u16(input: &mut &str) -> ModalResult<u16> {
hex_uint.parse_next(input)
}
fn parse_millis(input: &mut &str) -> ModalResult<u32> {
let word = parse_word.parse_next(input)?;
parse_millis_str(word).map_err(|_| winnow::error::ErrMode::Backtrack(ContextError::new()))
}
fn parse_word<'a>(input: &mut &'a str) -> ModalResult<&'a str> {
take_till(1.., |c: char| c == ' ').parse_next(input)
}
fn parse_millis_str(value: &str) -> Result<u32, &'static str> {
let (secs, millis) = value
.split_once('.')
.ok_or("missing milliseconds separator")?;
if millis.len() != 3 || !millis.bytes().all(|b| b.is_ascii_digit()) {
return Err("expected 3-digit millisecond suffix");
}
let secs = secs
.parse::<u32>()
.map_err(|_| "invalid seconds field in rtt value")?;
let millis = millis
.parse::<u32>()
.map_err(|_| "invalid milliseconds field in rtt value")?;
secs.checked_mul(1000)
.and_then(|s| s.checked_add(millis))
.ok_or("rtt value overflowed u32")
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::babel::line::parse::{ParsedLines, parse_line};
use std::str::FromStr;
#[test]
fn parse_header_banner() {
assert_eq!(
parse_line("BABEL 1.0").unwrap(),
BabelLine::Header(HeaderLine::Banner { major: 1, minor: 0 })
);
}
#[test]
fn parse_header_metadata() {
assert_eq!(
parse_line("version babeld-1.13.1").unwrap(),
BabelLine::Header(HeaderLine::Version("babeld-1.13.1".into()))
);
assert_eq!(
parse_line("host e2").unwrap(),
BabelLine::Header(HeaderLine::Host("e2".into()))
);
assert_eq!(
parse_line("my-id 02:00:00:00:00:00:00:01").unwrap(),
BabelLine::Header(HeaderLine::MyId(Eui64::new(2, 0, 0, 0, 0, 0, 0, 1)))
);
}
#[test]
fn parse_status_lines() {
assert_eq!(parse_line("ok").unwrap(), BabelLine::Status(Status::Ok));
assert_eq!(parse_line("bad").unwrap(), BabelLine::Status(Status::Bad));
assert_eq!(
parse_line("no No such interface").unwrap(),
BabelLine::Status(Status::No(Some("No such interface".into())))
);
}
#[test]
fn parse_interface_event() {
assert_eq!(
parse_line("add interface en2 up true ipv6 fe80::1 ipv4 169.254.1.2").unwrap(),
BabelLine::Event(Event::Interface(InterfaceEvent {
kind: EventKind::Add,
ifname: "en2".into(),
up: true,
ipv6: Some(IpAddr::from_str("fe80::1").unwrap()),
ipv4: Some(Ipv4Addr::new(169, 254, 1, 2)),
}))
);
assert_eq!(
parse_line("change interface en3 up false").unwrap(),
BabelLine::Event(Event::Interface(InterfaceEvent {
kind: EventKind::Change,
ifname: "en3".into(),
up: false,
ipv6: None,
ipv4: None,
}))
);
}
#[test]
fn parse_neighbour_event() {
assert_eq!(
parse_line(
"add neighbour 7ffdeadbeef address fe80::1 if en2 reach 00ff ureach 000f rxcost 256 txcost 96 rtt 0.123 rttcost 32 cost 128"
)
.unwrap(),
BabelLine::Event(Event::Neighbour(NeighbourEvent {
kind: EventKind::Add,
handle: 0x7ffdeadbeef,
address: IpAddr::from_str("fe80::1").unwrap(),
ifname: "en2".into(),
reach: 0x00ff,
ureach: 0x000f,
rxcost: 256,
txcost: 96,
rtt_millis: Some(123),
rttcost: Some(32),
external_bias_256: 0,
external_coef_256: 256,
cost: 128,
}))
);
assert_eq!(
parse_line(
"change neighbour 7ffdeadbeef address fe80::1 if en2 reach ffff ureach 000f rxcost 96 txcost 96 external-bias-256 -512 external-coef-256 128 cost 48"
)
.unwrap(),
BabelLine::Event(Event::Neighbour(NeighbourEvent {
kind: EventKind::Change,
handle: 0x7ffdeadbeef,
address: IpAddr::from_str("fe80::1").unwrap(),
ifname: "en2".into(),
reach: 0xffff,
ureach: 0x000f,
rxcost: 96,
txcost: 96,
rtt_millis: None,
rttcost: None,
external_bias_256: -512,
external_coef_256: 128,
cost: 48,
}))
);
assert_eq!(
parse_line(
"change neighbour 7ffdeadbeef address fe80::1 if en2 reach ffff ureach 000f rxcost 96 txcost 96 rtt 0.123 rttcost 32 external-bias-256 512 external-coef-256 512 cost 224"
)
.unwrap(),
BabelLine::Event(Event::Neighbour(NeighbourEvent {
kind: EventKind::Change,
handle: 0x7ffdeadbeef,
address: IpAddr::from_str("fe80::1").unwrap(),
ifname: "en2".into(),
reach: 0xffff,
ureach: 0x000f,
rxcost: 96,
txcost: 96,
rtt_millis: Some(123),
rttcost: Some(32),
external_bias_256: 512,
external_coef_256: 512,
cost: 224,
}))
);
}
#[test]
fn parse_neighbour_event_rejects_out_of_range_external_cost() {
assert!(
parse_line(
"change neighbour 7ffdeadbeef address fe80::1 if en2 reach ffff ureach 000f rxcost 96 txcost 96 external-bias-256 16776705 external-coef-256 256 cost 96"
)
.is_err()
);
assert!(
parse_line(
"change neighbour 7ffdeadbeef address fe80::1 if en2 reach ffff ureach 000f rxcost 96 txcost 96 external-bias-256 0 external-coef-256 65536 cost 96"
)
.is_err()
);
}
#[test]
fn parse_xroute_event() {
assert_eq!(
parse_line(
"add xroute fd00::1/128-fd00::/64 prefix fd00::1/128 from fd00::/64 metric 0"
)
.unwrap(),
BabelLine::Event(Event::XRoute(XRouteEvent {
kind: EventKind::Add,
prefix: IpNet::from_str("fd00::1/128").unwrap(),
from: IpNet::from_str("fd00::/64").unwrap(),
metric: 0,
}))
);
}
#[test]
fn parse_route_event() {
assert_eq!(
parse_line(
"change route 7ffdeadbeef prefix fd00::1/128 from fd00::/64 installed yes id 02:00:00:00:00:00:00:01 metric 96 refmetric 0 via fe80::2 if en2"
)
.unwrap(),
BabelLine::Event(Event::Route(RouteEvent {
kind: EventKind::Change,
handle: 0x7ffdeadbeef,
prefix: IpNet::from_str("fd00::1/128").unwrap(),
from: IpNet::from_str("fd00::/64").unwrap(),
installed: true,
id: Eui64::new(2, 0, 0, 0, 0, 0, 0, 1),
metric: 96,
refmetric: 0,
via: IpAddr::from_str("fe80::2").unwrap(),
ifname: "en2".into(),
}))
);
}
#[test]
fn raw_lines_uses_memchr_framing() {
let bytes = b"BABEL 1.0\nok\nadd interface en2 up false\n";
let parsed = ParsedLines::new(bytes)
.collect::<Result<Vec<_>, _>>()
.unwrap();
assert_eq!(
parsed,
vec![
BabelLine::Header(HeaderLine::Banner { major: 1, minor: 0 }),
BabelLine::Status(Status::Ok),
BabelLine::Event(Event::Interface(InterfaceEvent {
kind: EventKind::Add,
ifname: "en2".into(),
up: false,
ipv6: None,
ipv4: None,
})),
]
);
}
}
-197
View File
@@ -1,197 +0,0 @@
//! Temporary link-selection policy for broad macOS interface admission.
//!
//! The current MVP admits every usable interface and then steers Babel by
//! assigning each `enN` neighbour an absolute synthetic base cost of `N * 100`,
//! except that `en0` and `en1` are assigned the largest finite Babel cost. This is
//! intentionally a stopgap until measured link scoring lands.
use std::net::IpAddr;
use crate::babel::command::{
NEIGHBOUR_COST_BIAS_256_MAX, NeighbourCostBias256, NeighbourCostCoef256,
NeighbourCostCommand,
};
use crate::babel::line::{EventKind, NeighbourEvent};
use crate::babel::state::NeighbourState;
const EN_INDEX_COST_UNITS: u64 = 100;
const FIXED_POINT_SCALE: u64 = 256;
const ABSOLUTE_COST_COEF_256: u32 = 0;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
struct DesiredNeighbourCost {
bias_256: NeighbourCostBias256,
coef_256: NeighbourCostCoef256,
}
pub(crate) fn command_for_neighbour_event(event: &NeighbourEvent) -> Option<NeighbourCostCommand> {
if !matches!(event.kind, EventKind::Add | EventKind::Change) {
return None;
}
command_for_neighbour(
&event.ifname,
event.address,
event.external_bias_256,
event.external_coef_256,
)
}
pub(crate) fn command_for_neighbour_state(
neighbour: &NeighbourState,
) -> Option<NeighbourCostCommand> {
command_for_neighbour(
&neighbour.ifname,
neighbour.address,
neighbour.external_bias_256,
neighbour.external_coef_256,
)
}
fn command_for_neighbour(
ifname: &str,
address: IpAddr,
external_bias_256: i32,
external_coef_256: u32,
) -> Option<NeighbourCostCommand> {
let desired = desired_en_index_cost(ifname)?;
if external_bias_256 == desired.bias_256.raw() && external_coef_256 == desired.coef_256.raw() {
return None;
}
let IpAddr::V6(link_local_neighbour) = address else {
return None;
};
NeighbourCostCommand::new(
ifname,
link_local_neighbour,
desired.bias_256,
desired.coef_256,
)
.ok()
}
fn desired_en_index_cost(ifname: &str) -> Option<DesiredNeighbourCost> {
let index = parse_en_index(ifname)?;
let bias_256 = if index <= 1 {
NEIGHBOUR_COST_BIAS_256_MAX
} else {
let bias_256 = u64::from(index)
.checked_mul(EN_INDEX_COST_UNITS)?
.checked_mul(FIXED_POINT_SCALE)?;
i32::try_from(bias_256).ok()?
};
let bias_256 = NeighbourCostBias256::new(bias_256)?;
let coef_256 = NeighbourCostCoef256::new(ABSOLUTE_COST_COEF_256)
.expect("absolute-cost coefficient is within babeld's accepted range");
Some(DesiredNeighbourCost { bias_256, coef_256 })
}
fn parse_en_index(ifname: &str) -> Option<u32> {
let suffix = ifname.strip_prefix("en")?;
if suffix.is_empty() || !suffix.bytes().all(|byte| byte.is_ascii_digit()) {
return None;
}
suffix.parse().ok()
}
#[cfg(test)]
mod tests {
use std::net::{IpAddr, Ipv4Addr};
use super::{command_for_neighbour_event, command_for_neighbour_state};
use crate::babel::line::{EventKind, NeighbourEvent};
use crate::babel::state::NeighbourState;
fn neighbour_event(ifname: &str, bias: i32, coef: u32) -> NeighbourEvent {
NeighbourEvent {
kind: EventKind::Change,
handle: 0x42,
address: "fe80::1".parse().unwrap(),
ifname: ifname.into(),
reach: 0xffff,
ureach: 0,
rxcost: 96,
txcost: 96,
rtt_millis: None,
rttcost: None,
external_bias_256: bias,
external_coef_256: coef,
cost: 96,
}
}
fn neighbour_state(ifname: &str, bias: i32, coef: u32) -> NeighbourState {
NeighbourState {
handle: 0x42,
address: "fe80::1".parse().unwrap(),
ifname: ifname.into(),
reach: 0xffff,
ureach: 0,
rxcost: 96,
txcost: 96,
rtt_millis: None,
rttcost: None,
external_bias_256: bias,
external_coef_256: coef,
cost: 96,
}
}
#[test]
fn en_index_policy_sets_absolute_cost() {
let command = command_for_neighbour_event(&neighbour_event("en18", 0, 256)).unwrap();
assert_eq!(
command.to_string(),
"neighbour-cost en18 fe80::1 bias-256 460800 coef-256 0"
);
}
#[test]
fn en_index_policy_deprioritizes_en0_and_en1() {
let command = command_for_neighbour_event(&neighbour_event("en0", 0, 256)).unwrap();
assert_eq!(
command.to_string(),
"neighbour-cost en0 fe80::1 bias-256 16776704 coef-256 0"
);
let command = command_for_neighbour_event(&neighbour_event("en1", 0, 256)).unwrap();
assert_eq!(
command.to_string(),
"neighbour-cost en1 fe80::1 bias-256 16776704 coef-256 0"
);
}
#[test]
fn en_index_policy_skips_already_configured_neighbour() {
assert!(command_for_neighbour_state(&neighbour_state("en2", 51_200, 0)).is_none());
assert!(
command_for_neighbour_state(&neighbour_state("en0", 16_776_704, 0)).is_none()
);
assert!(
command_for_neighbour_state(&neighbour_state("en1", 16_776_704, 0)).is_none()
);
}
#[test]
fn en_index_policy_ignores_non_en_or_non_link_local_neighbours() {
assert!(command_for_neighbour_event(&neighbour_event("awdl0", 0, 256)).is_none());
let mut ipv4 = neighbour_event("en2", 0, 256);
ipv4.address = IpAddr::V4(Ipv4Addr::new(169, 254, 1, 2));
assert!(command_for_neighbour_event(&ipv4).is_none());
let mut non_link_local = neighbour_event("en2", 0, 256);
non_link_local.address = "2001:db8::1".parse().unwrap();
assert!(command_for_neighbour_event(&non_link_local).is_none());
}
#[test]
fn en_index_policy_ignores_flush_events() {
let mut event = neighbour_event("en2", 0, 256);
event.kind = EventKind::Flush;
assert!(command_for_neighbour_event(&event).is_none());
}
}
-38
View File
@@ -1,38 +0,0 @@
use ipnet::Ipv6Net;
use std::sync::Arc;
use tokio::sync::{mpsc, watch};
use crate::Result;
pub mod command;
pub mod line;
pub mod link_policy;
pub mod runtime;
pub mod state;
use runtime::BabelRuntime;
/// An EUI-64 type aliased to [`macaddr::MacAddr8`].
pub type Eui64 = macaddr::MacAddr8;
pub use command::{
NeighbourCostBias256, NeighbourCostCoef256, NeighbourCostCommand, NeighbourCostCommandError,
};
pub use state::BabelState;
#[derive(Debug)]
pub enum Babble {
AddIface(Box<str>),
SetNeighbourCost(NeighbourCostCommand),
}
#[tracing::instrument(skip(state_send, recv))]
pub async fn babel(
advertised: Ipv6Net,
recv: mpsc::Receiver<Babble>,
state_send: watch::Sender<Arc<BabelState>>,
) -> Result<()> {
let mut runtime = BabelRuntime::spawn(advertised, state_send).await?;
let res1 = runtime.run(recv).await;
let res2 = runtime.shutdown().await;
res1.and(res2)
}
-487
View File
@@ -1,487 +0,0 @@
//! Managed `babeld` runtime for `babblerd`.
//!
//! This module owns the full lifecycle of the private `babeld` instance:
//!
//! - spawn-time configuration of the child process
//! - the private Unix socket path used for the local control connection
//! - connecting to that socket and speaking the local Babel protocol
//! - running the monitor-driven control loop
//! - shutdown and cleanup of the child process and socket
//!
//! Unlike the old `process` / `session` split, this is intended to model the real runtime unit:
//! a single managed `babeld` process together with its single local control session.
use std::fs::Permissions;
use std::io;
use std::os::unix::fs::PermissionsExt;
use std::sync::Arc;
use ipnet::Ipv6Net;
use nix::errno::Errno;
use nix::sys::signal::{Signal, kill};
use nix::unistd::Pid;
use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader, Lines};
use tokio::net::UnixStream;
use tokio::net::unix::{OwnedReadHalf, OwnedWriteHalf};
use tokio::process::{Child, Command};
use tokio::sync::{mpsc, watch};
use tokio::time::{Duration, MissedTickBehavior, timeout};
use crate::babel::Babble;
use crate::babel::command::BabelCommand;
use crate::babel::line::parse::ParseError;
use crate::babel::line::{self, BabelLine, Event, HeaderLine, NeighbourEvent, Status};
use crate::babel::link_policy;
use crate::babel::state::BabelState;
use crate::{BabbleError, Result};
#[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";
const STARTUP_SOCKET_TIMEOUT: Duration = Duration::from_secs(10);
const STARTUP_SOCKET_POLL_INTERVAL: Duration = Duration::from_millis(50);
const SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(5);
pub(crate) struct BabelRuntime {
proc: Child,
read: Lines<BufReader<OwnedReadHalf>>,
write: OwnedWriteHalf,
state_send: watch::Sender<Arc<BabelState>>,
state: BabelState,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum StartupStage {
Banner,
Version,
Host,
MyId,
Ready,
}
impl StartupStage {
fn advance(self, line: BabelLine) -> Result<Option<Self>> {
match (self, line) {
(Self::Banner, BabelLine::Header(HeaderLine::Banner { major: 1, minor: 0 })) => {
Ok(Some(Self::Version))
}
(Self::Version, BabelLine::Header(HeaderLine::Version(_))) => Ok(Some(Self::Host)),
(Self::Host, BabelLine::Header(HeaderLine::Host(_))) => Ok(Some(Self::MyId)),
(Self::MyId, BabelLine::Header(HeaderLine::MyId(_))) => Ok(Some(Self::Ready)),
(Self::Ready, BabelLine::Status(Status::Ok)) => Ok(None),
(stage, other) => Err(BabbleError::Other(format!(
"unexpected babeld startup line while waiting for {stage:?}: {other:?}"
))),
}
}
}
impl Drop for BabelRuntime {
#[inline]
fn drop(&mut self) {
// Emergency SIGKILL to avoid leaking an unmanaged babeld subprocess.
match self.proc.try_wait() {
Ok(None) => {}
Ok(Some(sc)) => {
if !sc.success() {
_ = self.proc.start_kill();
}
}
_ => {
_ = self.proc.start_kill();
}
}
}
}
impl BabelRuntime {
#[tracing::instrument(skip(state_send))]
pub(crate) async fn spawn(
advertised: Ipv6Net,
state_send: watch::Sender<Arc<BabelState>>,
) -> Result<Self> {
tokio::fs::create_dir_all(PRIVATE_DIR).await?;
// TODO: remove this magic constant (and magic constants in general)
tokio::fs::set_permissions(PRIVATE_DIR, Permissions::from_mode(0o0700)).await?;
tracing::info!("spawning babeld socket in {PRIVATE_SOCK_PATH}");
let mut proc = match Command::new("babeld")
.arg("-G")
.arg(PRIVATE_SOCK_PATH)
.arg("-I")
.arg(format!("{PRIVATE_DIR}/babeld.pid"))
.arg("-C")
.arg("kernel-install false")
.arg("-C")
.arg(format!("redistribute local ip {advertised}"))
.arg("-C")
.arg("redistribute local deny")
.spawn()
{
Ok(proc) => {
tracing::info!(
"babeld spawned PID={}",
proc.id().expect("babeld process shouldn't die this early")
);
proc
}
Err(e) => {
tracing::warn!(error=%e, "failed to spawn babeld");
return Err(e.into());
}
};
if let Err(err) = Self::wait_for_socket(&mut proc).await {
Self::abort_child(&mut proc).await;
return Err(err);
}
// TODO: magic undocumented number
if let Err(err) =
std::fs::set_permissions(PRIVATE_SOCK_PATH, Permissions::from_mode(0o0600))
{
Self::abort_child(&mut proc).await;
return Err(err.into());
}
let (reader, write) = match UnixStream::connect(PRIVATE_SOCK_PATH).await {
Ok(stream) => stream.into_split(),
Err(err) => {
Self::abort_child(&mut proc).await;
return Err(err.into());
}
};
let mut runtime = Self {
proc,
read: BufReader::new(reader).lines(),
write,
state_send,
state: BabelState::new(),
};
if let Err(err) = runtime.await_ready().await {
let _ = runtime.shutdown().await;
return Err(err);
}
Ok(runtime)
}
async fn wait_for_socket(proc: &mut Child) -> Result<()> {
timeout(STARTUP_SOCKET_TIMEOUT, async {
let mut poll = tokio::time::interval(STARTUP_SOCKET_POLL_INTERVAL);
poll.set_missed_tick_behavior(MissedTickBehavior::Delay);
loop {
poll.tick().await;
match tokio::fs::try_exists(PRIVATE_SOCK_PATH).await {
Ok(true) => return Ok(()),
Ok(false) => {}
Err(err) => return Err(err.into()),
}
if let Some(status) = proc.try_wait()? {
return Err(BabbleError::BabeldCrashed(status.code()));
}
}
})
.await
.unwrap_or_else(|_| {
Err(BabbleError::Other(format!(
"timed out after {}s waiting for babeld socket {PRIVATE_SOCK_PATH}",
STARTUP_SOCKET_TIMEOUT.as_secs()
)))
})
}
async fn abort_child(proc: &mut Child) {
let _ = proc.kill().await;
}
#[tracing::instrument(skip_all)]
async fn await_ready(&mut self) -> Result<()> {
let mut stage = StartupStage::Banner;
while let Some(line) = self.read.next_line().await? {
match self.observe_line(line)? {
Ok(parsed) => match stage.advance(parsed)? {
Some(next) => stage = next,
None => {
tracing::info!("babeld ok");
return Ok(());
}
},
Err(err) => {
return Err(BabbleError::Other(format!(
"failed to parse babeld startup prelude: {err}"
)));
}
}
}
Err(BabbleError::Other(
"babeld closed before completing startup prelude".into(),
))
}
#[tracing::instrument(skip(self))]
async fn query(&mut self, cmd: &BabelCommand) -> io::Result<Option<Status>> {
self.write.write_all(cmd.encode().as_bytes()).await?;
loop {
let Some(line) = self.read.next_line().await? else {
tracing::warn!("babeld closed unexpectedly");
return Ok(None);
};
match self.observe_line(line)? {
Ok(parsed) => {
let status = self.reduce_live_line(parsed)?;
let Some(status) = status else {
continue;
};
match &status {
Status::Ok => {}
Status::Bad => tracing::warn!("malformed message sent to babeld"),
Status::No(rest) => tracing::warn!("message rejected: {rest:?}"),
}
return Ok(Some(status));
}
Err(err) => {
return Err(io::Error::new(
io::ErrorKind::InvalidData,
format!("failed to parse babeld command output: {err}"),
));
}
}
}
}
#[tracing::instrument(skip(self))]
fn observe_line(&self, line: String) -> io::Result<std::result::Result<BabelLine, ParseError>> {
tracing::info!("[babel] {:?}", line);
let observed = match line::parse::parse_line(&line) {
Ok(parsed) => {
tracing::info!("[parsed] {:?}", parsed);
Ok(parsed)
}
Err(err) => {
tracing::error!(error=%err, "failed to parse babeld line");
Err(err)
}
};
Ok(observed)
}
#[tracing::instrument(skip(self))]
async fn start_monitoring(&mut self) -> io::Result<Option<Status>> {
let mut snapshot = BabelState::new();
self.write
.write_all(BabelCommand::Monitor.encode().as_bytes())
.await?;
loop {
let Some(line) = self.read.next_line().await? else {
tracing::warn!("babeld closed unexpectedly");
return Ok(None);
};
match self.observe_line(line)? {
Ok(BabelLine::Event(event)) => {
snapshot.apply(event);
}
Ok(BabelLine::Status(status)) => {
match &status {
Status::Ok => {
self.state = snapshot;
self.publish_state();
}
Status::Bad => tracing::warn!("malformed message sent to babeld"),
Status::No(rest) => tracing::warn!("message rejected: {rest:?}"),
}
return Ok(Some(status));
}
Ok(BabelLine::Header(header)) => {
return Err(io::Error::new(
io::ErrorKind::InvalidData,
format!("unexpected header line during monitor bootstrap: {header:?}"),
));
}
Err(err) => {
return Err(io::Error::new(
io::ErrorKind::InvalidData,
format!("failed to parse babeld monitor bootstrap output: {err}"),
));
}
}
}
}
fn publish_state(&self) {
self.state_send.send_replace(Arc::new(self.state.clone()));
}
fn reduce_live_line(&mut self, line: BabelLine) -> io::Result<Option<Status>> {
match line {
BabelLine::Event(event) => {
self.state.apply(event);
self.publish_state();
Ok(None)
}
BabelLine::Status(status) => Ok(Some(status)),
BabelLine::Header(header) => Err(io::Error::new(
io::ErrorKind::InvalidData,
format!("unexpected header line after startup: {header:?}"),
)),
}
}
#[tracing::instrument(skip(self))]
async fn reconcile_neighbour_cost_policy(&mut self) -> io::Result<()> {
let commands = self
.state
.neighbours
.values()
.filter_map(link_policy::command_for_neighbour_state)
.collect::<Vec<_>>();
for command in commands {
tracing::info!(%command, "applying en-index neighbour-cost policy");
let command = BabelCommand::NeighbourCost(command);
self.query(&command).await?;
}
Ok(())
}
#[tracing::instrument(skip(self, neighbour))]
async fn apply_neighbour_cost_policy(&mut self, neighbour: &NeighbourEvent) -> io::Result<()> {
let Some(command) = link_policy::command_for_neighbour_event(neighbour) else {
return Ok(());
};
tracing::info!(%command, "applying en-index neighbour-cost policy");
let command = BabelCommand::NeighbourCost(command);
self.query(&command).await?;
Ok(())
}
#[tracing::instrument(skip_all)]
pub(crate) async fn run(&mut self, mut recv: mpsc::Receiver<Babble>) -> Result<()> {
match self.start_monitoring().await? {
Some(Status::Ok) => {}
Some(Status::Bad) => {
return Err(BabbleError::Other(
"babeld rejected monitor command as malformed".into(),
));
}
Some(Status::No(reason)) => {
return Err(BabbleError::Other(format!(
"babeld rejected monitor command: {reason:?}"
)));
}
None => {
return Err(BabbleError::Other(
"babeld control socket closed during monitor bootstrap".into(),
));
}
}
self.reconcile_neighbour_cost_policy().await?;
loop {
tokio::select! {
babble = recv.recv() => {
tracing::debug!("[babble] {:?}", babble);
let Some(babble) = babble else {
break;
};
match babble {
Babble::AddIface(iface) => {
let cmd = BabelCommand::Interface(iface);
self.query(&cmd).await?;
self.reconcile_neighbour_cost_policy().await?;
}
Babble::SetNeighbourCost(neighbour_cost) => {
let cmd = BabelCommand::NeighbourCost(neighbour_cost);
self.query(&cmd).await?;
}
}
},
line = self.read.next_line() => {
let line = match line {
Ok(Some(line)) => line,
Ok(None) => {
return Err(BabbleError::Other(
"babeld control socket closed during live monitoring".into(),
));
}
Err(err) => {
return Err(BabbleError::Other(format!(
"failed to read babeld control socket during live monitoring: {err}"
)));
}
};
match self.observe_line(line)? {
Ok(parsed) => {
let neighbour = match &parsed {
BabelLine::Event(Event::Neighbour(neighbour)) => {
Some(neighbour.clone())
}
_ => None,
};
if let Some(status) = self.reduce_live_line(parsed)? {
tracing::debug!(?status, "ignoring unsolicited status line from babeld");
}
if let Some(neighbour) = neighbour {
self.apply_neighbour_cost_policy(&neighbour).await?;
}
}
Err(err) => {
return Err(io::Error::new(
io::ErrorKind::InvalidData,
format!("failed to parse babeld monitor output: {err}"),
)
.into());
}
}
},
}
}
Ok(())
}
pub(crate) async fn shutdown(mut self) -> Result<()> {
let kill_res = if let Some(pid) = self.proc.id() {
let pid: i32 = pid.try_into().expect("pid overflow");
let rc_err = match kill(Pid::from_raw(pid), Signal::SIGINT) {
Ok(()) | Err(Errno::ESRCH) => Ok(()),
Err(err) => Err(io::Error::from_raw_os_error(err as i32).into()),
};
match timeout(SHUTDOWN_TIMEOUT, self.proc.wait()).await {
Ok(Ok(code)) => {
if code.success() {
rc_err
} else {
rc_err.and_then(|()| Err(BabbleError::BabeldCrashed(code.code())))
}
}
Ok(Err(e)) => Err(e.into()),
Err(_) => {
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)
}
}
-411
View File
@@ -1,411 +0,0 @@
//! Reduced in-memory state derived from `babeld` event lines.
//!
//! This module is the consumer-side counterpart to [`crate::babel::line`]:
//!
//! - [`Event`] is the wire/domain event stream emitted by `babeld`
//! - [`BabelState`] is the current snapshot obtained by reducing those events
//!
//! The reducer model is intentionally simple:
//!
//! - `add` inserts the entity into the relevant table
//! - `change` upserts the entity into the relevant table
//! - `flush` removes the entity from the relevant table
//!
//! The stored state types do **not** retain [`EventKind`], because the event
//! kind is transport/update metadata rather than persistent object state.
use std::collections::HashMap;
use std::net::{IpAddr, Ipv4Addr};
use crate::babel::Eui64;
use crate::babel::line::{
Event, EventKind, InterfaceEvent, NeighbourEvent, RouteEvent, XRouteEvent,
};
use ipnet::IpNet;
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct BabelState {
pub interfaces: HashMap<Box<str>, InterfaceState>,
pub neighbours: HashMap<u64, NeighbourState>,
pub xroutes: HashMap<XRouteKey, XRouteState>,
pub routes: HashMap<u64, RouteState>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct InterfaceState {
pub ifname: Box<str>,
pub up: bool,
pub ipv6: Option<IpAddr>,
pub ipv4: Option<Ipv4Addr>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct NeighbourState {
pub handle: u64,
pub address: IpAddr,
pub ifname: Box<str>,
pub reach: u16,
pub ureach: u16,
pub rxcost: u32,
pub txcost: u32,
pub rtt_millis: Option<u32>,
pub rttcost: Option<u32>,
pub external_bias_256: i32,
pub external_coef_256: u32,
pub cost: u32,
}
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct XRouteKey {
pub prefix: IpNet,
pub from: IpNet,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct XRouteState {
pub prefix: IpNet,
pub from: IpNet,
pub metric: u32,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct RouteState {
pub handle: u64,
pub prefix: IpNet,
pub from: IpNet,
pub installed: bool,
pub id: Eui64,
pub metric: u32,
pub refmetric: u32,
pub via: IpAddr,
pub ifname: Box<str>,
}
impl BabelState {
#[must_use]
pub fn new() -> Self {
Self::default()
}
pub fn apply(&mut self, event: Event) {
match event {
Event::Interface(event) => self.apply_interface(event),
Event::Neighbour(event) => self.apply_neighbour(event),
Event::XRoute(event) => self.apply_xroute(event),
Event::Route(event) => self.apply_route(event),
}
}
pub fn extend<I>(&mut self, events: I)
where
I: IntoIterator<Item = Event>,
{
for event in events {
self.apply(event);
}
}
fn apply_interface(&mut self, event: InterfaceEvent) {
let key = event.ifname.clone();
match event.kind {
EventKind::Add | EventKind::Change => {
self.interfaces.insert(key, event.into());
}
EventKind::Flush => {
self.interfaces.remove(&key);
}
}
}
fn apply_neighbour(&mut self, event: NeighbourEvent) {
let key = event.handle;
match event.kind {
EventKind::Add | EventKind::Change => {
self.neighbours.insert(key, event.into());
}
EventKind::Flush => {
self.neighbours.remove(&key);
}
}
}
fn apply_xroute(&mut self, event: XRouteEvent) {
let key = XRouteKey {
prefix: event.prefix,
from: event.from,
};
match event.kind {
EventKind::Add | EventKind::Change => {
self.xroutes.insert(key, event.into());
}
EventKind::Flush => {
self.xroutes.remove(&key);
}
}
}
fn apply_route(&mut self, event: RouteEvent) {
let key = event.handle;
match event.kind {
EventKind::Add | EventKind::Change => {
self.routes.insert(key, event.into());
}
EventKind::Flush => {
self.routes.remove(&key);
}
}
}
}
impl From<InterfaceEvent> for InterfaceState {
fn from(event: InterfaceEvent) -> Self {
Self {
ifname: event.ifname,
up: event.up,
ipv6: event.ipv6,
ipv4: event.ipv4,
}
}
}
impl From<NeighbourEvent> for NeighbourState {
fn from(event: NeighbourEvent) -> Self {
Self {
handle: event.handle,
address: event.address,
ifname: event.ifname,
reach: event.reach,
ureach: event.ureach,
rxcost: event.rxcost,
txcost: event.txcost,
rtt_millis: event.rtt_millis,
rttcost: event.rttcost,
external_bias_256: event.external_bias_256,
external_coef_256: event.external_coef_256,
cost: event.cost,
}
}
}
impl From<XRouteEvent> for XRouteState {
fn from(event: XRouteEvent) -> Self {
Self {
prefix: event.prefix,
from: event.from,
metric: event.metric,
}
}
}
impl From<RouteEvent> for RouteState {
fn from(event: RouteEvent) -> Self {
Self {
handle: event.handle,
prefix: event.prefix,
from: event.from,
installed: event.installed,
id: event.id,
metric: event.metric,
refmetric: event.refmetric,
via: event.via,
ifname: event.ifname,
}
}
}
#[cfg(test)]
mod tests {
use std::net::{IpAddr, Ipv4Addr, Ipv6Addr};
use super::{BabelState, InterfaceState, XRouteKey};
use crate::babel::Eui64;
use crate::babel::line::{
Event, EventKind, InterfaceEvent, NeighbourEvent, RouteEvent, XRouteEvent,
};
use ipnet::IpNet;
fn net(s: &str) -> IpNet {
s.parse().unwrap()
}
#[test]
fn interface_add_change_flush() {
let mut state = BabelState::new();
state.apply(Event::Interface(InterfaceEvent {
kind: EventKind::Add,
ifname: "en2".into(),
up: true,
ipv6: Some(IpAddr::V6(Ipv6Addr::LOCALHOST)),
ipv4: Some(Ipv4Addr::new(169, 254, 1, 2)),
}));
assert_eq!(
state.interfaces.get("en2"),
Some(&InterfaceState {
ifname: "en2".into(),
up: true,
ipv6: Some(IpAddr::V6(Ipv6Addr::LOCALHOST)),
ipv4: Some(Ipv4Addr::new(169, 254, 1, 2)),
})
);
state.apply(Event::Interface(InterfaceEvent {
kind: EventKind::Change,
ifname: "en2".into(),
up: false,
ipv6: None,
ipv4: None,
}));
assert_eq!(
state.interfaces.get("en2"),
Some(&InterfaceState {
ifname: "en2".into(),
up: false,
ipv6: None,
ipv4: None,
})
);
state.apply(Event::Interface(InterfaceEvent {
kind: EventKind::Flush,
ifname: "en2".into(),
up: false,
ipv6: None,
ipv4: None,
}));
assert!(!state.interfaces.contains_key("en2"));
}
#[test]
fn neighbour_add_and_flush() {
let mut state = BabelState::new();
state.apply(Event::Neighbour(NeighbourEvent {
kind: EventKind::Add,
handle: 0xabc,
address: IpAddr::V6("fe80::1".parse().unwrap()),
ifname: "en3".into(),
reach: 0x00ff,
ureach: 0x000f,
rxcost: 96,
txcost: 128,
rtt_millis: Some(42),
rttcost: Some(10),
external_bias_256: 4096,
external_coef_256: 128,
cost: 224,
}));
assert_eq!(state.neighbours.len(), 1);
assert_eq!(state.neighbours.get(&0xabc).unwrap().ifname.as_ref(), "en3");
assert_eq!(
state.neighbours.get(&0xabc).unwrap().external_bias_256,
4096
);
assert_eq!(state.neighbours.get(&0xabc).unwrap().external_coef_256, 128);
state.apply(Event::Neighbour(NeighbourEvent {
kind: EventKind::Flush,
handle: 0xabc,
address: IpAddr::V6("fe80::1".parse().unwrap()),
ifname: "en3".into(),
reach: 0,
ureach: 0,
rxcost: 0,
txcost: 0,
rtt_millis: None,
rttcost: None,
external_bias_256: 0,
external_coef_256: 256,
cost: 0,
}));
assert!(state.neighbours.is_empty());
}
#[test]
fn xroute_change_upserts_by_prefix_pair() {
let mut state = BabelState::new();
state.apply(Event::XRoute(XRouteEvent {
kind: EventKind::Add,
prefix: net("fde0:20c6:1fa7:ffff::/128"),
from: net("::/0"),
metric: 256,
}));
state.apply(Event::XRoute(XRouteEvent {
kind: EventKind::Change,
prefix: net("fde0:20c6:1fa7:ffff::/128"),
from: net("::/0"),
metric: 42,
}));
assert_eq!(state.xroutes.len(), 1);
assert_eq!(
state
.xroutes
.get(&XRouteKey {
prefix: net("fde0:20c6:1fa7:ffff::/128"),
from: net("::/0"),
})
.unwrap()
.metric,
42
);
}
#[test]
fn route_add_and_flush_by_handle() {
let mut state = BabelState::new();
state.apply(Event::Route(RouteEvent {
kind: EventKind::Add,
handle: 0xdeadbeef,
prefix: net("fde0:20c6:1fa7:ffff::/128"),
from: net("::/0"),
installed: true,
id: Eui64::new(0, 1, 2, 3, 4, 5, 6, 7),
metric: 96,
refmetric: 96,
via: IpAddr::V6("fe80::1234".parse().unwrap()),
ifname: "en2".into(),
}));
assert_eq!(state.routes.len(), 1);
assert!(state.routes.get(&0xdeadbeef).unwrap().installed);
state.apply(Event::Route(RouteEvent {
kind: EventKind::Flush,
handle: 0xdeadbeef,
prefix: net("fde0:20c6:1fa7:ffff::/128"),
from: net("::/0"),
installed: false,
id: Eui64::new(0, 1, 2, 3, 4, 5, 6, 7),
metric: 0,
refmetric: 0,
via: IpAddr::V6("fe80::1234".parse().unwrap()),
ifname: "en2".into(),
}));
assert!(state.routes.is_empty());
}
#[test]
fn extend_applies_multiple_events() {
let mut state = BabelState::new();
state.extend([
Event::Interface(InterfaceEvent {
kind: EventKind::Add,
ifname: "en2".into(),
up: true,
ipv6: None,
ipv4: None,
}),
Event::XRoute(XRouteEvent {
kind: EventKind::Add,
prefix: net("fde0:20c6:1fa7:ffff::/128"),
from: net("::/0"),
metric: 123,
}),
]);
assert_eq!(state.interfaces.len(), 1);
assert_eq!(state.xroutes.len(), 1);
}
}
-301
View File
@@ -1,301 +0,0 @@
//! Process configuration and shared constants for `babblerd`.
//!
//! This module centralizes:
//!
//! - default runtime paths
//! - environment variable overrides
//! - protocol/application constants such as the mesh prefix
//! - coarse daemon defaults such as the router UDP port
use color_eyre::eyre::{self, eyre};
use ipnet::Ipv6Net;
use std::collections::HashSet;
use std::env;
use std::fmt::{Display, Formatter};
use std::net::Ipv6Addr;
use std::path::{Path, PathBuf};
use std::str::FromStr;
pub const PUBLIC_SOCKET_PATH_ENV: &str = "BABBLER_SOCKET_PATH";
pub const NODE_ID_FILE_ENV: &str = "BABBLER_NODE_ID_FILE";
pub const ROUTER_UDP_PORT_ENV: &str = "BABBLER_ROUTER_UDP_PORT";
pub const ROUTER_TRANSPORT_ENV: &str = "BABBLER_ROUTER_TRANSPORT";
pub const TUN_MTU_ENV: &str = "BABBLER_TUN_MTU";
pub const TCP_BATCH_TARGET_BYTES_ENV: &str = "BABBLER_TCP_BATCH_TARGET_BYTES";
pub const TCP_SOCKET_BUFFER_BYTES_ENV: &str = "BABBLER_TCP_SOCKET_BUFFER_BYTES";
pub const INTERFACE_ALLOWLIST_ENV: &str = "BABBLER_INTERFACE_ALLOWLIST";
pub const DEFAULT_PUBLIC_SOCKET_PATH: &str = {
#[cfg(target_os = "macos")]
{
"/var/run/babbler/babblerd.sock"
}
#[cfg(target_os = "linux")]
{
"/run/babbler/babblerd.sock"
}
};
pub const DEFAULT_NODE_ID_FILE: &str = {
#[cfg(target_os = "macos")]
{
"/var/db/babbler/node-id"
}
#[cfg(target_os = "linux")]
{
"/var/lib/babbler/node-id"
}
};
// TODO: just picked a random one that didn't seem occupied, there is probably a better way
// to do this in the future :)
pub const DEFAULT_ROUTER_UDP_PORT: u16 = 41897;
pub const PHYSICAL_LINK_MTU: u16 = 1500;
pub const OUTER_IPV6_HEADER_BYTES: u16 = 40;
pub const OUTER_UDP_HEADER_BYTES: u16 = 8;
pub const UDP_TUN_MTU: u16 = PHYSICAL_LINK_MTU - OUTER_IPV6_HEADER_BYTES - OUTER_UDP_HEADER_BYTES;
pub const TCP_TUN_MTU: u16 = u16::MAX;
pub const MIN_TUN_MTU: u16 = 1280;
pub const MAX_TUN_MTU: u16 = u16::MAX;
pub const TUN_MTU: u16 = UDP_TUN_MTU;
pub const DEFAULT_TCP_BATCH_TARGET_BYTES: usize = 256 * 1024;
pub const TCP_PENDING_LIMIT_BYTES: usize = 4 * 1024 * 1024;
pub const DEFAULT_TCP_SOCKET_BUFFER_BYTES: usize = 4 * 1024 * 1024;
pub const MAX_TCP_SOCKET_BUFFER_BYTES: usize = 512 * 1024 * 1024;
pub const EXO_ULA_PREFIX: Ipv6Net = Ipv6Net::new_assert(
// TODO: break out into "fd" for ULA
// e0_20c61fa7 for EXO address-space
// ffff for anything else we want, like maybe versioning and so on (but for now its not used)
//
// NOTE: spell the hextets explicitly here. A previous `u128` bit-shift
// construction accidentally truncated the leading `fde0` and produced
// `20c6:1fa7:ffff::/64`, which is not ULA.
Ipv6Addr::new(0xfde0, 0x20c6, 0x1fa7, 0xffff, 0, 0, 0, 0),
64,
);
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum TransportMode {
Udp,
Tcp,
}
impl Display for TransportMode {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
match self {
Self::Udp => write!(f, "udp"),
Self::Tcp => write!(f, "tcp"),
}
}
}
impl FromStr for TransportMode {
type Err = String;
fn from_str(value: &str) -> Result<Self, Self::Err> {
match value.trim().to_ascii_lowercase().as_str() {
"udp" => Ok(Self::Udp),
"tcp" => Ok(Self::Tcp),
other => Err(format!("expected udp or tcp, got {other:?}")),
}
}
}
#[derive(Debug, Clone)]
pub struct Config {
pub public_socket_path: PathBuf,
pub public_dir: PathBuf,
pub node_id_file: PathBuf,
pub router_udp_port: u16,
pub router_transport: TransportMode,
pub tun_mtu: u16,
pub tcp_batch_target_bytes: usize,
pub tcp_socket_buffer_bytes: usize,
pub exo_ula_prefix: Ipv6Net,
}
impl Config {
pub fn from_env() -> eyre::Result<Self> {
let public_socket_path = env::var_os(PUBLIC_SOCKET_PATH_ENV)
.map(PathBuf::from)
.unwrap_or_else(|| PathBuf::from(DEFAULT_PUBLIC_SOCKET_PATH));
let Some(public_dir) = public_socket_path.parent().map(Path::to_path_buf) else {
return Err(eyre!(
"public socket path has no parent directory: {}",
public_socket_path.display()
));
};
let node_id_file = env::var_os(NODE_ID_FILE_ENV)
.map(PathBuf::from)
.unwrap_or_else(|| PathBuf::from(DEFAULT_NODE_ID_FILE));
let router_udp_port = match env::var(ROUTER_UDP_PORT_ENV) {
Ok(raw) => raw
.parse::<u16>()
.map_err(|e| eyre!("invalid {ROUTER_UDP_PORT_ENV} value {raw:?}: {e}"))?,
Err(_) => DEFAULT_ROUTER_UDP_PORT,
};
let router_transport = match env::var(ROUTER_TRANSPORT_ENV) {
Ok(raw) => raw
.parse::<TransportMode>()
.map_err(|e| eyre!("invalid {ROUTER_TRANSPORT_ENV} value {raw:?}: {e}"))?,
Err(_) => TransportMode::Udp,
};
let tun_mtu = match env::var(TUN_MTU_ENV) {
Ok(raw) => parse_tun_mtu(&raw)
.map_err(|e| eyre!("invalid {TUN_MTU_ENV} value {raw:?}: {e}"))?,
Err(_) => default_tun_mtu(router_transport),
};
let tcp_batch_target_bytes = match env::var(TCP_BATCH_TARGET_BYTES_ENV) {
Ok(raw) => parse_tcp_batch_target_bytes(&raw)
.map_err(|e| eyre!("invalid {TCP_BATCH_TARGET_BYTES_ENV} value {raw:?}: {e}"))?,
Err(_) => DEFAULT_TCP_BATCH_TARGET_BYTES,
};
let tcp_socket_buffer_bytes = match env::var(TCP_SOCKET_BUFFER_BYTES_ENV) {
Ok(raw) => parse_tcp_socket_buffer_bytes(&raw)
.map_err(|e| eyre!("invalid {TCP_SOCKET_BUFFER_BYTES_ENV} value {raw:?}: {e}"))?,
Err(_) => DEFAULT_TCP_SOCKET_BUFFER_BYTES,
};
Ok(Self {
public_socket_path,
public_dir,
node_id_file,
router_udp_port,
router_transport,
tun_mtu,
tcp_batch_target_bytes,
tcp_socket_buffer_bytes,
exo_ula_prefix: EXO_ULA_PREFIX,
})
}
}
pub fn default_tun_mtu(transport: TransportMode) -> u16 {
match transport {
TransportMode::Udp => UDP_TUN_MTU,
TransportMode::Tcp => TCP_TUN_MTU,
}
}
pub fn parse_tun_mtu(raw: &str) -> Result<u16, String> {
let mtu = raw
.trim()
.parse::<u16>()
.map_err(|err| format!("expected integer MTU: {err}"))?;
if !(MIN_TUN_MTU..=MAX_TUN_MTU).contains(&mtu) {
return Err(format!(
"expected MTU in {MIN_TUN_MTU}..={MAX_TUN_MTU}, got {mtu}"
));
}
Ok(mtu)
}
pub fn parse_tcp_batch_target_bytes(raw: &str) -> Result<usize, String> {
parse_usize_in_range(raw, 1024, TCP_PENDING_LIMIT_BYTES, "TCP batch target bytes")
}
pub fn parse_tcp_socket_buffer_bytes(raw: &str) -> Result<usize, String> {
parse_usize_in_range(
raw,
1024,
MAX_TCP_SOCKET_BUFFER_BYTES,
"TCP socket buffer bytes",
)
}
fn parse_usize_in_range(raw: &str, min: usize, max: usize, name: &str) -> Result<usize, String> {
let value = raw
.trim()
.parse::<usize>()
.map_err(|err| format!("expected integer {name}: {err}"))?;
if !(min..=max).contains(&value) {
return Err(format!("expected {name} in {min}..={max}, got {value}"));
}
Ok(value)
}
pub fn interface_allowlist_from_env() -> eyre::Result<Option<HashSet<Box<str>>>> {
let Ok(raw) = env::var(INTERFACE_ALLOWLIST_ENV) else {
return Ok(None);
};
let allowlist = raw
.split(',')
.map(str::trim)
.filter(|name| !name.is_empty())
.map(|name| name.into())
.collect::<HashSet<Box<str>>>();
if allowlist.is_empty() {
return Err(eyre!(
"{INTERFACE_ALLOWLIST_ENV} was set but contained no interface names"
));
}
Ok(Some(allowlist))
}
#[cfg(test)]
mod tests {
use super::{
DEFAULT_TCP_BATCH_TARGET_BYTES, DEFAULT_TCP_SOCKET_BUFFER_BYTES, EXO_ULA_PREFIX,
TCP_TUN_MTU, TransportMode, UDP_TUN_MTU, default_tun_mtu, parse_tcp_batch_target_bytes,
parse_tcp_socket_buffer_bytes, parse_tun_mtu,
};
use std::net::Ipv6Addr;
#[test]
fn exo_ula_prefix_keeps_fde0_high_bits() {
assert_eq!(
EXO_ULA_PREFIX.addr(),
Ipv6Addr::new(0xfde0, 0x20c6, 0x1fa7, 0xffff, 0, 0, 0, 0)
);
assert_eq!(EXO_ULA_PREFIX.prefix_len(), 64);
}
#[test]
fn transport_mode_parses_udp_and_tcp() {
assert_eq!("udp".parse::<TransportMode>().unwrap(), TransportMode::Udp);
assert_eq!("tcp".parse::<TransportMode>().unwrap(), TransportMode::Tcp);
assert_eq!("TCP".parse::<TransportMode>().unwrap(), TransportMode::Tcp);
assert!("quic".parse::<TransportMode>().is_err());
}
#[test]
fn tun_mtu_defaults_follow_transport_and_validate_bounds() {
assert_eq!(default_tun_mtu(TransportMode::Udp), UDP_TUN_MTU);
assert_eq!(default_tun_mtu(TransportMode::Tcp), TCP_TUN_MTU);
assert_eq!(parse_tun_mtu("9000").unwrap(), 9000);
assert_eq!(parse_tun_mtu("65535").unwrap(), 65535);
assert!(parse_tun_mtu("1279").is_err());
assert!(parse_tun_mtu("65536").is_err());
}
#[test]
fn tcp_tuning_knobs_validate_bounds() {
assert_eq!(
parse_tcp_batch_target_bytes(&DEFAULT_TCP_BATCH_TARGET_BYTES.to_string()).unwrap(),
DEFAULT_TCP_BATCH_TARGET_BYTES
);
assert_eq!(
parse_tcp_batch_target_bytes("2097152").unwrap(),
2 * 1024 * 1024
);
assert!(parse_tcp_batch_target_bytes("512").is_err());
assert!(parse_tcp_batch_target_bytes("4194305").is_err());
assert_eq!(
parse_tcp_socket_buffer_bytes(&DEFAULT_TCP_SOCKET_BUFFER_BYTES.to_string()).unwrap(),
DEFAULT_TCP_SOCKET_BUFFER_BYTES
);
assert_eq!(
parse_tcp_socket_buffer_bytes("33554432").unwrap(),
32 * 1024 * 1024
);
assert!(parse_tcp_socket_buffer_bytes("512").is_err());
}
}
-398
View File
@@ -1,398 +0,0 @@
use color_eyre::eyre::{Result, WrapErr, eyre};
use ipnet::Ipv6Net;
use std::fmt::{Display, Formatter};
use std::{future::pending, sync::Arc};
use tokio::{
io::{AsyncBufReadExt, AsyncWriteExt, BufReader},
net::UnixStream,
sync::{mpsc, oneshot, watch},
task::JoinHandle,
time::{Duration, Instant},
};
use crate::config::TransportMode;
use crate::route_ctl;
use crate::routing_stack::RoutingStack;
use crate::{babel::BabelState, tun::TunDevice};
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ServiceState {
Off,
Starting,
On,
Stopping,
}
impl Display for ServiceState {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
match self {
Self::Off => write!(f, "off"),
Self::Starting => write!(f, "starting"),
Self::On => write!(f, "on"),
Self::Stopping => write!(f, "stopping"),
}
}
}
#[derive(Debug, Clone)]
pub struct DaemonStatus {
pub service_state: ServiceState,
pub node_id: u64,
pub node_addr: Ipv6Net,
pub tun_ifname: Arc<str>,
// realistically should always have one?? right??
pub keepalive_deadline: Option<Instant>,
pub last_error: Option<Arc<str>>,
}
impl DaemonStatus {
fn new(node_id: u64, node_addr: Ipv6Net, tun_ifname: Arc<str>) -> Self {
Self {
service_state: ServiceState::Off,
node_id,
node_addr,
tun_ifname,
keepalive_deadline: None,
last_error: None,
}
}
pub fn render(&self) -> String {
let keepalive_remaining_ms = self
.keepalive_deadline
.and_then(|deadline| deadline.checked_duration_since(Instant::now()))
.map(|remaining| remaining.as_millis().to_string())
.unwrap_or_else(|| "none".to_owned());
let mut line = format!(
"state {} node_id={:#018x} node_addr={} tun={} keepalive_remaining_ms={}",
self.service_state,
self.node_id,
self.node_addr,
self.tun_ifname,
keepalive_remaining_ms
);
if let Some(err) = &self.last_error {
line.push_str(&format!(" last_error={err:?}"));
}
line
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum StackTaskKind {
Babel,
Watcher,
FibPublisher,
Dataplane,
}
#[derive(Debug)]
pub enum RoutingStackEvent {
Exited {
kind: StackTaskKind,
error: Option<String>,
},
}
#[derive(Debug)]
enum DaemonCommand {
KeepAlive {
ttl: Duration,
reply: oneshot::Sender<Result<DaemonStatus>>,
},
GetState {
reply: oneshot::Sender<DaemonStatus>,
},
}
#[derive(Clone)]
pub struct DaemonHandle {
send: mpsc::Sender<DaemonCommand>,
}
impl DaemonHandle {
pub async fn keep_alive(&self, ttl: Duration) -> Result<DaemonStatus> {
let (reply_send, reply_recv) = oneshot::channel();
self.send
.send(DaemonCommand::KeepAlive {
ttl,
reply: reply_send,
})
.await
.map_err(|_| eyre!("daemon core stopped"))?;
reply_recv.await.map_err(|_| eyre!("daemon core stopped"))?
}
pub async fn get_state(&self) -> Result<DaemonStatus> {
let (reply_send, reply_recv) = oneshot::channel();
self.send
.send(DaemonCommand::GetState { reply: reply_send })
.await
.map_err(|_| eyre!("daemon core stopped"))?;
reply_recv.await.map_err(|_| eyre!("daemon core stopped"))
}
}
pub struct DaemonCore {
status: DaemonStatus,
overlay_prefix: Ipv6Net,
router_udp_port: u16,
router_transport: TransportMode,
tun_mtu: u16,
tcp_batch_target_bytes: usize,
tcp_socket_buffer_bytes: usize,
_tun: TunDevice,
routing_stack: Option<RoutingStack>,
babel_state_send: watch::Sender<Arc<BabelState>>,
command_recv: mpsc::Receiver<DaemonCommand>,
event_send: mpsc::Sender<RoutingStackEvent>,
event_recv: mpsc::Receiver<RoutingStackEvent>,
}
impl DaemonCore {
pub fn spawn(
node_id: u64,
overlay_prefix: Ipv6Net,
router_udp_port: u16,
router_transport: TransportMode,
tun_mtu: u16,
tcp_batch_target_bytes: usize,
tcp_socket_buffer_bytes: usize,
node_addr: Ipv6Net,
tun: TunDevice,
babel_state_send: watch::Sender<Arc<BabelState>>,
) -> (DaemonHandle, JoinHandle<Result<()>>) {
let (command_send, command_recv) = mpsc::channel(32);
let (event_send, event_recv) = mpsc::channel(8);
let status = DaemonStatus::new(node_id, node_addr, Arc::from(tun.ifname().to_owned()));
let core = Self {
status,
overlay_prefix,
router_udp_port,
router_transport,
tun_mtu,
tcp_batch_target_bytes,
tcp_socket_buffer_bytes,
_tun: tun,
routing_stack: None,
babel_state_send,
command_recv,
event_send,
event_recv,
};
let handle = DaemonHandle { send: command_send };
let task = tokio::spawn(core.run());
(handle, task)
}
async fn run(mut self) -> Result<()> {
loop {
tokio::select! {
command = self.command_recv.recv() => {
let Some(command) = command else {
break;
};
self.handle_command(command).await?;
}
event = self.event_recv.recv() => {
let Some(event) = event else {
break;
};
self.handle_stack_event(event).await?;
}
_ = lease_timer(self.status.keepalive_deadline) => {
self.handle_lease_expiry().await?;
}
}
}
self.stop_stack().await?;
Ok(())
}
async fn handle_command(&mut self, command: DaemonCommand) -> Result<()> {
match command {
DaemonCommand::KeepAlive { ttl, reply } => {
self.status.keepalive_deadline = Some(Instant::now() + ttl);
if self.routing_stack.is_none() {
let result = self.start_stack().await.map(|()| self.status.clone());
let _ = reply.send(result);
return Ok(());
}
let _ = reply.send(Ok(self.status.clone()));
Ok(())
}
DaemonCommand::GetState { reply } => {
let _ = reply.send(self.status.clone());
Ok(())
}
}
}
async fn handle_stack_event(&mut self, event: RoutingStackEvent) -> Result<()> {
let RoutingStackEvent::Exited { kind, error } = event;
if self.routing_stack.is_none() {
return Ok(());
}
tracing::warn!(?kind, ?error, "routing stack task exited");
self.status.last_error =
Some(Arc::from(error.unwrap_or_else(|| {
format!("{kind:?} task exited unexpectedly")
})));
if let Err(err) = self.stop_stack().await {
self.status.last_error = Some(Arc::from(err.to_string()));
}
Ok(())
}
async fn handle_lease_expiry(&mut self) -> Result<()> {
let expired = self
.status
.keepalive_deadline
.is_some_and(|deadline| deadline <= Instant::now());
if expired {
tracing::info!("keepalive expired, transitioning routing stack off");
self.status.keepalive_deadline = None;
if let Err(err) = self.stop_stack().await {
self.status.last_error = Some(Arc::from(err.to_string()));
}
}
Ok(())
}
async fn start_stack(&mut self) -> Result<()> {
if self.routing_stack.is_some() {
return Ok(());
}
self.status.service_state = ServiceState::Starting;
self.status.last_error = None;
match RoutingStack::start(
self.status.node_addr,
&self._tun,
self.router_udp_port,
self.router_transport,
self.tun_mtu,
self.tcp_batch_target_bytes,
self.tcp_socket_buffer_bytes,
self.babel_state_send.clone(),
self.event_send.clone(),
) {
Ok(stack) => {
self.routing_stack = Some(stack);
if let Err(err) = route_ctl::ensure_overlay_route(
self.overlay_prefix,
self.status.tun_ifname.as_ref(),
) {
let _ = self.stop_stack().await;
self.status.service_state = ServiceState::Off;
self.status.last_error = Some(Arc::from(err.to_string()));
return Err(err.into());
}
self.status.service_state = ServiceState::On;
Ok(())
}
Err(err) => {
self.status.service_state = ServiceState::Off;
self.status.last_error = Some(Arc::from(err.to_string()));
Err(err)
}
}
}
async fn stop_stack(&mut self) -> Result<()> {
if let Err(err) = route_ctl::remove_overlay_route(self.overlay_prefix) {
tracing::warn!(error=%err, "failed to remove overlay route");
}
let Some(stack) = self.routing_stack.take() else {
self.status.service_state = ServiceState::Off;
self.babel_state_send
.send_replace(Arc::new(BabelState::new()));
return Ok(());
};
self.status.service_state = ServiceState::Stopping;
let stop_result = stack.stop().await;
self.babel_state_send
.send_replace(Arc::new(BabelState::new()));
self.status.service_state = ServiceState::Off;
if let Err(err) = stop_result {
self.status.last_error = Some(Arc::from(err.to_string()));
return Err(err);
}
Ok(())
}
}
async fn lease_timer(deadline: Option<Instant>) {
if let Some(deadline) = deadline {
tokio::time::sleep_until(deadline).await;
} else {
pending::<()>().await;
}
}
pub async fn handle_client(sock: UnixStream, daemon: DaemonHandle) {
tracing::info!("new socket conn");
let (reader, mut write) = sock.into_split();
let mut reader = BufReader::new(reader).lines();
if let Ok(state) = daemon.get_state().await {
let _ = write
.write_all(format!("{}\n", state.render()).as_bytes())
.await;
}
loop {
let Ok(Some(line)) = reader.next_line().await else {
break;
};
let trimmed = line.trim();
if trimmed.is_empty() {
continue;
}
let response = match handle_command_line(trimmed, &daemon).await {
Ok(response) => response,
Err(err) => format!("error {err}"),
};
if let Err(err) = write.write_all(format!("{response}\n").as_bytes()).await {
tracing::warn!(error=%err, "failed to write command response");
break;
}
}
tracing::info!("closing socket conn");
let _ = write.shutdown().await;
}
async fn handle_command_line(line: &str, daemon: &DaemonHandle) -> Result<String> {
let mut parts = line.split_whitespace();
let Some(command) = parts.next() else {
return Ok("error empty-command".to_owned());
};
match command {
"get-state" => Ok(daemon.get_state().await?.render()),
"keepalive" => {
let Some(ttl_ms) = parts.next() else {
return Err(eyre!("keepalive requires ttl_ms"));
};
let ttl_ms = ttl_ms
.parse::<u64>()
.wrap_err_with(|| format!("invalid ttl_ms: {ttl_ms:?}"))?;
Ok(daemon
.keep_alive(Duration::from_millis(ttl_ms))
.await?
.render())
}
"help" => Ok("commands: get-state | keepalive <ttl_ms>".to_owned()),
other => Err(eyre!("unknown command: {other}")),
}
}
File diff suppressed because it is too large. Load diff
-430
View File
@@ -1,430 +0,0 @@
//! Immutable forwarding snapshots derived from [`crate::babel::BabelState`].
//!
//! `BabelState` mirrors the control-plane view emitted by `babeld`.
//! `FibSnapshot` is the reduced dataplane view:
//!
//! - exact-match IPv6 host routes only for now,
//! - admitted interface ownership alongside those routes,
//! - one immutable snapshot swapped wholesale into the dataplane,
//! - keyed for fast lookup rather than protocol fidelity.
//!
//! The v1 forwarding model is intentionally narrow:
//!
//! - local addresses are explicit inputs, not inferred from every xroute,
//! - only interfaces with a live Babel neighbour are exposed to the dataplane,
//! - only installed IPv6 `/128` routes are considered,
//! - only destination-based forwarding is modeled,
//! - routes with non-link-local next hops are ignored.
use std::net::{IpAddr, Ipv6Addr};
use ahash::RandomState;
use hashbrown::{HashMap, HashSet, hash_map::Entry};
use ipnet::{IpNet, Ipv6Net};
use crate::babel::BabelState;
use crate::babel::state::RouteState;
pub type HostKey = u128;
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct FibEntry {
pub next_hop_ll: Ipv6Addr,
pub ifname: Box<str>,
pub mtu: u16,
}
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct AdmittedNeighbour {
pub ifname: Box<str>,
pub link_local: Ipv6Addr,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct FibSnapshot {
pub locals: HashSet<HostKey, RandomState>,
pub admitted_interfaces: HashSet<Box<str>, RandomState>,
pub admitted_neighbours: HashSet<AdmittedNeighbour, RandomState>,
pub interface_link_locals: HashMap<Box<str>, Ipv6Addr, RandomState>,
pub routes: HashMap<HostKey, FibEntry, RandomState>,
}
#[derive(Debug, Clone)]
pub struct FibBuilder {
local_addrs: Vec<Ipv6Addr>,
route_mtu: u16,
}
impl FibBuilder {
pub fn new<I>(local_addrs: I, route_mtu: u16) -> Self
where
I: IntoIterator<Item = Ipv6Addr>,
{
Self {
local_addrs: local_addrs.into_iter().collect(),
route_mtu,
}
}
pub fn derive(&self, state: &BabelState) -> FibSnapshot {
let mut locals =
HashSet::with_capacity_and_hasher(self.local_addrs.len(), RandomState::new());
for addr in &self.local_addrs {
locals.insert(host_key(*addr));
}
let up_interfaces: HashSet<&str, RandomState> = state
.interfaces
.values()
.filter(|interface| interface.up)
.map(|interface| interface.ifname.as_ref())
.collect::<HashSet<_, _>>();
let mut interface_link_locals = HashMap::with_hasher(RandomState::new());
for interface in state.interfaces.values().filter(|interface| interface.up) {
let Some(IpAddr::V6(link_local)) = interface.ipv6 else {
continue;
};
if link_local.is_unicast_link_local() {
interface_link_locals.insert(interface.ifname.clone(), link_local);
}
}
let mut admitted_interfaces = HashSet::with_hasher(RandomState::new());
let mut admitted_neighbours = HashSet::with_hasher(RandomState::new());
for neighbour in state.neighbours.values() {
if !up_interfaces.contains(neighbour.ifname.as_ref()) {
continue;
}
let IpAddr::V6(link_local) = neighbour.address else {
continue;
};
if !link_local.is_unicast_link_local() {
continue;
}
admitted_interfaces.insert(neighbour.ifname.clone());
admitted_neighbours.insert(AdmittedNeighbour {
ifname: neighbour.ifname.clone(),
link_local,
});
}
let mut routes = HashMap::with_hasher(RandomState::new());
let mut route_scores = HashMap::with_hasher(RandomState::new());
let mut candidates: Vec<&RouteState> = state.routes.values().collect();
candidates.sort_by(|left, right| {
left.ifname
.cmp(&right.ifname)
.then_with(|| left.prefix.to_string().cmp(&right.prefix.to_string()))
.then_with(|| left.metric.cmp(&right.metric))
.then_with(|| left.refmetric.cmp(&right.refmetric))
.then_with(|| left.handle.cmp(&right.handle))
});
for route in candidates {
let Some((dst, next_hop_ll)) = route_to_host(route) else {
continue;
};
if !admitted_interfaces.contains(route.ifname.as_ref()) {
continue;
}
if locals.contains(&dst) {
continue;
}
let candidate = FibEntry {
next_hop_ll,
ifname: route.ifname.clone(),
mtu: self.route_mtu,
};
let candidate_score = (route.metric, route.refmetric, route.handle);
match routes.entry(dst) {
Entry::Vacant(slot) => {
slot.insert(candidate);
route_scores.insert(dst, candidate_score);
}
Entry::Occupied(mut slot) => {
let Some(existing_score) = route_scores.get(&dst).copied() else {
slot.insert(candidate);
route_scores.insert(dst, candidate_score);
continue;
};
if candidate_score < existing_score {
slot.insert(candidate);
route_scores.insert(dst, candidate_score);
}
}
}
}
FibSnapshot {
locals,
admitted_interfaces,
admitted_neighbours,
interface_link_locals,
routes,
}
}
}
impl FibSnapshot {
pub fn empty() -> Self {
Self {
locals: HashSet::with_hasher(RandomState::new()),
admitted_interfaces: HashSet::with_hasher(RandomState::new()),
admitted_neighbours: HashSet::with_hasher(RandomState::new()),
interface_link_locals: HashMap::with_hasher(RandomState::new()),
routes: HashMap::with_hasher(RandomState::new()),
}
}
pub fn from_node_addr(node_addr: Ipv6Net, state: &BabelState, route_mtu: u16) -> Self {
FibBuilder::new([node_addr.addr()], route_mtu).derive(state)
}
pub fn is_local(&self, addr: Ipv6Addr) -> bool {
self.locals.contains(&host_key(addr))
}
pub fn lookup(&self, addr: Ipv6Addr) -> Option<&FibEntry> {
self.routes.get(&host_key(addr))
}
}
pub fn host_key(addr: Ipv6Addr) -> HostKey {
u128::from(addr)
}
fn route_to_host(route: &RouteState) -> Option<(HostKey, Ipv6Addr)> {
if !route.installed {
return None;
}
let IpNet::V6(prefix) = route.prefix else {
return None;
};
if prefix.prefix_len() != 128 {
return None;
}
let IpNet::V6(from) = route.from else {
return None;
};
if from.prefix_len() != 0 || from.addr() != Ipv6Addr::UNSPECIFIED {
return None;
}
let std::net::IpAddr::V6(via) = route.via else {
return None;
};
if !via.is_unicast_link_local() {
return None;
}
Some((host_key(prefix.addr()), via))
}
#[cfg(test)]
mod tests {
use std::net::{IpAddr, Ipv6Addr};
use crate::babel::Eui64;
use crate::babel::line::{Event, EventKind, InterfaceEvent, NeighbourEvent, RouteEvent};
use crate::babel::state::BabelState;
use super::{AdmittedNeighbour, FibBuilder};
fn route(
handle: u64,
prefix: &str,
from: &str,
installed: bool,
via: Ipv6Addr,
ifname: &str,
metric: u32,
refmetric: u32,
) -> Event {
Event::Route(RouteEvent {
kind: EventKind::Add,
handle,
prefix: prefix.parse().unwrap(),
from: from.parse().unwrap(),
installed,
id: Eui64::new(0, 1, 2, 3, 4, 5, 6, 7),
metric,
refmetric,
via: IpAddr::V6(via),
ifname: ifname.into(),
})
}
fn interface(ifname: &str, up: bool) -> Event {
Event::Interface(InterfaceEvent {
kind: EventKind::Add,
ifname: ifname.into(),
up,
ipv6: None,
ipv4: None,
})
}
fn neighbour(handle: u64, ifname: &str, address: &str) -> Event {
Event::Neighbour(NeighbourEvent {
kind: EventKind::Add,
handle,
address: IpAddr::V6(address.parse().unwrap()),
ifname: ifname.into(),
reach: 0xffff,
ureach: 0xffff,
rxcost: 96,
txcost: 96,
rtt_millis: Some(1),
rttcost: Some(0),
external_bias_256: 0,
external_coef_256: 256,
cost: 96,
})
}
#[test]
fn derives_local_and_host_routes() {
let mut state = BabelState::new();
state.apply(interface("en2", true));
state.apply(neighbour(1, "en2", "fe80::1"));
state.apply(route(
1,
"fde0::1234/128",
"::/0",
true,
"fe80::1".parse().unwrap(),
"en2",
96,
32,
));
let fib = FibBuilder::new(["fde0::1".parse().unwrap()], 1452).derive(&state);
assert!(fib.is_local("fde0::1".parse().unwrap()));
assert!(fib.admitted_interfaces.contains("en2"));
assert!(fib.admitted_neighbours.contains(&AdmittedNeighbour {
ifname: "en2".into(),
link_local: "fe80::1".parse().unwrap(),
}));
let entry = fib.lookup("fde0::1234".parse().unwrap()).unwrap();
assert_eq!(entry.next_hop_ll, "fe80::1".parse::<Ipv6Addr>().unwrap());
assert_eq!(entry.ifname.as_ref(), "en2");
assert_eq!(entry.mtu, 1452);
}
#[test]
fn skips_non_installed_or_non_host_routes() {
let mut state = BabelState::new();
state.apply(interface("en2", true));
state.apply(interface("en3", true));
state.apply(neighbour(1, "en2", "fe80::1"));
state.apply(neighbour(2, "en3", "fe80::2"));
state.apply(route(
1,
"fde0::abcd/128",
"::/0",
false,
"fe80::1".parse().unwrap(),
"en2",
96,
32,
));
state.apply(route(
2,
"fde0::/64",
"::/0",
true,
"fe80::2".parse().unwrap(),
"en3",
96,
32,
));
let fib = FibBuilder::new(["fde0::1".parse().unwrap()], 1452).derive(&state);
assert!(fib.routes.is_empty());
}
#[test]
fn prefers_lower_metric_when_multiple_installed_routes_exist() {
let mut state = BabelState::new();
state.apply(interface("en2", true));
state.apply(interface("en3", true));
state.apply(neighbour(1, "en2", "fe80::1"));
state.apply(neighbour(2, "en3", "fe80::2"));
state.apply(route(
1,
"fde0::beef/128",
"::/0",
true,
"fe80::1".parse().unwrap(),
"en2",
200,
20,
));
state.apply(route(
2,
"fde0::beef/128",
"::/0",
true,
"fe80::2".parse().unwrap(),
"en3",
100,
10,
));
let fib = FibBuilder::new(["fde0::1".parse().unwrap()], 1452).derive(&state);
let entry = fib.lookup("fde0::beef".parse().unwrap()).unwrap();
assert_eq!(entry.next_hop_ll, "fe80::2".parse::<Ipv6Addr>().unwrap());
assert_eq!(entry.ifname.as_ref(), "en3");
}
#[test]
fn skips_routes_for_interfaces_without_live_neighbours() {
let mut state = BabelState::new();
state.apply(interface("en2", true));
state.apply(route(
1,
"fde0::cafe/128",
"::/0",
true,
"fe80::1".parse().unwrap(),
"en2",
96,
32,
));
let fib = FibBuilder::new(["fde0::1".parse().unwrap()], 1452).derive(&state);
assert!(fib.admitted_interfaces.is_empty());
assert!(fib.lookup("fde0::cafe".parse().unwrap()).is_none());
}
#[test]
fn skips_neighbour_interfaces_that_are_not_up() {
let mut state = BabelState::new();
state.apply(interface("en2", false));
state.apply(neighbour(1, "en2", "fe80::1"));
state.apply(route(
1,
"fde0::cafe/128",
"::/0",
true,
"fe80::1".parse().unwrap(),
"en2",
96,
32,
));
let fib = FibBuilder::new(["fde0::1".parse().unwrap()], 1452).derive(&state);
assert!(fib.admitted_interfaces.is_empty());
assert!(fib.lookup("fde0::cafe".parse().unwrap()).is_none());
}
}
-154
View File
@@ -1,154 +0,0 @@
//! Persistent node identity for `babblerd`.
//!
//! The node ID occupies the full low 64 bits of the EXO ULA space.
use std::fs::{self, OpenOptions};
use std::io::Write;
use std::os::unix::fs::{MetadataExt, PermissionsExt};
use std::path::Path;
use color_eyre::eyre::{self, WrapErr, eyre};
use ipnet::Ipv6Net;
use nix::unistd::geteuid;
use std::net::Ipv6Addr;
pub fn load_or_create_node_id(path: &Path) -> eyre::Result<u64> {
match read_node_id(path) {
Ok(node_id) => Ok(node_id),
Err(err) if is_not_found(&err) => create_node_id(path),
Err(err) => Err(err),
}
}
pub fn node_addr(prefix: Ipv6Net, node_id: u64) -> eyre::Result<Ipv6Net> {
if prefix.prefix_len() != 64 {
return Err(eyre!(
"expected EXO ULA prefix to be /64, got {prefix} with /{}",
prefix.prefix_len()
));
}
Ok(Ipv6Net::new_assert(
Ipv6Addr::from_bits(prefix.trunc().addr().to_bits() | u128::from(node_id)),
128,
))
}
fn create_node_id(path: &Path) -> eyre::Result<u64> {
let Some(parent) = path.parent() else {
return Err(eyre!(
"node id file has no parent directory: {}",
path.display()
));
};
fs::create_dir_all(parent)
.wrap_err_with(|| format!("creating node id directory {}", parent.display()))?;
let node_id = generate_node_id();
let mut file = match OpenOptions::new().write(true).create_new(true).open(path) {
Ok(file) => file,
Err(err) if err.kind() == std::io::ErrorKind::AlreadyExists => {
return read_node_id(path);
}
Err(err) => {
return Err(err).wrap_err_with(|| format!("creating node id file {}", path.display()));
}
};
file.set_permissions(fs::Permissions::from_mode(0o600))
.wrap_err_with(|| format!("setting permissions on {}", path.display()))?;
writeln!(file, "{node_id:016x}")
.wrap_err_with(|| format!("writing node id file {}", path.display()))?;
file.sync_all()
.wrap_err_with(|| format!("syncing node id file {}", path.display()))?;
drop(file);
read_node_id(path)
}
fn read_node_id(path: &Path) -> eyre::Result<u64> {
let metadata =
fs::metadata(path).wrap_err_with(|| format!("reading metadata for {}", path.display()))?;
ensure_owner(path, &metadata)?;
let raw = fs::read_to_string(path)
.wrap_err_with(|| format!("reading node id file {}", path.display()))?;
let trimmed = raw.trim();
if trimmed.is_empty() {
return Err(eyre!("node id file is empty: {}", path.display()));
}
let node_id = u64::from_str_radix(trimmed.trim_start_matches("0x"), 16)
.wrap_err_with(|| format!("invalid node id in {}: {:?}", path.display(), trimmed))?;
Ok(node_id)
}
fn ensure_owner(path: &Path, metadata: &fs::Metadata) -> eyre::Result<()> {
let expected_uid = geteuid().as_raw();
let actual_uid = metadata.uid();
if actual_uid != expected_uid {
return Err(eyre!(
"node id file {} is owned by uid {}, expected {}",
path.display(),
actual_uid,
expected_uid
));
}
Ok(())
}
fn generate_node_id() -> u64 {
rand::random::<u64>()
}
fn is_not_found(err: &eyre::Report) -> bool {
err.downcast_ref::<std::io::Error>()
.is_some_and(|e| e.kind() == std::io::ErrorKind::NotFound)
}
#[cfg(test)]
mod tests {
use super::*;
fn temp_path(name: &str) -> std::path::PathBuf {
let nonce = rand::random::<u64>();
std::env::temp_dir().join(format!(
"babblerd-identity-{name}-{}-{nonce}",
std::process::id()
))
}
#[test]
fn creates_and_reloads_same_node_id() {
let dir = temp_path("create");
let path = dir.join("node-id");
let first = load_or_create_node_id(&path).expect("create node id");
let second = load_or_create_node_id(&path).expect("reload node id");
assert_eq!(first, second);
let _ = fs::remove_file(&path);
let _ = fs::remove_dir_all(&dir);
}
#[test]
fn node_addr_uses_full_low_64_bits() {
let addr = node_addr(
Ipv6Net::new_assert(
Ipv6Addr::new(0xfde0, 0x20c6, 0x1fa7, 0xffff, 0, 0, 0, 0),
64,
),
0x1234_5678_9abc_def0,
)
.expect("node address should be constructed");
assert_eq!(addr.prefix_len(), 128);
assert_eq!(
addr.addr(),
Ipv6Addr::new(
0xfde0, 0x20c6, 0x1fa7, 0xffff, 0x1234, 0x5678, 0x9abc, 0xdef0,
)
);
}
}
-337
View File
@@ -1,337 +0,0 @@
#[cfg(not(any(target_os = "macos", target_os = "linux")))]
compile_error!("babblerd is mac/linux-only");
pub mod babel;
pub mod config;
pub mod daemon;
pub mod dataplane;
pub mod fib;
pub mod identity;
pub mod profiling;
pub(crate) mod route_ctl;
pub mod routing_stack;
pub mod tun;
pub use babel::babel;
pub use config::EXO_ULA_PREFIX as PREFIX;
pub use error::{BabbleError, Result};
pub use if_watcher::watch;
pub mod error {
use std::io;
use thiserror::Error;
pub type Result<T> = core::result::Result<T, BabbleError>;
#[derive(Error, Debug)]
pub enum BabbleError {
#[error("An IO error occurred: {0}")]
Io(#[from] io::Error),
#[error("Unspecified error")]
Unspecified,
#[error("Babeld crashed unexpectedly with code: {0:?}")]
BabeldCrashed(Option<i32>),
#[error("Failed to set IP address")]
FailedToSetIp,
#[error("Other error: {0}")]
Other(String),
}
}
pub mod if_watcher {
#[cfg(target_os = "linux")]
use std::path::PathBuf;
use std::{collections::HashSet, net::IpAddr};
use futures_lite::StreamExt;
use n0_watcher::Watcher;
use netwatch::interfaces::{Interface, IpNet};
use tokio::sync::mpsc;
use crate::config::{EXO_ULA_PREFIX, PHYSICAL_LINK_MTU, interface_allowlist_from_env};
use crate::ip_manager::remove_ip;
use crate::{BabbleError, Result, babel::Babble};
pub const LOCALHOST_INTERFACE_NAMES: [&'static str; 2] = ["lo", "lo0"];
trait IfaceExt {
fn has_link_local_v6(&self) -> bool;
fn has_required_mtu(&self) -> bool;
fn is_real_interface(&self) -> bool;
fn will_babel(&self) -> bool;
}
impl IfaceExt for Interface {
fn will_babel(&self) -> bool {
self.has_link_local_v6()
&& self.has_required_mtu()
&& 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 has_required_mtu(&self) -> bool {
let Some(mtu) = interface_mtu(self.name()) else {
tracing::debug!(
"skipping interface {} because MTU could not be determined",
self.name()
);
return false;
};
if mtu < u32::from(PHYSICAL_LINK_MTU) {
tracing::debug!(
"skipping interface {} because mtu {} is below required {}",
self.name(),
mtu,
PHYSICAL_LINK_MTU
);
return false;
}
true
}
fn is_real_interface(&self) -> bool {
// macOS interface names are only a broad bootstrap signal here:
// Thunderbolt links are not limited to en2/en3, and high-numbered
// en* devices can be unrelated USB/Ethernet adapters.
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 interface_mtu(name: &str) -> Option<u32> {
netdev::get_interfaces()
.into_iter()
.find(|iface| iface.name == name)
.and_then(|iface| iface.mtu)
}
#[tracing::instrument(skip(send))]
pub async fn watch(send: mpsc::Sender<Babble>) -> Result<()> {
let mut ready_ifaces = HashSet::new();
let interface_allowlist = interface_allowlist_from_env()
.map_err(|e| BabbleError::Other(format!("invalid interface allowlist: {e}")))?;
if let Some(allowlist) = &interface_allowlist {
tracing::info!(?allowlist, "interface allowlist active");
}
tracing::info!("starting interface monitor");
let mon = netwatch::netmon::Monitor::new()
.await
.map_err(|_| BabbleError::Unspecified)?;
// TODD: this should never really be a thing thats the case, BUT I like the idea of having
// "heuristic" scripts that can help resolve issues but not necessarily gurantee success;
// I like the idea of generalising this concept into a framework where we have "heuristic tasks"
// that run to aid in tyring to fix some system ale-ment or whatever
//
// one-shot cleanup:
// - remove any stale app-prefix addresses from lo0
// - remove any app-prefix addresses that accidentally landed on physical links
{
let state = mon.interface_state();
for iface in state.peek().interfaces.values() {
let cleanup_target =
LOCALHOST_INTERFACE_NAMES.contains(&iface.name()) || iface.is_real_interface();
if !cleanup_target {
continue;
}
for addr in iface.addrs() {
if let IpNet::V6 { net: v6, .. } = addr
&& EXO_ULA_PREFIX.contains(&v6.addr())
{
tracing::info!("removing stale app ip {v6} from {}", iface.name());
if let Err(e) = remove_ip(v6, iface).await {
tracing::warn!(%e, "failed to remove stale app ip");
}
}
}
}
}
// stream updates
let mut mon_stream = mon.interface_state().stream();
while let Some(s) = mon_stream.next().await {
for iface in s.interfaces.values() {
if let Some(allowlist) = &interface_allowlist
&& !allowlist.contains(iface.name())
{
tracing::debug!(
"skipping interface {} because it is not in {}",
iface.name(),
crate::config::INTERFACE_ALLOWLIST_ENV
);
continue;
}
if !iface.is_real_interface() {
continue;
}
// physical links should not carry babbler application-space addresses
for addr in iface.addrs() {
if let IpNet::V6 { net: v6, .. } = addr
&& EXO_ULA_PREFIX.contains(&v6.addr())
{
tracing::info!("removing app ip {v6} from {}", iface.name());
if let Err(e) = remove_ip(v6, iface).await {
tracing::warn!(%e, "failed to remove ip");
}
}
}
if !iface.will_babel() {
continue;
}
if ready_ifaces.insert(iface.name().to_owned()) {
tracing::info!("telling babeld to watch {}", iface.name());
let Ok(()) = send.send(Babble::AddIface(iface.name().into())).await else {
return Ok(());
};
}
}
}
tracing::info!("stopping interface monitor");
Ok(())
}
}
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)
}
}
}
}
-251
View File
@@ -1,251 +0,0 @@
// Major TODO: at some point don't call it "babbler" because that is a silly name that makes no sense
// but this is at the very bottom of my concerns right now :)
#[cfg(not(any(target_os = "macos", target_os = "linux")))]
compile_error!("babblerd is mac/linux-only");
use std::{env, fs::Permissions, io, os::unix::fs::PermissionsExt, sync::Arc};
use babblerd::{
babel::BabelState,
config::{Config, TUN_MTU_ENV, TransportMode, default_tun_mtu},
daemon, identity,
tun::TunDevice,
};
use clap::Parser;
use color_eyre::eyre::{self, WrapErr, eyre};
use tokio::{
io::{AsyncBufReadExt, AsyncWriteExt, BufReader},
net::UnixListener,
net::UnixStream,
signal,
sync::watch,
task::JoinSet,
time::{Duration, sleep},
};
const INTERNAL_KEEPALIVE_TTL_MS: u64 = 30_000;
const INTERNAL_KEEPALIVE_INTERVAL_MS: u64 = 10_000;
#[derive(Debug, Parser)]
struct Cli {
#[arg(long, value_parser = parse_transport_mode)]
router_transport: Option<TransportMode>,
#[arg(long, conflicts_with = "router_transport")]
force_tcp: bool,
#[arg(long, value_parser = parse_tun_mtu)]
tun_mtu: Option<u16>,
}
#[tokio::main]
async fn main() -> eyre::Result<()> {
color_eyre::install()?;
tracing_subscriber::fmt()
.with_env_filter(tracing_subscriber::EnvFilter::from_default_env())
.init();
let cli = Cli::parse();
let mut config = Config::from_env()?;
let tun_mtu_from_env = env::var_os(TUN_MTU_ENV).is_some();
let transport_overridden = cli.router_transport.is_some() || cli.force_tcp;
if let Some(router_transport) = cli.router_transport {
config.router_transport = router_transport;
}
if cli.force_tcp {
config.router_transport = TransportMode::Tcp;
}
if transport_overridden && !tun_mtu_from_env && cli.tun_mtu.is_none() {
config.tun_mtu = default_tun_mtu(config.router_transport);
}
if let Some(tun_mtu) = cli.tun_mtu {
config.tun_mtu = tun_mtu;
}
// cleanup old public socket path
match std::fs::remove_file(&config.public_socket_path) {
Err(e) if e.kind() != io::ErrorKind::NotFound => return Err(e.into()),
Ok(()) => {
tracing::info!(
"cleaned up old file at {}",
config.public_socket_path.display()
);
}
_ => {}
}
// create new public directory
std::fs::create_dir_all(&config.public_dir)?;
if let Err(e) = std::fs::set_permissions(&config.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(&config).await;
let _ = std::fs::remove_file(&config.public_socket_path);
res
}
async fn inner_main(config: &Config) -> eyre::Result<()> {
let node_id = identity::load_or_create_node_id(&config.node_id_file)?;
let node_addr = identity::node_addr(config.exo_ula_prefix, node_id)?;
let tun = TunDevice::create(node_addr.addr(), config.tun_mtu)
.wrap_err("creating tun for node address")?;
tracing::info!("creating socket at {}", config.public_socket_path.display());
tracing::info!(
"router defaults: transport={} port={} tun_mtu={} tcp_batch_target_bytes={} tcp_socket_buffer_bytes={} node_id_file={} app_prefix={} node_id={:#018x} node_addr={} tun={}",
config.router_transport,
config.router_udp_port,
config.tun_mtu,
config.tcp_batch_target_bytes,
config.tcp_socket_buffer_bytes,
config.node_id_file.display(),
config.exo_ula_prefix,
node_id,
node_addr,
tun.ifname(),
);
if config.router_transport == TransportMode::Tcp {
tracing::warn!(
tun_mtu = config.tun_mtu,
tcp_batch_target_bytes = config.tcp_batch_target_bytes,
tcp_socket_buffer_bytes = config.tcp_socket_buffer_bytes,
"forced TCP transport requires matching TUN MTU on every peer; larger received frames are rejected"
);
}
let public_socket = UnixListener::bind(&config.public_socket_path)?;
// make our socket world accessible
std::fs::set_permissions(&config.public_socket_path, Permissions::from_mode(0o0666))?;
let (babel_state_send, _) = watch::channel(Arc::new(BabelState::new()));
let (daemon, mut core_task) = daemon::DaemonCore::spawn(
node_id,
config.exo_ula_prefix,
config.router_udp_port,
config.router_transport,
config.tun_mtu,
config.tcp_batch_target_bytes,
config.tcp_socket_buffer_bytes,
node_addr,
tun,
babel_state_send,
);
// TEMP: keep the daemon alive without an external client until the real
// frontend/test harness exists. This should be removed later.
let mut internal_keepalive =
tokio::spawn(internal_keepalive_client(config.public_socket_path.clone()));
let mut listeners = JoinSet::new();
loop {
tokio::select! {
sig = signal::ctrl_c() => {
sig?;
internal_keepalive.abort();
let _ = (&mut internal_keepalive).await;
listeners.abort_all();
while let Some(res) = listeners.join_next().await {
res.wrap_err("while ctrl-c")?;
}
drop(daemon);
core_task.await??;
break;
}
sock = public_socket.accept() => {
let sock = sock?.0;
listeners.spawn(daemon::handle_client(sock, daemon.clone()));
}
res = &mut core_task => {
res??;
internal_keepalive.abort();
let _ = (&mut internal_keepalive).await;
listeners.abort_all();
while let Some(res2) = listeners.join_next().await {
res2.wrap_err("while closing daemon core")?;
}
break;
}
res = &mut internal_keepalive => {
return Err(eyre!("internal keepalive client exited unexpectedly: {res:?}"));
}
next_join_result = listeners.join_next(), if !listeners.is_empty() => {
next_join_result.expect("checked")?;
tracing::info!("dropped a listener");
}
}
}
Ok(())
}
fn parse_transport_mode(raw: &str) -> Result<TransportMode, String> {
raw.parse()
}
fn parse_tun_mtu(raw: &str) -> Result<u16, String> {
babblerd::config::parse_tun_mtu(raw)
}
async fn internal_keepalive_client(socket_path: std::path::PathBuf) {
loop {
match UnixStream::connect(&socket_path).await {
Ok(stream) => {
tracing::info!(
socket=%socket_path.display(),
"internal keepalive client connected"
);
let (reader, mut writer) = stream.into_split();
let mut reader = BufReader::new(reader).lines();
match reader.next_line().await {
Ok(Some(line)) => {
tracing::debug!(?line, "internal keepalive initial state");
}
Ok(None) => {
tracing::warn!("internal keepalive connection closed before initial state");
sleep(Duration::from_secs(1)).await;
continue;
}
Err(err) => {
tracing::warn!(error=%err, "internal keepalive failed to read initial state");
sleep(Duration::from_secs(1)).await;
continue;
}
}
loop {
let command = format!("keepalive {INTERNAL_KEEPALIVE_TTL_MS}\n");
if let Err(err) = writer.write_all(command.as_bytes()).await {
tracing::warn!(error=%err, "internal keepalive failed to send keepalive");
break;
}
match reader.next_line().await {
Ok(Some(line)) => {
tracing::debug!(?line, "internal keepalive response");
}
Ok(None) => {
tracing::warn!("internal keepalive connection closed");
break;
}
Err(err) => {
tracing::warn!(error=%err, "internal keepalive failed to read response");
break;
}
}
sleep(Duration::from_millis(INTERNAL_KEEPALIVE_INTERVAL_MS)).await;
}
}
Err(err) => {
tracing::warn!(error=%err, socket=%socket_path.display(), "internal keepalive failed to connect");
}
}
sleep(Duration::from_secs(1)).await;
}
}
-124
View File
@@ -1,124 +0,0 @@
use std::time::Duration;
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct LatencyStats {
pub sent: u32,
pub received: u32,
pub loss_ratio: f64,
pub min: Duration,
pub avg: Duration,
pub max: Duration,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct CapacitySample {
pub sent_packets: u32,
pub received_packets: u32,
pub received_bytes: u64,
pub span: Duration,
}
impl CapacitySample {
pub fn loss_ratio(self) -> f64 {
if self.sent_packets == 0 {
return 0.0;
}
let lost = self.sent_packets.saturating_sub(self.received_packets);
f64::from(lost) / f64::from(self.sent_packets)
}
pub fn mbps(self) -> Option<f64> {
capacity_mbps(self.received_bytes, self.span)
}
}
pub fn latency_stats(sent: u32, samples: &[Duration]) -> Option<LatencyStats> {
let received = u32::try_from(samples.len()).ok()?;
if sent == 0 || samples.is_empty() {
return None;
}
let mut min = samples.first().copied()?;
let mut max = min;
let mut total_nanos = 0_u128;
for sample in samples {
min = min.min(*sample);
max = max.max(*sample);
total_nanos = total_nanos.saturating_add(sample.as_nanos());
}
let avg_nanos = total_nanos / u128::from(received);
let avg = Duration::from_nanos(u64_saturating_from_u128(avg_nanos));
let lost = sent.saturating_sub(received);
Some(LatencyStats {
sent,
received,
loss_ratio: f64::from(lost) / f64::from(sent),
min,
avg,
max,
})
}
pub fn capacity_mbps(received_bytes: u64, span: Duration) -> Option<f64> {
let nanos = span.as_nanos();
if received_bytes == 0 || nanos == 0 {
return None;
}
let bits = received_bytes.saturating_mul(8);
Some((bits as f64) * 1_000.0 / (nanos as f64))
}
pub fn duration_nanos_u64(duration: Duration) -> u64 {
u64_saturating_from_u128(duration.as_nanos())
}
fn u64_saturating_from_u128(value: u128) -> u64 {
u64::try_from(value).unwrap_or(u64::MAX)
}
#[cfg(test)]
mod tests {
use std::time::Duration;
use super::{CapacitySample, capacity_mbps, latency_stats};
#[test]
fn latency_summary_reports_loss_and_bounds() {
let samples = [
Duration::from_millis(3),
Duration::from_millis(1),
Duration::from_millis(2),
];
let Some(stats) = latency_stats(4, &samples) else {
panic!("expected latency stats");
};
assert_eq!(stats.sent, 4);
assert_eq!(stats.received, 3);
assert_eq!(stats.loss_ratio, 0.25);
assert_eq!(stats.min, Duration::from_millis(1));
assert_eq!(stats.avg, Duration::from_millis(2));
assert_eq!(stats.max, Duration::from_millis(3));
}
#[test]
fn capacity_summary_reports_mbps() {
let sample = CapacitySample {
sent_packets: 10,
received_packets: 10,
received_bytes: 125_000,
span: Duration::from_millis(1),
};
assert_eq!(sample.loss_ratio(), 0.0);
assert_eq!(
capacity_mbps(sample.received_bytes, sample.span),
Some(1000.0)
);
}
}
-16
View File
@@ -1,16 +0,0 @@
//! Link-local profiling support.
//!
//! This module is intentionally independent of the Babel control plane. The
//! standalone example uses it to measure one physical link directly; the daemon
//! can later consume the same types and estimators when route scoring is wired
//! in.
pub mod estimator;
pub mod pbprobe;
pub mod protocol;
pub mod socket;
pub mod standalone;
pub mod types;
pub use estimator::{CapacitySample, LatencyStats, capacity_mbps, latency_stats};
pub use types::{DEFAULT_PROFILE_PORT, LinkKey, ProbeConfig};
@@ -1,195 +0,0 @@
use std::time::Duration;
use crate::config::{OUTER_IPV6_HEADER_BYTES, OUTER_UDP_HEADER_BYTES, PHYSICAL_LINK_MTU};
pub const DEFAULT_PBPROBE_PORT: u16 = 41_902;
pub const DEFAULT_SAMPLE_COUNT: u32 = 200;
pub const DEFAULT_UTILIZATION: f64 = 0.01;
pub const DEFAULT_DISPERSION_THRESHOLD_MS: u64 = 1;
pub const DEFAULT_DISPERSION_THRESHOLD: Duration =
Duration::from_millis(DEFAULT_DISPERSION_THRESHOLD_MS);
pub const DEFAULT_MAX_BULK_LEN: u32 = 10_000;
pub const DEFAULT_RTS_TIMEOUT_MS: u64 = 750;
pub const DEFAULT_RTS_TIMEOUT: Duration = Duration::from_millis(DEFAULT_RTS_TIMEOUT_MS);
pub const DEFAULT_START_TIMEOUT_MS: u64 = 750;
pub const DEFAULT_START_TIMEOUT: Duration = Duration::from_millis(DEFAULT_START_TIMEOUT_MS);
pub const DEFAULT_CONTROL_RETRIES: u32 = 5;
#[derive(Debug, Clone)]
pub struct PbProbeConfig {
pub port: u16,
pub sample_count: u32,
pub utilization: f64,
pub dispersion_threshold: Duration,
pub initial_bulk_len: u32,
pub max_bulk_len: u32,
pub ip_packet_bytes: usize,
pub start_timeout: Duration,
pub rts_timeout: Duration,
pub control_retries: u32,
}
impl Default for PbProbeConfig {
fn default() -> Self {
Self {
port: DEFAULT_PBPROBE_PORT,
sample_count: DEFAULT_SAMPLE_COUNT,
utilization: DEFAULT_UTILIZATION,
dispersion_threshold: DEFAULT_DISPERSION_THRESHOLD,
initial_bulk_len: 1,
max_bulk_len: DEFAULT_MAX_BULK_LEN,
ip_packet_bytes: usize::from(PHYSICAL_LINK_MTU),
start_timeout: DEFAULT_START_TIMEOUT,
rts_timeout: DEFAULT_RTS_TIMEOUT,
control_retries: DEFAULT_CONTROL_RETRIES,
}
}
}
impl PbProbeConfig {
pub fn udp_payload_bytes(&self) -> usize {
let overhead = usize::from(OUTER_IPV6_HEADER_BYTES + OUTER_UDP_HEADER_BYTES);
self.ip_packet_bytes.saturating_sub(overhead)
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct AcceptedSample {
pub sample_id: u32,
pub bulk_len: u32,
pub delay_first: Duration,
pub delay_last: Duration,
pub dispersion: Duration,
pub server_issue_duration: Option<Duration>,
}
impl AcceptedSample {
pub fn delay_sum(self) -> Duration {
self.delay_first.saturating_add(self.delay_last)
}
}
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct SelectedSample {
pub sample: AcceptedSample,
pub capacity_mbps: f64,
}
#[derive(Debug, Clone, PartialEq)]
pub struct Estimate {
pub bulk_len: u32,
pub sample_count: u32,
pub attempts: u32,
pub lost_samples: u32,
pub ip_packet_bytes: usize,
pub selected: SelectedSample,
pub min_dispersion: Duration,
pub server_issue_samples: u32,
pub min_server_issue_duration: Option<Duration>,
}
#[derive(Debug, Clone, PartialEq)]
pub enum EstimateOutcome {
Complete(Estimate),
IncreaseBulk {
previous_bulk_len: u32,
next_bulk_len: u32,
observed_dispersion: Duration,
},
}
pub fn select_capacity_sample(
samples: &[AcceptedSample],
ip_packet_bytes: usize,
) -> Option<SelectedSample> {
let sample = samples
.iter()
.copied()
.min_by_key(|sample| sample.delay_sum())?;
let capacity_mbps = capacity_mbps(sample.bulk_len, ip_packet_bytes, sample.dispersion)?;
Some(SelectedSample {
sample,
capacity_mbps,
})
}
pub fn capacity_mbps(bulk_len: u32, ip_packet_bytes: usize, dispersion: Duration) -> Option<f64> {
let nanos = dispersion.as_nanos();
if bulk_len == 0 || ip_packet_bytes == 0 || nanos == 0 {
return None;
}
let bits = f64::from(bulk_len) * (ip_packet_bytes as f64) * 8.0;
Some(bits * 1_000.0 / (nanos as f64))
}
pub fn next_bulk_len(current: u32, max: u32) -> Option<u32> {
let next = current.checked_mul(10)?;
if next > max || next == current {
return None;
}
Some(next)
}
pub fn pacing_interval(dispersion: Duration, utilization: f64) -> Option<Duration> {
if dispersion.is_zero() || !utilization.is_finite() || utilization <= 0.0 {
return None;
}
Some(Duration::from_secs_f64(
(2.0 * dispersion.as_secs_f64()) / utilization,
))
}
#[cfg(test)]
mod tests {
use std::time::Duration;
use super::{
AcceptedSample, capacity_mbps, next_bulk_len, pacing_interval, select_capacity_sample,
};
#[test]
fn capacity_uses_bulk_length_not_packet_count() {
let estimate = capacity_mbps(100, 1500, Duration::from_micros(1200));
assert_eq!(estimate, Some(1000.0));
}
#[test]
fn selector_uses_minimum_delay_sum() {
let samples = [
AcceptedSample {
sample_id: 1,
bulk_len: 10,
delay_first: Duration::from_millis(3),
delay_last: Duration::from_millis(4),
dispersion: Duration::from_micros(900),
server_issue_duration: None,
},
AcceptedSample {
sample_id: 2,
bulk_len: 10,
delay_first: Duration::from_millis(1),
delay_last: Duration::from_millis(2),
dispersion: Duration::from_micros(1200),
server_issue_duration: None,
},
];
let selected = select_capacity_sample(&samples, 1500).expect("sample should be selected");
assert_eq!(selected.sample.sample_id, 2);
assert_eq!(selected.capacity_mbps, 100.0);
}
#[test]
fn bulk_growth_is_tenfold_and_capped() {
assert_eq!(next_bulk_len(1, 1000), Some(10));
assert_eq!(next_bulk_len(1000, 1000), None);
}
#[test]
fn pacing_follows_paper_formula() {
let interval = pacing_interval(Duration::from_millis(1), 0.01);
assert_eq!(interval, Some(Duration::from_millis(200)));
}
}
@@ -1,14 +0,0 @@
//! Paper-faithful PBProbe implementation.
//!
//! PBProbe is a CapProbe-derived capacity estimator that uses packet bulks
//! instead of a single packet pair. This module follows the paper algorithm
//! rather than the old C implementation's process/control structure.
pub mod estimator;
pub mod protocol;
pub mod standalone;
pub use estimator::{
AcceptedSample, Estimate, EstimateOutcome, PbProbeConfig, SelectedSample, next_bulk_len,
pacing_interval, select_capacity_sample,
};
@@ -1,317 +0,0 @@
use std::mem::size_of;
use std::time::Duration;
use thiserror::Error;
use zerocopy::byteorder::{NetworkEndian, U16, U32, U64};
use zerocopy::{FromBytes, Immutable, IntoBytes, KnownLayout};
type U16Be = U16<NetworkEndian>;
type U32Be = U32<NetworkEndian>;
type U64Be = U64<NetworkEndian>;
pub const HEADER_LEN: usize = size_of::<WireHeader>();
pub const RESULT_BODY_LEN: usize = size_of::<WireResultBody>();
pub const RESULT_PACKET_LEN: usize = HEADER_LEN + RESULT_BODY_LEN;
const MAGIC: &[u8; 4] = b"BBPB";
const VERSION: u8 = 1;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[repr(u8)]
pub enum PacketKind {
Start = 1,
StartAck = 2,
Rts = 3,
Bulk = 4,
Result = 5,
End = 6,
ErrorMessage = 7,
}
impl TryFrom<u8> for PacketKind {
type Error = ProtocolError;
fn try_from(value: u8) -> Result<Self, Self::Error> {
match value {
1 => Ok(Self::Start),
2 => Ok(Self::StartAck),
3 => Ok(Self::Rts),
4 => Ok(Self::Bulk),
5 => Ok(Self::Result),
6 => Ok(Self::End),
7 => Ok(Self::ErrorMessage),
other => Err(ProtocolError::UnknownKind(other)),
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Header {
pub kind: PacketKind,
pub run_id: u64,
pub sample_id: u32,
pub seq: u32,
pub bulk_len: u32,
pub sample_count: u32,
pub ip_packet_bytes: u32,
}
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct ResultBody {
pub attempts: u32,
pub lost_samples: u32,
pub selected_sample_id: u32,
pub accepted_samples: u32,
pub delay_sum: Duration,
pub dispersion: Duration,
pub min_dispersion: Duration,
pub capacity_mbps: f64,
}
#[repr(C)]
#[derive(Debug, Clone, Copy, FromBytes, Immutable, IntoBytes, KnownLayout)]
struct WireHeader {
magic: [u8; 4],
version: u8,
kind: u8,
flags: U16Be,
run_id: U64Be,
sample_id: U32Be,
seq: U32Be,
bulk_len: U32Be,
sample_count: U32Be,
ip_packet_bytes: U32Be,
reserved: U32Be,
}
#[repr(C)]
#[derive(Debug, Clone, Copy, FromBytes, Immutable, IntoBytes, KnownLayout)]
struct WireResultBody {
attempts: U32Be,
lost_samples: U32Be,
selected_sample_id: U32Be,
accepted_samples: U32Be,
delay_sum_nanos: U64Be,
dispersion_nanos: U64Be,
min_dispersion_nanos: U64Be,
capacity_mbps_bits: U64Be,
}
#[derive(Debug, Error, Clone, PartialEq, Eq)]
pub enum ProtocolError {
#[error("packet is too short")]
TooShort,
#[error("packet buffer is too small")]
BufferTooSmall,
#[error("bad PBProbe packet magic")]
BadMagic,
#[error("unsupported PBProbe protocol version {0}")]
BadVersion(u8),
#[error("unknown PBProbe packet kind {0}")]
UnknownKind(u8),
}
pub fn encode_header(dst: &mut [u8], header: Header) -> Result<usize, ProtocolError> {
encode_header_with_aux(dst, header, 0)
}
pub fn encode_header_with_aux(
dst: &mut [u8],
header: Header,
aux: u32,
) -> Result<usize, ProtocolError> {
write_bytes(dst, WireHeader::from_header(header, aux).as_bytes())
}
pub fn decode_header(src: &[u8]) -> Result<Header, ProtocolError> {
decode_header_with_aux(src).map(|(header, _aux)| header)
}
pub fn decode_header_with_aux(src: &[u8]) -> Result<(Header, u32), ProtocolError> {
let (wire, _) = WireHeader::read_from_prefix(src).map_err(|_| ProtocolError::TooShort)?;
wire.decode()
}
pub fn encode_result(
dst: &mut [u8],
header: Header,
body: ResultBody,
) -> Result<usize, ProtocolError> {
if dst.len() < RESULT_PACKET_LEN {
return Err(ProtocolError::BufferTooSmall);
}
let cursor = encode_header(dst, header)?;
write_bytes(&mut dst[cursor..], WireResultBody::from(body).as_bytes())?;
Ok(RESULT_PACKET_LEN)
}
pub fn decode_result_body(src: &[u8]) -> Result<ResultBody, ProtocolError> {
let body_src = src.get(HEADER_LEN..).ok_or(ProtocolError::TooShort)?;
let (wire, _) =
WireResultBody::read_from_prefix(body_src).map_err(|_| ProtocolError::TooShort)?;
Ok(ResultBody::from(wire))
}
pub fn duration_nanos(duration: Duration) -> u64 {
u64::try_from(duration.as_nanos()).unwrap_or(u64::MAX)
}
impl WireHeader {
fn from_header(header: Header, aux: u32) -> Self {
Self {
magic: *MAGIC,
version: VERSION,
kind: header.kind as u8,
flags: U16Be::ZERO,
run_id: U64Be::new(header.run_id),
sample_id: U32Be::new(header.sample_id),
seq: U32Be::new(header.seq),
bulk_len: U32Be::new(header.bulk_len),
sample_count: U32Be::new(header.sample_count),
ip_packet_bytes: U32Be::new(header.ip_packet_bytes),
reserved: U32Be::new(aux),
}
}
fn decode(self) -> Result<(Header, u32), ProtocolError> {
if self.magic != *MAGIC {
return Err(ProtocolError::BadMagic);
}
if self.version != VERSION {
return Err(ProtocolError::BadVersion(self.version));
}
Ok((
Header {
kind: PacketKind::try_from(self.kind)?,
run_id: self.run_id.get(),
sample_id: self.sample_id.get(),
seq: self.seq.get(),
bulk_len: self.bulk_len.get(),
sample_count: self.sample_count.get(),
ip_packet_bytes: self.ip_packet_bytes.get(),
},
self.reserved.get(),
))
}
}
impl From<ResultBody> for WireResultBody {
fn from(body: ResultBody) -> Self {
Self {
attempts: U32Be::new(body.attempts),
lost_samples: U32Be::new(body.lost_samples),
selected_sample_id: U32Be::new(body.selected_sample_id),
accepted_samples: U32Be::new(body.accepted_samples),
delay_sum_nanos: U64Be::new(duration_nanos(body.delay_sum)),
dispersion_nanos: U64Be::new(duration_nanos(body.dispersion)),
min_dispersion_nanos: U64Be::new(duration_nanos(body.min_dispersion)),
capacity_mbps_bits: U64Be::new(body.capacity_mbps.to_bits()),
}
}
}
impl From<WireResultBody> for ResultBody {
fn from(wire: WireResultBody) -> Self {
Self {
attempts: wire.attempts.get(),
lost_samples: wire.lost_samples.get(),
selected_sample_id: wire.selected_sample_id.get(),
accepted_samples: wire.accepted_samples.get(),
delay_sum: Duration::from_nanos(wire.delay_sum_nanos.get()),
dispersion: Duration::from_nanos(wire.dispersion_nanos.get()),
min_dispersion: Duration::from_nanos(wire.min_dispersion_nanos.get()),
capacity_mbps: f64::from_bits(wire.capacity_mbps_bits.get()),
}
}
}
fn write_bytes(dst: &mut [u8], src: &[u8]) -> Result<usize, ProtocolError> {
let Some(slot) = dst.get_mut(..src.len()) else {
return Err(ProtocolError::BufferTooSmall);
};
slot.copy_from_slice(src);
Ok(src.len())
}
#[cfg(test)]
mod tests {
use std::time::Duration;
use super::{
HEADER_LEN, Header, PacketKind, RESULT_PACKET_LEN, ResultBody, decode_header,
decode_header_with_aux, decode_result_body, encode_header, encode_header_with_aux,
encode_result,
};
#[test]
fn header_layout_is_stable() {
assert_eq!(HEADER_LEN, 40);
}
#[test]
fn header_round_trips() {
let header = Header {
kind: PacketKind::Bulk,
run_id: 7,
sample_id: 11,
seq: 3,
bulk_len: 100,
sample_count: 200,
ip_packet_bytes: 1500,
};
let mut buf = [0_u8; HEADER_LEN];
assert_eq!(encode_header(&mut buf, header), Ok(HEADER_LEN));
assert_eq!(decode_header(&buf), Ok(header));
}
#[test]
fn header_aux_round_trips() {
let header = Header {
kind: PacketKind::Bulk,
run_id: 7,
sample_id: 11,
seq: 100,
bulk_len: 100,
sample_count: 200,
ip_packet_bytes: 1500,
};
let mut buf = [0_u8; HEADER_LEN];
assert_eq!(
encode_header_with_aux(&mut buf, header, 12_345),
Ok(HEADER_LEN)
);
assert_eq!(decode_header_with_aux(&buf), Ok((header, 12_345)));
}
#[test]
fn result_round_trips() {
let header = Header {
kind: PacketKind::Result,
run_id: 9,
sample_id: 0,
seq: 0,
bulk_len: 100,
sample_count: 200,
ip_packet_bytes: 1500,
};
let body = ResultBody {
attempts: 210,
lost_samples: 10,
selected_sample_id: 42,
accepted_samples: 200,
delay_sum: Duration::from_micros(123),
dispersion: Duration::from_micros(1200),
min_dispersion: Duration::from_micros(1100),
capacity_mbps: 1000.25,
};
let mut buf = [0_u8; RESULT_PACKET_LEN];
assert_eq!(encode_result(&mut buf, header, body), Ok(RESULT_PACKET_LEN));
assert_eq!(decode_header(&buf), Ok(header));
assert_eq!(decode_result_body(&buf), Ok(body));
}
}
File diff suppressed because it is too large. Load diff
-228
View File
@@ -1,228 +0,0 @@
use std::mem::size_of;
use thiserror::Error;
use zerocopy::byteorder::{NetworkEndian, U16, U32, U64};
use zerocopy::{FromBytes, Immutable, IntoBytes, KnownLayout};
type U16Be = U16<NetworkEndian>;
type U32Be = U32<NetworkEndian>;
type U64Be = U64<NetworkEndian>;
pub const HEADER_LEN: usize = size_of::<WireHeader>();
pub const SUMMARY_BODY_LEN: usize = size_of::<WireSummaryBody>();
pub const SUMMARY_PACKET_LEN: usize = HEADER_LEN + SUMMARY_BODY_LEN;
const MAGIC: &[u8; 4] = b"BBLP";
const VERSION: u8 = 1;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[repr(u8)]
pub enum PacketKind {
EchoRequest = 1,
EchoReply = 2,
Train = 3,
SummaryRequest = 4,
SummaryReply = 5,
}
impl TryFrom<u8> for PacketKind {
type Error = ProtocolError;
fn try_from(value: u8) -> Result<Self, Self::Error> {
match value {
1 => Ok(Self::EchoRequest),
2 => Ok(Self::EchoReply),
3 => Ok(Self::Train),
4 => Ok(Self::SummaryRequest),
5 => Ok(Self::SummaryReply),
other => Err(ProtocolError::UnknownKind(other)),
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Header {
pub kind: PacketKind,
pub run_id: u64,
pub seq: u32,
pub count: u32,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct SummaryBody {
pub received_packets: u32,
pub received_bytes: u64,
pub span_nanos: u64,
}
#[repr(C)]
#[derive(Debug, Clone, Copy, FromBytes, Immutable, IntoBytes, KnownLayout)]
struct WireHeader {
magic: [u8; 4],
version: u8,
kind: u8,
flags: U16Be,
run_id: U64Be,
seq: U32Be,
count: U32Be,
}
#[repr(C)]
#[derive(Debug, Clone, Copy, FromBytes, Immutable, IntoBytes, KnownLayout)]
struct WireSummaryBody {
received_packets: U32Be,
received_bytes: U64Be,
span_nanos: U64Be,
}
#[derive(Debug, Error, Clone, PartialEq, Eq)]
pub enum ProtocolError {
#[error("packet is too short")]
TooShort,
#[error("packet buffer is too small")]
BufferTooSmall,
#[error("bad profiling packet magic")]
BadMagic,
#[error("unsupported profiling protocol version {0}")]
BadVersion(u8),
#[error("unknown profiling packet kind {0}")]
UnknownKind(u8),
}
pub fn encode_header(dst: &mut [u8], header: Header) -> Result<usize, ProtocolError> {
write_bytes(dst, WireHeader::from_header(header).as_bytes())
}
pub fn decode_header(src: &[u8]) -> Result<Header, ProtocolError> {
let (wire, _) = WireHeader::read_from_prefix(src).map_err(|_| ProtocolError::TooShort)?;
wire.decode()
}
pub fn encode_summary(
dst: &mut [u8],
header: Header,
body: SummaryBody,
) -> Result<usize, ProtocolError> {
if dst.len() < SUMMARY_PACKET_LEN {
return Err(ProtocolError::BufferTooSmall);
}
let cursor = encode_header(dst, header)?;
write_bytes(&mut dst[cursor..], WireSummaryBody::from(body).as_bytes())?;
Ok(SUMMARY_PACKET_LEN)
}
pub fn decode_summary_body(src: &[u8]) -> Result<SummaryBody, ProtocolError> {
let body_src = src.get(HEADER_LEN..).ok_or(ProtocolError::TooShort)?;
let (wire, _) =
WireSummaryBody::read_from_prefix(body_src).map_err(|_| ProtocolError::TooShort)?;
Ok(SummaryBody::from(wire))
}
impl WireHeader {
fn from_header(header: Header) -> Self {
Self {
magic: *MAGIC,
version: VERSION,
kind: header.kind as u8,
flags: U16Be::ZERO,
run_id: U64Be::new(header.run_id),
seq: U32Be::new(header.seq),
count: U32Be::new(header.count),
}
}
fn decode(self) -> Result<Header, ProtocolError> {
if self.magic != *MAGIC {
return Err(ProtocolError::BadMagic);
}
if self.version != VERSION {
return Err(ProtocolError::BadVersion(self.version));
}
Ok(Header {
kind: PacketKind::try_from(self.kind)?,
run_id: self.run_id.get(),
seq: self.seq.get(),
count: self.count.get(),
})
}
}
impl From<SummaryBody> for WireSummaryBody {
fn from(body: SummaryBody) -> Self {
Self {
received_packets: U32Be::new(body.received_packets),
received_bytes: U64Be::new(body.received_bytes),
span_nanos: U64Be::new(body.span_nanos),
}
}
}
impl From<WireSummaryBody> for SummaryBody {
fn from(wire: WireSummaryBody) -> Self {
Self {
received_packets: wire.received_packets.get(),
received_bytes: wire.received_bytes.get(),
span_nanos: wire.span_nanos.get(),
}
}
}
fn write_bytes(dst: &mut [u8], src: &[u8]) -> Result<usize, ProtocolError> {
let Some(slot) = dst.get_mut(..src.len()) else {
return Err(ProtocolError::BufferTooSmall);
};
slot.copy_from_slice(src);
Ok(src.len())
}
#[cfg(test)]
mod tests {
use super::{
HEADER_LEN, Header, PacketKind, SUMMARY_PACKET_LEN, SummaryBody, decode_header,
decode_summary_body, encode_header, encode_summary,
};
#[test]
fn layouts_are_stable() {
assert_eq!(HEADER_LEN, 24);
assert_eq!(SUMMARY_PACKET_LEN, 44);
}
#[test]
fn header_round_trips() {
let header = Header {
kind: PacketKind::Train,
run_id: 42,
seq: 7,
count: 64,
};
let mut buf = [0_u8; HEADER_LEN];
let encoded = encode_header(&mut buf, header);
assert_eq!(encoded, Ok(HEADER_LEN));
assert_eq!(decode_header(&buf), Ok(header));
}
#[test]
fn summary_round_trips() {
let header = Header {
kind: PacketKind::SummaryReply,
run_id: 99,
seq: 0,
count: 64,
};
let body = SummaryBody {
received_packets: 63,
received_bytes: 91_476,
span_nanos: 725_000,
};
let mut buf = [0_u8; SUMMARY_PACKET_LEN];
let encoded = encode_summary(&mut buf, header, body);
assert_eq!(encoded, Ok(SUMMARY_PACKET_LEN));
assert_eq!(decode_header(&buf), Ok(header));
assert_eq!(decode_summary_body(&buf), Ok(body));
}
}
-63
View File
@@ -1,63 +0,0 @@
use std::io;
use std::net::{Ipv6Addr, SocketAddr, SocketAddrV6, UdpSocket};
use std::num::NonZeroU32;
use std::time::Duration;
use nix::net::if_::if_nametoindex;
use socket2::{Domain, Protocol, Socket, Type};
const PROFILE_SOCKET_BUFFER_BYTES: usize = 4 * 1024 * 1024;
pub fn open_link_local_udp(
ifname: &str,
port: u16,
read_timeout: Option<Duration>,
) -> io::Result<(UdpSocket, u32)> {
let ifindex = if_nametoindex(ifname).map_err(io::Error::from)?;
let Some(nonzero_ifindex) = NonZeroU32::new(ifindex) else {
return Err(io::Error::new(
io::ErrorKind::InvalidInput,
format!("invalid ifindex for {ifname}"),
));
};
let socket = Socket::new(Domain::IPV6, Type::DGRAM, Some(Protocol::UDP))?;
socket.set_reuse_address(true)?;
socket.set_reuse_port(true)?;
socket.set_only_v6(true)?;
socket.set_recv_buffer_size(PROFILE_SOCKET_BUFFER_BYTES)?;
socket.set_send_buffer_size(PROFILE_SOCKET_BUFFER_BYTES)?;
#[cfg(target_os = "linux")]
socket.bind_device(Some(ifname.as_bytes()))?;
socket.bind_device_by_index_v6(Some(nonzero_ifindex))?;
socket.bind(&SocketAddrV6::new(Ipv6Addr::UNSPECIFIED, port, 0, 0).into())?;
let udp: UdpSocket = socket.into();
udp.set_read_timeout(read_timeout)?;
udp.set_write_timeout(read_timeout)?;
Ok((udp, ifindex))
}
pub fn scoped_peer_addr(peer: Ipv6Addr, port: u16, ifindex: u32) -> SocketAddr {
SocketAddr::V6(SocketAddrV6::new(peer, port, 0, ifindex))
}
pub fn with_default_scope(addr: SocketAddr, ifindex: u32) -> SocketAddr {
match addr {
SocketAddr::V6(v6) if v6.ip().is_unicast_link_local() && v6.scope_id() == 0 => {
SocketAddr::V6(SocketAddrV6::new(
*v6.ip(),
v6.port(),
v6.flowinfo(),
ifindex,
))
}
other => other,
}
}
pub fn parse_link_local_addr(raw: &str) -> Result<Ipv6Addr, std::net::AddrParseError> {
let addr = raw.split_once('%').map_or(raw, |(addr, _scope)| addr);
addr.parse()
}
-722
View File
@@ -1,722 +0,0 @@
use std::collections::HashMap;
use std::ffi::OsString;
use std::io::{self, ErrorKind};
use std::net::{Ipv6Addr, SocketAddr, UdpSocket};
use std::thread;
use std::time::{Duration, Instant};
use clap::{Args, Parser, Subcommand};
use color_eyre::eyre::{Result, WrapErr, eyre};
use super::estimator::{CapacitySample, duration_nanos_u64, latency_stats};
use super::protocol::{
HEADER_LEN, Header, PacketKind, SUMMARY_PACKET_LEN, SummaryBody, decode_header,
decode_summary_body, encode_header, encode_summary,
};
use super::socket::{
open_link_local_udp, parse_link_local_addr, scoped_peer_addr, with_default_scope,
};
use super::types::{
DEFAULT_CAPACITY_ROUNDS, DEFAULT_ECHO_COUNT, DEFAULT_ECHO_INTERVAL_MS, DEFAULT_ECHO_TIMEOUT_MS,
DEFAULT_PROFILE_PORT, DEFAULT_TRAIN_INTERVAL_MS, DEFAULT_TRAIN_PACKETS,
DEFAULT_TRAIN_SETTLE_MS, ProbeConfig,
};
const MAX_UDP_PACKET_BYTES: usize = 65_535;
const REFLECT_RECV_TIMEOUT: Duration = Duration::from_secs(1);
const STALE_TRAIN_AFTER: Duration = Duration::from_secs(60);
const SUMMARY_REQUEST_ATTEMPTS: u32 = 3;
pub fn run_from_env() -> Result<()> {
run_cli(Cli::parse())
}
pub fn run<I, S>(args: I) -> Result<()>
where
I: IntoIterator<Item = S>,
S: Into<OsString> + Clone,
{
run_cli(Cli::try_parse_from(args)?)
}
fn run_cli(cli: Cli) -> Result<()> {
match cli.command {
Command::Probe(args) => run_probe(args.into_probe_options()),
Command::Reflect(options) => run_reflect(options),
}
}
#[derive(Debug, Parser)]
#[command(name = "link_profile")]
#[command(about = "Run standalone link-local latency and packet-train probes")]
#[command(arg_required_else_help = true)]
struct Cli {
#[command(subcommand)]
command: Command,
}
#[derive(Debug, Subcommand)]
enum Command {
Probe(ProbeArgs),
Reflect(ReflectOptions),
}
#[derive(Debug, Clone)]
struct ProbeOptions {
ifname: String,
peer: Ipv6Addr,
config: ProbeConfig,
}
#[derive(Debug, Args, Clone)]
struct ReflectOptions {
#[arg(long)]
ifname: String,
#[arg(long, default_value_t = DEFAULT_PROFILE_PORT)]
port: u16,
}
#[derive(Debug, Args, Clone)]
struct ProbeArgs {
#[arg(long)]
ifname: String,
#[arg(long, value_parser = parse_link_local_addr_arg)]
peer: Ipv6Addr,
#[arg(long, default_value_t = DEFAULT_PROFILE_PORT)]
port: u16,
#[arg(long = "echo-count", default_value_t = DEFAULT_ECHO_COUNT)]
echo_count: u32,
#[arg(long = "echo-interval-ms", default_value_t = DEFAULT_ECHO_INTERVAL_MS)]
echo_interval_ms: u64,
#[arg(long = "timeout-ms", default_value_t = DEFAULT_ECHO_TIMEOUT_MS)]
timeout_ms: u64,
#[arg(long = "capacity-rounds", default_value_t = DEFAULT_CAPACITY_ROUNDS)]
capacity_rounds: u32,
#[arg(long = "train-packets", default_value_t = DEFAULT_TRAIN_PACKETS)]
train_packets: u32,
#[arg(long = "payload-bytes", default_value_t = usize::from(crate::config::TUN_MTU))]
payload_bytes: usize,
#[arg(long = "train-interval-ms", default_value_t = DEFAULT_TRAIN_INTERVAL_MS)]
train_interval_ms: u64,
#[arg(long = "settle-ms", default_value_t = DEFAULT_TRAIN_SETTLE_MS)]
settle_ms: u64,
}
impl ProbeArgs {
fn into_probe_options(self) -> ProbeOptions {
ProbeOptions {
ifname: self.ifname,
peer: self.peer,
config: ProbeConfig {
port: self.port,
echo_count: self.echo_count,
echo_interval: Duration::from_millis(self.echo_interval_ms),
echo_timeout: Duration::from_millis(self.timeout_ms),
capacity_rounds: self.capacity_rounds,
train_packets: self.train_packets,
train_payload_bytes: self.payload_bytes,
train_interval: Duration::from_millis(self.train_interval_ms),
train_settle: Duration::from_millis(self.settle_ms),
},
}
}
}
#[derive(Debug)]
struct TrainAccumulator {
received_packets: u32,
received_bytes: u64,
first_rx: Option<Instant>,
last_rx: Option<Instant>,
last_update: Instant,
seen: Vec<bool>,
}
impl TrainAccumulator {
fn new(expected_packets: u32, now: Instant) -> Self {
Self {
received_packets: 0,
received_bytes: 0,
first_rx: None,
last_rx: None,
last_update: now,
seen: vec![false; usize::try_from(expected_packets).unwrap_or(0)],
}
}
fn record(&mut self, seq: u32, packet_len: usize, now: Instant) {
self.last_update = now;
if !mark_seen(&mut self.seen, seq) {
return;
}
self.received_packets = self.received_packets.saturating_add(1);
self.received_bytes = self
.received_bytes
.saturating_add(u64::try_from(packet_len).unwrap_or(u64::MAX));
if self.first_rx.is_none() {
self.first_rx = Some(now);
}
self.last_rx = Some(now);
}
fn summary(&self) -> SummaryBody {
let span_nanos = match (self.first_rx, self.last_rx) {
(Some(first), Some(last)) => duration_nanos_u64(last.saturating_duration_since(first)),
_ => 0,
};
SummaryBody {
received_packets: self.received_packets,
received_bytes: self.received_bytes,
span_nanos,
}
}
}
fn mark_seen(seen: &mut [bool], seq: u32) -> bool {
let Ok(index) = usize::try_from(seq) else {
return false;
};
let Some(slot) = seen.get_mut(index) else {
return false;
};
if *slot {
return false;
}
*slot = true;
true
}
fn run_reflect(options: ReflectOptions) -> Result<()> {
let (socket, ifindex) =
open_link_local_udp(&options.ifname, options.port, Some(REFLECT_RECV_TIMEOUT))
.wrap_err_with(|| format!("opening profiling reflector on {}", options.ifname))?;
let local_addr = socket
.local_addr()
.wrap_err("reading reflector local address")?;
println!(
"reflecting profiling probes on {} ifindex={} local={}",
options.ifname, ifindex, local_addr
);
let mut buf = vec![0_u8; MAX_UDP_PACKET_BYTES];
let mut trains = HashMap::<u64, TrainAccumulator>::new();
let mut last_cleanup = Instant::now();
loop {
cleanup_stale_trains(&mut trains, &mut last_cleanup);
let (packet_len, from) = match socket.recv_from(&mut buf) {
Ok(received) => received,
Err(err) if is_timeout(&err) || err.kind() == ErrorKind::Interrupted => continue,
Err(err) => return Err(err).wrap_err("receiving profiling packet"),
};
let Some(packet) = buf.get(..packet_len) else {
continue;
};
let Ok(header) = decode_header(packet) else {
continue;
};
match header.kind {
PacketKind::EchoRequest => {
send_echo_reply(&socket, ifindex, from, header).wrap_err("sending echo reply")?;
}
PacketKind::Train => {
let now = Instant::now();
trains
.entry(header.run_id)
.or_insert_with(|| TrainAccumulator::new(header.count, now))
.record(header.seq, packet_len, now);
}
PacketKind::SummaryRequest => {
send_summary_reply(&socket, ifindex, from, header, &mut trains)
.wrap_err("sending train summary")?;
}
PacketKind::EchoReply | PacketKind::SummaryReply => {}
}
}
}
fn run_probe(options: ProbeOptions) -> Result<()> {
if options.config.train_payload_bytes < HEADER_LEN {
return Err(eyre!(
"train payload must be at least {HEADER_LEN} bytes, got {}",
options.config.train_payload_bytes
));
}
if options.config.train_packets == 0 {
return Err(eyre!("train packet count must be non-zero"));
}
let (socket, ifindex) =
open_link_local_udp(&options.ifname, 0, Some(options.config.echo_timeout))
.wrap_err_with(|| format!("opening profiling probe socket on {}", options.ifname))?;
let peer = scoped_peer_addr(options.peer, options.config.port, ifindex);
let local_addr = socket
.local_addr()
.wrap_err("reading probe local address")?;
let base_run_id = make_base_run_id();
println!(
"probing {} via {} ifindex={} local={} peer_port={}",
options.peer, options.ifname, ifindex, local_addr, options.config.port
);
println!(
"capacity probe: rounds={} train_packets={} payload_bytes={} interval_ms={}",
options.config.capacity_rounds,
options.config.train_packets,
options.config.train_payload_bytes,
options.config.train_interval.as_millis()
);
let latency_samples = run_echo_probes(&socket, peer, base_run_id, &options.config)
.wrap_err("running latency probes")?;
print_latency_summary(options.config.echo_count, &latency_samples);
let capacity_samples = run_capacity_probes(&socket, peer, base_run_id, &options.config)
.wrap_err("running capacity probes")?;
print_capacity_summary(&capacity_samples);
Ok(())
}
fn send_echo_reply(
socket: &UdpSocket,
ifindex: u32,
from: SocketAddr,
request: Header,
) -> Result<()> {
let reply = Header {
kind: PacketKind::EchoReply,
run_id: request.run_id,
seq: request.seq,
count: request.count,
};
let mut out = [0_u8; HEADER_LEN];
encode_header(&mut out, reply)?;
send_datagram(socket, &out, with_default_scope(from, ifindex))
}
fn send_summary_reply(
socket: &UdpSocket,
ifindex: u32,
from: SocketAddr,
request: Header,
trains: &mut HashMap<u64, TrainAccumulator>,
) -> Result<()> {
let body = trains
.get(&request.run_id)
.map_or_else(empty_summary, TrainAccumulator::summary);
let reply = Header {
kind: PacketKind::SummaryReply,
run_id: request.run_id,
seq: 0,
count: request.count,
};
let mut out = [0_u8; SUMMARY_PACKET_LEN];
let len = encode_summary(&mut out, reply, body)?;
send_datagram(
socket,
out.get(..len).unwrap_or(&out),
with_default_scope(from, ifindex),
)?;
trains.remove(&request.run_id);
Ok(())
}
fn run_echo_probes(
socket: &UdpSocket,
peer: SocketAddr,
base_run_id: u64,
config: &ProbeConfig,
) -> Result<Vec<Duration>> {
let mut samples = Vec::new();
let run_id = base_run_id ^ 0xe0c0_u64;
for seq in 0..config.echo_count {
let header = Header {
kind: PacketKind::EchoRequest,
run_id,
seq,
count: config.echo_count,
};
let mut out = [0_u8; HEADER_LEN];
encode_header(&mut out, header)?;
let start = Instant::now();
send_datagram(socket, &out, peer)?;
match receive_echo_reply(socket, run_id, seq, start, config.echo_timeout)? {
Some(sample) => {
println!("echo {:>3}: {}", seq + 1, format_duration(sample));
samples.push(sample);
}
None => {
println!("echo {:>3}: timeout", seq + 1);
}
}
thread::sleep(config.echo_interval);
}
Ok(samples)
}
fn run_capacity_probes(
socket: &UdpSocket,
peer: SocketAddr,
base_run_id: u64,
config: &ProbeConfig,
) -> Result<Vec<CapacitySample>> {
let mut samples = Vec::new();
let mut train = vec![0_u8; config.train_payload_bytes];
for round in 0..config.capacity_rounds {
let run_id = base_run_id ^ (0xc0_ffee_u64.wrapping_add(u64::from(round)));
let sender_start = Instant::now();
for seq in 0..config.train_packets {
let header = Header {
kind: PacketKind::Train,
run_id,
seq,
count: config.train_packets,
};
encode_header(&mut train, header)?;
send_datagram(socket, &train, peer)?;
}
let sender_span = sender_start.elapsed();
thread::sleep(config.train_settle);
let summary = request_summary(socket, peer, run_id, config)?;
match summary {
Some(body) => {
let sample = CapacitySample {
sent_packets: config.train_packets,
received_packets: body.received_packets,
received_bytes: body.received_bytes,
span: Duration::from_nanos(body.span_nanos),
};
print_capacity_round(round + 1, sample, sender_span);
samples.push(sample);
}
None => {
println!(
"capacity {:>3}: summary timeout after sender_burst={}",
round + 1,
format_duration(sender_span)
);
}
}
thread::sleep(config.train_interval);
}
Ok(samples)
}
fn request_summary(
socket: &UdpSocket,
peer: SocketAddr,
run_id: u64,
config: &ProbeConfig,
) -> Result<Option<SummaryBody>> {
let request = Header {
kind: PacketKind::SummaryRequest,
run_id,
seq: 0,
count: config.train_packets,
};
let mut out = [0_u8; HEADER_LEN];
encode_header(&mut out, request)?;
let attempt_timeout = div_duration(config.echo_timeout, SUMMARY_REQUEST_ATTEMPTS);
for _attempt in 0..SUMMARY_REQUEST_ATTEMPTS {
send_datagram(socket, &out, peer)?;
let deadline = Instant::now() + attempt_timeout;
if let Some(summary) = receive_summary_reply(socket, run_id, deadline)? {
return Ok(Some(summary));
}
}
Ok(None)
}
fn receive_echo_reply(
socket: &UdpSocket,
run_id: u64,
seq: u32,
start: Instant,
timeout: Duration,
) -> Result<Option<Duration>> {
let deadline = start + timeout;
let mut buf = vec![0_u8; MAX_UDP_PACKET_BYTES];
loop {
if !set_timeout_until(socket, deadline)? {
return Ok(None);
}
let (packet_len, _from) = match socket.recv_from(&mut buf) {
Ok(received) => received,
Err(err) if is_timeout(&err) => return Ok(None),
Err(err) if err.kind() == ErrorKind::Interrupted => continue,
Err(err) => return Err(err).wrap_err("receiving echo reply"),
};
let Some(packet) = buf.get(..packet_len) else {
continue;
};
let Ok(header) = decode_header(packet) else {
continue;
};
if header.kind == PacketKind::EchoReply && header.run_id == run_id && header.seq == seq {
return Ok(Some(start.elapsed()));
}
}
}
fn receive_summary_reply(
socket: &UdpSocket,
run_id: u64,
deadline: Instant,
) -> Result<Option<SummaryBody>> {
let mut buf = vec![0_u8; MAX_UDP_PACKET_BYTES];
loop {
if !set_timeout_until(socket, deadline)? {
return Ok(None);
}
let (packet_len, _from) = match socket.recv_from(&mut buf) {
Ok(received) => received,
Err(err) if is_timeout(&err) => return Ok(None),
Err(err) if err.kind() == ErrorKind::Interrupted => continue,
Err(err) => return Err(err).wrap_err("receiving train summary"),
};
let Some(packet) = buf.get(..packet_len) else {
continue;
};
let Ok(header) = decode_header(packet) else {
continue;
};
if header.kind == PacketKind::SummaryReply && header.run_id == run_id {
return Ok(Some(decode_summary_body(packet)?));
}
}
}
fn send_datagram(socket: &UdpSocket, buf: &[u8], target: SocketAddr) -> Result<()> {
let sent = socket
.send_to(buf, target)
.wrap_err_with(|| format!("sending profiling packet to {target}"))?;
if sent != buf.len() {
return Err(eyre!(
"short UDP send to {target}: sent {sent} of {} bytes",
buf.len()
));
}
Ok(())
}
fn set_timeout_until(socket: &UdpSocket, deadline: Instant) -> Result<bool> {
let now = Instant::now();
if now >= deadline {
return Ok(false);
}
socket
.set_read_timeout(Some(deadline.saturating_duration_since(now)))
.wrap_err("setting profiling socket read timeout")?;
Ok(true)
}
fn cleanup_stale_trains(trains: &mut HashMap<u64, TrainAccumulator>, last_cleanup: &mut Instant) {
if last_cleanup.elapsed() < Duration::from_secs(5) {
return;
}
trains.retain(|_run_id, train| train.last_update.elapsed() < STALE_TRAIN_AFTER);
*last_cleanup = Instant::now();
}
fn empty_summary() -> SummaryBody {
SummaryBody {
received_packets: 0,
received_bytes: 0,
span_nanos: 0,
}
}
fn print_latency_summary(sent: u32, samples: &[Duration]) {
match latency_stats(sent, samples) {
Some(stats) => println!(
"latency summary: sent={} received={} loss={:.1}% min={} avg={} max={}",
stats.sent,
stats.received,
stats.loss_ratio * 100.0,
format_duration(stats.min),
format_duration(stats.avg),
format_duration(stats.max)
),
None => println!("latency summary: no replies"),
}
}
fn print_capacity_round(round: u32, sample: CapacitySample, sender_span: Duration) {
let mbps = sample
.mbps()
.map_or_else(|| "n/a".to_owned(), |value| format!("{value:.1} Mbps"));
println!(
"capacity {:>3}: rx={}/{} loss={:.1}% span={} estimate={} sender_burst={}",
round,
sample.received_packets,
sample.sent_packets,
sample.loss_ratio() * 100.0,
format_duration(sample.span),
mbps,
format_duration(sender_span)
);
}
fn print_capacity_summary(samples: &[CapacitySample]) {
let mut estimates = samples
.iter()
.filter_map(|sample| sample.mbps())
.collect::<Vec<f64>>();
if estimates.is_empty() {
println!("capacity summary: no usable samples");
return;
}
estimates.sort_by(f64::total_cmp);
let median_index = estimates.len() / 2;
let median = estimates.get(median_index).copied().unwrap_or(0.0);
let best = estimates.last().copied().unwrap_or(median);
let received = samples
.iter()
.map(|sample| sample.received_packets)
.sum::<u32>();
let sent = samples
.iter()
.map(|sample| sample.sent_packets)
.sum::<u32>();
let loss = if sent == 0 {
0.0
} else {
f64::from(sent.saturating_sub(received)) / f64::from(sent)
};
println!(
"capacity summary: samples={} median={median:.1} Mbps best={best:.1} Mbps aggregate_loss={:.1}%",
estimates.len(),
loss * 100.0
);
}
fn parse_link_local_addr_arg(raw: &str) -> std::result::Result<Ipv6Addr, String> {
parse_link_local_addr(raw).map_err(|err| err.to_string())
}
fn is_timeout(err: &io::Error) -> bool {
matches!(err.kind(), ErrorKind::WouldBlock | ErrorKind::TimedOut)
}
fn div_duration(duration: Duration, divisor: u32) -> Duration {
if divisor == 0 {
return duration;
}
Duration::from_nanos(duration_nanos_u64(duration) / u64::from(divisor))
}
fn format_duration(duration: Duration) -> String {
if duration < Duration::from_millis(1) {
return format!("{:.3} us", duration.as_secs_f64() * 1_000_000.0);
}
format!("{:.3} ms", duration.as_secs_f64() * 1_000.0)
}
fn make_base_run_id() -> u64 {
rand::random()
}
#[cfg(test)]
mod tests {
use std::net::Ipv6Addr;
use clap::Parser;
use super::{Cli, Command, mark_seen};
#[test]
fn parses_probe_options() {
let cli = Cli::try_parse_from([
"link_profile",
"probe",
"--ifname",
"en3",
"--peer",
"fe80::1%en3",
"--train-packets",
"32",
])
.expect("probe options should parse");
match cli.command {
Command::Probe(args) => {
let options = args.into_probe_options();
assert_eq!(options.ifname, "en3");
assert_eq!(
options.peer,
"fe80::1".parse::<Ipv6Addr>().expect("valid IPv6")
);
assert_eq!(options.config.train_packets, 32);
}
Command::Reflect(_) => panic!("expected probe command"),
}
}
#[test]
fn parses_reflect_options() {
let cli = Cli::try_parse_from([
"link_profile",
"reflect",
"--ifname",
"en2",
"--port",
"42000",
])
.expect("reflect options should parse");
match cli.command {
Command::Reflect(options) => {
assert_eq!(options.ifname, "en2");
assert_eq!(options.port, 42_000);
}
Command::Probe(_) => panic!("expected reflect command"),
}
}
#[test]
fn mark_seen_accepts_each_sequence_once() {
let mut seen = vec![false; 2];
assert!(mark_seen(&mut seen, 0));
assert!(!mark_seen(&mut seen, 0));
assert!(mark_seen(&mut seen, 1));
assert!(!mark_seen(&mut seen, 2));
}
}
-49
View File
@@ -1,49 +0,0 @@
use std::net::Ipv6Addr;
use std::time::Duration;
use crate::config::TUN_MTU;
pub const DEFAULT_PROFILE_PORT: u16 = 41_901;
pub const DEFAULT_ECHO_COUNT: u32 = 10;
pub const DEFAULT_ECHO_INTERVAL_MS: u64 = 250;
pub const DEFAULT_ECHO_TIMEOUT_MS: u64 = 500;
pub const DEFAULT_CAPACITY_ROUNDS: u32 = 5;
pub const DEFAULT_TRAIN_PACKETS: u32 = 64;
pub const DEFAULT_TRAIN_INTERVAL_MS: u64 = 1_000;
pub const DEFAULT_TRAIN_SETTLE_MS: u64 = 25;
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct LinkKey {
pub ifname: Box<str>,
pub ifindex: u32,
pub peer_link_local: Ipv6Addr,
}
#[derive(Debug, Clone)]
pub struct ProbeConfig {
pub port: u16,
pub echo_count: u32,
pub echo_interval: Duration,
pub echo_timeout: Duration,
pub capacity_rounds: u32,
pub train_packets: u32,
pub train_payload_bytes: usize,
pub train_interval: Duration,
pub train_settle: Duration,
}
impl Default for ProbeConfig {
fn default() -> Self {
Self {
port: DEFAULT_PROFILE_PORT,
echo_count: DEFAULT_ECHO_COUNT,
echo_interval: Duration::from_millis(DEFAULT_ECHO_INTERVAL_MS),
echo_timeout: Duration::from_millis(DEFAULT_ECHO_TIMEOUT_MS),
capacity_rounds: DEFAULT_CAPACITY_ROUNDS,
train_packets: DEFAULT_TRAIN_PACKETS,
train_payload_bytes: usize::from(TUN_MTU),
train_interval: Duration::from_millis(DEFAULT_TRAIN_INTERVAL_MS),
train_settle: Duration::from_millis(DEFAULT_TRAIN_SETTLE_MS),
}
}
}
-59
View File
@@ -1,59 +0,0 @@
use std::io;
use std::net::IpAddr;
use ipnet::Ipv6Net;
use nix::net::if_::if_nametoindex;
use route_manager::{Route, RouteManager};
use crate::Result;
fn route_destination(prefix: Ipv6Net) -> IpAddr {
IpAddr::V6(prefix.trunc().addr())
}
fn is_overlay_route(route: &Route, prefix: Ipv6Net) -> bool {
route.destination() == route_destination(prefix) && route.prefix() == prefix.prefix_len()
}
fn tun_if_index(tun_ifname: &str) -> io::Result<u32> {
if_nametoindex(tun_ifname).map_err(io::Error::from)
}
pub fn ensure_overlay_route(prefix: Ipv6Net, tun_ifname: &str) -> Result<()> {
let tun_ifindex = tun_if_index(tun_ifname)?;
let desired =
Route::new(route_destination(prefix), prefix.prefix_len()).with_if_index(tun_ifindex);
let mut manager = RouteManager::new()?;
let existing: Vec<Route> = manager
.list()?
.into_iter()
.filter(|route| is_overlay_route(route, prefix))
.collect();
let already_present = existing
.iter()
.any(|route| route.if_index() == Some(tun_ifindex) && route.gateway().is_none());
if already_present {
return Ok(());
}
for route in existing {
manager.delete(&route)?;
}
manager.add(&desired)?;
Ok(())
}
pub fn remove_overlay_route(prefix: Ipv6Net) -> Result<()> {
let mut manager = RouteManager::new()?;
let existing: Vec<Route> = manager
.list()?
.into_iter()
.filter(|route| is_overlay_route(route, prefix))
.collect();
for route in existing {
manager.delete(&route)?;
}
Ok(())
}
-215
View File
@@ -1,215 +0,0 @@
use std::sync::Arc;
use color_eyre::eyre::{Result, WrapErr};
use ipnet::Ipv6Net;
use tokio::{
sync::{mpsc, watch},
task::JoinHandle,
time::{Duration, MissedTickBehavior},
};
use crate::babel::BabelState;
use crate::config::TransportMode;
use crate::daemon::{RoutingStackEvent, StackTaskKind};
use crate::dataplane::{Dataplane, DataplaneConfig, DataplanePublisher, PublishSnapshotError};
use crate::fib::FibBuilder;
use crate::tun::TunDevice;
pub struct RoutingStack {
babel: JoinHandle<crate::Result<()>>,
watcher: JoinHandle<crate::Result<()>>,
state_logger: JoinHandle<()>,
fib_publisher: JoinHandle<crate::Result<()>>,
dataplane_monitor: JoinHandle<()>,
dataplane: Dataplane,
}
impl RoutingStack {
pub fn start(
node_addr: Ipv6Net,
tun: &TunDevice,
udp_port: u16,
transport_mode: TransportMode,
tun_mtu: u16,
tcp_batch_target_bytes: usize,
tcp_socket_buffer_bytes: usize,
state_send: watch::Sender<Arc<BabelState>>,
event_send: mpsc::Sender<RoutingStackEvent>,
) -> Result<Self> {
let (iface_send, iface_recv) = mpsc::channel(32);
let mut state_recv = state_send.subscribe();
let fib_state_recv = state_send.subscribe();
let initial_state = state_send.borrow().clone();
let mut dataplane = Dataplane::spawn(DataplaneConfig {
tun_device: tun.shared_device(),
udp_port,
transport_mode,
tun_mtu,
tcp_batch_target_bytes,
tcp_socket_buffer_bytes,
initial_fib: Arc::new(
FibBuilder::new([node_addr.addr()], tun_mtu).derive(initial_state.as_ref()),
),
})?;
let dataplane_exit = dataplane
.take_exit_receiver()
.ok_or_else(|| color_eyre::eyre::eyre!("dataplane exit receiver missing"))?;
let state_logger = tokio::spawn(async move {
while state_recv.changed().await.is_ok() {
let snapshot = state_recv.borrow_and_update();
tracing::info!(state = ?*snapshot, "babel state snapshot updated");
}
tracing::info!("babel state stream closed");
});
let babel_events = event_send.clone();
let babel = tokio::spawn(async move {
let res = crate::babel(node_addr, iface_recv, state_send).await;
let _ = babel_events
.send(RoutingStackEvent::Exited {
kind: StackTaskKind::Babel,
error: res.as_ref().err().map(ToString::to_string),
})
.await;
res
});
let watcher_events = event_send.clone();
let watcher = tokio::spawn(async move {
let res = crate::watch(iface_send).await;
let _ = watcher_events
.send(RoutingStackEvent::Exited {
kind: StackTaskKind::Watcher,
error: res.as_ref().err().map(ToString::to_string),
})
.await;
res
});
let fib_events = event_send.clone();
let dataplane_publisher = dataplane.publisher();
let fib_publisher = tokio::spawn(async move {
let res =
publish_fib_updates(node_addr, tun_mtu, fib_state_recv, dataplane_publisher).await;
let _ = fib_events
.send(RoutingStackEvent::Exited {
kind: StackTaskKind::FibPublisher,
error: res.as_ref().err().map(ToString::to_string),
})
.await;
res
});
let dataplane_events = event_send;
let dataplane_monitor = tokio::spawn(async move {
let exit = dataplane_exit.await;
let (kind, error) = match exit {
Ok(Ok(())) => (StackTaskKind::Dataplane, None),
Ok(Err(err)) => (StackTaskKind::Dataplane, Some(err)),
Err(err) => (
StackTaskKind::Dataplane,
Some(format!("dataplane exit receiver dropped: {err}")),
),
};
let _ = dataplane_events
.send(RoutingStackEvent::Exited { kind, error })
.await;
});
Ok(Self {
babel,
watcher,
state_logger,
fib_publisher,
dataplane_monitor,
dataplane,
})
}
pub async fn stop(self) -> Result<()> {
let Self {
babel,
watcher,
state_logger,
fib_publisher,
dataplane_monitor,
dataplane,
} = self;
watcher.abort();
if let Ok(res) = watcher.await {
res.wrap_err("stopping interface watcher")?;
}
state_logger.abort();
let _ = state_logger.await;
fib_publisher.abort();
let _ = fib_publisher.await;
dataplane_monitor.abort();
let _ = dataplane_monitor.await;
dataplane.stop().wrap_err("stopping dataplane thread")?;
babel.await?.wrap_err("stopping babeld runtime")?;
Ok(())
}
}
async fn publish_fib_updates(
node_addr: Ipv6Net,
tun_mtu: u16,
mut state_recv: watch::Receiver<Arc<BabelState>>,
publisher: DataplanePublisher,
) -> crate::Result<()> {
let builder = FibBuilder::new([node_addr.addr()], tun_mtu);
let mut pending = Some(Arc::new(builder.derive(state_recv.borrow().as_ref())));
let mut published: Option<Arc<crate::fib::FibSnapshot>> = None;
let mut retry_tick = tokio::time::interval(Duration::from_millis(10));
retry_tick.set_missed_tick_behavior(MissedTickBehavior::Delay);
loop {
if let Some(snapshot) = pending.take() {
let published_snapshot = Arc::clone(&snapshot);
match publisher.try_publish(snapshot) {
Ok(()) => {
published = Some(published_snapshot);
}
Err(PublishSnapshotError::Full(snapshot)) => {
pending = Some(snapshot);
}
Err(PublishSnapshotError::Stopped) => {
return Err(crate::BabbleError::Other(
"dataplane thread stopped".to_owned(),
));
}
}
}
tokio::select! {
changed = state_recv.changed() => {
match changed {
Ok(()) => {
let snapshot = {
let state = state_recv.borrow_and_update();
Arc::new(builder.derive(state.as_ref()))
};
let matches_published = published
.as_ref()
.is_some_and(|current| current.as_ref() == snapshot.as_ref());
let matches_pending = pending
.as_ref()
.is_some_and(|current| current.as_ref() == snapshot.as_ref());
if !matches_published && !matches_pending {
pending = Some(snapshot);
}
}
Err(_) => return Ok(()),
}
}
_ = retry_tick.tick(), if pending.is_some() => {}
}
}
}
-61
View File
@@ -1,61 +0,0 @@
use ipnet::Ipv6Net;
use std::net::Ipv6Addr;
use std::sync::Arc;
use tun_rs::{DeviceBuilder, SyncDevice};
#[cfg(target_os = "linux")]
const DESIRED_TUN_NAME: &str = "exonet";
/// Holds the TUN device open for the lifetime of the daemon.
/// The interface disappears when this is dropped.
pub struct TunDevice {
dev: Arc<SyncDevice>,
ifname: String,
node_addr: Ipv6Net, // TODO: we are only ever gonna install /128 subnets, maybe change to Ipv6Addr in future??
}
impl TunDevice {
pub fn create(node_addr: Ipv6Addr, mtu: u16) -> crate::Result<Self> {
let builder = DeviceBuilder::new().ipv6(node_addr, 128u8).mtu(mtu);
#[cfg(target_os = "linux")]
let builder = builder.name(DESIRED_TUN_NAME);
let dev = builder
.with(|builder| {
builder.packet_information(false);
#[cfg(target_os = "macos")]
{
// Route ownership stays in userspace; do not let tun-rs auto-add routes.
// The IPv6 /128 address itself is still applied by tun-rs.
builder.associate_route(false);
}
})
.build_sync()?;
dev.set_nonblocking(true)?;
let ifname = dev.name()?;
Ok(Self {
dev: Arc::new(dev),
ifname,
node_addr: Ipv6Net::new_assert(node_addr, 128), // TODO: i dont't like the magic numbers, I also don't like wrapping and unwrapping
})
}
pub fn ifname(&self) -> &str {
&self.ifname
}
pub fn node_addr(&self) -> Ipv6Net {
self.node_addr
}
pub fn device(&self) -> &SyncDevice {
self.dev.as_ref()
}
pub fn shared_device(&self) -> Arc<SyncDevice> {
Arc::clone(&self.dev)
}
}
@@ -1,5 +1,5 @@
[package]
name = "exo_pyo3_bindings"
name = "exo_rs"
version = { workspace = true }
edition = { workspace = true }
publish = false
@@ -7,7 +7,7 @@ publish = false
[lib]
doctest = false
path = "src/lib.rs"
name = "exo_pyo3_bindings"
name = "exo_rs"
# "cdylib" needed to produce shared library for Python to import
# "rlib" needed for stub-gen to run
@@ -25,7 +25,7 @@ workspace = true
networking = { workspace = true }
# interop
pyo3 = { version = "0.27.2", features = [
pyo3 = { version = "0.28.3", features = [
# "abi3-py313", # tells pyo3 (and maturin) to build using the stable ABI with minimum Python version 3.13
# "nightly", # enables better-supported GIL integration
"experimental-async", # async support in #[pyfunction] & #[pymethods]
@@ -38,15 +38,15 @@ pyo3 = { version = "0.27.2", features = [
# "ordered-float", "rust_decimal", "smallvec",
# "anyhow", "chrono", "chrono-local", "chrono-tz", "eyre", "jiff-02", "lock_api", "parking-lot", "time", "serde",
] }
pyo3-stub-gen = { version = "0.17.2" }
pyo3-async-runtimes = { version = "0.27.0", features = [
pyo3-stub-gen = { version = "0.22.3" }
pyo3-async-runtimes = { version = "0.28.0", features = [
"attributes",
"tokio-runtime",
"testing",
] }
pyo3-log = "0.13.2"
pyo3-log = "0.13.3"
pidfile-rs = "0.3"
pidfile-rs = { git = "https://github.com/AndreiCravtov/pidfile-rs" }
# macro dependencies
extend = { workspace = true }
File renamed without changes.
@@ -1,10 +1,20 @@
# This file is automatically generated by pyo3_stub_gen
# ruff: noqa: E501, F401
# ruff: noqa: E501, F401, F403, F405
import builtins
import os
import pathlib
import typing
__all__ = [
"AllQueuesFullError",
"FromSwarm",
"Keypair",
"MessageTooLargeError",
"NetworkingHandle",
"NoPeersSubscribedToTopicError",
"Pidfile",
"PidfileError",
]
@typing.final
class AllQueuesFullError(builtins.Exception):
@@ -12,6 +22,29 @@ class AllQueuesFullError(builtins.Exception):
def __repr__(self) -> builtins.str: ...
def __str__(self) -> builtins.str: ...
class FromSwarm:
@typing.final
class Connection(FromSwarm):
__match_args__ = ("peer_id", "connected",)
@property
def peer_id(self) -> builtins.str: ...
@property
def connected(self) -> builtins.bool: ...
def __new__(cls, peer_id: builtins.str, connected: builtins.bool) -> FromSwarm.Connection: ...
@typing.final
class Message(FromSwarm):
__match_args__ = ("origin", "topic", "data",)
@property
def origin(self) -> builtins.str: ...
@property
def topic(self) -> builtins.str: ...
@property
def data(self) -> bytes: ...
def __new__(cls, origin: builtins.str, topic: builtins.str, data: bytes) -> FromSwarm.Message: ...
...
@typing.final
class Keypair:
r"""
@@ -45,6 +78,7 @@ class MessageTooLargeError(builtins.Exception):
@typing.final
class NetworkingHandle:
def __new__(cls, identity: Keypair, bootstrap_peers: typing.Sequence[builtins.str], listen_port: builtins.int) -> NetworkingHandle: ...
def recv(self) -> typing.Awaitable[FromSwarm]: ...
async def gossipsub_subscribe(self, topic: builtins.str) -> builtins.bool:
r"""
Subscribe to a `GossipSub` topic.
@@ -63,7 +97,6 @@ class NetworkingHandle:
If no peers are found that subscribe to this topic, throws `NoPeersSubscribedToTopicError` exception.
"""
async def recv(self) -> PyFromSwarm: ...
@typing.final
class NoPeersSubscribedToTopicError(builtins.Exception):
@@ -77,7 +110,7 @@ class Pidfile:
A PID file protected with a lock.
An instance of `Pidfile` can be used to manage a PID file: create it,
lock it, detect already running daemons. It is backed by [`pidfile`][]
lock it, detect already running daemons. It is backed by [`pidfile`]
functions of `libbsd`/`libutil` which use `flopen` to lock the PID
file.
@@ -107,32 +140,23 @@ class Pidfile:
The file is truncated before writing.
"""
def as_raw_fd(self) -> builtins.int:
r"""
Extracts the raw file descriptor.
This function is typically used to **borrow** an owned file descriptor.
When used in this way, this method does **not** pass ownership of the
raw file descriptor to the caller, and the file descriptor is only
guaranteed to be valid while the original object has not yet been
destroyed.
"""
def close(self) -> None:
r"""
Closes the PID file and releases associated resources.
"""
@typing.final
class PidfileError(builtins.Exception):
def __repr__(self) -> builtins.str: ...
def __str__(self) -> builtins.str: ...
class PyFromSwarm:
@typing.final
class Connection(PyFromSwarm):
__match_args__ = ("peer_id", "connected",)
@property
def peer_id(self) -> builtins.str: ...
@property
def connected(self) -> builtins.bool: ...
def __new__(cls, peer_id: builtins.str, connected: builtins.bool) -> PyFromSwarm.Connection: ...
@typing.final
class Message(PyFromSwarm):
__match_args__ = ("origin", "topic", "data",)
@property
def origin(self) -> builtins.str: ...
@property
def topic(self) -> builtins.str: ...
@property
def data(self) -> bytes: ...
def __new__(cls, origin: builtins.str, topic: builtins.str, data: bytes) -> PyFromSwarm.Message: ...
...
@@ -3,8 +3,8 @@ requires = ["maturin>=1.0,<2.0"]
build-backend = "maturin"
[project]
name = "exo_pyo3_bindings"
version = "0.2.2"
name = "exo_rs"
version = "0.2.16"
description = "Add your description here"
readme = "README.md"
authors = [
@@ -15,14 +15,17 @@ requires-python = ">=3.13"
dependencies = []
[dependency-groups]
dev = ["exo_pyo3_bindings", "pytest>=8.4.0", "pytest-asyncio>=1.0.0"]
dev = ["exo_rs", "pytest>=8.4.0", "pytest-asyncio>=1.0.0"]
[tool.maturin]
#purelib = true
#python-source = "python"
module-name = "exo_pyo3_bindings"
module-name = "exo_rs"
features = ["pyo3/extension-module", "pyo3/experimental-async"]
[tool.pyo3-stub-gen]
generate-init-py = true
[tool.pytest.ini_options]
log_cli = true
log_cli_level = "INFO"
@@ -2,7 +2,7 @@ use pyo3_stub_gen::Result;
fn main() -> Result<()> {
env_logger::Builder::from_env(env_logger::Env::default().filter_or("RUST_LOG", "info")).init();
let stub = exo_pyo3_bindings::stub_info()?;
let stub = exo_rs::stub_info()?;
stub.generate()?;
Ok(())
}
File renamed without changes.
@@ -153,7 +153,7 @@ pub(crate) mod ext {
/// A Python module implemented in Rust. The name of this function must match
/// the `lib.name` setting in the `Cargo.toml`, else Python will not be able to
/// import the module.
#[pymodule(name = "exo_pyo3_bindings")]
#[pymodule(name = "exo_rs", gil_used = true)]
fn main_module(m: &Bound<'_, PyModule>) -> PyResult<()> {
// install logger
pyo3_log::init();
@@ -16,9 +16,7 @@ use pyo3::exceptions::PyRuntimeError;
use pyo3::prelude::{PyModule, PyModuleMethods as _};
use pyo3::types::PyBytes;
use pyo3::{Bound, Py, PyAny, PyErr, PyResult, Python, pymethods};
use pyo3_stub_gen::derive::{
gen_methods_from_python, gen_stub_pyclass, gen_stub_pyclass_complex_enum, gen_stub_pymethods,
};
use pyo3_stub_gen::derive::{gen_stub_pyclass, gen_stub_pyclass_complex_enum, gen_stub_pymethods};
use tokio::sync::{Mutex, mpsc, oneshot};
mod exception {
@@ -138,7 +136,7 @@ struct PyNetworkingHandle {
}
#[gen_stub_pyclass_complex_enum]
#[pyclass]
#[pyclass(name = "FromSwarm")]
enum PyFromSwarm {
Connection {
peer_id: String,
@@ -204,9 +202,11 @@ impl PyNetworkingHandle {
})
}
#[gen_stub(skip)]
#[gen_stub(override_return_type(
type_repr="typing.Awaitable[FromSwarm]", imports=("typing")
))]
fn recv<'py>(&'py self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
let swarm = Arc::clone(&self.swarm);
let swarm = self.swarm.clone();
pyo3_async_runtimes::tokio::future_into_py(py, async move {
swarm
.try_lock()
@@ -297,15 +297,6 @@ impl PyNetworkingHandle {
}
}
pyo3_stub_gen::inventory::submit! {
gen_methods_from_python! {
r#"
class PyNetworkingHandle:
async def recv() -> PyFromSwarm: ...
"#
}
}
pub fn networking_submodule(m: &Bound<'_, PyModule>) -> PyResult<()> {
m.add_class::<exception::PyNoPeersSubscribedToTopicError>()?;
m.add_class::<exception::PyAllQueuesFullError>()?;
@@ -3,7 +3,9 @@ use pyo3::exceptions::PyException;
use pyo3::prelude::{PyModule, PyModuleMethods};
use pyo3::{Bound, PyErr, PyResult, Python, pyclass, pymethods};
use pyo3_stub_gen::derive::{gen_stub_pyclass, gen_stub_pymethods};
use std::fs;
use std::fs::Permissions;
use std::os::fd::{AsRawFd, RawFd};
use std::os::unix::prelude::PermissionsExt;
use std::path::PathBuf;
@@ -36,7 +38,7 @@ impl PyPidfileError {
/// A PID file protected with a lock.
///
/// An instance of `Pidfile` can be used to manage a PID file: create it,
/// lock it, detect already running daemons. It is backed by [`pidfile`][]
/// lock it, detect already running daemons. It is backed by [`pidfile`]
/// functions of `libbsd`/`libutil` which use `flopen` to lock the PID
/// file.
///
@@ -53,7 +55,23 @@ impl PyPidfileError {
/// [`daemon`(3)]: https://linux.die.net/man/3/daemon
#[gen_stub_pyclass]
#[pyclass(name = "Pidfile")]
pub struct PyPidfile(Pidfile);
pub struct PyPidfile(Option<Pidfile>);
impl PyPidfile {
#[inline(always)]
fn get(&self) -> &Pidfile {
self.0
.as_ref()
.expect("cannot use resource after exiting context")
}
#[inline(always)]
fn get_mut(&mut self) -> &mut Pidfile {
self.0
.as_mut()
.expect("cannot use resource after exiting context")
}
}
#[gen_stub_pymethods]
#[pymethods]
@@ -65,17 +83,40 @@ impl PyPidfile {
/// the PID file yet.
#[new]
fn py_new(py: Python, path: PathBuf, mode: u32) -> PyResult<Self> {
Ok(Self(
Pidfile::new(&path, Permissions::from_mode(mode))
.map_err(|e| PyPidfileError(e).into_pyerr(py))?,
))
// create all parent directories if don't exist
if let Some(parent) = path.parent() {
fs::create_dir_all(parent)
.map_err(|e| PyPidfileError(PidfileError::Io(e)).into_pyerr(py))?;
}
let pidfile = Pidfile::new(&path, Permissions::from_mode(mode))
.map_err(|e| PyPidfileError(e).into_pyerr(py))?;
Ok(Self(Some(pidfile)))
}
/// Writes the current process ID to the PID file.
///
/// The file is truncated before writing.
fn write<'py>(&mut self, py: Python<'py>) -> PyResult<()> {
self.0.write().map_err(|e| PyPidfileError(e).into_pyerr(py))
self.get_mut()
.write()
.map_err(|e| PyPidfileError(e).into_pyerr(py))
}
/// Extracts the raw file descriptor.
///
/// This function is typically used to **borrow** an owned file descriptor.
/// When used in this way, this method does **not** pass ownership of the
/// raw file descriptor to the caller, and the file descriptor is only
/// guaranteed to be valid while the original object has not yet been
/// destroyed.
fn as_raw_fd(&self) -> RawFd {
self.get().as_raw_fd()
}
/// Closes the PID file and releases associated resources.
fn close(&mut self) {
self.0 = None;
}
}
File renamed without changes.
@@ -2,12 +2,12 @@ import asyncio
import pytest
from _pytest.capture import CaptureFixture
from exo_pyo3_bindings import (
from exo_rs import (
Keypair,
NetworkingHandle,
NoPeersSubscribedToTopicError,
Pidfile,
PyFromSwarm,
FromSwarm,
)
@@ -39,9 +39,9 @@ async def _await_recv(h: NetworkingHandle):
while True:
event = await h.recv()
match event:
case PyFromSwarm.Connection() as c:
case FromSwarm.Connection() as c:
print(f"PYTHON: connection update: {c}")
case PyFromSwarm.Message() as m:
case FromSwarm.Message() as m:
print(f"PYTHON: message: {m}")
+6 -18
View File
@@ -1,7 +1,7 @@
{ inputs, ... }:
{
perSystem =
{ inputs', self', pkgs, lib, ... }:
{ inputs', pkgs, lib, ... }:
let
# Fenix nightly toolchain with all components
rustToolchain = inputs'.fenix.packages.stable.withComponents [
@@ -55,6 +55,7 @@
];
OPENSSL_NO_VENDOR = "1";
MATURIN_NO_INSTALL_RUST = "1";
# Required for pyo3 tests to find libpython
LD_LIBRARY_PATH = lib.makeLibraryPath [ pkgs.python313 ];
@@ -79,13 +80,13 @@
};
config = {
packages = rec {
packages = {
# Python bindings wheel via maturin
exo_pyo3_bindings = craneLib.buildPackage (
exo-rs = craneLib.buildPackage (
commonArgs
// {
inherit cargoArtifacts;
pname = "exo_pyo3_bindings";
pname = "exo-rs";
nativeBuildInputs = commonArgs.nativeBuildInputs ++ [
pkgs.maturin
@@ -95,7 +96,7 @@
maturin build \
--release \
--manylinux off \
--manifest-path rust/exo_pyo3_bindings/Cargo.toml \
--manifest-path rust/exo_rs/Cargo.toml \
--features "pyo3/extension-module,pyo3/experimental-async" \
--interpreter ${pkgs.python313}/bin/python \
--out dist
@@ -110,19 +111,6 @@
'';
}
);
babblerd-unwrapped = craneLib.buildPackage (
commonArgs // {
inherit cargoArtifacts;
pname = "babblerd-unwrapped";
}
);
babblerd = pkgs.writeShellApplication {
name = "babblerd";
runtimeInputs = [ pkgs.babeld pkgs.iperf3 ];
text = ''
exec ${babblerd-unwrapped}/bin/babblerd "$@"
'';
};
};
checks = {
-38
View File
@@ -1,38 +0,0 @@
#!/usr/bin/env python3
import argparse
import socket
import sys
def main() -> int:
p = argparse.ArgumentParser(
description="IPv6 UDP client with optional explicit source bind"
)
p.add_argument("--dest", required=True, help="Destination IPv6 address")
p.add_argument("--port", type=int, default=45679, help="Destination UDP port")
p.add_argument("--source", help="Optional source IPv6 address to bind to")
p.add_argument("--message", default="hello", help="Payload to send")
p.add_argument(
"--timeout", type=float, default=5.0, help="Receive timeout in seconds"
)
args = p.parse_args()
s = socket.socket(socket.AF_INET6, socket.SOCK_DGRAM)
s.setsockopt(socket.IPPROTO_IPV6, socket.IPV6_V6ONLY, 1)
s.settimeout(args.timeout)
if args.source:
s.bind((args.source, 0, 0, 0))
print(f"local-before-send={s.getsockname()}")
s.sendto(args.message.encode(), (args.dest, args.port, 0, 0))
print(f"sent to=[{args.dest}]:{args.port}")
print(f"local-after-send={s.getsockname()}")
data, peer = s.recvfrom(65535)
print(f"from={peer} data={data!r}")
return 0
if __name__ == "__main__":
raise SystemExit(main())
-35
View File
@@ -1,35 +0,0 @@
#!/usr/bin/env python3
import argparse
import socket
import sys
def main() -> int:
p = argparse.ArgumentParser(
description="IPv6 UDP server bound to a specific local address"
)
p.add_argument(
"--bind", required=True, help="Local IPv6 address to bind to, e.g. fde0:..."
)
p.add_argument("--port", type=int, default=45679, help="UDP port to listen on")
p.add_argument("--reply", default="ok", help="Reply prefix")
args = p.parse_args()
s = socket.socket(socket.AF_INET6, socket.SOCK_DGRAM)
s.setsockopt(socket.IPPROTO_IPV6, socket.IPV6_V6ONLY, 1)
s.bind((args.bind, args.port, 0, 0))
print(f"listening on [{args.bind}]:{args.port}")
print(f"sockname={s.getsockname()}")
data, peer = s.recvfrom(65535)
print(f"from={peer} data={data!r}")
out = args.reply.encode() + b":" + data
s.sendto(out, peer)
print(f"sent={out!r} to={peer}")
return 0
if __name__ == "__main__":
raise SystemExit(main())
+88 -42
View File
@@ -20,7 +20,7 @@ from fastapi.staticfiles import StaticFiles
from hypercorn.asyncio import serve # pyright: ignore[reportUnknownVariableType]
from hypercorn.config import Config
from hypercorn.typing import ASGIFramework
from hypercorn.utils import LifespanTimeoutError
from hypercorn.utils import LifespanTimeoutError, ShutdownError
from loguru import logger
from exo.api.adapters.chat_completions import (
@@ -50,6 +50,8 @@ from exo.api.keepalive import with_sse_keepalive
from exo.api.types import (
AddCustomModelParams,
AdvancedImageParams,
AwaitInstanceReadyMessage,
AwaitInstanceTimeoutMessage,
BenchChatCompletionRequest,
BenchChatCompletionResponse,
BenchImageGenerationResponse,
@@ -344,6 +346,7 @@ class API:
self.app.post("/place_instance")(self.place_instance)
self.app.get("/instance/placement")(self.get_placement)
self.app.get("/instance/previews")(self.get_placement_previews)
self.app.get("/instance/await", response_model=None)(self.await_instance)
self.app.get("/instance/{instance_id}")(self.get_instance)
self.app.delete("/instance/{instance_id}")(self.delete_instance)
self.app.get("/v1/instance-links")(self.list_instance_links)
@@ -633,6 +636,48 @@ class API:
raise HTTPException(status_code=404, detail="Instance not found")
return self.state.instances[instance_id]
async def await_instance(
self,
model_id: ModelId,
timeout_seconds: float = Query(default=0.0, ge=0.0, le=300.0),
) -> StreamingResponse:
_sleep = 0.1
async def _stream() -> AsyncGenerator[str, None]:
deadline = (
None if timeout_seconds == 0 else anyio.current_time() + timeout_seconds
)
while True:
for instance in self.state.instances.values():
if instance.shard_assignments.model_id == model_id:
payload = AwaitInstanceReadyMessage(instance=instance)
yield f"data: {payload.model_dump_json()}\n\n"
return
if deadline is None:
await anyio.sleep(_sleep)
else:
remaining = deadline - anyio.current_time()
if remaining <= 0:
payload = AwaitInstanceTimeoutMessage(
message=f"No instance found for model {model_id}"
)
yield f"data: {payload.model_dump_json()}\n\n"
return
await anyio.sleep(min(_sleep, remaining))
return StreamingResponse(
with_sse_keepalive(_stream()),
media_type="text/event-stream",
headers={
"Cache-Control": "no-cache",
"Connection": "close",
"X-Accel-Buffering": "no",
},
)
async def delete_instance(self, instance_id: InstanceId) -> DeleteInstanceResponse:
if instance_id not in self.state.instances:
raise HTTPException(status_code=404, detail="Instance not found")
@@ -761,6 +806,8 @@ class API:
if isinstance(chunk, PrefillProgressChunk):
continue
sampler.mark_prefill_done()
if chunk.finish_reason == "error":
raise HTTPException(
status_code=500,
@@ -871,10 +918,8 @@ class API:
) -> ChatCompletionResponse | StreamingResponse:
"""OpenAI Chat Completions API - adapter."""
task_params = await chat_request_to_text_generation(payload)
resolved_model = await self._resolve_and_validate_text_model(
ModelId(task_params.model)
)
task_params = task_params.model_copy(update={"model": resolved_model})
validated_model = await self._validate_model_has_instance(task_params.model)
task_params = task_params.model_copy(update={"model": validated_model})
command = await self._send_text_generation_with_images(task_params)
@@ -906,10 +951,10 @@ class API:
self, payload: BenchChatCompletionRequest
) -> BenchChatCompletionResponse | StreamingResponse:
task_params = await chat_request_to_text_generation(payload)
resolved_model = await self._resolve_and_validate_text_model(
validated_model = await self._validate_model_has_instance(
ModelId(task_params.model)
)
task_params = task_params.model_copy(update={"model": resolved_model})
task_params = task_params.model_copy(update={"model": validated_model})
task_params = task_params.model_copy(
update={
@@ -939,8 +984,10 @@ class API:
return await self._collect_text_generation_with_stats(command.command_id)
async def _resolve_and_validate_text_model(self, model_id: ModelId) -> ModelId:
"""Validate a text model exists and return the resolved model ID.
async def _validate_model_has_instance(self, model_id: ModelId) -> ModelId:
"""Validate a model has an active instance.
If the model isn't even downloaded, triggers notification to user to download model.
Raises HTTPException 404 if no instance is found for the model.
"""
@@ -948,30 +995,21 @@ class API:
instance.shard_assignments.model_id == model_id
for instance in self.state.instances.values()
):
await self._trigger_notify_user_to_download_model(model_id)
# Check if model is actually downloaded
model_is_downloaded = any(
isinstance(download, DownloadCompleted)
and download.shard_metadata.model_card.model_id == model_id
for node_downloads in self.state.downloads.values()
for download in node_downloads
)
if not model_is_downloaded:
await self._trigger_notify_user_to_download_model(model_id)
raise HTTPException(
status_code=404,
detail=f"No instance found for model {model_id}",
status_code=404, detail=f"No instance found for model {model_id}"
)
return model_id
async def _validate_image_model(self, model: ModelId) -> ModelId:
"""Validate model exists and return resolved model ID.
Raises HTTPException 404 if no instance is found for the model.
"""
model_card = await ModelCard.load(model)
resolved_model = model_card.model_id
if not any(
instance.shard_assignments.model_id == resolved_model
for instance in self.state.instances.values()
):
await self._trigger_notify_user_to_download_model(resolved_model)
raise HTTPException(
status_code=404, detail=f"No instance found for model {resolved_model}"
)
return resolved_model
def stream_events(self) -> StreamingResponse:
def _generate_json_array(events: Iterable[Event]) -> Iterable[str]:
yield "["
@@ -1024,7 +1062,9 @@ class API:
"""
payload = payload.model_copy(
update={
"model": await self._validate_image_model(ModelId(payload.model)),
"model": await self._validate_model_has_instance(
ModelId(payload.model)
),
"advanced_params": _ensure_seed(payload.advanced_params),
}
)
@@ -1292,7 +1332,9 @@ class API:
) -> BenchImageGenerationResponse:
payload = payload.model_copy(
update={
"model": await self._validate_image_model(ModelId(payload.model)),
"model": await self._validate_model_has_instance(
ModelId(payload.model)
),
"stream": False,
"partial_images": 0,
"advanced_params": _ensure_seed(payload.advanced_params),
@@ -1328,7 +1370,7 @@ class API:
advanced_params: AdvancedImageParams | None,
) -> ImageEdits:
"""Prepare and send an image edits command with chunked image upload."""
resolved_model = await self._validate_image_model(model)
validated_model = await self._validate_model_has_instance(model)
advanced_params = _ensure_seed(advanced_params)
image_content = await image.read()
@@ -1347,7 +1389,7 @@ class API:
image_data="",
total_input_chunks=total_chunks,
prompt=prompt,
model=resolved_model,
model=validated_model,
n=n,
size=size,
response_format=response_format,
@@ -1368,7 +1410,7 @@ class API:
await self._send(
SendInputChunk(
chunk=InputImageChunk(
model=resolved_model,
model=validated_model,
command_id=command.command_id,
data=chunk_data,
chunk_index=chunk_index,
@@ -1492,10 +1534,10 @@ class API:
) -> ClaudeMessagesResponse | StreamingResponse:
"""Claude Messages API - adapter."""
task_params = await claude_request_to_text_generation(payload)
resolved_model = await self._resolve_and_validate_text_model(
validated_model = await self._validate_model_has_instance(
ModelId(task_params.model)
)
task_params = task_params.model_copy(update={"model": resolved_model})
task_params = task_params.model_copy(update={"model": validated_model})
command = await self._send_text_generation_with_images(task_params)
@@ -1530,8 +1572,8 @@ class API:
) -> ResponsesResponse | StreamingResponse:
"""OpenAI Responses API."""
task_params = await responses_request_to_text_generation(payload)
resolved_model = await self._resolve_and_validate_text_model(task_params.model)
task_params = task_params.model_copy(update={"model": resolved_model})
validated_model = await self._validate_model_has_instance(task_params.model)
task_params = task_params.model_copy(update={"model": validated_model})
command = await self._send_text_generation_with_images(task_params)
@@ -1573,10 +1615,10 @@ class API:
body = await request.body()
payload = OllamaChatRequest.model_validate_json(body)
task_params = ollama_request_to_text_generation(payload)
resolved_model = await self._resolve_and_validate_text_model(
validated_model = await self._validate_model_has_instance(
ModelId(task_params.model)
)
task_params = task_params.model_copy(update={"model": resolved_model})
task_params = task_params.model_copy(update={"model": validated_model})
command = await self._send_text_generation_with_images(task_params)
@@ -1609,10 +1651,10 @@ class API:
body = await request.body()
payload = OllamaGenerateRequest.model_validate_json(body)
task_params = ollama_generate_request_to_text_generation(payload)
resolved_model = await self._resolve_and_validate_text_model(
validated_model = await self._validate_model_has_instance(
ModelId(task_params.model)
)
task_params = task_params.model_copy(update={"model": resolved_model})
task_params = task_params.model_copy(update={"model": validated_model})
command = await self._send_text_generation_with_images(task_params)
@@ -1914,6 +1956,10 @@ class API:
cfg,
shutdown_trigger=ev.wait,
)
if not ev.is_set():
raise ShutdownError(
"Server exited without shutdown trigger - exiting abnormally"
)
except LifespanTimeoutError as e:
logger.warning(
"Graceful server shutdown timed out, some connections forcebly closed"
+2
View File
@@ -1,5 +1,7 @@
from .api import AddCustomModelParams as AddCustomModelParams
from .api import AdvancedImageParams as AdvancedImageParams
from .api import AwaitInstanceReadyMessage as AwaitInstanceReadyMessage
from .api import AwaitInstanceTimeoutMessage as AwaitInstanceTimeoutMessage
from .api import BenchChatCompletionRequest as BenchChatCompletionRequest
from .api import BenchChatCompletionResponse as BenchChatCompletionResponse
from .api import BenchImageGenerationResponse as BenchImageGenerationResponse
+26
View File
@@ -186,6 +186,12 @@ class NodePowerStats(BaseModel, frozen=True):
node_id: NodeId
samples: int
avg_sys_power: float
# Per-phase breakdown. Populated only when the caller marks a phase
# boundary (e.g. prefill -> generation); None otherwise.
prefill_avg_sys_power: float | None = None
generation_avg_sys_power: float | None = None
prefill_energy_joules: float | None = None
generation_energy_joules: float | None = None
class PowerUsage(BaseModel, frozen=True):
@@ -193,6 +199,16 @@ class PowerUsage(BaseModel, frozen=True):
nodes: list[NodePowerStats]
total_avg_sys_power_watts: float
total_energy_joules: float
# Split between the prefill (prompt-processing) phase and the
# generation/decode phase. Populated only when the caller marks a phase
# boundary; None otherwise. The two phase energies should sum to
# approximately `total_energy_joules` (modulo interpolation rounding).
prefill_seconds: float | None = None
generation_seconds: float | None = None
prefill_energy_joules: float | None = None
generation_energy_joules: float | None = None
prefill_avg_sys_power_watts: float | None = None
generation_avg_sys_power_watts: float | None = None
class BenchChatCompletionResponse(ChatCompletionResponse):
@@ -291,6 +307,16 @@ class DeleteInstanceResponse(BaseModel):
instance_id: InstanceId
class AwaitInstanceReadyMessage(BaseModel):
type: Literal["ready"] = "ready"
instance: Instance
class AwaitInstanceTimeoutMessage(BaseModel):
type: Literal["timeout"] = "timeout"
message: str
class CancelCommandResponse(BaseModel):
message: str
command_id: CommandId
+11
View File
@@ -15,6 +15,10 @@ from exo.download.download_utils import (
resolve_existing_model,
)
from exo.download.shard_downloader import ShardDownloader
from exo.routing.event_router import (
EventRouterBrokenResourceError,
EventRouterClosedResourceError,
)
from exo.shared.constants import EXO_DEFAULT_MODELS_DIR, EXO_MODELS_READ_ONLY_DIRS
from exo.shared.models import model_cards
from exo.shared.models.model_cards import ModelId
@@ -139,7 +143,14 @@ class DownloadCoordinator:
async with self._tg as tg:
tg.start_soon(self._command_processor)
tg.start_soon(self._emit_existing_download_progress)
except* (EventRouterBrokenResourceError, EventRouterClosedResourceError):
# Event router has been closed (try-star syntax handles error groups)
pass
finally:
# don't forget to clean up resources
self.download_command_receiver.close()
self.event_sender.close()
self._stopped.set()
async def shutdown(self) -> None:
+66 -19
View File
@@ -8,6 +8,9 @@ from dataclasses import dataclass, field
from typing import Self
import anyio
from anyio.lowlevel import checkpoint as anyio_checkpoint
from daemon import DaemonContext # pyright: ignore[reportMissingTypeStubs]
from exo_rs import Pidfile, PidfileError
from loguru import logger
from pydantic import PositiveInt
@@ -18,13 +21,12 @@ from exo.download.impl_shard_downloader import exo_shard_downloader
from exo.master.main import Master
from exo.routing.event_router import EventRouter
from exo.routing.router import Router, get_node_id_keypair
from exo.shared.constants import EXO_DEFAULT_MODELS_DIR, EXO_LOG
from exo.shared.constants import EXO_DEFAULT_MODELS_DIR, EXO_LOG, EXO_PID_FILE
from exo.shared.election import Election, ElectionResult
from exo.shared.logging import logger_cleanup, logger_setup
from exo.shared.types.common import NodeId, SessionId
from exo.utils import STDIO_FDS
from exo.utils.channels import Receiver, channel
from exo.utils.daemon import detach_stdio_to_devnull
from exo.utils.pidfile import PidfileLockError, acquire_exo_pidfile
from exo.utils.pydantic_ext import FrozenModel
from exo.utils.task_group import TaskGroup
from exo.worker.main import Worker
@@ -190,7 +192,7 @@ class Node:
# - Shut down and re-create the API
if result.is_new_master:
await anyio.sleep(0)
await anyio_checkpoint()
self.event_router.shutdown()
self.event_router = EventRouter(
result.session_id,
@@ -203,7 +205,10 @@ class Node:
result.session_id.master_node_id == self.node_id
and self.master is not None
):
logger.info("Node elected Master")
assert not result.is_new_master, (
"cannot be new master if we remain master"
)
logger.info("Node elected Master - maintaining self")
elif (
result.session_id.master_node_id == self.node_id
and self.master is None
@@ -270,14 +275,60 @@ class Node:
def main():
# Exit early if no PID file (not compatible with double-for daemonization yet)
try:
pidfile = acquire_exo_pidfile()
except PidfileLockError as exception:
print(exception, file=sys.stderr)
raise SystemExit(1) from exception
# Parse args first => --help or bad args don't require PID-locking
args = Args.parse()
# Exit early if cannot acquire PID file
try:
pidfile = Pidfile(EXO_PID_FILE, 0o0600)
except PidfileError as e:
print(e, file=sys.stderr)
raise SystemExit(1) from e
try:
if args.legacy_daemon:
# keep stdio backed by explicit /dev/null streams. multiprocessing spawn expects
# valid stdio FDs; letting DaemonContext close/reopen them can break runner startup.
for stream in (sys.stdout, sys.stderr, sys.__stdout__, sys.__stderr__):
if stream is not None:
stream.flush()
stdin = open(os.devnull, "r") # noqa: SIM115
stdout = open(os.devnull, "w") # noqa: SIM115
stderr = open(os.devnull, "w") # noqa: SIM115
with DaemonContext(
detach_process=True,
files_preserve=[pidfile.as_raw_fd()],
stdin=stdin,
stdout=stdout,
stderr=stderr,
):
# cleanup loose file descriptors (as long as they aren't stdio)
for f in (
f for f in (stdin, stdout, stderr) if f.fileno() not in STDIO_FDS
):
f.close()
# 1) if daemonizing => fork then write PID
try:
pidfile.write()
except PidfileError as e:
print(e, file=sys.stderr)
raise SystemExit(1) from e
main_inner(args)
else:
# 2) otherwise => just write PID
try:
pidfile.write()
except PidfileError as e:
print(e, file=sys.stderr)
raise SystemExit(1) from e
main_inner(args)
finally:
pidfile.close()
def main_inner(args: "Args"):
soft, hard = resource.getrlimit(resource.RLIMIT_NOFILE)
target = min(max(soft, 65535), hard)
resource.setrlimit(resource.RLIMIT_NOFILE, (target, hard))
@@ -286,9 +337,6 @@ def main():
# TODO: Refactor the current verbosity system
logger_setup(EXO_LOG, args.verbosity)
if args.no_stdio:
detach_stdio_to_devnull()
logger.info("Detached stdio to /dev/null")
logger.info(f"{'=' * 40}")
logger.info(f"Starting EXO | pid={os.getpid()}")
@@ -324,7 +372,6 @@ def main():
finally:
logger.info("EXO Shutdown complete")
logger_cleanup()
del pidfile
class Args(FrozenModel):
@@ -338,7 +385,7 @@ class Args(FrozenModel):
offline: bool = os.getenv("EXO_OFFLINE", "false").lower() == "true"
no_batch: bool = False
fast_synch: bool | None = None # None = auto, True = force on, False = force off
no_stdio: bool = False
legacy_daemon: bool = False
bootstrap_peers: list[str] = []
libp2p_port: int
@@ -399,9 +446,9 @@ class Args(FrozenModel):
help="Disable continuous batching, use sequential generation",
)
parser.add_argument(
"--no-stdio",
"--legacy-daemon",
action="store_true",
help="Detach stdin/stdout/stderr to /dev/null after logging is configured",
help="Run as a legacy SysV-style background daemon using double-fork daemonization",
)
parser.add_argument(
"--bootstrap-peers",
+16 -3
View File
@@ -11,6 +11,10 @@ from exo.master.placement import (
place_instance,
)
from exo.master.placement_utils import find_ip_prioritised
from exo.routing.event_router import (
EventRouterBrokenResourceError,
EventRouterClosedResourceError,
)
from exo.shared.apply import apply
from exo.shared.constants import EXO_EVENT_LOG_DIR, EXO_TRACING_ENABLED
from exo.shared.types.commands import (
@@ -151,6 +155,9 @@ class Master:
tg.start_soon(self._event_processor)
tg.start_soon(self._command_processor)
tg.start_soon(self._plan)
except* (EventRouterBrokenResourceError, EventRouterClosedResourceError):
# Event router has been closed (try-star syntax handles error groups)
pass
finally:
self._event_log.close()
self.global_event_sender.close()
@@ -174,6 +181,7 @@ class Master:
case TestCommand():
pass
case TextGeneration():
# set-difference => prefill-only nodes
prefill_only: set[InstanceId] = set()
for link in self.state.instance_links.values():
prefill_only.update(link.prefill_instances)
@@ -181,11 +189,13 @@ class Master:
prefill_only.difference_update(link.decode_instances)
for instance in self.state.instances.values():
# NON-prefill-only instances matching the model ID
if (
instance.shard_assignments.model_id
== command.task_params.model
and instance.instance_id not in prefill_only
):
# count in-flight tasks of that instance
in_flight = {TaskStatus.Pending, TaskStatus.Running}
task_count = sum(
1
@@ -197,6 +207,7 @@ class Master:
task_count
)
# there are no NON-prefill-only instances matching this model ID
if not instance_task_counts:
raise ValueError(
f"No instance found for model {command.task_params.model}"
@@ -448,7 +459,9 @@ class Master:
self._event_log.read_range(command.since_idx, end),
start=command.since_idx,
):
await self._send_event(IndexedEvent(idx=i, event=event))
await self._send_indexed_event(
IndexedEvent(idx=i, event=event)
)
for event in generated_events:
await self.event_sender.send(event)
except ValueError as e:
@@ -506,10 +519,10 @@ class Master:
self.state = apply(self.state, indexed)
self._event_log.append(event)
await self._send_event(indexed)
await self._send_indexed_event(indexed)
# This function is re-entrant, take care!
async def _send_event(self, event: IndexedEvent):
async def _send_indexed_event(self, event: IndexedEvent):
# Convenience method since this line is ugly
await self.global_event_sender.send(
GlobalForwarderEvent(
+2 -2
View File
@@ -1,4 +1,4 @@
from exo_pyo3_bindings import PyFromSwarm
from exo_rs import FromSwarm
from exo.shared.types.common import NodeId
from exo.utils.pydantic_ext import FrozenModel
@@ -11,5 +11,5 @@ class ConnectionMessage(FrozenModel):
connected: bool
@classmethod
def from_update(cls, update: PyFromSwarm.Connection) -> "ConnectionMessage":
def from_update(cls, update: FromSwarm.Connection) -> "ConnectionMessage":
return cls(node_id=NodeId(update.peer_id), connected=update.connected)
+21 -2
View File
@@ -15,11 +15,30 @@ from exo.shared.types.events import (
IndexedEvent,
LocalForwarderEvent,
)
from exo.utils import channels
from exo.utils.channels import Receiver, Sender, channel
from exo.utils.event_buffer import OrderedBuffer
from exo.utils.task_group import TaskGroup
class EventRouterClosedResourceError(ClosedResourceError):
pass
class EventRouterBrokenResourceError(BrokenResourceError):
pass
# Event Router is created and destroyed before consumers of its channels are,
# hence its nice to have tagged errors for event-router channels being closed
#
# so consumers can catch specifically these errors, rather than the generic ones
_ERROR_CFG = channels.ErrorOverride(
closed_resource_error=EventRouterClosedResourceError,
broken_resource_error=EventRouterBrokenResourceError,
)
@dataclass
class EventRouter:
session_id: SessionId
@@ -64,7 +83,7 @@ class EventRouter:
await self.external_outbound.send(event)
def sender(self) -> Sender[Event]:
send, recv = channel[Event]()
send, recv = channel[Event](error_override_config=_ERROR_CFG)
if self._tg.is_running():
self._tg.start_soon(self._ingest, SystemId(), recv)
else:
@@ -73,7 +92,7 @@ class EventRouter:
def receiver(self) -> Receiver[IndexedEvent]:
assert not self._tg.is_running()
send, recv = channel[IndexedEvent]()
send, recv = channel[IndexedEvent](error_override_config=_ERROR_CFG)
self.internal_outbound.append(send)
return recv
+4 -4
View File
@@ -12,13 +12,13 @@ from anyio import (
move_on_after,
sleep_forever,
)
from exo_pyo3_bindings import (
from exo_rs import (
AllQueuesFullError,
FromSwarm,
Keypair,
MessageTooLargeError,
NetworkingHandle,
NoPeersSubscribedToTopicError,
PyFromSwarm,
)
from filelock import FileLock
from loguru import logger
@@ -191,7 +191,7 @@ class Router:
from_swarm = await self._net.recv()
logger.debug(from_swarm)
match from_swarm:
case PyFromSwarm.Message(origin, topic, data):
case FromSwarm.Message(origin, topic, data):
logger.trace(
f"Received message on {topic} from {origin} with payload {data}"
)
@@ -202,7 +202,7 @@ class Router:
continue
router = self.topic_routers[topic]
await router.publish_bytes(data)
case PyFromSwarm.Connection():
case FromSwarm.Connection():
message = ConnectionMessage.from_update(from_swarm)
logger.trace(
f"Received message on connection_messages with payload {message}"
+1 -1
View File
@@ -46,7 +46,7 @@ class _InterceptHandler(logging.Handler):
def logger_setup(log_file: Path | None, verbosity: int = 0):
"""Set up logging for this process - formatting, file handles, verbosity and output"""
logging.getLogger("exo_pyo3_bindings").setLevel(logging.WARNING)
logging.getLogger("exo_rs").setLevel(logging.WARNING)
logging.getLogger("httpx").setLevel(logging.WARNING)
logging.getLogger("httpcore").setLevel(logging.WARNING)
+5
View File
@@ -2,6 +2,11 @@ from typing import Any, Type
from .phantom import PhantomData
STDIN_FD = 0
STDOUT_FD = 1
STDERR_FD = 2
STDIO_FDS = (STDIN_FD, STDOUT_FD, STDERR_FD)
def ensure_type[T](obj: Any, expected_type: Type[T]) -> T: # type: ignore
if not isinstance(obj, expected_type):
+4 -5
View File
@@ -25,10 +25,9 @@ from anyio import (
from anyio.abc import TaskStatus
from loguru import logger
from exo.utils import STDERR_FD, STDIO_FDS, STDOUT_FD
from exo.utils.channels import Receiver, Sender, channel
_STDOUT_FD = 1
_STDERR_FD = 2
_READ_CHUNK_SIZE = 64 * 1024
_JOIN_GRACE_SECONDS = 3.0
_TERMINATE_GRACE_SECONDS = 5.0
@@ -256,11 +255,11 @@ def _run_with_captured_stdio(
stderr_fd = stderr.detach()
try:
os.dup2(stdout_fd, _STDOUT_FD)
os.dup2(stderr_fd, _STDERR_FD)
os.dup2(stdout_fd, STDOUT_FD)
os.dup2(stderr_fd, STDERR_FD)
finally:
for fd in (stdout_fd, stderr_fd):
if fd not in (_STDOUT_FD, _STDERR_FD):
if fd not in STDIO_FDS:
_close_fd(fd)
faulthandler.enable(file=sys.stderr, all_threads=True)
+155 -8
View File
@@ -1,13 +1,16 @@
import contextlib
import multiprocessing as mp
from dataclasses import dataclass, field
from functools import wraps
from inspect import iscoroutinefunction
from math import inf
from multiprocessing.synchronize import Event
from queue import Empty, Full
from types import TracebackType
from typing import Any, Self
from types import CoroutineType, TracebackType
from typing import Any, Callable, NoReturn, Self, cast, overload, override
from anyio import (
BrokenResourceError,
CapacityLimiter,
ClosedResourceError,
EndOfStream,
@@ -20,35 +23,172 @@ from anyio.streams.memory import (
from anyio.streams.memory import (
MemoryObjectSendStream as AnyioSender,
)
from anyio.streams.memory import (
MemoryObjectStreamState,
)
from anyio.streams.memory import (
MemoryObjectStreamState as AnyioState,
)
@dataclass(eq=False)
class ErrorOverride:
closed_resource_error: type[ClosedResourceError] = field(
default=ClosedResourceError,
)
broken_resource_error: type[BrokenResourceError] = field(
default=BrokenResourceError,
)
end_of_stream: type[EndOfStream] = field(
default=EndOfStream,
)
would_block: type[WouldBlock] = field(
default=WouldBlock,
)
@overload
def patch[**P, R](
self,
fn: Callable[P, CoroutineType[Any, Any, R]],
/,
) -> Callable[P, CoroutineType[Any, Any, R]]: ...
@overload
def patch[**P, R](
self,
fn: Callable[P, R],
/,
) -> Callable[P, R]: ...
def patch[**P, R](self, fn: Callable[P, Any], /) -> Callable[P, Any]:
"""
Returns a function with all these exceptions replaced by their overrides
"""
if iscoroutinefunction(fn):
async_fn = cast(Callable[P, CoroutineType[Any, Any, R]], fn)
@wraps(async_fn)
async def async_wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
try:
return await async_fn(*args, **kwargs)
except ClosedResourceError as e:
self._raise_replace(self.closed_resource_error, e)
except BrokenResourceError as e:
self._raise_replace(self.broken_resource_error, e)
except EndOfStream as e:
self._raise_replace(self.end_of_stream, e)
except WouldBlock as e:
self._raise_replace(self.would_block, e)
return async_wrapper
else:
sync_fn = cast(Callable[P, R], fn)
@wraps(sync_fn)
def sync_wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
try:
return sync_fn(*args, **kwargs)
except ClosedResourceError as e:
self._raise_replace(self.closed_resource_error, e)
except BrokenResourceError as e:
self._raise_replace(self.broken_resource_error, e)
except EndOfStream as e:
self._raise_replace(self.end_of_stream, e)
except WouldBlock as e:
self._raise_replace(self.would_block, e)
return sync_wrapper
@staticmethod
def _raise_replace(replacement: type[BaseException], e: BaseException) -> NoReturn:
if isinstance(e, replacement):
raise
raise replacement() from e
class Sender[T](AnyioSender[T]):
def __init__(
self,
state: MemoryObjectStreamState[T],
error_override_config: ErrorOverride | None,
):
super().__init__(_state=state)
# patch the methods we want to override errors for
#
# NOTE: it is very important that new methods which are added,
# and which can throw, are patched in this block
if (e := error_override_config) is not None:
# new methods of this class
self.clone_receiver = e.patch(self.clone_receiver)
# overridden methods
self.clone = e.patch(self.clone)
# parent methods
self.send_nowait = e.patch(self.send_nowait)
self.send = e.patch(self.send)
self.close = e.patch(self.close)
self.aclose = e.patch(self.aclose)
self.statistics = e.patch(self.statistics)
self.err_config = error_override_config
@override
def clone(self) -> "Sender[T]":
if self._closed:
raise ClosedResourceError
return Sender(_state=self._state)
return Sender(self._state, self.err_config)
def clone_receiver(self) -> "Receiver[T]":
"""Constructs a Receiver using a Senders shared state - similar to calling Receiver.clone() without needing the receiver"""
if self._closed:
raise ClosedResourceError
return Receiver(_state=self._state)
return Receiver(self._state, self.err_config)
class Receiver[T](AnyioReceiver[T]):
def __init__(
self,
state: MemoryObjectStreamState[T],
error_override_config: ErrorOverride | None,
):
super().__init__(_state=state)
# patch the methods we want to override errors for
#
# NOTE: it is very important that new methods which are added,
# and which can throw, are patched in this block
if (e := error_override_config) is not None:
# new methods of this class
self.clone_sender = e.patch(self.clone_sender)
self.collect = e.patch(self.collect)
self.receive_at_least = e.patch(self.receive_at_least)
# overridden methods
self.clone = e.patch(self.clone)
# parent methods
self.receive_nowait = e.patch(self.receive_nowait)
self.receive = e.patch(self.receive)
self.close = e.patch(self.close)
self.aclose = e.patch(self.aclose)
self.statistics = e.patch(self.statistics)
self.err_config = error_override_config
@override
def clone(self) -> "Receiver[T]":
if self._closed:
raise ClosedResourceError
return Receiver(_state=self._state)
return Receiver(self._state, self.err_config)
def clone_sender(self) -> Sender[T]:
"""Constructs a Sender using a Receivers shared state - similar to calling Sender.clone() without needing the sender"""
if self._closed:
raise ClosedResourceError
return Sender(_state=self._state)
return Sender(self._state, self.err_config)
def collect(self) -> list[T]:
"""Collect all currently available items from this receiver"""
@@ -70,6 +210,7 @@ class Receiver[T](AnyioReceiver[T]):
out.extend(self.collect())
return out
@override
def __enter__(self) -> Self:
return self
@@ -285,11 +426,17 @@ class MpReceiver[T]:
class channel[T]: # noqa: N801
"""Create a pair of asynchronous channels for communicating within the same process"""
def __new__(cls, max_buffer_size: float = inf) -> tuple[Sender[T], Receiver[T]]:
def __new__(
cls,
max_buffer_size: float = inf,
error_override_config: ErrorOverride | None = None,
) -> tuple[Sender[T], Receiver[T]]:
if max_buffer_size != inf and not isinstance(max_buffer_size, int):
raise ValueError("max_buffer_size must be either an integer or math.inf")
state = AnyioState[T](max_buffer_size)
return Sender(_state=state), Receiver(_state=state)
return Sender(state, error_override_config), Receiver(
state, error_override_config
)
class mp_channel[T]: # noqa: N801
-28
View File
@@ -1,28 +0,0 @@
import os
import sys
_STDIN_FD = 0
_STDOUT_FD = 1
_STDERR_FD = 2
def detach_stdio_to_devnull() -> None:
"""Redirect process stdio file descriptors to /dev/null."""
for stream in (sys.stdout, sys.stderr, sys.__stdout__, sys.__stderr__):
if stream is not None:
stream.flush()
stdin_fd = os.open(os.devnull, os.O_RDONLY)
stdout_fd = os.open(os.devnull, os.O_WRONLY)
stderr_fd = os.open(os.devnull, os.O_WRONLY)
try:
# dup2 closes the target fd first, but leaves the source fd open.
os.dup2(stdin_fd, _STDIN_FD)
os.dup2(stdout_fd, _STDOUT_FD)
os.dup2(stderr_fd, _STDERR_FD)
finally:
for fd in (stdin_fd, stdout_fd, stderr_fd):
if fd not in (_STDIN_FD, _STDOUT_FD, _STDERR_FD):
os.close(fd)
@@ -630,6 +630,14 @@ class InfoGatherer:
f"MacMon failed with return code {e.returncode}: {stderr_msg}"
)
self._tg.start_soon(self._monitor_memory_usage, 1)
except ProcessLookupError:
# usually throws by the process' context manager on exit
# when we ctrl+c, hence usually should be ignored;
# if anything else throws it, we explicitly don't care:
# process is dead anyways ;)
logger.warning(
"Macmon process not found - shutting down macmon monitor"
)
except Exception as e:
logger.opt(exception=e).warning("Error in macmon monitor")
self._tg.start_soon(self._monitor_memory_usage, 1)
-28
View File
@@ -1,28 +0,0 @@
from __future__ import annotations
import os
from typing import Final
from exo_pyo3_bindings import Pidfile, PidfileError
from exo.shared.constants import EXO_PID_FILE
_PIDFILE_MODE: Final = 0o600
class PidfileLockError(RuntimeError):
pass
def acquire_exo_pidfile() -> Pidfile:
path = EXO_PID_FILE
os.makedirs(os.path.dirname(path), exist_ok=True)
try:
pidfile = Pidfile(path, _PIDFILE_MODE)
pidfile.write()
except (OSError, PidfileError) as exception:
raise PidfileLockError(
f"Failed to acquire EXO pidfile at {path}: {exception}"
) from exception
return pidfile
+122
View File
@@ -24,6 +24,7 @@ class PowerSampler:
] = defaultdict(list)
self._start_time: float | None = None
self._stopped = False
self._prefill_done_at: float | None = None
def _take_sample(self, t_rel: float | None = None) -> None:
assert self._start_time is not None
@@ -38,14 +39,35 @@ class PowerSampler:
await anyio.sleep(self._interval)
self._take_sample()
def mark_prefill_done(self) -> None:
"""Anchor the prefill→generation boundary on a fresh sample.
Idempotent. Safe to call before `run()`; boundary then lands at t=0.
"""
if self._prefill_done_at is not None:
return
if self._start_time is None:
self._prefill_done_at = 0.0
return
t_rel = time.perf_counter() - self._start_time
self._take_sample(t_rel=t_rel)
self._prefill_done_at = t_rel
def result(self) -> PowerUsage:
self._stopped = True
assert self._start_time is not None, "result() called before run()"
elapsed = time.perf_counter() - self._start_time
self._take_sample(t_rel=elapsed)
# Clamp the split point to [0, elapsed] in case timing is weird (e.g.
# mark called after result, or sampler ran for < the prefill window).
split = self._prefill_done_at
if split is not None:
split = max(0.0, min(elapsed, split))
node_stats: list[NodePowerStats] = []
total_energy_j = 0.0
total_prefill_energy_j = 0.0
total_generation_energy_j = 0.0
for node_id, ts_profiles in self._samples.items():
n = len(ts_profiles)
if n == 0:
@@ -53,20 +75,68 @@ class PowerSampler:
node_energy_j = trapezoidal_energy(ts_profiles, elapsed)
avg_power_w = node_energy_j / elapsed if elapsed > 0 else 0.0
total_energy_j += node_energy_j
prefill_e: float | None = None
generation_e: float | None = None
prefill_avg: float | None = None
generation_avg: float | None = None
if split is not None:
prefill_e = trapezoidal_energy_range(ts_profiles, 0.0, split)
generation_e = trapezoidal_energy_range(ts_profiles, split, elapsed)
total_prefill_energy_j += prefill_e
total_generation_energy_j += generation_e
prefill_dt = split
generation_dt = elapsed - split
prefill_avg = prefill_e / prefill_dt if prefill_dt > 0 else 0.0
generation_avg = (
generation_e / generation_dt if generation_dt > 0 else 0.0
)
node_stats.append(
NodePowerStats(
node_id=node_id,
samples=n,
avg_sys_power=avg_power_w,
prefill_avg_sys_power=prefill_avg,
generation_avg_sys_power=generation_avg,
prefill_energy_joules=prefill_e,
generation_energy_joules=generation_e,
)
)
total_avg_sys_w = total_energy_j / elapsed if elapsed > 0 else 0.0
prefill_seconds: float | None = None
generation_seconds: float | None = None
prefill_energy_joules: float | None = None
generation_energy_joules: float | None = None
prefill_avg_w: float | None = None
generation_avg_w: float | None = None
if split is not None:
prefill_seconds = split
generation_seconds = elapsed - split
prefill_energy_joules = total_prefill_energy_j
generation_energy_joules = total_generation_energy_j
prefill_avg_w = (
total_prefill_energy_j / prefill_seconds if prefill_seconds > 0 else 0.0
)
generation_avg_w = (
total_generation_energy_j / generation_seconds
if generation_seconds > 0
else 0.0
)
return PowerUsage(
elapsed_seconds=elapsed,
nodes=node_stats,
total_avg_sys_power_watts=total_avg_sys_w,
total_energy_joules=total_energy_j,
prefill_seconds=prefill_seconds,
generation_seconds=generation_seconds,
prefill_energy_joules=prefill_energy_joules,
generation_energy_joules=generation_energy_joules,
prefill_avg_sys_power_watts=prefill_avg_w,
generation_avg_sys_power_watts=generation_avg_w,
)
@@ -89,3 +159,55 @@ def trapezoidal_energy(
continue
energy_j += (p_prev.sys_power + p_cur.sys_power) / 2.0 * dt
return energy_j
def trapezoidal_energy_range(
ts_profiles: list[tuple[float, SystemPerformanceProfile]],
t_start: float,
t_end: float,
) -> float:
"""Integrate sys_power(t) over [t_start, t_end] using the trapezoidal rule.
Linearly interpolates power at the endpoints when they fall between
existing samples, so callers can integrate over arbitrary sub-windows
(e.g. the prefill segment) without losing accuracy. Returns 0 for an
empty or zero-length window. Falls back to constant-power assumption
when only one sample exists.
"""
if t_end <= t_start:
return 0.0
if len(ts_profiles) == 0:
return 0.0
if len(ts_profiles) == 1:
return ts_profiles[0][1].sys_power * (t_end - t_start)
def power_at(t: float) -> float:
if t <= ts_profiles[0][0]:
return ts_profiles[0][1].sys_power
if t >= ts_profiles[-1][0]:
return ts_profiles[-1][1].sys_power
for i in range(1, len(ts_profiles)):
t_cur, p_cur = ts_profiles[i]
if t_cur >= t:
t_prev, p_prev = ts_profiles[i - 1]
span = t_cur - t_prev
if span <= 0:
return p_cur.sys_power
frac = (t - t_prev) / span
return p_prev.sys_power + frac * (p_cur.sys_power - p_prev.sys_power)
return ts_profiles[-1][1].sys_power
p_start = power_at(t_start)
p_end = power_at(t_end)
in_range: list[tuple[float, float]] = [
(t, profile.sys_power) for t, profile in ts_profiles if t_start < t < t_end
]
seq: list[tuple[float, float]] = [(t_start, p_start)] + in_range + [(t_end, p_end)]
energy_j = 0.0
for i in range(1, len(seq)):
dt = seq[i][0] - seq[i - 1][0]
if dt <= 0:
continue
energy_j += (seq[i - 1][1] + seq[i][1]) / 2.0 * dt
return energy_j
+121
View File
@@ -0,0 +1,121 @@
import multiprocessing as mp
import time
import pytest
from anyio import (
BrokenResourceError,
ClosedResourceError,
EndOfStream,
WouldBlock,
fail_after,
)
from loguru import logger
from exo.utils.channels import ErrorOverride, MpReceiver, MpSender, channel, mp_channel
class CustomClosedResourceError(ClosedResourceError):
pass
class CustomBrokenResourceError(BrokenResourceError):
pass
class CustomEndOfStream(EndOfStream):
pass
class CustomWouldBlock(WouldBlock):
pass
ERROR_OVERRIDE = ErrorOverride(
closed_resource_error=CustomClosedResourceError,
broken_resource_error=CustomBrokenResourceError,
end_of_stream=CustomEndOfStream,
would_block=CustomWouldBlock,
)
def foo(recv: MpReceiver[str]):
expected = ["hi", "hi 2", "bye"]
with recv as r:
for item in r:
assert item == expected.pop(0)
def bar(send: MpSender[str]):
logger.warning("hi")
send.send("hi")
time.sleep(0.1)
logger.warning("hi 2")
send.send("hi 2")
time.sleep(0.1)
logger.warning("bye")
send.send("bye")
time.sleep(0.1)
send.close()
@pytest.mark.anyio
async def test_channel_ipc():
with fail_after(0.5):
s, r = mp_channel[str]()
p1 = mp.Process(target=foo, args=(r,))
p2 = mp.Process(target=bar, args=(s,))
p1.start()
p2.start()
p1.join()
p2.join()
def test_channel_error_override_replaces_sync_errors_with_subclasses():
send, recv = channel[int](0, error_override_config=ERROR_OVERRIDE)
with pytest.raises(CustomWouldBlock) as would_block_info:
send.send_nowait(1)
assert type(would_block_info.value.__cause__) is WouldBlock
recv.close()
with pytest.raises(CustomBrokenResourceError) as broken_resource_info:
send.send_nowait(1)
assert type(broken_resource_info.value.__cause__) is BrokenResourceError
send.close()
with pytest.raises(CustomClosedResourceError) as closed_resource_info:
send.send_nowait(1)
assert type(closed_resource_info.value.__cause__) is ClosedResourceError
@pytest.mark.anyio
async def test_channel_error_override_replaces_async_errors_with_subclasses():
send, recv = channel[int](0, error_override_config=ERROR_OVERRIDE)
recv.close()
with pytest.raises(CustomBrokenResourceError) as broken_resource_info:
await send.send(1)
assert type(broken_resource_info.value.__cause__) is BrokenResourceError
send, recv = channel[int](error_override_config=ERROR_OVERRIDE)
send.close()
with pytest.raises(CustomEndOfStream) as end_of_stream_info:
await recv.receive()
assert type(end_of_stream_info.value.__cause__) is EndOfStream
@pytest.mark.anyio
async def test_channel_error_override_is_preserved_by_clones():
send, recv = channel[int](0, error_override_config=ERROR_OVERRIDE)
send_clone = send.clone()
recv.close()
with pytest.raises(CustomBrokenResourceError):
await send_clone.send(1)
send, recv = channel[int](0, error_override_config=ERROR_OVERRIDE)
cloned_send = recv.clone_sender()
recv.close()
with pytest.raises(CustomBrokenResourceError):
await cloned_send.send(1)
-168
View File
@@ -1,168 +0,0 @@
import contextlib
import os
from collections.abc import AsyncIterator
import anyio
import pytest
from anyio import EndOfStream, create_task_group, fail_after
from exo.utils.async_process import AsyncProcess
from exo.utils.channels import MpReceiver, MpSender, Receiver, mp_channel
from exo.utils.daemon import detach_stdio_to_devnull
def _write_before_and_after_detach() -> None:
os.write(1, b"before stdout\n")
os.write(2, b"before stderr\n")
detach_stdio_to_devnull()
os.write(1, b"after stdout\n")
os.write(2, b"after stderr\n")
def _write_grandchild_stdio(label: str) -> None:
os.write(1, f"{label} stdout\n".encode())
os.write(2, f"{label} stderr\n".encode())
async def _spawn_grandchild_and_report(
result_sender: MpSender[tuple[int, bytes, bytes]],
label: str,
) -> None:
result_sender.send(await _collect_spawned_child(label))
result_sender.close()
async def _collect_spawned_child(label: str) -> tuple[int, bytes, bytes]:
process = AsyncProcess(_write_grandchild_stdio, args=(label,))
async with _started_process(process):
return await _collect_process_output(process)
def _detach_stdio_then_spawn_captured_child(
result_sender: MpSender[tuple[int, bytes, bytes]],
) -> None:
detach_stdio_to_devnull()
anyio.run(_spawn_grandchild_and_report, result_sender, "grandchild")
def _detach_stdio_then_spawn_captured_children_sequentially(
result_sender: MpSender[list[tuple[int, bytes, bytes]]],
) -> None:
async def run_children() -> list[tuple[int, bytes, bytes]]:
results: list[tuple[int, bytes, bytes]] = []
for index in range(5):
results.append(await _collect_spawned_child(f"grandchild-{index}"))
return results
detach_stdio_to_devnull()
result_sender.send(anyio.run(run_children))
result_sender.close()
async def _collect_stream(stream: Receiver[bytes], output: bytearray) -> None:
while True:
try:
output.extend(await stream.receive())
except EndOfStream:
return
async def _collect_process_output(
process: AsyncProcess,
) -> tuple[int, bytes, bytes]:
stdout = bytearray()
stderr = bytearray()
exitcodes: list[int] = []
async with create_task_group() as collect_group:
collect_group.start_soon(_collect_stream, process.stdout, stdout)
collect_group.start_soon(_collect_stream, process.stderr, stderr)
exitcodes.append(await process.wait())
if not exitcodes:
raise RuntimeError("process exited without a return code")
return exitcodes[0], bytes(stdout), bytes(stderr)
@contextlib.asynccontextmanager
async def _started_process(process: AsyncProcess) -> AsyncIterator[None]:
async with create_task_group() as task_group:
await task_group.start(process.run)
try:
yield
finally:
await process.stop()
async def _run_process_and_receive[T](
process: AsyncProcess,
recv: MpReceiver[T],
*,
timeout: float,
) -> tuple[int, T]:
async with _started_process(process):
with fail_after(timeout):
result = await recv.receive_async()
exitcode = await process.wait()
return exitcode, result
@pytest.mark.anyio
async def test_detach_stdio_to_devnull_redirects_stdio_away_from_capture() -> None:
process = AsyncProcess(_write_before_and_after_detach)
async with _started_process(process):
exitcode, stdout, stderr = await _collect_process_output(process)
assert exitcode == 0
assert stdout == b"before stdout\n"
assert stderr == b"before stderr\n"
@pytest.mark.anyio
async def test_detached_stdio_process_can_spawn_and_capture_child_stdio() -> None:
send, recv = mp_channel[tuple[int, bytes, bytes]]()
process = AsyncProcess(_detach_stdio_then_spawn_captured_child, args=(send,))
try:
daemonized_parent_exitcode, result = await _run_process_and_receive(
process, recv, timeout=5
)
finally:
recv.close()
child_exitcode, child_stdout, child_stderr = result
assert daemonized_parent_exitcode == 0
assert child_exitcode == 0
assert child_stdout == b"grandchild stdout\n"
assert child_stderr == b"grandchild stderr\n"
@pytest.mark.anyio
async def test_detached_stdio_process_can_spawn_captured_children_sequentially() -> (
None
):
send, recv = mp_channel[list[tuple[int, bytes, bytes]]]()
process = AsyncProcess(
_detach_stdio_then_spawn_captured_children_sequentially,
args=(send,),
)
try:
daemonized_parent_exitcode, results = await _run_process_and_receive(
process, recv, timeout=10
)
finally:
recv.close()
assert daemonized_parent_exitcode == 0
assert results == [
(
0,
f"grandchild-{index} stdout\n".encode(),
f"grandchild-{index} stderr\n".encode(),
)
for index in range(5)
]
-40
View File
@@ -1,40 +0,0 @@
import multiprocessing as mp
import time
import pytest
from anyio import fail_after
from loguru import logger
from exo.utils.channels import MpReceiver, MpSender, mp_channel
def foo(recv: MpReceiver[str]):
expected = ["hi", "hi 2", "bye"]
with recv as r:
for item in r:
assert item == expected.pop(0)
def bar(send: MpSender[str]):
logger.warning("hi")
send.send("hi")
time.sleep(0.1)
logger.warning("hi 2")
send.send("hi 2")
time.sleep(0.1)
logger.warning("bye")
send.send("bye")
time.sleep(0.1)
send.close()
@pytest.mark.anyio
async def test_channel_ipc():
with fail_after(0.5):
s, r = mp_channel[str]()
p1 = mp.Process(target=foo, args=(r,))
p2 = mp.Process(target=bar, args=(s,))
p1.start()
p2.start()
p1.join()
p2.join()
+14 -24
View File
@@ -8,36 +8,28 @@ import textwrap
from pathlib import Path
from typing import Final
import pytest
import exo.utils.pidfile as pidfile
from exo.utils.pidfile import acquire_exo_pidfile
from exo_rs import Pidfile
_CHILD_ACQUIRE_PIDFILE_SCRIPT: Final = textwrap.dedent(
"""
import sys
from pathlib import Path
from unittest.mock import patch
import exo.utils.pidfile as pidfile
from exo.utils.pidfile import PidfileLockError, acquire_exo_pidfile
from exo_rs import Pidfile, PidfileError
with patch.object(pidfile, "EXO_PID_FILE", Path(sys.argv[1])):
try:
handle = acquire_exo_pidfile()
except PidfileLockError as exception:
print(str(exception))
raise SystemExit(73) from exception
path = Path(sys.argv[1])
try:
handle = Pidfile(path, 0o0600)
handle.write()
except (OSError, PidfileError) as exception:
print(f"Failed to acquire EXO pidfile at {path}: {exception}")
raise SystemExit(73) from exception
del handle
del handle
"""
)
def _use_pidfile_path(monkeypatch: pytest.MonkeyPatch, path: Path) -> None:
monkeypatch.setattr(pidfile, "EXO_PID_FILE", path)
def _run_child_acquire_pidfile(path: Path) -> subprocess.CompletedProcess[str]:
return subprocess.run(
[sys.executable, "-c", _CHILD_ACQUIRE_PIDFILE_SCRIPT, str(path)],
@@ -49,12 +41,11 @@ def _run_child_acquire_pidfile(path: Path) -> subprocess.CompletedProcess[str]:
def test_acquire_exo_pidfile_writes_current_pid_and_removes_on_drop(
tmp_path: Path,
monkeypatch: pytest.MonkeyPatch,
) -> None:
path = tmp_path / "exo.pid"
_use_pidfile_path(monkeypatch, path)
handle = acquire_exo_pidfile()
handle = Pidfile(path, 0o0600)
handle.write()
assert path.read_text() == str(os.getpid())
del handle
@@ -65,12 +56,11 @@ def test_acquire_exo_pidfile_writes_current_pid_and_removes_on_drop(
def test_acquire_exo_pidfile_rejects_second_process(
tmp_path: Path,
monkeypatch: pytest.MonkeyPatch,
) -> None:
path = tmp_path / "exo.pid"
_use_pidfile_path(monkeypatch, path)
handle = acquire_exo_pidfile()
handle = Pidfile(path, 0o0600)
handle.write()
try:
blocked_child = _run_child_acquire_pidfile(path)
assert blocked_child.returncode == 73
+136
View File
@@ -141,6 +141,142 @@ def test_trapezoidal_unit_single_sample() -> None:
assert trapezoidal_energy(samples, elapsed=3.0) == 42.0 * 3.0
def test_trapezoidal_range_interpolation() -> None:
"""Sub-window integration should linearly interpolate at the boundaries."""
from exo.utils.power_sampler import trapezoidal_energy_range
# Two samples: t=0 W=10, t=10 W=20 -> power(t) = 10 + t
samples = [
(0.0, _make_profile(10.0)),
(10.0, _make_profile(20.0)),
]
# Integral from t=4 to t=6: power goes 14 -> 16, mean 15, dt=2 -> 30 J
assert abs(trapezoidal_energy_range(samples, 4.0, 6.0) - 30.0) < 1e-9
# Integral over the full window matches the full trapezoidal integral.
full = trapezoidal_energy_range(samples, 0.0, 10.0)
assert abs(full - 150.0) < 1e-9
def test_trapezoidal_range_zero_window() -> None:
"""Zero-length or reversed windows integrate to zero."""
from exo.utils.power_sampler import trapezoidal_energy_range
samples = [(0.0, _make_profile(10.0)), (5.0, _make_profile(20.0))]
assert trapezoidal_energy_range(samples, 3.0, 3.0) == 0.0
assert trapezoidal_energy_range(samples, 5.0, 3.0) == 0.0
def test_trapezoidal_range_splits_sum_to_full() -> None:
"""Energy split at an arbitrary boundary should sum back to the full integral."""
from exo.utils.power_sampler import (
trapezoidal_energy,
trapezoidal_energy_range,
)
samples = [
(0.0, _make_profile(10.0)),
(1.0, _make_profile(20.0)),
(3.0, _make_profile(15.0)),
(5.0, _make_profile(25.0)),
]
full = trapezoidal_energy(samples, elapsed=5.0)
# Split at t=2.5 (between samples) — interpolation should be exact.
left = trapezoidal_energy_range(samples, 0.0, 2.5)
right = trapezoidal_energy_range(samples, 2.5, 5.0)
assert abs((left + right) - full) < 1e-9
async def test_prefill_generation_split() -> None:
"""When mark_prefill_done() is called, the result should split energy."""
state: dict[NodeId, SystemPerformanceProfile] = {
NODE_A: _make_profile(10.0),
}
sampler = PowerSampler(get_node_system=lambda: state, interval=0.02)
async with anyio.create_task_group() as tg:
tg.start_soon(sampler.run)
# "Prefill" phase: power = 10 W
await anyio.sleep(0.1)
# Mark the boundary BEFORE changing state — this matches what
# _collect_text_generation_with_stats does in production: the mark
# fires on the first non-prefill chunk, so the boundary sample is
# the genuine end-of-prefill reading rather than the new phase's.
sampler.mark_prefill_done()
state[NODE_A] = _make_profile(30.0)
# "Generation" phase: power = 30 W
await anyio.sleep(0.1)
tg.cancel_scope.cancel()
result = sampler.result()
assert result.prefill_seconds is not None
assert result.generation_seconds is not None
assert result.prefill_energy_joules is not None
assert result.generation_energy_joules is not None
assert result.prefill_avg_sys_power_watts is not None
assert result.generation_avg_sys_power_watts is not None
# Phase durations should sum to the elapsed seconds.
assert (
abs(
(result.prefill_seconds + result.generation_seconds)
- result.elapsed_seconds
)
< 1e-6
)
# Phase energies should sum to (approximately) the total.
assert (
abs(
(result.prefill_energy_joules + result.generation_energy_joules)
- result.total_energy_joules
)
< 1e-6
)
# With the boundary sample anchored at the genuine end-of-prefill (10 W),
# prefill avg should converge tightly on 10 W and generation on 30 W.
# 15 W cleanly separates the two and would catch any cross-contamination.
assert result.prefill_avg_sys_power_watts < 15.0
assert result.generation_avg_sys_power_watts > 15.0
assert result.nodes[0].prefill_avg_sys_power is not None
assert result.nodes[0].generation_avg_sys_power is not None
async def test_no_split_when_unmarked() -> None:
"""If mark_prefill_done() is never called, phase fields stay None."""
state: dict[NodeId, SystemPerformanceProfile] = {
NODE_A: _make_profile(10.0),
}
sampler = PowerSampler(get_node_system=lambda: state, interval=0.02)
async with anyio.create_task_group() as tg:
tg.start_soon(sampler.run)
await anyio.sleep(0.05)
tg.cancel_scope.cancel()
result = sampler.result()
assert result.prefill_seconds is None
assert result.generation_seconds is None
assert result.prefill_energy_joules is None
assert result.generation_energy_joules is None
assert result.nodes[0].prefill_energy_joules is None
assert result.nodes[0].generation_energy_joules is None
async def test_mark_prefill_done_is_idempotent() -> None:
"""Only the first call to mark_prefill_done() should take effect."""
state: dict[NodeId, SystemPerformanceProfile] = {
NODE_A: _make_profile(10.0),
}
sampler = PowerSampler(get_node_system=lambda: state, interval=0.02)
async with anyio.create_task_group() as tg:
tg.start_soon(sampler.run)
await anyio.sleep(0.05)
sampler.mark_prefill_done()
first_prefill_at = sampler._prefill_done_at # pyright: ignore[reportPrivateUsage]
await anyio.sleep(0.05)
sampler.mark_prefill_done()
assert sampler._prefill_done_at == first_prefill_at # pyright: ignore[reportPrivateUsage]
tg.cancel_scope.cancel()
async def test_result_stops_sampling() -> None:
"""Calling result() should stop the sampler's run loop."""
state: dict[NodeId, SystemPerformanceProfile] = {
+45 -2
View File
@@ -229,6 +229,47 @@ def has_non_kv_caches(cache: KVCacheType) -> bool:
return any(is_non_trimmable_cache_entry(c) for c in cache)
# Max snapshots retained per cache entry. Each CacheSnapshot pins detached GPU
# copies of every non-trimmable (SSM/ArraysCache, RotatingKVCache) layer, so
# retaining one per ~4096-token prefill chunk makes snapshot memory grow linearly
# with context — the dominant residual cost when a single entry is grown to long
# contexts on hybrid models (~56 MB/snapshot on Qwen3.5-122B, so a full 256K
# context = 64 snapshots ≈ 3.6 GB). A sliding window of the most-recent N caps
# this at N×per-snapshot (~0.9 GB here) while preserving the restore points
# in-place grows actually use (they always extend from the tip).
_MAX_RETAINED_SNAPSHOTS = 16
def _bounded_snapshots(snapshots: list[CacheSnapshot]) -> list[CacheSnapshot]:
"""Deduplicate snapshots by token position and bound the retained count.
Returned list is sorted ascending by ``token_count``.
"""
# Deduplicate by position, keeping the most-recently-appended snapshot per
# position. Repeated in-place grows re-snapshot positions the kept old
# snapshots already cover, which would otherwise grow `_snapshots`
# unbounded even at constant context.
# TODO: keying on token_count alone is safe only while a position uniquely
# identifies the prefix within an entry (grows are strict prefix-extensions).
# If edit-and-regenerate, sliding-window/prefix trimming, cross-entry
# snapshot sharing, per-request adapter/LoRA swap, or branchy decoding
# (beam/parallel/speculative) is added, enrich the key to
# (token_count, prefix_hash[, media/adapter id]) — else a stale snapshot
# could be restored for a different prefix (silent wrong output).
by_position: dict[int, CacheSnapshot] = {}
for snapshot in snapshots:
by_position[snapshot.token_count] = snapshot
deduped = [by_position[pos] for pos in sorted(by_position)]
# Sliding window: keep only the most-recent N positions. In-place grows
# always extend from the tip, so the newest snapshots are the ones future
# grows restore from — dropping the oldest is never incorrect: a later hit on
# a prefix older than the window finds no snapshot <= target, so get_kv_cache
# returns a fresh cache (matched_index=None) and the request takes a full cold
# prefill — correct, just slower than a partial-hit reuse for that one request.
return deduped[-_MAX_RETAINED_SNAPSHOTS:]
class KVPrefixCache:
def __init__(self, group: mx.distributed.Group | None):
self.prompts: list[mx.array] = [] # mx array of tokens (ints)
@@ -261,7 +302,9 @@ class KVPrefixCache:
self._evict_if_needed()
self.prompts.append(prompt_tokens)
self.caches.append(deepcopy(cache))
self._snapshots.append(ssm_snapshots)
self._snapshots.append(
_bounded_snapshots(ssm_snapshots) if ssm_snapshots else None
)
self._media_regions.append(media_regions or [])
self.prefill_tps.append(prefill_tps)
self._access_counter += 1
@@ -288,7 +331,7 @@ class KVPrefixCache:
self.prompts[index] = prompt_tokens
self.caches[index] = deepcopy(cache)
self._snapshots[index] = merged or None
self._snapshots[index] = _bounded_snapshots(merged) or None
self._media_regions[index] = media_regions or []
self.prefill_tps[index] = prefill_tps
self._access_counter += 1
+7 -1
View File
@@ -8,6 +8,10 @@ from loguru import logger
from exo.api.types import ImageEditsTaskParams
from exo.download.download_utils import is_read_only_model_dir, resolve_existing_model
from exo.routing.event_router import (
EventRouterBrokenResourceError,
EventRouterClosedResourceError,
)
from exo.shared.apply import apply
from exo.shared.constants import EXO_MAX_INSTANCE_RETRIES
from exo.shared.models.model_cards import ModelId, card_cache
@@ -109,7 +113,9 @@ class Worker:
tg.start_soon(self._event_applier)
tg.start_soon(self._poll_connection_updates)
tg.start_soon(self._reconcile_custom_cards)
except* (EventRouterBrokenResourceError, EventRouterClosedResourceError):
# Event router has been closed (try-star syntax handles error groups)
pass
finally:
# Actual shutdown code - waits for all tasks to complete before executing.
logger.info("Stopping Worker")
@@ -11,6 +11,7 @@ from mlx_lm.sample_utils import make_sampler
from exo.shared.types.common import ModelId
from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams
from exo.worker.engines.mlx.cache import (
CacheSnapshot,
KVPrefixCache,
cache_length,
encode_prompt,
@@ -77,6 +78,74 @@ class TestGetPrefixLength:
assert get_prefix_length(a, b) == 0
class TestSnapshotAccumulation:
"""Locks in the fix for the actual per-grow Metal leak on hybrid (SSM)
models: `update_kv_cache` must not let `_snapshots` grow without bound when
the same entry is grown in place many times."""
def test_repeated_update_does_not_accumulate_snapshots(self):
with patch(
"exo.worker.engines.mlx.cache.get_memory_used_percentage",
return_value=0.0,
):
kv_prefix_cache = KVPrefixCache(None)
initial = [
CacheSnapshot(states=[None], token_count=4096),
CacheSnapshot(states=[None], token_count=8192),
]
kv_prefix_cache.add_kv_cache(
mx.arange(10000), [KVCache()], ssm_snapshots=initial
)
# Each in-place grow re-prefills from restore_pos and produces a
# fresh snapshot at a position the retained old snapshots already
# cover. Pre-fix this appended one snapshot per grow forever.
for _ in range(50):
fresh = [CacheSnapshot(states=[None], token_count=8192)]
kv_prefix_cache.update_kv_cache(
0, mx.arange(10000), [KVCache()], fresh, restore_pos=8192
)
stored = kv_prefix_cache._snapshots[0]
assert stored is not None
# Bounded by the number of distinct snapshot positions (here 2),
# not by the 50 grows.
assert len(stored) == 2
assert sorted(s.token_count for s in stored) == [4096, 8192]
# The kept 8192 snapshot must be the most recently supplied one.
assert stored[1] is fresh[0]
def test_extension_caps_snapshots_to_sliding_window(self):
"""Extending a single entry to a long context (one snapshot per ~4096
tokens) must cap retained snapshots to a sliding window of the most-recent
N, not keep all of them that linear-in-context retention was the
residual OOM cause."""
from exo.worker.engines.mlx.cache import _MAX_RETAINED_SNAPSHOTS
with patch(
"exo.worker.engines.mlx.cache.get_memory_used_percentage",
return_value=0.0,
):
kv_prefix_cache = KVPrefixCache(None)
# 64 distinct positions = a 262144-token context at 4096/chunk.
num_positions = 64
snaps = [
CacheSnapshot(states=[None], token_count=4096 * (i + 1))
for i in range(num_positions)
]
kv_prefix_cache.add_kv_cache(
mx.arange(10), [KVCache()], ssm_snapshots=snaps
)
stored = kv_prefix_cache._snapshots[0]
assert stored is not None
# Capped at the window; the most-recent N positions are retained
# (in-place grows extend from the tip, so these are what get used).
assert len(stored) == _MAX_RETAINED_SNAPSHOTS
assert stored == snaps[-_MAX_RETAINED_SNAPSHOTS:]
assert stored[-1] is snaps[-1] # tip always kept
class TestKVPrefix:
@pytest.fixture
def mock_tokenizer(self):
+27 -3
View File
@@ -7,6 +7,9 @@ set -uo pipefail
HOST="${1:-localhost:52415}"
MODEL_ID="KevTheHermit/security-testing"
ENCODED_MODEL_ID=$(
python3 -c 'import sys, urllib.parse; print(urllib.parse.quote(sys.argv[1], safe=""))' "$MODEL_ID"
)
CUSTOM_CARDS_DIR="$HOME/.exo/custom_model_cards"
CARD_FILE="$CUSTOM_CARDS_DIR/KevTheHermit--security-testing.toml"
@@ -71,9 +74,30 @@ PLACE_BODY=$(echo "$PLACE_RESPONSE" | sed '$d')
echo " HTTP $PLACE_CODE"
echo " Response: $PLACE_BODY"
# Step 3b: Send a chat completion to actually trigger tokenizer loading
if [ "$PLACE_CODE" -ge 400 ]; then
echo " Placement failed; cannot trigger tokenizer loading."
exit 1
fi
# Step 3b: Wait for placement to materialize before inference.
echo ""
echo "[3b] Sending chat completion to trigger tokenizer load ..."
echo "[3b] Waiting for placed instance ..."
if ! AWAIT_RESPONSE=$(curl -fsS --max-time 65 \
"http://$HOST/instance/await?model_id=$ENCODED_MODEL_ID&timeout_seconds=60" |
awk '/^data: / { sub(/^data: /, ""); print; exit }'); then
echo " Timed out waiting for an instance for $MODEL_ID"
exit 1
fi
if ! printf '%s' "$AWAIT_RESPONSE" | grep -q '"type":"ready"'; then
echo " Timed out waiting for an instance for $MODEL_ID"
exit 1
fi
echo " Instance ready"
# Step 3c: Send a chat completion to actually trigger tokenizer loading
echo ""
echo "[3c] Sending chat completion to trigger tokenizer load ..."
CHAT_RESPONSE=$(curl -s -w "\n%{http_code}" --max-time 30 -X POST "http://$HOST/v1/chat/completions" \
-H "Content-Type: application/json" \
-d "{\"model\":\"$MODEL_ID\",\"messages\":[{\"role\":\"user\",\"content\":\"hello\"}],\"max_tokens\":1}")
@@ -82,7 +106,7 @@ CHAT_BODY=$(echo "$CHAT_RESPONSE" | sed '$d')
echo " HTTP $CHAT_CODE"
echo " Response: $CHAT_BODY"
echo ""
echo "[3c] Checking for RCE proof ..."
echo "[3d] Checking for RCE proof ..."
sleep 5
if [ -f /tmp/exo-rce-proof.txt ]; then
echo " VULNERABLE: Remote code executed!"
Generated
+804 -854
View File
File diff suppressed because it is too large. Load diff