crates/jera/src/agent_grpc.rs
//! `agent-grpc` — the tonic wire for the multi-machine [`Scheduler`](crate::agent::Scheduler).
//!
//! This is the drop-in `impl AgentTransport` the [`LoopbackTransport`] proved out:
//! a [`GrpcTransport`] (client) carries `state`/`dispatch`/`kill`/`go_offline` to a
//! remote [`serve`]d worker running over a [`LocalAgent`]. The scheduler, capability
//! match, and bidirectional kill are unchanged — going multi-machine just swaps the
//! transport. Feature-gated so jera's lean default never links tonic.
//!
//! The `AgentTransport` methods are synchronous (the scheduler is sync by design), so
//! the client bridges into tonic via a private multi-thread runtime + `block_on`
//! (safe from the sync control-plane; do not call from inside another runtime).
use std::collections::HashMap;
use std::sync::{Arc, Mutex};
use crate::agent::{Agent, AgentState, LocalAgent, WorkHandle, WorkStatus, WorkUnit};
/// Generated prost+tonic bindings (`package jera.agent.v1`).
pub mod pb {
tonic::include_proto!("jera.agent.v1");
}
fn state_to_str(s: AgentState) -> &'static str {
match s {
AgentState::Idle => "idle",
AgentState::Busy => "busy",
AgentState::Offline => "offline",
}
}
fn state_from_str(s: &str) -> AgentState {
match s {
"idle" => AgentState::Idle,
"busy" => AgentState::Busy,
_ => AgentState::Offline, // unknown ⇒ offline (never schedule to a mystery worker)
}
}
fn status_from_str(s: &str) -> WorkStatus {
match s {
"running" => WorkStatus::Running,
"done" => WorkStatus::Done,
"failed" => WorkStatus::Failed,
"killed" => WorkStatus::Killed,
_ => WorkStatus::Pending,
}
}
// ── server: an AgentWorker over a LocalAgent ─────────────────────────────────────
/// How many terminal verdicts a worker remembers after it has released the handle
/// that produced them. The memo exists so `Poll` stays truthful for an id whose
/// handle is gone; it is FIFO-bounded so the lifecycle fix is not a second unbounded
/// map. One entry is an id `String` plus a one-byte [`WorkStatus`] — orders of
/// magnitude smaller than the handle it replaces.
pub const TERMINAL_MEMO_CAP: usize = 1024;
/// The bounded id → terminal-[`WorkStatus`] memo. FIFO: the oldest verdict is
/// evicted once [`TERMINAL_MEMO_CAP`] is reached.
#[derive(Default)]
struct TerminalMemo {
by_id: HashMap<String, WorkStatus>,
order: std::collections::VecDeque<String>,
}
impl TerminalMemo {
fn insert(&mut self, id: &str, status: WorkStatus) {
if self.by_id.insert(id.to_string(), status).is_none() {
self.order.push_back(id.to_string());
while self.order.len() > TERMINAL_MEMO_CAP {
if let Some(old) = self.order.pop_front() {
self.by_id.remove(&old);
}
}
}
}
}
/// The worker-box service: runs dispatched units on a wrapped [`LocalAgent`] and
/// tracks their handles by id so `Poll`/`Kill` reach the right in-flight work.
///
/// **A lifecycle, not a cache.** `handles` used to be a registry with no deregister:
/// `dispatch` inserted, `poll`/`kill` only ever `get`, so a long-lived worker box
/// accumulated one `Arc<dyn WorkHandle>` per unit it had EVER been asked to run,
/// each pinning whatever that handle owns. A handle is now released the moment its
/// unit reaches a terminal state — on the `poll` that observes it, and by a sweep on
/// `dispatch` as the backstop for units nobody polls — and its terminal verdict
/// moves into the bounded [`TerminalMemo`] so `Poll` keeps answering truthfully
/// afterwards. That last part is load-bearing: the receipt path
/// (`WorkResult::observed`) polls a SECOND time after the caller has already seen
/// terminal, so a worker that simply forgot a released id would report a successful
/// job as `Failed`.
pub struct AgentWorkerSvc {
agent: Arc<LocalAgent>,
/// LIVE handles only — see the type docs.
handles: Mutex<HashMap<String, Arc<dyn WorkHandle>>>,
/// Terminal verdicts of released units, bounded at [`TERMINAL_MEMO_CAP`].
terminal: Mutex<TerminalMemo>,
}
impl AgentWorkerSvc {
pub fn new(agent: Arc<LocalAgent>) -> Self {
Self {
agent,
handles: Mutex::new(HashMap::new()),
terminal: Mutex::new(TerminalMemo::default()),
}
}
/// How many **live** work handles this worker is holding. Diagnostic, and the
/// oracle of the lifecycle tests: a worker that has served thousands of units
/// must not be holding thousands of handles.
pub fn live_handles(&self) -> usize {
self.handles.lock().unwrap().len()
}
/// How many terminal verdicts are remembered (≤ [`TERMINAL_MEMO_CAP`]).
pub fn remembered_terminals(&self) -> usize {
self.terminal.lock().unwrap().by_id.len()
}
fn remembered_status(&self, id: &str) -> Option<WorkStatus> {
self.terminal.lock().unwrap().by_id.get(id).copied()
}
fn remember_terminal(&self, id: &str, status: WorkStatus) {
self.terminal.lock().unwrap().insert(id, status);
}
/// Release the handle for `id` and record `status` as its verdict. The ONE place
/// a handle leaves the live map (LAW #5) — `poll` and the `dispatch` sweep both
/// route through it, so the drop and the memo write can never disagree.
fn release(&self, id: &str, status: WorkStatus) {
self.handles.lock().unwrap().remove(id);
self.remember_terminal(id, status);
}
/// Backstop: nothing obliges a client to poll a unit it dispatched, so
/// release-on-poll alone still lets the map grow. Sweep the live handles, release
/// every one that has reached a terminal state, and return how many were reaped.
/// The handle `Arc`s are cloned out before polling so no handle's `poll` runs
/// while the registry lock is held.
fn sweep_finished(&self) -> usize {
let live: Vec<(String, Arc<dyn WorkHandle>)> = {
let g = self.handles.lock().unwrap();
g.iter().map(|(k, v)| (k.clone(), v.clone())).collect()
};
let mut reaped = 0;
for (id, h) in live {
let s = h.poll();
if s.is_terminal() {
self.release(&id, s);
reaped += 1;
}
}
reaped
}
}
#[tonic::async_trait]
impl pb::agent_worker_server::AgentWorker for AgentWorkerSvc {
async fn state(
&self,
_req: tonic::Request<pb::Empty>,
) -> Result<tonic::Response<pb::StateReply>, tonic::Status> {
Ok(tonic::Response::new(pb::StateReply { state: state_to_str(self.agent.state()).into() }))
}
async fn dispatch(
&self,
req: tonic::Request<pb::DispatchReq>,
) -> Result<tonic::Response<pb::WorkId>, tonic::Status> {
let u = req.into_inner().unit.ok_or_else(|| tonic::Status::invalid_argument("no unit"))?;
let mut unit = WorkUnit::new(u.id.clone(), u.group, u.required_caps);
// Reconstruct the payload (the keystone) — "" ⇒ a legacy payload-less unit.
if !u.payload_json.is_empty() {
let p = serde_json::from_str::<crate::agent::WorkPayload>(&u.payload_json)
.map_err(|e| tonic::Status::invalid_argument(format!("bad payload_json: {e}")))?;
unit = unit.with_payload(p);
}
// Backstop reap BEFORE taking on new work: a client is not obliged to poll a
// unit it dispatched, so without this the live map grows by one per dispatch
// for ever. Bounded by the handles currently held, which is exactly what this
// call keeps small.
self.sweep_finished();
let handle = self.agent.dispatch(unit);
let id = handle.id().to_string();
self.handles.lock().unwrap().insert(id.clone(), handle);
Ok(tonic::Response::new(pb::WorkId { id }))
}
async fn poll(
&self,
req: tonic::Request<pb::WorkId>,
) -> Result<tonic::Response<pb::StatusReply>, tonic::Status> {
let id = req.into_inner().id;
// Clone the Arc out and drop the registry lock before polling — a handle's
// `poll` must never run under that lock.
let live = self.handles.lock().unwrap().get(&id).cloned();
let status = match live {
Some(h) => {
let s = h.poll();
// Terminal ⇒ the handle has nothing left to do: release it and keep
// only the verdict, so the receipt poll that follows still reads it.
if s.is_terminal() {
self.release(&id, s);
}
s
}
// No live handle: the remembered verdict of a released unit, else an id
// this worker never ran (or one whose verdict has aged out of the memo)
// ⇒ Failed, as before.
None => self.remembered_status(&id).unwrap_or(WorkStatus::Failed),
};
Ok(tonic::Response::new(pb::StatusReply { status: status.as_str().into() }))
}
async fn kill(
&self,
req: tonic::Request<pb::WorkId>,
) -> Result<tonic::Response<pb::Empty>, tonic::Status> {
let id = req.into_inner().id;
let live = self.handles.lock().unwrap().get(&id).cloned();
if let Some(h) = live {
h.kill();
// `kill` only *requests* cancellation — the runner resolves the status
// when it notices. If it has already settled, release now; otherwise the
// next `poll` (or the next dispatch sweep) does it.
let s = h.poll();
if s.is_terminal() {
self.release(&id, s);
}
}
Ok(tonic::Response::new(pb::Empty {}))
}
async fn go_offline(
&self,
_req: tonic::Request<pb::Empty>,
) -> Result<tonic::Response<pb::Empty>, tonic::Status> {
self.agent.go_offline();
Ok(tonic::Response::new(pb::Empty {}))
}
}
/// Serve an [`AgentWorkerSvc`] on `addr` until the process ends, **open** (no auth).
/// Equivalent to [`serve_authed`] with an empty token. The worker box runs this
/// (`<host> agent --group …`); the control-plane connects a [`GrpcTransport`] to its
/// address. Prefer [`serve_authed`] in production — an open `agent-grpc` port lets
/// anyone dispatch work to the box.
pub async fn serve(
addr: std::net::SocketAddr,
agent: Arc<LocalAgent>,
) -> Result<(), tonic::transport::Error> {
serve_authed(addr, agent, Arc::new(String::new())).await
}
/// Serve an [`AgentWorkerSvc`] behind the **shared bearer-auth interceptor** and the
/// shared serve loop (LAW #5: the wire auth is composed from `nornir-service`, never
/// re-rolled). This is what secures the otherwise-unauthenticated `agent-grpc` wire.
///
/// * `token` empty ⇒ **open** (dev / single-box bring-up); the interceptor passes
/// every request through.
/// * `token` non-empty ⇒ every RPC MUST carry `authorization: Bearer <token>`
/// (constant-time compared by `nornir_service::token_matches`), else the call is
/// rejected `UNAUTHENTICATED` before it reaches the worker — an unauthenticated
/// `Dispatch`/`Kill` never runs work on the box.
pub async fn serve_authed(
addr: std::net::SocketAddr,
agent: Arc<LocalAgent>,
token: Arc<String>,
) -> Result<(), tonic::transport::Error> {
let interceptor = nornir_service::auth_interceptor(token);
let svc = pb::agent_worker_server::AgentWorkerServer::with_interceptor(
AgentWorkerSvc::new(agent),
interceptor,
);
let routes = tonic::service::Routes::new(svc);
nornir_service::serve(routes, addr).await
}
/// The reusable **worker entrypoint** — the primitive behind a `<host> agent --group
/// <g>` command. Builds a [`LocalAgent`] from `spec` + `runner` and serves it over
/// the secured [`serve_authed`] loop, so a worker box is one call. Per LAW #5a the
/// leaf ships the entrypoint; the conductor (nornir/dwarves) wires the CLI flag
/// parsing (`agent --group …`) and picks the `runner` (its payload → job mapping).
///
/// ```ignore
/// // in nornir's `agent` subcommand, after parsing --group/--bind/--token:
/// let spec = AgentSpec::new(host, group, caps, cores, mem_mb);
/// let runner: LocalRunner = Arc::new(|unit, kill| run_payload(unit, kill));
/// jera::agent_grpc::run_worker(bind, spec, token, runner).await?;
/// ```
pub async fn run_worker(
bind: std::net::SocketAddr,
spec: crate::agent::AgentSpec,
token: Arc<String>,
runner: crate::agent::LocalRunner,
) -> Result<(), tonic::transport::Error> {
let agent = Arc::new(LocalAgent::new(spec, runner));
serve_authed(bind, agent, token).await
}
// ── client: GrpcTransport impl AgentTransport ────────────────────────────────────
/// Attach `authorization: Bearer <token>` to a request when `token` is non-empty
/// (empty ⇒ the open/dev path — no header, matching an open [`serve`]). This is the
/// client half of the shared bearer auth `nornir-service` enforces server-side.
fn with_bearer<T>(msg: T, token: &str) -> tonic::Request<T> {
let mut req = tonic::Request::new(msg);
if !token.is_empty() {
if let Ok(v) = format!("Bearer {token}").parse() {
req.metadata_mut().insert("authorization", v);
}
}
req
}
/// A remote worker reached over gRPC. Sync `AgentTransport` methods bridge to tonic
/// via a private multi-thread runtime. Carries the bearer `token` the secured
/// [`serve_authed`] worker requires (empty ⇒ open).
pub struct GrpcTransport {
rt: Arc<tokio::runtime::Runtime>,
channel: tonic::transport::Channel,
token: Arc<String>,
}
impl GrpcTransport {
/// Connect to a worker at `endpoint` (e.g. `http://127.0.0.1:7900`), **open**
/// (no bearer). Use [`connect_with_token`](Self::connect_with_token) for a
/// secured worker.
pub fn connect(endpoint: impl Into<String>) -> anyhow::Result<Self> {
Self::connect_with_token(endpoint, Arc::new(String::new()))
}
/// Connect to a secured worker, sending `authorization: Bearer <token>` on every
/// RPC. Must match the token the worker's [`serve_authed`] was started with, or
/// each call is rejected `UNAUTHENTICATED` (and the agent reads as
/// [`AgentState::Offline`], never scheduled). An empty token ⇒ open (dev).
pub fn connect_with_token(
endpoint: impl Into<String>,
token: Arc<String>,
) -> anyhow::Result<Self> {
let rt = Arc::new(tokio::runtime::Builder::new_multi_thread().enable_all().build()?);
let endpoint = endpoint.into();
let channel = rt.block_on(async move {
let ep = tonic::transport::Channel::from_shared(endpoint)
.map_err(|e| anyhow::anyhow!("bad agent endpoint uri: {e}"))?;
ep.connect().await.map_err(|e| anyhow::anyhow!("connect to agent: {e}"))
})?;
Ok(Self { rt, channel, token })
}
fn client(&self) -> pb::agent_worker_client::AgentWorkerClient<tonic::transport::Channel> {
pb::agent_worker_client::AgentWorkerClient::new(self.channel.clone())
}
}
impl crate::agent::AgentTransport for GrpcTransport {
fn remote_state(&self) -> AgentState {
let mut c = self.client();
let req = with_bearer(pb::Empty {}, &self.token);
self.rt
.block_on(async move { c.state(req).await })
.map(|r| state_from_str(&r.into_inner().state))
.unwrap_or(AgentState::Offline) // unreachable/unauthorized ⇒ offline, never scheduled
}
fn remote_dispatch(&self, w: WorkUnit) -> Arc<dyn WorkHandle> {
let mut c = self.client();
let unit = pb::WorkUnitMsg {
id: w.id.clone(),
group: w.group.clone(),
required_caps: w.required_caps.iter().cloned().collect(),
// Encode the payload (the keystone) so *what* to run crosses the wire;
// "" for a legacy payload-less unit.
payload_json: w
.payload()
.map(|p| serde_json::to_string(p).unwrap_or_default())
.unwrap_or_default(),
};
let req = with_bearer(pb::DispatchReq { unit: Some(unit) }, &self.token);
let res = self.rt.block_on(async move { c.dispatch(req).await });
// A dispatch that never reached the worker (unreachable / rejected
// UNAUTHENTICATED) has no live remote status → the handle reports Failed, so
// an unauthenticated dispatch is a visible failure, not a silent stall.
let dead = res.is_err();
let id = match res {
Ok(r) => r.into_inner().id,
Err(_) => w.id,
};
Arc::new(GrpcWorkHandle {
rt: self.rt.clone(),
channel: self.channel.clone(),
token: self.token.clone(),
id,
dead,
})
}
fn remote_go_offline(&self) {
let mut c = self.client();
let req = with_bearer(pb::Empty {}, &self.token);
let _ = self.rt.block_on(async move { c.go_offline(req).await });
}
}
/// Handle to remote work: `poll`/`kill` are the client half over the wire.
pub struct GrpcWorkHandle {
rt: Arc<tokio::runtime::Runtime>,
channel: tonic::transport::Channel,
token: Arc<String>,
id: String,
/// The dispatch RPC itself failed (unreachable / `UNAUTHENTICATED`): there is no
/// remote work to poll, so `poll` reports `Failed` without a wire round-trip.
dead: bool,
}
impl WorkHandle for GrpcWorkHandle {
fn id(&self) -> &str {
&self.id
}
fn poll(&self) -> WorkStatus {
if self.dead {
return WorkStatus::Failed;
}
let mut c = pb::agent_worker_client::AgentWorkerClient::new(self.channel.clone());
let req = with_bearer(pb::WorkId { id: self.id.clone() }, &self.token);
self.rt
.block_on(async move { c.poll(req).await })
.map(|r| status_from_str(&r.into_inner().status))
.unwrap_or(WorkStatus::Failed)
}
fn kill(&self) {
let mut c = pb::agent_worker_client::AgentWorkerClient::new(self.channel.clone());
let req = with_bearer(pb::WorkId { id: self.id.clone() }, &self.token);
let _ = self.rt.block_on(async move { c.kill(req).await });
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::agent::{AgentSpec, AgentTransport, GrpcAgent, Scheduler};
use std::sync::atomic::{AtomicBool, Ordering};
use std::time::{Duration, Instant};
fn spawn_server(agent: Arc<LocalAgent>) -> String {
// Grab an ephemeral port via std, drop it, then let tonic bind it on a
// background runtime. Returns the http endpoint.
let std_l = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
let addr = std_l.local_addr().unwrap();
drop(std_l);
std::thread::spawn(move || {
let rt = tokio::runtime::Builder::new_multi_thread().enable_all().build().unwrap();
rt.block_on(async move {
let _ = serve(addr, agent).await;
});
});
// give the server a moment to bind
std::thread::sleep(Duration::from_millis(250));
format!("http://{addr}")
}
fn spawn_server_authed(agent: Arc<LocalAgent>, token: Arc<String>) -> String {
let std_l = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
let addr = std_l.local_addr().unwrap();
drop(std_l);
std::thread::spawn(move || {
let rt = tokio::runtime::Builder::new_multi_thread().enable_all().build().unwrap();
rt.block_on(async move {
let _ = serve_authed(addr, agent, token).await;
});
});
std::thread::sleep(Duration::from_millis(250));
format!("http://{addr}")
}
fn poll_until<F: Fn() -> WorkStatus>(f: F, want: WorkStatus, timeout: Duration) -> WorkStatus {
let start = Instant::now();
loop {
let s = f();
if s == want || s.is_terminal() || start.elapsed() > timeout {
return s;
}
std::thread::sleep(Duration::from_millis(20));
}
}
#[test]
fn grpc_dispatch_runs_to_done_over_the_wire() {
let spec = AgentSpec::new("remote-1", "default", ["linux"], 4, 4096);
let server_agent = Arc::new(LocalAgent::immediate(spec.clone()));
let endpoint = spawn_server(server_agent);
let transport = GrpcTransport::connect(&endpoint).expect("connect");
assert_eq!(transport.remote_state(), AgentState::Idle);
let agent = GrpcAgent::with_transport(spec, &endpoint, Arc::new(transport));
// schedule through the real Scheduler → the remote agent → over gRPC.
let mut sched = Scheduler::new();
sched.register(Arc::new(agent));
let handle = sched
.dispatch(WorkUnit::new("w-1", "default", ["linux"]))
.expect("an eligible remote agent");
let status = poll_until(|| handle.poll(), WorkStatus::Done, Duration::from_secs(3));
assert_eq!(status, WorkStatus::Done, "unit ran to Done across the wire");
}
#[test]
fn grpc_dispatch_carries_payload_over_the_wire() {
// RED-when-broken: the WorkPayload (the keystone) must survive the wire and
// reach the worker's runner — proving a remote unit says *what* to run.
use crate::agent::{LocalRunner, WorkPayload};
let spec = AgentSpec::new("remote-payload", "default", ["linux"], 4, 4096);
let seen: Arc<Mutex<Option<String>>> = Arc::new(Mutex::new(None));
let seen_r = seen.clone();
let runner: LocalRunner = Arc::new(move |w: &WorkUnit, _kill: &AtomicBool| {
match w.payload() {
Some(WorkPayload::Process(p)) => {
*seen_r.lock().unwrap() = Some(p.program.clone());
p.program == "cargo"
}
_ => false,
}
});
let server_agent = Arc::new(LocalAgent::new(spec.clone(), runner));
let endpoint = spawn_server(server_agent);
let transport = GrpcTransport::connect(&endpoint).expect("connect");
let unit = WorkUnit::new("w-pl", "default", ["linux"])
.with_payload(WorkPayload::process("cargo", ["build", "--release"]));
let handle = transport.remote_dispatch(unit);
let status = poll_until(|| handle.poll(), WorkStatus::Done, Duration::from_secs(3));
assert_eq!(status, WorkStatus::Done, "payload unit ran to Done across the wire");
assert_eq!(
seen.lock().unwrap().as_deref(),
Some("cargo"),
"the remote runner saw the payload's ProcessSpec across the wire"
);
}
#[test]
fn wire_auth_rejects_unauthenticated_dispatch_and_admits_the_authed_one() {
// RED-when-broken: the secured worker MUST reject a tokenless dispatch and
// run an authed one — the whole point of securing the agent-grpc wire.
let token = Arc::new("agent-wire-secret-token-long-enough".to_string());
let spec = AgentSpec::new("secure-1", "default", ["linux"], 4, 4096);
let server_agent = Arc::new(LocalAgent::immediate(spec.clone()));
let endpoint = spawn_server_authed(server_agent, token.clone());
// 1) No token → the interceptor rejects every RPC. State reads Offline (so
// the scheduler never routes to it) and a dispatch is a visible Failure.
let open = GrpcTransport::connect(&endpoint).expect("connect (open client)");
assert_eq!(
open.remote_state(),
AgentState::Offline,
"an unauthenticated client cannot even read state → Offline, never scheduled"
);
let denied = open.remote_dispatch(WorkUnit::new("nope", "default", ["linux"]));
assert_eq!(
poll_until(|| denied.poll(), WorkStatus::Failed, Duration::from_secs(2)),
WorkStatus::Failed,
"unauthenticated dispatch is rejected (never runs on the box)"
);
// 2) Correct token → state is Idle and the unit runs to Done.
let authed =
GrpcTransport::connect_with_token(&endpoint, token.clone()).expect("connect (authed)");
assert_eq!(authed.remote_state(), AgentState::Idle, "authed client reads live state");
let ok = authed.remote_dispatch(WorkUnit::new("yes", "default", ["linux"]));
assert_eq!(
poll_until(|| ok.poll(), WorkStatus::Done, Duration::from_secs(3)),
WorkStatus::Done,
"an authenticated dispatch runs to Done over the secured wire"
);
// 3) A WRONG token is rejected exactly like no token.
let wrong = GrpcTransport::connect_with_token(
&endpoint,
Arc::new("the-wrong-token-entirely-here".to_string()),
)
.expect("connect (wrong token)");
assert_eq!(wrong.remote_state(), AgentState::Offline, "a wrong token is rejected");
}
#[test]
fn grpc_kill_propagates_to_killed_over_the_wire() {
// a runner that loops until its kill flag is set — proves cancel crosses the wire.
let spec = AgentSpec::new("remote-2", "default", ["linux"], 4, 4096);
let runner: crate::agent::LocalRunner = Arc::new(|_w, kill: &AtomicBool| {
while !kill.load(Ordering::SeqCst) {
std::thread::sleep(Duration::from_millis(10));
}
true // returns after kill; final status resolves to Killed
});
let server_agent = Arc::new(LocalAgent::new(spec.clone(), runner));
let endpoint = spawn_server(server_agent);
let transport = GrpcTransport::connect(&endpoint).expect("connect");
let handle = transport.remote_dispatch(WorkUnit::new("w-2", "default", ["linux"]));
// let it start, then kill over the wire.
std::thread::sleep(Duration::from_millis(100));
handle.kill();
let status = poll_until(|| handle.poll(), WorkStatus::Killed, Duration::from_secs(3));
assert_eq!(status, WorkStatus::Killed, "kill crossed the wire → Killed");
}
}
#[cfg(test)]
mod handle_lifecycle_tests {
use super::*;
use crate::agent::{AgentSpec, LocalRunner};
use std::sync::atomic::{AtomicBool, Ordering};
use std::time::{Duration, Instant};
use pb::agent_worker_server::AgentWorker as _;
fn rt() -> tokio::runtime::Runtime {
tokio::runtime::Builder::new_current_thread().enable_all().build().unwrap()
}
fn svc_immediate(name: &str) -> AgentWorkerSvc {
AgentWorkerSvc::new(Arc::new(LocalAgent::immediate(AgentSpec::new(
name,
"default",
["linux"],
4,
4096,
))))
}
fn dispatch(rt: &tokio::runtime::Runtime, svc: &AgentWorkerSvc, id: &str) -> String {
let req = tonic::Request::new(pb::DispatchReq {
unit: Some(pb::WorkUnitMsg {
id: id.to_string(),
group: "default".into(),
required_caps: vec!["linux".into()],
payload_json: String::new(),
}),
});
rt.block_on(svc.dispatch(req)).expect("dispatch").into_inner().id
}
fn poll_once(rt: &tokio::runtime::Runtime, svc: &AgentWorkerSvc, id: &str) -> WorkStatus {
let req = tonic::Request::new(pb::WorkId { id: id.to_string() });
status_from_str(&rt.block_on(svc.poll(req)).expect("poll").into_inner().status)
}
fn poll_to_terminal(
rt: &tokio::runtime::Runtime,
svc: &AgentWorkerSvc,
id: &str,
) -> WorkStatus {
let start = Instant::now();
loop {
let s = poll_once(rt, svc, id);
if s.is_terminal() || start.elapsed() > Duration::from_secs(5) {
return s;
}
std::thread::sleep(Duration::from_millis(10));
}
}
/// RED-when-broken: `handles` was a registry with **no deregister** — `poll` and
/// `kill` only `get`, so a long-lived worker box pinned one `Arc<dyn WorkHandle>`
/// per unit EVER dispatched, each holding whatever the handle owns. Twenty units
/// run to completion must leave ZERO live handles behind.
#[test]
fn a_terminal_unit_releases_its_handle_on_the_worker() {
let rt = rt();
let svc = svc_immediate("leak-1");
for i in 0..20 {
let id = format!("u-{i}");
let wid = dispatch(&rt, &svc, &id);
assert_eq!(poll_to_terminal(&rt, &svc, &wid), WorkStatus::Done, "{id} ran to Done");
}
assert_eq!(
svc.live_handles(),
0,
"every unit reached a terminal state ⇒ the worker holds no live handle"
);
}
/// RED-when-broken (the trap a naive `remove` walks into): the receipt path polls
/// a SECOND time after the caller already saw terminal —
/// `WorkResult::observed(job, handle, node)` calls `handle.poll()`. If the worker
/// simply forgot a released id, that second poll would hit the unknown-id branch
/// and report `Failed`, so a job that DID succeed would be recorded in the ledger
/// as failed. The terminal verdict must outlive the handle.
#[test]
fn a_released_units_verdict_survives_the_receipt_poll() {
let rt = rt();
let svc = svc_immediate("verdict-1");
let wid = dispatch(&rt, &svc, "u-receipt");
assert_eq!(poll_to_terminal(&rt, &svc, &wid), WorkStatus::Done);
assert_eq!(svc.live_handles(), 0, "handle released");
assert_eq!(
poll_once(&rt, &svc, &wid),
WorkStatus::Done,
"the remembered verdict answers the receipt poll — NOT Failed"
);
assert_eq!(
poll_once(&rt, &svc, &wid),
WorkStatus::Done,
"and it is stable across further polls"
);
}
/// RED-when-broken: nothing forces a client to poll a unit it dispatched, so
/// release-on-poll alone still lets the map grow without bound. `dispatch` sweeps
/// the finished handles as the backstop.
#[test]
fn dispatch_sweeps_the_finished_handles_nobody_polled() {
let rt = rt();
let svc = svc_immediate("sweep-1");
for i in 0..10 {
dispatch(&rt, &svc, &format!("unpolled-{i}"));
}
// let the immediate runner finish all ten without anyone polling them.
std::thread::sleep(Duration::from_millis(300));
let wid = dispatch(&rt, &svc, "the-eleventh");
assert!(
svc.live_handles() <= 1,
"dispatch reaped the ten finished handles; {} live",
svc.live_handles()
);
assert_eq!(poll_to_terminal(&rt, &svc, &wid), WorkStatus::Done);
}
/// A memo of terminal verdicts is itself unbounded unless it is capped, which
/// would only move the leak. It is FIFO-bounded at [`TERMINAL_MEMO_CAP`].
#[test]
fn the_terminal_verdict_memo_is_bounded() {
let svc = svc_immediate("memo-1");
for i in 0..(TERMINAL_MEMO_CAP * 2 + 7) {
svc.remember_terminal(&format!("m-{i}"), WorkStatus::Done);
}
assert_eq!(
svc.remembered_terminals(),
TERMINAL_MEMO_CAP,
"the memo is capped, so the fix is not a second unbounded map"
);
// The OLDEST ids were evicted; the newest are still answerable.
assert_eq!(svc.remembered_status("m-0"), None, "oldest evicted");
let last = format!("m-{}", TERMINAL_MEMO_CAP * 2 + 6);
assert_eq!(svc.remembered_status(&last), Some(WorkStatus::Done), "newest kept");
}
/// A killed unit is terminal too, so it must release like a Done one — and its
/// `Killed` verdict must survive, not decay to `Failed`.
#[test]
fn a_killed_unit_releases_and_keeps_its_killed_verdict() {
let rt = rt();
let runner: LocalRunner = Arc::new(|_w, kill: &AtomicBool| {
while !kill.load(Ordering::SeqCst) {
std::thread::sleep(Duration::from_millis(5));
}
true
});
let svc = AgentWorkerSvc::new(Arc::new(LocalAgent::new(
AgentSpec::new("kill-1", "default", ["linux"], 4, 4096),
runner,
)));
let wid = dispatch(&rt, &svc, "u-kill");
std::thread::sleep(Duration::from_millis(80));
let _ = rt
.block_on(svc.kill(tonic::Request::new(pb::WorkId { id: wid.clone() })))
.expect("kill");
assert_eq!(poll_to_terminal(&rt, &svc, &wid), WorkStatus::Killed);
assert_eq!(svc.live_handles(), 0, "a killed unit releases its handle too");
assert_eq!(
poll_once(&rt, &svc, &wid),
WorkStatus::Killed,
"the Killed verdict survives the handle"
);
}
/// An id the worker never saw is still `Failed` — that branch is unchanged, and
/// it must stay distinguishable from a released Done.
#[test]
fn an_unknown_id_is_still_failed() {
let rt = rt();
let svc = svc_immediate("unknown-1");
assert_eq!(poll_once(&rt, &svc, "never-dispatched"), WorkStatus::Failed);
}
}