From 73a687647cef0b62da34a8f0583640bf8c0c597f Mon Sep 17 00:00:00 2001 From: kartikeya-27 <67143288+kartikeya-27@users.noreply.github.com> Date: Fri, 31 Jul 2026 12:14:55 +0530 Subject: [PATCH 01/15] docs: correct the stale worker.rs paths in the contract The module became a directory; both references still pointed at the old file. --- crates/taskito-core/BINDING_CONTRACT.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/crates/taskito-core/BINDING_CONTRACT.md b/crates/taskito-core/BINDING_CONTRACT.md index b1b735740..091513c1b 100644 --- a/crates/taskito-core/BINDING_CONTRACT.md +++ b/crates/taskito-core/BINDING_CONTRACT.md @@ -68,7 +68,7 @@ handler-binding model. ## Dispatch call sequence 1. Shell constructs `Storage` (SQLite default; `postgres`/`redis` features) — `storage/traits.rs`. 2. Shell constructs `Scheduler::new(storage, queues, SchedulerConfig, namespace)` — `scheduler/mod.rs`. -3. Shell implements `WorkerDispatcher` — `worker.rs`. +3. Shell implements `WorkerDispatcher` — `worker/mod.rs`. 4. `Scheduler.run(job_tx)` polls + claims jobs and sends each `Job` over a `tokio::sync::mpsc::Sender` — `scheduler/poller.rs`, `scheduler/mod.rs`. 5. `WorkerDispatcher::run(job_rx, result_tx)` receives the `Job`, deserializes @@ -86,7 +86,7 @@ Scheduler.handle_result ─▶ ResultOutcome ─▶ shell emits events / middlew ``` ## What a shell MUST implement -### `WorkerDispatcher` — `worker.rs` +### `WorkerDispatcher` — `worker/mod.rs` | Method | Signature | Required | |--------|-----------|----------| | `run` | `async fn run(&self, job_rx: tokio::sync::mpsc::Receiver, result_tx: crossbeam_channel::Sender)` | yes | From ea29f4aabb92ee60aa35d29512bcec4bb7096e36 Mon Sep 17 00:00:00 2001 From: kartikeya-27 <67143288+kartikeya-27@users.noreply.github.com> Date: Fri, 31 Jul 2026 12:15:11 +0530 Subject: [PATCH 02/15] feat(core): add a shared worker frame protocol One JSON header line then the payload's raw bytes, so the same format serves a pipe and a socket and the blob stays the wire-envelope bytes. Closes #547. --- crates/taskito-core/BINDING_CONTRACT.md | 40 ++ crates/taskito-core/src/lib.rs | 3 +- crates/taskito-core/src/worker/mod.rs | 2 + crates/taskito-core/src/worker/protocol.rs | 664 +++++++++++++++++++++ 4 files changed, 708 insertions(+), 1 deletion(-) create mode 100644 crates/taskito-core/src/worker/protocol.rs diff --git a/crates/taskito-core/BINDING_CONTRACT.md b/crates/taskito-core/BINDING_CONTRACT.md index 091513c1b..ea48434d1 100644 --- a/crates/taskito-core/BINDING_CONTRACT.md +++ b/crates/taskito-core/BINDING_CONTRACT.md @@ -96,6 +96,46 @@ Scheduler.handle_result ─▶ ResultOutcome ─▶ shell emits events / middlew Channels: inbound `tokio::sync::mpsc::Receiver` (async); outbound `crossbeam_channel::Sender` (sync, cloneable). +## Worker frame protocol (out-of-process executors) — `worker/protocol.rs` + +A dispatcher that runs tasks in another process speaks this format over its +stream. The same format serves a pipe (the prefork pool's stdio children) and a +socket, so an executor written in any SDK attaches to any scheduler. + +A frame is a JSON header line, then exactly the number of raw payload bytes the +header declares: + +``` +{"type":"job","id":"018f…","task_name":"resize","payload_len":7,…}\n +<7 raw bytes> +``` + +The blob is **not** base64-encoded — the bytes on the wire are the wire-envelope +bytes of the section above, unchanged. `MAX_HEADER_BYTES` (64 KiB) and +`MAX_PAYLOAD_BYTES` (64 MiB) bound a desynced or hostile peer. + +| Frame | Direction | Payload | +|---|---|---| +| `hello` | executor → scheduler | `{executor_id, sdk, version, tasks[], slots, protocol_version}` | +| `hello_ack` | scheduler → executor | `{scheduler_id, protocol_version}` | +| `heartbeat` | executor → scheduler | `{free_slots}` | +| `job` | scheduler → executor | `{id, task_name, payload_len, retry_count, max_retries, queue, timeout_ms, namespace}` + blob | +| `cancel` | scheduler → executor | `{job_id}` | +| `shutdown` | scheduler → executor | — | +| `success` | executor → scheduler | `{job_id, result_len, task_name, wall_time_ns}` + blob | +| `failure` | executor → scheduler | `{job_id, error, retry_count, max_retries, task_name, wall_time_ns, should_retry, timed_out}` | +| `cancelled` | executor → scheduler | `{job_id, task_name, wall_time_ns}` | + +Rules: +- `hello` is the first frame on every connection; no `job` may precede its ack. +- Both sides announce `protocol_version` and both reject a mismatch. A version + is never silently downgraded. The scheduler sends `hello_ack` even when it is + rejecting, so both ends can log both versions. +- `result_len: null` means the task returned nothing; `0` means it returned an + empty value. They are distinct. +- `should_retry` is the executor's decision — only it can see the exception. The + core never inspects one. + ## Task errors (structured, cross-SDK) When a task raises, the shell reports the failure as a **canonical JSON object** serialized into `JobResult::Failure.error` (and thus into `jobs.error`, diff --git a/crates/taskito-core/src/lib.rs b/crates/taskito-core/src/lib.rs index 66126f7d8..fff2906db 100644 --- a/crates/taskito-core/src/lib.rs +++ b/crates/taskito-core/src/lib.rs @@ -46,5 +46,6 @@ pub use storage::Storage; pub use storage::StorageBackend; pub use storage::{DeadJob, QueueStats, SubscriptionBacklogStats}; pub use worker::{ - NativeDispatcher, TaskError, TaskRegistry, TaskResult, Worker, WorkerDispatcher, WorkerHandle, + ExecutorMessage, NativeDispatcher, ProtocolError, SchedulerMessage, TaskError, TaskRegistry, + TaskResult, Worker, WorkerDispatcher, WorkerHandle, PROTOCOL_VERSION, }; diff --git a/crates/taskito-core/src/worker/mod.rs b/crates/taskito-core/src/worker/mod.rs index 922eef608..2f028cee0 100644 --- a/crates/taskito-core/src/worker/mod.rs +++ b/crates/taskito-core/src/worker/mod.rs @@ -1,8 +1,10 @@ pub mod dispatcher; +pub mod protocol; pub mod registry; pub mod runner; pub use dispatcher::NativeDispatcher; +pub use protocol::{ExecutorMessage, ProtocolError, SchedulerMessage, PROTOCOL_VERSION}; pub use registry::{TaskError, TaskHandler, TaskRegistry, TaskResult}; pub use runner::{Worker, WorkerHandle}; diff --git a/crates/taskito-core/src/worker/protocol.rs b/crates/taskito-core/src/worker/protocol.rs new file mode 100644 index 000000000..7eb87570a --- /dev/null +++ b/crates/taskito-core/src/worker/protocol.rs @@ -0,0 +1,664 @@ +//! Wire protocol shared by every out-of-process worker transport. +//! +//! A frame is a JSON header line followed by exactly the number of raw payload +//! bytes it declares: +//! +//! ```text +//! {"type":"job","id":"018f…","task_name":"resize","payload_len":7,…}\n +//! <7 raw bytes> +//! ``` +//! +//! The blob stays raw instead of base64 inside the header so the bytes on the +//! wire *are* the wire-envelope bytes of `BINDING_CONTRACT.md`. Headers stay +//! JSON so every SDK can write an executor with its standard library alone. +//! The same format serves a pipe (prefork's stdio children) and a socket. + +use std::io::{BufRead, BufWriter, Read, Write}; + +use serde::de::DeserializeOwned; +use serde::{Deserialize, Serialize}; + +use crate::job::Job; +use crate::scheduler::JobResult; + +/// Frame format version. Both sides announce it in the handshake; a mismatch +/// is rejected rather than silently downgraded. +pub const PROTOCOL_VERSION: u32 = 1; + +/// Header cap, bounding a peer that never sends a newline. +pub const MAX_HEADER_BYTES: u64 = 64 * 1024; + +/// Payload cap. A header declares its own length, so without this a corrupt +/// length field would allocate unboundedly. +pub const MAX_PAYLOAD_BYTES: usize = 64 * 1024 * 1024; + +/// Errors raised while framing or parsing a message. +#[derive(Debug, thiserror::Error)] +pub enum ProtocolError { + /// The underlying stream failed. + #[error("protocol I/O error: {0}")] + Io(#[from] std::io::Error), + + /// Peer closed cleanly *between* frames — an orderly shutdown, not + /// corruption. A close mid-frame surfaces as [`ProtocolError::Io`]. + #[error("peer closed the connection")] + Eof, + + /// A header line hit [`MAX_HEADER_BYTES`] without a newline. + #[error("frame header exceeds {MAX_HEADER_BYTES} bytes")] + HeaderTooLarge, + + /// A header declared a payload over [`MAX_PAYLOAD_BYTES`]. + #[error("frame payload of {len} bytes exceeds the {MAX_PAYLOAD_BYTES} byte limit")] + PayloadTooLarge { + /// Length the header declared. + len: usize, + }, + + /// The header line was not valid JSON for the expected frame type. + #[error("malformed frame header: {0}")] + Json(#[from] serde_json::Error), + + /// Header length disagreed with the payload handed to the writer. Always a + /// caller bug — the reader would desync on it. + #[error("frame declared {declared} payload bytes but {actual} were supplied")] + PayloadLengthMismatch { + /// Length the header declared. + declared: usize, + /// Length actually supplied. + actual: usize, + }, + + /// The peer speaks a different version of this protocol. + #[error("protocol version mismatch: we speak {ours}, peer speaks {theirs}")] + VersionMismatch { + /// Version this build speaks. + ours: u32, + /// Version the peer announced. + theirs: u32, + }, + + /// A valid frame arrived where the state machine expected another kind. + #[error("expected a {expected} frame")] + UnexpectedFrame { + /// Frame the reader was waiting for. + expected: &'static str, + }, +} + +/// A message the scheduler sends to an executor. +#[derive(Debug, Serialize, Deserialize)] +#[serde(tag = "type", rename_all = "snake_case")] +pub enum SchedulerMessage { + /// Answer to [`ExecutorMessage::Hello`], completing the handshake. + HelloAck { + /// Identity of the scheduler that accepted the attach. + scheduler_id: String, + /// Version the scheduler speaks, so a rejected peer can log both. + protocol_version: u32, + }, + /// Run a job. The task payload follows the header as raw bytes. + Job { + /// Job id, echoed back on the result frame. + id: String, + /// Task to run. + task_name: String, + /// Length of the payload blob that follows. + payload_len: usize, + /// Retries already attempted. + retry_count: i32, + /// The job's retry cap. + max_retries: i32, + /// Queue the job came from. + queue: String, + /// Execution timeout in milliseconds; `<= 0` means none. + timeout_ms: i64, + /// Namespace the job belongs to. + namespace: Option, + }, + /// Cooperative-cancel request, so the executor observes a cancel without + /// either side polling storage. + Cancel { + /// Job to cancel. + job_id: String, + }, + /// Stop accepting work and exit once in-flight jobs finish. + Shutdown, +} + +/// A message an executor sends to the scheduler. +#[derive(Debug, Serialize, Deserialize)] +#[serde(tag = "type", rename_all = "snake_case")] +pub enum ExecutorMessage { + /// First frame on every connection: who is attaching and what it can run. + Hello { + /// Stable identity of this executor. + executor_id: String, + /// SDK the executor is built on, e.g. `"python"`. + sdk: String, + /// SDK version string, for logs and inventory. + version: String, + /// Tasks this executor has handlers for. Nothing else is sent to it. + tasks: Vec, + /// How many jobs it can run concurrently. + slots: u32, + /// Version the executor speaks. + protocol_version: u32, + }, + /// Liveness signal carrying current free capacity. + Heartbeat { + /// Slots free right now. + free_slots: u32, + }, + /// The task completed. Its serialized result, if any, follows the header. + Success { + /// Job that finished. + job_id: String, + /// Length of the result blob, or `None` when the task returned nothing. + /// `Some(0)` is an empty result, not a missing one. + result_len: Option, + /// Task that ran. + task_name: String, + /// Wall-clock execution time in nanoseconds. + wall_time_ns: i64, + }, + /// The task raised or timed out. + Failure { + /// Job that failed. + job_id: String, + /// Error message (canonical JSON `TaskError` when structured). + error: String, + /// Retries already attempted before this failure. + retry_count: i32, + /// The job's retry cap. + max_retries: i32, + /// Task that ran. + task_name: String, + /// Wall-clock execution time in nanoseconds. + wall_time_ns: i64, + /// Whether the failure is retryable. The executor decides — only it can + /// see the exception; the core never inspects one. + should_retry: bool, + /// True when the failure was an execution timeout. + timed_out: bool, + }, + /// The task observed its cancel request and stopped. + Cancelled { + /// Job that was cancelled. + job_id: String, + /// Task that ran. + task_name: String, + /// Wall-clock execution time in nanoseconds. + wall_time_ns: i64, + }, +} + +/// A frame header that may declare a trailing binary blob. Implemented by both +/// message enums so the reader and writer stay generic over direction. +pub trait Frame: Serialize + DeserializeOwned { + /// Bytes of payload that follow this header; zero for frames carrying none. + fn payload_len(&self) -> usize; +} + +impl Frame for SchedulerMessage { + fn payload_len(&self) -> usize { + match self { + Self::Job { payload_len, .. } => *payload_len, + _ => 0, + } + } +} + +impl Frame for ExecutorMessage { + fn payload_len(&self) -> usize { + match self { + Self::Success { result_len, .. } => result_len.unwrap_or(0), + _ => 0, + } + } +} + +impl From<&Job> for SchedulerMessage { + fn from(job: &Job) -> Self { + Self::Job { + id: job.id.clone(), + task_name: job.task_name.clone(), + payload_len: job.payload.len(), + retry_count: job.retry_count, + max_retries: job.max_retries, + queue: job.queue.clone(), + timeout_ms: job.timeout_ms, + namespace: job.namespace.clone(), + } + } +} + +impl ExecutorMessage { + /// Convert a result frame plus its payload into a [`JobResult`]. `None` for + /// non-result frames (`hello`, `heartbeat`). + pub fn into_job_result(self, payload: Vec) -> Option { + match self { + Self::Hello { .. } | Self::Heartbeat { .. } => None, + Self::Success { + job_id, + result_len, + task_name, + wall_time_ns, + } => Some(JobResult::Success { + job_id, + result: result_len.map(|_| payload), + task_name, + wall_time_ns, + }), + Self::Failure { + job_id, + error, + retry_count, + max_retries, + task_name, + wall_time_ns, + should_retry, + timed_out, + } => Some(JobResult::Failure { + job_id, + error, + retry_count, + max_retries, + task_name, + wall_time_ns, + should_retry, + timed_out, + }), + Self::Cancelled { + job_id, + task_name, + wall_time_ns, + } => Some(JobResult::Cancelled { + job_id, + task_name, + wall_time_ns, + }), + } + } +} + +/// Writes frames onto any byte sink — a child's stdin, a socket, a test buffer. +pub struct FrameWriter { + inner: BufWriter, +} + +impl FrameWriter { + /// Wrap a sink. Each frame is flushed as written, so a peer never waits on + /// a partially buffered header. + pub fn new(sink: W) -> Self { + Self { + inner: BufWriter::new(sink), + } + } + + /// Write one frame and its payload. A length disagreement would desync the + /// reader, so it is rejected before anything reaches the wire. + pub fn write(&mut self, frame: &F, payload: &[u8]) -> Result<(), ProtocolError> { + let declared = frame.payload_len(); + if declared != payload.len() { + return Err(ProtocolError::PayloadLengthMismatch { + declared, + actual: payload.len(), + }); + } + let header = serde_json::to_vec(frame)?; + self.inner.write_all(&header)?; + self.inner.write_all(b"\n")?; + self.inner.write_all(payload)?; + self.inner.flush()?; + Ok(()) + } + + /// Write a payload-free frame. + pub fn write_header(&mut self, frame: &F) -> Result<(), ProtocolError> { + self.write(frame, &[]) + } + + /// Dispatch a job, sending its payload as the frame's blob. + pub fn write_job(&mut self, job: &Job) -> Result<(), ProtocolError> { + self.write(&SchedulerMessage::from(job), &job.payload) + } + + /// Ask the peer to cancel a running job. + pub fn write_cancel(&mut self, job_id: &str) -> Result<(), ProtocolError> { + self.write_header(&SchedulerMessage::Cancel { + job_id: job_id.to_string(), + }) + } + + /// Ask the peer to drain and exit. + pub fn write_shutdown(&mut self) -> Result<(), ProtocolError> { + self.write_header(&SchedulerMessage::Shutdown) + } +} + +/// Reads frames from any buffered byte source. +pub struct FrameReader { + inner: R, +} + +impl FrameReader { + /// Wrap a buffered source. + pub fn new(source: R) -> Self { + Self { inner: source } + } + + /// Read one frame and its payload. Blocks until a whole frame arrives. + pub fn read(&mut self) -> Result<(F, Vec), ProtocolError> { + let header = self.read_header_line()?; + let frame: F = serde_json::from_slice(&header)?; + + let len = frame.payload_len(); + if len > MAX_PAYLOAD_BYTES { + return Err(ProtocolError::PayloadTooLarge { len }); + } + let mut payload = vec![0u8; len]; + if len > 0 { + self.inner.read_exact(&mut payload)?; + } + Ok((frame, payload)) + } + + /// Read through the header's newline, capped so a peer that never sends one + /// cannot grow the buffer without bound. + fn read_header_line(&mut self) -> Result, ProtocolError> { + let mut header = Vec::new(); + let read = (&mut self.inner) + .take(MAX_HEADER_BYTES) + .read_until(b'\n', &mut header)?; + if read == 0 { + return Err(ProtocolError::Eof); + } + if !header.ends_with(b"\n") { + return Err(ProtocolError::HeaderTooLarge); + } + Ok(header) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::job::JobStatus; + + fn round_trip(frame: &F, payload: &[u8]) -> (F, Vec) { + let mut buf = Vec::new(); + FrameWriter::new(&mut buf) + .write(frame, payload) + .expect("write"); + FrameReader::new(buf.as_slice()).read::().expect("read") + } + + fn sample_job(payload: &[u8]) -> Job { + Job { + id: "job-1".into(), + queue: "default".into(), + task_name: "resize".into(), + payload: payload.to_vec(), + status: JobStatus::Running, + priority: 0, + retry_count: 1, + max_retries: 3, + scheduled_at: 0, + created_at: 0, + started_at: None, + completed_at: None, + error: None, + result: None, + timeout_ms: 30_000, + unique_key: None, + progress: None, + metadata: None, + notes: None, + cancel_requested: false, + expires_at: None, + result_ttl_ms: None, + namespace: Some("tenant-a".into()), + has_deps: false, + } + } + + #[test] + fn job_frame_carries_payload_bytes_verbatim() { + // The CBOR envelope for f(1, "a") — the BINDING_CONTRACT test vector. + let payload = [0x02, 0x82, 0x82, 0x01, 0x61, 0x61, 0xa0]; + let job = sample_job(&payload); + + let mut buf = Vec::new(); + FrameWriter::new(&mut buf).write_job(&job).expect("write"); + assert!(buf.ends_with(&payload), "payload must be written raw"); + + let (frame, read_payload) = FrameReader::new(buf.as_slice()) + .read::() + .expect("read"); + assert_eq!(read_payload, payload); + match frame { + SchedulerMessage::Job { + id, + task_name, + payload_len, + retry_count, + max_retries, + queue, + timeout_ms, + namespace, + } => { + assert_eq!(id, "job-1"); + assert_eq!(task_name, "resize"); + assert_eq!(payload_len, payload.len()); + assert_eq!(retry_count, 1); + assert_eq!(max_retries, 3); + assert_eq!(queue, "default"); + assert_eq!(timeout_ms, 30_000); + assert_eq!(namespace.as_deref(), Some("tenant-a")); + } + other => panic!("expected a job frame, got {other:?}"), + } + } + + #[test] + fn handshake_frames_round_trip() { + let (hello, payload) = round_trip( + &ExecutorMessage::Hello { + executor_id: "exec-1".into(), + sdk: "python".into(), + version: "0.21.0".into(), + tasks: vec!["resize".into(), "thumbnail".into()], + slots: 4, + protocol_version: PROTOCOL_VERSION, + }, + &[], + ); + assert!(payload.is_empty()); + match hello { + ExecutorMessage::Hello { + tasks, + slots, + protocol_version, + .. + } => { + assert_eq!(tasks, ["resize", "thumbnail"]); + assert_eq!(slots, 4); + assert_eq!(protocol_version, PROTOCOL_VERSION); + } + other => panic!("expected hello, got {other:?}"), + } + + let (ack, _) = round_trip( + &SchedulerMessage::HelloAck { + scheduler_id: "scheduler-1".into(), + protocol_version: PROTOCOL_VERSION, + }, + &[], + ); + assert!(matches!(ack, SchedulerMessage::HelloAck { .. })); + } + + #[test] + fn control_frames_round_trip() { + let mut buf = Vec::new(); + let mut writer = FrameWriter::new(&mut buf); + writer.write_cancel("job-1").expect("cancel"); + writer.write_shutdown().expect("shutdown"); + drop(writer); + + let mut reader = FrameReader::new(buf.as_slice()); + assert!(matches!( + reader.read::().expect("read cancel").0, + SchedulerMessage::Cancel { job_id } if job_id == "job-1" + )); + assert!(matches!( + reader.read::().expect("read shutdown").0, + SchedulerMessage::Shutdown + )); + assert!(matches!( + reader.read::(), + Err(ProtocolError::Eof) + )); + } + + #[test] + fn empty_result_is_distinct_from_no_result() { + let empty = ExecutorMessage::Success { + job_id: "job-1".into(), + result_len: Some(0), + task_name: "t".into(), + wall_time_ns: 5, + }; + let (frame, payload) = round_trip(&empty, &[]); + match frame.into_job_result(payload) { + Some(JobResult::Success { result, .. }) => assert_eq!(result, Some(vec![])), + _ => panic!("expected a success result"), + } + + let none = ExecutorMessage::Success { + job_id: "job-1".into(), + result_len: None, + task_name: "t".into(), + wall_time_ns: 5, + }; + let (frame, payload) = round_trip(&none, &[]); + match frame.into_job_result(payload) { + Some(JobResult::Success { result, .. }) => assert_eq!(result, None), + _ => panic!("expected a success result"), + } + } + + #[test] + fn result_frames_map_onto_job_results() { + let (failure, payload) = round_trip( + &ExecutorMessage::Failure { + job_id: "job-1".into(), + error: r#"{"errtype":"ValueError","message":"boom","traceback":[]}"#.into(), + retry_count: 2, + max_retries: 3, + task_name: "t".into(), + wall_time_ns: 42, + should_retry: false, + timed_out: true, + }, + &[], + ); + match failure.into_job_result(payload) { + Some(JobResult::Failure { + should_retry, + timed_out, + retry_count, + .. + }) => { + assert!(!should_retry); + assert!(timed_out); + assert_eq!(retry_count, 2); + } + _ => panic!("expected a failure result"), + } + + let (cancelled, payload) = round_trip( + &ExecutorMessage::Cancelled { + job_id: "job-1".into(), + task_name: "t".into(), + wall_time_ns: 7, + }, + &[], + ); + assert!(matches!( + cancelled.into_job_result(payload), + Some(JobResult::Cancelled { .. }) + )); + } + + #[test] + fn non_result_frames_produce_no_job_result() { + let heartbeat = ExecutorMessage::Heartbeat { free_slots: 3 }; + assert!(heartbeat.into_job_result(vec![]).is_none()); + } + + #[test] + fn declared_length_must_match_the_payload() { + let job = sample_job(b"1234"); + let frame = SchedulerMessage::from(&job); + let mut buf = Vec::new(); + let err = FrameWriter::new(&mut buf) + .write(&frame, b"12") + .expect_err("length mismatch must be rejected"); + assert!(matches!( + err, + ProtocolError::PayloadLengthMismatch { + declared: 4, + actual: 2 + } + )); + assert!(buf.is_empty(), "nothing may reach the wire on a mismatch"); + } + + #[test] + fn oversized_header_is_rejected() { + let mut buf = vec![b'{'; MAX_HEADER_BYTES as usize + 1]; + buf.push(b'\n'); + assert!(matches!( + FrameReader::new(buf.as_slice()).read::(), + Err(ProtocolError::HeaderTooLarge) + )); + } + + #[test] + fn oversized_payload_is_rejected_before_allocating() { + let header = format!( + r#"{{"type":"job","id":"j","task_name":"t","payload_len":{},"retry_count":0,"max_retries":0,"queue":"q","timeout_ms":0,"namespace":null}}"#, + MAX_PAYLOAD_BYTES + 1 + ); + let mut buf = header.into_bytes(); + buf.push(b'\n'); + assert!(matches!( + FrameReader::new(buf.as_slice()).read::(), + Err(ProtocolError::PayloadTooLarge { .. }) + )); + } + + #[test] + fn truncated_payload_is_an_error_not_a_clean_eof() { + let job = sample_job(b"1234"); + let mut buf = Vec::new(); + FrameWriter::new(&mut buf).write_job(&job).expect("write"); + buf.truncate(buf.len() - 2); + + let err = FrameReader::new(buf.as_slice()) + .read::() + .expect_err("truncated frame must not read"); + assert!(matches!(err, ProtocolError::Io(_)), "got {err:?}"); + } + + #[test] + fn malformed_header_is_reported_as_json_error() { + let buf = b"not json\n"; + assert!(matches!( + FrameReader::new(&buf[..]).read::(), + Err(ProtocolError::Json(_)) + )); + } +} From 242ea629ea429195288cc0d8f2cb5ead77bf7c49 Mon Sep 17 00:00:00 2001 From: kartikeya-27 <67143288+kartikeya-27@users.noreply.github.com> Date: Fri, 31 Jul 2026 12:15:35 +0530 Subject: [PATCH 03/15] refactor(prefork): move onto the shared worker protocol Replaces the base64-in-JSON frames and the one-way ready signal with raw payload bytes and a versioned hello/hello_ack, so a mismatched TASKITO_PYTHON interpreter fails closed instead of silently. --- Cargo.lock | 1 - crates/taskito-python/Cargo.toml | 1 - crates/taskito-python/src/lib.rs | 5 + crates/taskito-python/src/prefork/child.rs | 120 +++++++--------- crates/taskito-python/src/prefork/mod.rs | 17 ++- crates/taskito-python/src/prefork/protocol.rs | 129 ------------------ .../guides/advanced-execution/prefork.mdx | 4 +- sdks/python/taskito/_taskito.pyi | 3 + sdks/python/taskito/prefork/child.py | 124 +++++++++++------ sdks/python/taskito/worker_protocol.py | 103 ++++++++++++++ .../tests/worker/prefork_apps/timeout_app.py | 13 ++ sdks/python/tests/worker/test_prefork.py | 25 ++++ 12 files changed, 288 insertions(+), 257 deletions(-) delete mode 100644 crates/taskito-python/src/prefork/protocol.rs create mode 100644 sdks/python/taskito/worker_protocol.py diff --git a/Cargo.lock b/Cargo.lock index 7f50fd250..7019d4fb0 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2579,7 +2579,6 @@ name = "taskito-python" version = "0.21.0" dependencies = [ "async-trait", - "base64 0.23.0", "crossbeam-channel", "gethostname", "log", diff --git a/crates/taskito-python/Cargo.toml b/crates/taskito-python/Cargo.toml index 75a66402c..c8ddc9377 100644 --- a/crates/taskito-python/Cargo.toml +++ b/crates/taskito-python/Cargo.toml @@ -34,7 +34,6 @@ taskito-workflows = { path = "../taskito-workflows", optional = true } taskito-mesh = { path = "../taskito-mesh", optional = true } serde_json = { workspace = true } serde = { workspace = true } -base64 = "0.23" log = { workspace = true } pyo3-log = "0.13" gethostname = "1.1.0" diff --git a/crates/taskito-python/src/lib.rs b/crates/taskito-python/src/lib.rs index 90d4f3d28..bdf4beb41 100644 --- a/crates/taskito-python/src/lib.rs +++ b/crates/taskito-python/src/lib.rs @@ -43,6 +43,11 @@ fn reserved_setting_prefixes() -> Vec { fn _taskito(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_function(wrap_pyfunction!(_init_rust_logging, m)?)?; m.add_function(wrap_pyfunction!(reserved_setting_prefixes, m)?)?; + // Sourced from the core so the executor side never mirrors the literal. + m.add( + "WORKER_PROTOCOL_VERSION", + taskito_core::worker::protocol::PROTOCOL_VERSION, + )?; m.add_class::()?; m.add_class::()?; m.add_class::()?; diff --git a/crates/taskito-python/src/prefork/child.rs b/crates/taskito-python/src/prefork/child.rs index a25332c0e..d4d5bff50 100644 --- a/crates/taskito-python/src/prefork/child.rs +++ b/crates/taskito-python/src/prefork/child.rs @@ -1,62 +1,29 @@ //! Child process handle — spawn, write jobs, read results. //! -//! A child is split into two halves after spawning: -//! - `ChildWriter`: sends jobs to the child's stdin (owned by dispatch thread) -//! - `ChildReader`: reads results from the child's stdout (owned by reader thread) +//! A child is split into three halves after spawning: +//! - `ChildWriter`: sends frames to the child's stdin (owned by dispatch thread) +//! - `ChildReader`: reads frames from the child's stdout (owned by reader thread) //! - `ChildProcess`: holds the process handle for lifecycle management +//! +//! The frames themselves are the shared worker protocol, so a pipe child and a +//! socket-attached executor speak the same wire format. -use std::io::{BufRead, BufReader, BufWriter, Write}; +use std::io::BufReader; use std::process::{Child, ChildStdin, ChildStdout, Command, Stdio}; -use super::protocol::{ChildMessage, ParentMessage}; +use taskito_core::worker::protocol::{ + ExecutorMessage, FrameReader, FrameWriter, SchedulerMessage, PROTOCOL_VERSION, +}; -/// Writer half — sends job messages to the child process via stdin. -pub struct ChildWriter { - writer: BufWriter, -} +/// Identity this pool announces in `hello_ack`. Informational — it only ever +/// reaches the child's logs. +const SCHEDULER_ID: &str = "prefork"; -impl ChildWriter { - /// Send a message to the child process. - pub fn send(&mut self, msg: &ParentMessage) -> std::io::Result<()> { - let json = serde_json::to_string(msg) - .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e.to_string()))?; - self.writer.write_all(json.as_bytes())?; - self.writer.write_all(b"\n")?; - self.writer.flush() - } +/// Writer half — sends frames to the child process via stdin. +pub type ChildWriter = FrameWriter; - /// Send a shutdown message. Errors are silently ignored (child may already be gone). - pub fn send_shutdown(&mut self) { - let _ = self.send(&ParentMessage::Shutdown); - } - - /// Send a cooperative-cancel request for `job_id`. Returns the underlying - /// I/O error if the pipe is broken so the caller can decide whether to - /// retry or drop the request. - pub fn send_cancel(&mut self, job_id: &str) -> std::io::Result<()> { - self.send(&ParentMessage::Cancel { - job_id: job_id.to_string(), - }) - } -} - -/// Reader half — reads result messages from the child process via stdout. -pub struct ChildReader { - reader: BufReader, -} - -impl ChildReader { - /// Read one message from the child's stdout. Blocks until a line is available. - pub fn read(&mut self) -> Result { - let mut line = String::new(); - match self.reader.read_line(&mut line) { - Ok(0) => Err("child process closed stdout".into()), - Ok(_) => serde_json::from_str(&line) - .map_err(|e| format!("failed to parse child message: {e}")), - Err(e) => Err(format!("failed to read from child stdout: {e}")), - } - } -} +/// Reader half — reads frames from the child process via stdout. +pub type ChildReader = FrameReader>; /// Process handle for lifecycle management. pub struct ChildProcess { @@ -99,9 +66,12 @@ impl ChildProcess { } } -/// Spawn a child worker process and wait for its `ready` signal. +/// Spawn a child worker process and complete the `hello`/`hello_ack` handshake. /// -/// Returns the three split halves: writer, reader, and process handle. +/// Returns the three split halves: writer, reader, and process handle. The ack +/// is sent even on a version mismatch so both sides can log both versions — +/// `TASKITO_PYTHON` lets the child run from a different interpreter, so a +/// mismatched taskito install is reachable in practice. pub fn spawn_child( python: &str, app_path: &str, @@ -117,26 +87,36 @@ pub fn spawn_child( let stdin = process.stdin.take().expect("stdin should be piped"); let stdout = process.stdout.take().expect("stdout should be piped"); - let mut reader = ChildReader { - reader: BufReader::new(stdout), + let mut reader = ChildReader::new(BufReader::new(stdout)); + let mut writer = ChildWriter::new(stdin); + + let hello = reader + .read::() + .map_err(|e| format!("child handshake failed: {e}"))? + .0; + let ExecutorMessage::Hello { + sdk, + version, + protocol_version, + .. + } = hello + else { + return Err("child sent a non-hello frame before the handshake completed".into()); }; - // Wait for ready signal - match reader.read()? { - ChildMessage::Ready => {} - other => { - return Err(format!( - "expected ready message, got: {:?}", - std::any::type_name_of_val(&other) - )); - } + writer + .write_header(&SchedulerMessage::HelloAck { + scheduler_id: SCHEDULER_ID.to_string(), + protocol_version: PROTOCOL_VERSION, + }) + .map_err(|e| format!("failed to acknowledge child handshake: {e}"))?; + + if protocol_version != PROTOCOL_VERSION { + return Err(format!( + "child speaks worker protocol {protocol_version}, we speak {PROTOCOL_VERSION} \ + (child is {sdk} {version}; check TASKITO_PYTHON points at the same install)" + )); } - Ok(( - ChildWriter { - writer: BufWriter::new(stdin), - }, - reader, - ChildProcess { process }, - )) + Ok((writer, reader, ChildProcess { process })) } diff --git a/crates/taskito-python/src/prefork/mod.rs b/crates/taskito-python/src/prefork/mod.rs index 61623b47f..6ef31e1b8 100644 --- a/crates/taskito-python/src/prefork/mod.rs +++ b/crates/taskito-python/src/prefork/mod.rs @@ -16,7 +16,6 @@ mod child; mod dispatch; -pub mod protocol; mod slot; mod watchdog; @@ -30,10 +29,10 @@ use crossbeam_channel::{Receiver, Sender, TrySendError}; use taskito_core::job::Job; use taskito_core::scheduler::JobResult; +use taskito_core::worker::protocol::ExecutorMessage; use taskito_core::worker::WorkerDispatcher; use child::{spawn_child, ChildProcess, ChildReader, ChildWriter}; -use protocol::ParentMessage; use slot::{ActiveJob, SlotState}; /// How long graceful shutdown will wait for each child to drain before @@ -180,7 +179,8 @@ impl WorkerDispatcher for PreforkPool { for idx in 0..num_workers { if let Ok(mut guard) = writers[idx].lock() { if let Some(w) = guard.as_mut() { - w.send_shutdown(); + // Best-effort: the child may already be gone. + let _ = w.write_shutdown(); log::info!("[taskito] sent shutdown to prefork child {idx}"); } } @@ -280,10 +280,9 @@ fn dispatch_job( }; slot::set(slots, idx, active); - let msg = ParentMessage::from(&job); let send_result = match writers[idx].lock() { Ok(mut guard) => match guard.as_mut() { - Some(writer) => writer.send(&msg), + Some(writer) => writer.write_job(&job), None => { drop(guard); let _ = slot::take(slots, idx); @@ -374,9 +373,9 @@ fn spawn_reader_thread( thread::Builder::new() .name(format!("taskito-prefork-reader-{idx}")) .spawn(move || loop { - match reader.read() { - Ok(msg) => { - let Some(job_result) = msg.into_job_result() else { + match reader.read::() { + Ok((msg, payload)) => { + let Some(job_result) = msg.into_job_result(payload) else { continue; }; if slot::take(&slots, idx).is_none() { @@ -425,7 +424,7 @@ fn spawn_cancel_router( let Some(writer) = guard.as_mut() else { continue; }; - if let Err(e) = writer.send_cancel(&job_id) { + if let Err(e) = writer.write_cancel(&job_id) { log::warn!( "[taskito] failed to forward cancel for {job_id} to child {idx}: {e}" ); diff --git a/crates/taskito-python/src/prefork/protocol.rs b/crates/taskito-python/src/prefork/protocol.rs deleted file mode 100644 index 3405a89e4..000000000 --- a/crates/taskito-python/src/prefork/protocol.rs +++ /dev/null @@ -1,129 +0,0 @@ -//! IPC message types for parent↔child communication. -//! -//! Uses JSON Lines (one JSON object per line) over stdio pipes. -//! The `payload` field is base64-encoded since it contains opaque bytes. - -use base64::Engine; -use serde::{Deserialize, Serialize}; - -use taskito_core::job::Job; -use taskito_core::scheduler::JobResult; - -/// Message sent from parent to child. -#[derive(Serialize)] -#[serde(tag = "type", rename_all = "snake_case")] -pub enum ParentMessage { - Job { - id: String, - task_name: String, - payload: String, // base64-encoded - retry_count: i32, - max_retries: i32, - queue: String, - timeout_ms: i64, - namespace: Option, - }, - /// Cooperative-cancel request for a job currently running (or queued in - /// the child's stdin buffer). The child marks `job_id` as cancelled so - /// `current_job.check_cancelled()` raises `TaskCancelledError`. - Cancel { - job_id: String, - }, - Shutdown, -} - -impl From<&Job> for ParentMessage { - fn from(job: &Job) -> Self { - Self::Job { - id: job.id.clone(), - task_name: job.task_name.clone(), - payload: base64::engine::general_purpose::STANDARD.encode(&job.payload), - retry_count: job.retry_count, - max_retries: job.max_retries, - queue: job.queue.clone(), - timeout_ms: job.timeout_ms, - namespace: job.namespace.clone(), - } - } -} - -/// Message sent from child to parent. -#[derive(Deserialize)] -#[serde(tag = "type", rename_all = "snake_case")] -pub enum ChildMessage { - Ready, - Success { - job_id: String, - result: Option, // base64-encoded - task_name: String, - wall_time_ns: i64, - }, - Failure { - job_id: String, - error: String, - retry_count: i32, - max_retries: i32, - task_name: String, - wall_time_ns: i64, - should_retry: bool, - timed_out: bool, - }, - Cancelled { - job_id: String, - task_name: String, - wall_time_ns: i64, - }, -} - -impl ChildMessage { - /// Convert a child message into a `JobResult` for the scheduler. - /// Returns `None` for non-result messages (e.g. `Ready`). - pub fn into_job_result(self) -> Option { - match self { - Self::Ready => None, - Self::Success { - job_id, - result, - task_name, - wall_time_ns, - } => { - let result_bytes = result - .and_then(|b64| base64::engine::general_purpose::STANDARD.decode(b64).ok()); - Some(JobResult::Success { - job_id, - result: result_bytes, - task_name, - wall_time_ns, - }) - } - Self::Failure { - job_id, - error, - retry_count, - max_retries, - task_name, - wall_time_ns, - should_retry, - timed_out, - } => Some(JobResult::Failure { - job_id, - error, - retry_count, - max_retries, - task_name, - wall_time_ns, - should_retry, - timed_out, - }), - Self::Cancelled { - job_id, - task_name, - wall_time_ns, - } => Some(JobResult::Cancelled { - job_id, - task_name, - wall_time_ns, - }), - } - } -} diff --git a/docs/content/docs/python/guides/advanced-execution/prefork.mdx b/docs/content/docs/python/guides/advanced-execution/prefork.mdx index a09067b12..b8356bafc 100644 --- a/docs/content/docs/python/guides/advanced-execution/prefork.mdx +++ b/docs/content/docs/python/guides/advanced-execution/prefork.mdx @@ -54,8 +54,8 @@ defaults to `thread`, while Celery defaults to `prefork`. /> 1. The Rust scheduler dequeues jobs from storage -2. `PreforkPool` serializes each job as JSON and writes it to the least-loaded child's stdin pipe -3. Each child deserializes the job, executes the task wrapper (with middleware, resources, proxies), and writes the result as JSON to stdout +2. `PreforkPool` frames each job as a JSON header line followed by the task payload's raw bytes, and writes it to the least-loaded child's stdin pipe +3. Each child decodes the frame, executes the task wrapper (with middleware, resources, proxies), and writes the result back as the same kind of frame on stdout 4. Reader threads parse results and feed them back to the scheduler 5. The scheduler updates job status in storage diff --git a/sdks/python/taskito/_taskito.pyi b/sdks/python/taskito/_taskito.pyi index 0a3a43918..43e861fa4 100644 --- a/sdks/python/taskito/_taskito.pyi +++ b/sdks/python/taskito/_taskito.pyi @@ -578,3 +578,6 @@ def _init_rust_logging() -> None: def reserved_setting_prefixes() -> list[str]: """Settings-key prefixes the dashboard's generic KV surface must hide.""" ... + +WORKER_PROTOCOL_VERSION: int +"""Frame-format version an executor must announce to attach.""" diff --git a/sdks/python/taskito/prefork/child.py b/sdks/python/taskito/prefork/child.py index 8fd28dd30..72e3d1bfe 100644 --- a/sdks/python/taskito/prefork/child.py +++ b/sdks/python/taskito/prefork/child.py @@ -3,21 +3,20 @@ Each child is an independent Python interpreter that: 1. Imports the app module and builds the task registry. 2. Initializes resources (if any). -3. Runs a stdin reader thread that demultiplexes ``job``, ``cancel``, and - ``shutdown`` messages from the parent. Jobs go on an internal queue; +3. Completes the ``hello``/``hello_ack`` handshake with the parent. +4. Runs a stdin reader thread that demultiplexes ``job``, ``cancel``, and + ``shutdown`` frames from the parent. Jobs go on an internal queue; cancels populate a local set that ``current_job.check_cancelled()`` reads via a registered hook. -4. Pulls jobs off the internal queue on the main thread, executes them, - and writes JSON results to stdout. +5. Pulls jobs off the internal queue on the main thread, executes them, + and writes result frames to stdout. Spawned by the Rust ``PreforkPool`` via ``python -m taskito.prefork ``. """ from __future__ import annotations -import base64 import importlib -import json import logging import os import queue as _queue_mod @@ -28,6 +27,7 @@ import traceback from typing import Any +from taskito import __version__ from taskito.async_support.helpers import run_maybe_async from taskito.context import ( _clear_context, @@ -39,12 +39,17 @@ from taskito.exceptions import TaskCancelledError from taskito.log_config import silence_asyncio_pipe_noise from taskito.task_errors import encode_task_error +from taskito.worker_protocol import ( + WORKER_PROTOCOL_VERSION, + ProtocolError, + read_frame, + write_frame, +) logger = logging.getLogger("taskito.prefork.child") -# Sentinel pushed onto the internal job queue when the parent requests -# shutdown so the main loop can terminate without polling. -_SHUTDOWN_SENTINEL: dict[str, Any] = {"__shutdown__": True} +# One job at a time per child: each is a whole interpreter. +_SLOTS = 1 def _import_queue(app_path: str) -> Any: @@ -57,10 +62,39 @@ def _import_queue(app_path: str) -> Any: return queue -def _write_message(msg: dict[str, Any]) -> None: - """Write a JSON message to stdout (one line, flushed).""" - sys.stdout.write(json.dumps(msg) + "\n") - sys.stdout.flush() +def _write_message(header: dict[str, Any], payload: bytes = b"") -> None: + """Write one frame to the parent.""" + write_frame(sys.stdout.buffer, header, payload) + + +def _handshake(queue: Any) -> None: + """Announce what this child can run and check the parent speaks our version. + + Runs before the stdin reader thread starts so the ack is not consumed by it. + """ + _write_message( + { + "type": "hello", + "executor_id": f"prefork-{os.getpid()}", + "sdk": "python", + "version": __version__, + "tasks": sorted(queue._task_registry), + "slots": _SLOTS, + "protocol_version": WORKER_PROTOCOL_VERSION, + } + ) + + ack, _ = read_frame(sys.stdin.buffer) + if ack.get("type") != "hello_ack": + raise ProtocolError(f"expected hello_ack, got {ack.get('type')!r}") + + theirs = ack.get("protocol_version") + if theirs != WORKER_PROTOCOL_VERSION: + raise ProtocolError( + f"worker protocol mismatch: parent speaks {theirs}, " + f"we speak {WORKER_PROTOCOL_VERSION} — check TASKITO_PYTHON points " + f"at the interpreter holding the same taskito install" + ) class _CancelSignal: @@ -91,11 +125,11 @@ def discard(self, job_id: str) -> None: def _execute_job( queue: Any, job: dict[str, Any], -) -> dict[str, Any]: - """Execute a single job and return the result message.""" + payload: bytes, +) -> tuple[dict[str, Any], bytes]: + """Execute a single job and return its result frame and result payload.""" task_name = job["task_name"] job_id = job["id"] - payload = base64.b64decode(job["payload"]) retry_count = job.get("retry_count", 0) max_retries = job.get("max_retries", 3) @@ -112,7 +146,7 @@ def _execute_job( "wall_time_ns": 0, "should_retry": False, "timed_out": False, - } + }, b"" _set_context(job_id, task_name, retry_count, job.get("queue", "default")) @@ -126,10 +160,10 @@ def _execute_job( return { "type": "success", "job_id": job_id, - "result": base64.b64encode(result_bytes).decode() if result_bytes else None, + "result_len": None if result_bytes is None else len(result_bytes), "task_name": task_name, "wall_time_ns": wall_time_ns, - } + }, result_bytes or b"" except TaskCancelledError: wall_time_ns = time.monotonic_ns() - start_ns @@ -138,7 +172,7 @@ def _execute_job( "job_id": job_id, "task_name": task_name, "wall_time_ns": wall_time_ns, - } + }, b"" except Exception: wall_time_ns = time.monotonic_ns() - start_ns @@ -170,53 +204,51 @@ def _execute_job( "wall_time_ns": wall_time_ns, "should_retry": should_retry, "timed_out": False, - } + }, b"" finally: _clear_context() def _spawn_stdin_reader( - job_queue: _queue_mod.Queue[dict[str, Any]], + job_queue: _queue_mod.Queue[tuple[dict[str, Any], bytes] | None], cancels: _CancelSignal, ) -> threading.Thread: - """Run a background thread that demultiplexes parent → child messages. + """Run a background thread that demultiplexes parent → child frames. The main thread is blocked inside ``_execute_job`` while a job is - running, so reading stdin must happen elsewhere. This thread converts - the line-delimited JSON stream into queue items + cancel-set updates. + running, so reading stdin must happen elsewhere. This thread turns the + frame stream into queue items + cancel-set updates. """ def reader() -> None: try: - for line in sys.stdin: - line = line.strip() - if not line: - continue + while True: try: - msg = json.loads(line) - except json.JSONDecodeError as e: - logger.warning("invalid IPC message from parent: %s", e) - continue + msg, payload = read_frame(sys.stdin.buffer) + except ProtocolError as e: + # A desynced stream cannot be resynchronised: the payload + # boundary is lost, so every later frame would be garbage. + logger.error("invalid frame from parent: %s", e) + return msg_type = msg.get("type") if msg_type == "shutdown": - job_queue.put(_SHUTDOWN_SENTINEL) return if msg_type == "job": - job_queue.put(msg) + job_queue.put((msg, payload)) elif msg_type == "cancel": job_id = msg.get("job_id") if isinstance(job_id, str): cancels.request(job_id) else: - logger.warning("unknown IPC message type: %r", msg_type) + logger.warning("unknown frame type from parent: %r", msg_type) except (BrokenPipeError, EOFError, KeyboardInterrupt): logger.debug("child stdin closed") finally: - # Ensure the main loop wakes up even if stdin closed without a - # shutdown message (e.g. the parent died). - job_queue.put(_SHUTDOWN_SENTINEL) + # Wake the main loop even if stdin closed without a shutdown + # frame (e.g. the parent died). + job_queue.put(None) thread = threading.Thread(target=reader, name="taskito-prefork-stdin", daemon=True) thread.start() @@ -264,21 +296,23 @@ def main() -> None: if runtime is not None: runtime.initialize() - job_queue: _queue_mod.Queue[dict[str, Any]] = _queue_mod.Queue() + _handshake(queue) + + job_queue: _queue_mod.Queue[tuple[dict[str, Any], bytes] | None] = _queue_mod.Queue() cancels = _CancelSignal() set_local_cancel_check(cancels.is_requested) _spawn_stdin_reader(job_queue, cancels) - _write_message({"type": "ready"}) logger.info("child ready (app=%s, pid=%d)", app_path, os.getpid()) try: while True: - msg = job_queue.get() - if msg is _SHUTDOWN_SENTINEL: + item = job_queue.get() + if item is None: break - result = _execute_job(queue, msg) - _write_message(result) + job, payload = item + result, result_payload = _execute_job(queue, job, payload) + _write_message(result, result_payload) # Drop the cancel marker once the result is written so a future # job with the same ID (extremely unlikely, but possible across # ID-reuse boundaries) does not auto-cancel. diff --git a/sdks/python/taskito/worker_protocol.py b/sdks/python/taskito/worker_protocol.py new file mode 100644 index 000000000..94c782011 --- /dev/null +++ b/sdks/python/taskito/worker_protocol.py @@ -0,0 +1,103 @@ +"""Frame codec for the worker protocol. + +A frame is a JSON header line followed by exactly the number of raw payload +bytes it declares. The blob stays raw rather than base64 inside the header so +the bytes on the wire are the wire-envelope bytes themselves. The same format +serves a pipe (prefork children) and a socket (an attached executor). + +Mirrors ``crates/taskito-core/src/worker/protocol.rs``; the version constant is +read from the native module so it is never restated here. +""" + +from __future__ import annotations + +import json +from typing import Any, BinaryIO + +from taskito._taskito import WORKER_PROTOCOL_VERSION + +__all__ = [ + "MAX_HEADER_BYTES", + "MAX_PAYLOAD_BYTES", + "WORKER_PROTOCOL_VERSION", + "ProtocolError", + "declared_payload_len", + "read_frame", + "write_frame", +] + +# Header cap, bounding a peer that never sends a newline. +MAX_HEADER_BYTES = 64 * 1024 + +# Payload cap, so a corrupt length field cannot allocate unboundedly. +MAX_PAYLOAD_BYTES = 64 * 1024 * 1024 + + +class ProtocolError(Exception): + """A frame could not be encoded or decoded.""" + + +def declared_payload_len(header: dict[str, Any]) -> int: + """Bytes of payload a header says follow it.""" + kind = header.get("type") + if kind == "job": + return int(header.get("payload_len") or 0) + if kind == "success": + result_len = header.get("result_len") + return 0 if result_len is None else int(result_len) + return 0 + + +def write_frame(stream: BinaryIO, header: dict[str, Any], payload: bytes = b"") -> None: + """Write one frame and flush it. + + A length disagreement would desync the reader, so it is rejected before + anything reaches the wire. + """ + declared = declared_payload_len(header) + if declared != len(payload): + raise ProtocolError( + f"frame declared {declared} payload bytes but {len(payload)} were supplied" + ) + stream.write(json.dumps(header, separators=(",", ":")).encode() + b"\n") + if payload: + stream.write(payload) + stream.flush() + + +def read_frame(stream: BinaryIO) -> tuple[dict[str, Any], bytes]: + """Read one frame. Raises ``EOFError`` when the peer closes between frames.""" + line = stream.readline(MAX_HEADER_BYTES + 1) + if not line: + raise EOFError("peer closed the connection") + if not line.endswith(b"\n"): + raise ProtocolError(f"frame header exceeds {MAX_HEADER_BYTES} bytes") + + try: + header = json.loads(line) + except json.JSONDecodeError as exc: + raise ProtocolError(f"malformed frame header: {exc}") from exc + if not isinstance(header, dict): + raise ProtocolError("frame header must be a JSON object") + + length = declared_payload_len(header) + if length > MAX_PAYLOAD_BYTES: + raise ProtocolError( + f"frame payload of {length} bytes exceeds the {MAX_PAYLOAD_BYTES} byte limit" + ) + return header, _read_exact(stream, length) + + +def _read_exact(stream: BinaryIO, length: int) -> bytes: + """Read exactly ``length`` bytes, looping because a raw stream may short-read.""" + if length == 0: + return b"" + chunks: list[bytes] = [] + remaining = length + while remaining: + chunk = stream.read(remaining) + if not chunk: + raise ProtocolError(f"truncated frame payload: wanted {length} bytes") + chunks.append(chunk) + remaining -= len(chunk) + return b"".join(chunks) diff --git a/sdks/python/tests/worker/prefork_apps/timeout_app.py b/sdks/python/tests/worker/prefork_apps/timeout_app.py index b546e3492..52060f381 100644 --- a/sdks/python/tests/worker/prefork_apps/timeout_app.py +++ b/sdks/python/tests/worker/prefork_apps/timeout_app.py @@ -33,6 +33,19 @@ def quick(x: int) -> int: return x * 2 +@queue.task() +def echo_bytes(blob: bytes) -> bytes: + """Echo a binary blob — frames carry payloads raw, so a payload containing + newlines must survive intact in both directions.""" + return blob + + +@queue.task() +def returns_nothing() -> None: + """Returns None — a result frame that declares no payload at all.""" + return None + + @queue.task(timeout=2, max_retries=0) def sleep_then_finish(seconds: float) -> str: """Sleeps for `seconds`, then finishes — used to verify the watchdog only diff --git a/sdks/python/tests/worker/test_prefork.py b/sdks/python/tests/worker/test_prefork.py index 3bbaa1e48..f25d05c8a 100644 --- a/sdks/python/tests/worker/test_prefork.py +++ b/sdks/python/tests/worker/test_prefork.py @@ -208,6 +208,31 @@ def test_prefork_no_timeout_unaffected(timeout_app: object) -> None: assert result == 42 +@prefork_unix_only +def test_prefork_round_trips_binary_payloads(timeout_app: object) -> None: + """Frames carry payloads as raw bytes after the header line, so a payload + containing newlines must not desync the stream in either direction.""" + queue: Queue = timeout_app.queue # type: ignore[attr-defined] + + blob = b'\n{"type":"success"}\n\x00\xff\n' + job = timeout_app.echo_bytes.delay(blob) # type: ignore[attr-defined] + _start_prefork_worker(queue) + + assert job.result(timeout=15) == blob + + +@prefork_unix_only +def test_prefork_task_returning_none_completes(timeout_app: object) -> None: + """A result frame that declares no payload still completes the job.""" + queue: Queue = timeout_app.queue # type: ignore[attr-defined] + + job = timeout_app.returns_nothing.delay() # type: ignore[attr-defined] + _start_prefork_worker(queue) + + assert _wait_for_terminal(job, timeout=15) == "complete" + assert job.result(timeout=5) is None + + @prefork_unix_only def test_prefork_finishes_before_deadline(timeout_app: object) -> None: """A task that completes well before its deadline returns normally — the From ad4d24375b085d5dd26e62b824abbe46735181f6 Mon Sep 17 00:00:00 2001 From: kartikeya-27 <67143288+kartikeya-27@users.noreply.github.com> Date: Fri, 31 Jul 2026 12:16:04 +0530 Subject: [PATCH 04/15] feat(core): let Worker run a custom dispatcher The registry reports the supplied pool_type so it never claims 'native' for a pool that isn't. --- crates/taskito-core/src/worker/runner.rs | 34 +++++++++++++++++++++--- 1 file changed, 30 insertions(+), 4 deletions(-) diff --git a/crates/taskito-core/src/worker/runner.rs b/crates/taskito-core/src/worker/runner.rs index 4021120ef..95384fd26 100644 --- a/crates/taskito-core/src/worker/runner.rs +++ b/crates/taskito-core/src/worker/runner.rs @@ -1,4 +1,4 @@ -//! Turn-key worker: wires a [`Scheduler`], a [`NativeDispatcher`], the result +//! Turn-key worker: wires a [`Scheduler`], a [`WorkerDispatcher`], the result //! drain loop, and the heartbeat/reap cadence into one `Worker::spawn()` call — //! the zero-to-executed-task path for a Rust consumer. //! @@ -42,6 +42,7 @@ pub struct Worker { queue_configs: Vec<(String, QueueConfig)>, worker_id: Option, on_outcome: Option, + dispatcher: Option<(String, Arc)>, } impl Worker { @@ -58,6 +59,7 @@ impl Worker { queue_configs: Vec::new(), worker_id: None, on_outcome: None, + dispatcher: None, } } @@ -110,6 +112,23 @@ impl Worker { self } + /// Run tasks on `dispatcher` instead of the built-in native pool — e.g. a + /// [`RemoteDispatcher`](super::RemoteDispatcher) feeding attached + /// executors. Registered handlers are then unused, and `num_workers` + /// should match the dispatcher's own concurrency so `max_in_flight` bounds + /// dispatch correctly. + /// + /// `pool_type` is what the worker registry reports, so it must describe the + /// pool that is actually running. + pub fn dispatcher( + mut self, + pool_type: impl Into, + dispatcher: Arc, + ) -> Self { + self.dispatcher = Some((pool_type.into(), dispatcher)); + self + } + /// Register a blocking handler. See [`TaskRegistry::register`]. pub fn register( mut self, @@ -145,10 +164,18 @@ impl Worker { queue_configs, worker_id, on_outcome, + dispatcher, } = self; let worker_id = worker_id.unwrap_or_else(|| format!("rust-worker-{}", uuid::Uuid::now_v7())); + let (pool_type, dispatcher): (String, Arc) = dispatcher + .unwrap_or_else(|| { + ( + "native".to_string(), + Arc::new(NativeDispatcher::new(registry, num_workers)), + ) + }); storage.register_worker( &worker_id, @@ -159,7 +186,7 @@ impl Worker { num_workers as i32, None, Some(std::process::id() as i32), - Some("native"), + Some(&pool_type), )?; // Bound dispatch to the pool size so this scheduler never claims more @@ -182,7 +209,6 @@ impl Worker { let (job_tx, job_rx) = tokio::sync::mpsc::channel(num_workers * 2); let (result_tx, result_rx) = crossbeam_channel::bounded(num_workers * 2); - let dispatcher = Arc::new(NativeDispatcher::new(registry, num_workers)); let runtime_done = Arc::new(AtomicBool::new(false)); // Runtime thread: scheduler dispatch + task execution. `result_tx` @@ -300,7 +326,7 @@ pub struct WorkerHandle { worker_id: String, storage: StorageBackend, shutdown: Arc, - dispatcher: Arc, + dispatcher: Arc, stop_tx: Option>, threads: Vec>, } From 9e91a932cc6725295337c5a40bb2537432822416 Mon Sep 17 00:00:00 2001 From: kartikeya-27 <67143288+kartikeya-27@users.noreply.github.com> Date: Fri, 31 Jul 2026 12:16:21 +0530 Subject: [PATCH 05/15] feat(core): add a Transport abstraction for executors split() hands back the read/write halves plus a Connection control, because the read timeout and the close both have to outlive the split. --- crates/taskito-core/src/lib.rs | 2 +- crates/taskito-core/src/worker/mod.rs | 4 + crates/taskito-core/src/worker/transport.rs | 379 ++++++++++++++++++++ 3 files changed, 384 insertions(+), 1 deletion(-) create mode 100644 crates/taskito-core/src/worker/transport.rs diff --git a/crates/taskito-core/src/lib.rs b/crates/taskito-core/src/lib.rs index fff2906db..9e1ff1801 100644 --- a/crates/taskito-core/src/lib.rs +++ b/crates/taskito-core/src/lib.rs @@ -47,5 +47,5 @@ pub use storage::StorageBackend; pub use storage::{DeadJob, QueueStats, SubscriptionBacklogStats}; pub use worker::{ ExecutorMessage, NativeDispatcher, ProtocolError, SchedulerMessage, TaskError, TaskRegistry, - TaskResult, Worker, WorkerDispatcher, WorkerHandle, PROTOCOL_VERSION, + TaskResult, Transport, Worker, WorkerDispatcher, WorkerHandle, PROTOCOL_VERSION, }; diff --git a/crates/taskito-core/src/worker/mod.rs b/crates/taskito-core/src/worker/mod.rs index 2f028cee0..e00e73f4b 100644 --- a/crates/taskito-core/src/worker/mod.rs +++ b/crates/taskito-core/src/worker/mod.rs @@ -2,11 +2,15 @@ pub mod dispatcher; pub mod protocol; pub mod registry; pub mod runner; +pub mod transport; pub use dispatcher::NativeDispatcher; pub use protocol::{ExecutorMessage, ProtocolError, SchedulerMessage, PROTOCOL_VERSION}; pub use registry::{TaskError, TaskHandler, TaskRegistry, TaskResult}; pub use runner::{Worker, WorkerHandle}; +#[cfg(unix)] +pub use transport::UnixTransport; +pub use transport::{MemoryTransport, TcpTransport, Transport}; use async_trait::async_trait; use crossbeam_channel::Sender; diff --git a/crates/taskito-core/src/worker/transport.rs b/crates/taskito-core/src/worker/transport.rs new file mode 100644 index 000000000..3931e4921 --- /dev/null +++ b/crates/taskito-core/src/worker/transport.rs @@ -0,0 +1,379 @@ +//! Byte streams an executor can attach over. +//! +//! The frame protocol only needs `Read + Write`, so a transport's whole job is +//! to hand back owned halves — a reader thread and the dispatch thread use the +//! connection concurrently — plus a read timeout so a silent peer cannot pin +//! the handshake, and a peer label for logs. + +use std::collections::VecDeque; +use std::io::{self, BufRead, BufReader, Read, Write}; +use std::net::TcpStream; +#[cfg(unix)] +use std::os::unix::net::UnixStream; +use std::sync::{Arc, Condvar, Mutex}; +use std::time::Duration; + +/// Owned read half of a split transport. +pub type ReadHalf = Box; +/// Owned write half of a split transport. +pub type WriteHalf = Box; + +/// Controls a connection whose halves have already been split. +/// +/// Both knobs have to outlive the split. The read timeout bounds the handshake +/// so a silent peer cannot pin an attach, then clears so an executor idling +/// between jobs is not dropped. [`Connection::close`] is what makes a blocked +/// reader return, so shutdown cannot hang on a peer that stops responding. +pub struct Connection { + set_read_timeout: Box) -> io::Result<()> + Send + Sync>, + close: Box, +} + +impl Connection { + /// Bound how long a read may block. `None` blocks indefinitely. + /// + /// A timed-out read reports [`io::ErrorKind::WouldBlock`], matching the + /// platform socket behaviour the socket transports inherit. + pub fn set_read_timeout(&self, timeout: Option) -> io::Result<()> { + (self.set_read_timeout)(timeout) + } + + /// Tear the connection down so a reader blocked on it returns. Idempotent + /// and best-effort: the peer may already be gone. + pub fn close(&self) { + (self.close)(); + } +} + +impl std::fmt::Debug for Connection { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str("Connection") + } +} + +/// A bidirectional stream carrying the worker frame protocol. +pub trait Transport: Send { + /// Split into owned halves — so a reader thread and the dispatch thread can + /// use the connection at the same time — plus its lifetime controls. + fn split(self: Box) -> io::Result<(ReadHalf, WriteHalf, Connection)>; + + /// Peer label for logs. Never carries credentials. + fn peer(&self) -> String; +} + +/// Attach over a Unix domain socket — the same-pod sidecar case. +#[cfg(unix)] +pub struct UnixTransport(UnixStream); + +#[cfg(unix)] +impl UnixTransport { + /// Wrap an accepted or connected stream. + pub fn new(stream: UnixStream) -> Self { + Self(stream) + } +} + +#[cfg(unix)] +impl Transport for UnixTransport { + fn split(self: Box) -> io::Result<(ReadHalf, WriteHalf, Connection)> { + // Every clone shares one fd, so the control reaches the read half. + let read = self.0.try_clone()?; + let control = Arc::new(self.0.try_clone()?); + let closer = control.clone(); + Ok(( + Box::new(BufReader::new(read)), + Box::new(self.0), + Connection { + set_read_timeout: Box::new(move |timeout| control.set_read_timeout(timeout)), + close: Box::new(move || { + let _ = closer.shutdown(std::net::Shutdown::Both); + }), + }, + )) + } + + fn peer(&self) -> String { + match self.0.peer_addr().ok().and_then(|a| { + a.as_pathname() + .map(|p| p.to_string_lossy().into_owned()) + .filter(|p| !p.is_empty()) + }) { + Some(path) => format!("unix:{path}"), + None => "unix:unnamed".to_string(), + } + } +} + +/// Attach over TCP — an executor in another pod or host. +pub struct TcpTransport(TcpStream); + +impl TcpTransport { + /// Wrap an accepted or connected stream. Nagle is disabled: frames are + /// small and latency-sensitive, and a job dispatch must not wait on + /// coalescing. + pub fn new(stream: TcpStream) -> io::Result { + stream.set_nodelay(true)?; + Ok(Self(stream)) + } +} + +impl Transport for TcpTransport { + fn split(self: Box) -> io::Result<(ReadHalf, WriteHalf, Connection)> { + let read = self.0.try_clone()?; + let control = Arc::new(self.0.try_clone()?); + let closer = control.clone(); + Ok(( + Box::new(BufReader::new(read)), + Box::new(self.0), + Connection { + set_read_timeout: Box::new(move |timeout| control.set_read_timeout(timeout)), + close: Box::new(move || { + let _ = closer.shutdown(std::net::Shutdown::Both); + }), + }, + )) + } + + fn peer(&self) -> String { + match self.0.peer_addr() { + Ok(addr) => format!("tcp:{addr}"), + Err(_) => "tcp:unknown".to_string(), + } + } +} + +/// In-process transport pair, for tests and for embedding an executor in the +/// scheduler process without a socket. +pub struct MemoryTransport { + incoming: Arc, + outgoing: Arc, + label: String, +} + +impl MemoryTransport { + /// Build two ends wired to each other. Dropping one end's write half + /// signals EOF to the other's reader. + pub fn pair() -> (Self, Self) { + let left = Arc::new(Channel::default()); + let right = Arc::new(Channel::default()); + ( + Self { + incoming: left.clone(), + outgoing: right.clone(), + label: "memory:a".to_string(), + }, + Self { + incoming: right, + outgoing: left, + label: "memory:b".to_string(), + }, + ) + } +} + +impl Transport for MemoryTransport { + fn split(self: Box) -> io::Result<(ReadHalf, WriteHalf, Connection)> { + self.outgoing.open_writer(); + let control = self.incoming.clone(); + let closer = self.incoming.clone(); + Ok(( + Box::new(BufReader::new(ChannelReader { + channel: self.incoming, + })), + Box::new(ChannelWriter { + channel: self.outgoing, + }), + Connection { + set_read_timeout: Box::new(move |timeout| { + *control.read_timeout.lock().unwrap_or_else(recover) = timeout; + Ok(()) + }), + close: Box::new(move || closer.close_reader()), + }, + )) + } + + fn peer(&self) -> String { + self.label.clone() + } +} + +/// Recover a guard from a poisoned lock rather than cascading the panic — the +/// state behind it is a plain buffer, so reading it is always safe. +fn recover(poisoned: std::sync::PoisonError) -> T { + poisoned.into_inner() +} + +/// One direction of a [`MemoryTransport`] pair. +#[derive(Default)] +struct Channel { + state: Mutex, + ready: Condvar, + read_timeout: Mutex>, +} + +#[derive(Default)] +struct ChannelState { + buffer: VecDeque, + writer_open: bool, + writer_ever_opened: bool, +} + +impl Channel { + fn open_writer(&self) { + let mut state = self.state.lock().unwrap_or_else(recover); + state.writer_open = true; + state.writer_ever_opened = true; + } + + fn close_writer(&self) { + let mut state = self.state.lock().unwrap_or_else(recover); + state.writer_open = false; + drop(state); + self.ready.notify_all(); + } + + /// Force this direction to EOF, the in-memory analogue of a socket + /// shutdown, so a blocked reader returns. + fn close_reader(&self) { + let mut state = self.state.lock().unwrap_or_else(recover); + state.writer_open = false; + state.writer_ever_opened = true; + state.buffer.clear(); + drop(state); + self.ready.notify_all(); + } +} + +struct ChannelReader { + channel: Arc, +} + +impl Read for ChannelReader { + fn read(&mut self, out: &mut [u8]) -> io::Result { + if out.is_empty() { + return Ok(0); + } + let timeout = *self.channel.read_timeout.lock().unwrap_or_else(recover); + let mut state = self.channel.state.lock().unwrap_or_else(recover); + + while state.buffer.is_empty() { + // Before the peer splits, `writer_open` is false but no data can + // have been lost — treating that as EOF would race the handshake. + if state.writer_ever_opened && !state.writer_open { + return Ok(0); + } + state = match timeout { + None => self.channel.ready.wait(state).unwrap_or_else(recover), + Some(limit) => { + let (guard, result) = self + .channel + .ready + .wait_timeout(state, limit) + .unwrap_or_else(recover); + if result.timed_out() && guard.buffer.is_empty() { + return Err(io::Error::new(io::ErrorKind::WouldBlock, "read timed out")); + } + guard + } + }; + } + + let count = out.len().min(state.buffer.len()); + for (slot, byte) in out.iter_mut().zip(state.buffer.drain(..count)) { + *slot = byte; + } + Ok(count) + } +} + +struct ChannelWriter { + channel: Arc, +} + +impl Write for ChannelWriter { + fn write(&mut self, data: &[u8]) -> io::Result { + let mut state = self.channel.state.lock().unwrap_or_else(recover); + state.buffer.extend(data); + drop(state); + self.channel.ready.notify_all(); + Ok(data.len()) + } + + fn flush(&mut self) -> io::Result<()> { + Ok(()) + } +} + +impl Drop for ChannelWriter { + fn drop(&mut self) { + self.channel.close_writer(); + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn memory_pair_carries_bytes_both_ways() { + let (a, b) = MemoryTransport::pair(); + let (mut a_read, mut a_write, _) = Box::new(a).split().expect("split a"); + let (mut b_read, mut b_write, _) = Box::new(b).split().expect("split b"); + + a_write.write_all(b"ping\n").expect("write"); + a_write.flush().expect("flush"); + let mut line = String::new(); + b_read.read_line(&mut line).expect("read"); + assert_eq!(line, "ping\n"); + + b_write.write_all(b"pong\n").expect("write"); + b_write.flush().expect("flush"); + line.clear(); + a_read.read_line(&mut line).expect("read"); + assert_eq!(line, "pong\n"); + } + + #[test] + fn dropping_the_writer_signals_eof() { + let (a, b) = MemoryTransport::pair(); + let (_a_read, a_write, _) = Box::new(a).split().expect("split a"); + let (mut b_read, _b_write, _) = Box::new(b).split().expect("split b"); + + drop(a_write); + let mut buf = Vec::new(); + assert_eq!(b_read.read_to_end(&mut buf).expect("read"), 0); + } + + #[test] + fn read_timeout_applies_and_clears_after_the_split() { + let (a, b) = MemoryTransport::pair(); + let (_a_read, mut a_write, _) = Box::new(a).split().expect("split a"); + let (mut b_read, _b_write, timeout) = Box::new(b).split().expect("split b"); + + timeout + .set_read_timeout(Some(Duration::from_millis(20))) + .expect("set timeout"); + let mut byte = [0u8; 1]; + let err = b_read.read(&mut byte).expect_err("must time out"); + assert_eq!(err.kind(), io::ErrorKind::WouldBlock); + + // Clearing must reach the already-split read half — otherwise an idle + // executor would be dropped the moment its handshake budget elapsed. + timeout.set_read_timeout(None).expect("clear timeout"); + let reader = std::thread::spawn(move || { + let mut buf = [0u8; 4]; + b_read.read_exact(&mut buf).expect("blocking read"); + buf + }); + std::thread::sleep(Duration::from_millis(60)); + assert!( + !reader.is_finished(), + "a cleared timeout must block, not expire" + ); + + a_write.write_all(b"ping").expect("write"); + a_write.flush().expect("flush"); + assert_eq!(&reader.join().expect("reader thread")[..], b"ping"); + } +} From 4e1882c5fb42d02be8a8c9a510ef1f47233355e7 Mon Sep 17 00:00:00 2001 From: kartikeya-27 <67143288+kartikeya-27@users.noreply.github.com> Date: Fri, 31 Jul 2026 12:16:56 +0530 Subject: [PATCH 06/15] feat(core): add RemoteDispatcher for attached executors Dispatches only task names an executor advertised in hello; a dropped executor leaves its in-flight jobs to the scheduler's reaper. Closes #548. --- crates/taskito-core/src/lib.rs | 5 +- crates/taskito-core/src/worker/mod.rs | 4 + crates/taskito-core/src/worker/remote.rs | 735 ++++++++++++++++++ .../taskito-core/src/worker/remote_tests.rs | 565 ++++++++++++++ 4 files changed, 1307 insertions(+), 2 deletions(-) create mode 100644 crates/taskito-core/src/worker/remote.rs create mode 100644 crates/taskito-core/src/worker/remote_tests.rs diff --git a/crates/taskito-core/src/lib.rs b/crates/taskito-core/src/lib.rs index 9e1ff1801..ff5fe95c4 100644 --- a/crates/taskito-core/src/lib.rs +++ b/crates/taskito-core/src/lib.rs @@ -46,6 +46,7 @@ pub use storage::Storage; pub use storage::StorageBackend; pub use storage::{DeadJob, QueueStats, SubscriptionBacklogStats}; pub use worker::{ - ExecutorMessage, NativeDispatcher, ProtocolError, SchedulerMessage, TaskError, TaskRegistry, - TaskResult, Transport, Worker, WorkerDispatcher, WorkerHandle, PROTOCOL_VERSION, + AttachError, AttachedExecutor, ExecutorMessage, NativeDispatcher, ProtocolError, RemoteConfig, + RemoteDispatcher, SchedulerMessage, TaskError, TaskRegistry, TaskResult, Transport, Worker, + WorkerDispatcher, WorkerHandle, PROTOCOL_VERSION, }; diff --git a/crates/taskito-core/src/worker/mod.rs b/crates/taskito-core/src/worker/mod.rs index e00e73f4b..fd8fcd043 100644 --- a/crates/taskito-core/src/worker/mod.rs +++ b/crates/taskito-core/src/worker/mod.rs @@ -1,12 +1,14 @@ pub mod dispatcher; pub mod protocol; pub mod registry; +pub mod remote; pub mod runner; pub mod transport; pub use dispatcher::NativeDispatcher; pub use protocol::{ExecutorMessage, ProtocolError, SchedulerMessage, PROTOCOL_VERSION}; pub use registry::{TaskError, TaskHandler, TaskRegistry, TaskResult}; +pub use remote::{AttachError, AttachedExecutor, Capacity, RemoteConfig, RemoteDispatcher}; pub use runner::{Worker, WorkerHandle}; #[cfg(unix)] pub use transport::UnixTransport; @@ -39,5 +41,7 @@ pub trait WorkerDispatcher: Send + Sync { fn notify_cancel(&self, _job_id: &str) {} } +#[cfg(test)] +mod remote_tests; #[cfg(test)] mod tests; diff --git a/crates/taskito-core/src/worker/remote.rs b/crates/taskito-core/src/worker/remote.rs new file mode 100644 index 000000000..de7fa015a --- /dev/null +++ b/crates/taskito-core/src/worker/remote.rs @@ -0,0 +1,735 @@ +//! Dispatch to executors attached over a [`Transport`]. +//! +//! An executor dials in, announces the tasks it can run, and receives jobs for +//! those tasks only. The scheduler is untouched: this is a [`WorkerDispatcher`] +//! like any other, so the same claim, retry, and reaper machinery applies. + +use std::collections::{HashMap, HashSet}; +use std::sync::atomic::{AtomicBool, AtomicU32, Ordering}; +use std::sync::{Arc, Mutex, PoisonError}; +use std::thread::{self, JoinHandle}; +use std::time::{Duration, Instant}; + +use async_trait::async_trait; +use crossbeam_channel::{Receiver, Sender, TrySendError}; +use tokio::sync::Notify; + +use super::protocol::{ + ExecutorMessage, FrameReader, FrameWriter, ProtocolError, SchedulerMessage, PROTOCOL_VERSION, +}; +use super::transport::{Connection, ReadHalf, Transport, WriteHalf}; +use super::WorkerDispatcher; +use crate::job::Job; +use crate::scheduler::JobResult; + +/// Tuning for a [`RemoteDispatcher`]. +#[derive(Debug, Clone)] +pub struct RemoteConfig { + /// Identity announced in `hello_ack`. + pub scheduler_id: String, + /// How long the handshake may take before the connection is dropped, so a + /// peer that connects and says nothing cannot pin an attach. + pub handshake_timeout: Duration, + /// How long a job waits for a slot on an executor advertising its task + /// before it is failed back retryably. + pub placement_timeout: Duration, + /// How long shutdown waits for attached executors to finish in-flight + /// jobs before their connections are closed. + pub shutdown_drain: Duration, + /// Capacity of the cancel side-channel. + pub cancel_capacity: usize, +} + +impl Default for RemoteConfig { + fn default() -> Self { + Self { + scheduler_id: format!("scheduler-{}", uuid::Uuid::now_v7()), + handshake_timeout: Duration::from_secs(10), + placement_timeout: Duration::from_secs(30), + shutdown_drain: Duration::from_secs(30), + cancel_capacity: 1024, + } + } +} + +/// Why an executor could not attach. +#[derive(Debug, thiserror::Error)] +pub enum AttachError { + /// The transport could not be split or configured. + #[error("attach transport failed: {0}")] + Transport(#[from] std::io::Error), + + /// The handshake was malformed, timed out, or announced a version we do + /// not speak ([`ProtocolError::VersionMismatch`]). + #[error(transparent)] + Protocol(#[from] ProtocolError), + + /// Another executor is already attached under this id. + #[error("executor {0} is already attached")] + DuplicateId(String), +} + +/// A snapshot of one attached executor. +#[derive(Debug, Clone)] +pub struct AttachedExecutor { + /// Identity the executor announced. + pub executor_id: String, + /// SDK it is built on. + pub sdk: String, + /// SDK version string. + pub version: String, + /// Tasks it advertised, sorted. + pub tasks: Vec, + /// Concurrency it advertised. + pub slots: u32, + /// Slots free right now. + pub free_slots: u32, + /// Jobs currently running on it. + pub in_flight: usize, + /// Peer label, for logs. + pub peer: String, + /// Milliseconds since its last frame. + pub idle_ms: u32, +} + +/// Total advertised capacity across attached executors. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub struct Capacity { + /// Executors attached. + pub executors: usize, + /// Slots advertised in total. + pub total_slots: u32, + /// Slots free right now. + pub free_slots: u32, +} + +/// How often the shutdown drain re-checks whether executors have finished. +const DRAIN_POLL: Duration = Duration::from_millis(50); + +/// Recover a guard from a poisoned lock instead of cascading the panic. The +/// state behind these locks is plain bookkeeping, so reading it stays safe. +fn recover(poisoned: PoisonError) -> T { + poisoned.into_inner() +} + +/// A [`WorkerDispatcher`] that runs jobs on executors attached over a socket. +/// +/// Executors dial out, so the app needs no inbound port. Binding and accepting +/// belong to the caller: hand each accepted connection to [`Self::attach`]. +/// Cloning is cheap and shares one registry, so a listener thread can hold a +/// handle while the scheduler owns the dispatcher. +#[derive(Clone)] +pub struct RemoteDispatcher { + shared: Arc, +} + +impl RemoteDispatcher { + /// Build a dispatcher with no executors attached yet. + pub fn new(config: RemoteConfig) -> Self { + Self { + shared: Arc::new(Shared { + config, + executors: Mutex::new(HashMap::new()), + capacity_changed: Notify::new(), + result_tx: Mutex::new(None), + cancel_tx: Mutex::new(None), + readers: Mutex::new(Vec::new()), + shutdown: AtomicBool::new(false), + started_at: Instant::now(), + }), + } + } + + /// Complete the handshake on `transport` and register the executor. + pub fn attach(&self, transport: Box) -> Result { + self.shared.attach(transport) + } + + /// Snapshot every attached executor. + pub fn executors(&self) -> Vec { + self.shared.snapshot() + } + + /// Advertised capacity across all attached executors. + /// + /// A scheduler sizes `SchedulerConfig::max_in_flight` from this rather than + /// running a second, parallel limiter. + pub fn capacity(&self) -> Capacity { + self.shared.capacity() + } +} + +#[async_trait] +impl WorkerDispatcher for RemoteDispatcher { + async fn run(&self, job_rx: tokio::sync::mpsc::Receiver, result_tx: Sender) { + self.shared.clone().run(job_rx, result_tx).await; + } + + fn shutdown(&self) { + self.shared.shutdown.store(true, Ordering::SeqCst); + } + + fn notify_cancel(&self, job_id: &str) { + self.shared.notify_cancel(job_id); + } +} + +/// One attached executor: what it can run, what it is running, and how to +/// reach it. +struct Executor { + id: String, + sdk: String, + version: String, + tasks: HashSet, + slots: u32, + free: AtomicU32, + /// Job id → task name. Taking an entry is the exactly-once token for + /// emitting that job's single `JobResult`. + in_flight: Mutex>, + writer: Mutex>, + connection: Connection, + peer: String, + /// Milliseconds since the dispatcher started, at the last frame read. + last_seen_ms: AtomicU32, +} + +impl Executor { + fn is_busy(executor: &Arc) -> bool { + !executor.in_flight.lock().unwrap_or_else(recover).is_empty() + } + + fn snapshot(&self, now_ms: u32) -> AttachedExecutor { + let mut tasks: Vec = self.tasks.iter().cloned().collect(); + tasks.sort(); + AttachedExecutor { + executor_id: self.id.clone(), + sdk: self.sdk.clone(), + version: self.version.clone(), + tasks, + slots: self.slots, + free_slots: self.free.load(Ordering::Relaxed), + in_flight: self.in_flight.lock().unwrap_or_else(recover).len(), + peer: self.peer.clone(), + idle_ms: now_ms.saturating_sub(self.last_seen_ms.load(Ordering::Relaxed)), + } + } +} + +/// Where a job can go right now. +enum Placement { + /// An executor advertising the task has a free slot, now reserved. + Ready(Arc), + /// Some executor advertises the task but all of them are busy. + Saturated, + /// No attached executor advertises the task at all. + Unadvertised, +} + +/// State shared by the dispatcher handle, its reader threads, and its router. +struct Shared { + config: RemoteConfig, + executors: Mutex>>, + /// Woken when an executor attaches, frees a slot, or detaches — the signal + /// a job waiting for placement is parked on. + capacity_changed: Notify, + /// Installed by `run`. Reader threads may start before it exists, but no + /// job can be dispatched until it does. + result_tx: Mutex>>, + cancel_tx: Mutex>>, + readers: Mutex>>, + shutdown: AtomicBool, + started_at: Instant, +} + +impl Shared { + /// Handshake and register. The ack is sent even when the version is + /// rejected, so both ends log both versions instead of one side guessing. + fn attach(self: &Arc, transport: Box) -> Result { + let peer = transport.peer(); + let (read, write, connection) = transport.split()?; + let mut reader = FrameReader::new(read); + let mut writer = FrameWriter::new(write); + + // Bound the handshake only. An attached executor waiting between jobs + // must block indefinitely, or it would be dropped every time it idled + // past this budget. + connection.set_read_timeout(Some(self.config.handshake_timeout))?; + let hello = reader.read::()?.0; + connection.set_read_timeout(None)?; + let ExecutorMessage::Hello { + executor_id, + sdk, + version, + tasks, + slots, + protocol_version, + } = hello + else { + return Err(ProtocolError::UnexpectedFrame { expected: "hello" }.into()); + }; + + writer.write_header(&SchedulerMessage::HelloAck { + scheduler_id: self.config.scheduler_id.clone(), + protocol_version: PROTOCOL_VERSION, + })?; + + if protocol_version != PROTOCOL_VERSION { + log::warn!( + "[taskito] rejecting executor {executor_id} ({sdk} {version}, {peer}): \ + speaks worker protocol {protocol_version}, we speak {PROTOCOL_VERSION}" + ); + return Err(ProtocolError::VersionMismatch { + ours: PROTOCOL_VERSION, + theirs: protocol_version, + } + .into()); + } + + let executor = Arc::new(Executor { + id: executor_id.clone(), + sdk, + version, + tasks: tasks.into_iter().collect(), + slots, + free: AtomicU32::new(slots), + in_flight: Mutex::new(HashMap::new()), + writer: Mutex::new(writer), + connection, + peer: peer.clone(), + last_seen_ms: AtomicU32::new(self.elapsed_ms()), + }); + + { + let mut executors = self.executors.lock().unwrap_or_else(recover); + if executors.contains_key(&executor_id) { + return Err(AttachError::DuplicateId(executor_id)); + } + executors.insert(executor_id.clone(), executor.clone()); + } + + let handle = Arc::clone(self).spawn_reader(executor, reader); + self.readers.lock().unwrap_or_else(recover).push(handle); + self.capacity_changed.notify_waiters(); + + log::info!("[taskito] executor {executor_id} attached from {peer} with {slots} slot(s)"); + Ok(executor_id) + } + + fn snapshot(&self) -> Vec { + let now = self.elapsed_ms(); + let executors = self.executors.lock().unwrap_or_else(recover); + let mut snapshot: Vec = + executors.values().map(|e| e.snapshot(now)).collect(); + snapshot.sort_by(|a, b| a.executor_id.cmp(&b.executor_id)); + snapshot + } + + fn capacity(&self) -> Capacity { + let executors = self.executors.lock().unwrap_or_else(recover); + executors + .values() + .fold(Capacity::default(), |mut acc, executor| { + acc.executors += 1; + acc.total_slots += executor.slots; + acc.free_slots += executor.free.load(Ordering::Relaxed); + acc + }) + } + + /// Milliseconds since the dispatcher was created — a monotonic clock for + /// liveness that a wall-clock jump cannot move backwards. + fn elapsed_ms(&self) -> u32 { + self.started_at + .elapsed() + .as_millis() + .min(u128::from(u32::MAX)) as u32 + } + + async fn run( + self: Arc, + mut job_rx: tokio::sync::mpsc::Receiver, + result_tx: Sender, + ) { + *self.result_tx.lock().unwrap_or_else(recover) = Some(result_tx); + + let (cancel_tx, cancel_rx) = crossbeam_channel::bounded(self.config.cancel_capacity); + self.set_cancel_sender(Some(cancel_tx)); + let cancel_router = Arc::clone(&self).spawn_cancel_router(cancel_rx); + + while let Some(job) = job_rx.recv().await { + if self.shutdown.load(Ordering::Relaxed) { + break; + } + self.place(job).await; + } + + // Stop accepting cancels so the router drains and exits while the + // writers it uses are still alive. + self.set_cancel_sender(None); + self.drain_and_close(); + + let readers = std::mem::take(&mut *self.readers.lock().unwrap_or_else(recover)); + for handle in readers { + let _ = handle.join(); + } + let _ = cancel_router.join(); + } + + /// Place one job, waiting for a slot if every advertising executor is busy. + /// + /// A job for a task nobody advertises, or one that waits out + /// `placement_timeout`, comes back as a retryable failure: it reschedules + /// under the normal retry policy and surfaces the misconfiguration rather + /// than hiding it. One unplaceable job blocks the ones behind it, since the + /// scheduler hands over a single job stream — the mitigation is per-task + /// `max_concurrent`, which gates before the job is ever dequeued. + async fn place(&self, job: Job) { + let deadline = Instant::now() + self.config.placement_timeout; + loop { + // Register the waiter *before* checking capacity: `notify_waiters` + // only wakes waiters already registered, so subscribing lazily + // would lose a slot freed between the check and the await. + let mut changed = std::pin::pin!(self.capacity_changed.notified()); + changed.as_mut().enable(); + + let reason = match self.try_acquire(&job.task_name) { + Placement::Ready(executor) => { + self.dispatch_to(&executor, job); + return; + } + Placement::Saturated => "every executor advertising it is busy", + Placement::Unadvertised => "no attached executor advertises it", + }; + + let remaining = deadline.saturating_duration_since(Instant::now()); + if remaining.is_zero() || tokio::time::timeout(remaining, changed).await.is_err() { + self.fail_unplaceable(&job, reason); + return; + } + } + } + + /// Pick the executor with the most free slots that advertises `task_name`, + /// reserving one of its slots. + fn try_acquire(&self, task_name: &str) -> Placement { + let executors = self.executors.lock().unwrap_or_else(recover); + let mut advertised = false; + let mut best: Option<&Arc> = None; + + for executor in executors.values() { + if !executor.tasks.contains(task_name) { + continue; + } + advertised = true; + let free = executor.free.load(Ordering::Relaxed); + if free == 0 { + continue; + } + if best.is_none_or(|current| free > current.free.load(Ordering::Relaxed)) { + best = Some(executor); + } + } + + match best { + // Reserve under the registry lock so a later placement already + // sees this job's slot as taken. + Some(executor) => { + executor.free.fetch_sub(1, Ordering::Relaxed); + Placement::Ready(executor.clone()) + } + None if advertised => Placement::Saturated, + None => Placement::Unadvertised, + } + } + + /// Send a reserved job to its executor. + /// + /// Registers the job before writing so a fast executor cannot return a + /// result the reader can't pair with an in-flight entry. A failed write + /// means the connection is gone: the slot is released, the executor is + /// dropped, and the job is left to the scheduler's reaper — the same + /// recovery path a mid-job executor crash takes. + fn dispatch_to(&self, executor: &Arc, job: Job) { + executor + .in_flight + .lock() + .unwrap_or_else(recover) + .insert(job.id.clone(), job.task_name.clone()); + + let write = executor + .writer + .lock() + .unwrap_or_else(recover) + .write_job(&job); + + if let Err(e) = write { + executor + .in_flight + .lock() + .unwrap_or_else(recover) + .remove(&job.id); + executor.free.fetch_add(1, Ordering::Relaxed); + log::error!( + "[taskito] failed to send job {} to executor {}: {e}; will be reaped", + job.id, + executor.id + ); + self.deregister(&executor.id); + } + } + + /// Hand a job that never reached an executor back as a retryable failure. + fn fail_unplaceable(&self, job: &Job, reason: &str) { + let error = format!("task '{}' was not dispatched: {reason}", job.task_name); + log::warn!("[taskito] {error} (job {})", job.id); + self.emit(JobResult::Failure { + job_id: job.id.clone(), + error, + retry_count: job.retry_count, + max_retries: job.max_retries, + task_name: job.task_name.clone(), + wall_time_ns: 0, + should_retry: true, + timed_out: false, + }); + } + + /// Send a result to the scheduler, if `run` has installed the channel. + fn emit(&self, result: JobResult) { + let sender = self + .result_tx + .lock() + .unwrap_or_else(recover) + .as_ref() + .cloned(); + match sender { + Some(tx) => { + if tx.send(result).is_err() { + log::debug!("[taskito] result channel closed; dropping executor result"); + } + } + None => log::warn!("[taskito] result arrived before the dispatcher started; dropping"), + } + } + + /// Reader thread for one executor: results in, capacity updates, and + /// deregistration when the connection ends. + fn spawn_reader( + self: Arc, + executor: Arc, + mut reader: FrameReader, + ) -> JoinHandle<()> { + thread::Builder::new() + .name(format!("taskito-executor-{}", executor.id)) + .spawn(move || { + loop { + match reader.read::() { + Ok((message, payload)) => { + executor + .last_seen_ms + .store(self.elapsed_ms(), Ordering::Relaxed); + self.handle_frame(&executor, message, payload); + } + Err(ProtocolError::Eof) => { + log::info!("[taskito] executor {} disconnected", executor.id); + break; + } + Err(e) => { + log::warn!("[taskito] executor {} read error: {e}", executor.id); + break; + } + } + } + self.abandon(&executor); + }) + .expect("failed to spawn executor reader thread") + } + + /// Process one frame from an executor. + fn handle_frame(&self, executor: &Arc, message: ExecutorMessage, payload: Vec) { + if let ExecutorMessage::Heartbeat { free_slots } = message { + // Local accounting is exact, so a heartbeat may only *shrink* + // capacity — an executor shedding slots — never invent it. + let _ = executor + .free + .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |current| { + Some(current.min(free_slots)) + }); + return; + } + + let Some(result) = message.into_job_result(payload) else { + log::warn!( + "[taskito] executor {} sent a handshake frame mid-stream", + executor.id + ); + return; + }; + + // Taking the in-flight entry is the exactly-once token: a duplicate or + // unknown result has no entry to take and is dropped. + let known = executor + .in_flight + .lock() + .unwrap_or_else(recover) + .remove(result.job_id()) + .is_some(); + if !known { + log::warn!( + "[taskito] executor {} returned a result for unknown job {}", + executor.id, + result.job_id() + ); + return; + } + + executor.free.fetch_add(1, Ordering::Relaxed); + self.capacity_changed.notify_waiters(); + self.emit(result); + } + + /// Drop an executor whose connection ended, leaving its in-flight jobs to + /// the scheduler's reaper — the same recovery a crashed worker gets, and + /// the only correct answer when a lost result may still have run. + fn abandon(&self, executor: &Arc) { + self.deregister(&executor.id); + let abandoned: Vec = executor + .in_flight + .lock() + .unwrap_or_else(recover) + .drain() + .map(|(job_id, _)| job_id) + .collect(); + if !abandoned.is_empty() { + log::warn!( + "[taskito] executor {} ({}) left {} job(s) in flight; they will be reaped: {}", + executor.id, + executor.peer, + abandoned.len(), + abandoned.join(", ") + ); + } + } + + /// Remove an executor from the registry. Idempotent — the writer and the + /// reader can both discover a broken connection. + fn deregister(&self, executor_id: &str) { + let removed = self + .executors + .lock() + .unwrap_or_else(recover) + .remove(executor_id); + if removed.is_some() { + self.capacity_changed.notify_waiters(); + } + } + + /// Cancel router: forwards requests to whichever executor holds the job. + /// + /// Runs on its own thread so the synchronous, infallible `notify_cancel` + /// never blocks the caller on a wedged peer. + fn spawn_cancel_router(self: Arc, cancel_rx: Receiver) -> JoinHandle<()> { + thread::Builder::new() + .name("taskito-executor-cancel-router".into()) + .spawn(move || { + for job_id in cancel_rx.iter() { + // No executor holds it: already finished, never dispatched, + // or gone. The storage cancel flag covers those. + let Some(executor) = self.executor_running(&job_id) else { + continue; + }; + let sent = executor + .writer + .lock() + .unwrap_or_else(recover) + .write_cancel(&job_id); + if let Err(e) = sent { + log::warn!( + "[taskito] failed to forward cancel for {job_id} to executor {}: {e}", + executor.id + ); + } + } + }) + .expect("failed to spawn executor cancel-router thread") + } + + /// The executor currently running `job_id`, if any. + fn executor_running(&self, job_id: &str) -> Option> { + let executors = self.executors.lock().unwrap_or_else(recover); + executors + .values() + .find(|executor| { + executor + .in_flight + .lock() + .unwrap_or_else(recover) + .contains_key(job_id) + }) + .cloned() + } + + /// Ask every attached executor to finish, wait out the drain budget, then + /// close the connections. + /// + /// Closing is what bounds shutdown: a reader thread is parked on a blocking + /// read, and an executor that stops responding would otherwise keep it — + /// and the join below it — parked forever. + fn drain_and_close(&self) { + let executors: Vec> = self + .executors + .lock() + .unwrap_or_else(recover) + .drain() + .map(|(_, executor)| executor) + .collect(); + + for executor in &executors { + // Best-effort: the executor may already be gone. + let _ = executor + .writer + .lock() + .unwrap_or_else(recover) + .write_shutdown(); + } + + let deadline = Instant::now() + self.config.shutdown_drain; + while Instant::now() < deadline && executors.iter().any(Executor::is_busy) { + std::thread::sleep(DRAIN_POLL); + } + + for executor in &executors { + let still_running = executor.in_flight.lock().unwrap_or_else(recover).len(); + if still_running > 0 { + log::warn!( + "[taskito] executor {} did not drain {still_running} job(s) within the \ + shutdown budget; closing — they will be reaped", + executor.id + ); + } + executor.connection.close(); + } + } + + fn set_cancel_sender(&self, tx: Option>) { + *self.cancel_tx.lock().unwrap_or_else(recover) = tx; + } + + fn notify_cancel(&self, job_id: &str) { + let sender = self + .cancel_tx + .lock() + .unwrap_or_else(recover) + .as_ref() + .cloned(); + let Some(tx) = sender else { + return; + }; + match tx.try_send(job_id.to_string()) { + Ok(()) | Err(TrySendError::Disconnected(_)) => {} + Err(TrySendError::Full(_)) => { + log::warn!("[taskito] executor cancel channel full, dropping cancel for {job_id}"); + } + } + } +} diff --git a/crates/taskito-core/src/worker/remote_tests.rs b/crates/taskito-core/src/worker/remote_tests.rs new file mode 100644 index 000000000..cd010be28 --- /dev/null +++ b/crates/taskito-core/src/worker/remote_tests.rs @@ -0,0 +1,565 @@ +//! Tests for [`RemoteDispatcher`], driven over [`MemoryTransport`] so no +//! socket is bound. A `FakeExecutor` plays the far end of the connection. + +use std::sync::Arc; +use std::time::{Duration, Instant}; + +use crossbeam_channel::{Receiver, RecvTimeoutError}; + +use super::protocol::{ + ExecutorMessage, FrameReader, FrameWriter, ProtocolError, SchedulerMessage, PROTOCOL_VERSION, +}; +use super::remote::{AttachError, RemoteConfig, RemoteDispatcher}; +use super::transport::{MemoryTransport, ReadHalf, Transport, WriteHalf}; +use super::WorkerDispatcher; +use crate::job::{now_millis, Job, JobStatus, NewJob}; +use crate::scheduler::{JobResult, SchedulerConfig}; +use crate::storage::sqlite::SqliteStorage; +use crate::storage::{Storage, StorageBackend}; +use crate::worker::Worker; + +const SETTLE: Duration = Duration::from_secs(5); + +/// The executor side of an attached connection. +struct FakeExecutor { + reader: FrameReader, + writer: FrameWriter, +} + +impl FakeExecutor { + /// Attach to `dispatcher`, announcing `tasks` and `slots`. + fn attach( + dispatcher: &RemoteDispatcher, + executor_id: &str, + tasks: &[&str], + slots: u32, + ) -> Result { + let (executor, attached) = + Self::attach_with_version(dispatcher, executor_id, tasks, slots, PROTOCOL_VERSION); + attached.map(|_| executor) + } + + /// Attach announcing `protocol_version`, keeping the executor end whatever + /// the outcome so a rejected handshake can still be inspected. + fn attach_with_version( + dispatcher: &RemoteDispatcher, + executor_id: &str, + tasks: &[&str], + slots: u32, + protocol_version: u32, + ) -> (Self, Result) { + let (scheduler_end, executor_end) = MemoryTransport::pair(); + let (read, write, _timeout) = Box::new(executor_end).split().expect("split executor end"); + let mut executor = Self { + reader: FrameReader::new(read), + writer: FrameWriter::new(write), + }; + + executor + .writer + .write_header(&ExecutorMessage::Hello { + executor_id: executor_id.to_string(), + sdk: "test".to_string(), + version: "0.0.0".to_string(), + tasks: tasks.iter().map(|t| (*t).to_string()).collect(), + slots, + protocol_version, + }) + .expect("send hello"); + + let attached = dispatcher.attach(Box::new(scheduler_end)); + (executor, attached) + } + + fn read(&mut self) -> Result<(SchedulerMessage, Vec), ProtocolError> { + self.reader.read::() + } + + fn expect_hello_ack(&mut self) -> u32 { + match self.read().expect("read ack").0 { + SchedulerMessage::HelloAck { + protocol_version, .. + } => protocol_version, + other => panic!("expected hello_ack, got {other:?}"), + } + } + + fn expect_shutdown(&mut self) { + loop { + match self.read().expect("read frame").0 { + SchedulerMessage::HelloAck { .. } => continue, + SchedulerMessage::Shutdown => return, + other => panic!("expected shutdown, got {other:?}"), + } + } + } + + /// Read the next job frame, skipping the handshake ack if still queued. + fn expect_job(&mut self) -> (String, String, Vec) { + loop { + match self.read().expect("read frame") { + (SchedulerMessage::HelloAck { .. }, _) => continue, + (SchedulerMessage::Job { id, task_name, .. }, payload) => { + return (id, task_name, payload) + } + (other, _) => panic!("expected a job frame, got {other:?}"), + } + } + } + + fn succeed(&mut self, job_id: &str, task_name: &str, result: Option<&[u8]>) { + self.writer + .write( + &ExecutorMessage::Success { + job_id: job_id.to_string(), + result_len: result.map(<[u8]>::len), + task_name: task_name.to_string(), + wall_time_ns: 1, + }, + result.unwrap_or(&[]), + ) + .expect("send success"); + } +} + +fn dispatcher_with(placement_timeout: Duration) -> RemoteDispatcher { + RemoteDispatcher::new(RemoteConfig { + scheduler_id: "scheduler-test".to_string(), + placement_timeout, + // Keep teardown snappy: no test leaves work in flight on purpose + // except the one that asserts the drain budget is honoured. + shutdown_drain: Duration::from_millis(200), + ..RemoteConfig::default() + }) +} + +fn make_job(id: &str, task_name: &str, payload: &[u8]) -> Job { + Job { + id: id.to_string(), + queue: "default".to_string(), + task_name: task_name.to_string(), + payload: payload.to_vec(), + status: JobStatus::Running, + priority: 0, + created_at: 0, + scheduled_at: 0, + started_at: None, + completed_at: None, + retry_count: 0, + max_retries: 3, + result: None, + error: None, + timeout_ms: 30_000, + unique_key: None, + progress: None, + metadata: None, + notes: None, + cancel_requested: false, + expires_at: None, + result_ttl_ms: None, + namespace: None, + has_deps: false, + } +} + +/// Run `body` with the dispatcher's `run` loop live on a current-thread runtime. +fn with_running(dispatcher: &RemoteDispatcher, capacity: usize, body: F) +where + F: FnOnce(&tokio::sync::mpsc::Sender, &Receiver), +{ + let runtime = tokio::runtime::Builder::new_multi_thread() + .worker_threads(2) + .enable_all() + .build() + .expect("runtime"); + + let (job_tx, job_rx) = tokio::sync::mpsc::channel(capacity); + let (result_tx, result_rx) = crossbeam_channel::bounded(capacity); + + let running = { + let dispatcher = dispatcher.clone(); + runtime.spawn(async move { dispatcher.run(job_rx, result_tx).await }) + }; + + body(&job_tx, &result_rx); + + drop(job_tx); + runtime.block_on(async { running.await.expect("run loop") }); +} + +fn expect_result(results: &Receiver) -> JobResult { + results.recv_timeout(SETTLE).expect("a result") +} + +/// `JobResult` deliberately has no `Debug` — it carries task payloads — so +/// assertion messages name the variant instead of dumping it. +fn kind(result: &JobResult) -> &'static str { + match result { + JobResult::Success { .. } => "success", + JobResult::Failure { .. } => "failure", + JobResult::Cancelled { .. } => "cancelled", + } +} + +#[test] +fn handshake_registers_the_executor() { + let dispatcher = dispatcher_with(Duration::from_millis(200)); + let mut executor = + FakeExecutor::attach(&dispatcher, "exec-1", &["resize", "thumbnail"], 3).expect("attach"); + + assert_eq!(executor.expect_hello_ack(), PROTOCOL_VERSION); + + let attached = dispatcher.executors(); + assert_eq!(attached.len(), 1); + assert_eq!(attached[0].executor_id, "exec-1"); + assert_eq!(attached[0].tasks, ["resize", "thumbnail"]); + assert_eq!(attached[0].slots, 3); + assert_eq!(attached[0].free_slots, 3); + + let capacity = dispatcher.capacity(); + assert_eq!(capacity.executors, 1); + assert_eq!(capacity.total_slots, 3); + assert_eq!(capacity.free_slots, 3); +} + +#[test] +fn version_mismatch_is_rejected_but_still_acked() { + let dispatcher = dispatcher_with(Duration::from_millis(200)); + let (mut executor, attached) = FakeExecutor::attach_with_version( + &dispatcher, + "exec-old", + &["resize"], + 1, + PROTOCOL_VERSION + 1, + ); + + match attached.expect_err("mismatched version must be rejected") { + AttachError::Protocol(ProtocolError::VersionMismatch { ours, theirs }) => { + assert_eq!(ours, PROTOCOL_VERSION); + assert_eq!(theirs, PROTOCOL_VERSION + 1); + } + other => panic!("expected a version mismatch, got {other:?}"), + } + // The ack still went out, so the executor can name both versions too. + assert_eq!(executor.expect_hello_ack(), PROTOCOL_VERSION); + assert!(dispatcher.executors().is_empty()); +} + +#[test] +fn duplicate_executor_id_is_rejected() { + let dispatcher = dispatcher_with(Duration::from_millis(200)); + let _first = FakeExecutor::attach(&dispatcher, "exec-1", &["resize"], 1).expect("attach"); + let err = FakeExecutor::attach(&dispatcher, "exec-1", &["resize"], 1) + .err() + .expect("a duplicate id must be rejected"); + assert!(matches!(err, AttachError::DuplicateId(id) if id == "exec-1")); + assert_eq!(dispatcher.executors().len(), 1); +} + +#[test] +fn job_round_trips_to_an_executor_and_back() { + let dispatcher = dispatcher_with(Duration::from_secs(5)); + let mut executor = FakeExecutor::attach(&dispatcher, "exec-1", &["resize"], 1).expect("attach"); + + with_running(&dispatcher, 4, |jobs, results| { + jobs.blocking_send(make_job("job-1", "resize", b"in")) + .expect("send job"); + + let (job_id, task_name, payload) = executor.expect_job(); + assert_eq!(job_id, "job-1"); + assert_eq!(task_name, "resize"); + assert_eq!(payload, b"in"); + + executor.succeed("job-1", "resize", Some(b"out")); + + match expect_result(results) { + JobResult::Success { + job_id, + result, + task_name, + .. + } => { + assert_eq!(job_id, "job-1"); + assert_eq!(task_name, "resize"); + assert_eq!(result.as_deref(), Some(&b"out"[..])); + } + ref other => panic!("expected success, got {}", kind(other)), + } + }); +} + +#[test] +fn unadvertised_task_is_never_sent_and_fails_retryably() { + let dispatcher = dispatcher_with(Duration::from_millis(100)); + let mut executor = FakeExecutor::attach(&dispatcher, "exec-1", &["resize"], 1).expect("attach"); + assert_eq!(executor.expect_hello_ack(), PROTOCOL_VERSION); + + with_running(&dispatcher, 4, |jobs, results| { + jobs.blocking_send(make_job("job-1", "transcode", b"")) + .expect("send job"); + + match expect_result(results) { + JobResult::Failure { + job_id, + should_retry, + timed_out, + error, + .. + } => { + assert_eq!(job_id, "job-1"); + assert!(should_retry, "an unplaced job must be retryable"); + assert!(!timed_out, "never dispatched is not a timeout"); + assert!(error.contains("transcode"), "error names the task: {error}"); + } + ref other => panic!("expected a retryable failure, got {}", kind(other)), + } + }); + + // The next frame after the ack is the shutdown — no job was ever written. + executor.expect_shutdown(); +} + +#[test] +fn a_single_slot_admits_one_job_at_a_time() { + let dispatcher = dispatcher_with(Duration::from_secs(5)); + let mut executor = FakeExecutor::attach(&dispatcher, "exec-1", &["resize"], 1).expect("attach"); + + with_running(&dispatcher, 4, |jobs, results| { + jobs.blocking_send(make_job("job-1", "resize", b"")) + .expect("send first"); + jobs.blocking_send(make_job("job-2", "resize", b"")) + .expect("send second"); + + let (first, _, _) = executor.expect_job(); + assert_eq!(first, "job-1"); + assert_eq!(dispatcher.capacity().free_slots, 0, "slot must be reserved"); + + // The second job stays unplaced for as long as the slot is occupied — + // asserted against the dispatcher's own bookkeeping rather than a + // read timeout, so the test cannot pass on a slow reader. + let watch = Instant::now() + Duration::from_millis(300); + while Instant::now() < watch { + assert_eq!( + dispatcher.executors()[0].in_flight, + 1, + "only one job may be in flight on a single-slot executor" + ); + std::thread::sleep(Duration::from_millis(10)); + } + + executor.succeed("job-1", "resize", None); + assert!(matches!(expect_result(results), JobResult::Success { .. })); + + let (second, _, _) = executor.expect_job(); + assert_eq!(second, "job-2"); + executor.succeed("job-2", "resize", None); + assert!(matches!(expect_result(results), JobResult::Success { .. })); + }); +} + +#[test] +fn cancel_reaches_the_executor_running_the_job() { + let dispatcher = dispatcher_with(Duration::from_secs(5)); + let mut executor = FakeExecutor::attach(&dispatcher, "exec-1", &["resize"], 1).expect("attach"); + + with_running(&dispatcher, 4, |jobs, results| { + jobs.blocking_send(make_job("job-1", "resize", b"")) + .expect("send job"); + let (job_id, _, _) = executor.expect_job(); + assert_eq!(job_id, "job-1"); + + dispatcher.notify_cancel("job-1"); + match executor.read().expect("read cancel").0 { + SchedulerMessage::Cancel { job_id } => assert_eq!(job_id, "job-1"), + other => panic!("expected a cancel frame, got {other:?}"), + } + + executor + .writer + .write_header(&ExecutorMessage::Cancelled { + job_id: "job-1".to_string(), + task_name: "resize".to_string(), + wall_time_ns: 1, + }) + .expect("send cancelled"); + assert!(matches!( + expect_result(results), + JobResult::Cancelled { .. } + )); + }); +} + +#[test] +fn executor_drop_leaves_its_in_flight_job_to_the_reaper() { + let dispatcher = dispatcher_with(Duration::from_millis(100)); + let mut executor = FakeExecutor::attach(&dispatcher, "exec-1", &["resize"], 1).expect("attach"); + + with_running(&dispatcher, 4, |jobs, results| { + jobs.blocking_send(make_job("job-1", "resize", b"")) + .expect("send job"); + let (job_id, _, _) = executor.expect_job(); + assert_eq!(job_id, "job-1"); + + // The executor dies mid-job. + drop(executor); + + // No result is synthesised: the job may have run, so recovery belongs + // to the scheduler's reaper, not to the dispatcher. + match results.recv_timeout(Duration::from_millis(500)) { + Err(RecvTimeoutError::Timeout) => {} + Ok(ref unexpected) => panic!( + "the dispatcher must not synthesise a {} for an abandoned job", + kind(unexpected) + ), + Err(RecvTimeoutError::Disconnected) => panic!("result channel closed"), + } + + let deadline = Instant::now() + SETTLE; + while !dispatcher.executors().is_empty() { + assert!(Instant::now() < deadline, "executor was never deregistered"); + std::thread::sleep(Duration::from_millis(10)); + } + assert_eq!(dispatcher.capacity().executors, 0); + }); +} + +#[test] +fn heartbeat_can_shrink_capacity_but_not_invent_it() { + let dispatcher = dispatcher_with(Duration::from_secs(5)); + let mut executor = FakeExecutor::attach(&dispatcher, "exec-1", &["resize"], 4).expect("attach"); + assert_eq!(executor.expect_hello_ack(), PROTOCOL_VERSION); + + executor + .writer + .write_header(&ExecutorMessage::Heartbeat { free_slots: 1 }) + .expect("send heartbeat"); + wait_until(|| dispatcher.capacity().free_slots == 1, "capacity shrinks"); + + executor + .writer + .write_header(&ExecutorMessage::Heartbeat { free_slots: 9 }) + .expect("send heartbeat"); + std::thread::sleep(Duration::from_millis(100)); + assert_eq!( + dispatcher.capacity().free_slots, + 1, + "a heartbeat must never invent capacity" + ); +} + +#[test] +fn a_dropped_executor_leaves_its_job_for_the_scheduler_to_recover() { + let storage = StorageBackend::Sqlite(SqliteStorage::in_memory().expect("in-memory sqlite")); + let dispatcher = dispatcher_with(Duration::from_millis(200)); + let mut executor = FakeExecutor::attach(&dispatcher, "exec-1", &["resize"], 1).expect("attach"); + + let handle = Worker::new(storage.clone()) + .num_workers(1) + .scheduler_config(SchedulerConfig { + poll_interval: Duration::from_millis(10), + reap_interval: 1, + ..SchedulerConfig::default() + }) + .dispatcher("remote", Arc::new(dispatcher.clone())) + .spawn() + .expect("spawn"); + + let job = storage + .enqueue(NewJob { + queue: "default".to_string(), + task_name: "resize".to_string(), + payload: b"in".to_vec(), + priority: 0, + scheduled_at: now_millis(), + max_retries: 3, + // Short enough that the stale-job sweep recovers it promptly. + timeout_ms: 500, + unique_key: None, + metadata: None, + notes: None, + depends_on: vec![], + expires_at: None, + result_ttl_ms: None, + namespace: None, + }) + .expect("enqueue"); + + let (job_id, _, _) = executor.expect_job(); + assert_eq!(job_id, job.id); + + // The executor dies holding the job; the scheduler must recover it. + drop(executor); + wait_until( + || { + storage + .get_job(&job.id) + .expect("get_job") + .is_some_and(|recovered| recovered.retry_count > 0) + }, + "job was never recovered after the executor dropped", + ); + + handle.shutdown().expect("shutdown"); +} + +#[test] +fn an_idle_executor_outlives_the_handshake_budget() { + // Regression: the handshake read timeout used to leak onto the attached + // connection, so every executor was dropped once it idled past it. + let dispatcher = RemoteDispatcher::new(RemoteConfig { + scheduler_id: "scheduler-test".to_string(), + handshake_timeout: Duration::from_millis(50), + placement_timeout: Duration::from_secs(5), + shutdown_drain: Duration::from_millis(200), + ..RemoteConfig::default() + }); + let mut executor = FakeExecutor::attach(&dispatcher, "exec-1", &["resize"], 1).expect("attach"); + + std::thread::sleep(Duration::from_millis(250)); + assert_eq!( + dispatcher.executors().len(), + 1, + "an idle executor must not be dropped" + ); + + with_running(&dispatcher, 4, |jobs, results| { + jobs.blocking_send(make_job("job-1", "resize", b"")) + .expect("send job"); + let (job_id, _, _) = executor.expect_job(); + assert_eq!(job_id, "job-1"); + executor.succeed("job-1", "resize", None); + assert!(matches!(expect_result(results), JobResult::Success { .. })); + }); +} + +#[test] +fn shutdown_is_bounded_by_the_drain_budget() { + // Regression: shutdown joined reader threads parked on a blocking read, so + // an executor that stopped responding hung the worker forever. + let dispatcher = dispatcher_with(Duration::from_secs(5)); + let mut executor = FakeExecutor::attach(&dispatcher, "exec-1", &["resize"], 1).expect("attach"); + + let started = Instant::now(); + with_running(&dispatcher, 4, |jobs, _results| { + jobs.blocking_send(make_job("job-1", "resize", b"")) + .expect("send job"); + // The executor takes the job and then goes silent — it never replies + // and never closes its end. + let (job_id, _, _) = executor.expect_job(); + assert_eq!(job_id, "job-1"); + }); + + assert!( + started.elapsed() < SETTLE, + "shutdown must not wait on an unresponsive executor (took {:?})", + started.elapsed() + ); +} + +fn wait_until(mut condition: impl FnMut() -> bool, message: &str) { + let deadline = Instant::now() + SETTLE; + while !condition() { + assert!(Instant::now() < deadline, "{message}"); + std::thread::sleep(Duration::from_millis(10)); + } +} From 4e7265927fb6112963faffdc10aa417023f41d62 Mon Sep 17 00:00:00 2001 From: kartikeya-27 <67143288+kartikeya-27@users.noreply.github.com> Date: Fri, 31 Jul 2026 13:19:57 +0530 Subject: [PATCH 07/15] fix(core): re-export Capacity at the crate root RemoteDispatcher::capacity() returns it, so a root-level consumer could not name its own return type. --- crates/taskito-core/src/lib.rs | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/crates/taskito-core/src/lib.rs b/crates/taskito-core/src/lib.rs index ff5fe95c4..596d67842 100644 --- a/crates/taskito-core/src/lib.rs +++ b/crates/taskito-core/src/lib.rs @@ -46,7 +46,7 @@ pub use storage::Storage; pub use storage::StorageBackend; pub use storage::{DeadJob, QueueStats, SubscriptionBacklogStats}; pub use worker::{ - AttachError, AttachedExecutor, ExecutorMessage, NativeDispatcher, ProtocolError, RemoteConfig, - RemoteDispatcher, SchedulerMessage, TaskError, TaskRegistry, TaskResult, Transport, Worker, - WorkerDispatcher, WorkerHandle, PROTOCOL_VERSION, + AttachError, AttachedExecutor, Capacity, ExecutorMessage, NativeDispatcher, ProtocolError, + RemoteConfig, RemoteDispatcher, SchedulerMessage, TaskError, TaskRegistry, TaskResult, + Transport, Worker, WorkerDispatcher, WorkerHandle, PROTOCOL_VERSION, }; From 3c69a66b4990a635ffe89a25a6b23493ef65182d Mon Sep 17 00:00:00 2001 From: kartikeya-27 <67143288+kartikeya-27@users.noreply.github.com> Date: Fri, 31 Jul 2026 13:20:20 +0530 Subject: [PATCH 08/15] fix(core): report a truncated header as I/O, not oversized A peer closing mid-header reported HeaderTooLarge, contradicting the Eof contract and mislabelling a plain disconnect as an oversized frame. --- crates/taskito-core/src/worker/protocol.rs | 23 +++++++++++++++++++++- 1 file changed, 22 insertions(+), 1 deletion(-) diff --git a/crates/taskito-core/src/worker/protocol.rs b/crates/taskito-core/src/worker/protocol.rs index 7eb87570a..c3a93eeb7 100644 --- a/crates/taskito-core/src/worker/protocol.rs +++ b/crates/taskito-core/src/worker/protocol.rs @@ -375,7 +375,17 @@ impl FrameReader { return Err(ProtocolError::Eof); } if !header.ends_with(b"\n") { - return Err(ProtocolError::HeaderTooLarge); + // Only a header that actually reached the cap is oversized. Fewer + // bytes with no newline means the peer closed mid-header — a + // truncated frame, which the `Eof` contract says surfaces as I/O. + if read as u64 == MAX_HEADER_BYTES { + return Err(ProtocolError::HeaderTooLarge); + } + return Err(std::io::Error::new( + std::io::ErrorKind::UnexpectedEof, + "peer closed mid-header", + ) + .into()); } Ok(header) } @@ -653,6 +663,17 @@ mod tests { assert!(matches!(err, ProtocolError::Io(_)), "got {err:?}"); } + #[test] + fn truncated_header_is_an_io_error_not_an_oversized_header() { + // A short header with no newline is a mid-frame disconnect, not a peer + // that blew the size cap — the two must not report the same way. + let buf = br#"{"type":"job","id":"j""#; + let err = FrameReader::new(&buf[..]) + .read::() + .expect_err("truncated header must not read"); + assert!(matches!(err, ProtocolError::Io(_)), "got {err:?}"); + } + #[test] fn malformed_header_is_reported_as_json_error() { let buf = b"not json\n"; From aa1ff74998e722c298b0bf7e7ffe511b789ad3c5 Mon Sep 17 00:00:00 2001 From: kartikeya-27 <67143288+kartikeya-27@users.noreply.github.com> Date: Fri, 31 Jul 2026 13:20:49 +0530 Subject: [PATCH 09/15] fix(prefork): reap the child on a failed handshake Child::drop does not terminate the process, so a version-mismatched child survived as an orphan the restart path never notices. --- crates/taskito-python/src/prefork/child.rs | 18 +++++++++++++++++- 1 file changed, 17 insertions(+), 1 deletion(-) diff --git a/crates/taskito-python/src/prefork/child.rs b/crates/taskito-python/src/prefork/child.rs index d4d5bff50..ce2972ec0 100644 --- a/crates/taskito-python/src/prefork/child.rs +++ b/crates/taskito-python/src/prefork/child.rs @@ -89,7 +89,23 @@ pub fn spawn_child( let mut reader = ChildReader::new(BufReader::new(stdout)); let mut writer = ChildWriter::new(stdin); + // `Child::drop` does not terminate the process, so every failure below has + // to reap it or the child survives as an orphan the restart path never + // notices (it is only ever reached via `is_alive()` on a live handle). + let mut child = ChildProcess { process }; + + let handshake = handshake(&mut reader, &mut writer); + match handshake { + Ok(()) => Ok((writer, reader, child)), + Err(e) => { + child.kill_and_reap(); + Err(e) + } + } +} +/// Read the child's `hello`, acknowledge it, and check the protocol version. +fn handshake(reader: &mut ChildReader, writer: &mut ChildWriter) -> Result<(), String> { let hello = reader .read::() .map_err(|e| format!("child handshake failed: {e}"))? @@ -118,5 +134,5 @@ pub fn spawn_child( )); } - Ok((writer, reader, ChildProcess { process })) + Ok(()) } From 1907b11976a60cbd5e7b5a9b7fce405b0c936481 Mon Sep 17 00:00:00 2001 From: kartikeya-27 <67143288+kartikeya-27@users.noreply.github.com> Date: Fri, 31 Jul 2026 13:20:56 +0530 Subject: [PATCH 10/15] fix(python): validate a declared frame length before reading A negative length reached stream.read(), which drains the connection to EOF rather than reading a payload. --- sdks/python/taskito/worker_protocol.py | 31 +++++++++---- .../tests/worker/test_worker_protocol.py | 46 +++++++++++++++++++ 2 files changed, 68 insertions(+), 9 deletions(-) create mode 100644 sdks/python/tests/worker/test_worker_protocol.py diff --git a/sdks/python/taskito/worker_protocol.py b/sdks/python/taskito/worker_protocol.py index 94c782011..baa22ab74 100644 --- a/sdks/python/taskito/worker_protocol.py +++ b/sdks/python/taskito/worker_protocol.py @@ -38,16 +38,34 @@ class ProtocolError(Exception): def declared_payload_len(header: dict[str, Any]) -> int: - """Bytes of payload a header says follow it.""" + """Bytes of payload a header says follow it. + + Validated here rather than at the call sites: a negative or non-integer + length would otherwise reach ``stream.read(n)``, where a negative count + drains the connection to EOF. + """ kind = header.get("type") if kind == "job": - return int(header.get("payload_len") or 0) + return _checked_len(header.get("payload_len") or 0, "payload_len") if kind == "success": result_len = header.get("result_len") - return 0 if result_len is None else int(result_len) + return 0 if result_len is None else _checked_len(result_len, "result_len") return 0 +def _checked_len(value: Any, field: str) -> int: + """Coerce a declared length, rejecting anything unusable as a read count.""" + if isinstance(value, bool) or not isinstance(value, int): + raise ProtocolError(f"{field} must be an integer, got {type(value).__name__}") + if value < 0: + raise ProtocolError(f"{field} must not be negative, got {value}") + if value > MAX_PAYLOAD_BYTES: + raise ProtocolError( + f"frame payload of {value} bytes exceeds the {MAX_PAYLOAD_BYTES} byte limit" + ) + return int(value) + + def write_frame(stream: BinaryIO, header: dict[str, Any], payload: bytes = b"") -> None: """Write one frame and flush it. @@ -80,12 +98,7 @@ def read_frame(stream: BinaryIO) -> tuple[dict[str, Any], bytes]: if not isinstance(header, dict): raise ProtocolError("frame header must be a JSON object") - length = declared_payload_len(header) - if length > MAX_PAYLOAD_BYTES: - raise ProtocolError( - f"frame payload of {length} bytes exceeds the {MAX_PAYLOAD_BYTES} byte limit" - ) - return header, _read_exact(stream, length) + return header, _read_exact(stream, declared_payload_len(header)) def _read_exact(stream: BinaryIO, length: int) -> bytes: diff --git a/sdks/python/tests/worker/test_worker_protocol.py b/sdks/python/tests/worker/test_worker_protocol.py new file mode 100644 index 000000000..8f1f6c926 --- /dev/null +++ b/sdks/python/tests/worker/test_worker_protocol.py @@ -0,0 +1,46 @@ +"""Frame codec guards: a malformed length must never become a read count.""" + +from __future__ import annotations + +import io + +import pytest + +from taskito.worker_protocol import ProtocolError, read_frame, write_frame + + +def _frame(header: bytes, payload: bytes = b"") -> io.BytesIO: + return io.BytesIO(header + b"\n" + payload) + + +def test_negative_payload_len_is_rejected() -> None: + # read(-1) would drain the connection to EOF rather than read a payload. + stream = _frame(b'{"type":"job","payload_len":-1}', b"trailing") + with pytest.raises(ProtocolError, match="must not be negative"): + read_frame(stream) + + +def test_non_integer_payload_len_is_rejected() -> None: + stream = _frame(b'{"type":"job","payload_len":"12"}') + with pytest.raises(ProtocolError, match="must be an integer"): + read_frame(stream) + + +def test_oversized_payload_len_is_rejected() -> None: + stream = _frame(b'{"type":"job","payload_len":99999999999}') + with pytest.raises(ProtocolError, match="exceeds"): + read_frame(stream) + + +def test_write_rejects_a_length_that_disagrees_with_the_payload() -> None: + with pytest.raises(ProtocolError, match="declared"): + write_frame(io.BytesIO(), {"type": "job", "payload_len": 4}, b"ab") + + +def test_round_trip_preserves_bytes_containing_newlines() -> None: + payload = b'\n{"type":"success"}\n\x00\xff' + sink = io.BytesIO() + write_frame(sink, {"type": "job", "payload_len": len(payload)}, payload) + header, read_back = read_frame(io.BytesIO(sink.getvalue())) + assert header["type"] == "job" + assert read_back == payload From 9a7f4cc088ace799f9430ff2892db436ce7da22e Mon Sep 17 00:00:00 2001 From: kartikeya-27 <67143288+kartikeya-27@users.noreply.github.com> Date: Fri, 31 Jul 2026 13:21:39 +0530 Subject: [PATCH 11/15] fix(core): reap finished executor reader handles The vector only shrank at shutdown, so a reconnecting executor grew it for the life of the process. --- crates/taskito-core/src/worker/remote.rs | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/crates/taskito-core/src/worker/remote.rs b/crates/taskito-core/src/worker/remote.rs index de7fa015a..cab16ed38 100644 --- a/crates/taskito-core/src/worker/remote.rs +++ b/crates/taskito-core/src/worker/remote.rs @@ -308,7 +308,13 @@ impl Shared { } let handle = Arc::clone(self).spawn_reader(executor, reader); - self.readers.lock().unwrap_or_else(recover).push(handle); + { + let mut readers = self.readers.lock().unwrap_or_else(recover); + // Reap handles of already-detached executors so a reconnect loop + // cannot grow this vector for the life of the process. + readers.retain(|reader| !reader.is_finished()); + readers.push(handle); + } self.capacity_changed.notify_waiters(); log::info!("[taskito] executor {executor_id} attached from {peer} with {slots} slot(s)"); From 587d2bee2ca5458dd8e728cdc3aa37af3e99d24a Mon Sep 17 00:00:00 2001 From: kartikeya-27 <67143288+kartikeya-27@users.noreply.github.com> Date: Fri, 31 Jul 2026 13:22:02 +0530 Subject: [PATCH 12/15] fix(core): refuse an attach once shutdown has started An executor attaching after the registry drained was never sent a shutdown and never joined, leaking its reader thread. --- crates/taskito-core/src/worker/remote.rs | 13 ++++++++++++ .../taskito-core/src/worker/remote_tests.rs | 21 +++++++++++++++++++ 2 files changed, 34 insertions(+) diff --git a/crates/taskito-core/src/worker/remote.rs b/crates/taskito-core/src/worker/remote.rs index cab16ed38..ebb5af0ea 100644 --- a/crates/taskito-core/src/worker/remote.rs +++ b/crates/taskito-core/src/worker/remote.rs @@ -67,6 +67,10 @@ pub enum AttachError { /// Another executor is already attached under this id. #[error("executor {0} is already attached")] DuplicateId(String), + + /// The dispatcher is shutting down and accepts no new executors. + #[error("dispatcher is shutting down")] + ShuttingDown, } /// A snapshot of one attached executor. @@ -245,6 +249,9 @@ impl Shared { /// Handshake and register. The ack is sent even when the version is /// rejected, so both ends log both versions instead of one side guessing. fn attach(self: &Arc, transport: Box) -> Result { + if self.shutdown.load(Ordering::SeqCst) { + return Err(AttachError::ShuttingDown); + } let peer = transport.peer(); let (read, write, connection) = transport.split()?; let mut reader = FrameReader::new(read); @@ -301,6 +308,12 @@ impl Shared { { let mut executors = self.executors.lock().unwrap_or_else(recover); + // Re-check under the registry lock: `drain_and_close` empties this + // map, so an attach racing it would leave an executor nobody ever + // shuts down or joins. + if self.shutdown.load(Ordering::SeqCst) { + return Err(AttachError::ShuttingDown); + } if executors.contains_key(&executor_id) { return Err(AttachError::DuplicateId(executor_id)); } diff --git a/crates/taskito-core/src/worker/remote_tests.rs b/crates/taskito-core/src/worker/remote_tests.rs index cd010be28..ba91e7c2f 100644 --- a/crates/taskito-core/src/worker/remote_tests.rs +++ b/crates/taskito-core/src/worker/remote_tests.rs @@ -502,6 +502,27 @@ fn a_dropped_executor_leaves_its_job_for_the_scheduler_to_recover() { handle.shutdown().expect("shutdown"); } +#[test] +fn attach_is_refused_once_shutdown_has_started() { + // An executor attaching after `drain_and_close` drained the registry would + // never be told to stop and never be joined — a leaked reader thread. + let dispatcher = dispatcher_with(Duration::from_millis(200)); + WorkerDispatcher::shutdown(&dispatcher); + + let (_executor, attached) = FakeExecutor::attach_with_version( + &dispatcher, + "exec-late", + &["resize"], + 1, + PROTOCOL_VERSION, + ); + assert!(matches!( + attached.expect_err("attach must be refused during shutdown"), + AttachError::ShuttingDown + )); + assert!(dispatcher.executors().is_empty()); +} + #[test] fn an_idle_executor_outlives_the_handshake_budget() { // Regression: the handshake read timeout used to leak onto the attached From 3bdb2ed040950f95f81204224a605a012bdb55b1 Mon Sep 17 00:00:00 2001 From: kartikeya-27 <67143288+kartikeya-27@users.noreply.github.com> Date: Fri, 31 Jul 2026 13:22:25 +0530 Subject: [PATCH 13/15] fix(core): stop waiting for a slot once shutdown starts A job parked on placement ignored the flag, so shutdown blocked for a full placement_timeout before the drain even began. --- crates/taskito-core/src/worker/remote.rs | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/crates/taskito-core/src/worker/remote.rs b/crates/taskito-core/src/worker/remote.rs index ebb5af0ea..954205aa2 100644 --- a/crates/taskito-core/src/worker/remote.rs +++ b/crates/taskito-core/src/worker/remote.rs @@ -420,6 +420,11 @@ impl Shared { Placement::Unadvertised => "no attached executor advertises it", }; + if self.shutdown.load(Ordering::Relaxed) { + self.fail_unplaceable(&job, "the dispatcher is shutting down"); + return; + } + let remaining = deadline.saturating_duration_since(Instant::now()); if remaining.is_zero() || tokio::time::timeout(remaining, changed).await.is_err() { self.fail_unplaceable(&job, reason); From e156ce1a7c4d10be8a2b12b434d18c969655c6db Mon Sep 17 00:00:00 2001 From: kartikeya-27 <67143288+kartikeya-27@users.noreply.github.com> Date: Fri, 31 Jul 2026 13:22:32 +0530 Subject: [PATCH 14/15] fix(core): await the shutdown drain instead of sleeping run shares a runtime with the scheduler task, so a blocking sleep starved it for up to the whole drain budget. --- crates/taskito-core/src/worker/remote.rs | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/crates/taskito-core/src/worker/remote.rs b/crates/taskito-core/src/worker/remote.rs index 954205aa2..ee1449c1f 100644 --- a/crates/taskito-core/src/worker/remote.rs +++ b/crates/taskito-core/src/worker/remote.rs @@ -385,7 +385,7 @@ impl Shared { // Stop accepting cancels so the router drains and exits while the // writers it uses are still alive. self.set_cancel_sender(None); - self.drain_and_close(); + self.drain_and_close().await; let readers = std::mem::take(&mut *self.readers.lock().unwrap_or_else(recover)); for handle in readers { @@ -699,7 +699,7 @@ impl Shared { /// Closing is what bounds shutdown: a reader thread is parked on a blocking /// read, and an executor that stops responding would otherwise keep it — /// and the join below it — parked forever. - fn drain_and_close(&self) { + async fn drain_and_close(&self) { let executors: Vec> = self .executors .lock() @@ -717,9 +717,11 @@ impl Shared { .write_shutdown(); } + // Awaited, not slept: `run` shares a runtime with the scheduler task, + // and a blocking sleep here starves it for the whole drain budget. let deadline = Instant::now() + self.config.shutdown_drain; while Instant::now() < deadline && executors.iter().any(Executor::is_busy) { - std::thread::sleep(DRAIN_POLL); + tokio::time::sleep(DRAIN_POLL).await; } for executor in &executors { From 6d405b685d57f3d82802c57a3427df6dd279af63 Mon Sep 17 00:00:00 2001 From: kartikeya-27 <67143288+kartikeya-27@users.noreply.github.com> Date: Fri, 31 Jul 2026 13:23:03 +0530 Subject: [PATCH 15/15] fix(core): bound the dispatch write to an executor An executor that stops reading fills the send buffer; the write had no timeout and could park a runtime thread indefinitely. --- crates/taskito-core/src/worker/remote.rs | 5 +++++ crates/taskito-core/src/worker/transport.rs | 19 +++++++++++++++++++ 2 files changed, 24 insertions(+) diff --git a/crates/taskito-core/src/worker/remote.rs b/crates/taskito-core/src/worker/remote.rs index ee1449c1f..f6e061bfd 100644 --- a/crates/taskito-core/src/worker/remote.rs +++ b/crates/taskito-core/src/worker/remote.rs @@ -33,6 +33,9 @@ pub struct RemoteConfig { /// How long a job waits for a slot on an executor advertising its task /// before it is failed back retryably. pub placement_timeout: Duration, + /// How long a dispatch write may block before the executor is treated as + /// wedged. Bounds the dispatch thread against a peer that stops reading. + pub write_timeout: Duration, /// How long shutdown waits for attached executors to finish in-flight /// jobs before their connections are closed. pub shutdown_drain: Duration, @@ -46,6 +49,7 @@ impl Default for RemoteConfig { scheduler_id: format!("scheduler-{}", uuid::Uuid::now_v7()), handshake_timeout: Duration::from_secs(10), placement_timeout: Duration::from_secs(30), + write_timeout: Duration::from_secs(30), shutdown_drain: Duration::from_secs(30), cancel_capacity: 1024, } @@ -254,6 +258,7 @@ impl Shared { } let peer = transport.peer(); let (read, write, connection) = transport.split()?; + connection.set_write_timeout(Some(self.config.write_timeout))?; let mut reader = FrameReader::new(read); let mut writer = FrameWriter::new(write); diff --git a/crates/taskito-core/src/worker/transport.rs b/crates/taskito-core/src/worker/transport.rs index 3931e4921..93f1ef6e6 100644 --- a/crates/taskito-core/src/worker/transport.rs +++ b/crates/taskito-core/src/worker/transport.rs @@ -26,6 +26,7 @@ pub type WriteHalf = Box; /// reader return, so shutdown cannot hang on a peer that stops responding. pub struct Connection { set_read_timeout: Box) -> io::Result<()> + Send + Sync>, + set_write_timeout: Box) -> io::Result<()> + Send + Sync>, close: Box, } @@ -38,6 +39,14 @@ impl Connection { (self.set_read_timeout)(timeout) } + /// Bound how long a write may block. `None` blocks indefinitely. + /// + /// A peer that stops reading fills the kernel send buffer, and an unbounded + /// write would then park the dispatch thread for good. + pub fn set_write_timeout(&self, timeout: Option) -> io::Result<()> { + (self.set_write_timeout)(timeout) + } + /// Tear the connection down so a reader blocked on it returns. Idempotent /// and best-effort: the peer may already be gone. pub fn close(&self) { @@ -79,12 +88,16 @@ impl Transport for UnixTransport { // Every clone shares one fd, so the control reaches the read half. let read = self.0.try_clone()?; let control = Arc::new(self.0.try_clone()?); + let writer_control = control.clone(); let closer = control.clone(); Ok(( Box::new(BufReader::new(read)), Box::new(self.0), Connection { set_read_timeout: Box::new(move |timeout| control.set_read_timeout(timeout)), + set_write_timeout: Box::new(move |timeout| { + writer_control.set_write_timeout(timeout) + }), close: Box::new(move || { let _ = closer.shutdown(std::net::Shutdown::Both); }), @@ -121,12 +134,16 @@ impl Transport for TcpTransport { fn split(self: Box) -> io::Result<(ReadHalf, WriteHalf, Connection)> { let read = self.0.try_clone()?; let control = Arc::new(self.0.try_clone()?); + let writer_control = control.clone(); let closer = control.clone(); Ok(( Box::new(BufReader::new(read)), Box::new(self.0), Connection { set_read_timeout: Box::new(move |timeout| control.set_read_timeout(timeout)), + set_write_timeout: Box::new(move |timeout| { + writer_control.set_write_timeout(timeout) + }), close: Box::new(move || { let _ = closer.shutdown(std::net::Shutdown::Both); }), @@ -188,6 +205,8 @@ impl Transport for MemoryTransport { *control.read_timeout.lock().unwrap_or_else(recover) = timeout; Ok(()) }), + // The in-memory buffer is unbounded, so a write never blocks. + set_write_timeout: Box::new(|_| Ok(())), close: Box::new(move || closer.close_reader()), }, ))