Build: RemoteWorker and the Worker Binary

Maps to: Phase 6 (cluster). Kind: Build.

Objective

Close the seam. You have a wire (MessageStream) and, from the concept chapter, the shape of the machine that drives it. Here you build both ends: on the coordinator, a RemoteWorker that is a WorkerHandle like any other, a connection_actor that multiplexes one socket, a SharedPool the scheduler snapshots, and an accept loop that turns each incoming connection into a registered worker. On the far side, the panoptes-worker binary that connects, registers, and serves eval jobs. When this compiles green, the scheduler can dispatch a job to a process on another machine and never know the difference — the promise Part I made, paid in full.

This is the largest build in the course, and it spans two crates. Take it one test at a time; each named test locks one piece.

Scaffold

Create — coordinator side (panoptes-control):

• crates/panoptes-control/src/remote.rs — RemoteWorker, SharedPool, connection_actor, serve_workers, handle_connection, and two tests.

Edit:

• crates/panoptes-control/src/lib.rs — add pub mod remote; and re-export SharedPool, serve_workers.
• crates/panoptes-control/Cargo.toml — no new external deps; tokio (with net), async-trait, and control-core are already present. The tokio::sync::{mpsc, oneshot} primitives are part of tokio's sync/full features you already enabled.

Create — worker side (new crate panoptes-worker; add crates/panoptes-worker to the workspace members):

• crates/panoptes-worker/Cargo.toml
  • [dependencies]: control-core, control-eval (both { path = ... }), plus tokio (features ["full"]), async-trait, anyhow, clap (features ["derive"]), tracing, tracing-subscriber (features ["env-filter"]) — all { workspace = true } where the workspace pins them.
    • control-eval gives you ModelClient, HttpModelClient, and run_eval — the workload the worker runs.
    • clap parses the binary's flags; tracing-subscriber sets up logging in main.
  • [dev-dependencies]: async-trait, tokio (for the StubClient and #[tokio::test]).
• crates/panoptes-worker/src/lib.rs — run_session, serve_worker, run_worker_forever, DEFAULT_HEARTBEAT, and one test.
• crates/panoptes-worker/src/main.rs — the binary: Cli, #[tokio::main], wire up run_worker_forever.

Expected result:

• cargo test -p panoptes-control → the scheduler-arc tests still pass, plus 2 new: remote_worker_dispatches_over_the_wire, worker_death_makes_dispatch_retryable.
• cargo test -p panoptes-worker → 1 test: worker_registers_runs_a_job_and_returns_a_result.

The spec (givens)

RemoteWorker and SharedPool

/// A job handed to the connection actor plus the channel to answer it on.
type Dispatch = (Job, oneshot::Sender<Result<JobOutcome, ControlError>>);

/// Coordinator-side handle to a worker across a TCP connection. To the scheduler
/// it is an ordinary `WorkerHandle`; `dispatch` hides the round-trip.
pub struct RemoteWorker {
    id: String,
    tx: mpsc::Sender<Dispatch>,
}

#[async_trait]
impl WorkerHandle for RemoteWorker {
    fn id(&self) -> &str;
    async fn dispatch(&self, job: Job) -> Result<JobOutcome, ControlError>;
}

/// A live, shared set of workers. The accept loop adds one on register and
/// removes it on disconnect; the scheduler snapshots it per run. Cheap to clone.
#[derive(Clone, Default)]
pub struct SharedPool {
    workers: Arc<Mutex<Vec<Arc<dyn WorkerHandle>>>>,
}

• RemoteWorker::dispatch is the whole illusion: make a oneshot, self.tx.send((job, reply_tx)).await, then reply_rx.await. Both failure paths map to ControlError::Worker(...) — a send error means the actor (and thus the connection) is already gone; a dropped reply_rx means the actor died with the job in flight. Both are retryable, which is exactly what makes the scheduler redeliver. Note the double ?/unwrap shape: reply_rx.await yields Result<Result<JobOutcome, ControlError>, RecvError> — the outer error is the dropped channel, the inner is the actor's own answer.
• SharedPool needs add(Arc<dyn WorkerHandle>), remove(&str) (retain by id()), len(), is_empty(), and an impl WorkerSource whose snapshot(&self) -> Vec<Arc<dyn WorkerHandle>> clones the inner vec. WorkerSource is the scheduler-arc trait — this is the second implementation of it, the remote counterpart to the local pool, and the scheduler cannot tell them apart.

→ Answer key

connection_actor

async fn connection_actor(
    mut conn: MessageStream,
    mut jobs: mpsc::Receiver<Dispatch>,
    capacity: usize,
);

Own the socket; multiplex it. Keep a HashMap<JobId, oneshot::Sender<Result<JobOutcome, ControlError>>> of in-flight jobs. Loop over a tokio::select! of two arms:

• maybe = jobs.recv(), if pending.len() < capacity — the guarded arm. On Some((job, reply)): capture job.id, send Message::Assign { job } over conn; if that send fails, answer reply with the error and break (the wire is broken); otherwise insert reply into pending under the job id. On None (all RemoteWorker handles dropped), break.
• frame = conn.recv() — the socket arm. Match: Ok(Some(Message::Result { outcome })) → remove outcome.job_id from pending and send Ok(outcome) to its waiter; Ok(Some(Message::Heartbeat { .. })) → ignore (liveness only); Ok(Some(_)) → ignore (an inbound Register/Assign is a protocol slip, not fatal); Ok(None) | Err(_) → break (closed or malformed).

After the loop, the crucial cleanup: drain pending and send each waiter Err(ControlError::Worker("worker connection lost".into())). Retryable, so every job the dead worker was holding gets redelivered. This is the line the concept chapter called the pivot of the arc.

→ Answer key

serve_workers and handle_connection

/// Accept worker connections forever, registering each into `pool`.
pub async fn serve_workers(listener: TcpListener, pool: SharedPool);

/// Register one connection and run its actor until the socket closes.
async fn handle_connection(sock: TcpStream, pool: SharedPool);

• serve_workers loops on listener.accept(); for each socket, clone the pool and tokio::spawn(handle_connection(sock, pool)). An accept error should not kill the loop — continue.
• handle_connection runs the handshake: wrap the socket in a MessageStream, recv the first frame, and require it to be Message::Register { worker_id, capacity }. Anything else (or a closed/errored socket) → return, dropping the connection. On a valid register: make mpsc::channel(capacity) (clamp capacity to at least 1), build Arc::new(RemoteWorker { id: worker_id, tx }), pool.add(...), run connection_actor(conn, jobs, capacity).await, and when it returns, pool.remove(&worker_id). Add-then-actor-then-remove is the worker's whole lifecycle in the pool.

→ Answer key

The worker binary — run_session, serve_worker, run_worker_forever, main

pub const DEFAULT_HEARTBEAT: Duration = Duration::from_secs(5);

/// Register, then serve jobs until the coordinator closes (`Ok`) or the wire fails (`Err`).
pub async fn run_session(
    mut conn: MessageStream,
    worker_id: &str,
    capacity: u32,
    client: &dyn ModelClient,
    heartbeat: Duration,
) -> Result<(), ControlError>;

/// Connect to the coordinator and run one session.
pub async fn serve_worker(
    coordinator: &str,
    worker_id: &str,
    capacity: u32,
    client: Arc<dyn ModelClient>,
    heartbeat: Duration,
) -> Result<(), ControlError>;

/// Serve forever, reconnecting after any disconnect.
pub async fn run_worker_forever(
    coordinator: &str,
    worker_id: &str,
    capacity: u32,
    client: Arc<dyn ModelClient>,
    heartbeat: Duration,
    reconnect_delay: Duration,
);

• run_session first sends Message::Register { worker_id, capacity }. Then it sets up a tokio::time::interval(heartbeat) and skips the immediate first tick. Then it loops a select! of two arms: beat.tick() → send Message::Heartbeat { worker_id }; conn.recv() → match Some(Message::Assign { job }) (capture job.id, run_eval(client, &job.spec).await?, then send Message::Result { outcome: JobOutcome { job_id, records } }), Some(_) (ignore — the coordinator only ever sends Assign), None (clean close → return Ok(())). A failed eval ?-propagates and ends the session; the coordinator will notice the dropped connection and redeliver. The worker is deliberately simple: no map, no multiplexing — that all lives on the coordinator.
• serve_worker connects a TcpStream, wraps it in MessageStream::new, and calls run_session.
• run_worker_forever loops serve_worker forever, logging the outcome (Ok → coordinator closed; Err → session failed) and sleeping reconnect_delay between attempts. Workers are cattle; supervision is just a loop.
• main is #[tokio::main]: parse a clap Cli (a --coordinator address, --model-api base URL, --model, --id, --capacity), init a tracing_subscriber with an EnvFilter (default info), build Arc::new(HttpModelClient::new(model_api, model)), log a startup line, and call run_worker_forever(..., DEFAULT_HEARTBEAT, Duration::from_secs(1)).

→ Answer key

Concepts exercised

• A second impl WorkerHandle (RemoteWorker) and a second impl WorkerSource (SharedPool) — the distributed layer arriving as new implementations of existing seams, the scheduler untouched.
• The connection actor: select! with a capacity guard, an id→waiter HashMap, and retryable failure of all pending jobs on close.
• A TCP accept loop that spawns one task per connection and a register handshake that gates the connection.
• The worker session loop: interleaving heartbeats with job service over one MessageStream via select!.
• A supervised reconnect loop and a clap binary that composes the whole worker.

The build loop (you drive)

Test 1 — remote_worker_dispatches_over_the_wire (in remote.rs, #[tokio::test])

1. Write the failing test. Stand up a coordinator: bind 127.0.0.1:0, make a SharedPool, spawn serve_workers(listener, pool.clone()). Spawn a fake worker task: connect, send Register { worker_id: "w1", capacity: 2 }, then loop recv and for each Assign { job } build one ResponseRecord per vignette and send back Result { outcome: JobOutcome { job_id: job.id, records } }. Poll the pool until it holds 1 worker, snapshot().pop() it, dispatch(a_job()).await.unwrap(), and assert the outcome has one record.
2. Predict: the fake worker answers on the same socket the actor is reading. Which task inserts into pending, and which removes from it? Trace the one job's oneshot from dispatch to the assert.
3. Run — fails to compile (nothing built yet).
4. Implement RemoteWorker, SharedPool, connection_actor, serve_workers, handle_connection.
5. Run green, commit.

Test 2 — worker_death_makes_dispatch_retryable (in remote.rs, #[tokio::test])

1. Write the failing test. Same coordinator. This fake worker registers (capacity: 1), receives its one Assign, then drops the socket without answering — a crash mid-job. Poll for 1 worker, snapshot-pop it, dispatch(a_job()).await.unwrap_err(), and assert err.is_retryable().
2. Predict: the worker never sends a Result, so the pending entry is never removed by the socket arm. What removes it, and what error does the waiter's oneshot receive? Which cleanup path in connection_actor fires?
3. Run, confirm the assertion pins the retryable error.
4. Run green, commit.

Test 3 — worker_registers_runs_a_job_and_returns_a_result (in panoptes-worker, #[tokio::test])

1. Write the failing test. This time the coordinator is faked and the worker is real. Define a StubClient: ModelClient (no network — generate returns ModelResponse { text: format!("re: {prompt}"), usage: .. }). Bind 127.0.0.1:0; spawn a fake coordinator that accepts, asserts the first frame is Register, sends one Assign { job } (a job with one vignette "a"/"hi"), then loops recv ignoring heartbeats until it sees Result { outcome } and asserts outcome.job_id matches, one record, records[0].response == "re: hi". In the main task, run serve_worker(addr, "w1", 1, Arc::new(StubClient), Duration::from_secs(30)) under a tokio::time::timeout (it serves forever, so time it out).
2. Predict: the coordinator sends one Assign and then never closes. Why does wrapping serve_worker in a timeout matter — what would happen without it, and is that a bug or expected?
3. Run — fails to compile (worker crate not built).
4. Implement run_session, serve_worker, run_worker_forever, and main.rs.
5. Run green, commit.

Done when

cargo test -p panoptes-control shows the scheduler-arc tests plus remote_worker_dispatches_over_the_wire and worker_death_makes_dispatch_retryable green; cargo test -p panoptes-worker shows worker_registers_runs_a_job_and_returns_a_result green; a whole-workspace cargo test is green; and you can trace one job from RemoteWorker::dispatch, through the actor's Assign, across the socket to the worker's run_eval, back as a Result, and out the job's oneshot — and say why a worker's death turns that same path into a retryable error. Commit. The scheduler now dispatches across machines through the exact seam it has always used.

The cluster is wired but not yet safe: a frozen worker that never closes its socket would stall a job forever, and a redelivered job could be recorded twice. Those are the capstone — heartbeat reaping and at-least-once delivery made safe by the store's idempotency — where the whole course pays off.