crates/jera/src/agent.rs
//! **Multi-machine agents + worker-groups** — the distributed-runner seam.
//!
//! jera's job ledger ([`crate::jobs`]) tracks *that* work ran; this module is *where*
//! it runs: an [`Agent`] is a worker (a machine / process) that belongs to a named
//! **worker-group** and advertises its capabilities. A [`Scheduler`] matches a ready
//! [`WorkUnit`] to a **capable idle agent in the target group** and dispatches it,
//! returning a [`WorkHandle`] with **bidirectional kill** (cancel the handle → the
//! agent kills the work; an agent going offline → its in-flight work → `Killed`).
//!
//! The whole point is the seam: [`LocalAgent`] (in-process, single-machine) ships
//! REAL and tested now; [`GrpcAgent`] is the remote skeleton. Going multi-machine is
//! **adding the `GrpcAgent` transport, not rewriting the scheduler** — nornir, dwarves,
//! and holger all consume this one `jera::agent` layer, per the reuse-first law
//! (never a per-app runner twin). This is the Buildkite split: one control-plane
//! ([`Scheduler`]) + N `<host> agent --group …` workers over gRPC.
//!
//! Pure std (threads + atomics); no new deps, additive to jera's public surface.
use std::collections::{BTreeMap, BTreeSet};
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex};
use serde::{Deserialize, Serialize};
use crate::policy::JobPolicy;
/// Durable node **roster** + heartbeat/liveness — the fleet the [`Scheduler`] draws
/// its remote workers from (Loki/Odin/NUC + Njord). See [`roster::Roster`].
pub mod roster;
pub use roster::{NodeEntry, Roster};
/// Stable id of a worker.
pub type AgentId = String;
/// Name of a worker-group (the pool a work-unit targets, e.g. `"loki"`, `"gpu"`).
pub type WorkerGroup = String;
/// What a worker advertises: identity, its group, capability labels, and size.
/// Capabilities are opaque labels (`"linux"`, `"cargo"`, `"gpu"`, `"kvm"`); a work
/// unit's `required_caps` must be a subset for the agent to be eligible.
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct AgentSpec {
pub id: AgentId,
pub group: WorkerGroup,
pub caps: BTreeSet<String>,
pub cores: u32,
pub mem_mb: u64,
}
impl AgentSpec {
/// Convenience: build a spec from id/group + capability strings.
pub fn new(
id: impl Into<String>,
group: impl Into<String>,
caps: impl IntoIterator<Item = impl Into<String>>,
cores: u32,
mem_mb: u64,
) -> Self {
Self {
id: id.into(),
group: group.into(),
caps: caps.into_iter().map(Into::into).collect(),
cores,
mem_mb,
}
}
}
/// A worker's liveness/availability as the scheduler sees it.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum AgentState {
/// Ready to accept a work-unit.
Idle,
/// Currently running a work-unit.
Busy,
/// Dead / unregistered — never scheduled; its in-flight work was killed.
Offline,
}
// ── WorkPayload — *what* a unit runs (the keystone) ──────────────────────────────
/// A built-binary run: program + args (+ optional cwd/env/input-refs). The
/// [`WorkPayload::Process`] leg — the worker spawns this as a local child through
/// jera's own Exe runner ([`crate::boot`]).
#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct ProcessSpec {
pub program: String,
#[serde(default)]
pub args: Vec<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub cwd: Option<String>,
#[serde(default)]
pub env: BTreeMap<String, String>,
/// Opaque input refs (artifact ids / repo urls) the run needs; the worker
/// resolves them (e.g. via `nornir-flight` `do_get`) before it runs.
#[serde(default)]
pub inputs: Vec<String>,
}
impl ProcessSpec {
pub fn new(program: impl Into<String>, args: impl IntoIterator<Item = impl Into<String>>) -> Self {
Self {
program: program.into(),
args: args.into_iter().map(Into::into).collect(),
..Default::default()
}
}
}
/// An OCI container to bring up on the worker — the [`WorkPayload::Container`] leg,
/// mapped onto jera's no-shell `run_container` (draupnir OCI) at the far end.
#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct ContainerRun {
pub image: String,
#[serde(default)]
pub cmd: Vec<String>,
#[serde(default)]
pub env: BTreeMap<String, String>,
/// A SEALED run (a Runa `Sandbox` step): only these mounts, no network unless
/// asked, a read-only root, every capability dropped, the image never pulled.
/// `None` = an ordinary container run, serialised exactly as before.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub sandbox: Option<ContainerSandbox>,
}
/// What a sealed container run may see (see [`ContainerRun::sandbox`]).
#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct ContainerSandbox {
/// Bind mounts, podman `-v` form (`/host:/path[:ro]`).
#[serde(default)]
pub binds: Vec<String>,
/// tmpfs mounts, podman `--tmpfs` form (`/path:rw,size=1024m,mode=1777`).
#[serde(default)]
pub tmpfs: Vec<String>,
/// Network on; `false` = `--network=none`.
#[serde(default)]
pub network: bool,
/// The working directory inside the container (empty = the image's).
#[serde(default)]
pub workdir: String,
}
/// A KVM guest to boot on the worker — the [`WorkPayload::Vm`] leg, mapped onto
/// jera's `vm::BootSpec` (draupnir → tunnr) at the far end.
#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct VmBoot {
pub image: String,
#[serde(default)]
pub cores: u32,
#[serde(default)]
pub mem_mb: u64,
}
/// The serializable **job spec** a remote worker executes — the keystone that turns
/// a [`WorkUnit`] from "*where* to run" into "*what* to run". Mirrors the three
/// [`JobPolicy`] kinds but is self-contained + serde-ready (no draupnir type crosses
/// the wire): the worker maps a payload back onto jera's own runners
/// ([`crate::boot`] / [`crate::container`] / [`crate::vm`]) at the far end, so
/// dispatching "build repo X" remotely no longer needs a fixed pre-bound closure.
///
/// A [`WorkUnit`] without a payload keeps the legacy closure-based behaviour (the
/// [`LocalRunner`] decides from id/group) — this is purely additive.
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case", tag = "kind")]
pub enum WorkPayload {
/// Run a built binary on the worker (maps to [`JobPolicy::Process`]).
Process(ProcessSpec),
/// Bring up an OCI container on the worker (maps to [`JobPolicy::Container`]).
Container(ContainerRun),
/// Boot a KVM guest on the worker (maps to [`JobPolicy::Vm`]).
Vm(VmBoot),
}
impl WorkPayload {
/// Convenience: a Process payload from a program + args.
pub fn process(program: impl Into<String>, args: impl IntoIterator<Item = impl Into<String>>) -> Self {
WorkPayload::Process(ProcessSpec::new(program, args))
}
/// The [`JobPolicy`] this payload runs as — the single link back to jera's
/// policy classification (so the worker picks the right runner).
pub fn policy(&self) -> JobPolicy {
match self {
WorkPayload::Process(_) => JobPolicy::Process,
WorkPayload::Container(_) => JobPolicy::Container,
WorkPayload::Vm(_) => JobPolicy::Vm,
}
}
/// Stable label of the payload kind (`"process"`/`"container"`/`"vm"`).
pub fn kind_str(&self) -> &'static str {
self.policy().as_str()
}
}
// ── ResourceReq — *how big* a unit is (Stage 5: resource-aware placement) ─────────
/// The resource floor a [`WorkUnit`] demands of the node that runs it — **min
/// cores, min memory, and required resource tags**. This is the control-plane's
/// placement input: [`Scheduler::match_agent`] only routes a unit to an
/// [`AgentSpec`] that clears the floor, so a **heavy bench never lands on the
/// 6-core NUC** — it goes to Loki/Odin.
///
/// Distinct from [`WorkUnit::required_caps`] (opaque capability *labels*): this is
/// the *sizing* dimension (numbers + resource tags). Both must be satisfied.
/// `#[serde(default)]` on every field so an older-encoded unit decodes as "no
/// floor" (additive). Placement is a control-plane concern, so `ResourceReq` never
/// crosses the wire — the worker just runs what it is handed.
#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct ResourceReq {
/// Minimum core count the node must have (a heavy job floors this above the NUC).
#[serde(default)]
pub min_cores: u32,
/// Minimum memory (MiB) the node must have.
#[serde(default)]
pub min_mem_mb: u64,
/// Resource tags the node's [`AgentSpec::caps`] must be a superset of (e.g.
/// `"heavy"`, `"bench"`) — a sizing-oriented sibling of `required_caps`.
#[serde(default)]
pub tags: BTreeSet<String>,
}
impl ResourceReq {
/// A floor of `min_cores`/`min_mem_mb` (no tags).
pub fn new(min_cores: u32, min_mem_mb: u64) -> Self {
Self { min_cores, min_mem_mb, tags: BTreeSet::new() }
}
/// The **HEAVY bench** floor — the sizing that keeps a heavy run off a
/// NUC-class box (6 cores). Requires **≥ 8 cores** and ≥ 16 GiB, so it places
/// only on Loki (12) / Odin (32), never the NUC — the "serial HEAVY on a quiet,
/// capable box" law expressed as placement data. Memory alone can't
/// discriminate (the NUC also has 64 GiB), so the core floor is the guard.
pub fn heavy() -> Self {
Self { min_cores: 8, min_mem_mb: 16384, tags: BTreeSet::new() }
}
/// Attach required resource tags (builder style).
pub fn with_tags(mut self, tags: impl IntoIterator<Item = impl Into<String>>) -> Self {
self.tags = tags.into_iter().map(Into::into).collect();
self
}
/// Whether `spec` clears this floor: enough cores, enough memory, and its caps
/// ⊇ the required resource tags.
pub fn satisfied_by(&self, spec: &AgentSpec) -> bool {
spec.cores >= self.min_cores
&& spec.mem_mb >= self.min_mem_mb
&& self.tags.is_subset(&spec.caps)
}
}
/// `true` for a `false` bool — the `skip_serializing_if` predicate that keeps the
/// wire/JSON lean for the common non-exclusive unit (a bare `false` is elided).
fn is_false(b: &bool) -> bool {
!*b
}
/// One schedulable unit of work: which group it must run in, the caps it needs, and
/// — the keystone — an optional [`WorkPayload`] describing *what* to run.
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct WorkUnit {
pub id: String,
pub group: WorkerGroup,
#[serde(default)]
pub required_caps: BTreeSet<String>,
/// The serializable spec the worker executes (the keystone). `#[serde(default)]`
/// so a unit encoded before this field decodes as `None` — no migration, and a
/// payload-less unit keeps the legacy closure-based path (additive).
#[serde(default, skip_serializing_if = "Option::is_none")]
pub payload: Option<WorkPayload>,
/// The **resource floor** for placement (Stage 5). `None` = no floor (runs on
/// any capable idle agent in the group). A heavy unit sets [`ResourceReq::heavy`]
/// so it can only place on Loki/Odin, never the NUC. Control-plane placement
/// metadata — not sent to the worker. Additive (`#[serde(default)]`).
#[serde(default, skip_serializing_if = "Option::is_none")]
pub resources: Option<ResourceReq>,
/// **Fleet-exclusive** within its group: at most one exclusive unit per group
/// runs at a time across the whole fleet — a scheduler-level generalization of
/// the process-global [`crate::jobs`] exclusive gate, so a HEAVY bench serializes
/// *across machines* ([`Scheduler::dispatch_scheduled`]). Additive; a bare
/// `false` is elided from the wire/JSON.
#[serde(default, skip_serializing_if = "is_false")]
pub exclusive: bool,
}
impl WorkUnit {
pub fn new(
id: impl Into<String>,
group: impl Into<String>,
required_caps: impl IntoIterator<Item = impl Into<String>>,
) -> Self {
Self {
id: id.into(),
group: group.into(),
required_caps: required_caps.into_iter().map(Into::into).collect(),
payload: None,
resources: None,
exclusive: false,
}
}
/// Attach a [`WorkPayload`] (builder style) — the unit now carries *what* to
/// run, not just where. Chains onto [`new`](Self::new).
pub fn with_payload(mut self, payload: WorkPayload) -> Self {
self.payload = Some(payload);
self
}
/// Attach a [`ResourceReq`] placement floor (builder style) — the unit now says
/// *how big* a node it needs (Stage 5). Chains onto [`new`](Self::new).
pub fn with_resources(mut self, resources: ResourceReq) -> Self {
self.resources = Some(resources);
self
}
/// Mark this unit **fleet-exclusive** within its group (builder style) — it
/// serializes against every other exclusive unit in the same group via
/// [`Scheduler::dispatch_scheduled`]. Chains onto [`new`](Self::new).
pub fn exclusive(mut self) -> Self {
self.exclusive = true;
self
}
/// The payload this unit carries, if any.
pub fn payload(&self) -> Option<&WorkPayload> {
self.payload.as_ref()
}
/// The resource floor this unit demands, if any.
pub fn resources(&self) -> Option<&ResourceReq> {
self.resources.as_ref()
}
}
/// Lifecycle of a dispatched work-unit. `Killed` is distinct from `Failed`: it means
/// the kill was requested (cancel or agent-offline), never a silent stall (RAGNARÖK).
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum WorkStatus {
Pending,
Running,
Done,
Failed,
Killed,
}
impl WorkStatus {
pub fn is_terminal(self) -> bool {
matches!(self, WorkStatus::Done | WorkStatus::Failed | WorkStatus::Killed)
}
pub fn as_str(self) -> &'static str {
match self {
WorkStatus::Pending => "pending",
WorkStatus::Running => "running",
WorkStatus::Done => "done",
WorkStatus::Failed => "failed",
WorkStatus::Killed => "killed",
}
}
}
// ── WorkResult — the terminal outcome the worker returns (Stage 4) ────────────────
/// The **result a worker returns** for a dispatched [`WorkUnit`] — its terminal
/// [`WorkStatus`], **which node ran it**, an optional process exit code, a
/// `result_ref`, and free-form detail. This is the return leg of the multi-machine
/// flow: a control-plane dispatches a unit → the worker runs it → it sends a
/// `WorkResult` back to the central ledger, which records it against the right
/// [`crate::jobs::JobRecord::node`] via [`crate::jobs::JobStore::submit`] (the thin
/// `Jobs.Submit` path — reused, not re-invented).
///
/// **Artifacts stay out of this record.** Bulk output (a build's `target/`, a bench
/// CSV) is not inlined here — it travels over the **nornir-flight** transport, and
/// `result_ref` carries its *flight ticket* (e.g. `"flight://loki/bench_runs/abc"`)
/// or a domain row ref (`"bench_runs:abc"`). The control-plane `do_get`s the bytes
/// by that ticket; jera stays a lean leaf and never links Arrow. Fully serde-ready
/// so it round-trips over the wire intact.
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct WorkResult {
/// The ledger key this result belongs to — the `job_id` the control-plane
/// assigned when it recorded the dispatch, so the return overwrites the right row.
pub job_id: String,
/// The [`WorkUnit::id`] this result is for (human-readable trace).
pub unit_id: String,
/// Terminal outcome (`Done`/`Failed`/`Killed`).
pub status: WorkStatus,
/// Which fleet node ran it (roster node id / host) — stamped onto the recorded
/// [`crate::jobs::JobRecord::node`]. `None` = unknown / local.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub node: Option<String>,
/// The process exit code, when the payload was a [`ProcessSpec`] (`None` for a
/// container/VM or when not applicable).
#[serde(default, skip_serializing_if = "Option::is_none")]
pub exit_code: Option<i32>,
/// Pointer to the artifact/domain result — a **nornir-flight ticket** for bulk
/// output, or a domain row ref (`"bench_runs:<uuid>"`), or `""`.
#[serde(default)]
pub result_ref: String,
/// Free-form JSON detail (log tail / counts / error chain), or `""`.
#[serde(default)]
pub detail_json: String,
}
impl WorkResult {
/// A result for `job_id`/`unit_id` with a terminal `status` and the `node` that
/// ran it. Chain [`with_result_ref`](Self::with_result_ref) /
/// [`with_exit_code`](Self::with_exit_code) / [`with_detail`](Self::with_detail).
pub fn new(
job_id: impl Into<String>,
unit_id: impl Into<String>,
status: WorkStatus,
node: Option<String>,
) -> Self {
Self {
job_id: job_id.into(),
unit_id: unit_id.into(),
status,
node,
exit_code: None,
result_ref: String::new(),
detail_json: String::new(),
}
}
/// Observe a dispatched handle's *current* status into a result for `job_id`,
/// stamping the `node` it ran on. The worker calls this once the handle is
/// terminal to form the return record.
pub fn observed(job_id: impl Into<String>, handle: &dyn WorkHandle, node: Option<String>) -> Self {
Self::new(job_id, handle.id().to_string(), handle.poll(), node)
}
/// Attach the artifact/domain `result_ref` (a nornir-flight ticket or row ref).
pub fn with_result_ref(mut self, result_ref: impl Into<String>) -> Self {
self.result_ref = result_ref.into();
self
}
/// Attach the process exit code.
pub fn with_exit_code(mut self, code: i32) -> Self {
self.exit_code = Some(code);
self
}
/// Attach free-form JSON detail.
pub fn with_detail(mut self, detail: impl Into<String>) -> Self {
self.detail_json = detail.into();
self
}
/// The [`crate::jobs`] lifecycle status this result maps onto (`Killed` folds to
/// `failed` — the ledger has no distinct killed state).
fn jobs_status(&self) -> &'static str {
match self.status {
WorkStatus::Done => crate::jobs::status::DONE,
WorkStatus::Failed | WorkStatus::Killed => crate::jobs::status::FAILED,
WorkStatus::Pending => crate::jobs::status::QUEUED,
WorkStatus::Running => crate::jobs::status::RUNNING,
}
}
/// Build the [`crate::jobs::JobRecord`] the central ledger records for this
/// result — status mapped, **`node` stamped**, `result_ref`/`detail` carried,
/// timing filled from `ts_start_micros` (the dispatch time). Hand the returned
/// record to [`crate::jobs::JobStore::submit`] (the thin-client write path) to
/// record it against the node that ran the job. Keeps the jobs module free of any
/// reverse dependency on `agent` — the mapping lives here.
pub fn to_job_record(
&self,
kind: &str,
target: &str,
workspace: &str,
ts_start_micros: i64,
) -> crate::jobs::JobRecord {
let now = chrono::Utc::now().timestamp_micros();
let terminal = self.status.is_terminal();
crate::jobs::JobRecord {
job_id: self.job_id.clone(),
kind: kind.to_string(),
target: target.to_string(),
workspace: workspace.to_string(),
status: self.jobs_status().to_string(),
ts_start_micros,
ts_end_micros: terminal.then_some(now),
elapsed_ms: terminal.then_some((now - ts_start_micros) / 1000),
detail_json: self.detail_json.clone(),
result_ref: self.result_ref.clone(),
parent_id: None,
node: self.node.clone(),
phase: None,
actor: None,
actor_method: None,
}
}
}
/// A handle to dispatched work. `kill()` is the client→remote half of the
/// bidirectional lifecycle; `poll()` reads the current status.
pub trait WorkHandle: Send + Sync {
fn id(&self) -> &str;
fn poll(&self) -> WorkStatus;
/// Request cancellation. The running work observes the flag and terminates as
/// `Killed`. Idempotent.
fn kill(&self);
}
/// A worker that can run [`WorkUnit`]s. Implemented by [`LocalAgent`] (now) and
/// [`GrpcAgent`] (the remote seam).
pub trait Agent: Send + Sync {
fn spec(&self) -> &AgentSpec;
fn state(&self) -> AgentState;
/// Eligible iff idle, same group, its caps ⊇ the work-unit's required caps, and
/// (Stage 5) it clears the unit's [`ResourceReq`] floor — so a heavy unit never
/// matches a too-small node (the NUC).
fn can_run(&self, w: &WorkUnit) -> bool {
self.state() == AgentState::Idle
&& self.spec().group == w.group
&& w.required_caps.is_subset(&self.spec().caps)
&& w.resources().is_none_or(|r| r.satisfied_by(self.spec()))
}
/// Start `w`; returns its handle. (Caller should have checked [`Agent::can_run`].)
fn dispatch(&self, w: WorkUnit) -> Arc<dyn WorkHandle>;
/// The agent-death half of the bidirectional lifecycle: mark offline and kill
/// every in-flight unit (each → `Killed`), so no job is left orphaned.
fn go_offline(&self);
}
// ── LocalAgent — the in-process, single-machine worker (ships now) ───────────────
/// The runner an in-process agent executes per unit. It receives the unit and a kill
/// flag it must poll; returning `true` = success, `false` = failure. Long-running
/// runners MUST check `kill` so cancellation is honoured.
pub type LocalRunner = Arc<dyn Fn(&WorkUnit, &AtomicBool) -> bool + Send + Sync>;
struct LocalWorkHandle {
id: String,
status: Arc<Mutex<WorkStatus>>,
kill: Arc<AtomicBool>,
}
impl WorkHandle for LocalWorkHandle {
fn id(&self) -> &str {
&self.id
}
fn poll(&self) -> WorkStatus {
*self.status.lock().unwrap()
}
fn kill(&self) {
self.kill.store(true, Ordering::SeqCst);
}
}
/// A single-machine agent that runs each unit on its own std thread. This is the
/// concrete impl the whole distributed design is proven against before any wire is
/// involved — the scheduler, capability match, and bidirectional kill are all real
/// and tested here.
pub struct LocalAgent {
spec: AgentSpec,
runner: LocalRunner,
busy: Arc<AtomicBool>,
offline: Arc<AtomicBool>,
inflight: Mutex<Vec<Arc<LocalWorkHandle>>>,
}
impl LocalAgent {
pub fn new(spec: AgentSpec, runner: LocalRunner) -> Self {
Self {
spec,
runner,
busy: Arc::new(AtomicBool::new(false)),
offline: Arc::new(AtomicBool::new(false)),
inflight: Mutex::new(Vec::new()),
}
}
/// An agent whose work always succeeds immediately — for scheduling/topology
/// tests that don't care about the body.
pub fn immediate(spec: AgentSpec) -> Self {
Self::new(spec, Arc::new(|_w, _kill| true))
}
}
impl Agent for LocalAgent {
fn spec(&self) -> &AgentSpec {
&self.spec
}
fn state(&self) -> AgentState {
if self.offline.load(Ordering::SeqCst) {
AgentState::Offline
} else if self.busy.load(Ordering::SeqCst) {
AgentState::Busy
} else {
AgentState::Idle
}
}
fn dispatch(&self, w: WorkUnit) -> Arc<dyn WorkHandle> {
let handle = Arc::new(LocalWorkHandle {
id: w.id.clone(),
status: Arc::new(Mutex::new(WorkStatus::Pending)),
kill: Arc::new(AtomicBool::new(false)),
});
self.inflight.lock().unwrap().push(handle.clone());
self.busy.store(true, Ordering::SeqCst);
let runner = self.runner.clone();
let busy = self.busy.clone();
let status = handle.status.clone();
let kill = handle.kill.clone();
std::thread::spawn(move || {
*status.lock().unwrap() = WorkStatus::Running;
let ok = runner(&w, &kill);
let final_status = if kill.load(Ordering::SeqCst) {
WorkStatus::Killed
} else if ok {
WorkStatus::Done
} else {
WorkStatus::Failed
};
*status.lock().unwrap() = final_status;
busy.store(false, Ordering::SeqCst);
});
handle
}
fn go_offline(&self) {
self.offline.store(true, Ordering::SeqCst);
for h in self.inflight.lock().unwrap().iter() {
h.kill();
}
}
}
// ── GrpcAgent + the AgentTransport seam — the remote path (loopback REAL now) ─────
/// The wire seam behind a remote [`GrpcAgent`]. A transport carries
/// `state`/`dispatch`/`go_offline` across a process (or machine) boundary. Two impls:
/// - [`LoopbackTransport`] — in-process, **REAL + tested**: proves the seam
/// end-to-end (a unit dispatched *through* the transport runs to `Done`, and a
/// `kill` propagates to `Killed`) with no socket — so the whole distributed path
/// is exercised in code before any wire exists.
/// - a tonic `GrpcTransport` — the remaining wire, a drop-in `impl AgentTransport`
/// behind an `agent-grpc` feature (so jera's lean default never pulls tonic). That
/// is the whole point of the seam: going multi-machine ADDS a transport;
/// [`Scheduler`] and everything above stay unchanged. See
/// `.nornir/agents-multimachine.md`.
pub trait AgentTransport: Send + Sync {
/// Liveness of the remote worker (`Idle` when it can accept a unit).
fn remote_state(&self) -> AgentState;
/// Send a unit across the boundary; the handle tracks remote status and its
/// `kill()` is the client→remote cancel half of the bidirectional lifecycle.
fn remote_dispatch(&self, w: WorkUnit) -> Arc<dyn WorkHandle>;
/// Agent-death half: drop the remote registration, killing in-flight work.
fn remote_go_offline(&self);
}
/// A worker reached over an [`AgentTransport`]. `state`/`dispatch`/`go_offline`
/// delegate to the transport, so the [`Scheduler`] treats it exactly like a
/// [`LocalAgent`] — the seam is invisible above this line.
pub struct GrpcAgent {
spec: AgentSpec,
/// The remote worker's endpoint (e.g. `http://loki:7900`) — informational for
/// loopback, the connect target for the gRPC transport.
pub endpoint: String,
transport: Arc<dyn AgentTransport>,
}
impl GrpcAgent {
/// An **unwired** remote agent: reports [`AgentState::Offline`] so the scheduler
/// never routes to it (safe — no panic, no hang). The back-compat default until a
/// transport is attached.
pub fn new(spec: AgentSpec, endpoint: impl Into<String>) -> Self {
Self { spec, endpoint: endpoint.into(), transport: Arc::new(UnwiredTransport) }
}
/// A remote agent driven by a concrete [`AgentTransport`] (loopback in tests, the
/// gRPC transport in production).
pub fn with_transport(
spec: AgentSpec,
endpoint: impl Into<String>,
transport: Arc<dyn AgentTransport>,
) -> Self {
Self { spec, endpoint: endpoint.into(), transport }
}
/// A remote agent over an in-process [`LoopbackTransport`] running `runner` — the
/// real, testable proof of the transport seam (used before any wire exists).
pub fn loopback(spec: AgentSpec, runner: LocalRunner) -> Self {
let t = Arc::new(LoopbackTransport::new(spec.clone(), runner));
Self::with_transport(spec, "loopback://local", t)
}
}
struct DeadHandle {
id: String,
}
impl WorkHandle for DeadHandle {
fn id(&self) -> &str {
&self.id
}
fn poll(&self) -> WorkStatus {
WorkStatus::Failed
}
fn kill(&self) {}
}
/// The unwired transport: reports `Offline` and fails any dispatch — the safe
/// not-yet-connected state a bare [`GrpcAgent::new`] carries.
struct UnwiredTransport;
impl AgentTransport for UnwiredTransport {
fn remote_state(&self) -> AgentState {
AgentState::Offline
}
fn remote_dispatch(&self, w: WorkUnit) -> Arc<dyn WorkHandle> {
Arc::new(DeadHandle { id: w.id })
}
fn remote_go_offline(&self) {}
}
/// An in-process transport that runs the unit on a real [`LocalAgent`] *behind the
/// transport boundary* — proving `dispatch → Done` and `kill → Killed` cross the seam
/// intact, with no network. The gRPC transport is the same shape over a socket.
pub struct LoopbackTransport {
inner: LocalAgent,
}
impl LoopbackTransport {
pub fn new(spec: AgentSpec, runner: LocalRunner) -> Self {
Self { inner: LocalAgent::new(spec, runner) }
}
/// A loopback whose work always succeeds immediately.
pub fn immediate(spec: AgentSpec) -> Self {
Self { inner: LocalAgent::immediate(spec) }
}
}
impl AgentTransport for LoopbackTransport {
fn remote_state(&self) -> AgentState {
self.inner.state()
}
fn remote_dispatch(&self, w: WorkUnit) -> Arc<dyn WorkHandle> {
self.inner.dispatch(w)
}
fn remote_go_offline(&self) {
self.inner.go_offline()
}
}
impl Agent for GrpcAgent {
fn spec(&self) -> &AgentSpec {
&self.spec
}
fn state(&self) -> AgentState {
self.transport.remote_state()
}
fn dispatch(&self, w: WorkUnit) -> Arc<dyn WorkHandle> {
self.transport.remote_dispatch(w)
}
fn go_offline(&self) {
self.transport.remote_go_offline()
}
}
// ── Scheduler — the control-plane: match a ready unit to a capable idle agent ─────
/// A dispatched handle that also **holds a fleet-exclusive group permit** until it
/// is dropped (Stage 5). Delegates `id`/`poll`/`kill` to the real handle; releasing
/// the permit on drop admits the next queued exclusive unit in the same group, so
/// only one exclusive unit per group runs at a time across the fleet.
struct ScheduledWorkHandle {
inner: Arc<dyn WorkHandle>,
_permit: tokio::sync::OwnedSemaphorePermit,
}
impl WorkHandle for ScheduledWorkHandle {
fn id(&self) -> &str {
self.inner.id()
}
fn poll(&self) -> WorkStatus {
self.inner.poll()
}
fn kill(&self) {
self.inner.kill()
}
}
/// The control-plane. Holds the registered agents (of any [`Agent`] impl — local or
/// gRPC) and matches a ready [`WorkUnit`] to the first eligible idle agent in its
/// target worker-group, in deterministic registration order.
#[derive(Default)]
pub struct Scheduler {
agents: Vec<Arc<dyn Agent>>,
/// Per-group **fleet-exclusive gate** (Stage 5): one single-permit semaphore per
/// worker-group, created on first use. An exclusive unit acquires its group's
/// permit before it dispatches and holds it (via [`ScheduledWorkHandle`]) until
/// its handle drops — so a fleet-wide HEAVY bench serializes *across machines*,
/// the scheduler-level generalization of the process-global [`crate::jobs`] gate.
group_gates: Mutex<BTreeMap<WorkerGroup, Arc<tokio::sync::Semaphore>>>,
}
impl Scheduler {
pub fn new() -> Self {
Self { agents: Vec::new(), group_gates: Mutex::new(BTreeMap::new()) }
}
/// The single-permit fleet-exclusive gate for `group`, lazily created. Every
/// exclusive unit in the same group contends for this one permit.
fn group_gate(&self, group: &str) -> Arc<tokio::sync::Semaphore> {
self.group_gates
.lock()
.unwrap()
.entry(group.to_string())
.or_insert_with(|| Arc::new(tokio::sync::Semaphore::new(1)))
.clone()
}
/// Register a worker (a `<host> agent --group …` in the distributed picture).
pub fn register(&mut self, agent: Arc<dyn Agent>) {
self.agents.push(agent);
}
pub fn agents(&self) -> &[Arc<dyn Agent>] {
&self.agents
}
/// The first capable idle agent in `w`'s group, or `None` (queue it, with a
/// reason — never a silent drop).
pub fn match_agent(&self, w: &WorkUnit) -> Option<Arc<dyn Agent>> {
self.agents.iter().find(|a| a.can_run(w)).cloned()
}
/// Match then dispatch. Returns the handle, or `None` if no agent is eligible
/// right now (caller re-queues).
///
/// When the unit carries a [`WorkPayload`], a functional-status row is emitted
/// through jera's [`crate::jobs::selftest`] bridge (a no-op unless a consumer —
/// nornir/dwarves — installed a `nornir-testmatrix` hook), so the payload
/// dispatch is observable as a matrix cell at the consumer.
pub fn dispatch(&self, w: WorkUnit) -> Option<Arc<dyn WorkHandle>> {
let agent = self.match_agent(&w)?;
if let Some(p) = w.payload() {
crate::jobs::selftest::emit(
"jera-agent",
&format!("dispatch_payload_{}", p.kind_str()),
true,
&format!("unit '{}' ({}) → agent '{}'", w.id, p.kind_str(), agent.spec().id),
);
}
Some(agent.dispatch(w))
}
/// **Fleet-exclusive dispatch** (Stage 5). Like [`dispatch`](Self::dispatch),
/// but a unit marked [`WorkUnit::exclusive`] first **acquires its group's single
/// permit** — so at most one exclusive unit per group runs at a time *across the
/// whole fleet*, even when several idle agents are free. The returned handle
/// holds the permit until it drops, admitting the next queued exclusive unit.
///
/// A non-exclusive unit is never gated (delegates straight to [`dispatch`]). This
/// is the scheduler-level generalization of the process-global
/// [`crate::jobs`] exclusive gate: a HEAVY bench dispatched to the `bench` group
/// serializes across Loki/Odin instead of only within one control-plane process
/// (the "serial HEAVY on a quiet box" law, made fleet-wide).
///
/// Returns `None` (releasing the permit) if no eligible agent exists right now —
/// the caller re-queues, exactly like [`dispatch`].
pub async fn dispatch_scheduled(&self, w: WorkUnit) -> Option<Arc<dyn WorkHandle>> {
if !w.exclusive {
return self.dispatch(w);
}
let gate = self.group_gate(&w.group);
let permit = gate
.acquire_owned()
.await
.expect("a group's exclusive gate is never closed");
let inner = self.dispatch(w)?; // permit drops here if no eligible agent
crate::jobs::selftest::emit(
"jera-agent",
"exclusive_gate_acquired",
true,
&format!("exclusive unit '{}' holds the group permit", inner.id()),
);
Some(Arc::new(ScheduledWorkHandle { inner, _permit: permit }))
}
/// Observable fleet state — feeds a UI "Agents" pane / a `nornir-testmatrix`
/// cell at the consumer (nornir/dwarves), so the robot/matrix can assert the
/// pool as data (the remote-control-in-code law).
pub fn state_json(&self) -> serde_json::Value {
serde_json::json!({
"agents": self.agents.iter().map(|a| {
let s = a.spec();
serde_json::json!({
"id": s.id,
"group": s.group,
"caps": s.caps,
"cores": s.cores,
"mem_mb": s.mem_mb,
"state": a.state(),
})
}).collect::<Vec<_>>(),
})
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::time::{Duration, Instant};
fn poll_until_terminal(h: &Arc<dyn WorkHandle>) -> WorkStatus {
let deadline = Instant::now() + Duration::from_secs(3);
loop {
let s = h.poll();
if s.is_terminal() || Instant::now() > deadline {
return s;
}
std::thread::sleep(Duration::from_millis(2));
}
}
#[test]
fn scheduler_matches_group_and_caps_then_runs_to_done() {
let mut sch = Scheduler::new();
// Two groups; only the loki/linux+cargo agent is eligible.
sch.register(Arc::new(LocalAgent::immediate(AgentSpec::new(
"odin-1", "odin", ["linux"], 32, 65536,
))));
sch.register(Arc::new(LocalAgent::immediate(AgentSpec::new(
"loki-1", "loki", ["linux", "cargo"], 12, 62000,
))));
let w = WorkUnit::new("build-1", "loki", ["cargo"]);
let agent = sch.match_agent(&w).expect("a capable idle loki agent exists");
assert_eq!(agent.spec().id, "loki-1", "matched the right group+caps");
let h = sch.dispatch(w).expect("dispatched");
assert_eq!(poll_until_terminal(&h), WorkStatus::Done);
}
#[test]
fn no_eligible_agent_returns_none_never_a_silent_drop() {
let mut sch = Scheduler::new();
sch.register(Arc::new(LocalAgent::immediate(AgentSpec::new(
"loki-1", "loki", ["linux"], 12, 62000,
))));
// Wrong group.
assert!(sch.match_agent(&WorkUnit::new("x", "odin", Vec::<String>::new())).is_none());
// Missing capability.
assert!(sch.match_agent(&WorkUnit::new("y", "loki", ["gpu"])).is_none());
// Dispatch to no match → None, so the caller re-queues (no panic, no hang).
assert!(sch.dispatch(WorkUnit::new("z", "odin", Vec::<String>::new())).is_none());
}
#[test]
fn kill_handle_terminates_work_as_killed() {
// A runner that blocks until the kill flag is set.
let runner: LocalRunner = Arc::new(|_w, kill| {
while !kill.load(Ordering::SeqCst) {
std::thread::sleep(Duration::from_millis(1));
}
false // returned value is ignored once kill was requested → Killed
});
let agent = LocalAgent::new(AgentSpec::new("loki-1", "loki", ["linux"], 12, 62000), runner);
let h = agent.dispatch(WorkUnit::new("long", "loki", Vec::<String>::new()));
// Give it a moment to enter Running, then cancel.
std::thread::sleep(Duration::from_millis(10));
h.kill();
assert_eq!(poll_until_terminal(&h), WorkStatus::Killed);
}
#[test]
fn agent_going_offline_kills_inflight_work() {
let runner: LocalRunner = Arc::new(|_w, kill| {
while !kill.load(Ordering::SeqCst) {
std::thread::sleep(Duration::from_millis(1));
}
false
});
let agent = LocalAgent::new(AgentSpec::new("loki-1", "loki", ["linux"], 12, 62000), runner);
let h = agent.dispatch(WorkUnit::new("long", "loki", Vec::<String>::new()));
std::thread::sleep(Duration::from_millis(10));
agent.go_offline(); // agent death → in-flight work killed
assert_eq!(agent.state(), AgentState::Offline);
assert_eq!(poll_until_terminal(&h), WorkStatus::Killed);
}
#[test]
fn busy_agent_is_not_rescheduled_until_idle() {
let runner: LocalRunner = Arc::new(|_w, kill| {
while !kill.load(Ordering::SeqCst) {
std::thread::sleep(Duration::from_millis(1));
}
true
});
let agent: Arc<dyn Agent> =
Arc::new(LocalAgent::new(AgentSpec::new("loki-1", "loki", ["linux"], 12, 62000), runner));
let mut sch = Scheduler::new();
sch.register(agent.clone());
let w1 = WorkUnit::new("w1", "loki", Vec::<String>::new());
let h1 = sch.dispatch(w1).expect("first dispatch");
std::thread::sleep(Duration::from_millis(10));
assert_eq!(agent.state(), AgentState::Busy);
// Now busy → not eligible.
assert!(sch.match_agent(&WorkUnit::new("w2", "loki", Vec::<String>::new())).is_none());
h1.kill();
assert_eq!(poll_until_terminal(&h1), WorkStatus::Killed);
}
#[test]
fn grpc_agent_is_offline_skeleton_never_scheduled() {
let mut sch = Scheduler::new();
sch.register(Arc::new(GrpcAgent::new(
AgentSpec::new("remote-1", "loki", ["linux", "cargo"], 12, 62000),
"http://loki:7900",
)));
// Skeleton reports Offline → never matched (no panic, safe seam).
assert!(sch.match_agent(&WorkUnit::new("w", "loki", ["cargo"])).is_none());
assert_eq!(sch.agents()[0].state(), AgentState::Offline);
}
#[test]
fn grpc_agent_over_loopback_transport_runs_to_done() {
// A remote agent whose transport is the in-process loopback: the scheduler
// routes to it exactly like a local agent, and the unit crosses the transport
// boundary and runs to Done — the whole distributed path, proven in code.
let mut sch = Scheduler::new();
sch.register(Arc::new(GrpcAgent::loopback(
AgentSpec::new("remote-1", "loki", ["linux", "cargo"], 12, 62000),
Arc::new(|_w, _kill| true),
)));
let agent = sch
.match_agent(&WorkUnit::new("w", "loki", ["cargo"]))
.expect("the loopback remote is Idle → schedulable");
assert_eq!(agent.spec().id, "remote-1");
let h = sch.dispatch(WorkUnit::new("w", "loki", ["cargo"])).expect("dispatched over the transport");
assert_eq!(poll_until_terminal(&h), WorkStatus::Done);
}
#[test]
fn grpc_agent_over_loopback_transport_kill_propagates() {
// kill() on the client handle must cross the transport seam → Killed.
let runner: LocalRunner = Arc::new(|_w, kill| {
while !kill.load(Ordering::SeqCst) {
std::thread::sleep(Duration::from_millis(1));
}
false
});
let agent = GrpcAgent::loopback(
AgentSpec::new("remote-1", "loki", ["cargo"], 8, 32000),
runner,
);
let h = agent.dispatch(WorkUnit::new("w", "loki", ["cargo"]));
h.kill();
assert_eq!(poll_until_terminal(&h), WorkStatus::Killed);
}
#[test]
fn work_payload_round_trips_through_json() {
// RED-when-broken: the keystone must survive serde intact (this is what
// crosses the wire as `payload_json`). Cover all three legs.
for payload in [
WorkPayload::process("cargo", ["build", "--release"]),
WorkPayload::Container(ContainerRun {
image: "docker.io/library/alpine:3".into(),
cmd: vec!["echo".into(), "hi".into()],
env: BTreeMap::from([("K".to_string(), "V".to_string())]),
sandbox: None,
}),
WorkPayload::Vm(VmBoot { image: "loki.qcow2".into(), cores: 4, mem_mb: 8192 }),
] {
let w = WorkUnit::new("job-1", "loki", ["cargo"]).with_payload(payload.clone());
let json = serde_json::to_string(&w).expect("encode");
let back: WorkUnit = serde_json::from_str(&json).expect("decode");
assert_eq!(back, w, "WorkUnit+payload round-trips byte-for-byte");
assert_eq!(back.payload(), Some(&payload));
assert_eq!(back.payload().unwrap().policy(), payload.policy());
}
}
#[test]
fn payload_less_unit_still_decodes_and_legacy_path_unchanged() {
// Additive guarantee: a unit encoded WITHOUT a payload field decodes as
// None (no migration), and the classic constructor yields None.
let legacy_json = r#"{"id":"x","group":"loki","required_caps":["cargo"]}"#;
let w: WorkUnit = serde_json::from_str(legacy_json).expect("legacy decodes");
assert_eq!(w.payload(), None, "missing payload field ⇒ None");
assert_eq!(WorkUnit::new("y", "loki", Vec::<String>::new()).payload(), None);
}
#[test]
fn payload_carrying_unit_dispatches_and_reaches_the_runner() {
// The payload must cross the dispatch boundary and be visible to the runner
// — proving a unit now says *what* to run, not just where. The runner
// asserts it received the exact ProcessSpec.program.
let seen: Arc<Mutex<Option<String>>> = Arc::new(Mutex::new(None));
let seen_w = seen.clone();
let runner: LocalRunner = Arc::new(move |w, _kill| {
if let Some(WorkPayload::Process(p)) = w.payload() {
*seen_w.lock().unwrap() = Some(p.program.clone());
p.program == "cargo"
} else {
false
}
});
let mut sch = Scheduler::new();
sch.register(Arc::new(LocalAgent::new(
AgentSpec::new("loki-1", "loki", ["cargo"], 12, 62000),
runner,
)));
let w = WorkUnit::new("build-1", "loki", ["cargo"])
.with_payload(WorkPayload::process("cargo", ["build"]));
let h = sch.dispatch(w).expect("dispatched a payload-carrying unit");
assert_eq!(poll_until_terminal(&h), WorkStatus::Done);
assert_eq!(
seen.lock().unwrap().as_deref(),
Some("cargo"),
"the runner saw the payload's ProcessSpec across the dispatch boundary"
);
}
#[test]
fn resource_req_round_trips_and_heavy_floors_the_cores() {
// RED-when-broken: the placement floor must survive serde intact and the
// heavy() floor must sit above a 6-core NUC but below Loki (12) / Odin (32).
let heavy = ResourceReq::heavy();
let json = serde_json::to_string(&heavy).unwrap();
let back: ResourceReq = serde_json::from_str(&json).unwrap();
assert_eq!(back, heavy, "ResourceReq round-trips");
assert!(heavy.min_cores >= 8, "heavy floors cores above the NUC");
let nuc = AgentSpec::new("nuc-1", "bench", ["linux"], 6, 65536);
let loki = AgentSpec::new("loki-1", "bench", ["linux"], 12, 65536);
let odin = AgentSpec::new("odin-1", "bench", ["linux"], 32, 131072);
assert!(!heavy.satisfied_by(&nuc), "6-core NUC does NOT clear the heavy floor");
assert!(heavy.satisfied_by(&loki), "Loki clears it");
assert!(heavy.satisfied_by(&odin), "Odin clears it");
}
#[test]
fn heavy_resource_req_never_places_on_a_nuc_class_node() {
// The headline Stage-5 guarantee: a heavy job routes to Loki/Odin, NEVER the
// 6-core NUC — even though the NUC is idle and has the caps + memory.
let mut sch = Scheduler::new();
sch.register(Arc::new(LocalAgent::immediate(AgentSpec::new(
"nuc-1", "bench", ["linux", "cargo"], 6, 65536, // 6 cores, 64 GiB
))));
let heavy = WorkUnit::new("heavy-bench", "bench", ["cargo"])
.with_resources(ResourceReq::heavy());
// RED-when-broken: no match — the NUC is too small for a heavy job.
assert!(
sch.match_agent(&heavy).is_none(),
"a 6-core NUC must never take a heavy job"
);
assert!(sch.dispatch(heavy.clone()).is_none(), "and dispatch refuses it too");
// A NON-heavy unit still places on the NUC (additive: no floor ⇒ any capable node).
let light = WorkUnit::new("light", "bench", ["cargo"]);
assert_eq!(
sch.match_agent(&light).map(|a| a.spec().id.clone()),
Some("nuc-1".to_string()),
"an un-floored unit still uses the NUC"
);
// Add Loki (12 cores): now the heavy job places, and only on Loki.
sch.register(Arc::new(LocalAgent::immediate(AgentSpec::new(
"loki-1", "bench", ["linux", "cargo"], 12, 65536,
))));
assert_eq!(
sch.match_agent(&heavy).map(|a| a.spec().id.clone()),
Some("loki-1".to_string()),
"heavy routes to the capable Loki node"
);
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn exclusive_group_units_do_not_run_concurrently_across_the_fleet() {
// RED-when-broken: two exclusive units in the same group must serialize even
// when TWO idle agents are free — without the group gate both would dispatch
// at once. This is the fleet-wide HEAVY-bench serialization.
let mut sch = Scheduler::new();
sch.register(Arc::new(LocalAgent::immediate(AgentSpec::new(
"loki-1", "bench", ["cargo"], 12, 65536,
))));
sch.register(Arc::new(LocalAgent::immediate(AgentSpec::new(
"odin-1", "bench", ["cargo"], 32, 131072,
))));
let sch = Arc::new(sch);
// First exclusive unit acquires the group permit and holds it (handle alive).
let h1 = sch
.dispatch_scheduled(WorkUnit::new("bench-a", "bench", ["cargo"]).exclusive())
.await
.expect("first exclusive dispatched");
// Second exclusive unit in the same group MUST block on the group gate.
let sch2 = sch.clone();
let mut second = tokio::spawn(async move {
sch2.dispatch_scheduled(WorkUnit::new("bench-b", "bench", ["cargo"]).exclusive())
.await
.map(|h| h.id().to_string())
});
assert!(
tokio::time::timeout(Duration::from_millis(200), &mut second).await.is_err(),
"second exclusive unit is blocked while the first holds the group permit"
);
// Release the first's permit → the second is admitted and dispatches.
drop(h1);
let got = tokio::time::timeout(Duration::from_secs(2), second)
.await
.expect("second no longer blocked once the permit is free")
.expect("join")
.expect("second exclusive dispatched after the gate freed");
assert_eq!(got, "bench-b", "the queued exclusive unit is the one that ran");
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn non_exclusive_unit_is_never_gated() {
// A parallel (non-exclusive) unit dispatches immediately even while an
// exclusive unit holds its group's permit — the gate is exclusive-only.
let mut sch = Scheduler::new();
// Two idle agents so the parallel unit always has a free node even while the
// exclusive unit's agent is momentarily Busy — isolating "is it gated?" from
// "is an agent free?".
sch.register(Arc::new(LocalAgent::immediate(AgentSpec::new(
"loki-1", "bench", ["cargo"], 12, 65536,
))));
sch.register(Arc::new(LocalAgent::immediate(AgentSpec::new(
"odin-1", "bench", ["cargo"], 32, 131072,
))));
let sch = Arc::new(sch);
let _held = sch
.dispatch_scheduled(WorkUnit::new("excl", "bench", ["cargo"]).exclusive())
.await
.expect("exclusive holds the permit");
// Non-exclusive still dispatches right away (agent is idle again after the
// immediate exclusive finished; the point is it is not gate-blocked).
let par = tokio::time::timeout(
Duration::from_millis(500),
sch.dispatch_scheduled(WorkUnit::new("par", "bench", ["cargo"])),
)
.await
.expect("non-exclusive dispatch is never blocked by the gate");
assert!(par.is_some(), "parallel unit dispatched immediately");
}
#[test]
fn dispatched_units_result_round_trips_back_and_records_against_its_node() {
// Stage 4 RED-when-broken: a dispatched WorkPayload's result must survive the
// return wire (serde) AND land in the central ledger stamped with the node
// that ran it — the whole point of result return.
use crate::jobs::{kind, JobSelector, JobStore};
let mut sch = Scheduler::new();
sch.register(Arc::new(LocalAgent::immediate(AgentSpec::new(
"odin-1", "bench", ["cargo"], 32, 131072,
))));
let start = chrono::Utc::now().timestamp_micros();
let unit = WorkUnit::new("bench-znippy", "bench", ["cargo"])
.with_payload(WorkPayload::process("nornir", ["bench", "run"]));
let h = sch.dispatch(unit).expect("dispatched to odin");
assert_eq!(poll_until_terminal(&h), WorkStatus::Done);
// The worker forms the result, stamping the node it ran on + the artifact ref.
let result = WorkResult::observed("job-42", h.as_ref(), Some("odin".to_string()))
.with_result_ref("bench_runs:abc")
.with_exit_code(0)
.with_detail(r#"{"rows":3}"#);
assert_eq!(result.status, WorkStatus::Done);
// It round-trips over the return wire intact.
let wire = serde_json::to_string(&result).expect("encode result");
let back: WorkResult = serde_json::from_str(&wire).expect("decode result");
assert_eq!(back, result, "WorkResult round-trips byte-for-byte");
// The central ledger records it — against the RIGHT node.
let dir = std::env::temp_dir().join(format!("jera-result-{}", std::process::id()));
std::fs::remove_dir_all(&dir).ok();
let store = JobStore::open(&dir).unwrap();
store
.submit(&back.to_job_record(kind::BENCH_RUN, "znippy-zoomies", "nordisk", start))
.expect("submit to central ledger");
let rec = store
.list(&JobSelector::All)
.unwrap()
.into_iter()
.find(|r| r.job_id == "job-42")
.expect("the result was recorded");
assert_eq!(rec.node.as_deref(), Some("odin"), "recorded against the node that ran it");
assert_eq!(rec.status, "done", "terminal status carried");
assert_eq!(rec.result_ref, "bench_runs:abc", "artifact ref carried");
assert_eq!(rec.kind, kind::BENCH_RUN);
assert!(rec.elapsed_ms.is_some(), "timing filled from dispatch time");
std::fs::remove_dir_all(&dir).ok();
}
#[test]
fn state_json_exposes_the_fleet_as_data() {
let mut sch = Scheduler::new();
sch.register(Arc::new(LocalAgent::immediate(AgentSpec::new(
"loki-1", "loki", ["linux", "cargo"], 12, 62000,
))));
let v = sch.state_json();
let a = &v["agents"][0];
assert_eq!(a["id"], "loki-1");
assert_eq!(a["group"], "loki");
assert_eq!(a["state"], "idle");
assert_eq!(a["cores"], 12);
}
}