ive harnessed the harness
Something went wrong. Try again.
32 kB · 908 lines
Rust
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909//! Bounded process actor. Only this actor owns the control channel and child//! lifetime.//!//! Every interrupted/expired exchange invalidates the stream and terminates the//! worker. Reusing a half-read frame or accepting its late effects is forbidden.//!//! Transport lives behind `WorkerBackend`: the local backend spawns a child and//! hands over its accepted unix stream; a remote backend hands over whatever//! private framed channel its relay demultiplexed. Handshake validation, target//! negotiation, startup deadline, per-cell exchange, host-request budget,//! bounded diagnostics, interrupt and generation invalidation are shared and//! identical for both backends.use crate::{behavior::Release, protocol::*, skills::SkillRoots, target::TargetIdentity};use anyhow::{anyhow, bail, ensure, Context, Result};use std::{ collections::BTreeMap, fmt, future::Future, path::PathBuf, pin::Pin, sync::{ atomic::{AtomicBool, Ordering}, Arc, Mutex, }, time::Duration,};use tokio::{ io::{AsyncRead, AsyncReadExt, AsyncWrite}, net::UnixListener, process::{Child, Command}, sync::{mpsc, oneshot, watch}, time::{timeout, timeout_at, Instant},};
#[derive(Clone)]pub struct WorkerConfig { pub python: PathBuf, pub runner: PathBuf, pub cwd: PathBuf, pub environment: BTreeMap<String, String>, pub startup_timeout: Duration, /// Skill roots this session resolves ids against. They belong to the daemon /// rather than to one checkout, because the contract a session receives is /// already the daemon's: a peer running elsewhere must not be handed a skills /// directory that checkout never published. pub skills: SkillRoots,}impl WorkerConfig { /// Resolve a worker from explicit installation paths. /// /// The runner and workspace have independent lifetimes: a packaged runner can /// execute against a writable source checkout without importing Python from it. pub fn installed(python: PathBuf, runner: PathBuf, workspace: PathBuf) -> Self { let environment = [ "PATH", "HOME", "LANG", "LC_ALL", "TMPDIR", "LD_LIBRARY_PATH", "DYLD_LIBRARY_PATH", ] .into_iter() .filter_map(|name| std::env::var(name).ok().map(|value| (name.into(), value))) .collect(); let skills = SkillRoots::for_workspace(&workspace); Self { python, runner, cwd: workspace, environment, startup_timeout: Duration::from_secs(10), skills, } }
/// Point this session at the daemon's skill roots. Absence of a managed root /// is meaningful: builtin skills alone are a complete namespace. pub fn with_skills(mut self, skills: SkillRoots) -> Self { self.skills = skills; self }
pub fn source_tree(root: &std::path::Path) -> Self { Self::installed( std::env::var_os("KLBR_PYTHON") .map(PathBuf::from) .unwrap_or_else(|| "python3".into()), root.join("python/klbr-runtime/src/klbr_runtime/__main__.py"), root.into(), ) }}
#[derive(Clone, Debug)]pub struct HostCall { pub id: RequestId, pub execution_id: OperationId, pub request: HostRequest,}pub type HostFuture = Pin<Box<dyn Future<Output = HostReply> + Send>>;pub type HostService = Arc<dyn Fn(HostCall) -> HostFuture + Send + Sync>;
pub fn no_host_operations() -> HostService { Arc::new(|_| { Box::pin(async { HostReply::error("forbidden", "this role has no host-effect capability") }) })}
#[derive(Clone, Debug)]pub enum WorkCommand { Execute { id: OperationId, code: String, purpose: Purpose, }, Invoke { id: OperationId, slot: HookSlot, input: serde_json::Value, }, RemoteDispatch { id: OperationId, operation: String, payload: serde_json::Value, },}impl WorkCommand { fn into_message(self) -> (OperationId, Message) { match self { Self::Execute { id, code, purpose } => { (id.clone(), Message::Execute { id, code, purpose }) } Self::Invoke { id, slot, input } => (id.clone(), Message::Invoke { id, slot, input }), Self::RemoteDispatch { id, operation, payload, } => ( id.clone(), Message::RemoteDispatch { id, operation, payload, }, ), } }}struct Work { command: WorkCommand, deadline: Instant, service: HostService, reply: oneshot::Sender<Result<Outcome>>,}
#[derive(Default)]pub struct Tails { stdout: Vec<u8>, stderr: Vec<u8>,}impl Tails { pub fn stderr_tail(&self) -> Option<String> { Some(String::from_utf8_lossy(&self.stderr).into_owned()) }}#[derive(Clone, Debug)]pub struct Diagnostics { pub stdout_tail: String, pub stderr_tail: String,}
pub type BoxFuture<'a, T> = Pin<Box<dyn Future<Output = T> + Send + 'a>>;
/// Framed duplex worker control channel. Local backends hand over an accepted/// unix stream; remote backends hand over the demultiplexed relay stream.pub trait ControlChannel: AsyncRead + AsyncWrite + Unpin + Send {}impl<T: AsyncRead + AsyncWrite + Unpin + Send> ControlChannel for T {}
/// Terminal states of a launched process. Exit, signal and unknown are distinct/// outcomes: a failed transport read is not proof that a remote process did not/// run, so it must stay representable as unknown.#[derive(Clone, Debug)]pub enum ProcessEnd { Exited(i64), Signaled(String), Unknown,}impl fmt::Display for ProcessEnd { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { match self { Self::Exited(code) => write!(f, "exited({code})"), Self::Signaled(signal) => write!(f, "signaled({signal})"), Self::Unknown => write!(f, "unknown"), } }}
pub trait ProcessHandle: Send { /// Idempotent termination; must not fail on an already-dead process. fn kill(&mut self) -> BoxFuture<'_, ()>; /// Resolves once; after kill it resolves promptly. fn wait(&mut self) -> BoxFuture<'_, ProcessEnd>;}
impl ProcessEnd { pub fn from_exit_status(status: &std::process::ExitStatus) -> Self { if let Some(code) = status.code() { Self::Exited(i64::from(code)) } else if let Some(signal) = unix_signal(status) { Self::Signaled(signal) } else { Self::Unknown } }}
#[cfg(unix)]fn unix_signal(status: &std::process::ExitStatus) -> Option<String> { use std::os::unix::process::ExitStatusExt; status.signal().map(|s| s.to_string())}#[cfg(not(unix))]fn unix_signal(_status: &std::process::ExitStatus) -> Option<String> { None}
#[derive(Clone, Debug)]pub struct BehaviorSource { pub revision: String, pub location: BehaviorLocation,}#[derive(Clone, Debug)]pub enum BehaviorLocation { LocalPath(PathBuf), RemoteBundle(String),}
/// Everything needed to launch ONE worker. Identity and binding are decided by/// the shared launch path, never by the backend.pub struct WorkerLaunch { pub session_id: SessionId, pub role: Role, pub generation: Generation, pub token: String, pub behavior: Option<BehaviorSource>, pub binding: Option<TargetStamp>, pub startup_deadline: Duration,}
/// Typed startup refusals. Each is an explicit failure BEFORE any execution;/// none of them has a fallback to local execution, home directories, or/// another transport.#[derive(Debug)]pub enum BackendStartError { UnavailableTarget { target: String, reason: String }, MissingWorkspace { workspace: String }, PlatformMismatch { required: String, actual: String }, ExitedBeforeReady { detail: String }, ChannelFailed { detail: String }, Deadline,}impl fmt::Display for BackendStartError { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { match self { Self::UnavailableTarget { target, reason } => { write!(f, "unavailable target {target}: {reason}") } Self::MissingWorkspace { workspace } => { write!(f, "configured workspace does not exist: {workspace}") } Self::PlatformMismatch { required, actual } => { write!(f, "platform mismatch: required {required}, actual {actual}") } Self::ExitedBeforeReady { detail } => { write!(f, "worker exited before ready: {detail}") } Self::ChannelFailed { detail } => write!(f, "control channel failed: {detail}"), Self::Deadline => write!(f, "worker startup timed out"), } }}impl std::error::Error for BackendStartError {}
pub struct WorkerStreams { pub stdout: Box<dyn AsyncRead + Unpin + Send>, pub stderr: Box<dyn AsyncRead + Unpin + Send>,}
/// One launched worker handed to the shared actor: private control channel,/// kill/wait handle, resolved target identity (None = local), and diagnostics/// already draining into bounded tails via the shared `drain` fn. The backend/// owns its demultiplexing; the shared actor owns everything else.pub struct LaunchedWorker { pub control: Box<dyn ControlChannel>, pub process: Box<dyn ProcessHandle>, pub tails: Arc<std::sync::Mutex<Tails>>, pub drains: DrainTasks, /// The identity the connect boundary actually verified, once per launch. /// None for a local launch. pub resolved: Option<TargetIdentity>,}
/// The single transport boundary. Local and remote backends feed the same/// execution, interruption, output, callback and completion contracts; there is/// no fallback backend.pub trait WorkerBackend: Send + Sync { /// Acquire the private control channel and start the process. Typed /// refusals precede any execution; the future must resolve within /// `launch.startup_deadline` or report `BackendStartError::Deadline`. fn launch( &self, launch: WorkerLaunch, ) -> BoxFuture<'_, std::result::Result<LaunchedWorker, BackendStartError>>;}
/// The local backend: today's spawn behavior (temp directory, unix listener,/// child process, piped diagnostics) as a thin adapter.pub struct LocalBackend { config: WorkerConfig,}impl LocalBackend { pub fn new(config: WorkerConfig) -> Self { Self { config } }}impl WorkerBackend for LocalBackend { fn launch( &self, launch: WorkerLaunch, ) -> BoxFuture<'_, std::result::Result<LaunchedWorker, BackendStartError>> { Box::pin(self.launch_local(launch)) }}impl LocalBackend { async fn launch_local( &self, launch: WorkerLaunch, ) -> std::result::Result<LaunchedWorker, BackendStartError> { if !self.config.cwd.exists() { return Err(BackendStartError::MissingWorkspace { workspace: self.config.cwd.display().to_string(), }); } if launch.binding.is_some() { return Err(BackendStartError::UnavailableTarget { target: "local".into(), reason: "local launch cannot carry a remote target binding".into(), }); } if let Some(behavior) = &launch.behavior { if matches!(behavior.location, BehaviorLocation::RemoteBundle(_)) { return Err(BackendStartError::UnavailableTarget { target: "local".into(), reason: "behavior source is a remote bundle; local backend cannot serve it" .into(), }); } } let directory = tempfile::Builder::new() .prefix("klbr-") .tempdir() .map_err(|e| BackendStartError::ChannelFailed { detail: format!("temp directory: {e}"), })?; let socket_path = directory.path().join("control.sock"); // The worker is handed the skill bodies rather than the roots: one reader, // one validator, and a set that cannot change under a running turn. let skills_path = crate::skills::materialize(&self.config.skills, directory.path()) .map_err(|e| BackendStartError::ChannelFailed { detail: format!("skill payload: {e}"), })?; let listener = UnixListener::bind(&socket_path).map_err(|e| BackendStartError::ChannelFailed { detail: format!("{e}"), })?; let mut command = Command::new(&self.config.python); command .args(["-I", "-B", "-u"]) .arg(&self.config.runner) .arg("--socket") .arg(&socket_path) .arg("--role") .arg(launch.role.argument()) .arg("--session") .arg(launch.session_id.as_str()) .arg("--generation") .arg(launch.generation.as_str()) .arg("--token") .arg(&launch.token) .arg("--skills") .arg(&skills_path) .current_dir(&self.config.cwd) .env_clear() .envs(&self.config.environment) .stdin(std::process::Stdio::null()) .stdout(std::process::Stdio::piped()) .stderr(std::process::Stdio::piped()) .kill_on_drop(true); if let Some(behavior) = &launch.behavior { match &behavior.location { BehaviorLocation::LocalPath(path) => { command .arg("--behavior") .arg(path) .arg("--revision") .arg(&behavior.revision); } BehaviorLocation::RemoteBundle(_) => unreachable!("rejected above"), } } let mut child = command .spawn() .map_err(|e| BackendStartError::ChannelFailed { detail: format!("launch Python worker: {e}"), })?; let stdout = child .stdout .take() .ok_or(BackendStartError::ChannelFailed { detail: "missing stdout".into(), })?; let stderr = child .stderr .take() .ok_or(BackendStartError::ChannelFailed { detail: "missing stderr".into(), })?; let tails = Arc::new(std::sync::Mutex::new(Tails::default())); let drains = DrainTasks(vec![ tokio::spawn(drain(stdout, tails.clone(), false)), tokio::spawn(drain(stderr, tails.clone(), true)), ]); let channel = tokio::select! { result = timeout(launch.startup_deadline, async { let (stream, _) = listener .accept() .await .map_err(|e| BackendStartError::ChannelFailed { detail: format!("{e}") })?; Ok::<_, BackendStartError>(stream) }) => result.map_err(|_| BackendStartError::Deadline).and_then(|r| r), status = child.wait() => return Err(BackendStartError::ExitedBeforeReady { detail: format!("{status:?}") }), }; let channel = channel?; Ok(LaunchedWorker { control: Box::new(channel), process: Box::new(LocalProcess { child, _directory: directory, }), tails, drains, resolved: None, }) }}
struct LocalProcess { child: Child, _directory: tempfile::TempDir,}impl ProcessHandle for LocalProcess { fn kill(&mut self) -> BoxFuture<'_, ()> { Box::pin(async move { let _ = self.child.kill().await; }) } fn wait(&mut self) -> BoxFuture<'_, ProcessEnd> { Box::pin(async move { match self.child.wait().await { Ok(status) => ProcessEnd::from_exit_status(&status), Err(_) => ProcessEnd::Unknown, } }) }}
/// What the worker reported about itself at hello, after validation.struct Handshake { identity: Option<WorkerIdentity>, skills: Option<String>,}
struct Inner { queue: mpsc::Sender<Work>, stop: watch::Sender<bool>, stopped: watch::Receiver<bool>, alive: Arc<AtomicBool>, tails: Arc<Mutex<Tails>>, generation: Generation, identity: Option<WorkerIdentity>, /// Skill set the worker echoed from the spawn payload. Recorded so a durable /// reset names the instruction set that generation loaded. skills: Option<String>, /// Target identity MEASURED by the launcher's own probe (None for local). /// /// This is never the worker's echoed stamp: the echo is a binding check, and a /// consumer needing account/machine/cwd must read the measured value or none. resolved: Option<TargetIdentity>,}impl Drop for Inner { fn drop(&mut self) { let _ = self.stop.send(true); }}#[derive(Clone)]pub struct Worker(Arc<Inner>);
struct Life { alive: Arc<AtomicBool>, stopped: watch::Sender<bool>,}impl Drop for Life { fn drop(&mut self) { self.alive.store(false, Ordering::Release); let _ = self.stopped.send(true); }}pub struct DrainTasks(pub Vec<tokio::task::JoinHandle<()>>);impl Drop for DrainTasks { fn drop(&mut self) { for task in &self.0 { task.abort(); } }}
pub async fn drain<R: AsyncRead + Unpin>(mut reader: R, tails: Arc<Mutex<Tails>>, stderr: bool) { let mut buffer = [0; 4096]; loop { let count = match reader.read(&mut buffer).await { Ok(0) | Err(_) => break, Ok(n) => n, }; let Ok(mut tails) = tails.lock() else { break }; let tail = if stderr { &mut tails.stderr } else { &mut tails.stdout }; tail.extend_from_slice(&buffer[..count]); let excess = tail.len().saturating_sub(8192); tail.drain(..excess); }}
impl Worker { pub async fn spawn( config: &WorkerConfig, session_id: SessionId, role: Role, release: Option<&Release>, ) -> Result<Self> { let behavior = release.map(|r| BehaviorSource { revision: r.id.clone(), location: BehaviorLocation::LocalPath(r.path.clone()), }); Self::launch_with( &LocalBackend::new(config.clone()), config.startup_timeout, session_id, role, behavior, None, ) .await }
/// Shared launch path for every backend. The backend only produces the /// private channel, a kill/wait handle and its verified identity; handshake /// validation, negotiation, deadline and failure handling are identical. pub async fn launch_with( backend: &dyn WorkerBackend, startup_timeout: Duration, session_id: SessionId, role: Role, behavior: Option<BehaviorSource>, binding: Option<TargetStamp>, ) -> Result<Self> { ensure!( role != Role::Hooks || behavior.is_some(), "hook role requires a release" ); let behavior_revision = behavior.as_ref().map(|b| b.revision.clone()); let generation = Generation::fresh(); let token = uuid::Uuid::new_v4().to_string(); let LaunchedWorker { mut control, mut process, tails, drains, resolved, } = backend .launch(WorkerLaunch { session_id: session_id.clone(), role, generation: generation.clone(), token: token.clone(), behavior, binding: binding.clone(), startup_deadline: startup_timeout, }) .await .context("worker backend refused launch")?; let expected_binding: Option<HelloBinding> = match (binding.as_ref(), resolved.as_ref()) { (Some(stamp), Some(identity)) => Some(HelloBinding { stamp: stamp.clone(), identity: Some(identity.clone()), }), (Some(stamp), None) => { process.kill().await; return Err(anyhow!( "target {stamp}: launch did not resolve a verified target identity" )); } (None, _) => None, }; let handshake = async { let envelope = read_frame(&mut control).await?; ensure!( envelope.session_id == session_id && envelope.generation == generation, "foreign handshake identity" ); let handshake = match envelope.message { Message::Hello { token: actual_token, role: actual_role, revision, slots, identity, skills, } => { ensure!( actual_token == token && actual_role == role, "worker handshake mismatch" ); ensure!(revision == behavior_revision, "worker revision mismatch"); let expected = if role == Role::Hooks { vec![HookSlot::Attention, HookSlot::Delivery] } else { vec![] }; ensure!(slots == expected, "worker hook contract mismatch"); if let Some(reported) = &identity { WorkerIdentity::validate(reported, expected_binding.as_ref())?; } else if expected_binding.is_some() { bail!(NegotiationError::TargetBindingMissing); } Handshake { identity, skills } } _ => bail!("expected hello"), }; Ok(handshake) }; let startup = tokio::select! { result = timeout(startup_timeout, handshake) => result.map_err(|_| anyhow!("worker handshake timed out")).and_then(|result| result), end = process.wait() => Err(anyhow!("worker exited before ready: {end}")), }; let handshake = match startup { Ok(handshake) => handshake, Err(error) => { process.kill().await; let diagnostic = tails .lock() .ok() .and_then(|t| t.stderr_tail()) .unwrap_or_default(); return Err( error.context(format!("worker startup failed; stderr tail: {diagnostic}")) ); } }; let (queue, mut requests) = mpsc::channel::<Work>(16); let (stop, mut stopping) = watch::channel(false); let (stopped_tx, stopped) = watch::channel(false); let alive = Arc::new(AtomicBool::new(true)); let life = Life { alive: alive.clone(), stopped: stopped_tx, }; let exchange_identity = (session_id, generation.clone()); tokio::spawn(async move { let life_guard = life; let _drains = drains; loop { let work = tokio::select! { biased; _ = stopping.changed() => break, end = process.wait() => { let _ = end; break; } work = requests.recv() => match work { Some(work) => work, None => break }, }; let Work { command, deadline, service, mut reply, } = work; if reply.is_closed() { continue; } if Instant::now() >= deadline { let _ = reply.send(Err(anyhow!("worker queue deadline exceeded"))); continue; } let result = tokio::select! { biased; _ = stopping.changed() => Err(anyhow!("worker stopped by operator")), _ = reply.closed() => Err(anyhow!("caller cancelled; worker generation invalidated")), result = timeout_at(deadline, exchange(&mut control, &exchange_identity, role, command, &service, binding.as_ref())) => result.map_err(|_| anyhow!("worker deadline exceeded; generation invalidated")).and_then(|r| r), }; let failed = result.is_err(); if failed { life_guard.alive.store(false, Ordering::Release); } let _ = reply.send(result); if failed { break; } } // Stop admissions before draining the queue, including process-exit // paths that did not finish a request. life_guard.alive.store(false, Ordering::Release); requests.close(); // No reuse of a channel with an unknown framing/cancellation state. process.kill().await; while let Ok(work) = requests.try_recv() { let _ = work.reply.send(Err(anyhow!("worker generation has ended"))); } }); Ok(Self(Arc::new(Inner { queue, stop, stopped, alive, tails, generation, identity: handshake.identity, skills: handshake.skills, resolved, }))) }
pub fn generation(&self) -> &Generation { &self.0.generation } /// The validated runtime identity the worker reported at hello, if any. pub fn worker_identity(&self) -> Option<WorkerIdentity> { self.0.identity.clone() } /// The skill set this generation loaded, as the worker echoed it. pub fn skills_revision(&self) -> Option<String> { self.0.skills.clone() } /// The target identity MEASURED by the launcher's probe, if this is a remote worker. /// /// Returns `None` for a local worker. Callers needing account/machine/cwd must use /// this rather than the runtime fields of [`Self::worker_identity`]: those are /// reported by the worker, while this was measured about it. An absent measurement /// must surface as absent, never be substituted with self-reported values. pub fn target_identity(&self) -> Option<TargetIdentity> { self.0.resolved.clone() } pub fn is_alive(&self) -> bool { self.0.alive.load(Ordering::Acquire) && !*self.0.stop.borrow() } pub fn diagnostics(&self) -> Diagnostics { match self.0.tails.lock() { Ok(t) => Diagnostics { stdout_tail: String::from_utf8_lossy(&t.stdout).into(), stderr_tail: String::from_utf8_lossy(&t.stderr).into(), }, Err(_) => Diagnostics { stdout_tail: String::new(), stderr_tail: "diagnostic lock poisoned".into(), }, } }
pub async fn call( &self, command: WorkCommand, budget: Duration, service: HostService, ) -> Result<Outcome> { ensure!(self.is_alive(), "worker generation is unavailable"); let (reply, receive) = oneshot::channel(); self.0 .queue .try_send(Work { command, deadline: Instant::now() + budget, service, reply, }) .map_err(|_| anyhow!("worker queue full or closed"))?; receive .await .context("worker ended before completing request")? }
/// Independent of the request queue. Initially uses a hard process stop; /// cooperative in-place cancellation is supported by Python but not trusted /// as the host's recovery guarantee in this slice. pub async fn stop(&self) { let _ = self.0.stop.send(true); let mut stopped = self.0.stopped.clone(); let _ = stopped.wait_for(|value| *value).await; }}
async fn exchange( stream: &mut Box<dyn ControlChannel>, identity: &(SessionId, Generation), role: Role, command: WorkCommand, service: &HostService, binding: Option<&TargetStamp>,) -> Result<Outcome> { ensure!( matches!( (role, &command), (Role::Workbench, WorkCommand::Execute { .. }) | (Role::Hooks, WorkCommand::Invoke { .. }) | (Role::Remote, WorkCommand::RemoteDispatch { .. }) ), "command is not allowed for worker role" ); match &command { WorkCommand::Execute { code, .. } => ensure!(code.len() <= MAX_CODE, "cell too large"), WorkCommand::RemoteDispatch { operation, payload, .. } => { ensure!(operation.len() <= 128, "remote operation name too large"); ensure!( serde_json::to_vec(payload)?.len() <= MAX_REMOTE_VALUE, "remote payload too large" ); } WorkCommand::Invoke { .. } => {} } let (operation_id, message) = command.into_message(); let envelope = |message| Envelope { version: VERSION, session_id: identity.0.clone(), generation: identity.1.clone(), message, target: binding.cloned(), }; write_frame(stream, &envelope(message)).await?;
let mut requests = 0; let mut yielded = false; loop { let frame = read_frame(stream).await?; ensure!( frame.session_id == identity.0 && frame.generation == identity.1, "stale worker envelope" ); validate_envelope_target(&frame, binding)?; match frame.message { Message::Completed { id, outcome } => { ensure!(id == operation_id, "wrong completed operation"); outcome.validate(role)?; return Ok(outcome); } Message::HostRequest { id, execution_id, request, } => { ensure!( role == Role::Workbench && execution_id == operation_id, "foreign host request" ); requests += 1; ensure!(requests <= 128, "per-cell host-request budget exhausted"); let is_wait = matches!(request, HostRequest::Wait { .. }); let reply = if yielded { HostReply::error("scope_closed", "execution has already yielded") } else { service(HostCall { id: id.clone(), execution_id: execution_id.clone(), request, }) .await }; if is_wait && matches!(reply, HostReply::Ok { .. }) { yielded = true; } write_frame( stream, &envelope(Message::HostReply { id, execution_id, reply, }), ) .await?; } _ => bail!("unexpected worker message"), } }}