diff --git a/crates/taskito-core/src/lib.rs b/crates/taskito-core/src/lib.rs index 48ad81577..4ab05e5e7 100644 --- a/crates/taskito-core/src/lib.rs +++ b/crates/taskito-core/src/lib.rs @@ -46,7 +46,8 @@ pub use storage::Storage; pub use storage::StorageBackend; pub use storage::{DeadJob, QueueStats, SubscriptionBacklogStats}; pub use worker::{ - AttachError, AttachedExecutor, Capacity, ExecutorMessage, NativeDispatcher, ProtocolError, - RemoteConfig, RemoteDispatcher, SchedulerMessage, Secret, TaskError, TaskRegistry, TaskResult, - Transport, Worker, WorkerDispatcher, WorkerHandle, PROTOCOL_VERSION, + AttachAddress, AttachError, AttachedExecutor, Capacity, ExecutorClient, ExecutorConfig, + ExecutorError, ExecutorHandle, ExecutorMessage, ExecutorSession, NativeDispatcher, + ProtocolError, RemoteConfig, RemoteDispatcher, SchedulerMessage, Secret, TaskError, + TaskRegistry, TaskResult, Transport, Worker, WorkerDispatcher, WorkerHandle, PROTOCOL_VERSION, }; diff --git a/crates/taskito-core/src/worker/cancel.rs b/crates/taskito-core/src/worker/cancel.rs new file mode 100644 index 000000000..8f213e0c5 --- /dev/null +++ b/crates/taskito-core/src/worker/cancel.rs @@ -0,0 +1,145 @@ +//! How a dispatcher learns that a running job was cancelled. +//! +//! An in-process worker reads the storage flag `Storage::request_cancel` sets. +//! An attached executor has no storage — it is the whole point of #546 that it +//! carries no database credentials — and learns instead from the `cancel` frame +//! the scheduler sends, which arrives as +//! [`WorkerDispatcher::notify_cancel`](super::WorkerDispatcher::notify_cancel). +//! +//! Both sources answer one question, so both dispatchers ask it here rather +//! than growing two cancel paths each. + +use std::collections::HashSet; +use std::sync::{Mutex, PoisonError}; + +use crate::storage::{Storage, StorageBackend}; + +/// The cancel sources available to one dispatcher. +pub struct CancelSignals { + /// Present only for a dispatcher running inside a worker. + storage: Option, + /// Ids delivered out of band. Kept until the job reports, so a cancel that + /// races a job's start still fires rather than being missed. + signalled: Mutex>, +} + +impl CancelSignals { + /// Read cancels from storage, and from `notify_cancel` when it is called. + pub fn from_storage(storage: StorageBackend) -> Self { + Self { + storage: Some(storage), + signalled: Mutex::new(HashSet::new()), + } + } + + /// Read cancels only from `notify_cancel` — the attached-executor case. + pub fn detached() -> Self { + Self { + storage: None, + signalled: Mutex::new(HashSet::new()), + } + } + + /// Record a cancel request for `job_id`. + pub fn signal(&self, job_id: &str) { + self.lock().insert(job_id.to_string()); + } + + /// Whether `job_id` has been cancelled. + /// + /// The out-of-band set is checked first: it needs no I/O, and for a + /// detached executor it is the only answer there is. + pub fn is_cancelled(&self, job_id: &str) -> bool { + if self.lock().contains(job_id) { + return true; + } + self.storage + .as_ref() + .is_some_and(|storage| storage.is_cancel_requested(job_id).unwrap_or(false)) + } + + /// Drop the record for a finished job, so the set cannot grow for the life + /// of the process. + pub fn forget(&self, job_id: &str) { + self.lock().remove(job_id); + } + + fn lock(&self) -> std::sync::MutexGuard<'_, HashSet> { + // The state behind the lock is a plain set, so reading it stays safe + // even if a holder panicked. + self.signalled + .lock() + .unwrap_or_else(PoisonError::into_inner) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn a_detached_signal_is_observed_without_storage() { + let signals = CancelSignals::detached(); + assert!(!signals.is_cancelled("job-1")); + + signals.signal("job-1"); + assert!(signals.is_cancelled("job-1")); + } + + #[test] + fn forgetting_a_job_clears_its_signal() { + // Ids are held until the job reports, so they have to be released or + // the set grows for the life of the process. + let signals = CancelSignals::detached(); + signals.signal("job-1"); + signals.forget("job-1"); + assert!(!signals.is_cancelled("job-1")); + } + + #[test] + fn an_unknown_job_is_not_cancelled() { + let signals = CancelSignals::detached(); + signals.signal("job-1"); + assert!(!signals.is_cancelled("job-2")); + } + + #[test] + fn a_storage_flag_is_honoured_too() { + use crate::job::{now_millis, NewJob}; + use crate::storage::sqlite::SqliteStorage; + + let storage = StorageBackend::Sqlite(SqliteStorage::in_memory().expect("in-memory")); + let job = storage + .enqueue(NewJob { + queue: "default".to_string(), + task_name: "resize".to_string(), + payload: Vec::new(), + priority: 0, + scheduled_at: now_millis(), + max_retries: 0, + timeout_ms: 0, + unique_key: None, + metadata: None, + notes: None, + depends_on: vec![], + expires_at: None, + result_ttl_ms: None, + namespace: None, + }) + .expect("enqueue"); + + // `request_cancel` only flags a *running* job — a pending one is + // cancelled outright — so the job has to be dequeued first. + let running = storage + .dequeue("default", now_millis() + 1_000, None) + .expect("dequeue") + .expect("the enqueued job"); + assert_eq!(running.id, job.id); + + let signals = CancelSignals::from_storage(storage.clone()); + assert!(!signals.is_cancelled(&job.id)); + + assert!(storage.request_cancel(&job.id).expect("request cancel")); + assert!(signals.is_cancelled(&job.id)); + } +} diff --git a/crates/taskito-core/src/worker/dial.rs b/crates/taskito-core/src/worker/dial.rs new file mode 100644 index 000000000..126580e92 --- /dev/null +++ b/crates/taskito-core/src/worker/dial.rs @@ -0,0 +1,219 @@ +//! Dial the address an executor was pointed at. +//! +//! The listener parses the same grammar on the bind side +//! (`taskito-server`'s `config::listen`), so the two stay readable against each +//! other: whatever `TASKITO_LISTEN` accepts, `TASKITO_ATTACH` dials. Every SDK +//! shares this rather than reimplementing the grammar in its own language, +//! where `unix:` support would inevitably drift. + +use std::io; +use std::net::{TcpStream, ToSocketAddrs}; +#[cfg(unix)] +use std::os::unix::net::UnixStream; +use std::time::Duration; + +#[cfg(unix)] +use super::transport::UnixTransport; +use super::transport::{TcpTransport, Transport}; + +/// Where an executor attaches. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum AttachAddress { + /// TCP, for a scheduler in another container or host. + Tcp(String), + /// Unix domain socket, the same-pod sidecar case. + #[cfg(unix)] + Unix(std::path::PathBuf), +} + +impl std::fmt::Display for AttachAddress { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::Tcp(target) => write!(f, "tcp://{target}"), + #[cfg(unix)] + Self::Unix(path) => write!(f, "unix:{}", path.display()), + } + } +} + +impl AttachAddress { + /// Parse one attach spec: `unix:/path`, `tcp://host:port`, `host:port`, or + /// `:port`. + /// + /// A bare `:port` means loopback, matching the listener's reading of the + /// same ambiguous value. + pub fn parse(spec: &str) -> io::Result { + let spec = spec.trim(); + if spec.is_empty() { + return Err(invalid("an attach address must not be empty")); + } + + if let Some(path) = spec.strip_prefix("unix:") { + #[cfg(unix)] + { + if path.is_empty() { + return Err(invalid( + "a unix attach address needs a socket path, e.g. unix:/run/taskito.sock", + )); + } + return Ok(Self::Unix(std::path::PathBuf::from(path))); + } + #[cfg(not(unix))] + { + let _ = path; + return Err(invalid( + "unix socket attach addresses are not supported on this platform", + )); + } + } + + // The listener prints itself as `tcp://host:port`, so an operator who + // copies that line out of the logs must get a working address back. + let target = spec.strip_prefix("tcp://").unwrap_or(spec); + let target = match target.strip_prefix(':') { + Some(port) => format!("127.0.0.1:{port}"), + None => target.to_string(), + }; + if !target.contains(':') { + return Err(invalid(format!( + "'{spec}' has no port — an attach address looks like host:port or \ + unix:/run/taskito.sock" + ))); + } + Ok(Self::Tcp(target)) + } + + /// Open a connection to this address. + /// + /// `timeout` bounds the TCP connect so an unreachable scheduler fails + /// promptly instead of sitting in the platform's default retry window, + /// which can be minutes. + pub fn connect(&self, timeout: Duration) -> io::Result> { + match self { + Self::Tcp(target) => { + let addresses = target.to_socket_addrs().map_err(|error| { + invalid(format!("'{target}' is not a valid host:port: {error}")) + })?; + // Every resolved address is tried, not just the first: a + // dual-stack scheduler resolves to both an AAAA and an A + // record, and a host that cannot route one still reaches the + // other. The last failure is what gets reported. + let mut last_error = None; + for address in addresses { + match TcpStream::connect_timeout(&address, timeout) { + Ok(stream) => return Ok(Box::new(TcpTransport::new(stream)?)), + Err(error) => last_error = Some(error), + } + } + Err(last_error + .unwrap_or_else(|| invalid(format!("'{target}' resolved to no address")))) + } + #[cfg(unix)] + Self::Unix(path) => { + // No connect timeout exists for a Unix socket, and none is + // needed: the peer is on this host, so a connect either + // succeeds or fails at once. + let stream = UnixStream::connect(path)?; + Ok(Box::new(UnixTransport::new(stream))) + } + } + } +} + +fn invalid(message: impl Into) -> io::Error { + io::Error::new(io::ErrorKind::InvalidInput, message.into()) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn a_host_and_port_parses_as_tcp() { + assert_eq!( + AttachAddress::parse("scheduler:7749").expect("parse"), + AttachAddress::Tcp("scheduler:7749".to_string()) + ); + } + + #[test] + fn the_tcp_scheme_the_listener_prints_is_accepted() { + // The listener logs `attach listener on tcp://127.0.0.1:7749`; pasting + // that back must work. + assert_eq!( + AttachAddress::parse("tcp://127.0.0.1:7749").expect("parse"), + AttachAddress::Tcp("127.0.0.1:7749".to_string()) + ); + } + + #[test] + fn a_bare_port_means_loopback() { + assert_eq!( + AttachAddress::parse(":7749").expect("parse"), + AttachAddress::Tcp("127.0.0.1:7749".to_string()) + ); + } + + #[test] + fn surrounding_whitespace_is_ignored() { + // Shell heredocs and Kubernetes manifests both leak trailing newlines. + assert_eq!( + AttachAddress::parse(" 127.0.0.1:7749\n").expect("parse"), + AttachAddress::Tcp("127.0.0.1:7749".to_string()) + ); + } + + #[test] + fn an_address_without_a_port_is_rejected_with_the_shape_it_wanted() { + let error = AttachAddress::parse("scheduler").expect_err("must be rejected"); + assert_eq!(error.kind(), io::ErrorKind::InvalidInput); + assert!( + error.to_string().contains("host:port"), + "the message must show the expected shape: {error}" + ); + } + + #[test] + fn an_empty_address_is_rejected() { + assert!(AttachAddress::parse("").is_err()); + assert!(AttachAddress::parse(" ").is_err()); + } + + #[cfg(unix)] + #[test] + fn a_unix_path_parses_and_prints_back() { + let address = AttachAddress::parse("unix:/run/taskito.sock").expect("parse"); + assert_eq!( + address, + AttachAddress::Unix(std::path::PathBuf::from("/run/taskito.sock")) + ); + assert_eq!(address.to_string(), "unix:/run/taskito.sock"); + } + + #[cfg(unix)] + #[test] + fn a_unix_scheme_without_a_path_is_rejected() { + let error = AttachAddress::parse("unix:").expect_err("must be rejected"); + assert!(error.to_string().contains("socket path"), "{error}"); + } + + #[test] + fn connecting_to_a_closed_port_fails_rather_than_hanging() { + // Port 1 on loopback: reserved, and nothing listens there. + let address = AttachAddress::parse("127.0.0.1:1").expect("parse"); + assert!(address.connect(Duration::from_millis(500)).is_err()); + } + + #[test] + fn a_dialed_address_round_trips_through_a_real_listener() { + let listener = std::net::TcpListener::bind("127.0.0.1:0").expect("bind"); + let port = listener.local_addr().expect("addr").port(); + let accepting = std::thread::spawn(move || listener.accept().expect("accept")); + + let address = AttachAddress::parse(&format!(":{port}")).expect("parse"); + let transport = address.connect(Duration::from_secs(5)).expect("connect"); + assert!(transport.peer().starts_with("tcp:")); + + let _ = accepting.join(); + } +} diff --git a/crates/taskito-core/src/worker/executor.rs b/crates/taskito-core/src/worker/executor.rs new file mode 100644 index 000000000..c8288cccb --- /dev/null +++ b/crates/taskito-core/src/worker/executor.rs @@ -0,0 +1,728 @@ +//! The executor side of an attach: dial a scheduler, run its jobs locally. +//! +//! This is [`Worker`](super::runner::Worker) with storage swapped for a socket. +//! A worker pulls jobs from a [`Scheduler`](crate::scheduler::Scheduler) and +//! pushes results back into it; an executor pulls jobs from a [`FrameReader`] +//! and pushes results out through a [`FrameWriter`]. Everything between is the +//! same [`WorkerDispatcher`] every SDK already implements, so the prefork pool, +//! the Node dispatcher and the Java dispatcher all attach unchanged. +//! +//! Nothing here touches [`Storage`](crate::storage::Storage) — that is the +//! point. The executor image carries app code and no database credentials; the +//! scheduler image carries credentials and no app code. A job frame already +//! holds everything running a task needs. +//! +//! ```no_run +//! # use std::sync::Arc; +//! # use taskito_core::worker::{ExecutorClient, ExecutorConfig, TcpTransport, WorkerDispatcher}; +//! # fn run(dispatcher: Arc) -> Result<(), Box> { +//! let stream = std::net::TcpStream::connect("scheduler:7749")?; +//! let client = ExecutorClient::connect( +//! Box::new(TcpTransport::new(stream)?), +//! ExecutorConfig { +//! tasks: vec!["resize".to_string()], +//! slots: 4, +//! ..ExecutorConfig::new("python", "0.21.0") +//! }, +//! )?; +//! client.spawn(dispatcher).wait(); // until the scheduler ends the session +//! # Ok(()) +//! # } +//! ``` + +use std::sync::atomic::{AtomicBool, AtomicU32, Ordering}; +use std::sync::{Arc, Mutex, PoisonError}; +use std::thread::{self, JoinHandle}; +use std::time::{Duration, Instant}; + +use crossbeam_channel::{Receiver, RecvTimeoutError, Sender}; +use tokio::sync::mpsc::error::TrySendError; + +use super::auth::Secret; +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; + +/// How often a waiting loop wakes to re-check its condition. +const POLL: Duration = Duration::from_millis(20); + +/// Channel for jobs on their way to the pool. +type JobSender = tokio::sync::mpsc::Sender; + +/// Tuning for an [`ExecutorClient`]. +#[derive(Debug, Clone)] +pub struct ExecutorConfig { + /// Stable identity announced in `hello`. Two executors attached under one + /// id is an error on the scheduler, so this must be unique per process. + pub executor_id: String, + /// SDK this executor is built on, e.g. `"python"`. + pub sdk: String, + /// SDK version string, for the scheduler's inventory and logs. + pub version: String, + /// Tasks this executor has handlers for. The scheduler sends it nothing + /// else, so a name missing here is a job that never arrives. + pub tasks: Vec, + /// Jobs this executor can run concurrently. + pub slots: u32, + /// Shared secret, when the scheduler requires one. + pub token: Option, + /// How long the handshake may take before the attach is abandoned. + pub handshake_timeout: Duration, + /// How long a frame write may block before the scheduler is treated as + /// wedged. Bounds the result loop against a peer that stops reading. + pub write_timeout: Duration, + /// How often to send a liveness heartbeat. + pub heartbeat_interval: Duration, + /// How long a drain waits for in-flight jobs to finish and their results to + /// reach the scheduler before the connection is closed anyway. + pub shutdown_drain: Duration, +} + +impl ExecutorConfig { + /// Defaults for `sdk`/`version`, with a generated executor id. + /// + /// `tasks` is empty and `slots` is 1 — both are the caller's to set, and an + /// executor advertising no tasks is deliberately inert rather than a peer + /// that quietly receives everything. + pub fn new(sdk: impl Into, version: impl Into) -> Self { + let sdk = sdk.into(); + Self { + executor_id: format!("{sdk}-executor-{}", uuid::Uuid::now_v7()), + sdk, + version: version.into(), + tasks: Vec::new(), + slots: 1, + token: None, + handshake_timeout: Duration::from_secs(10), + write_timeout: Duration::from_secs(30), + heartbeat_interval: Duration::from_secs(5), + shutdown_drain: Duration::from_secs(30), + } + } +} + +/// Why an executor could not attach. +#[derive(Debug, thiserror::Error)] +pub enum ExecutorError { + /// The transport could not be dialled, split, or configured. + #[error("attach transport failed: {0}")] + Transport(#[from] std::io::Error), + + /// A frame was malformed, or the peer announced a version we do not speak. + #[error(transparent)] + Protocol(#[from] ProtocolError), + + /// The scheduler closed the connection instead of acknowledging the + /// handshake. + /// + /// A refused peer never receives an ack, so this is what a rejected + /// credential looks like from this side. Named rather than surfaced as a + /// bare I/O error because a wrong or missing attach token is the likeliest + /// deployment mistake, and "connection reset" would send the operator + /// looking at the network instead. + #[error("the scheduler refused the attach (check the attach token)")] + Refused, +} + +/// A completed handshake, not yet running. +/// +/// Split from [`ExecutorClient::spawn`] so a caller can report a failed attach +/// — a bad token, an unreachable scheduler — before it builds an execution pool +/// it would only tear down again. +pub struct ExecutorClient { + config: ExecutorConfig, + scheduler_id: String, + peer: String, + reader: FrameReader, + link: Link, +} + +impl ExecutorClient { + /// Dial, handshake, and register with the scheduler on the far end. + /// + /// Writes `hello`, then requires `hello_ack` before anything else. The read + /// is bounded by `handshake_timeout`; the bound is cleared afterwards, or an + /// executor idling longer than the budget between jobs would tear itself + /// down. + pub fn connect( + transport: Box, + config: ExecutorConfig, + ) -> Result { + let peer = transport.peer(); + let (read, write, connection) = transport.split()?; + connection.set_write_timeout(Some(config.write_timeout))?; + connection.set_read_timeout(Some(config.handshake_timeout))?; + + let mut reader = FrameReader::new(read); + let mut writer = FrameWriter::new(write); + + writer.write_header(&ExecutorMessage::Hello { + executor_id: config.executor_id.clone(), + sdk: config.sdk.clone(), + version: config.version.clone(), + tasks: config.tasks.clone(), + slots: config.slots, + protocol_version: PROTOCOL_VERSION, + token: config.token.clone(), + })?; + + let scheduler_id = read_ack(&mut reader)?; + connection.set_read_timeout(None)?; + + log::info!( + "[taskito] executor {} attached to scheduler {scheduler_id} at {peer} with {} slot(s)", + config.executor_id, + config.slots + ); + + Ok(Self { + scheduler_id, + peer, + reader, + link: Link { + writer: Mutex::new(writer), + connection, + }, + config, + }) + } + + /// Identity the scheduler announced in its `hello_ack`. + pub fn scheduler_id(&self) -> &str { + &self.scheduler_id + } + + /// Peer label of the scheduler connection, for logs. + pub fn peer(&self) -> &str { + &self.peer + } + + /// Start running jobs on `dispatcher`. + /// + /// Returns immediately; the returned handle is how the caller waits for the + /// scheduler to end the session, or asks for a drain of its own. + pub fn spawn(self, dispatcher: Arc) -> ExecutorHandle { + let Self { + config, + reader, + link, + .. + } = self; + + // Sized off the slot count for the same reason `Worker::spawn` does: + // enough buffer that a free slot is never left idle waiting on the + // channel, small enough that the executor is not hoarding jobs it has + // no capacity to run. + let capacity = (config.slots as usize).max(1) * 2; + let (job_tx, job_rx) = tokio::sync::mpsc::channel(capacity); + let (result_tx, result_rx) = crossbeam_channel::bounded(capacity); + + let shared = Arc::new(Shared { + link, + executor_id: config.executor_id, + slots: config.slots, + free_slots: AtomicU32::new(config.slots), + in_flight: AtomicU32::new(0), + draining: AtomicBool::new(false), + session_over: AtomicBool::new(false), + results_flushed: AtomicBool::new(false), + job_tx: Mutex::new(Some(job_tx)), + }); + + let threads = vec![ + spawn_runtime(dispatcher.clone(), job_rx, result_tx), + spawn_result_loop(shared.clone(), result_rx), + spawn_heartbeat(shared.clone(), config.heartbeat_interval), + spawn_reader(shared.clone(), reader, dispatcher.clone()), + ]; + + ExecutorHandle { + shared, + dispatcher, + shutdown_drain: config.shutdown_drain, + threads, + } + } +} + +/// Read the frame that completes the handshake. +fn read_ack(reader: &mut FrameReader) -> Result { + match reader.read::() { + Ok(( + SchedulerMessage::HelloAck { + scheduler_id, + protocol_version, + }, + _, + )) => { + if protocol_version != PROTOCOL_VERSION { + return Err(ProtocolError::VersionMismatch { + ours: PROTOCOL_VERSION, + theirs: protocol_version, + } + .into()); + } + Ok(scheduler_id) + } + Ok(_) => Err(ProtocolError::UnexpectedFrame { + expected: "hello_ack", + } + .into()), + // A refused peer is closed on without an ack, so a clean EOF here is a + // rejection rather than a transport fault. + Err(ProtocolError::Eof) => Err(ExecutorError::Refused), + Err(error) => Err(error.into()), + } +} + +/// Handle to a running executor. +pub struct ExecutorHandle { + shared: Arc, + dispatcher: Arc, + shutdown_drain: Duration, + threads: Vec>, +} + +/// A cheap, cloneable view of whether an executor's session is still open. +/// +/// [`ExecutorHandle::wait`] consumes the handle, which a shell that needs to +/// observe the session from elsewhere — an async runtime resolving a promise, +/// say — cannot do while it also holds the handle to shut it down. +#[derive(Clone)] +pub struct ExecutorSession { + shared: Arc, +} + +impl ExecutorSession { + /// Whether this executor is still accepting work. + /// + /// False once the session ends *or* a local drain starts: a caller parked + /// in [`ExecutorSession::wait`] has to be released by its own `stop()` too, + /// and after a drain there is nothing left to wait for. + pub fn is_running(&self) -> bool { + !self.shared.session_over.load(Ordering::Acquire) + } + + /// Block until this executor stops accepting work. Does not drain or join — + /// that is [`ExecutorHandle::shutdown`]'s job. + pub fn wait(&self) { + while self.is_running() { + thread::sleep(POLL); + } + } +} + +impl ExecutorHandle { + /// Id this executor attached under. + pub fn executor_id(&self) -> &str { + &self.shared.executor_id + } + + /// A view another thread can watch the session through. + pub fn session(&self) -> ExecutorSession { + ExecutorSession { + shared: self.shared.clone(), + } + } + + /// Whether this executor is still accepting work. + pub fn is_running(&self) -> bool { + !self.shared.session_over.load(Ordering::Acquire) + } + + /// Block until this executor stops accepting work, then drain and join. + pub fn wait(self) { + while self.is_running() { + thread::sleep(POLL); + } + self.teardown(); + } + + /// Block for at most `timeout`, returning whether the session has ended. + /// + /// The bounded form exists for shells whose signal handling needs the + /// calling thread back periodically — a Python `SIGTERM` handler only runs + /// when the main thread reacquires the GIL, which it cannot do while parked + /// in [`ExecutorHandle::wait`]. + pub fn wait_timeout(&self, timeout: Duration) -> bool { + let deadline = Instant::now() + timeout; + while self.is_running() { + let remaining = deadline.saturating_duration_since(Instant::now()); + if remaining.is_zero() { + return false; + } + thread::sleep(POLL.min(remaining)); + } + true + } + + /// Ask the scheduler to stop sending work and finish what is in flight. + /// + /// Returns without waiting, so it is safe on a signal-handling path. + /// Idempotent. + pub fn stop(&self) { + self.shared.begin_drain(); + self.dispatcher.shutdown(); + } + + /// Drain, disconnect, and join every thread. + pub fn shutdown(self) { + self.stop(); + self.teardown(); + } + + /// Wait out the drain, close the connection, and join. + /// + /// Closing is what unparks the reader, which is otherwise blocked on a read + /// only a scheduler frame would satisfy. It must not happen before results + /// have flushed, or a finished job's outcome is lost and the job waits for + /// the reaper instead of being recorded. + fn teardown(mut self) { + self.stop(); + + let deadline = Instant::now() + self.shutdown_drain; + while Instant::now() < deadline && !self.shared.results_flushed.load(Ordering::Acquire) { + thread::sleep(POLL); + } + + let stranded = self.shared.in_flight.load(Ordering::Relaxed); + if stranded > 0 { + log::warn!( + "[taskito] executor {} did not drain {stranded} job(s) within the shutdown \ + budget; disconnecting — they will be reaped", + self.shared.executor_id + ); + } + + self.shared.link.connection.close(); + // Joining is only safe once the results are out. A task that ignores + // its cancel keeps the pool's `run` from returning, so its thread never + // exits — joining it would hang the very shutdown it was asked to + // bound. Those threads are dropped instead, and the process is free to + // exit. + if self.shared.results_flushed.load(Ordering::Acquire) { + for thread in self.threads.drain(..) { + if thread.join().is_err() { + log::error!( + "[taskito] an executor thread panicked during shutdown of {}", + self.shared.executor_id + ); + } + } + } + log::info!("[taskito] executor {} detached", self.shared.executor_id); + } +} + +/// The scheduler connection: one writer, shared, plus its lifetime controls. +struct Link { + /// Behind a lock because the result loop and the heartbeat both write. + writer: Mutex>, + connection: Connection, +} + +/// State every thread of a running executor shares. +struct Shared { + link: Link, + executor_id: String, + /// Concurrency announced at handshake; the ceiling `free_slots` counts down + /// from. + slots: u32, + /// Slots not currently occupied, published on each heartbeat. + free_slots: AtomicU32, + /// Jobs handed to the pool and not yet answered. + in_flight: AtomicU32, + /// Set once no further jobs will be accepted. + draining: AtomicBool, + /// Set when this executor stops accepting work — the reader's conversation + /// ending, or a local drain. + session_over: AtomicBool, + /// Set once every result the pool produced has been written. + results_flushed: AtomicBool, + /// Dropped by `begin_drain`, which is what lets the pool's `run` return + /// once it has finished the jobs it already holds. Held in an `Option` so a + /// local shutdown can release it without waiting on the parked reader. + job_tx: Mutex>, +} + +impl Shared { + /// Send one frame, logging rather than propagating: every caller is a + /// worker thread whose only remedy is to stop, which the reader's own EOF + /// already handles. + fn send(&self, frame: &ExecutorMessage, payload: &[u8]) -> bool { + let sent = self + .link + .writer + .lock() + .unwrap_or_else(recover) + .write(frame, payload); + match sent { + Ok(()) => true, + Err(error) => { + log::warn!( + "[taskito] executor {} failed to send a frame: {error}", + self.executor_id + ); + false + } + } + } + + /// Stop accepting work, announcing it in-protocol first. + /// + /// A heartbeat may only *shrink* the scheduler's view of capacity + /// (`remote.rs`), so zeroing it is a standing "send me nothing more" that + /// needs no new frame type. Dropping the job sender then lets the pool + /// finish what it holds and return. Idempotent — a local signal and a + /// `shutdown` frame can both arrive. + fn begin_drain(&self) { + if self.draining.swap(true, Ordering::AcqRel) { + return; + } + self.free_slots.store(0, Ordering::Relaxed); + self.job_tx.lock().unwrap_or_else(recover).take(); + // Release anyone parked in `wait`: the reader is still blocked on a read + // only the scheduler could satisfy, so nothing else would wake them. + self.session_over.store(true, Ordering::Release); + // Announced last, because a scheduler that stopped reading blocks this + // write for up to `write_timeout`. `stop` is documented as safe on a + // signal-handling path, so the local drain must not wait on the peer. + self.send(&ExecutorMessage::Heartbeat { free_slots: 0 }, &[]); + log::info!( + "[taskito] executor {} draining; no further jobs will be accepted", + self.executor_id + ); + } + + /// A sender for the pool, or `None` once draining. + fn job_sender(&self) -> Option { + self.job_tx.lock().unwrap_or_else(recover).clone() + } + + fn job_started(&self) { + self.in_flight.fetch_add(1, Ordering::Relaxed); + self.publish_free(); + } + + /// Saturating, so a result for a job this executor never counted cannot + /// wrap the counter and strand the drain forever. + fn job_finished(&self) { + let _ = self + .in_flight + .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |running| { + Some(running.saturating_sub(1)) + }); + self.publish_free(); + } + + /// Recompute free capacity from the in-flight count, unless draining — a + /// drain has pinned it to zero, and a job finishing must not undo that. + fn publish_free(&self) { + if self.draining.load(Ordering::Acquire) { + return; + } + let running = self.in_flight.load(Ordering::Relaxed); + self.free_slots + .store(self.slots.saturating_sub(running), Ordering::Relaxed); + } +} + +/// Recover a guard from a poisoned lock instead of cascading the panic. The +/// state behind it is a frame writer, which stays usable. +fn recover(poisoned: PoisonError) -> T { + poisoned.into_inner() +} + +/// Runtime thread: drives the pool, exactly as `Worker::spawn` does. +/// +/// `result_tx` moves in, so the result loop sees a disconnect the moment +/// execution is finished. +fn spawn_runtime( + dispatcher: Arc, + job_rx: tokio::sync::mpsc::Receiver, + result_tx: Sender, +) -> JoinHandle<()> { + thread::Builder::new() + .name("taskito-executor-runtime".to_string()) + .spawn(move || { + let runtime = tokio::runtime::Builder::new_multi_thread() + .worker_threads(2) + .enable_all() + .build() + .expect("tokio runtime construction cannot fail with these settings"); + runtime.block_on(async move { dispatcher.run(job_rx, result_tx).await }); + }) + .expect("spawning the executor runtime thread cannot fail with a valid name") +} + +/// Reader thread: the scheduler's half of the conversation. +fn spawn_reader( + shared: Arc, + mut reader: FrameReader, + dispatcher: Arc, +) -> JoinHandle<()> { + thread::Builder::new() + .name("taskito-executor-reader".to_string()) + .spawn(move || { + loop { + match reader.read::() { + Ok((SchedulerMessage::Shutdown, _)) => { + log::info!( + "[taskito] scheduler asked executor {} to shut down", + shared.executor_id + ); + break; + } + Ok((SchedulerMessage::Cancel { job_id }, _)) => { + dispatcher.notify_cancel(&job_id); + } + Ok((frame, payload)) => accept_job(&shared, frame, payload), + Err(ProtocolError::Eof) => { + log::info!( + "[taskito] scheduler closed the connection to executor {}", + shared.executor_id + ); + break; + } + // A closed connection during teardown lands here; the + // session is over either way. + Err(error) => { + if !shared.draining.load(Ordering::Acquire) { + log::warn!( + "[taskito] executor {} read error: {error}", + shared.executor_id + ); + } + break; + } + } + } + + // Whatever ended the loop, no more jobs are coming. + shared.begin_drain(); + dispatcher.shutdown(); + shared.session_over.store(true, Ordering::Release); + }) + .expect("spawning the executor reader thread cannot fail with a valid name") +} + +/// Hand one dispatched job to the pool. +/// +/// A job that cannot be run — the executor is draining, or the pool has already +/// stopped — is answered with a retryable failure rather than dropped. The +/// scheduler's reaper would recover it either way, but only after a reap cycle, +/// and the race is expected: a `job` already in flight when the zero-capacity +/// heartbeat lands is normal, not a fault. +fn accept_job(shared: &Arc, frame: SchedulerMessage, payload: Vec) { + let Some(job) = frame.into_job(payload) else { + // `hello_ack` is the only frame left, and it is handshake-only. + log::warn!( + "[taskito] executor {} received a handshake frame mid-stream", + shared.executor_id + ); + return; + }; + + let Some(sender) = shared.job_sender() else { + decline(shared, &job, "the executor is draining"); + return; + }; + + shared.job_started(); + // `try_send` rather than a blocking send: the reader also carries cancels, + // and parking it on a full channel would stall them behind the very jobs + // they target. The channel holds twice the advertised slots and the + // scheduler reserves a slot before dispatching, so a correct peer cannot + // fill it. + match sender.try_send(job) { + Ok(()) => {} + Err(TrySendError::Full(job)) => { + shared.job_finished(); + decline(shared, &job, "the executor pool is saturated"); + } + Err(TrySendError::Closed(job)) => { + shared.job_finished(); + decline(shared, &job, "the executor pool has stopped"); + } + } +} + +/// Answer a job this executor will not run with a retryable failure, so the +/// scheduler reschedules it now instead of waiting for a reap. +fn decline(shared: &Arc, job: &Job, reason: &str) { + log::warn!( + "[taskito] executor {} declining job {}: {reason}", + shared.executor_id, + job.id + ); + let (frame, payload) = ExecutorMessage::from_job_result(JobResult::Failure { + job_id: job.id.clone(), + error: format!("executor did not run '{}': {reason}", job.task_name), + 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, + }); + shared.send(&frame, &payload); +} + +/// Result thread: every outcome the pool produces, framed back to the +/// scheduler. +/// +/// Ends when the pool has dropped its sender and the queue is empty, which is +/// what releases the teardown — a result written after the socket closed would +/// be a job silently lost. +fn spawn_result_loop(shared: Arc, result_rx: Receiver) -> JoinHandle<()> { + thread::Builder::new() + .name("taskito-executor-results".to_string()) + .spawn(move || { + loop { + match result_rx.recv_timeout(POLL) { + Ok(result) => { + let (frame, payload) = ExecutorMessage::from_job_result(result); + let sent = shared.send(&frame, &payload); + shared.job_finished(); + if !sent { + break; + } + } + Err(RecvTimeoutError::Timeout) => {} + Err(RecvTimeoutError::Disconnected) => break, + } + } + shared.results_flushed.store(true, Ordering::Release); + }) + .expect("spawning the executor result thread cannot fail with a valid name") +} + +/// Heartbeat thread: liveness plus current free capacity. +/// +/// Stops once draining — the zero-capacity heartbeat `begin_drain` sent is the +/// last thing the scheduler needs to hear, and repeating it would only contend +/// for the writer while results are trying to flush. +fn spawn_heartbeat(shared: Arc, interval: Duration) -> JoinHandle<()> { + thread::Builder::new() + .name("taskito-executor-heartbeat".to_string()) + .spawn(move || loop { + // Slept in slices so a drain is observed promptly rather than one + // full interval late. + let deadline = Instant::now() + interval; + while Instant::now() < deadline { + if shared.draining.load(Ordering::Acquire) { + return; + } + thread::sleep(POLL.min(deadline.saturating_duration_since(Instant::now()))); + } + let free_slots = shared.free_slots.load(Ordering::Relaxed); + if !shared.send(&ExecutorMessage::Heartbeat { free_slots }, &[]) { + return; + } + }) + .expect("spawning the executor heartbeat thread cannot fail with a valid name") +} diff --git a/crates/taskito-core/src/worker/mod.rs b/crates/taskito-core/src/worker/mod.rs index 386f6161a..4e258013f 100644 --- a/crates/taskito-core/src/worker/mod.rs +++ b/crates/taskito-core/src/worker/mod.rs @@ -1,5 +1,8 @@ pub mod auth; +pub mod cancel; +pub mod dial; pub mod dispatcher; +pub mod executor; pub mod protocol; pub mod registry; pub mod remote; @@ -7,7 +10,12 @@ pub mod runner; pub mod transport; pub use auth::Secret; +pub use cancel::CancelSignals; +pub use dial::AttachAddress; pub use dispatcher::NativeDispatcher; +pub use executor::{ + ExecutorClient, ExecutorConfig, ExecutorError, ExecutorHandle, ExecutorSession, +}; pub use protocol::{ExecutorMessage, ProtocolError, SchedulerMessage, PROTOCOL_VERSION}; pub use registry::{TaskError, TaskHandler, TaskRegistry, TaskResult}; pub use remote::{AttachError, AttachedExecutor, Capacity, RemoteConfig, RemoteDispatcher}; @@ -42,8 +50,3 @@ pub trait WorkerDispatcher: Send + Sync { /// storage. fn notify_cancel(&self, _job_id: &str) {} } - -#[cfg(test)] -mod remote_tests; -#[cfg(test)] -mod tests; diff --git a/crates/taskito-core/src/worker/protocol.rs b/crates/taskito-core/src/worker/protocol.rs index c770beb8e..fdc8e8ac5 100644 --- a/crates/taskito-core/src/worker/protocol.rs +++ b/crates/taskito-core/src/worker/protocol.rs @@ -19,7 +19,7 @@ use serde::de::DeserializeOwned; use serde::{Deserialize, Serialize}; use super::auth::Secret; -use crate::job::Job; +use crate::job::{Job, JobStatus}; use crate::scheduler::JobResult; /// Frame format version. Both sides announce it in the handshake; a mismatch @@ -241,7 +241,132 @@ impl From<&Job> for SchedulerMessage { } } +impl SchedulerMessage { + /// Rebuild the [`Job`] a dispatch frame describes. `None` for control + /// frames (`hello_ack`, `cancel`, `shutdown`). + /// + /// The inverse of [`SchedulerMessage::from`]. A frame carries only what + /// running a task needs, so the columns an executor never reads — timing, + /// dedup key, archived result — take their defaults rather than being put + /// on the wire. `status` is [`JobStatus::Running`] because that is what the + /// job is by the time a frame describing it has been dispatched. + /// + /// Those defaults are not purely internal: the Node and Python SDKs build + /// their handler-visible job objects straight from this `Job`, so a task + /// reading `created_at`, `scheduled_at`, `priority`, `metadata`, + /// `unique_key` or `notes` sees zeros and nulls on an attached executor + /// where an in-process worker would show the stored values. Carrying them + /// would mean widening the frame for fields no dispatch decision uses, so + /// the difference is documented rather than papered over — see the + /// `detached` module in each SDK for the other side of the same trade. + pub fn into_job(self, payload: Vec) -> Option { + match self { + Self::HelloAck { .. } | Self::Cancel { .. } | Self::Shutdown => None, + Self::Job { + id, + task_name, + retry_count, + max_retries, + queue, + timeout_ms, + namespace, + payload_len: _, + } => Some(Job { + id, + queue, + task_name, + payload, + status: JobStatus::Running, + priority: 0, + created_at: 0, + scheduled_at: 0, + started_at: None, + completed_at: None, + retry_count, + max_retries, + result: None, + error: None, + timeout_ms, + unique_key: None, + progress: None, + metadata: None, + notes: None, + cancel_requested: false, + expires_at: None, + result_ttl_ms: None, + namespace, + has_deps: false, + }), + } + } +} + impl ExecutorMessage { + /// Build the result frame and payload for a finished job. + /// + /// The inverse of [`ExecutorMessage::into_job_result`]. A success carries + /// its serialized result as the frame's blob, so the payload is returned + /// alongside the header rather than inside it. + pub fn from_job_result(result: JobResult) -> (Self, Vec) { + match result { + JobResult::Success { + job_id, + result, + task_name, + wall_time_ns, + } => { + // `Some(vec![])` is an empty result and `None` is no result, so + // the length is read off the `Option` itself — deriving it from + // the flattened payload would collapse the two. + let result_len = result.as_ref().map(Vec::len); + let payload = result.unwrap_or_default(); + ( + Self::Success { + job_id, + result_len, + task_name, + wall_time_ns, + }, + payload, + ) + } + JobResult::Failure { + job_id, + error, + retry_count, + max_retries, + task_name, + wall_time_ns, + should_retry, + timed_out, + } => ( + Self::Failure { + job_id, + error, + retry_count, + max_retries, + task_name, + wall_time_ns, + should_retry, + timed_out, + }, + Vec::new(), + ), + JobResult::Cancelled { + job_id, + task_name, + wall_time_ns, + } => ( + Self::Cancelled { + job_id, + task_name, + wall_time_ns, + }, + Vec::new(), + ), + } + } + /// 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 { @@ -479,6 +604,138 @@ mod tests { } } + #[test] + fn a_job_survives_a_round_trip_through_a_frame() { + let payload = [0x02, 0x82, 0x82, 0x01, 0x61, 0x61, 0xa0]; + let original = sample_job(&payload); + + let (frame, read_payload) = round_trip(&SchedulerMessage::from(&original), &payload); + let rebuilt = frame.into_job(read_payload).expect("a job frame"); + + // Everything a task body can observe. The columns left out of the frame + // are storage bookkeeping the executor never reads. + assert_eq!(rebuilt.id, original.id); + assert_eq!(rebuilt.queue, original.queue); + assert_eq!(rebuilt.task_name, original.task_name); + assert_eq!(rebuilt.payload, original.payload); + assert_eq!(rebuilt.retry_count, original.retry_count); + assert_eq!(rebuilt.max_retries, original.max_retries); + assert_eq!(rebuilt.timeout_ms, original.timeout_ms); + assert_eq!(rebuilt.namespace, original.namespace); + assert_eq!(rebuilt.status, JobStatus::Running); + } + + #[test] + fn control_frames_describe_no_job() { + assert!(SchedulerMessage::Shutdown.into_job(vec![]).is_none()); + assert!(SchedulerMessage::Cancel { + job_id: "job-1".into() + } + .into_job(vec![]) + .is_none()); + assert!(SchedulerMessage::HelloAck { + scheduler_id: "s".into(), + protocol_version: PROTOCOL_VERSION, + } + .into_job(vec![]) + .is_none()); + } + + #[test] + fn a_result_survives_a_round_trip_through_a_frame() { + for (label, original) in [ + ( + "a result", + JobResult::Success { + job_id: "job-1".into(), + result: Some(b"out".to_vec()), + task_name: "resize".into(), + wall_time_ns: 42, + }, + ), + ( + "an empty result", + JobResult::Success { + job_id: "job-1".into(), + result: Some(Vec::new()), + task_name: "resize".into(), + wall_time_ns: 42, + }, + ), + ( + "no result", + JobResult::Success { + job_id: "job-1".into(), + result: None, + task_name: "resize".into(), + wall_time_ns: 42, + }, + ), + ] { + let JobResult::Success { result: before, .. } = &original else { + unreachable!("the table holds successes only") + }; + let expected = before.clone(); + + let (frame, payload) = ExecutorMessage::from_job_result(original); + let (frame, payload) = round_trip(&frame, &payload); + + match frame.into_job_result(payload) { + Some(JobResult::Success { result, .. }) => { + assert_eq!(result, expected, "{label} must survive the round trip") + } + _ => panic!("expected a success for {label}"), + } + } + } + + #[test] + fn a_failure_round_trips_with_its_verdict_intact() { + let (frame, payload) = ExecutorMessage::from_job_result(JobResult::Failure { + job_id: "job-1".into(), + error: "boom".into(), + retry_count: 2, + max_retries: 5, + task_name: "resize".into(), + wall_time_ns: 7, + should_retry: false, + timed_out: true, + }); + let (frame, payload) = round_trip(&frame, &payload); + + match frame.into_job_result(payload) { + Some(JobResult::Failure { + error, + retry_count, + max_retries, + should_retry, + timed_out, + .. + }) => { + assert_eq!(error, "boom"); + assert_eq!(retry_count, 2); + assert_eq!(max_retries, 5); + assert!(!should_retry, "only the executor can judge retryability"); + assert!(timed_out); + } + _ => panic!("expected a failure"), + } + } + + #[test] + fn a_cancellation_round_trips() { + let (frame, payload) = ExecutorMessage::from_job_result(JobResult::Cancelled { + job_id: "job-1".into(), + task_name: "resize".into(), + wall_time_ns: 9, + }); + let (frame, payload) = round_trip(&frame, &payload); + assert!(matches!( + frame.into_job_result(payload), + Some(JobResult::Cancelled { job_id, .. }) if job_id == "job-1" + )); + } + #[test] fn handshake_frames_round_trip() { let (hello, payload) = round_trip( diff --git a/crates/taskito-core/tests/rust.rs b/crates/taskito-core/tests/rust.rs index c846075d1..806316134 100644 --- a/crates/taskito-core/tests/rust.rs +++ b/crates/taskito-core/tests/rust.rs @@ -1,3 +1,6 @@ mod rust { + mod executor_tests; + mod remote_tests; mod storage_tests; + mod worker_tests; } diff --git a/crates/taskito-core/tests/rust/executor_tests.rs b/crates/taskito-core/tests/rust/executor_tests.rs new file mode 100644 index 000000000..ada157a62 --- /dev/null +++ b/crates/taskito-core/tests/rust/executor_tests.rs @@ -0,0 +1,988 @@ +//! Tests for [`ExecutorClient`], driven against a real [`RemoteDispatcher`] +//! over [`MemoryTransport`] so no socket is bound. +//! +//! Both halves of the attach are the shipping implementation — a fake on either +//! side could only prove it agrees with itself. + +use std::collections::HashMap; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::{Arc, Mutex}; +use std::thread; +use std::time::{Duration, Instant}; + +use async_trait::async_trait; +use crossbeam_channel::{Receiver, Sender}; + +use taskito_core::job::{Job, JobStatus}; +use taskito_core::scheduler::JobResult; +use taskito_core::worker::auth::Secret; +use taskito_core::worker::executor::{ + ExecutorClient, ExecutorConfig, ExecutorError, ExecutorHandle, +}; +use taskito_core::worker::protocol::{ + ExecutorMessage, FrameReader, FrameWriter, ProtocolError, SchedulerMessage, PROTOCOL_VERSION, +}; +use taskito_core::worker::remote::{RemoteConfig, RemoteDispatcher}; +use taskito_core::worker::transport::{MemoryTransport, ReadHalf, Transport, WriteHalf}; +use taskito_core::worker::WorkerDispatcher; + +const SETTLE: Duration = Duration::from_secs(5); + +/// The drain budget every attached executor in these tests runs with. +const SHUTDOWN_DRAIN: Duration = Duration::from_secs(2); + +/// What a [`TestPool`] should do with a job. +enum Behaviour { + Succeed(Option>), + Fail { + should_retry: bool, + }, + /// Park until released, so a test can hold a job in flight. + Block(Receiver<()>), + /// Park until the test drops its sender. Unlike [`Behaviour::Block`] no + /// timeout releases it, so a shutdown that waits on the job never returns — + /// which is what makes the drain budget observable. + Wedge(Receiver<()>), +} + +/// A minimal [`WorkerDispatcher`]: one job at a time, scripted per task name. +/// +/// The SDK pools are the real consumers, but each needs a language runtime. +/// This stands in for them and, unlike [`NativeDispatcher`](super::NativeDispatcher), +/// records `notify_cancel` so the cancel path is observable. +struct TestPool { + behaviours: Mutex>, + /// Every job the pool was handed, as rebuilt from the wire. + seen: Mutex>, + cancels: Mutex>, + started: Sender, + shutdown: AtomicBool, +} + +impl TestPool { + fn new(started: Sender) -> Arc { + Arc::new(Self { + behaviours: Mutex::new(HashMap::new()), + seen: Mutex::new(Vec::new()), + cancels: Mutex::new(Vec::new()), + started, + shutdown: AtomicBool::new(false), + }) + } + + /// Script `task_name`. Consumed on first use, so a repeat of the same task + /// falls through to the default success. + fn on(&self, task_name: &str, behaviour: Behaviour) { + self.behaviours + .lock() + .expect("behaviours") + .insert(task_name.to_string(), behaviour); + } + + fn cancelled(&self) -> Vec { + self.cancels.lock().expect("cancels").clone() + } + + /// The job as the executor rebuilt it from the frame. + fn received(&self, job_id: &str) -> Option { + self.seen + .lock() + .expect("seen") + .iter() + .find(|job| job.id == job_id) + .cloned() + } + + /// Run one job, blocking the pool loop for as long as the task would. + fn execute(&self, job: &Job) -> JobResult { + self.seen.lock().expect("seen").push(job.clone()); + let _ = self.started.send(job.id.clone()); + + // Taken out of the map rather than read under it: a blocking task holds + // its release channel for the whole wait. + let behaviour = self + .behaviours + .lock() + .expect("behaviours") + .remove(&job.task_name); + + match behaviour { + Some(Behaviour::Block(release)) => { + let _ = release.recv_timeout(SETTLE); + success(job, None) + } + Some(Behaviour::Wedge(release)) => { + let _ = release.recv(); + success(job, None) + } + Some(Behaviour::Fail { should_retry }) => JobResult::Failure { + job_id: job.id.clone(), + error: "deliberate failure".to_string(), + retry_count: job.retry_count, + max_retries: job.max_retries, + task_name: job.task_name.clone(), + wall_time_ns: 1, + should_retry, + timed_out: false, + }, + Some(Behaviour::Succeed(result)) => success(job, result), + // Unscripted tasks succeed: most tests care about the transport. + None => success(job, None), + } + } +} + +fn success(job: &Job, result: Option>) -> JobResult { + JobResult::Success { + job_id: job.id.clone(), + result, + task_name: job.task_name.clone(), + wall_time_ns: 1, + } +} + +#[async_trait] +impl WorkerDispatcher for TestPool { + async fn run( + &self, + mut job_rx: tokio::sync::mpsc::Receiver, + result_tx: Sender, + ) { + while let Some(job) = job_rx.recv().await { + // Executing on the runtime thread would block the reactor, which is + // what a real pool avoids by handing work to processes or threads. + let outcome = tokio::task::block_in_place(|| self.execute(&job)); + if result_tx.send(outcome).is_err() { + break; + } + } + } + + fn shutdown(&self) { + self.shutdown.store(true, Ordering::SeqCst); + } + + fn notify_cancel(&self, job_id: &str) { + self.cancels + .lock() + .expect("cancels") + .push(job_id.to_string()); + } +} + +/// A scheduler and an executor wired to each other. +struct Attached { + dispatcher: RemoteDispatcher, + handle: ExecutorHandle, + pool: Arc, + started: Receiver, +} + +/// Attach an executor advertising `tasks` to a fresh scheduler. +fn attach(tasks: &[&str], slots: u32) -> Attached { + let (started_tx, started) = crossbeam_channel::unbounded(); + let pool = TestPool::new(started_tx); + let dispatcher = scheduler(None); + let handle = dial(&dispatcher, tasks, slots, None) + .expect("attach") + .spawn(pool.clone()); + Attached { + dispatcher, + handle, + pool, + started, + } +} + +fn scheduler(auth_token: Option<&str>) -> RemoteDispatcher { + RemoteDispatcher::new(RemoteConfig { + scheduler_id: "scheduler-test".to_string(), + auth_token: auth_token.map(Secret::new), + placement_timeout: Duration::from_secs(5), + shutdown_drain: Duration::from_millis(200), + ..RemoteConfig::default() + }) +} + +/// Complete a handshake against `dispatcher`. +/// +/// `attach` blocks reading `hello` while `connect` blocks reading the ack, so +/// the scheduler side has to run concurrently — the same shape as a listener +/// thread accepting a connection. +fn dial( + dispatcher: &RemoteDispatcher, + tasks: &[&str], + slots: u32, + token: Option<&str>, +) -> Result { + let (scheduler_end, executor_end) = MemoryTransport::pair(); + let accepting = { + let dispatcher = dispatcher.clone(); + thread::spawn(move || dispatcher.attach(Box::new(scheduler_end))) + }; + + let connected = ExecutorClient::connect( + Box::new(executor_end), + ExecutorConfig { + executor_id: "exec-1".to_string(), + tasks: tasks.iter().map(|task| (*task).to_string()).collect(), + slots, + token: token.map(Secret::new), + // Fast enough that a capacity assertion does not wait on a + // production-cadence heartbeat. + heartbeat_interval: Duration::from_millis(50), + shutdown_drain: SHUTDOWN_DRAIN, + ..ExecutorConfig::new("test", "0.0.0") + }, + ); + let _ = accepting.join(); + connected +} + +/// The scheduler end of the wire, hand-driven. +/// +/// [`RemoteDispatcher`] is the right peer for most of these tests, but it is +/// also a correct one: it will not dispatch to an executor advertising no free +/// slots. Forcing that race needs a peer that writes whatever it is told to. +struct FakeScheduler { + reader: FrameReader, + writer: FrameWriter, +} + +impl FakeScheduler { + /// Handshake with an executor and return both ends live. + fn attach(tasks: &[&str], slots: u32) -> (Self, ExecutorHandle, Arc) { + let (scheduler_end, executor_end) = MemoryTransport::pair(); + + // `connect` blocks on the ack, so the scheduler side runs concurrently. + let accepting = thread::spawn(move || { + let (read, write, _connection) = Box::new(scheduler_end) + .split() + .expect("split scheduler end"); + let mut scheduler = Self { + reader: FrameReader::new(read), + writer: FrameWriter::new(write), + }; + match scheduler.reader.read::().expect("hello").0 { + ExecutorMessage::Hello { .. } => {} + other => panic!("expected hello, got {other:?}"), + } + scheduler + .writer + .write_header(&SchedulerMessage::HelloAck { + scheduler_id: "scheduler-fake".to_string(), + protocol_version: PROTOCOL_VERSION, + }) + .expect("send ack"); + scheduler + }); + + let client = ExecutorClient::connect( + Box::new(executor_end), + ExecutorConfig { + executor_id: "exec-1".to_string(), + tasks: tasks.iter().map(|task| (*task).to_string()).collect(), + slots, + heartbeat_interval: Duration::from_millis(50), + shutdown_drain: Duration::from_secs(2), + ..ExecutorConfig::new("test", "0.0.0") + }, + ) + .expect("attach"); + + let (started_tx, _started) = crossbeam_channel::unbounded(); + let pool = TestPool::new(started_tx); + let handle = client.spawn(pool.clone()); + (accepting.join().expect("handshake thread"), handle, pool) + } + + fn send_job(&mut self, id: &str, task_name: &str, payload: &[u8]) { + self.writer + .write( + &SchedulerMessage::Job { + id: id.to_string(), + task_name: task_name.to_string(), + payload_len: payload.len(), + retry_count: 1, + max_retries: 3, + queue: "default".to_string(), + timeout_ms: 30_000, + namespace: None, + }, + payload, + ) + .expect("send job"); + } + + /// Block until the executor reports the given free-slot count. + fn expect_heartbeat(&mut self, free: u32) { + let deadline = Instant::now() + SETTLE; + loop { + assert!( + Instant::now() < deadline, + "no heartbeat reporting {free} slots" + ); + if let ExecutorMessage::Heartbeat { free_slots } = self.next_frame().0 { + if free_slots == free { + return; + } + } + } + } + + /// The next frame that is not a heartbeat. + fn expect_result(&mut self) -> ExecutorMessage { + let deadline = Instant::now() + SETTLE; + loop { + assert!(Instant::now() < deadline, "no result frame arrived"); + let frame = self.next_frame().0; + if !matches!(frame, ExecutorMessage::Heartbeat { .. }) { + return frame; + } + } + } + + fn next_frame(&mut self) -> (ExecutorMessage, Vec) { + self.reader.read::().expect("read a frame") + } +} + +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 scheduler's dispatch loop live. +fn with_running(dispatcher: &RemoteDispatcher, 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(4); + let (result_tx, result_rx) = crossbeam_channel::bounded(4); + 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", + } +} + +fn wait_until(mut condition: impl FnMut() -> bool, message: &str) { + let deadline = Instant::now() + SETTLE; + while !condition() { + assert!(Instant::now() < deadline, "{message}"); + thread::sleep(Duration::from_millis(10)); + } +} + +#[test] +fn the_handshake_registers_the_executor_with_the_scheduler() { + let attached = attach(&["resize", "thumbnail"], 3); + + let executors = attached.dispatcher.executors(); + assert_eq!(executors.len(), 1); + assert_eq!(executors[0].executor_id, "exec-1"); + assert_eq!(executors[0].tasks, ["resize", "thumbnail"]); + assert_eq!(executors[0].sdk, "test"); + assert_eq!(executors[0].slots, 3); + assert_eq!(attached.dispatcher.capacity().free_slots, 3); + assert_eq!(attached.handle.executor_id(), "exec-1"); + + attached.handle.shutdown(); +} + +#[test] +fn a_job_runs_on_the_executor_and_its_result_comes_back() { + let attached = attach(&["resize"], 1); + attached + .pool + .on("resize", Behaviour::Succeed(Some(b"out".to_vec()))); + + with_running(&attached.dispatcher, |jobs, results| { + jobs.blocking_send(make_job("job-1", "resize", b"in")) + .expect("send job"); + + 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)), + } + }); + + attached.handle.shutdown(); +} + +#[test] +fn the_payload_reaches_the_pool_verbatim() { + // The CBOR envelope for f(1, "a") — the BINDING_CONTRACT test vector. A + // wire that mangled it would break every cross-SDK attach. + const ENVELOPE: &[u8] = &[0x02, 0x82, 0x82, 0x01, 0x61, 0x61, 0xa0]; + + let attached = attach(&["resize"], 1); + + with_running(&attached.dispatcher, |jobs, results| { + jobs.blocking_send(make_job("job-1", "resize", ENVELOPE)) + .expect("send job"); + assert!(matches!(expect_result(results), JobResult::Success { .. })); + }); + + let received = attached + .pool + .received("job-1") + .expect("the pool saw the job"); + assert_eq!( + received.payload, ENVELOPE, + "the wire-envelope bytes must survive the hop unchanged" + ); + + attached.handle.shutdown(); +} + +#[test] +fn the_job_frame_carries_every_field_a_task_needs() { + // Rebuilt on the far side from the frame alone — the executor never reads + // storage, so a field missing here is invisible to the task that runs. + let attached = attach(&["resize"], 1); + + with_running(&attached.dispatcher, |jobs, results| { + let mut job = make_job("job-1", "resize", b""); + job.retry_count = 2; + job.max_retries = 7; + job.timeout_ms = 1_234; + job.queue = "images".to_string(); + job.namespace = Some("tenant-a".to_string()); + + jobs.blocking_send(job).expect("send job"); + assert!(matches!(expect_result(results), JobResult::Success { .. })); + }); + + let received = attached + .pool + .received("job-1") + .expect("the pool saw the job"); + assert_eq!(received.task_name, "resize"); + assert_eq!(received.queue, "images"); + assert_eq!( + received.retry_count, 2, + "retry_count drives backoff reporting" + ); + assert_eq!(received.max_retries, 7); + assert_eq!(received.timeout_ms, 1_234, "the pool enforces the timeout"); + assert_eq!(received.namespace.as_deref(), Some("tenant-a")); + + attached.handle.shutdown(); +} + +#[test] +fn a_task_failure_crosses_the_wire_with_its_retry_verdict() { + let attached = attach(&["flaky"], 1); + attached + .pool + .on("flaky", Behaviour::Fail { should_retry: true }); + + with_running(&attached.dispatcher, |jobs, results| { + jobs.blocking_send(make_job("job-1", "flaky", 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, "the executor's verdict must survive the hop"); + assert!(!timed_out); + assert_eq!(error, "deliberate failure"); + } + ref other => panic!("expected a failure, got {}", kind(other)), + } + }); + + attached.handle.shutdown(); +} + +#[test] +fn a_non_retryable_failure_stays_non_retryable() { + // Only the executor can see the exception, so its verdict is the one that + // counts; a wire defaulting this to `true` would retry poison jobs forever. + let attached = attach(&["fatal"], 1); + attached.pool.on( + "fatal", + Behaviour::Fail { + should_retry: false, + }, + ); + + with_running(&attached.dispatcher, |jobs, results| { + jobs.blocking_send(make_job("job-1", "fatal", b"")) + .expect("send job"); + match expect_result(results) { + JobResult::Failure { should_retry, .. } => assert!(!should_retry), + ref other => panic!("expected a failure, got {}", kind(other)), + } + }); + + attached.handle.shutdown(); +} + +#[test] +fn an_empty_result_stays_distinct_from_no_result() { + let attached = attach(&["empty"], 1); + attached.pool.on("empty", Behaviour::Succeed(Some(vec![]))); + + with_running(&attached.dispatcher, |jobs, results| { + jobs.blocking_send(make_job("job-1", "empty", b"")) + .expect("send job"); + match expect_result(results) { + JobResult::Success { result, .. } => { + assert_eq!(result, Some(vec![]), "Some(empty) must not become None") + } + ref other => panic!("expected success, got {}", kind(other)), + } + }); + + attached.handle.shutdown(); +} + +#[test] +fn a_task_returning_nothing_reports_no_result() { + let attached = attach(&["nothing"], 1); + attached.pool.on("nothing", Behaviour::Succeed(None)); + + with_running(&attached.dispatcher, |jobs, results| { + jobs.blocking_send(make_job("job-1", "nothing", b"")) + .expect("send job"); + match expect_result(results) { + JobResult::Success { result, .. } => { + assert_eq!(result, None, "None must not become Some(empty)") + } + ref other => panic!("expected success, got {}", kind(other)), + } + }); + + attached.handle.shutdown(); +} + +#[test] +fn a_cancel_from_the_scheduler_reaches_the_pool() { + let (release, released) = crossbeam_channel::bounded(1); + let attached = attach(&["slow"], 1); + attached.pool.on("slow", Behaviour::Block(released)); + + with_running(&attached.dispatcher, |jobs, results| { + jobs.blocking_send(make_job("job-1", "slow", b"")) + .expect("send job"); + assert_eq!( + attached.started.recv_timeout(SETTLE).expect("job started"), + "job-1" + ); + + attached.dispatcher.notify_cancel("job-1"); + wait_until( + || attached.pool.cancelled() == ["job-1"], + "the cancel never reached the pool", + ); + + let _ = release.send(()); + assert!(matches!(expect_result(results), JobResult::Success { .. })); + }); + + attached.handle.shutdown(); +} + +#[test] +fn free_capacity_falls_while_a_job_is_running() { + let (release, released) = crossbeam_channel::bounded(1); + let attached = attach(&["slow"], 2); + attached.pool.on("slow", Behaviour::Block(released)); + + with_running(&attached.dispatcher, |jobs, results| { + jobs.blocking_send(make_job("job-1", "slow", b"")) + .expect("send job"); + attached.started.recv_timeout(SETTLE).expect("job started"); + + // The executor's heartbeat reports one slot occupied. The scheduler + // reserved that slot itself, so this asserts the two agree rather than + // that either alone is right. + wait_until( + || attached.dispatcher.capacity().free_slots == 1, + "the heartbeat never reported the occupied slot", + ); + + let _ = release.send(()); + assert!(matches!(expect_result(results), JobResult::Success { .. })); + }); + + attached.handle.shutdown(); +} + +#[test] +fn a_shutdown_frame_ends_the_session() { + let attached = attach(&["resize"], 1); + assert!(attached.handle.is_running()); + + // `drain_and_close` writes `shutdown` to every attached executor. + with_running(&attached.dispatcher, |_jobs, _results| {}); + + wait_until( + || !attached.handle.is_running(), + "a shutdown frame must end the session", + ); + attached.handle.wait(); +} + +#[test] +fn stop_finishes_in_flight_work_before_disconnecting() { + // The SIGTERM path: an executor asked to stop must still report the job it + // is holding, or that job waits for a reap it never needed. + let (release, released) = crossbeam_channel::bounded(1); + let attached = attach(&["slow"], 1); + attached.pool.on("slow", Behaviour::Block(released)); + + with_running(&attached.dispatcher, |jobs, results| { + jobs.blocking_send(make_job("job-1", "slow", b"")) + .expect("send job"); + attached.started.recv_timeout(SETTLE).expect("job started"); + + attached.handle.stop(); + + // The zero-capacity heartbeat is what tells the scheduler to stop + // dispatching; it must land before the connection goes away. + wait_until( + || attached.dispatcher.capacity().free_slots == 0, + "the drain never announced zero capacity", + ); + + let _ = release.send(()); + match expect_result(results) { + JobResult::Success { job_id, .. } => assert_eq!(job_id, "job-1"), + ref other => panic!("in-flight work must still report, got {}", kind(other)), + } + }); + + attached.handle.shutdown(); +} + +#[test] +fn a_drain_announces_zero_capacity_before_anything_else() { + let (mut scheduler, handle, _pool) = FakeScheduler::attach(&["resize"], 2); + scheduler.expect_heartbeat(2); + + handle.stop(); + + // This is what makes the drain clean rather than a race: the scheduler is + // told to stop dispatching in-protocol, before the connection goes away. + scheduler.expect_heartbeat(0); + handle.shutdown(); +} + +#[test] +fn a_job_arriving_after_the_drain_is_declined_retryably() { + // A job already on the wire when the zero-capacity heartbeat landed. The + // reaper would recover it either way, but declining reschedules it now + // instead of costing a whole reap cycle. + let (mut scheduler, handle, pool) = FakeScheduler::attach(&["resize"], 1); + handle.stop(); + scheduler.expect_heartbeat(0); + + scheduler.send_job("job-late", "resize", b""); + + match scheduler.expect_result() { + ExecutorMessage::Failure { + job_id, + should_retry, + error, + retry_count, + timed_out, + .. + } => { + assert_eq!(job_id, "job-late"); + assert!(should_retry, "a declined job must be retryable"); + assert!(!timed_out, "never started is not a timeout"); + assert_eq!( + retry_count, 1, + "the frame's retry count must be echoed back" + ); + assert!(error.contains("draining"), "error explains why: {error}"); + } + other => panic!("expected a retryable failure, got {other:?}"), + } + + assert!( + pool.received("job-late").is_none(), + "a declined job must never reach the pool" + ); + handle.shutdown(); +} + +#[test] +fn a_cancel_for_an_unknown_job_is_harmless() { + // Cancels race completion by nature; one for a job that already finished + // must not desync the stream or take the executor down. + let (mut scheduler, handle, pool) = FakeScheduler::attach(&["resize"], 1); + scheduler + .writer + .write_header(&SchedulerMessage::Cancel { + job_id: "job-gone".to_string(), + }) + .expect("send cancel"); + + wait_until( + || pool.cancelled() == ["job-gone"], + "the cancel never reached the pool", + ); + assert!( + handle.is_running(), + "a stray cancel must not end the session" + ); + + // The stream still works afterwards. + scheduler.send_job("job-1", "resize", b""); + assert!(matches!( + scheduler.expect_result(), + ExecutorMessage::Success { .. } + )); + handle.shutdown(); +} + +#[test] +fn a_missing_token_is_reported_as_a_refusal() { + // The likeliest deployment mistake. It must not surface as a transport + // error, or the operator goes looking at the network. + let dispatcher = scheduler(Some("attach-token-0123456789abcdef")); + let refused = dial(&dispatcher, &["resize"], 1, None) + .err() + .expect("an unauthenticated attach must be refused"); + + assert!( + matches!(refused, ExecutorError::Refused), + "expected a refusal, got {refused}" + ); + assert!(dispatcher.executors().is_empty()); +} + +#[test] +fn a_wrong_token_is_reported_as_a_refusal() { + let dispatcher = scheduler(Some("attach-token-0123456789abcdef")); + let refused = dial(&dispatcher, &["resize"], 1, Some("wrong-token")) + .err() + .expect("a bad credential must be refused"); + + assert!( + matches!(refused, ExecutorError::Refused), + "expected a refusal, got {refused}" + ); + assert!(dispatcher.executors().is_empty()); +} + +#[test] +fn the_right_token_attaches() { + const TOKEN: &str = "attach-token-0123456789abcdef"; + let dispatcher = scheduler(Some(TOKEN)); + let client = dial(&dispatcher, &["resize"], 1, Some(TOKEN)).expect("attach"); + + assert_eq!(client.scheduler_id(), "scheduler-test"); + assert_eq!(dispatcher.executors().len(), 1); + + let (started_tx, _started) = crossbeam_channel::unbounded(); + client.spawn(TestPool::new(started_tx)).shutdown(); +} + +#[test] +fn a_dead_peer_does_not_attach() { + // Nothing answers the hello and the far end is gone. The executor must fail + // rather than sit waiting for a job that is never coming. + let (scheduler_end, executor_end) = MemoryTransport::pair(); + drop(scheduler_end); + + let error = ExecutorClient::connect( + Box::new(executor_end), + ExecutorConfig { + tasks: vec!["resize".to_string()], + handshake_timeout: Duration::from_millis(200), + ..ExecutorConfig::new("test", "0.0.0") + }, + ) + .err() + .expect("a dead peer must not attach"); + + assert!( + matches!( + error, + ExecutorError::Refused | ExecutorError::Transport(_) | ExecutorError::Protocol(_) + ), + "expected a failed handshake, got {error}" + ); +} + +#[test] +fn a_garbled_ack_is_a_protocol_error() { + let (scheduler_end, executor_end) = MemoryTransport::pair(); + let responder = thread::spawn(move || { + let (_read, mut write, _connection) = Box::new(scheduler_end) + .split() + .expect("split scheduler end"); + use std::io::Write; + let _ = write.write_all(b"this is not a frame header\n"); + let _ = write.flush(); + // Hold the write half open so the reader sees the bytes, not an EOF. + thread::sleep(Duration::from_millis(300)); + }); + + let error = ExecutorClient::connect( + Box::new(executor_end), + ExecutorConfig { + tasks: vec!["resize".to_string()], + handshake_timeout: Duration::from_secs(1), + ..ExecutorConfig::new("test", "0.0.0") + }, + ) + .err() + .expect("a garbled ack must not attach"); + + assert!( + matches!(error, ExecutorError::Protocol(ProtocolError::Json(_))), + "expected a protocol error, got {error}" + ); + let _ = responder.join(); +} + +#[test] +fn the_executor_detaches_from_the_scheduler_when_it_stops() { + let attached = attach(&["resize"], 1); + assert_eq!(attached.dispatcher.executors().len(), 1); + + attached.handle.shutdown(); + + wait_until( + || attached.dispatcher.executors().is_empty(), + "the scheduler never saw the executor leave", + ); +} + +#[test] +fn shutdown_is_bounded_when_a_job_never_finishes() { + // A task that ignores its cancel must not hang the process forever: the + // drain budget expires and the executor disconnects, leaving the job to the + // scheduler's reaper. The job stays wedged for as long as `release` is held, + // so nothing but the budget can end the shutdown. + let (release, wedged) = crossbeam_channel::bounded::<()>(1); + let attached = attach(&["stuck"], 1); + attached.pool.on("stuck", Behaviour::Wedge(wedged)); + + let started = Instant::now(); + with_running(&attached.dispatcher, |jobs, _results| { + jobs.blocking_send(make_job("job-1", "stuck", b"")) + .expect("send job"); + attached.started.recv_timeout(SETTLE).expect("job started"); + }); + attached.handle.shutdown(); + + assert!( + started.elapsed() < SHUTDOWN_DRAIN * 3, + "shutdown must be bounded by the drain budget (took {:?})", + started.elapsed() + ); + + // Held until here so the job could not have finished on its own. + drop(release); +} + +#[test] +fn a_local_stop_releases_a_parked_waiter() { + // Regression: `stop()` cannot unpark the reader, which is blocked on a read + // only the scheduler could satisfy. If the session were only ended by the + // reader, a shell that called `stop()` from a signal handler and then waited + // would hang forever instead of shutting down. + let (_scheduler, handle, _pool) = FakeScheduler::attach(&["resize"], 1); + let session = handle.session(); + assert!(session.is_running()); + + let waiting = thread::spawn(move || session.wait()); + thread::sleep(Duration::from_millis(50)); + assert!(!waiting.is_finished(), "the waiter must park while running"); + + handle.stop(); + + let deadline = Instant::now() + SETTLE; + while !waiting.is_finished() { + assert!( + Instant::now() < deadline, + "stop() never released the waiter" + ); + thread::sleep(Duration::from_millis(10)); + } + waiting.join().expect("waiter thread"); + handle.shutdown(); +} + +#[test] +fn wait_timeout_reports_a_session_that_is_still_open() { + let attached = attach(&["resize"], 1); + + assert!( + !attached.handle.wait_timeout(Duration::from_millis(100)), + "a live session must not report as finished" + ); + assert!(attached.handle.is_running()); + + attached.handle.shutdown(); +} diff --git a/crates/taskito-core/src/worker/remote_tests.rs b/crates/taskito-core/tests/rust/remote_tests.rs similarity index 97% rename from crates/taskito-core/src/worker/remote_tests.rs rename to crates/taskito-core/tests/rust/remote_tests.rs index 0710b5708..bd64300df 100644 --- a/crates/taskito-core/src/worker/remote_tests.rs +++ b/crates/taskito-core/tests/rust/remote_tests.rs @@ -6,18 +6,18 @@ use std::time::{Duration, Instant}; use crossbeam_channel::{Receiver, RecvTimeoutError}; -use super::auth::Secret; -use super::protocol::{ +use taskito_core::job::{now_millis, Job, JobStatus, NewJob}; +use taskito_core::scheduler::{JobResult, SchedulerConfig}; +use taskito_core::storage::sqlite::SqliteStorage; +use taskito_core::storage::{Storage, StorageBackend}; +use taskito_core::worker::auth::Secret; +use taskito_core::worker::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; +use taskito_core::worker::remote::{AttachError, RemoteConfig, RemoteDispatcher}; +use taskito_core::worker::transport::{MemoryTransport, ReadHalf, Transport, WriteHalf}; +use taskito_core::worker::Worker; +use taskito_core::worker::WorkerDispatcher; const SETTLE: Duration = Duration::from_secs(5); diff --git a/crates/taskito-core/src/worker/tests.rs b/crates/taskito-core/tests/rust/worker_tests.rs similarity index 93% rename from crates/taskito-core/src/worker/tests.rs rename to crates/taskito-core/tests/rust/worker_tests.rs index 134aaaa9b..97d756c10 100644 --- a/crates/taskito-core/src/worker/tests.rs +++ b/crates/taskito-core/tests/rust/worker_tests.rs @@ -5,13 +5,13 @@ use std::sync::atomic::{AtomicU32, Ordering}; use std::sync::Arc; use std::time::{Duration, Instant}; -use super::registry::TaskError; -use super::runner::Worker; -use crate::job::{now_millis, Job, JobStatus, NewJob}; -use crate::resilience::retry::RetryPolicy; -use crate::scheduler::TaskConfig; -use crate::storage::sqlite::SqliteStorage; -use crate::storage::{Storage, StorageBackend}; +use taskito_core::job::{now_millis, Job, JobStatus, NewJob}; +use taskito_core::resilience::retry::RetryPolicy; +use taskito_core::scheduler::TaskConfig; +use taskito_core::storage::sqlite::SqliteStorage; +use taskito_core::storage::{Storage, StorageBackend}; +use taskito_core::worker::registry::TaskError; +use taskito_core::worker::runner::Worker; fn test_backend() -> StorageBackend { StorageBackend::Sqlite(SqliteStorage::in_memory().expect("in-memory sqlite")) diff --git a/crates/taskito-java/src/dispatcher.rs b/crates/taskito-java/src/dispatcher.rs index cc69323d8..f7c4a297e 100644 --- a/crates/taskito-java/src/dispatcher.rs +++ b/crates/taskito-java/src/dispatcher.rs @@ -17,8 +17,8 @@ use crossbeam_channel::Sender; use jni::objects::{GlobalRef, JValue}; use taskito_core::job::Job; use taskito_core::scheduler::JobResult; -use taskito_core::worker::WorkerDispatcher; -use taskito_core::{Storage, StorageBackend}; +use taskito_core::worker::{CancelSignals, WorkerDispatcher}; +use taskito_core::StorageBackend; use tokio::sync::oneshot; use crate::jvm; @@ -63,7 +63,9 @@ impl Registry { pub struct JavaDispatcher { callbacks: GlobalRef, registry: Arc, - storage: StorageBackend, + /// Where a cancel comes from: the storage flag for a worker, the scheduler's + /// `cancel` frame for an attached executor, which has no storage at all. + cancels: Arc, } impl JavaDispatcher { @@ -71,7 +73,17 @@ impl JavaDispatcher { Self { callbacks, registry, - storage, + cancels: Arc::new(CancelSignals::from_storage(storage)), + } + } + + /// A dispatcher with no storage, for an attached executor. Cancels arrive + /// only through [`WorkerDispatcher::notify_cancel`]. + pub fn detached(callbacks: GlobalRef, registry: Arc) -> Self { + Self { + callbacks, + registry, + cancels: Arc::new(CancelSignals::detached()), } } } @@ -86,23 +98,31 @@ impl WorkerDispatcher for JavaDispatcher { while let Some(job) = job_rx.recv().await { let callbacks = self.callbacks.clone(); let registry = self.registry.clone(); - let storage = self.storage.clone(); + let cancels = self.cancels.clone(); let result_tx = result_tx.clone(); tokio::spawn(async move { - let result = run_one(&callbacks, ®istry, &storage, job).await; + let job_id = job.id.clone(); + let result = run_one(&callbacks, ®istry, &cancels, job).await; + // Release the cancel record now the job has reported, so a + // long-lived process does not accumulate ids. + cancels.forget(&job_id); let _ = result_tx.send(result); }); } } fn shutdown(&self) {} + + fn notify_cancel(&self, job_id: &str) { + self.cancels.signal(job_id); + } } /// Submit one job to Java, await its completion, and translate to a [`JobResult`]. async fn run_one( callbacks: &GlobalRef, registry: &Registry, - storage: &StorageBackend, + cancels: &CancelSignals, job: Job, ) -> JobResult { let started = Instant::now(); @@ -139,7 +159,7 @@ async fn run_one( Ok(TaskOutcome::Cancelled) => cancelled(job, wall), Ok(TaskOutcome::Failure(error, retryable)) => { // A failure on a cancel-requested job is a cancellation, not a fault. - if storage.is_cancel_requested(&job.id).unwrap_or(false) { + if cancels.is_cancelled(&job.id) { cancelled(job, wall) } else { failure(job, error, wall, false, retryable) diff --git a/crates/taskito-java/src/executor.rs b/crates/taskito-java/src/executor.rs new file mode 100644 index 000000000..2b7031fcc --- /dev/null +++ b/crates/taskito-java/src/executor.rs @@ -0,0 +1,329 @@ +//! `taskito executor` — attach to a detached scheduler and run its jobs here. +//! +//! The mirror of [`crate::worker`] with storage swapped for a socket. Jobs run +//! on the same [`JavaDispatcher`] and the same `WorkerBridge` callback, so a +//! handler behaves identically whichever way its job arrived; only the +//! transport differs. +//! +//! No `QueueHandle` is involved, and deliberately so: an executor opens no +//! storage. That is the point of the split — the scheduler image holds the +//! database credentials, the app image holds the task bodies. + +use std::sync::Arc; +use std::time::Duration; + +use jni::objects::{GlobalRef, JByteArray, JClass, JObject, JString}; +use jni::sys::{jboolean, jlong, JNI_FALSE}; +use jni::JNIEnv; +use serde::Deserialize; + +use taskito_core::worker::{ + AttachAddress, ExecutorClient, ExecutorConfig, ExecutorError, ExecutorHandle, ExecutorSession, + WorkerDispatcher, +}; + +use crate::convert::parse_json; +use crate::dispatcher::{JavaDispatcher, Registry, TaskOutcome}; +use crate::error::BindingError; +use crate::ffi::{guard, new_string, read_bytes, read_string}; +use crate::handle::{self, into_handle}; + +/// How an executor attaches. Durations are milliseconds, matching the Java API. +#[derive(Debug, Deserialize)] +#[serde(rename_all = "camelCase")] +struct ExecutorOptions { + /// Scheduler address: `host:port`, `:port`, or `unix:/run/taskito.sock`. + address: String, + /// Task names this executor can run. The scheduler sends it nothing else. + tasks: Vec, + #[serde(default)] + slots: Option, + #[serde(default)] + token: Option, + #[serde(default)] + executor_id: Option, + #[serde(default)] + connect_timeout_ms: Option, + #[serde(default)] + heartbeat_interval_ms: Option, + #[serde(default)] + shutdown_drain_ms: Option, +} + +/// How long to wait for the connection when the caller gave no budget. +const DEFAULT_CONNECT_TIMEOUT_MS: u64 = 10_000; + +/// A running attachment. Holds the completion registry for the executor's life, +/// so a token handed to Java stays resolvable until [`close`] runs. +pub struct AttachedHandle { + handle: Option, + session: ExecutorSession, + /// Resolves the tokens handed to Java, so a handler can complete its job. + registry: Arc, + scheduler_id: String, + executor_id: String, + peer: String, +} + +impl AttachedHandle { + /// Drain in-flight work, disconnect, and join. Idempotent. + fn shutdown(&mut self) { + if let Some(handle) = self.handle.take() { + handle.shutdown(); + } + } +} + +/// Dial, handshake, and start running jobs on the Java bridge. +fn attach(options: ExecutorOptions, callbacks: GlobalRef) -> Result { + if options.tasks.is_empty() { + // This would attach successfully and then sit idle forever, because the + // scheduler only dispatches task names an executor advertises. + return Err(BindingError::new( + "no handlers were found, so the executor would never be sent any work", + )); + } + let slots = options.slots.unwrap_or(1).max(1); + + let mut config = ExecutorConfig { + tasks: options.tasks, + slots, + token: options.token.map(taskito_core::Secret::new), + ..ExecutorConfig::new("java", env!("CARGO_PKG_VERSION")) + }; + if let Some(id) = options.executor_id { + config.executor_id = id; + } + if let Some(interval) = options.heartbeat_interval_ms { + config.heartbeat_interval = Duration::from_millis(interval); + } + if let Some(drain) = options.shutdown_drain_ms { + config.shutdown_drain = Duration::from_millis(drain); + } + + let target = AttachAddress::parse(&options.address) + .map_err(|error| BindingError::new(format!("invalid attach address: {error}")))?; + let connect_timeout = Duration::from_millis( + options + .connect_timeout_ms + .unwrap_or(DEFAULT_CONNECT_TIMEOUT_MS), + ); + let transport = target.connect(connect_timeout).map_err(|error| { + BindingError::new(format!( + "could not reach the scheduler at {target}: {error}" + )) + })?; + let client = ExecutorClient::connect(transport, config).map_err(|error| match error { + // Named so a wrong token reads as a refusal rather than a network fault. + ExecutorError::Refused => BindingError::new(error.to_string()), + other => BindingError::new(format!("could not attach to {target}: {other}")), + })?; + + let scheduler_id = client.scheduler_id().to_string(); + let peer = client.peer().to_string(); + + let registry = Arc::new(Registry::default()); + let pool: Arc = + Arc::new(JavaDispatcher::detached(callbacks, registry.clone())); + let handle = client.spawn(pool); + + Ok(AttachedHandle { + executor_id: handle.executor_id().to_string(), + session: handle.session(), + handle: Some(handle), + registry, + scheduler_id, + peer, + }) +} + +/// Borrow the value behind an executor handle. +/// +/// # Safety +/// `handle` must be a live `AttachedHandle` pointer from [`attach`]. +unsafe fn borrow<'a>(handle: jlong) -> &'a AttachedHandle { + handle::borrow::(handle) +} + +/// `long attach(Object bridge, String optionsJson)` — dial and start; returns a +/// handle. `bridge` is a Java `WorkerBridge`. +#[no_mangle] +pub extern "system" fn Java_org_byteveda_taskito_internal_NativeExecutor_attach<'local>( + mut env: JNIEnv<'local>, + _class: JClass<'local>, + bridge: JObject<'local>, + options_json: JString<'local>, +) -> jlong { + guard(&mut env, 0, |env| { + let raw = read_string(env, &options_json)?; + let options: ExecutorOptions = parse_json(&raw, "executor options")?; + let callbacks = env + .new_global_ref(&bridge) + .map_err(|e| BindingError::new(format!("global ref failed: {e}")))?; + Ok(into_handle(attach(options, callbacks)?)) + }) +} + +/// `void completeJob(long handle, long token, byte[] result)`. +#[no_mangle] +pub extern "system" fn Java_org_byteveda_taskito_internal_NativeExecutor_completeJob<'local>( + mut env: JNIEnv<'local>, + _class: JClass<'local>, + handle: jlong, + token: jlong, + result: JByteArray<'local>, +) { + guard(&mut env, (), |env| { + let executor = unsafe { borrow(handle) }; + let bytes = read_bytes(env, &result)?; + executor + .registry + .complete(token as u64, TaskOutcome::Success(bytes)); + Ok(()) + }) +} + +/// `void failJob(long handle, long token, String error, boolean retryable)`. +#[no_mangle] +pub extern "system" fn Java_org_byteveda_taskito_internal_NativeExecutor_failJob<'local>( + mut env: JNIEnv<'local>, + _class: JClass<'local>, + handle: jlong, + token: jlong, + error: JString<'local>, + retryable: jboolean, +) { + guard(&mut env, (), |env| { + let executor = unsafe { borrow(handle) }; + let message = read_string(env, &error)?; + executor.registry.complete( + token as u64, + TaskOutcome::Failure(message, retryable != JNI_FALSE), + ); + Ok(()) + }) +} + +/// `void cancelJob(long handle, long token)`. +#[no_mangle] +pub extern "system" fn Java_org_byteveda_taskito_internal_NativeExecutor_cancelJob( + mut env: JNIEnv, + _class: JClass, + handle: jlong, + token: jlong, +) { + guard(&mut env, (), |_env| { + let executor = unsafe { borrow(handle) }; + executor + .registry + .complete(token as u64, TaskOutcome::Cancelled); + Ok(()) + }) +} + +/// `String schedulerId(long handle)`. +#[no_mangle] +pub extern "system" fn Java_org_byteveda_taskito_internal_NativeExecutor_schedulerId<'local>( + mut env: JNIEnv<'local>, + _class: JClass<'local>, + handle: jlong, +) -> jni::sys::jstring { + guard(&mut env, std::ptr::null_mut(), |env| { + let executor = unsafe { borrow(handle) }; + new_string(env, executor.scheduler_id.clone()) + }) +} + +/// `String executorId(long handle)`. +#[no_mangle] +pub extern "system" fn Java_org_byteveda_taskito_internal_NativeExecutor_executorId<'local>( + mut env: JNIEnv<'local>, + _class: JClass<'local>, + handle: jlong, +) -> jni::sys::jstring { + guard(&mut env, std::ptr::null_mut(), |env| { + let executor = unsafe { borrow(handle) }; + new_string(env, executor.executor_id.clone()) + }) +} + +/// `String peer(long handle)`. +#[no_mangle] +pub extern "system" fn Java_org_byteveda_taskito_internal_NativeExecutor_peer<'local>( + mut env: JNIEnv<'local>, + _class: JClass<'local>, + handle: jlong, +) -> jni::sys::jstring { + guard(&mut env, std::ptr::null_mut(), |env| { + let executor = unsafe { borrow(handle) }; + new_string(env, executor.peer.clone()) + }) +} + +/// `boolean isRunning(long handle)` — whether the scheduler session is open. +#[no_mangle] +pub extern "system" fn Java_org_byteveda_taskito_internal_NativeExecutor_isRunning<'local>( + mut env: JNIEnv<'local>, + _class: JClass<'local>, + handle: jlong, +) -> jboolean { + guard(&mut env, 0, |_env| { + let executor = unsafe { borrow(handle) }; + Ok(jboolean::from(executor.session.is_running())) + }) +} + +/// `void awaitSession(long handle)` — block until the scheduler ends the +/// session. Java calls it from a thread it is willing to park. +#[no_mangle] +pub extern "system" fn Java_org_byteveda_taskito_internal_NativeExecutor_awaitSession<'local>( + mut env: JNIEnv<'local>, + _class: JClass<'local>, + handle: jlong, +) { + guard(&mut env, (), |_env| { + let executor = unsafe { borrow(handle) }; + // Cloned so the wait holds no borrow of the handle: `close` may run as + // soon as this returns. + let session = executor.session.clone(); + session.wait(); + Ok(()) + }) +} + +/// `void stop(long handle)` — stop accepting work; returns at once. +#[no_mangle] +pub extern "system" fn Java_org_byteveda_taskito_internal_NativeExecutor_stop<'local>( + mut env: JNIEnv<'local>, + _class: JClass<'local>, + handle: jlong, +) { + guard(&mut env, (), |_env| { + if handle != 0 { + let executor = unsafe { borrow(handle) }; + if let Some(running) = executor.handle.as_ref() { + running.stop(); + } + } + Ok(()) + }) +} + +/// `void close(long handle)` — drain, disconnect, and reclaim the handle. +#[no_mangle] +pub extern "system" fn Java_org_byteveda_taskito_internal_NativeExecutor_close<'local>( + mut env: JNIEnv<'local>, + _class: JClass<'local>, + handle: jlong, +) { + // Route through `guard` so a panic in teardown cannot unwind across FFI. + guard(&mut env, (), |_env| { + if handle != 0 { + // Reclaimed by hand rather than via `drop_handle`: the drain has to + // run before the box drops, and it needs `&mut`. + let mut executor = unsafe { Box::from_raw(handle as *mut AttachedHandle) }; + executor.shutdown(); + } + Ok(()) + }) +} diff --git a/crates/taskito-java/src/lib.rs b/crates/taskito-java/src/lib.rs index f65d42bd9..bb22ebdf2 100644 --- a/crates/taskito-java/src/lib.rs +++ b/crates/taskito-java/src/lib.rs @@ -12,6 +12,7 @@ mod backend; mod convert; mod dispatcher; mod error; +mod executor; mod ffi; mod ffi_c; mod handle; diff --git a/crates/taskito-node/src/dispatcher.rs b/crates/taskito-node/src/dispatcher.rs index 16af79d49..fc995f352 100644 --- a/crates/taskito-node/src/dispatcher.rs +++ b/crates/taskito-node/src/dispatcher.rs @@ -14,8 +14,8 @@ use napi::threadsafe_function::{ThreadsafeFunction, ThreadsafeFunctionCallMode}; use napi::Env; use taskito_core::job::Job; use taskito_core::scheduler::JobResult; -use taskito_core::worker::WorkerDispatcher; -use taskito_core::{Storage, StorageBackend}; +use taskito_core::worker::{CancelSignals, WorkerDispatcher}; +use taskito_core::StorageBackend; use tokio::sync::{oneshot, Semaphore}; use crate::convert::{JsTaskInvocation, JsTaskOutcome}; @@ -31,7 +31,9 @@ pub struct NodeDispatcher { /// `Arc` because napi 3's `ThreadsafeFunction` is not `Clone` — every /// spawned job needs its own handle to the same callback. callback: Arc, - storage: StorageBackend, + /// Where a cancel comes from: the storage flag for a worker, the scheduler's + /// `cancel` frame for an attached executor, which has no storage at all. + cancels: Arc, /// Caps jobs running at once. Without it the loop spawns every job it is /// handed and immediately takes the next, so nothing bounds concurrency. /// @@ -48,13 +50,38 @@ impl NodeDispatcher { callback: TaskCallback, storage: StorageBackend, concurrency: Option, + ) -> Self { + Self::with_cancels( + callback, + Arc::new(CancelSignals::from_storage(storage)), + concurrency, + ) + } + + /// A dispatcher with no storage, for an attached executor. Cancels arrive + /// only through [`WorkerDispatcher::notify_cancel`]. + pub fn detached(callback: TaskCallback, concurrency: Option) -> Self { + Self::with_cancels(callback, Arc::new(CancelSignals::detached()), concurrency) + } + + /// The cancel source this dispatcher reads. An attached executor hands it to + /// its JS side, which has no storage to poll and so cannot otherwise learn + /// that a running job was cancelled. + pub fn cancels(&self) -> Arc { + self.cancels.clone() + } + + fn with_cancels( + callback: TaskCallback, + cancels: Arc, + concurrency: Option, ) -> Self { let permits = concurrency .map(|c| c.max(1)) .unwrap_or(Semaphore::MAX_PERMITS); Self { callback: Arc::new(callback), - storage, + cancels, concurrency: Arc::new(Semaphore::new(permits)), } } @@ -80,12 +107,15 @@ impl WorkerDispatcher for NodeDispatcher { Err(_) => break, }; let callback = self.callback.clone(); - let storage = self.storage.clone(); + let cancels = self.cancels.clone(); let result_tx = result_tx.clone(); spawn(async move { let _permit = permit; let job_id = job.id.clone(); - let result = run_one(&callback, &storage, job).await; + let result = run_one(&callback, &cancels, job).await; + // Release the cancel record now the job has reported, so a + // long-lived process does not accumulate ids. + cancels.forget(&job_id); // A full bounded channel parks the sender — do it on the // blocking pool, never on the shared async runtime. match spawn_blocking(move || result_tx.send(result)).await { @@ -102,10 +132,14 @@ impl WorkerDispatcher for NodeDispatcher { } fn shutdown(&self) {} + + fn notify_cancel(&self, job_id: &str) { + self.cancels.signal(job_id); + } } /// Invoke the JS task for one job and translate the outcome into a [`JobResult`]. -async fn run_one(callback: &TaskCallback, storage: &StorageBackend, mut job: Job) -> JobResult { +async fn run_one(callback: &TaskCallback, cancels: &CancelSignals, mut job: Job) -> JobResult { let started = Instant::now(); let invocation = JsTaskInvocation { id: job.id.clone(), @@ -163,13 +197,11 @@ async fn run_one(callback: &TaskCallback, storage: &StorageBackend, mut job: Job }, // A failed task that was cancel-requested is a cancellation, not a // failure (the JS side aborts via the cancel signal). - Some(_) if storage.is_cancel_requested(&job.id).unwrap_or(false) => { - JobResult::Cancelled { - job_id: job.id, - task_name: job.task_name, - wall_time_ns, - } - } + Some(_) if cancels.is_cancelled(&job.id) => JobResult::Cancelled { + job_id: job.id, + task_name: job.task_name, + wall_time_ns, + }, Some(error) => failure(job, error, wall_time_ns, false, outcome.retryable), }, // The promise rejected rather than resolving an outcome: the shell threw diff --git a/crates/taskito-node/src/executor.rs b/crates/taskito-node/src/executor.rs new file mode 100644 index 000000000..198972a22 --- /dev/null +++ b/crates/taskito-node/src/executor.rs @@ -0,0 +1,228 @@ +//! `taskito executor` — attach to a detached scheduler and run its jobs here. +//! +//! A free function rather than a `JsQueue` method, and deliberately so: an +//! executor opens no storage. That is the point of the split — the scheduler +//! image holds the database credentials, the app image holds the task bodies, +//! and everything a job needs to run arrives on the wire. +//! +//! Task execution is the same [`NodeDispatcher`] the in-process worker uses, so +//! concurrency, timeouts and the cancel signal behave identically; only the +//! transport differs. + +use std::sync::{Arc, Mutex, PoisonError}; +use std::time::Duration; + +use napi::bindgen_prelude::{spawn_blocking, Promise, Result, Status}; +use napi::threadsafe_function::ThreadsafeFunction; +use napi_derive::napi; +use taskito_core::worker::{ + AttachAddress, CancelSignals, ExecutorClient, ExecutorConfig, ExecutorError, ExecutorHandle, + ExecutorSession, WorkerDispatcher, +}; + +use crate::convert::{JsTaskInvocation, JsTaskOutcome}; +use crate::dispatcher::NodeDispatcher; +use crate::error::invalid_arg; + +/// Defaults chosen to match the other SDKs rather than to be tuned here. +const DEFAULT_SLOTS: u32 = 1; +const DEFAULT_CONNECT_TIMEOUT_MS: u32 = 10_000; + +/// How an executor attaches. Durations are milliseconds, per Node convention. +#[napi(object)] +pub struct ExecutorOptions { + /// Scheduler address: `host:port`, `:port`, or `unix:/run/taskito.sock`. + pub address: String, + /// Task names this executor can run. The scheduler sends it nothing else, + /// so a name missing here is a job that never arrives. + pub tasks: Vec, + /// Jobs to run at once (default 1). + pub slots: Option, + /// Shared secret, when the scheduler requires one. + pub token: Option, + /// Identity announced to the scheduler (default: generated per process). + pub executor_id: Option, + /// How long to wait for the connection (default 10000). + pub connect_timeout_ms: Option, + /// How often to send a liveness heartbeat (default 5000). + pub heartbeat_interval_ms: Option, + /// How long a drain waits for in-flight jobs before disconnecting anyway + /// (default 30000). + pub shutdown_drain_ms: Option, +} + +/// A running attachment to a scheduler. +#[napi] +pub struct JsExecutor { + /// Taken by `shutdown`, which consumes the handle to join its threads. + handle: Arc>>, + session: ExecutorSession, + /// Cancels delivered as protocol frames. The JS side polls this instead of + /// a storage flag, which a detached executor does not have. + cancels: Arc, + scheduler_id: String, + executor_id: String, + peer: String, +} + +#[napi] +impl JsExecutor { + /// Identity the scheduler announced when it accepted this attach. + #[napi(getter)] + pub fn scheduler_id(&self) -> String { + self.scheduler_id.clone() + } + + /// Identity this executor attached under. + #[napi(getter)] + pub fn executor_id(&self) -> String { + self.executor_id.clone() + } + + /// Peer label of the scheduler connection. + #[napi(getter)] + pub fn peer(&self) -> String { + self.peer.clone() + } + + /// Whether the scheduler session is still open. + #[napi] + pub fn is_running(&self) -> bool { + self.session.is_running() + } + + /// Whether the scheduler has asked for `job_id` to be cancelled. + /// + /// The cancel arrives as a frame, not a storage flag, so this is the only + /// way a running handler can observe one. + #[napi] + pub fn is_cancel_requested(&self, job_id: String) -> bool { + self.cancels.is_cancelled(&job_id) + } + + /// Resolve once the scheduler ends the session — a `shutdown` frame, or the + /// connection dropping. Does not drain; call `shutdown()` for that. + #[napi] + pub async fn wait(&self) -> Result<()> { + let session = self.session.clone(); + spawn_blocking(move || session.wait()) + .await + .map_err(|error| napi::Error::from_reason(error.to_string())) + } + + /// Ask the scheduler to stop sending work and finish what is in flight. + /// Returns at once, so it is safe from a signal handler. + #[napi] + pub fn stop(&self) { + if let Some(handle) = self.locked().as_ref() { + handle.stop(); + } + } + + /// Drain in-flight work, disconnect, and join. Idempotent. + #[napi] + pub async fn shutdown(&self) -> Result<()> { + let taken = self.locked().take(); + let Some(handle) = taken else { + return Ok(()); + }; + // Joining blocks on the drain budget; never park the JS event loop on it. + spawn_blocking(move || handle.shutdown()) + .await + .map_err(|error| napi::Error::from_reason(error.to_string())) + } +} + +impl JsExecutor { + fn locked(&self) -> std::sync::MutexGuard<'_, Option> { + self.handle.lock().unwrap_or_else(PoisonError::into_inner) + } +} + +/// Attach to a scheduler and run `callback` for each job it dispatches. +/// +/// The handshake happens here, so a bad token or an unreachable scheduler +/// rejects before any pool is built. +/// +/// Async because dialling and the handshake block: run on the JS thread they +/// would freeze the event loop until the scheduler answers, or for the whole +/// handshake timeout when it never does. +#[napi] +pub async fn start_executor( + // Spelled out rather than via the `TaskCallback` alias: napi-derive resolves + // these generics syntactically, and an alias reaches the generated + // `index.d.ts` as an undefined type name. + callback: ThreadsafeFunction< + JsTaskInvocation, + Promise, + JsTaskInvocation, + Status, + false, + >, + options: ExecutorOptions, +) -> Result { + if options.tasks.is_empty() { + // This would attach successfully and then sit idle forever, because the + // scheduler only dispatches task names an executor advertises. + return Err(invalid_arg( + "no tasks are registered on this app, so the executor would never be sent any work", + )); + } + let slots = options.slots.unwrap_or(DEFAULT_SLOTS).max(1); + + let mut config = ExecutorConfig { + tasks: options.tasks, + slots, + token: options.token.map(taskito_core::Secret::new), + ..ExecutorConfig::new("node", env!("CARGO_PKG_VERSION")) + }; + if let Some(id) = options.executor_id { + config.executor_id = id; + } + if let Some(interval) = options.heartbeat_interval_ms { + config.heartbeat_interval = Duration::from_millis(interval.into()); + } + if let Some(drain) = options.shutdown_drain_ms { + config.shutdown_drain = Duration::from_millis(drain.into()); + } + + let connect_timeout = Duration::from_millis( + options + .connect_timeout_ms + .unwrap_or(DEFAULT_CONNECT_TIMEOUT_MS) + .into(), + ); + let target = AttachAddress::parse(&options.address) + .map_err(|error| invalid_arg(format!("invalid attach address: {error}")))?; + + spawn_blocking(move || { + let transport = target.connect(connect_timeout).map_err(|error| { + napi::Error::from_reason(format!( + "could not reach the scheduler at {target}: {error}" + )) + })?; + let client = ExecutorClient::connect(transport, config).map_err(|error| match error { + // Named so a wrong token reads as a refusal rather than a network fault. + ExecutorError::Refused => napi::Error::from_reason(error.to_string()), + other => napi::Error::from_reason(format!("could not attach to {target}: {other}")), + })?; + + let scheduler_id = client.scheduler_id().to_string(); + let peer = client.peer().to_string(); + + let dispatcher = NodeDispatcher::detached(callback, Some(slots as usize)); + let cancels = dispatcher.cancels(); + let handle = client.spawn(Arc::new(dispatcher) as Arc); + + Ok(JsExecutor { + executor_id: handle.executor_id().to_string(), + session: handle.session(), + cancels, + handle: Arc::new(Mutex::new(Some(handle))), + scheduler_id, + peer, + }) + }) + .await + .map_err(|error| napi::Error::from_reason(error.to_string()))? +} diff --git a/crates/taskito-node/src/lib.rs b/crates/taskito-node/src/lib.rs index fdf8a0a2d..7f5aab61f 100644 --- a/crates/taskito-node/src/lib.rs +++ b/crates/taskito-node/src/lib.rs @@ -10,9 +10,11 @@ mod config; mod convert; mod dispatcher; mod error; +mod executor; mod queue; mod worker; +pub use executor::{start_executor, JsExecutor}; pub use queue::JsQueue; pub use worker::JsWorker; diff --git a/crates/taskito-python/src/executor.rs b/crates/taskito-python/src/executor.rs new file mode 100644 index 000000000..e34b61b47 --- /dev/null +++ b/crates/taskito-python/src/executor.rs @@ -0,0 +1,183 @@ +//! `taskito executor` — attach to a detached scheduler and run its jobs here. +//! +//! Jobs run on the same [`PreforkPool`] the in-process worker uses, so the +//! timeout watchdog, least-loaded dispatch, restart-on-crash and `SIGKILL` +//! after the drain budget all come from shipped code. `taskito.prefork.child` +//! is untouched: it already speaks these frames over stdio, so attaching is a +//! second hop of the same protocol rather than a new one. +//! +//! The lifecycle is split into `wait`/`stop` rather than one blocking `run` +//! because Python signal handlers only run when the main thread holds the GIL. +//! A single blocking call with the GIL released would make the process +//! unkillable by `SIGTERM`, which is exactly the signal a container gets. + +use std::sync::{Arc, Mutex}; +use std::time::Duration; + +use pyo3::exceptions::{PyRuntimeError, PyValueError}; +use pyo3::prelude::*; + +use taskito_core::worker::{ + AttachAddress, ExecutorClient, ExecutorConfig, ExecutorError, ExecutorHandle, +}; + +use crate::prefork::PreforkPool; + +/// How long to wait for the TCP connect before giving up on the scheduler. +const CONNECT_TIMEOUT: Duration = Duration::from_secs(10); + +/// A running attachment to a scheduler. +/// +/// Constructed already connected: the handshake happens in `__new__`, so a bad +/// token or an unreachable scheduler raises here rather than after the pool has +/// been built. +#[pyclass(name = "Executor", module = "taskito._taskito")] +pub struct PyExecutor { + /// Taken by `shutdown`, which consumes the handle to join its threads. + handle: Mutex>, + scheduler_id: String, + executor_id: String, + peer: String, +} + +#[pymethods] +impl PyExecutor { + /// Attach to `address` and start running `tasks` from `app_path`. + /// + /// `slots` is the number of prefork children, so it is also the number of + /// jobs that can run at once. + #[new] + #[pyo3(signature = (address, app_path, tasks, slots, token=None, executor_id=None))] + fn new( + py: Python<'_>, + address: &str, + app_path: &str, + tasks: Vec, + slots: u32, + token: Option, + executor_id: Option, + ) -> PyResult { + if slots == 0 { + return Err(PyValueError::new_err("slots must be at least 1")); + } + if tasks.is_empty() { + // The scheduler only sends tasks an executor advertises, so this + // would attach successfully and then sit idle forever. + return Err(PyValueError::new_err( + "no tasks are registered on this app, so the executor would never \ + be sent any work", + )); + } + + let mut config = ExecutorConfig { + tasks, + slots, + token: token.map(taskito_core::Secret::new), + ..ExecutorConfig::new("python", env!("CARGO_PKG_VERSION")) + }; + if let Some(id) = executor_id { + config.executor_id = id; + } + + let pool: Arc = + Arc::new(PreforkPool::new(slots as usize, app_path.to_string())); + + // Dialling and the handshake both block on the network; holding the GIL + // across them would freeze every other Python thread in the process. + let client = py + .detach(|| -> Result { + let target = AttachAddress::parse(address).map_err(|error| error.to_string())?; + let transport = target.connect(CONNECT_TIMEOUT).map_err(|error| { + format!("could not reach the scheduler at {target}: {error}") + })?; + ExecutorClient::connect(transport, config).map_err(|error| match error { + // Named so a wrong token reads as a refusal rather than as + // a network fault. + ExecutorError::Refused => error.to_string(), + other => format!("could not attach to {target}: {other}"), + }) + }) + .map_err(PyRuntimeError::new_err)?; + + let scheduler_id = client.scheduler_id().to_string(); + let peer = client.peer().to_string(); + let handle = client.spawn(pool); + + Ok(Self { + executor_id: handle.executor_id().to_string(), + handle: Mutex::new(Some(handle)), + scheduler_id, + peer, + }) + } + + /// Identity the scheduler announced when it accepted this attach. + #[getter] + fn scheduler_id(&self) -> &str { + &self.scheduler_id + } + + /// Identity this executor attached under. + #[getter] + fn executor_id(&self) -> &str { + &self.executor_id + } + + /// Peer label of the scheduler connection. + #[getter] + fn peer(&self) -> &str { + &self.peer + } + + /// Whether the scheduler session is still open. + fn is_running(&self) -> bool { + self.with_handle(|handle| handle.is_running()) + .unwrap_or(false) + } + + /// Block for at most `timeout_ms`, returning whether the session has ended. + /// + /// The caller loops on this so that each return hands the GIL back, which + /// is the only moment a pending Python signal handler can run. + fn wait(&self, py: Python<'_>, timeout_ms: u64) -> bool { + let waited = self.with_handle(|handle| { + py.detach(|| handle.wait_timeout(Duration::from_millis(timeout_ms))) + }); + // A shut-down executor has no handle left, and is certainly finished. + waited.unwrap_or(true) + } + + /// Ask the scheduler to stop sending work and finish what is in flight. + /// + /// Returns immediately when it can take the handle lock, which is what + /// makes it safe on the signal-handling path CPython actually uses: the + /// handler runs on the main thread only between `wait` calls, when the lock + /// is free. A caller on any *other* thread can instead block for up to the + /// current `wait`'s `timeout_ms`, since `wait` holds the lock across it. + fn stop(&self) { + self.with_handle(ExecutorHandle::stop); + } + + /// Drain in-flight work, disconnect, and join. Idempotent. + fn shutdown(&self, py: Python<'_>) { + let handle = self + .handle + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .take(); + if let Some(handle) = handle { + py.detach(|| handle.shutdown()); + } + } +} + +impl PyExecutor { + /// Run `body` against the live handle, or `None` once shut down. + fn with_handle(&self, body: impl FnOnce(&ExecutorHandle) -> T) -> Option { + self.handle + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .as_ref() + .map(body) + } +} diff --git a/crates/taskito-python/src/lib.rs b/crates/taskito-python/src/lib.rs index bdf4beb41..c006f400a 100644 --- a/crates/taskito-python/src/lib.rs +++ b/crates/taskito-python/src/lib.rs @@ -2,6 +2,7 @@ use pyo3::prelude::*; #[cfg(not(feature = "native-async"))] mod async_worker; +mod executor; #[cfg(feature = "native-async")] mod native_async; mod prefork; @@ -51,6 +52,7 @@ fn _taskito(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; m.add_class::()?; m.add_class::()?; + m.add_class::()?; #[cfg(feature = "native-async")] { m.add_class::()?; diff --git a/sdks/java/processor/src/main/java/org/byteveda/taskito/processor/TaskHandlerProcessor.java b/sdks/java/processor/src/main/java/org/byteveda/taskito/processor/TaskHandlerProcessor.java index bf1e8a478..031edc4ba 100644 --- a/sdks/java/processor/src/main/java/org/byteveda/taskito/processor/TaskHandlerProcessor.java +++ b/sdks/java/processor/src/main/java/org/byteveda/taskito/processor/TaskHandlerProcessor.java @@ -3,6 +3,7 @@ import java.io.IOException; import java.io.Writer; import java.util.LinkedHashMap; +import java.util.LinkedHashSet; import java.util.List; import java.util.Map; import java.util.Objects; @@ -20,7 +21,9 @@ import javax.lang.model.element.TypeElement; import javax.lang.model.element.VariableElement; import javax.tools.Diagnostic; +import javax.tools.FileObject; import javax.tools.JavaFileObject; +import javax.tools.StandardLocation; import org.jspecify.annotations.Nullable; /** @@ -28,10 +31,19 @@ * companion holding a typed {@code Task} constant per handler plus a static * {@code bind(Worker.Builder, )}. Reads the annotation structurally, so it * needs no dependency on the runtime module. + * + *

When the annotated class is top-level with an accessible no-arg constructor, + * the companion also gets a nested {@code Provider} class and is listed in + * {@code META-INF/services}, which is what lets {@code taskito executor} find + * handlers on the classpath with no user {@code main}. */ @SupportedAnnotationTypes(TaskHandlerProcessor.ANNOTATION) public final class TaskHandlerProcessor extends AbstractProcessor { static final String ANNOTATION = "org.byteveda.taskito.annotation.TaskHandler"; + + /** Service the generated providers are registered under, for {@link java.util.ServiceLoader}. */ + static final String SERVICE = "org.byteveda.taskito.worker.HandlerRegistryProvider"; + static final String RESOURCE = "org.byteveda.taskito.annotation.Resource"; static final String COMPRESSED = "org.byteveda.taskito.annotation.Compressed"; static final String ENCRYPTED = "org.byteveda.taskito.annotation.Encrypted"; @@ -43,6 +55,15 @@ public SourceVersion getSupportedSourceVersion() { return SourceVersion.latestSupported(); } + /** + * Companions that can be service-loaded, accumulated across rounds. + * + *

The service file has to be written once, in the final round: the Filer + * rejects reopening a resource it has already created, so a per-round write + * would fail as soon as generated code triggered a second round. + */ + private final Set providers = new LinkedHashSet<>(); + @Override public boolean process(Set annotations, RoundEnvironment roundEnv) { TypeElement marker = processingEnv.getElementUtils().getTypeElement(ANNOTATION); @@ -63,9 +84,72 @@ public boolean process(Set annotations, RoundEnvironment byClass.computeIfAbsent(owner, key -> new java.util.ArrayList<>()).add(method); } byClass.forEach(this::generate); + if (roundEnv.processingOver()) { + writeServiceFile(); + } return true; } + /** + * List the generated companions in {@code META-INF/services} so + * {@code ServiceLoader.load(HandlerRegistryProvider.class)} finds them. + */ + private void writeServiceFile() { + if (providers.isEmpty()) { + return; + } + try { + FileObject file = processingEnv + .getFiler() + .createResource(StandardLocation.CLASS_OUTPUT, "", "META-INF/services/" + SERVICE); + try (Writer writer = file.openWriter()) { + for (String provider : providers) { + writer.write(provider); + writer.write("\n"); + } + } + } catch (IOException e) { + processingEnv + .getMessager() + .printMessage( + Diagnostic.Kind.ERROR, + "failed to write META-INF/services/" + SERVICE + ": " + e.getMessage()); + } + } + + /** + * Whether {@code owner} can be built by {@link java.util.ServiceLoader} — a + * top-level, non-abstract type with an accessible no-arg constructor. + * + *

A class with only injected dependencies cannot be, and only the user's + * own code knows how to build it; that stays the explicit + * {@code register(...)} path. + * + *

Every nested owner is excluded, static ones included: the companion is + * written at package level and names the handler by its simple name, which + * does not resolve for an {@code Outer.Inner}. + */ + private boolean isServiceLoadable(TypeElement owner) { + if (owner.getModifiers().contains(Modifier.ABSTRACT) + || owner.getNestingKind().isNested()) { + return false; + } + boolean declaresAny = false; + for (Element enclosed : owner.getEnclosedElements()) { + if (enclosed.getKind() != ElementKind.CONSTRUCTOR) { + continue; + } + declaresAny = true; + ExecutableElement constructor = (ExecutableElement) enclosed; + if (constructor.getParameters().isEmpty() + && !constructor.getModifiers().contains(Modifier.PRIVATE)) { + return true; + } + } + // No declared constructor at all means the implicit public no-arg one. + return !declaresAny; + } + private boolean validate(ExecutableElement method) { List params = method.getParameters(); if (params.isEmpty()) { @@ -114,6 +198,7 @@ private void generate(TypeElement owner, List methods) { .append("import org.byteveda.taskito.task.Task;\n") .append("import org.byteveda.taskito.worker.Handler;\n") .append("import org.byteveda.taskito.worker.HandlerRegistry;\n") + .append("import org.byteveda.taskito.worker.HandlerRegistryProvider;\n") .append("import org.byteveda.taskito.worker.Worker;\n") .append(anyResources ? "import org.byteveda.taskito.resources.Resources;\n" : "") .append("\n") @@ -172,7 +257,34 @@ private void generate(TypeElement owner, List methods) { .append(")"); out.append(i + 1 < methods.size() ? ",\n" : ");\n"); } - out.append(" }\n}\n"); + out.append(" }\n"); + + // The ServiceLoader entry point. On the classpath a listed class must be + // a subtype of the service with a public no-arg constructor — the static + // `provider()` form only applies to modules — so this is a real nested + // implementation rather than a factory method on the companion. + if (isServiceLoadable(owner)) { + out.append("\n /** Discovered by {@code taskito executor} via ServiceLoader. */\n") + .append(" public static final class Provider implements HandlerRegistryProvider {\n") + .append(" public Provider() {}\n\n") + .append(" @Override\n") + .append(" public HandlerRegistry registry() {\n return handlers(new ") + .append(ownerSimple) + .append("());\n }\n }\n"); + // Binary name: ServiceLoader resolves the nested class by `$` form. + providers.add(qualified + "$Provider"); + } else { + processingEnv + .getMessager() + .printMessage( + Diagnostic.Kind.NOTE, + ownerSimple + + " is not a top-level type with an accessible no-arg constructor, so it" + + " cannot be discovered by `taskito executor`; register its handlers" + + " explicitly with " + companion + ".handlers(impl)", + owner); + } + out.append("}\n"); write(owner, qualified, out.toString()); } diff --git a/sdks/java/src/main/java/org/byteveda/taskito/cli/Cli.java b/sdks/java/src/main/java/org/byteveda/taskito/cli/Cli.java index 481a9a877..ce2998ecd 100644 --- a/sdks/java/src/main/java/org/byteveda/taskito/cli/Cli.java +++ b/sdks/java/src/main/java/org/byteveda/taskito/cli/Cli.java @@ -4,6 +4,7 @@ import java.util.List; import java.util.concurrent.Callable; import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; import org.byteveda.taskito.Taskito; import org.byteveda.taskito.dashboard.DashboardServer; import org.byteveda.taskito.model.DeadJob; @@ -29,7 +30,8 @@ Cli.Pause.class, Cli.Resume.class, Cli.Dlq.class, - Cli.Dashboard.class + Cli.Dashboard.class, + Cli.Executor.class }) public final class Cli { static final ObjectMapper JSON = new ObjectMapper(); @@ -252,6 +254,147 @@ public Integer call() { } } + @Command(name = "executor", description = "Run tasks for a detached scheduler instead of polling storage.") + static final class Executor implements Callable { + /** How long a shutdown hook waits for the drain before letting the JVM go. */ + private static final int DRAIN_WAIT_SECONDS = 40; + + @CommandLine.Spec + CommandLine.Model.CommandSpec spec; + + @Option( + names = "--attach", + description = "Scheduler address: host:port, :port, or unix:/path (env: TASKITO_ATTACH).") + @Nullable + String attach; + + @Option(names = "--slots", description = "Jobs to run concurrently (env: TASKITO_SLOTS).") + @Nullable + Integer slots; + + @Option(names = "--executor-id", description = "Identity announced to the scheduler.") + @Nullable + String executorId; + + @Override + public Integer call() throws Exception { + String address = attach != null ? attach : System.getenv("TASKITO_ATTACH"); + if (address == null || address.isBlank()) { + System.err.println("--attach is required (or set TASKITO_ATTACH), e.g. --attach scheduler:7749"); + return 1; + } + + int slotCount = resolveSlots(); + org.byteveda.taskito.worker.Executor.Builder builder = org.byteveda.taskito.worker.Executor.builder() + // Handlers come from META-INF/services, so no application + // code has to run to register them. + .discover() + .attach(address) + .slots(slotCount) + // Env only, never a flag: a token in argv is visible in `ps` + // output and lands in shell history. + .token(envOrNull("TASKITO_ATTACH_TOKEN")) + .executorId(executorId); + + if (builder.tasks().isEmpty()) { + System.err.println("no handlers found on the classpath. Annotate methods with @TaskHandler and " + + "make sure the annotation processor ran, or register them from your own main."); + return 1; + } + + try (org.byteveda.taskito.worker.Executor executor = builder.start()) { + System.out.printf( + "taskito executor %s attached to %s at %s (%d slot(s), %d task(s)) — Ctrl-C to stop%n", + executor.executorId(), + executor.schedulerId(), + executor.peer(), + slotCount, + builder.tasks().size()); + // A shutdown hook does not hold the JVM open, so stopping alone + // would let the process exit while the drain was still running + // and strand in-flight jobs for the reaper. Signal, then wait + // for the main thread to finish close(). + Thread main = Thread.currentThread(); + Thread hook = new Thread(() -> { + executor.stop(); + try { + main.join(TimeUnit.SECONDS.toMillis(DRAIN_WAIT_SECONDS)); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } + }); + Runtime.getRuntime().addShutdownHook(hook); + try { + executor.awaitSession(); + } finally { + // Dropped before the normal exit path reaches it. `main` + // calls `System.exit`, which runs hooks while the main + // thread is still inside `Runtime.exit` — so a hook left + // registered would wait out the whole drain budget for a + // join that cannot complete, then call `stop()` on an + // already-closed control. + dropHook(hook); + } + } + System.out.println("taskito executor detached"); + return 0; + } + + /** + * Unregister the drain hook, tolerating the one case that cannot: the + * JVM is already shutting down, which means the hook is running and has + * the drain in hand. + */ + private static void dropHook(Thread hook) { + try { + Runtime.getRuntime().removeShutdownHook(hook); + } catch (IllegalStateException alreadyShuttingDown) { + // Nothing to undo — the hook is doing its job right now. + } + } + + /** + * An environment value, or null when unset or blank. Compose + * files and container runtimes set empty strings freely, and an empty + * token must read as "no token" rather than be presented as one. + */ + private static @Nullable String envOrNull(String name) { + String value = System.getenv(name); + return value == null || value.isBlank() ? null : value; + } + + /** Slots from the flag, then the env, then one. */ + private int resolveSlots() { + if (slots != null) { + return atLeastOne(slots, "--slots"); + } + String raw = System.getenv("TASKITO_SLOTS"); + if (raw == null || raw.isBlank()) { + return 1; + } + int parsed; + try { + parsed = Integer.parseInt(raw.trim()); + } catch (NumberFormatException e) { + throw new CommandLine.ParameterException( + spec.commandLine(), "TASKITO_SLOTS must be an integer, got '" + raw + "'"); + } + return atLeastOne(parsed, "TASKITO_SLOTS"); + } + + /** + * Reject a count the executor would silently clamp, so the banner cannot + * announce a concurrency the executor does not run with. + */ + private int atLeastOne(int value, String source) { + if (value < 1) { + throw new CommandLine.ParameterException( + spec.commandLine(), source + " must be at least 1, got " + value); + } + return value; + } + } + @Command(name = "dashboard", description = "Serve the dashboard until interrupted.") static final class Dashboard implements Callable { @ParentCommand diff --git a/sdks/java/src/main/java/org/byteveda/taskito/internal/JniExecutorControl.java b/sdks/java/src/main/java/org/byteveda/taskito/internal/JniExecutorControl.java new file mode 100644 index 000000000..f1fb97359 --- /dev/null +++ b/sdks/java/src/main/java/org/byteveda/taskito/internal/JniExecutorControl.java @@ -0,0 +1,155 @@ +package org.byteveda.taskito.internal; + +import java.util.concurrent.locks.ReentrantReadWriteLock; +import java.util.function.Supplier; +import org.byteveda.taskito.spi.WorkerControl; + +/** + * JNI-backed {@link WorkerControl} over an attached-executor handle. + * + *

The mirror of {@link JniWorkerControl}, with the same locking contract: + * every call holds the read lock, {@code close()} takes the write lock so it + * waits out in-flight calls, and later calls throw instead of touching freed + * native memory. {@link #awaitSession()} is the one exception — it parks for as + * long as the session lasts, so it counts itself instead of holding the lock. + */ +public final class JniExecutorControl implements WorkerControl { + private final long handle; + private final ReentrantReadWriteLock stateLock = new ReentrantReadWriteLock(); + private boolean closed; // guarded by stateLock + + /** Monitor for {@link #waiting}, held only for the count. */ + private final Object waiters = new Object(); + + /** Threads inside the native session wait; the handle is freed at zero. */ + private int waiting; // guarded by waiters + + public JniExecutorControl(long handle) { + this.handle = handle; + } + + private T withOpenHandle(Supplier nativeCall) { + stateLock.readLock().lock(); + try { + if (closed) { + throw new IllegalStateException("executor control is closed"); + } + return nativeCall.get(); + } finally { + stateLock.readLock().unlock(); + } + } + + @Override + public void completeJob(long token, byte[] result) { + withOpenHandle(() -> { + NativeExecutor.completeJob(handle, token, result); + return null; + }); + } + + @Override + public void failJob(long token, String error, boolean retryable) { + withOpenHandle(() -> { + NativeExecutor.failJob(handle, token, error, retryable); + return null; + }); + } + + @Override + public void cancelJob(long token) { + withOpenHandle(() -> { + NativeExecutor.cancelJob(handle, token); + return null; + }); + } + + @Override + public void stop() { + withOpenHandle(() -> { + NativeExecutor.stop(handle); + return null; + }); + } + + /** Identity the scheduler announced when it accepted the attach. */ + public String schedulerId() { + return withOpenHandle(() -> NativeExecutor.schedulerId(handle)); + } + + /** Identity this executor attached under. */ + public String executorId() { + return withOpenHandle(() -> NativeExecutor.executorId(handle)); + } + + /** Peer label of the scheduler connection. */ + public String peer() { + return withOpenHandle(() -> NativeExecutor.peer(handle)); + } + + /** Whether the scheduler session is still open. */ + public boolean isRunning() { + return withOpenHandle(() -> NativeExecutor.isRunning(handle)); + } + + /** + * Block until the scheduler ends the session. + * + *

The read lock is released before the native wait rather than held across + * it. {@link ReentrantReadWriteLock} blocks a new reader once a writer is + * queued, so a parked waiter holding the lock would leave a concurrent + * {@code stop()} stuck behind a {@code close()} that is itself waiting on the + * park — and only {@code stop()} could have ended it. The handle stays alive + * instead by counting waiters, which {@code close} drains before freeing. + */ + public void awaitSession() { + withOpenHandle(() -> { + synchronized (waiters) { + waiting++; + } + return null; + }); + try { + NativeExecutor.awaitSession(handle); + } finally { + synchronized (waiters) { + waiting--; + waiters.notifyAll(); + } + } + } + + /** Idempotent: drains and frees the native handle exactly once. */ + @Override + public void close() { + stateLock.writeLock().lock(); + try { + if (closed) { + return; + } + closed = true; + } finally { + stateLock.writeLock().unlock(); + } + + // Ends the session so anyone parked in `awaitSession` returns, then waits + // them out: the handle must not be freed while a native call still holds + // it. Interrupts are deferred rather than obeyed — leaving early would + // free the handle under a parked thread. + NativeExecutor.stop(handle); + boolean interrupted = false; + synchronized (waiters) { + while (waiting > 0) { + try { + waiters.wait(); + } catch (InterruptedException e) { + interrupted = true; + } + } + } + NativeExecutor.close(handle); + if (interrupted) { + Thread.currentThread().interrupt(); + } + } +} diff --git a/sdks/java/src/main/java/org/byteveda/taskito/internal/MiddlewareDisables.java b/sdks/java/src/main/java/org/byteveda/taskito/internal/MiddlewareDisables.java index d217e6293..a7f260840 100644 --- a/sdks/java/src/main/java/org/byteveda/taskito/internal/MiddlewareDisables.java +++ b/sdks/java/src/main/java/org/byteveda/taskito/internal/MiddlewareDisables.java @@ -7,6 +7,7 @@ import org.byteveda.taskito.logging.TaskitoLogger; import org.byteveda.taskito.middleware.Middleware; import org.byteveda.taskito.spi.QueueBackend; +import org.jspecify.annotations.Nullable; /** * Per-task middleware disable list, persisted under @@ -26,9 +27,14 @@ public final class MiddlewareDisables { private static final String KEY_PREFIX = "middleware:disabled:"; - private final QueueBackend backend; + /** + * Absent for an attached executor, which has no storage to read toggles + * from; nothing is then disabled, which is the same answer an unconfigured + * queue gives. + */ + private final @Nullable QueueBackend backend; - public MiddlewareDisables(QueueBackend backend) { + public MiddlewareDisables(@Nullable QueueBackend backend) { this.backend = backend; } @@ -48,6 +54,9 @@ public static String nameOf(Middleware middleware) { /** Names disabled for {@code taskName}; empty when none are, or the list is unreadable. */ public List disabledFor(String taskName) { + if (backend == null) { + return List.of(); + } Optional raw = backend.getSetting(key(taskName)); if (raw.isEmpty()) { return List.of(); diff --git a/sdks/java/src/main/java/org/byteveda/taskito/internal/NativeExecutor.java b/sdks/java/src/main/java/org/byteveda/taskito/internal/NativeExecutor.java new file mode 100644 index 000000000..d551ed86b --- /dev/null +++ b/sdks/java/src/main/java/org/byteveda/taskito/internal/NativeExecutor.java @@ -0,0 +1,46 @@ +package org.byteveda.taskito.internal; + +/** + * JNI surface for an executor attached to a detached scheduler. + * + *

The mirror of {@link NativeWorker}: the same token-based completion, over a + * socket to a scheduler rather than over this process's own storage. The + * {@code handle} comes from {@link #attach}; it stays valid until {@link #close}. + */ +public final class NativeExecutor { + static { + NativeLoader.load(); + } + + private NativeExecutor() {} + + /** Dial, handshake, and start running jobs. {@code bridge} is a {@code WorkerBridge}. */ + public static native long attach(Object bridge, String optionsJson); + + public static native void completeJob(long handle, long token, byte[] result); + + public static native void failJob(long handle, long token, String error, boolean retryable); + + public static native void cancelJob(long handle, long token); + + /** Identity the scheduler announced when it accepted the attach. */ + public static native String schedulerId(long handle); + + /** Identity this executor attached under. */ + public static native String executorId(long handle); + + /** Peer label of the scheduler connection. */ + public static native String peer(long handle); + + /** Whether the scheduler session is still open. */ + public static native boolean isRunning(long handle); + + /** Block until the scheduler ends the session. Parks the calling thread. */ + public static native void awaitSession(long handle); + + /** Stop accepting work and finish what is in flight. Returns at once. */ + public static native void stop(long handle); + + /** Drain, disconnect, and reclaim the handle. */ + public static native void close(long handle); +} diff --git a/sdks/java/src/main/java/org/byteveda/taskito/worker/Executor.java b/sdks/java/src/main/java/org/byteveda/taskito/worker/Executor.java new file mode 100644 index 000000000..1636b7811 --- /dev/null +++ b/sdks/java/src/main/java/org/byteveda/taskito/worker/Executor.java @@ -0,0 +1,360 @@ +package org.byteveda.taskito.worker; + +import com.fasterxml.jackson.databind.ObjectMapper; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.ServiceLoader; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.TimeUnit; +import java.util.function.Consumer; +import org.byteveda.taskito.events.Emitter; +import org.byteveda.taskito.events.EventName; +import org.byteveda.taskito.events.TaskitoEvent; +import org.byteveda.taskito.events.WorkerEvent; +import org.byteveda.taskito.internal.JniExecutorControl; +import org.byteveda.taskito.internal.NativeExecutor; +import org.byteveda.taskito.logging.TaskitoLogger; +import org.byteveda.taskito.middleware.Middleware; +import org.byteveda.taskito.resources.ResourceRuntime; +import org.byteveda.taskito.serialization.JsonSerializer; +import org.byteveda.taskito.serialization.PayloadCodec; +import org.byteveda.taskito.serialization.Serializer; +import org.jspecify.annotations.Nullable; + +/** + * Runs tasks for a detached scheduler instead of polling storage. + * + *

The inverse of {@link Worker}: the scheduler holds the database connection + * and dispatches jobs over a socket, so this process runs task bodies without + * any database credentials of its own. Handlers, middleware, codecs and the + * cancel signal behave exactly as they do in a worker — only the transport + * differs. + * + *

{@code
+ * try (Executor executor = Executor.builder()
+ *         .discover()                 // handlers from META-INF/services
+ *         .attach("scheduler:7749")
+ *         .slots(4)
+ *         .start()) {
+ *     executor.awaitSession();        // until the scheduler shuts down
+ * }
+ * }
+ */ +public final class Executor implements AutoCloseable { + private static final TaskitoLogger LOG = TaskitoLogger.create("executor"); + private static final int SHUTDOWN_TIMEOUT_SECONDS = 30; + + private final JniExecutorControl control; + private final ExecutorService handlerPool; + private final ResourceRuntime resources; + private final Emitter emitter; + private boolean closed; + + private Executor( + JniExecutorControl control, ExecutorService handlerPool, ResourceRuntime resources, Emitter emitter) { + this.control = control; + this.handlerPool = handlerPool; + this.resources = resources; + this.emitter = emitter; + } + + /** A builder for an executor. Register handlers, point it at a scheduler, start. */ + public static Builder builder() { + return new Builder(); + } + + /** Identity the scheduler announced when it accepted this attach. */ + public String schedulerId() { + return control.schedulerId(); + } + + /** Identity this executor attached under. */ + public String executorId() { + return control.executorId(); + } + + /** Peer label of the scheduler connection. */ + public String peer() { + return control.peer(); + } + + /** Whether this executor is still accepting work. */ + public boolean isRunning() { + return control.isRunning(); + } + + /** + * Block until this executor stops accepting work — the scheduler ending the + * session, or a local {@link #stop()}. Does not drain; {@link #close()} does. + */ + public void awaitSession() { + control.awaitSession(); + } + + /** + * Ask the scheduler to stop sending work and finish what is in flight. + * Returns at once, so it is safe from a shutdown hook. + */ + public void stop() { + control.stop(); + } + + /** + * Drain in-flight work, disconnect, and release handler threads. + * + *

The handler pool is drained before the native handle is freed: a handler + * may still be completing its job through {@code control}, and the handle has + * to outlive every such call. Idempotent. + */ + @Override + public void close() { + if (closed) { + return; + } + closed = true; + control.stop(); + handlerPool.shutdown(); // stop accepting; let running handlers finish + try { + if (!handlerPool.awaitTermination(SHUTDOWN_TIMEOUT_SECONDS, TimeUnit.SECONDS)) { + handlerPool.shutdownNow(); + if (!handlerPool.awaitTermination(SHUTDOWN_TIMEOUT_SECONDS, TimeUnit.SECONDS)) { + LOG.warn("handler threads still running after " + (2 * SHUTDOWN_TIMEOUT_SECONDS) + + "s; closing the executor handle anyway"); + } + } + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } + try { + control.close(); + } finally { + resources.teardownWorker(); + emitter.emit(new WorkerEvent(EventName.WORKER_STOPPED, List.of())); + } + } + + /** Registers handlers and attach options, then starts the executor. */ + public static final class Builder { + private static final ObjectMapper JSON = new ObjectMapper(); + + private final Map handlers = new HashMap<>(); + private final Map>> listeners = new LinkedHashMap<>(); + private Serializer serializer = new JsonSerializer(); + private List middleware = List.of(); + private Map codecs = Map.of(); + private ResourceRuntime resources = new ResourceRuntime(); + private @Nullable String address; + private @Nullable String token; + private @Nullable String executorId; + private int slots = 1; + private @Nullable Long connectTimeoutMs; + private @Nullable Long heartbeatIntervalMs; + private @Nullable Long shutdownDrainMs; + + private Builder() {} + + /** + * Register every handler discoverable on the classpath. + * + *

The {@code @TaskHandler} processor lists a provider per annotated + * class in {@code META-INF/services}, so this needs no application code + * at all. A class the executor cannot construct is skipped at build time + * with a compiler note — register those explicitly. + */ + public Builder discover() { + return discover(Thread.currentThread().getContextClassLoader()); + } + + /** {@link #discover()} against a specific class loader. */ + public Builder discover(ClassLoader loader) { + for (HandlerRegistryProvider provider : ServiceLoader.load(HandlerRegistryProvider.class, loader)) { + register(provider.registry()); + } + return this; + } + + /** Register a generated {@code Tasks.handlers(impl)} bundle. */ + public Builder register(HandlerRegistry registry) { + registry.handlers().forEach(this::register); + return this; + } + + /** Register one handler. */ + public Builder register(Handler handler) { + handlers.put( + handler.task().name(), + new RegisteredTask( + handler.task().payloadType(), + cast(handler.function()), + handler.task().codecNames(), + handler.task().retryOn())); + return this; + } + + /** Scheduler address: {@code host:port}, {@code :port}, or {@code unix:/path}. */ + public Builder attach(String address) { + this.address = address; + return this; + } + + /** Jobs to run at once (default 1). */ + public Builder slots(int slots) { + this.slots = Math.max(1, slots); + return this; + } + + /** Shared secret, when the scheduler requires one. */ + public Builder token(@Nullable String token) { + this.token = token; + return this; + } + + /** Identity announced to the scheduler (default: generated per process). */ + public Builder executorId(@Nullable String executorId) { + this.executorId = executorId; + return this; + } + + /** How long to wait for the connection (default 10000ms). */ + public Builder connectTimeoutMs(long millis) { + this.connectTimeoutMs = millis; + return this; + } + + /** How often to send a liveness heartbeat (default 5000ms). */ + public Builder heartbeatIntervalMs(long millis) { + this.heartbeatIntervalMs = millis; + return this; + } + + /** How long a drain waits for in-flight jobs (default 30000ms). */ + public Builder shutdownDrainMs(long millis) { + this.shutdownDrainMs = millis; + return this; + } + + /** Payload serializer (default JSON). Must match the enqueuing side. */ + public Builder serializer(Serializer serializer) { + this.serializer = serializer; + return this; + } + + /** Middleware applied around every handler. */ + public Builder middleware(List middleware) { + this.middleware = List.copyOf(middleware); + return this; + } + + /** Named payload codecs, for tasks declaring {@code codecs}. */ + public Builder codecs(Map codecs) { + this.codecs = Map.copyOf(codecs); + return this; + } + + /** Injectable resources available to handlers. */ + public Builder resources(ResourceRuntime resources) { + this.resources = resources; + return this; + } + + /** Subscribe to a worker lifecycle or job event. */ + public Builder on(EventName name, Consumer listener) { + listeners.computeIfAbsent(name, key -> new ArrayList<>()).add(listener); + return this; + } + + /** Task names this executor will advertise. */ + public List tasks() { + return List.copyOf(handlers.keySet()); + } + + /** + * Dial the scheduler and start running jobs. + * + *

The handshake happens here, so a bad token or an unreachable + * scheduler throws before any handler thread exists. + */ + public Executor start() { + if (address == null || address.isBlank()) { + throw new IllegalStateException( + "no scheduler address: call attach(...) or set TASKITO_ATTACH (e.g. scheduler:7749)"); + } + if (handlers.isEmpty()) { + // The scheduler only dispatches task names an executor + // advertises, so this would attach and then sit idle forever. + throw new IllegalStateException("no handlers registered: call discover() or register(...)"); + } + + ExecutorService pool = Executors.newFixedThreadPool(slots); + Emitter emitter = new Emitter(); + listeners.forEach((name, bound) -> bound.forEach(listener -> emitter.onEvent(name, listener))); + + // No QueueBackend: an executor reads no storage, which is the point + // of the split. Job metadata and dashboard middleware toggles are + // storage-backed and so unavailable here. + WorkerDispatchBridge bridge = + new WorkerDispatchBridge(null, handlers, serializer, pool, emitter, middleware, resources, codecs); + + long handle; + try { + handle = NativeExecutor.attach(bridge, encodeOptions()); + } catch (RuntimeException e) { + pool.shutdownNow(); + throw e; + } + JniExecutorControl control = new JniExecutorControl(handle); + try { + bridge.bind(control); + emitter.emit(new WorkerEvent(EventName.WORKER_STARTED, List.of())); + // Lease worker resources only after the attach succeeded, so a + // refused handshake leaks nothing. + resources.acquireWorker(); + } catch (RuntimeException e) { + // The attach already succeeded, so the socket, the reader thread + // and the native handle are live; a throwing listener or + // resource factory must not strand them for the process's life. + control.close(); + pool.shutdownNow(); + throw e; + } + return new Executor(control, pool, resources, emitter); + } + + /** The attach options, as the JSON the native side parses. */ + private String encodeOptions() { + Map options = new LinkedHashMap<>(); + options.put("address", address); + options.put("tasks", tasks()); + options.put("slots", slots); + if (token != null) { + options.put("token", token); + } + if (executorId != null) { + options.put("executorId", executorId); + } + if (connectTimeoutMs != null) { + options.put("connectTimeoutMs", connectTimeoutMs); + } + if (heartbeatIntervalMs != null) { + options.put("heartbeatIntervalMs", heartbeatIntervalMs); + } + if (shutdownDrainMs != null) { + options.put("shutdownDrainMs", shutdownDrainMs); + } + try { + return JSON.writeValueAsString(options); + } catch (Exception e) { + throw new IllegalStateException("failed to encode executor options", e); + } + } + + @SuppressWarnings("unchecked") + private static org.byteveda.taskito.task.TaskFunction cast(Object function) { + return (org.byteveda.taskito.task.TaskFunction) function; + } + } +} diff --git a/sdks/java/src/main/java/org/byteveda/taskito/worker/HandlerRegistryProvider.java b/sdks/java/src/main/java/org/byteveda/taskito/worker/HandlerRegistryProvider.java new file mode 100644 index 000000000..c41cc90d0 --- /dev/null +++ b/sdks/java/src/main/java/org/byteveda/taskito/worker/HandlerRegistryProvider.java @@ -0,0 +1,25 @@ +package org.byteveda.taskito.worker; + +/** + * A {@link HandlerRegistry} discoverable on the classpath. + * + *

The {@code @TaskHandler} processor generates one implementation per + * annotated class and lists it in {@code META-INF/services}, which is what lets + * {@code taskito executor} find handlers with no user {@code main} calling + * {@code register(...)}. + * + *

The indirection is not incidental. {@link java.util.ServiceLoader} can only + * instantiate a listed class when it is a subtype of the service with a public + * no-arg constructor — the static {@code provider()} form works for modules on + * the module path, not for a classpath {@code java -cp app.jar}. {@link + * HandlerRegistry} is a final value type built from a user instance, so it can + * be neither; this interface is the thing that can. + * + *

Implement it by hand when a handler class needs constructor arguments the + * executor cannot supply — the processor skips those, since only your own code + * knows how to build them. + */ +public interface HandlerRegistryProvider { + /** The handlers this provider contributes. Called once per executor start. */ + HandlerRegistry registry(); +} diff --git a/sdks/java/src/main/java/org/byteveda/taskito/worker/WorkerDispatchBridge.java b/sdks/java/src/main/java/org/byteveda/taskito/worker/WorkerDispatchBridge.java index 7bc14891b..738f135f9 100644 --- a/sdks/java/src/main/java/org/byteveda/taskito/worker/WorkerDispatchBridge.java +++ b/sdks/java/src/main/java/org/byteveda/taskito/worker/WorkerDispatchBridge.java @@ -38,7 +38,12 @@ final class WorkerDispatchBridge implements WorkerBridge { private static final ObjectMapper JSON = new ObjectMapper(); private static final TypeReference> MAP = new TypeReference>() {}; - private final QueueBackend backend; + /** + * Absent for an attached executor, which reads no storage: job metadata is + * then unavailable to middleware, and the toggle list is empty. + */ + private final @Nullable QueueBackend backend; + private final Map handlers; private final Serializer serializer; private final ExecutorService executor; @@ -51,7 +56,7 @@ final class WorkerDispatchBridge implements WorkerBridge { private final CompletableFuture control = new CompletableFuture<>(); WorkerDispatchBridge( - QueueBackend backend, + @Nullable QueueBackend backend, Map handlers, Serializer serializer, ExecutorService executor, @@ -188,6 +193,9 @@ private static void dispatch(Middleware m, EventName name, OutcomeEvent event) { /** Lazily load a job's metadata blob into a map (empty on absence/parse failure). */ private Map loadMetadata(String jobId) { + if (backend == null) { + return Collections.emptyMap(); + } try { JsonNode view = backend.getJobJson(jobId) .map(WorkerDispatchBridge::readTree) diff --git a/sdks/java/src/test/java/org/byteveda/taskito/worker/ExecutorAttachTest.java b/sdks/java/src/test/java/org/byteveda/taskito/worker/ExecutorAttachTest.java new file mode 100644 index 000000000..c690e833b --- /dev/null +++ b/sdks/java/src/test/java/org/byteveda/taskito/worker/ExecutorAttachTest.java @@ -0,0 +1,367 @@ +package org.byteveda.taskito.worker; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; +import java.io.ByteArrayOutputStream; +import java.io.IOException; +import java.io.InputStream; +import java.io.OutputStream; +import java.net.InetAddress; +import java.net.ServerSocket; +import java.net.Socket; +import java.nio.charset.StandardCharsets; +import java.util.ArrayDeque; +import java.util.Deque; +import java.util.Map; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicReference; +import org.byteveda.taskito.task.Task; +import org.jspecify.annotations.Nullable; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.Timeout; + +/** + * End-to-end tests for {@link Executor} against a socket speaking the frame + * protocol, so they need no `taskito-server` build. + * + *

The wire is the contract: a job frame in, a result frame out. Asserting on + * the frames rather than on storage is what makes these runnable anywhere. + */ +class ExecutorAttachTest { + private static final ObjectMapper JSON = new ObjectMapper(); + private static final int PROTOCOL_VERSION = 1; + private static final long SETTLE_MS = 20_000; + + private @Nullable FakeScheduler scheduler; + private @Nullable Executor executor; + + @AfterEach + void tearDown() throws Exception { + if (executor != null) { + executor.close(); + executor = null; + } + if (scheduler != null) { + scheduler.close(); + scheduler = null; + } + } + + /** The scheduler end of an attach, driven frame by frame. */ + private static final class FakeScheduler implements AutoCloseable { + private final ServerSocket server; + private final Thread accepting; + private final CountDownLatch connected = new CountDownLatch(1); + private final AtomicReference<@Nullable Socket> socket = new AtomicReference<>(); + private final AtomicReference<@Nullable JsonNode> hello = new AtomicReference<>(); + private final Deque results = new ArrayDeque<>(); + private final boolean refuse; + private @Nullable InputStream in; + private @Nullable OutputStream out; + + FakeScheduler(boolean refuse) throws IOException { + this.refuse = refuse; + this.server = new ServerSocket(0, 1, InetAddress.getLoopbackAddress()); + this.accepting = new Thread(this::accept, "fake-scheduler"); + this.accepting.setDaemon(true); + this.accepting.start(); + } + + int port() { + return server.getLocalPort(); + } + + private void accept() { + try { + Socket client = server.accept(); + socket.set(client); + in = client.getInputStream(); + out = client.getOutputStream(); + JsonNode frame = readFrame(); + hello.set(frame); + if (refuse) { + client.close(); + } else { + send(Map.of( + "type", + "hello_ack", + "scheduler_id", + "fake-scheduler", + "protocol_version", + PROTOCOL_VERSION)); + } + connected.countDown(); + } catch (IOException e) { + connected.countDown(); + } + } + + /** Read one frame: a JSON header line, then exactly the payload it declares. */ + private @Nullable JsonNode readFrame() throws IOException { + InputStream stream = in; + if (stream == null) { + return null; + } + ByteArrayOutputStream header = new ByteArrayOutputStream(); + int b; + while ((b = stream.read()) != -1 && b != '\n') { + header.write(b); + } + if (header.size() == 0) { + return null; + } + JsonNode node = JSON.readTree(header.toString(StandardCharsets.UTF_8)); + int declared = declaredPayloadLength(node); + if (declared > 0) { + stream.readNBytes(declared); + } + return node; + } + + private static int declaredPayloadLength(JsonNode header) { + String type = header.path("type").asText(""); + if (type.equals("job")) { + return header.path("payload_len").asInt(0); + } + if (type.equals("success")) { + JsonNode len = header.get("result_len"); + return len == null || len.isNull() ? 0 : len.asInt(0); + } + return 0; + } + + JsonNode awaitHello() throws InterruptedException { + assertTrue(connected.await(SETTLE_MS, TimeUnit.MILLISECONDS), "the executor never attached"); + JsonNode frame = hello.get(); + assertNotNull(frame, "no hello frame arrived"); + return frame; + } + + void send(Map header) throws IOException { + OutputStream stream = out; + if (stream == null) { + return; + } + stream.write(JSON.writeValueAsBytes(header)); + stream.write('\n'); + stream.flush(); + } + + void sendJob(String id, String taskName, byte[] payload) throws IOException { + OutputStream stream = out; + if (stream == null) { + return; + } + Map header = Map.of( + "type", + "job", + "id", + id, + "task_name", + taskName, + "payload_len", + payload.length, + "retry_count", + 0, + "max_retries", + 3, + "queue", + "default", + "timeout_ms", + 30_000); + stream.write(JSON.writeValueAsBytes(header)); + stream.write('\n'); + stream.write(payload); + stream.flush(); + } + + /** The next frame that is not a heartbeat. */ + JsonNode nextResult() throws IOException { + long deadline = System.nanoTime() + TimeUnit.MILLISECONDS.toNanos(SETTLE_MS); + while (System.nanoTime() < deadline) { + if (!results.isEmpty()) { + return results.removeFirst(); + } + JsonNode frame = readFrame(); + if (frame == null) { + break; + } + if (!frame.path("type").asText("").equals("heartbeat")) { + return frame; + } + } + throw new AssertionError("no result frame arrived"); + } + + @Override + public void close() throws IOException { + Socket client = socket.get(); + if (client != null) { + client.close(); + } + server.close(); + accepting.interrupt(); + } + } + + private Executor attach(Executor.Builder builder, int port) { + Executor started = builder.attach("127.0.0.1:" + port) + .heartbeatIntervalMs(50) + .shutdownDrainMs(5_000) + .start(); + executor = started; + return started; + } + + private static Executor.Builder greeter() { + Task greet = Task.of("greet", String.class); + return Executor.builder().register(Handler.of(greet, (String who) -> "hello " + who)); + } + + @Test + @Timeout(60) + void announcesItselfAndTheTasksItCanRun() throws Exception { + FakeScheduler fake = new FakeScheduler(false); + scheduler = fake; + + attach(greeter().executorId("exec-test").slots(2), fake.port()); + JsonNode hello = fake.awaitHello(); + + assertEquals("exec-test", hello.path("executor_id").asText()); + assertEquals("java", hello.path("sdk").asText()); + assertEquals(2, hello.path("slots").asInt()); + assertEquals(PROTOCOL_VERSION, hello.path("protocol_version").asInt()); + // Only advertised tasks are ever dispatched, so a missing name here is a + // job that silently never runs. + assertEquals("greet", hello.path("tasks").get(0).asText()); + // A token that was never configured must not appear on the wire. + assertTrue(hello.get("token") == null || hello.get("token").isNull()); + } + + @Test + @Timeout(60) + void runsADispatchedJobAndReturnsItsResult() throws Exception { + FakeScheduler fake = new FakeScheduler(false); + scheduler = fake; + + attach(greeter(), fake.port()); + fake.awaitHello(); + fake.sendJob("job-1", "greet", JSON.writeValueAsBytes("ada")); + + JsonNode result = fake.nextResult(); + assertEquals("success", result.path("type").asText()); + assertEquals("job-1", result.path("job_id").asText()); + } + + @Test + @Timeout(60) + void reportsAFailureWithItsRetryVerdict() throws Exception { + FakeScheduler fake = new FakeScheduler(false); + scheduler = fake; + + Task boom = Task.of("boom", String.class); + attach( + Executor.builder().register(Handler.of(boom, (String ignored) -> { + throw new IllegalStateException("deliberate failure"); + })), + fake.port()); + fake.awaitHello(); + fake.sendJob("job-1", "boom", JSON.writeValueAsBytes("x")); + + JsonNode result = fake.nextResult(); + assertEquals("failure", result.path("type").asText()); + assertEquals("job-1", result.path("job_id").asText()); + assertTrue(result.path("should_retry").asBoolean(), "an unclassified failure is retryable"); + assertFalse(result.path("timed_out").asBoolean()); + assertTrue(result.path("error").asText().contains("deliberate failure")); + } + + @Test + @Timeout(60) + void aShutdownFrameEndsTheSession() throws Exception { + FakeScheduler fake = new FakeScheduler(false); + scheduler = fake; + + Executor running = attach(greeter(), fake.port()); + fake.awaitHello(); + assertTrue(running.isRunning()); + + fake.send(Map.of("type", "shutdown")); + running.awaitSession(); + assertFalse(running.isRunning(), "a shutdown frame must end the session"); + } + + @Test + @Timeout(60) + void stopReleasesAParkedWaiter() throws Exception { + // `stop()` cannot unpark the frame reader, so the session has to end + // locally too — otherwise a shutdown hook that stops and then waits + // would hang instead of draining. + FakeScheduler fake = new FakeScheduler(false); + scheduler = fake; + + Executor running = attach(greeter(), fake.port()); + fake.awaitHello(); + + Thread waiter = new Thread(running::awaitSession, "await-session"); + waiter.setDaemon(true); + waiter.start(); + + running.stop(); + waiter.join(SETTLE_MS); + assertFalse(waiter.isAlive(), "stop() never released the waiter"); + } + + @Test + @Timeout(60) + void aRefusedAttachIsReportedRatherThanRetried() throws Exception { + FakeScheduler fake = new FakeScheduler(true); + scheduler = fake; + + // A wrong token is the likeliest deployment mistake; it must not surface + // as a bare network error. + RuntimeException error = assertThrows(RuntimeException.class, () -> greeter() + .attach("127.0.0.1:" + fake.port()) + .token("wrong-token") + .start()); + assertTrue( + error.getMessage().toLowerCase(java.util.Locale.ROOT).contains("refused") + || error.getMessage().toLowerCase(java.util.Locale.ROOT).contains("token"), + "expected a refusal, got: " + error.getMessage()); + } + + @Test + @Timeout(60) + void anAddressIsRequired() { + IllegalStateException error = + assertThrows(IllegalStateException.class, () -> greeter().start()); + assertTrue(error.getMessage().contains("TASKITO_ATTACH"), error.getMessage()); + } + + @Test + @Timeout(60) + void handlersAreRequired() { + IllegalStateException error = assertThrows( + IllegalStateException.class, + () -> Executor.builder().attach("127.0.0.1:1").start()); + assertTrue(error.getMessage().contains("no handlers"), error.getMessage()); + } + + @Test + @Timeout(60) + void anUnreachableSchedulerFailsFast() { + // Port 1 on loopback is reserved and nothing listens there. + RuntimeException error = assertThrows( + RuntimeException.class, + () -> greeter().attach("127.0.0.1:1").connectTimeoutMs(500).start()); + assertTrue(error.getMessage().contains("could not reach the scheduler"), error.getMessage()); + } +} diff --git a/sdks/node/src/cli/commands/executor.ts b/sdks/node/src/cli/commands/executor.ts new file mode 100644 index 000000000..4363b8df2 --- /dev/null +++ b/sdks/node/src/cli/commands/executor.ts @@ -0,0 +1,72 @@ +import type { Command } from "commander"; +import { DETACHED_ENV } from "../../detached"; +import { loadApp } from "../load-app"; +import { positiveIntFlag } from "../parse"; + +interface ExecutorOptions { + attach?: string; + slots?: string; + executorId?: string; + connectTimeout?: string; + drainTimeout?: string; +} + +export function registerExecutor(program: Command): void { + program + .command("executor ") + .description( + "Run tasks for a detached scheduler. is a module exporting a configured Queue " + + "(default export or `queue`).", + ) + .option( + "--attach

", + "scheduler address: host:port, :port, or unix:/path (env: TASKITO_ATTACH)", + ) + .option("--slots ", "jobs to run concurrently (env: TASKITO_SLOTS)") + .option("--executor-id ", "identity announced to the scheduler") + .option("--connect-timeout ", "how long to wait for the connection") + .option("--drain-timeout ", "how long a drain waits for in-flight jobs") + .action(async (appPath: string, options: ExecutorOptions) => { + // Set before the app is imported: building a Queue is what opens storage, + // and an executor must not. + process.env[DETACHED_ENV] = "1"; + + const app = await loadApp(appPath); + // The token is read from the environment inside `runExecutor`, never as a + // flag: in argv it would show up in `ps` output and shell history. + const executor = await app.runExecutor({ + attach: options.attach, + slots: positiveIntFlag(options.slots, "slots"), + executorId: options.executorId, + connectTimeoutMs: positiveIntFlag(options.connectTimeout, "connect-timeout"), + shutdownDrainMs: positiveIntFlag(options.drainTimeout, "drain-timeout"), + }); + + process.stdout.write( + `taskito executor ${executor.executorId} attached to ${executor.schedulerId} ` + + `at ${executor.peer} — Ctrl-C to stop\n`, + ); + + // `stop()` drains in-flight work and disconnects; it is memoized, so the + // signal path and a scheduler-initiated shutdown cannot tear down twice. + // The handler cannot be `async`: nothing awaits what a signal listener + // returns, so a rejected `stop()` would surface as an unhandled rejection + // and abort the process mid-drain instead of finishing it. + const stop = (): void => { + executor.stop().then( + () => process.exit(0), + (error: unknown) => { + process.stderr.write(`taskito executor failed to stop cleanly: ${String(error)}\n`); + process.exit(1); + }, + ); + }; + process.once("SIGINT", stop); + process.once("SIGTERM", stop); + + // Resolves when the scheduler ends the session, so a `taskito-server` + // shutting down takes its executors with it rather than stranding them. + await executor.wait(); + await executor.stop(); + }); +} diff --git a/sdks/node/src/cli/commands/index.ts b/sdks/node/src/cli/commands/index.ts index b6473ba6d..7f1649829 100644 --- a/sdks/node/src/cli/commands/index.ts +++ b/sdks/node/src/cli/commands/index.ts @@ -3,6 +3,7 @@ export { registerCancel } from "./cancel"; export { registerDashboard } from "./dashboard"; export { registerDlq } from "./dlq"; export { registerEnqueue } from "./enqueue"; +export { registerExecutor } from "./executor"; export { registerJobs } from "./jobs"; export { registerQueues } from "./queues"; export { registerRun } from "./run"; diff --git a/sdks/node/src/cli/commands/run.ts b/sdks/node/src/cli/commands/run.ts index 75c755461..9cdccb68a 100644 --- a/sdks/node/src/cli/commands/run.ts +++ b/sdks/node/src/cli/commands/run.ts @@ -1,7 +1,5 @@ -import { resolve } from "node:path"; -import { pathToFileURL } from "node:url"; import type { Command } from "commander"; -import type { Worker } from "../../index"; +import { loadApp } from "../load-app"; import { positiveIntFlag } from "../parse"; /** Grace period for in-flight results to drain after a stop signal. */ @@ -12,11 +10,6 @@ interface RunOptions { batchSize?: string; } -/** The minimal surface `run` needs from a user's app module. */ -interface WorkerApp { - runWorker(options?: { queues?: string[]; batchSize?: number }): Worker; -} - export function registerRun(program: Command): void { program .command("run ") @@ -50,13 +43,3 @@ export function registerRun(program: Command): void { await new Promise(() => {}); }); } - -/** Import the user's app module and return its configured queue. */ -async function loadApp(appPath: string): Promise { - const module = (await import(pathToFileURL(resolve(appPath)).href)) as Record; - const candidate = module.default ?? module.queue; - if (!candidate || typeof (candidate as WorkerApp).runWorker !== "function") { - throw new Error(`module "${appPath}" must export a Queue (default export or \`queue\`)`); - } - return candidate as WorkerApp; -} diff --git a/sdks/node/src/cli/index.ts b/sdks/node/src/cli/index.ts index d47b5b768..4639aaad2 100644 --- a/sdks/node/src/cli/index.ts +++ b/sdks/node/src/cli/index.ts @@ -9,6 +9,7 @@ import { registerDashboard, registerDlq, registerEnqueue, + registerExecutor, registerJobs, registerQueues, registerRun, @@ -38,6 +39,7 @@ registerCancel(program); registerQueues(program); registerDlq(program); registerRun(program); +registerExecutor(program); registerDashboard(program); registerScaler(program); registerAutoscale(program); diff --git a/sdks/node/src/cli/load-app.ts b/sdks/node/src/cli/load-app.ts new file mode 100644 index 000000000..d29dfd122 --- /dev/null +++ b/sdks/node/src/cli/load-app.ts @@ -0,0 +1,18 @@ +import { resolve } from "node:path"; +import { pathToFileURL } from "node:url"; +import type { Queue } from "../queue"; + +/** + * Import the user's app module and return its configured queue. + * + * Shared by `run` and `executor`: both need the same registry of tasks, and + * only differ in where the jobs come from. + */ +export async function loadApp(appPath: string): Promise { + const module = (await import(pathToFileURL(resolve(appPath)).href)) as Record; + const candidate = module.default ?? module.queue; + if (!candidate || typeof (candidate as Queue).runWorker !== "function") { + throw new Error(`module "${appPath}" must export a Queue (default export or \`queue\`)`); + } + return candidate as Queue; +} diff --git a/sdks/node/src/detached.ts b/sdks/node/src/detached.ts new file mode 100644 index 000000000..520a4a2ad --- /dev/null +++ b/sdks/node/src/detached.ts @@ -0,0 +1,139 @@ +import type { NativeQueue } from "./native"; +import { createLogger } from "./utils"; + +const log = createLogger("executor"); + +/** + * Marks this process as an executor, so a {@link Queue} built here opens no + * storage. Set by `taskito executor` before it imports the app. Internal: + * applications should not set it. + */ +export const DETACHED_ENV = "TASKITO_DETACHED_EXECUTOR"; + +/** Whether this process runs task bodies without any storage of its own. */ +export function isDetached(): boolean { + return process.env[DETACHED_ENV] === "1"; +} + +/** An executor was asked for something only a database could answer. */ +export class DetachedStorageError extends Error { + constructor(operation: string) { + super( + `'${operation}' needs a database, and an attached executor has none. ` + + "Only running tasks is supported here — the scheduler owns storage. " + + "Use an in-process worker (`runWorker`) if this app needs to reach the queue itself.", + ); + this.name = "DetachedStorageError"; + } +} + +/** + * Properties the JavaScript runtime itself probes on any object. + * + * They must answer "absent" rather than throwing: `await`ing a value reads + * `then`, `JSON.stringify` reads `toJSON`, and Node's inspector reads several + * more. A throwing getter would turn a harmless probe into a crash. + */ +const RUNTIME_PROBES = new Set([ + "then", + "toJSON", + "inspect", + "constructor", + "valueOf", + "toString", + "nodeType", +]); + +/** + * The native queue's job-scoped conveniences, degraded. + * + * Reads answer empty, because that is exactly what a queue with no such row + * returns and callers already handle it. Progress and task logs are dropped + * with one warning rather than throwing: a task that only wanted to report + * progress must not fail because it happens to be running detached. + */ +function degraded(warnOnce: (what: string) => void): Record { + return { + updateProgress(): void { + warnOnce("setProgress"); + }, + writeTaskLog(): void { + warnOnce("publish"); + }, + // A cancel reaches an executor as a protocol frame; the executor overrides + // this check with its own native state (see `Executor.start`). + isCancelRequested(): boolean { + return false; + }, + getSetting(): string | null { + return null; + }, + listSettings(): Record { + return {}; + }, + }; +} + +/** + * A stand-in for the native queue in an executor. + * + * An attached executor exists so the app image needs no database credentials: + * the scheduler holds the connection and dispatches over a socket. But an + * executor still imports the user's app module to find its handlers, and that + * module builds a `Queue` — which would otherwise connect the moment it is + * constructed, putting the credentials right back in the app image. + * + * Everything outside the degraded set throws, because an enqueue that quietly + * vanished would be worse than one that failed. + * + * The same split shows up in the job a handler receives. A dispatch frame + * carries what running the task needs, so `createdAt`, `scheduledAt`, + * `priority`, `metadata`, `uniqueKey` and `notes` arrive as zeros and nulls on + * an executor where an in-process worker would show the stored values. A task + * that needs them wants a worker, not an executor. + */ +export function createDetachedNative(): NativeQueue { + const warned = new Set(); + const warnOnce = (what: string): void => { + // Once per process, not per call: a progress-reporting loop would other- + // wise bury the log it is trying to be useful in. + if (!warned.has(what)) { + warned.add(what); + log.warn( + () => + `${what} is unavailable on an attached executor, which has no storage; ignoring. ` + + "Run an in-process worker if you need it.", + ); + } + }; + + const supported = degraded(warnOnce); + const proxy = new Proxy(supported, { + get(target, property): unknown { + if (typeof property !== "string") { + return undefined; + } + if (property in target) { + return target[property]; + } + if (RUNTIME_PROBES.has(property)) { + return undefined; + } + // Returned rather than thrown here: the caller wanted the method, and + // failing at the call site keeps the stack pointing at their code. + return () => { + throw new DetachedStorageError(property); + }; + }, + has(target, property): boolean { + return typeof property === "string" && !RUNTIME_PROBES.has(property) + ? true + : property in target; + }, + }); + + // The stand-in answers the calls a running task makes and throws on the rest, + // so a union type would force every storage call site to handle a case only + // an executor ever sees. + return proxy as unknown as NativeQueue; +} diff --git a/sdks/node/src/executor.ts b/sdks/node/src/executor.ts new file mode 100644 index 000000000..006d1fd78 --- /dev/null +++ b/sdks/node/src/executor.ts @@ -0,0 +1,201 @@ +import type { Emitter } from "./events"; +import type { Middleware } from "./middleware"; +import { + type NativeExecutor, + type NativeQueue, + startExecutor as startNativeExecutor, +} from "./native"; +import type { ResourceRuntime } from "./resources"; +import type { PayloadCodec, Serializer } from "./serializers"; +import { createTaskCallback } from "./task-callback"; +import type { RegisteredTask } from "./types"; +import { createLogger } from "./utils"; + +const log = createLogger("executor"); + +/** How an executor attaches. Durations are milliseconds, per Node convention. */ +export interface ExecutorRunOptions { + /** + * Scheduler address: `host:port`, `:port`, or `unix:/run/taskito.sock`. + * Defaults to `$TASKITO_ATTACH`. + */ + attach?: string; + /** Jobs to run at once. Defaults to `$TASKITO_SLOTS`, then 1. */ + slots?: number; + /** Shared secret, when the scheduler requires one. Defaults to `$TASKITO_ATTACH_TOKEN`. */ + token?: string; + /** Identity announced to the scheduler. Defaults to one generated per process. */ + executorId?: string; + /** How long to wait for the connection (default 10000). */ + connectTimeoutMs?: number; + /** How often to send a liveness heartbeat (default 5000). */ + heartbeatIntervalMs?: number; + /** How long a drain waits for in-flight jobs before disconnecting (default 30000). */ + shutdownDrainMs?: number; + /** Only advertise these tasks. Defaults to every registered task. */ + tasks?: readonly string[]; +} + +/** Inputs assembled by {@link Queue.runExecutor}. @internal */ +export interface ExecutorStartParams { + /** Called once the executor has stopped, so the queue can forget it. */ + onStopped?: () => void; + tasks: ReadonlyMap; + serializer: Serializer; + codecs?: ReadonlyMap; + middlewareFor: (taskName: string) => readonly Middleware[]; + emitter: Emitter; + resources: ResourceRuntime; + run?: ExecutorRunOptions; +} + +/** + * A running attachment to a detached scheduler. + * + * The inverse of a {@link Worker}: instead of polling storage for work, it + * dials a scheduler that already holds the database connection and runs + * whatever it is sent. Task execution is identical — same middleware, codecs, + * resources and cancel signal — because only the transport differs. + */ +export class Executor { + private stopped?: Promise; + + private constructor( + private readonly native: NativeExecutor, + private readonly resources: ResourceRuntime, + private readonly emitter: Emitter, + private readonly onStopped?: () => void, + ) {} + + /** + * Attach and start running jobs. Use {@link Queue.runExecutor} rather than + * calling this directly. + * + * @internal + */ + static async start(queue: NativeQueue, params: ExecutorStartParams): Promise { + const { tasks, serializer, codecs, middlewareFor, emitter, resources, run, onStopped } = params; + + const address = run?.attach ?? process.env.TASKITO_ATTACH; + if (!address) { + throw new Error( + "no scheduler address: pass `attach` or set TASKITO_ATTACH (e.g. scheduler:7749)", + ); + } + const advertised = [...(run?.tasks ?? tasks.keys())]; + + // The executor does not exist yet, and the callback it needs must already + // be able to reach it — a cancel frame lands in native state that a running + // handler polls. Resolved through this holder, assigned once the attach + // succeeds; until then nothing is running, so nothing can be cancelled. + let attached: NativeExecutor | undefined; + + const taskCallback = createTaskCallback({ + tasks, + serializer, + codecs, + middlewareFor, + emitter, + resources, + queue, + isCancelled: (jobId) => attached?.isCancelRequested(jobId) ?? false, + }); + + const native = await startNativeExecutor(taskCallback, { + address, + tasks: advertised, + slots: run?.slots ?? envInt("TASKITO_SLOTS"), + // Env by preference for the token: in argv it shows up in `ps` output and + // shell history. + token: run?.token ?? process.env.TASKITO_ATTACH_TOKEN, + executorId: run?.executorId, + connectTimeoutMs: run?.connectTimeoutMs, + heartbeatIntervalMs: run?.heartbeatIntervalMs, + shutdownDrainMs: run?.shutdownDrainMs, + }); + + attached = native; + try { + // Only lease the resource runtime once the attach actually succeeded, so a + // refused handshake leaks nothing. + resources.acquireWorker(); + emitter.emit("worker.started", { workerId: native.executorId }); + } catch (error) { + // The session is live by now and no caller holds an `Executor` to stop + // it, so a throwing resource factory or `worker.started` listener would + // leak the attach until the process exits. + await native.shutdown().catch((failure) => { + log.debug(() => "releasing the attach after a failed start failed", failure); + }); + throw error; + } + + return new Executor(native, resources, emitter, onStopped); + } + + /** Identity the scheduler announced when it accepted this attach. */ + get schedulerId(): string { + return this.native.schedulerId; + } + + /** Identity this executor attached under. */ + get executorId(): string { + return this.native.executorId; + } + + /** Peer label of the scheduler connection. */ + get peer(): string { + return this.native.peer; + } + + /** Whether the scheduler session is still open. */ + get running(): boolean { + return this.native.isRunning(); + } + + /** + * Resolve once the scheduler ends the session — a shutdown frame, or the + * connection dropping. Does not drain; call {@link Executor.stop} for that. + */ + wait(): Promise { + return this.native.wait(); + } + + /** + * Drain in-flight work, disconnect, and release worker-scoped resources. + * + * Idempotent, and memoized like {@link Worker.stop} so a signal handler + * racing a scheduler shutdown does not tear down twice. + */ + stop(): Promise { + this.stopped ??= this.teardown(); + return this.stopped; + } + + private async teardown(): Promise { + try { + await this.native.shutdown(); + } finally { + try { + await this.resources.teardownWorker(); + } catch (error) { + log.debug(() => "resource release during executor shutdown failed", error); + } + this.emitter.emit("worker.stopped", { workerId: this.executorId }); + this.onStopped?.(); + } + } +} + +/** Read a positive integer from the environment, or `undefined` when unusable. */ +function envInt(name: string): number | undefined { + const raw = process.env[name]; + if (raw === undefined || raw === "") { + return undefined; + } + const value = Number(raw); + if (!Number.isInteger(value) || value < 1) { + throw new RangeError(`${name} must be a positive integer, got "${raw}"`); + } + return value; +} diff --git a/sdks/node/src/index.ts b/sdks/node/src/index.ts index e1c1ea0c6..952705b98 100644 --- a/sdks/node/src/index.ts +++ b/sdks/node/src/index.ts @@ -29,6 +29,7 @@ export { type Role, serveDashboard, } from "./dashboard"; +export { DetachedStorageError } from "./detached"; export { CryptoError, EnqueueSkippedError, @@ -69,6 +70,7 @@ export { type WorkerUnhealthyEvent, type WorkflowEvent, } from "./events"; +export { Executor, type ExecutorRunOptions } from "./executor"; export { checkHealth, checkReadiness, diff --git a/sdks/node/src/native.ts b/sdks/node/src/native.ts index 09d6bf55f..07dd72042 100644 --- a/sdks/node/src/native.ts +++ b/sdks/node/src/native.ts @@ -9,16 +9,18 @@ const require = createRequire(import.meta.url); const bindingPath = fileURLToPath(new URL("../native/index.js", import.meta.url)); const binding = require(bindingPath) as typeof import("../native/index"); -export const { JsQueue, JsWorker, reservedSettingPrefixes } = binding; +export const { JsQueue, JsWorker, startExecutor, reservedSettingPrefixes } = binding; /** Instance types of the native classes, for typing fields and parameters. */ export type NativeQueue = InstanceType; export type NativeWorker = InstanceType; +export type NativeExecutor = Awaited>; export type { CircuitBreakerInput, DetailedJobFilter, EnqueueOptions, + ExecutorOptions as NativeExecutorOptions, JobFilter, JsCircuitBreaker, JsDagEdge, diff --git a/sdks/node/src/queue.ts b/sdks/node/src/queue.ts index 538170d95..043747942 100644 --- a/sdks/node/src/queue.ts +++ b/sdks/node/src/queue.ts @@ -8,6 +8,7 @@ import { type QueueOverride, type TaskOverride, } from "./dashboard/stores"; +import { createDetachedNative, isDetached } from "./detached"; import { EnqueueSkippedError, InterceptionError, @@ -24,6 +25,7 @@ import { TaskitoError, } from "./errors"; import { Emitter, type EventMap, type EventName, type PredicateEvent } from "./events"; +import { Executor, type ExecutorRunOptions } from "./executor"; import { type Interception, type InterceptionAnalysis, @@ -181,6 +183,7 @@ export class Queue { private readonly resources = new ResourceRuntime(); /** Workers started from this queue and not yet stopped — the shutdown set. */ private readonly liveWorkers = new Set(); + private readonly liveExecutors = new Set(); private readonly webhookManager: WebhookManager; /** Built lazily — its constructor throws on addons lacking the `workflows` feature. */ private workflowManager?: WorkflowManager; @@ -188,7 +191,10 @@ export class Queue { private workflowTracker?: WorkflowTracker; constructor(options: QueueOptions = {}) { - this.native = JsQueue.open(toOpenOptions(options)); + // An executor imports this app only to find its handlers; connecting here + // would put the database credentials back in the app image that the attach + // split exists to keep them out of. + this.native = isDetached() ? createDetachedNative() : JsQueue.open(toOpenOptions(options)); const chain = options.codec === undefined ? [] : [options.codec].flat(); const baseSerializer = options.serializer ?? new JsonSerializer(); this.serializer = @@ -218,7 +224,9 @@ export class Queue { /** The shared workflow tracker, or `undefined` on addons without workflows. */ private trackerIfSupported(): WorkflowTracker | undefined { - if (typeof this.native.markWorkflowNodeResult !== "function") { + // Workflow tracking is storage-backed, so a detached executor has none — + // answered here rather than by probing the stand-in, which would throw. + if (isDetached() || typeof this.native.markWorkflowNodeResult !== "function") { return undefined; } this.workflowTracker ??= new WorkflowTracker( @@ -1643,17 +1651,50 @@ export class Queue { } /** - * Stop every worker started from this queue — the programmatic equivalent of - * SIGINT/SIGTERM. Dispatch halts at once and the promise resolves once - * worker-scoped resources are disposed. Handlers already mid-flight are not - * awaited: like {@link Worker.stop}, this stops dispatch rather than draining - * the invocations in progress. + * Attach to a detached scheduler and run its jobs in this process. * - * A no-op when no worker is running, and safe alongside a direct - * {@link Worker.stop} — stopping twice does nothing the second time. + * The inverse of {@link Queue.runWorker}: the scheduler holds the database + * connection and dispatches over a socket, so this process runs task bodies + * without polling storage itself. Hold the returned {@link Executor}. + */ + async runExecutor(options?: ExecutorRunOptions): Promise { + const disables = new MiddlewareDisableStore(this.native); + const executor: Executor = await Executor.start(this.native, { + onStopped: () => this.liveExecutors.delete(executor), + tasks: this.tasks, + serializer: this.serializer, + codecs: this.codecs, + middlewareFor: (taskName) => { + const disabled = disables.getFor(taskName); + return disabled.length === 0 + ? this.middleware + : this.middleware.filter((mw, index) => !disabled.includes(middlewareKey(mw, index))); + }, + emitter: this.emitter, + resources: this.resources, + run: options, + }); + this.liveExecutors.add(executor); + return executor; + } + + /** + * Stop every worker and executor started from this queue — the programmatic + * equivalent of SIGINT/SIGTERM. Dispatch halts at once and the promise + * resolves once worker-scoped resources are disposed. Handlers already + * mid-flight are not awaited: like {@link Worker.stop}, this stops dispatch + * rather than draining the invocations in progress. An executor is the + * exception — {@link Executor.stop} drains before it disconnects. + * + * A no-op when nothing is running, and safe alongside a direct + * {@link Worker.stop} or {@link Executor.stop} — stopping twice does nothing + * the second time. */ async shutdown(): Promise { - await Promise.all([...this.liveWorkers].map((worker) => worker.stop())); + await Promise.all([ + ...[...this.liveWorkers].map((worker) => worker.stop()), + ...[...this.liveExecutors].map((executor) => executor.stop()), + ]); } } diff --git a/sdks/node/src/task-callback.ts b/sdks/node/src/task-callback.ts new file mode 100644 index 000000000..809b6dfc1 --- /dev/null +++ b/sdks/node/src/task-callback.ts @@ -0,0 +1,169 @@ +import { type JobContext, runInContext } from "./context"; +import { SerializationError, TaskNotRegisteredError } from "./errors"; +import type { Emitter } from "./events"; +import type { Middleware, TaskContext } from "./middleware"; +import type { JsTaskInvocation, JsTaskOutcome, NativeQueue } from "./native"; +import { type ResourceRuntime, runWithResolver } from "./resources"; +import { deserializeCall, type PayloadCodec, type Serializer } from "./serializers"; +import { encodeTaskError } from "./task-error"; +import type { RegisteredTask } from "./types"; +import { createLogger } from "./utils"; +import { CACHE_TASK } from "./workflows/cache"; + +const log = createLogger("task"); + +/** How often a running job polls the storage cancel flag. */ +const CANCEL_POLL_INTERVAL_MS = 200; + +/** What running one task needs, independent of how the job arrived. */ +export interface TaskCallbackDeps { + tasks: ReadonlyMap; + serializer: Serializer; + /** Named codec registry for per-task payload decode (see `TaskOptions.codecs`). */ + codecs?: ReadonlyMap; + /** The middleware chain for a task, after dashboard disables are applied. */ + middlewareFor: (taskName: string) => readonly Middleware[]; + emitter: Emitter; + resources: ResourceRuntime; + /** Backs progress, published partials, and the cancel-flag poll. */ + queue: NativeQueue; + /** + * Overrides how a running job learns it was cancelled. + * + * An attached executor reads no storage, so the flag `queue` would poll is + * never set; its cancels arrive as protocol frames instead, and this reads + * the native state those land in. + */ + isCancelled?: (jobId: string) => boolean; +} + +/** + * Build the function the native layer calls for each dispatched job. + * + * Shared by the in-process worker and the attached executor: a job is a job + * however it arrived, and duplicating this would mean two places to fix a + * middleware-ordering or codec bug. + */ +export function createTaskCallback( + deps: TaskCallbackDeps, +): (invocation: JsTaskInvocation) => Promise { + const { tasks, serializer, codecs, middlewareFor, emitter, resources, queue } = deps; + const isCancelled = deps.isCancelled ?? ((jobId: string) => queue.isCancelRequested(jobId)); + + return async (invocation: JsTaskInvocation): Promise => { + // Built-in workflow cache-return: echo the single (cached) arg as the result. + if (invocation.taskName === CACHE_TASK) { + const [value] = deserializeCall(serializer, invocation.payload); + return { result: Buffer.from(serializer.serialize(value)) }; + } + const task = tasks.get(invocation.taskName); + if (!task) { + throw new TaskNotRegisteredError(invocation.taskName); + } + // Reverse the task's named codecs (see `TaskOptions.codecs`) before decode. + let payload: Uint8Array = invocation.payload; + for (const codecName of [...(task.options?.codecs ?? [])].reverse()) { + const codec = codecs?.get(codecName); + if (!codec) { + throw new SerializationError(`no codec registered named "${codecName}"`); + } + payload = codec.decode(payload); + } + const args = deserializeCall(serializer, payload); + const ctx: TaskContext = { jobId: invocation.id, taskName: invocation.taskName, args }; + // Resolve the middleware chain BEFORE allocating the cancel poller and + // task scope — it reads storage and may throw, and nothing would clean + // those up yet. + const chain = middlewareFor(invocation.taskName); + + // Cooperative cancel signal + job context exposed to the handler. + const controller = new AbortController(); + const context: JobContext = { + jobId: invocation.id, + signal: controller.signal, + setProgress: (progress) => queue.updateProgress(invocation.id, progress), + publish: (value) => + queue.writeTaskLog(invocation.id, invocation.taskName, "result", "", JSON.stringify(value)), + }; + const poller = setInterval(() => { + try { + if (isCancelled(invocation.id)) { + controller.abort(); + } + } catch (error) { + // transient storage error — retry on the next tick + log.debug(() => `cancel poll for ${invocation.id} failed`, error); + } + }, CANCEL_POLL_INTERVAL_MS); + poller.unref(); + + // Per-invocation resource scope; `useResource`/`inject` resolve against it. + const scope = resources.createTaskScope(); + const invoke = async (): Promise => { + const inject = task.options?.inject; + if (inject && inject.length > 0) { + const deps: Record = {}; + for (const name of inject) { + deps[name] = await scope.resolver(name); + } + return task.handler(...args, deps); + } + return task.handler(...args); + }; + + const startedAt = performance.now(); + try { + for (const mw of chain) { + await mw.before?.(ctx); + } + const result = await runWithResolver(scope.resolver, () => runInContext(context, invoke)); + for (const mw of chain) { + await mw.after?.(ctx, result); + } + return { result: Buffer.from(serializer.serialize(result)) }; + } catch (error) { + for (const mw of chain) { + try { + await mw.onError?.(ctx, error); + } catch { + // onError hooks must not mask the original task failure. + } + } + // Resolve rather than reject: a rejection carries only a string, and the + // native layer needs the retry verdict alongside the canonical + // structured-error JSON it stores as the job's error. + const encoded = encodeTaskError(error); + // One `job.failed` per failed attempt (the retry/dead verdict follows + // as its own outcome event once the scheduler settles the job). + emitter.emit("job.failed", { + jobId: invocation.id, + taskName: invocation.taskName, + error: encoded, + durationMs: performance.now() - startedAt, + }); + return { error: encoded, retryable: isRetryable(task, error) }; + } finally { + clearInterval(poller); + try { + await scope.teardown(); + } catch (error) { + // dispose errors must not fail an already-settled job + log.debug(() => `task-scope teardown for ${invocation.id} failed`, error); + } + } + }; +} + +/** Whether a task's `retryOn` predicate accepts this failure. */ +export function isRetryable(task: RegisteredTask, error: unknown): boolean { + const predicate = task.options?.retryOn; + if (!predicate) { + return true; + } + try { + return predicate(error); + } catch (predicateError) { + log.error(() => "retryOn predicate threw; retrying the failure", predicateError); + return true; + } +} diff --git a/sdks/node/src/worker.ts b/sdks/node/src/worker.ts index d803aa889..f40ebbc2c 100644 --- a/sdks/node/src/worker.ts +++ b/sdks/node/src/worker.ts @@ -1,4 +1,3 @@ -import { type JobContext, runInContext } from "./context"; import { applyQueueOverrides, applyTaskOverrides, @@ -6,22 +5,19 @@ import { middlewareKey, OverridesStore, } from "./dashboard/stores"; -import { SerializationError, TaskNotRegisteredError } from "./errors"; import { type Emitter, OUTCOME_KIND_EVENTS, type OutcomeEvent } from "./events"; -import type { Middleware, TaskContext } from "./middleware"; +import type { Middleware } from "./middleware"; import type { JsOutcome, - JsTaskInvocation, - JsTaskOutcome, NativeQueue, NativeWorker, WorkerOptions as NativeWorkerOptions, QueueConfigInput, TaskConfigInput, } from "./native"; -import { type ResourceRuntime, runWithResolver } from "./resources"; +import type { ResourceRuntime } from "./resources"; import { deserializeCall, type PayloadCodec, type Serializer } from "./serializers"; -import { encodeTaskError } from "./task-error"; +import { createTaskCallback } from "./task-callback"; import type { AnyHandler, QueueLimits, @@ -31,13 +27,9 @@ import type { } from "./types"; import { createLogger } from "./utils"; import type { WorkflowTracker } from "./workflows"; -import { CACHE_TASK } from "./workflows/cache"; const log = createLogger("worker"); -/** How often a running job polls the storage cancel flag. */ -const CANCEL_POLL_INTERVAL_MS = 200; - /** How often the worker heartbeats (with resource health) to storage. */ const HEARTBEAT_INTERVAL_MS = 5000; @@ -122,114 +114,15 @@ export class Worker { // Advance workflow runs as node-jobs settle, unless disabled or unsupported. const tracker = (run?.advanceWorkflows ?? true) ? (params.workflowTracker ?? null) : null; - const taskCallback = async (invocation: JsTaskInvocation): Promise => { - // Built-in workflow cache-return: echo the single (cached) arg as the result. - if (invocation.taskName === CACHE_TASK) { - const [value] = deserializeCall(serializer, invocation.payload); - return { result: Buffer.from(serializer.serialize(value)) }; - } - const task = tasks.get(invocation.taskName); - if (!task) { - throw new TaskNotRegisteredError(invocation.taskName); - } - // Reverse the task's named codecs (see `TaskOptions.codecs`) before decode. - let payload: Uint8Array = invocation.payload; - for (const codecName of [...(task.options?.codecs ?? [])].reverse()) { - const codec = codecs?.get(codecName); - if (!codec) { - throw new SerializationError(`no codec registered named "${codecName}"`); - } - payload = codec.decode(payload); - } - const args = deserializeCall(serializer, payload); - const ctx: TaskContext = { jobId: invocation.id, taskName: invocation.taskName, args }; - // Resolve the middleware chain BEFORE allocating the cancel poller and - // task scope — it reads storage and may throw, and nothing would clean - // those up yet. - const chain = middlewareFor(invocation.taskName); - - // Cooperative cancel signal + job context exposed to the handler. - const controller = new AbortController(); - const context: JobContext = { - jobId: invocation.id, - signal: controller.signal, - setProgress: (progress) => queue.updateProgress(invocation.id, progress), - publish: (value) => - queue.writeTaskLog( - invocation.id, - invocation.taskName, - "result", - "", - JSON.stringify(value), - ), - }; - const poller = setInterval(() => { - try { - if (queue.isCancelRequested(invocation.id)) { - controller.abort(); - } - } catch (error) { - // transient storage error — retry on the next tick - log.debug(() => `cancel poll for ${invocation.id} failed`, error); - } - }, CANCEL_POLL_INTERVAL_MS); - poller.unref(); - - // Per-invocation resource scope; `useResource`/`inject` resolve against it. - const scope = resources.createTaskScope(); - const invoke = async (): Promise => { - const inject = task.options?.inject; - if (inject && inject.length > 0) { - const deps: Record = {}; - for (const name of inject) { - deps[name] = await scope.resolver(name); - } - return task.handler(...args, deps); - } - return task.handler(...args); - }; - - const startedAt = performance.now(); - try { - for (const mw of chain) { - await mw.before?.(ctx); - } - const result = await runWithResolver(scope.resolver, () => runInContext(context, invoke)); - for (const mw of chain) { - await mw.after?.(ctx, result); - } - return { result: Buffer.from(serializer.serialize(result)) }; - } catch (error) { - for (const mw of chain) { - try { - await mw.onError?.(ctx, error); - } catch { - // onError hooks must not mask the original task failure. - } - } - // Resolve rather than reject: a rejection carries only a string, and the - // native layer needs the retry verdict alongside the canonical - // structured-error JSON it stores as the job's error. - const encoded = encodeTaskError(error); - // One `job.failed` per failed attempt (the retry/dead verdict follows - // as its own outcome event once the scheduler settles the job). - emitter.emit("job.failed", { - jobId: invocation.id, - taskName: invocation.taskName, - error: encoded, - durationMs: performance.now() - startedAt, - }); - return { error: encoded, retryable: isRetryable(task, error) }; - } finally { - clearInterval(poller); - try { - await scope.teardown(); - } catch (error) { - // dispose errors must not fail an already-settled job - log.debug(() => `task-scope teardown for ${invocation.id} failed`, error); - } - } - }; + const taskCallback = createTaskCallback({ + tasks, + serializer, + codecs, + middlewareFor, + emitter, + resources, + queue, + }); const outcomeCallback = (outcome: JsOutcome): void => { const kind = outcome.kind as keyof typeof OUTCOME_KIND_EVENTS; @@ -444,19 +337,6 @@ export class Worker { * predicate means retry, and so does one that throws — a broken classifier must * not silently turn transient failures into dead letters. */ -function isRetryable(task: RegisteredTask, error: unknown): boolean { - const predicate = task.options?.retryOn; - if (!predicate) { - return true; - } - try { - return predicate(error); - } catch (predicateError) { - log.error(() => "retryOn predicate threw; retrying the failure", predicateError); - return true; - } -} - /** Start one poll loop per managed consumer; return their timers to clear on stop. */ function startLogConsumers( queue: NativeQueue, diff --git a/sdks/node/test/worker/executorAttach.test.ts b/sdks/node/test/worker/executorAttach.test.ts new file mode 100644 index 000000000..29fb63005 --- /dev/null +++ b/sdks/node/test/worker/executorAttach.test.ts @@ -0,0 +1,619 @@ +/** + * End-to-end tests for the attached executor. + * + * A scheduler is played by a plain socket speaking the frame protocol, so these + * run without building the Rust server binary. `executorAttachServer.test.ts` + * runs the same assertions against the real `taskito-server` when one is + * available — that pairing is what keeps this fake honest. + */ + +import { mkdtempSync } from "node:fs"; +import { createServer, type Server, type Socket } from "node:net"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { afterEach, expect, it } from "vitest"; +import { DETACHED_ENV } from "../../src/detached"; +import { currentJob, DetachedStorageError, type Executor, Queue } from "../../src/index"; + +/** Frame-format version this build speaks; mirrored from the core. */ +const PROTOCOL_VERSION = 1; + +const SETTLE_MS = 15_000; + +let executor: Executor | undefined; +let scheduler: FakeScheduler | undefined; + +afterEach(async () => { + await executor?.stop(); + executor = undefined; + scheduler?.close(); + scheduler = undefined; +}); + +interface Frame { + header: Record; + payload: Buffer; +} + +/** + * The scheduler end of an attach, driven frame by frame. + * + * A frame is a JSON header line followed by exactly the number of raw payload + * bytes it declares, so decoding has to be length-driven rather than + * line-driven — a payload can contain newlines. + */ +class FakeScheduler { + private server: Server; + private socket?: Socket; + private buffer = Buffer.alloc(0); + private frames: Frame[] = []; + private waiters: (() => void)[] = []; + private connected: Promise; + port = 0; + /** The `hello` the executor opened with, once it has attached. */ + hello?: Record; + /** Set when the handshake should be refused rather than acked. */ + refuse = false; + + private constructor(server: Server, port: number, connected: Promise) { + this.server = server; + this.port = port; + this.connected = connected; + } + + static async listen(options?: { refuse?: boolean }): Promise { + const server = createServer(); + await new Promise((resolve) => server.listen(0, "127.0.0.1", resolve)); + const address = server.address(); + if (address === null || typeof address === "string") { + throw new Error("expected a TCP address"); + } + + let onConnected: () => void = () => {}; + const connected = new Promise((resolve) => { + onConnected = resolve; + }); + const fake = new FakeScheduler(server, address.port, connected); + if (options?.refuse) { + fake.refuse = true; + } + + server.on("connection", (socket) => { + fake.socket = socket; + socket.on("data", (chunk: Buffer) => { + fake.buffer = Buffer.concat([fake.buffer, chunk]); + fake.drain(); + onConnected(); + }); + socket.on("error", () => {}); + }); + return fake; + } + + /** Decode every whole frame currently buffered. */ + private drain(): void { + for (;;) { + const newline = this.buffer.indexOf(0x0a); + if (newline < 0) { + return; + } + const header = JSON.parse(this.buffer.subarray(0, newline).toString()) as Record< + string, + unknown + >; + const declared = declaredPayloadLength(header); + const total = newline + 1 + declared; + if (this.buffer.length < total) { + return; + } + const payload = this.buffer.subarray(newline + 1, total); + this.buffer = this.buffer.subarray(total); + + if (header.type === "hello") { + this.hello = header; + if (!this.refuse) { + this.send({ + type: "hello_ack", + schedulerId: undefined, + scheduler_id: "fake-scheduler", + protocol_version: PROTOCOL_VERSION, + }); + } else { + this.socket?.destroy(); + } + continue; + } + this.frames.push({ header, payload: Buffer.from(payload) }); + for (const wake of this.waiters.splice(0)) { + wake(); + } + } + } + + send(header: Record, payload: Buffer = Buffer.alloc(0)): void { + const clean = Object.fromEntries(Object.entries(header).filter(([, v]) => v !== undefined)); + this.socket?.write(`${JSON.stringify(clean)}\n`); + if (payload.length > 0) { + this.socket?.write(payload); + } + } + + sendJob( + id: string, + taskName: string, + payload: Buffer, + options?: { retryCount?: number; maxRetries?: number; timeoutMs?: number }, + ): void { + this.send( + { + type: "job", + id, + task_name: taskName, + payload_len: payload.length, + retry_count: options?.retryCount ?? 0, + max_retries: options?.maxRetries ?? 3, + queue: "default", + timeout_ms: options?.timeoutMs ?? 30_000, + namespace: null, + }, + payload, + ); + } + + /** Wait for the executor's `hello` to arrive. */ + async attached(): Promise> { + await this.connected; + const deadline = Date.now() + SETTLE_MS; + while (this.hello === undefined && Date.now() < deadline) { + await sleep(10); + } + if (this.hello === undefined) { + throw new Error("the executor never sent a hello"); + } + return this.hello; + } + + /** The next frame that is not a heartbeat. */ + async nextResult(): Promise { + const deadline = Date.now() + SETTLE_MS; + for (;;) { + const found = this.frames.findIndex((frame) => frame.header.type !== "heartbeat"); + const frame = found >= 0 ? this.frames.splice(found, 1)[0] : undefined; + if (frame !== undefined) { + return frame; + } + if (Date.now() > deadline) { + throw new Error("no result frame arrived"); + } + await new Promise((resolve) => { + this.waiters.push(resolve); + setTimeout(resolve, 50); + }); + } + } + + /** Block until the executor reports exactly `free` slots. */ + async heartbeat(free: number): Promise { + const deadline = Date.now() + SETTLE_MS; + for (;;) { + const found = this.frames.findIndex( + (frame) => frame.header.type === "heartbeat" && frame.header.free_slots === free, + ); + if (found >= 0) { + this.frames.splice(0, found + 1); + return; + } + if (Date.now() > deadline) { + throw new Error(`no heartbeat reporting ${free} free slots`); + } + await sleep(20); + } + } + + close(): void { + this.socket?.destroy(); + this.server.close(); + } +} + +/** Bytes of payload a header says follow it — the reader's framing rule. */ +function declaredPayloadLength(header: Record): number { + if (header.type === "job") { + return typeof header.payload_len === "number" ? header.payload_len : 0; + } + if (header.type === "success") { + return typeof header.result_len === "number" ? header.result_len : 0; + } + return 0; +} + +function sleep(ms: number): Promise { + return new Promise((resolve) => setTimeout(resolve, ms)); +} + +function newQueue(): Queue { + return new Queue({ dbPath: join(mkdtempSync(join(tmpdir(), "taskito-exec-")), "q.db") }); +} + +/** Encode a call the way the enqueue path does. */ +function payloadFor(queue: Queue, args: unknown[]): Buffer { + // biome-ignore lint/complexity/useLiteralKeys: reaching the internal serializer + const serializer = (queue as unknown as { serializer: { serialize(v: unknown): Uint8Array } })[ + "serializer" + ]; + return Buffer.from(serializer.serialize([args, {}])); +} + +it("announces itself and the tasks it can run", async () => { + scheduler = await FakeScheduler.listen(); + const queue = newQueue(); + queue.task("echo", (value: string) => `echo:${value}`); + queue.task("boom", () => { + throw new Error("deliberate failure"); + }); + + executor = await queue.runExecutor({ + attach: `127.0.0.1:${scheduler.port}`, + slots: 2, + executorId: "exec-test", + }); + const hello = await scheduler.attached(); + + expect(hello.executor_id).toBe("exec-test"); + expect(hello.sdk).toBe("node"); + expect(hello.slots).toBe(2); + expect(hello.protocol_version).toBe(PROTOCOL_VERSION); + // Only advertised tasks are ever dispatched, so a missing name here is a job + // that silently never runs. + expect(hello.tasks).toEqual(expect.arrayContaining(["echo", "boom"])); + // A token that was never configured must not appear on the wire. + expect(hello).not.toHaveProperty("token"); + expect(executor.schedulerId).toBe("fake-scheduler"); +}); + +it("runs a dispatched job and returns its result", async () => { + scheduler = await FakeScheduler.listen(); + const queue = newQueue(); + queue.task("echo", (value: string) => `echo:${value}`); + + executor = await queue.runExecutor({ attach: `127.0.0.1:${scheduler.port}` }); + await scheduler.attached(); + scheduler.sendJob("job-1", "echo", payloadFor(queue, ["hello"])); + + const frame = await scheduler.nextResult(); + expect(frame.header.type).toBe("success"); + expect(frame.header.job_id).toBe("job-1"); + expect(frame.payload.length).toBeGreaterThan(0); +}); + +it("reports a task failure with its retry verdict", async () => { + scheduler = await FakeScheduler.listen(); + const queue = newQueue(); + queue.task("boom", () => { + throw new Error("deliberate failure"); + }); + + executor = await queue.runExecutor({ attach: `127.0.0.1:${scheduler.port}` }); + await scheduler.attached(); + scheduler.sendJob("job-1", "boom", payloadFor(queue, []), { retryCount: 2 }); + + const frame = await scheduler.nextResult(); + expect(frame.header.type).toBe("failure"); + expect(frame.header.job_id).toBe("job-1"); + expect(frame.header.should_retry).toBe(true); + expect(frame.header.timed_out).toBe(false); + // The frame's retry count is echoed back so the scheduler's backoff is right. + expect(frame.header.retry_count).toBe(2); + expect(String(frame.header.error)).toContain("deliberate failure"); +}); + +it("honours a task's retryOn predicate over the wire", async () => { + // Only the executor sees the exception, so its verdict is the one that + // counts; a wire defaulting this to true would retry poison jobs forever. + scheduler = await FakeScheduler.listen(); + const queue = newQueue(); + queue.task( + "fatal", + () => { + throw new Error("do not retry me"); + }, + { retryOn: () => false }, + ); + + executor = await queue.runExecutor({ attach: `127.0.0.1:${scheduler.port}` }); + await scheduler.attached(); + scheduler.sendJob("job-1", "fatal", payloadFor(queue, [])); + + const frame = await scheduler.nextResult(); + expect(frame.header.type).toBe("failure"); + expect(frame.header.should_retry).toBe(false); +}); + +it("runs jobs on separate slots concurrently", async () => { + scheduler = await FakeScheduler.listen(); + const queue = newQueue(); + let running = 0; + let peak = 0; + let release = (): void => {}; + const released = new Promise((resolve) => { + release = resolve; + }); + queue.task("slow", async () => { + running += 1; + peak = Math.max(peak, running); + await released; + running -= 1; + return null; + }); + + executor = await queue.runExecutor({ attach: `127.0.0.1:${scheduler.port}`, slots: 2 }); + await scheduler.attached(); + scheduler.sendJob("job-1", "slow", payloadFor(queue, [])); + scheduler.sendJob("job-2", "slow", payloadFor(queue, [])); + + const deadline = Date.now() + SETTLE_MS; + while (peak < 2 && Date.now() < deadline) { + await sleep(10); + } + release(); + + const first = await scheduler.nextResult(); + const second = await scheduler.nextResult(); + expect([first.header.job_id, second.header.job_id].sort()).toEqual(["job-1", "job-2"]); + expect(peak).toBe(2); +}); + +it("announces zero capacity before disconnecting on stop", async () => { + // This is what makes the drain clean rather than a race: the scheduler is + // told to stop dispatching in-protocol, before the connection goes away. + scheduler = await FakeScheduler.listen(); + const queue = newQueue(); + queue.task("echo", (value: string) => value); + + const running = await queue.runExecutor({ + attach: `127.0.0.1:${scheduler.port}`, + slots: 2, + heartbeatIntervalMs: 50, + }); + await scheduler.attached(); + + const stopping = running.stop(); + await scheduler.heartbeat(0); + await stopping; + executor = undefined; +}); + +it("finishes in-flight work before disconnecting", async () => { + scheduler = await FakeScheduler.listen(); + const queue = newQueue(); + let release = (): void => {}; + const released = new Promise((resolve) => { + release = resolve; + }); + let started = false; + queue.task("slow", async () => { + started = true; + await released; + return "done"; + }); + + const running = await queue.runExecutor({ + attach: `127.0.0.1:${scheduler.port}`, + heartbeatIntervalMs: 50, + }); + await scheduler.attached(); + scheduler.sendJob("job-1", "slow", payloadFor(queue, [])); + + const deadline = Date.now() + SETTLE_MS; + while (!started && Date.now() < deadline) { + await sleep(10); + } + + // Stop while the job is still running: it must still report, or the job + // waits for a reap it never needed. + const stopping = running.stop(); + release(); + const frame = await scheduler.nextResult(); + expect(frame.header.type).toBe("success"); + expect(frame.header.job_id).toBe("job-1"); + + await stopping; + executor = undefined; +}); + +it("ends the session when the scheduler sends shutdown", async () => { + scheduler = await FakeScheduler.listen(); + const queue = newQueue(); + queue.task("echo", (value: string) => value); + + const running = await queue.runExecutor({ attach: `127.0.0.1:${scheduler.port}` }); + await scheduler.attached(); + expect(running.running).toBe(true); + + scheduler.send({ type: "shutdown" }); + await running.wait(); + + expect(running.running).toBe(false); + await running.stop(); + executor = undefined; +}); + +it("reports a refused attach rather than a network error", async () => { + // A wrong token is the likeliest deployment mistake; it must not surface as + // "connection reset" and send the operator looking at the network. + scheduler = await FakeScheduler.listen({ refuse: true }); + const queue = newQueue(); + queue.task("echo", (value: string) => value); + + await expect( + queue.runExecutor({ attach: `127.0.0.1:${scheduler?.port}`, token: "wrong-token" }), + ).rejects.toThrow(/refused|token/i); +}); + +it("fails fast when no scheduler is listening", async () => { + const queue = newQueue(); + queue.task("echo", (value: string) => value); + + // Port 1 on loopback is reserved and nothing listens there. + await expect(queue.runExecutor({ attach: "127.0.0.1:1", connectTimeoutMs: 500 })).rejects.toThrow( + /could not reach the scheduler/, + ); +}); + +it("requires an attach address", async () => { + const queue = newQueue(); + queue.task("echo", (value: string) => value); + + const previous = process.env.TASKITO_ATTACH; + process.env.TASKITO_ATTACH = undefined; + delete process.env.TASKITO_ATTACH; + try { + await expect(queue.runExecutor()).rejects.toThrow(/TASKITO_ATTACH/); + } finally { + if (previous !== undefined) { + process.env.TASKITO_ATTACH = previous; + } + } +}); + +it("takes the attach address and slot count from the environment", async () => { + scheduler = await FakeScheduler.listen(); + const queue = newQueue(); + queue.task("echo", (value: string) => value); + + const previousAttach = process.env.TASKITO_ATTACH; + const previousSlots = process.env.TASKITO_SLOTS; + process.env.TASKITO_ATTACH = `127.0.0.1:${scheduler.port}`; + process.env.TASKITO_SLOTS = "3"; + try { + executor = await queue.runExecutor(); + const hello = await scheduler.attached(); + expect(hello.slots).toBe(3); + } finally { + restoreEnv("TASKITO_ATTACH", previousAttach); + restoreEnv("TASKITO_SLOTS", previousSlots); + } +}); + +it("rejects a non-numeric slot count from the environment", async () => { + const queue = newQueue(); + queue.task("echo", (value: string) => value); + + const previous = process.env.TASKITO_SLOTS; + process.env.TASKITO_SLOTS = "many"; + try { + await expect(queue.runExecutor({ attach: "127.0.0.1:1" })).rejects.toThrow(RangeError); + } finally { + restoreEnv("TASKITO_SLOTS", previous); + } +}); + +function restoreEnv(name: string, previous: string | undefined): void { + if (previous === undefined) { + delete process.env[name]; + } else { + process.env[name] = previous; + } +} + +it("opens no storage", async () => { + // The point of the attach split: app code without database credentials. + // Pointed at a Postgres DSN nothing is listening on — a Queue that connected + // could not even be constructed. + scheduler = await FakeScheduler.listen(); + process.env[DETACHED_ENV] = "1"; + try { + const queue = new Queue({ backend: "postgres", dsn: "postgres://x:y@127.0.0.1:1/absent" }); + queue.task("echo", (value: string) => `echo:${value}`); + + executor = await queue.runExecutor({ attach: `127.0.0.1:${scheduler.port}` }); + await scheduler.attached(); + + // And it still runs jobs. + scheduler.sendJob("job-1", "echo", payloadFor(queue, ["detached"])); + const frame = await scheduler.nextResult(); + expect(frame.header.type).toBe("success"); + } finally { + delete process.env[DETACHED_ENV]; + } +}); + +it("degrades progress and publish rather than failing the job", async () => { + // Losing the progress bar is a degradation; failing the job over it would be + // a regression for anyone moving a worker to an executor. + scheduler = await FakeScheduler.listen(); + process.env[DETACHED_ENV] = "1"; + try { + const queue = new Queue({ backend: "postgres", dsn: "postgres://x:y@127.0.0.1:1/absent" }); + queue.task("reports", async () => { + const job = currentJob(); + job?.setProgress(50); + job?.publish({ stage: "halfway" }); + job?.setProgress(100); + return "reported"; + }); + + executor = await queue.runExecutor({ attach: `127.0.0.1:${scheduler.port}` }); + await scheduler.attached(); + scheduler.sendJob("job-1", "reports", payloadFor(queue, [])); + + const frame = await scheduler.nextResult(); + expect(frame.header.type).toBe("success"); + expect(frame.header.job_id).toBe("job-1"); + } finally { + delete process.env[DETACHED_ENV]; + } +}); + +it("refuses a storage operation instead of silently dropping it", async () => { + // An enqueue that quietly vanished would be worse than one that threw. + process.env[DETACHED_ENV] = "1"; + try { + const queue = new Queue({ backend: "postgres", dsn: "postgres://x:y@127.0.0.1:1/absent" }); + queue.task("echo", (value: string) => value); + expect(() => queue.enqueue("echo", ["x"])).toThrow(DetachedStorageError); + } finally { + delete process.env[DETACHED_ENV]; + } +}); + +it("aborts a running handler when a cancel frame arrives", async () => { + // A detached executor has no storage flag to poll, so the cancel has to reach + // the handler's AbortSignal from the frame the scheduler sent. + scheduler = await FakeScheduler.listen(); + process.env[DETACHED_ENV] = "1"; + try { + const queue = new Queue({ backend: "postgres", dsn: "postgres://x:y@127.0.0.1:1/absent" }); + let started = false; + queue.task("slow", async () => { + started = true; + const job = currentJob(); + for (let i = 0; i < 400; i += 1) { + if (job?.signal.aborted) { + throw new Error("cancelled"); + } + await sleep(25); + } + return "never"; + }); + + executor = await queue.runExecutor({ attach: `127.0.0.1:${scheduler.port}` }); + await scheduler.attached(); + scheduler.sendJob("job-1", "slow", payloadFor(queue, [])); + + const deadline = Date.now() + SETTLE_MS; + while (!started && Date.now() < deadline) { + await sleep(10); + } + scheduler.send({ type: "cancel", job_id: "job-1" }); + + const frame = await scheduler.nextResult(); + // `cancelled`, not `failure`: the handler observed the signal and threw, + // and the native side reclassified that throw as the cancellation it was. + // Both halves of the frame-driven cancel path in one assertion. + expect(frame.header.type).toBe("cancelled"); + expect(frame.header.job_id).toBe("job-1"); + } finally { + delete process.env[DETACHED_ENV]; + } +}); diff --git a/sdks/node/test/worker/executorAttachServer.test.ts b/sdks/node/test/worker/executorAttachServer.test.ts new file mode 100644 index 000000000..0cf654316 --- /dev/null +++ b/sdks/node/test/worker/executorAttachServer.test.ts @@ -0,0 +1,182 @@ +/** + * The same attach assertions, against the real `taskito-server` binary. + * + * Gated on `TASKITO_SERVER_BIN` so the default suite needs no Rust build. Its + * job is to keep `executorAttach.test.ts`'s hand-rolled scheduler honest: that + * file proves the executor speaks the protocol it was told to, this one proves + * the protocol it was told to is the one the server actually speaks. + * + * cargo build -p taskito-server + * TASKITO_SERVER_BIN=../../target/debug/taskito-server npx vitest run + */ + +import { type ChildProcess, spawn } from "node:child_process"; +import { mkdtempSync } from "node:fs"; +import { createConnection, createServer } from "node:net"; +import { tmpdir } from "node:os"; +import { join, resolve } from "node:path"; +import { afterEach, describe, expect, it } from "vitest"; +import { type Executor, Queue } from "../../src/index"; + +const SERVER_BIN = process.env.TASKITO_SERVER_BIN; +const SETTLE_MS = 60_000; + +let executor: Executor | undefined; +let server: ChildProcess | undefined; + +afterEach(async () => { + await executor?.stop(); + executor = undefined; + server?.kill("SIGTERM"); + server = undefined; +}); + +function sleep(ms: number): Promise { + return new Promise((done) => setTimeout(done, ms)); +} + +async function freePort(): Promise { + const probe = createServer(); + await new Promise((done) => probe.listen(0, "127.0.0.1", done)); + const address = probe.address(); + if (address === null || typeof address === "string") { + throw new Error("expected a TCP address"); + } + const { port } = address; + await new Promise((done) => probe.close(() => done())); + return port; +} + +/** Start a real scheduler over a temp SQLite database. */ +async function startScheduler(options?: { token?: string }): Promise<{ + port: number; + dbPath: string; +}> { + const dbPath = join(mkdtempSync(join(tmpdir(), "taskito-server-")), "server.db"); + const port = await freePort(); + + const env: NodeJS.ProcessEnv = { + ...process.env, + TASKITO_BACKEND: "sqlite", + TASKITO_DSN: dbPath, + TASKITO_LISTEN: `127.0.0.1:${port}`, + }; + // Unset, not "off": the dashboard is disabled by having no bind address. + delete env.TASKITO_DASHBOARD; + delete env.TASKITO_ATTACH_TOKEN; + if (options?.token) { + env.TASKITO_ATTACH_TOKEN = options.token; + } + + // Discarded, not piped: nothing reads these, and a scheduler that logs per + // job fills the pipe buffer and then blocks on its next write. + server = spawn(resolve(SERVER_BIN as string), { env, stdio: "ignore" }); + await waitForPort(port); + return { port, dbPath }; +} + +async function waitForPort(port: number): Promise { + const deadline = Date.now() + SETTLE_MS; + while (Date.now() < deadline) { + const open = await new Promise((done) => { + const socket = createConnection({ port, host: "127.0.0.1" }); + socket.once("connect", () => { + socket.destroy(); + done(true); + }); + socket.once("error", () => done(false)); + }); + if (open) { + return; + } + await sleep(50); + } + throw new Error(`the server never bound port ${port}`); +} + +async function waitFor(predicate: () => Promise, what: string): Promise { + const deadline = Date.now() + SETTLE_MS; + while (Date.now() < deadline) { + if (await predicate()) { + return; + } + await sleep(50); + } + throw new Error(what); +} + +describe.skipIf(!SERVER_BIN)("against a real taskito-server", () => { + it("runs a job the real scheduler dispatched", async () => { + const { port, dbPath } = await startScheduler(); + const queue = new Queue({ dbPath }); + queue.task("echo", (value: string) => `echo:${value}`); + + executor = await queue.runExecutor({ attach: `127.0.0.1:${port}` }); + // The scheduler starts on the first attach, so enqueue after attaching. + await sleep(1000); + const jobId = queue.enqueue("echo", ["hello"]); + + await waitFor( + async () => queue.getJob(String(jobId))?.status === "complete", + "the job never completed on the attached executor", + ); + }, 120_000); + + it("retries a failure through the real scheduler", async () => { + const { port, dbPath } = await startScheduler(); + const queue = new Queue({ dbPath }); + queue.task("boom", () => { + throw new Error("deliberate failure"); + }); + + executor = await queue.runExecutor({ attach: `127.0.0.1:${port}` }); + await sleep(1000); + const jobId = queue.enqueue("boom", []); + + // The error reaching storage proves the executor's failure crossed the wire + // and the scheduler applied it. + await waitFor( + async () => Boolean(queue.getJob(String(jobId))?.error?.includes("deliberate failure")), + "the failure never reached storage", + ); + }, 120_000); + + it("refuses an attach with the wrong token", async () => { + const { port, dbPath } = await startScheduler({ token: "correct-token-0123456789" }); + const queue = new Queue({ dbPath }); + queue.task("echo", (value: string) => value); + + await expect( + queue.runExecutor({ attach: `127.0.0.1:${port}`, token: "wrong-token-0123456789" }), + ).rejects.toThrow(/refused|token/i); + }, 120_000); + + it("drains in-flight work when stopped", async () => { + const { port, dbPath } = await startScheduler(); + const queue = new Queue({ dbPath }); + let release = (): void => {}; + const released = new Promise((done) => { + release = done; + }); + let started = false; + queue.task("slow", async () => { + started = true; + await released; + return "done"; + }); + + const running = await queue.runExecutor({ attach: `127.0.0.1:${port}` }); + await sleep(1000); + const jobId = queue.enqueue("slow", []); + await waitFor(async () => started, "the job never started"); + + const stopping = running.stop(); + release(); + await stopping; + + await waitFor( + async () => queue.getJob(String(jobId))?.status === "complete", + "in-flight work was not drained before disconnecting", + ); + }, 120_000); +}); diff --git a/sdks/python/taskito/_taskito.pyi b/sdks/python/taskito/_taskito.pyi index 43e861fa4..df7c20ec6 100644 --- a/sdks/python/taskito/_taskito.pyi +++ b/sdks/python/taskito/_taskito.pyi @@ -571,6 +571,58 @@ class PyResultSender: wall_time_ns: int, ) -> bool: ... +class Executor: + """A running attachment to a detached scheduler. + + Constructing one performs the handshake, so a bad token or an unreachable + scheduler raises here rather than after the pool has been built. Jobs run on + the prefork pool — one child per slot. + """ + + def __init__( + self, + address: str, + app_path: str, + tasks: list[str], + slots: int, + token: str | None = None, + executor_id: str | None = None, + ) -> None: ... + @property + def scheduler_id(self) -> str: + """Identity the scheduler announced when it accepted this attach.""" + ... + + @property + def executor_id(self) -> str: + """Identity this executor attached under.""" + ... + + @property + def peer(self) -> str: + """Peer label of the scheduler connection.""" + ... + + def is_running(self) -> bool: + """Whether the scheduler session is still open.""" + ... + + def wait(self, timeout_ms: int) -> bool: + """Block up to `timeout_ms`, returning whether the session has ended. + + Loop on this rather than blocking indefinitely: each return hands the + GIL back, which is the only moment a pending signal handler can run. + """ + ... + + def stop(self) -> None: + """Stop accepting work and finish what is in flight. Returns at once.""" + ... + + def shutdown(self) -> None: + """Drain in-flight work, disconnect, and join. Idempotent.""" + ... + def _init_rust_logging() -> None: """Activate the Rust → Python `logging` bridge (idempotent).""" ... diff --git a/sdks/python/taskito/app.py b/sdks/python/taskito/app.py index 910fe4d24..bfc440264 100644 --- a/sdks/python/taskito/app.py +++ b/sdks/python/taskito/app.py @@ -24,12 +24,13 @@ from collections import Counter from collections.abc import Callable, Sequence from concurrent.futures import ThreadPoolExecutor -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, cast from taskito._taskito import PyQueue from taskito.async_support.mixins import AsyncQueueMixin from taskito.batching import BatchAccumulator, BatchConfig from taskito.codecs import CodecSerializer, PayloadCodec +from taskito.detached import DetachedNative, is_detached from taskito.enums import StorageBackend, coerce_enum from taskito.events import EventBus, EventType from taskito.exceptions import QueueFullError, SerializationError @@ -250,32 +251,44 @@ def __init__( if isinstance(backend, StorageBackend): backend = backend.value - if backend == "sqlite": + # An executor imports this app only to find its handlers; opening + # storage here would put the database credentials back in the app image + # that the attach split exists to keep them out of. + detached = is_detached() + + if backend == "sqlite" and not detached: # Ensure parent directory exists for SQLite db_dir = os.path.dirname(db_path) if db_dir: os.makedirs(db_dir, exist_ok=True) - self._inner = PyQueue( - db_path=db_path, - workers=workers, - default_retry=default_retry, - default_timeout=default_timeout, - default_priority=default_priority, - result_ttl=result_ttl, - backend=backend, - db_url=db_url, - schema=schema, - pool_size=pool_size, - scheduler_poll_interval_ms=scheduler_poll_interval_ms, - scheduler_reap_interval=scheduler_reap_interval, - scheduler_cleanup_interval=scheduler_cleanup_interval, - scheduler_batch_size=scheduler_batch_size, - namespace=namespace, - push_dispatch=push_dispatch, - dlq_auto_retry_delay=dlq_auto_retry_delay, - dlq_auto_retry_max=dlq_auto_retry_max, - retention=retention._as_map() if retention is not None else None, + # `cast`, not a union: the stand-in answers the job-scoped calls a task + # makes and raises on the rest, so widening the type would force all ~140 + # storage call sites to handle a case only an executor ever sees. + self._inner: PyQueue = ( + cast("PyQueue", DetachedNative()) + if detached + else PyQueue( + db_path=db_path, + workers=workers, + default_retry=default_retry, + default_timeout=default_timeout, + default_priority=default_priority, + result_ttl=result_ttl, + backend=backend, + db_url=db_url, + schema=schema, + pool_size=pool_size, + scheduler_poll_interval_ms=scheduler_poll_interval_ms, + scheduler_reap_interval=scheduler_reap_interval, + scheduler_cleanup_interval=scheduler_cleanup_interval, + scheduler_batch_size=scheduler_batch_size, + namespace=namespace, + push_dispatch=push_dispatch, + dlq_auto_retry_delay=dlq_auto_retry_delay, + dlq_auto_retry_max=dlq_auto_retry_max, + retention=retention._as_map() if retention is not None else None, + ) ) self._backend = backend self._namespace = namespace diff --git a/sdks/python/taskito/cli.py b/sdks/python/taskito/cli.py index fae8ba92a..bbc719167 100644 --- a/sdks/python/taskito/cli.py +++ b/sdks/python/taskito/cli.py @@ -7,14 +7,24 @@ import os import signal as sig import sys +import threading import time +from taskito._taskito import Executor as _Executor from taskito.app import Queue from taskito.autoscale import AutoscaleConfig, serve_autoscaler from taskito.dashboard import serve_dashboard +from taskito.detached import DETACHED_ENV from taskito.log_config import configure as configure_logging from taskito.scaler import serve_scaler +# How long the executor's main loop blocks between signal checks. Short enough +# that Ctrl-C feels immediate, long enough not to spin. +_EXECUTOR_POLL_MS = 200 + +# One job at a time unless asked otherwise: each slot is a whole interpreter. +_EXECUTOR_DEFAULT_SLOTS = 1 + def _build_parser() -> argparse.ArgumentParser: """Build the CLI argument parser with all subcommands.""" @@ -49,6 +59,40 @@ def _build_parser() -> argparse.ArgumentParser: help="Worker pool: 'thread' (default) or 'prefork' for true CPU parallelism", ) + # executor subcommand + exec_parser = subparsers.add_parser( + "executor", + help="Run tasks for a detached scheduler instead of polling storage", + description=( + "Attach to a taskito-server scheduler and run its jobs in this process. " + "Needs no database credentials: everything a task needs arrives on the wire." + ), + ) + exec_parser.add_argument( + "--app", + required=True, + help="Python path to the Queue instance (e.g., 'myapp.tasks:queue')", + ) + exec_parser.add_argument( + "--attach", + default=None, + help=( + "Scheduler address: host:port, :port, or unix:/run/taskito.sock " + "(default: $TASKITO_ATTACH)" + ), + ) + exec_parser.add_argument( + "--slots", + type=int, + default=None, + help=("Jobs to run concurrently, one prefork child each (default: $TASKITO_SLOTS, or 1)"), + ) + exec_parser.add_argument( + "--executor-id", + default=None, + help="Identity announced to the scheduler (default: generated per process)", + ) + # dashboard subcommand dash_parser = subparsers.add_parser("dashboard", help="Start the web dashboard") dash_parser.add_argument( @@ -219,6 +263,8 @@ def main() -> None: if args.command == "worker": run_worker(args) + elif args.command == "executor": + run_executor(args) elif args.command == "dashboard": run_dashboard(args) elif args.command == "info": @@ -293,6 +339,96 @@ def run_worker(args: argparse.Namespace) -> None: queue.run_worker(queues=queues, pool=args.pool, app=args.app) +def run_executor(args: argparse.Namespace) -> None: + """Attach to a detached scheduler and run its jobs in this process. + + The inverse of ``run_worker``: instead of polling storage for work, the + executor dials a scheduler that already holds the database connection and + runs whatever it is sent. Only the transport differs — jobs execute on the + same prefork pool, one child per slot. + """ + address = args.attach or os.environ.get("TASKITO_ATTACH") + if not address: + print( + "Error: --attach is required (or set TASKITO_ATTACH), e.g. --attach scheduler:7749", + file=sys.stderr, + ) + sys.exit(1) + + slots = _executor_slots(args.slots) + + # Set before the app is imported: building a Queue is what opens storage, + # and an executor must not. Inherited by the prefork children, which import + # the same app module in their own interpreters. + os.environ[DETACHED_ENV] = "1" + + queue = _load_queue(args.app) + tasks = sorted(queue._task_registry) + + try: + executor = _Executor( + address, + args.app, + tasks, + slots, + # Env only, never a flag: a token in argv is visible in `ps` output + # and lands in shell history. + os.environ.get("TASKITO_ATTACH_TOKEN"), + args.executor_id, + ) + except (RuntimeError, ValueError) as exc: + print(f"Error: {exc}", file=sys.stderr) + sys.exit(1) + + print( + f"taskito executor {executor.executor_id} attached to " + f"{executor.scheduler_id} at {executor.peer} " + f"({slots} slot(s), {len(tasks)} task(s)) — Ctrl-C to stop" + ) + + stopping = threading.Event() + + def request_stop(signum: int, frame: object) -> None: + # Returns immediately; the drain runs on the executor's own threads. + stopping.set() + executor.stop() + + sig.signal(sig.SIGTERM, request_stop) + sig.signal(sig.SIGINT, request_stop) + + # Polled rather than blocked on: a Python signal handler only runs when the + # main thread holds the GIL, which it cannot do inside a blocking call. + while not stopping.is_set(): + if executor.wait(_EXECUTOR_POLL_MS): + break + + executor.shutdown() + print("taskito executor detached") + + +def _executor_slots(flag: int | None) -> int: + """Resolve the slot count from the flag, then the env, then the default.""" + if flag is not None: + value = flag + else: + raw = os.environ.get("TASKITO_SLOTS") + if raw is None: + return _EXECUTOR_DEFAULT_SLOTS + try: + value = int(raw) + except ValueError: + print( + f"Error: TASKITO_SLOTS must be an integer, got '{raw}'", + file=sys.stderr, + ) + sys.exit(1) + + if value < 1: + print(f"Error: slots must be at least 1, got {value}", file=sys.stderr) + sys.exit(1) + return value + + def run_dashboard(args: argparse.Namespace) -> None: """Start the web dashboard.""" queue = _load_queue(args.app) diff --git a/sdks/python/taskito/detached.py b/sdks/python/taskito/detached.py new file mode 100644 index 000000000..feb107fa7 --- /dev/null +++ b/sdks/python/taskito/detached.py @@ -0,0 +1,124 @@ +"""Running a ``Queue`` with no storage behind it. + +An attached executor exists so the app image needs no database credentials: +the scheduler holds the connection and dispatches over a socket, and everything +a task needs to run arrives on the wire. But an executor still imports the +user's app module to find its handlers, and that module builds a ``Queue`` — +which would otherwise open storage the moment it is constructed, putting the +credentials right back in the app image. + +So in an executor the native queue is replaced by :class:`DetachedNative`. Task +execution never touches storage, so nothing on the hot path notices. The +operations that genuinely need a database fail loudly rather than silently +doing nothing, because an enqueue that quietly vanished would be worse than one +that raised. Reads are the exception: ``None`` is what a queue with no such row +returns anyway, and ``Queue.__init__`` performs some. + +The same split shows up in the job a handler receives. A dispatch frame carries +what running the task needs, so ``created_at``, ``scheduled_at``, ``priority``, +``metadata``, ``unique_key`` and ``notes`` arrive as zeros and ``None`` on an +executor where an in-process worker would show the stored values. A task that +needs them wants a worker, not an executor. + +Set by ``taskito executor`` before it imports the app, and inherited by the +prefork children it spawns. Internal: applications should not set it. +""" + +from __future__ import annotations + +import logging +import os +from typing import NoReturn + +logger = logging.getLogger("taskito.executor") + +__all__ = ["DETACHED_ENV", "DetachedNative", "DetachedStorageError", "is_detached"] + +#: Marks this process as an executor, so a ``Queue`` built here opens no storage. +DETACHED_ENV = "TASKITO_DETACHED_EXECUTOR" + + +class DetachedStorageError(RuntimeError, AttributeError): + """An executor was asked for something only a database could answer. + + Deliberately both: ``RuntimeError`` so it reads as the operational fault it + is, and ``AttributeError`` so a capability probe (``hasattr(queue._inner, + "submit_workflow")``) answers "no" instead of exploding. A detached queue + genuinely does not have those capabilities. + """ + + +def is_detached() -> bool: + """Whether this process runs task bodies without any storage of its own.""" + return os.environ.get(DETACHED_ENV) == "1" + + +class DetachedNative: + """Stands in for the native queue in an executor. + + Degrades the three job-scoped conveniences that only observability depends + on, and refuses everything else. A task calling ``update_progress`` must not + fail merely because it happens to be running detached, but a task calling + ``enqueue`` must not appear to succeed. + """ + + __slots__ = ("_warned",) + + def __init__(self) -> None: + # One warning per process, not per call: a progress-reporting loop would + # otherwise bury the log it is trying to be useful in. + self._warned: set[str] = set() + + def _warn_once(self, what: str) -> None: + if what not in self._warned: + self._warned.add(what) + logger.warning( + "%s is unavailable on an attached executor, which has no storage; " + "ignoring. Run an in-process worker if you need it.", + what, + ) + + def update_progress(self, job_id: str, progress: int) -> None: + """Ignored: progress lives in storage, and there is none here.""" + self._warn_once("update_progress") + + def write_task_log( + self, + job_id: str, + task_name: str, + level: str, + message: str, + extra: str | None = None, + ) -> None: + """Ignored: task logs and published partials live in storage.""" + self._warn_once("current_job.log/publish") + + def is_cancel_requested(self, job_id: str) -> bool: + """Always false: a cancel reaches an executor as a protocol frame. + + ``check_cancelled`` consults the local signal the prefork child installs + before it ever gets here, so answering false loses nothing. + """ + return False + + def get_setting(self, key: str) -> None: + """No settings without storage. + + Reads degrade, writes do not: ``None`` is exactly what a queue with no + such setting returns, so callers already handle it. ``Queue.__init__`` + reads settings — webhook subscriptions, dashboard overrides — so this + has to answer rather than raise, or an executor could not build a Queue + at all. + """ + return None + + def __getattr__(self, name: str) -> NoReturn: + raise DetachedStorageError( + f"'{name}' needs a database, and an attached executor has none. " + "Only running tasks is supported here — the scheduler owns storage. " + "Use an in-process worker (`taskito worker`) if this app needs to " + "reach the queue itself." + ) + + def __repr__(self) -> str: + return "" diff --git a/sdks/python/tests/worker/executor_apps/__init__.py b/sdks/python/tests/worker/executor_apps/__init__.py new file mode 100644 index 000000000..01993a4e5 --- /dev/null +++ b/sdks/python/tests/worker/executor_apps/__init__.py @@ -0,0 +1 @@ +"""App modules imported by the executor attach tests.""" diff --git a/sdks/python/tests/worker/executor_apps/attach_app.py b/sdks/python/tests/worker/executor_apps/attach_app.py new file mode 100644 index 000000000..fc5d29bd4 --- /dev/null +++ b/sdks/python/tests/worker/executor_apps/attach_app.py @@ -0,0 +1,85 @@ +"""Module-level Queue + tasks for the executor attach tests. + +The Queue must be importable both in the ``taskito executor`` process and +inside each prefork child interpreter, so it lives at module scope and takes +its DB path from the environment — the same shape ``prefork_apps`` uses. + +Storage is only touched because importing a ``Queue`` opens one; the executor +itself never reads it, since everything a task needs arrives on the wire. +""" + +from __future__ import annotations + +import os +import time +from pathlib import Path + +from taskito import Queue +from taskito.context import current_job + +# The backend is configurable so one test can point this at a database that +# does not exist, proving an executor never connects to it. +queue = Queue( + backend=os.environ.get("TASKITO_EXECUTOR_TEST_BACKEND", "sqlite"), + db_path=os.environ.get("TASKITO_EXECUTOR_TEST_DB", "/tmp/taskito-executor.db"), + db_url=os.environ.get("TASKITO_EXECUTOR_TEST_DB") + if os.environ.get("TASKITO_EXECUTOR_TEST_BACKEND") + else None, +) + + +def _markers() -> Path | None: + """Directory the test uses to rendezvous with a running task. + + Heartbeats are seconds apart, so they cannot tell a test that a sub-second + job has *started*. A file can, without adding a protocol frame that exists + only for tests. + """ + configured = os.environ.get("TASKITO_EXECUTOR_MARKERS") + return Path(configured) if configured else None + + +@queue.task(max_retries=3) +def echo(value: str) -> str: + """Return its argument, proving the payload survived the hop.""" + return f"echo:{value}" + + +@queue.task(max_retries=3) +def boom() -> None: + """Always fail, so the retry verdict can be asserted.""" + raise RuntimeError("deliberate failure") + + +@queue.task(max_retries=0, timeout=60) +def slow(max_iters: int = 600) -> int: + """Announce that it started, then loop until released or cancelled. + + ``check_cancelled()`` is polled each tick so a cancel lands promptly, and + the release file lets a test end the job deterministically instead of + waiting out a sleep. + """ + markers = _markers() + if markers is not None: + markers.joinpath(f"{current_job.id}.started").write_text("1") + release = None if markers is None else markers / "release" + + for completed in range(max_iters): + current_job.check_cancelled() + if release is not None and release.exists(): + return completed + time.sleep(0.05) + return max_iters + + +@queue.task(max_retries=0) +def reports() -> str: + """Use the job-scoped conveniences that need storage in a worker. + + On an executor there is none, so these degrade; the job must still finish. + """ + current_job.update_progress(50) + current_job.log("halfway") + current_job.publish({"stage": "halfway"}) + current_job.update_progress(100) + return "reported" diff --git a/sdks/python/tests/worker/test_executor_attach.py b/sdks/python/tests/worker/test_executor_attach.py new file mode 100644 index 000000000..d170c5c47 --- /dev/null +++ b/sdks/python/tests/worker/test_executor_attach.py @@ -0,0 +1,544 @@ +"""End-to-end tests for ``taskito executor``. + +A scheduler is played by a plain socket speaking the frame protocol, so these +run in every CI job without building the Rust server binary. ``taskito-server`` +is the real peer, and `test_executor_attach_server.py` runs the same assertions +against it when one is available — that pairing is what keeps this fake honest. + +The executor runs as a real subprocess rather than in-process: prefork children +and ``SIGTERM`` handling are most of what is under test, and neither is +observable from inside the test interpreter. +""" + +from __future__ import annotations + +import contextlib +import os +import signal +import socket +import subprocess +import sys +import time +from collections.abc import Iterator +from pathlib import Path +from typing import Any + +import pytest + +from taskito.worker_protocol import WORKER_PROTOCOL_VERSION, read_frame, write_frame + +pytestmark = pytest.mark.skipif( + sys.platform == "win32", + reason="the executor runs tasks on prefork children, which Windows does not support", +) + +APP_DIR = Path(__file__).parent / "executor_apps" +APP_PATH = "attach_app:queue" + +# Task names are module-qualified in the registry, and it is those names the +# scheduler routes on. +ECHO = "attach_app.echo" +BOOM = "attach_app.boom" +SLOW = "attach_app.slow" +REPORTS = "attach_app.reports" + +# Generous: a cold subprocess import of the app plus a prefork child spawn. +ATTACH_TIMEOUT = 60.0 +FRAME_TIMEOUT = 60.0 + + +class FakeScheduler: + """The scheduler end of an attach, driven frame by frame.""" + + def __init__(self) -> None: + self._listener = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + self._listener.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + self._listener.bind(("127.0.0.1", 0)) + self._listener.listen(1) + self.port: int = self._listener.getsockname()[1] + self._conn: socket.socket | None = None + self._rfile: Any = None + self._wfile: Any = None + + def accept(self, timeout: float = ATTACH_TIMEOUT) -> dict[str, Any]: + """Accept the attach and complete the handshake, returning the hello.""" + self._listener.settimeout(timeout) + self._conn, _ = self._listener.accept() + self._conn.settimeout(FRAME_TIMEOUT) + self._rfile = self._conn.makefile("rb") + self._wfile = self._conn.makefile("wb") + + hello, _ = read_frame(self._rfile) + assert hello["type"] == "hello", f"expected hello, got {hello}" + self.send( + { + "type": "hello_ack", + "scheduler_id": "fake-scheduler", + "protocol_version": WORKER_PROTOCOL_VERSION, + } + ) + return hello + + def refuse(self, timeout: float = ATTACH_TIMEOUT) -> dict[str, Any]: + """Accept, read the hello, then close without acking — a rejected peer.""" + self._listener.settimeout(timeout) + conn, _ = self._listener.accept() + conn.settimeout(FRAME_TIMEOUT) + with conn.makefile("rb") as rfile: + hello, _ = read_frame(rfile) + conn.close() + return hello + + def send(self, header: dict[str, Any], payload: bytes = b"") -> None: + write_frame(self._wfile, header, payload) + + def send_job( + self, + job_id: str, + task_name: str, + payload: bytes, + *, + retry_count: int = 0, + max_retries: int = 3, + timeout_ms: int = 30_000, + ) -> None: + self.send( + { + "type": "job", + "id": job_id, + "task_name": task_name, + "payload_len": len(payload), + "retry_count": retry_count, + "max_retries": max_retries, + "queue": "default", + "timeout_ms": timeout_ms, + "namespace": None, + }, + payload, + ) + + def next_result(self, timeout: float = FRAME_TIMEOUT) -> tuple[dict[str, Any], bytes]: + """The next frame that is not a heartbeat.""" + deadline = time.monotonic() + timeout + while True: + assert time.monotonic() < deadline, "no result frame arrived" + header, payload = read_frame(self._rfile) + if header.get("type") != "heartbeat": + return header, payload + + def next_heartbeat(self, free_slots: int, timeout: float = FRAME_TIMEOUT) -> None: + """Block until the executor reports exactly ``free_slots`` free.""" + deadline = time.monotonic() + timeout + while True: + assert time.monotonic() < deadline, f"no heartbeat reporting {free_slots} slots" + header, _ = read_frame(self._rfile) + if header.get("type") == "heartbeat" and header.get("free_slots") == free_slots: + return + + def close(self) -> None: + # A test that already tore the connection down leaves these broken; + # teardown must not turn that into a second, confusing failure. + for handle in (self._rfile, self._wfile, self._conn, self._listener): + if handle is not None: + with contextlib.suppress(OSError): + handle.close() + + +@pytest.fixture +def scheduler() -> Iterator[FakeScheduler]: + fake = FakeScheduler() + try: + yield fake + finally: + fake.close() + + +def spawn_executor( + port: int, + db_path: Path, + *, + slots: int = 1, + token: str | None = None, + executor_id: str | None = None, + markers: Path | None = None, +) -> subprocess.Popen[str]: + """Run ``taskito executor`` against ``port`` as a real subprocess.""" + env = dict(os.environ) + env["TASKITO_EXECUTOR_TEST_DB"] = str(db_path) + if markers is not None: + markers.mkdir(parents=True, exist_ok=True) + env["TASKITO_EXECUTOR_MARKERS"] = str(markers) + # Prefork children default to `python` on PATH; point them at this + # interpreter so they import the same taskito build the test does. + env["TASKITO_PYTHON"] = sys.executable + env.pop("TASKITO_ATTACH_TOKEN", None) + if token is not None: + env["TASKITO_ATTACH_TOKEN"] = token + + command = [ + sys.executable, + "-m", + "taskito.cli", + "executor", + "--app", + APP_PATH, + "--attach", + f"127.0.0.1:{port}", + "--slots", + str(slots), + ] + if executor_id is not None: + command += ["--executor-id", executor_id] + + return subprocess.Popen( + command, + # The CLI puts the working directory on `sys.path`, and so does each + # prefork child, so this is what makes `attach_app` importable in both. + cwd=str(APP_DIR), + env=env, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + ) + + +def read_stderr(process: subprocess.Popen[str]) -> str: + """Drain a process's stderr, which every spawn here pipes.""" + assert process.stderr is not None, "the process was spawned without a stderr pipe" + text: str = process.stderr.read() + return text + + +def terminate(process: subprocess.Popen[str]) -> None: + """Best-effort teardown for a process a test did not stop itself.""" + if process.poll() is None: + process.terminate() + try: + process.wait(timeout=30) + except subprocess.TimeoutExpired: + process.kill() + process.wait(timeout=10) + + +def wait_started(markers: Path, job_id: str, timeout: float = FRAME_TIMEOUT) -> None: + """Block until the task for ``job_id`` reports that it is running. + + A heartbeat cannot serve here: they are seconds apart, so a sub-second job + starts and finishes between two of them. + """ + deadline = time.monotonic() + timeout + marker = markers / f"{job_id}.started" + while not marker.exists(): + assert time.monotonic() < deadline, f"{job_id} never started" + time.sleep(0.02) + + +def release_tasks(markers: Path) -> None: + """Let every parked task return.""" + (markers / "release").write_text("1") + + +def payload_for(task_name: str, *args: Any, **kwargs: Any) -> bytes: + """Encode a call the way the enqueue path does.""" + from attach_app import queue # type: ignore[import-not-found] + + payload: bytes = queue._get_serializer(task_name).dumps((args, kwargs)) + return payload + + +@pytest.fixture(autouse=True) +def _app_importable() -> Iterator[None]: + """Put the app dir on `sys.path` so the test can use its serializer.""" + sys.path.insert(0, str(APP_DIR)) + try: + yield + finally: + if str(APP_DIR) in sys.path: + sys.path.remove(str(APP_DIR)) + + +def test_executor_announces_itself_and_its_tasks(scheduler: FakeScheduler, tmp_path: Path) -> None: + """The handshake carries what the scheduler routes on.""" + process = spawn_executor(scheduler.port, tmp_path / "t.db", slots=2, executor_id="exec-test") + try: + hello = scheduler.accept() + + assert hello["executor_id"] == "exec-test" + assert hello["sdk"] == "python" + assert hello["slots"] == 2 + assert hello["protocol_version"] == WORKER_PROTOCOL_VERSION + # Only advertised tasks are ever dispatched, so a missing name here is + # a job that silently never runs. + assert set(hello["tasks"]) >= {ECHO, BOOM, SLOW} + # A token that was never configured must not appear on the wire. + assert "token" not in hello + finally: + terminate(process) + + +def test_a_job_runs_on_the_executor_and_returns_its_result( + scheduler: FakeScheduler, tmp_path: Path +) -> None: + process = spawn_executor(scheduler.port, tmp_path / "t.db") + try: + scheduler.accept() + scheduler.send_job("job-1", ECHO, payload_for(ECHO, "hello")) + + header, payload = scheduler.next_result() + assert header["type"] == "success", header + assert header["job_id"] == "job-1" + assert header["task_name"] == ECHO + + from attach_app import queue + + assert queue._get_serializer(ECHO).loads(payload) == "echo:hello" + finally: + terminate(process) + + +def test_a_failing_task_reports_a_retryable_failure( + scheduler: FakeScheduler, tmp_path: Path +) -> None: + """The retry verdict is the executor's to make — only it sees the exception.""" + process = spawn_executor(scheduler.port, tmp_path / "t.db") + try: + scheduler.accept() + scheduler.send_job("job-1", BOOM, payload_for(BOOM), retry_count=1) + + header, _ = scheduler.next_result() + assert header["type"] == "failure", header + assert header["job_id"] == "job-1" + assert header["should_retry"] is True + assert header["timed_out"] is False + assert header["retry_count"] == 1, "the frame's retry count is echoed back" + assert "deliberate failure" in header["error"] + finally: + terminate(process) + + +def test_a_retry_is_dispatched_to_the_same_executor( + scheduler: FakeScheduler, tmp_path: Path +) -> None: + """A second attempt reuses the live attachment rather than needing a reattach.""" + process = spawn_executor(scheduler.port, tmp_path / "t.db") + try: + scheduler.accept() + + scheduler.send_job("job-1", BOOM, payload_for(BOOM), retry_count=0) + first, _ = scheduler.next_result() + assert first["type"] == "failure" + + scheduler.send_job("job-1", ECHO, payload_for(ECHO, "retried"), retry_count=1) + second, _ = scheduler.next_result() + assert second["type"] == "success" + assert second["job_id"] == "job-1" + finally: + terminate(process) + + +def test_a_cancel_stops_a_running_task(scheduler: FakeScheduler, tmp_path: Path) -> None: + markers = tmp_path / "markers" + process = spawn_executor(scheduler.port, tmp_path / "t.db", markers=markers) + try: + scheduler.accept() + scheduler.send_job("job-1", SLOW, payload_for(SLOW, 600)) + wait_started(markers, "job-1") + + # The task polls `check_cancelled()`, so the cancel lands within a tick + # rather than waiting out the whole loop. + scheduler.send({"type": "cancel", "job_id": "job-1"}) + + header, _ = scheduler.next_result() + assert header["type"] == "cancelled", header + assert header["job_id"] == "job-1" + finally: + terminate(process) + + +def test_sigterm_drains_in_flight_work_before_exiting( + scheduler: FakeScheduler, tmp_path: Path +) -> None: + """The container-shutdown path: a held job must still report its result.""" + markers = tmp_path / "markers" + process = spawn_executor(scheduler.port, tmp_path / "t.db", markers=markers) + try: + scheduler.accept() + scheduler.send_job("job-1", SLOW, payload_for(SLOW, 600)) + wait_started(markers, "job-1") + + process.send_signal(signal.SIGTERM) + + # The drain announces zero capacity in-protocol before disconnecting, so + # the scheduler stops dispatching rather than racing the close. + scheduler.next_heartbeat(free_slots=0) + + # The job is still running at this point; finishing it must still be + # reported, or it would wait for a reap it never needed. + release_tasks(markers) + header, _ = scheduler.next_result() + assert header["type"] == "success", header + assert header["job_id"] == "job-1" + + assert process.wait(timeout=60) == 0, "a drained executor exits cleanly" + finally: + terminate(process) + + +def test_a_shutdown_frame_ends_the_session(scheduler: FakeScheduler, tmp_path: Path) -> None: + """The scheduler's own teardown stops the executor without a signal.""" + process = spawn_executor(scheduler.port, tmp_path / "t.db") + try: + scheduler.accept() + scheduler.send({"type": "shutdown"}) + assert process.wait(timeout=60) == 0 + finally: + terminate(process) + + +def test_a_refused_attach_exits_nonzero_with_a_token_hint( + scheduler: FakeScheduler, tmp_path: Path +) -> None: + """A refusal must name the likely cause, not surface as a network error.""" + process = spawn_executor(scheduler.port, tmp_path / "t.db", token="attach-token-0123456789") + try: + hello = scheduler.refuse() + # The token is presented for the scheduler to check, and read from the + # environment rather than argv so it stays out of `ps`. + assert hello["token"] == "attach-token-0123456789" + + assert process.wait(timeout=60) != 0 + assert "token" in read_stderr(process).lower() + finally: + terminate(process) + + +def test_an_unreachable_scheduler_exits_nonzero(tmp_path: Path) -> None: + listener = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + listener.bind(("127.0.0.1", 0)) + closed_port = listener.getsockname()[1] + listener.close() + + process = spawn_executor(closed_port, tmp_path / "t.db") + try: + assert process.wait(timeout=60) != 0 + assert "could not reach the scheduler" in read_stderr(process) + finally: + terminate(process) + + +def test_missing_attach_address_is_reported(tmp_path: Path) -> None: + env = dict(os.environ) + env["TASKITO_EXECUTOR_TEST_DB"] = str(tmp_path / "t.db") + env.pop("TASKITO_ATTACH", None) + + result = subprocess.run( + [sys.executable, "-m", "taskito.cli", "executor", "--app", APP_PATH], + cwd=str(APP_DIR), + env=env, + capture_output=True, + text=True, + timeout=60, + ) + assert result.returncode != 0 + assert "TASKITO_ATTACH" in result.stderr + + +def test_the_attach_address_can_come_from_the_environment( + scheduler: FakeScheduler, tmp_path: Path +) -> None: + """Deployments configure by env, not flags — the contract shared with the other SDKs.""" + env = dict(os.environ) + env["TASKITO_EXECUTOR_TEST_DB"] = str(tmp_path / "t.db") + env["TASKITO_PYTHON"] = sys.executable + env["TASKITO_ATTACH"] = f"127.0.0.1:{scheduler.port}" + env["TASKITO_SLOTS"] = "3" + env.pop("TASKITO_ATTACH_TOKEN", None) + + process = subprocess.Popen( + [sys.executable, "-m", "taskito.cli", "executor", "--app", APP_PATH], + cwd=str(APP_DIR), + env=env, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + ) + try: + hello = scheduler.accept() + assert hello["slots"] == 3, "TASKITO_SLOTS must be honoured" + finally: + terminate(process) + + +def test_slots_run_jobs_concurrently(scheduler: FakeScheduler, tmp_path: Path) -> None: + """Two slots means two prefork children, so two jobs run at once.""" + markers = tmp_path / "markers" + process = spawn_executor(scheduler.port, tmp_path / "t.db", slots=2, markers=markers) + try: + scheduler.accept() + scheduler.send_job("job-1", SLOW, payload_for(SLOW, 600)) + scheduler.send_job("job-2", SLOW, payload_for(SLOW, 600)) + + # Both parked at once. A pool that serialized them would never let the + # second start while the first is still holding its child. + wait_started(markers, "job-1") + wait_started(markers, "job-2") + + release_tasks(markers) + finished = {scheduler.next_result()[0]["job_id"] for _ in range(2)} + assert finished == {"job-1", "job-2"} + finally: + terminate(process) + + +def test_the_executor_opens_no_storage(scheduler: FakeScheduler, tmp_path: Path) -> None: + """The point of the attach split: app code without database credentials. + + Pointed at a Postgres DSN nothing is listening on. An executor that opened + storage could not even start; one that does not never notices. + """ + env = dict(os.environ) + env["TASKITO_PYTHON"] = sys.executable + env["TASKITO_ATTACH"] = f"127.0.0.1:{scheduler.port}" + env.pop("TASKITO_ATTACH_TOKEN", None) + # Port 1 on loopback is reserved and nothing listens there. + env["TASKITO_EXECUTOR_TEST_BACKEND"] = "postgres" + env["TASKITO_EXECUTOR_TEST_DB"] = "postgres://taskito:nope@127.0.0.1:1/absent" + + process = subprocess.Popen( + [sys.executable, "-m", "taskito.cli", "executor", "--app", APP_PATH], + cwd=str(APP_DIR), + env=env, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + ) + try: + hello = scheduler.accept() + assert set(hello["tasks"]) >= {ECHO, BOOM, SLOW} + + # And it still runs jobs, on a prefork child that also opened nothing. + scheduler.send_job("job-1", ECHO, payload_for(ECHO, "detached")) + header, _ = scheduler.next_result() + assert header["type"] == "success", header + finally: + terminate(process) + + +def test_progress_and_logs_degrade_rather_than_failing_the_job( + scheduler: FakeScheduler, tmp_path: Path +) -> None: + """A task calling `update_progress` must not fail for want of storage. + + Losing the progress bar is a degradation; failing the job over it would be + a regression for anyone moving a worker to an executor. + """ + process = spawn_executor(scheduler.port, tmp_path / "t.db") + try: + scheduler.accept() + scheduler.send_job("job-1", REPORTS, payload_for(REPORTS)) + + header, _ = scheduler.next_result() + assert header["type"] == "success", header + assert header["job_id"] == "job-1" + finally: + terminate(process) diff --git a/sdks/python/tests/worker/test_executor_attach_server.py b/sdks/python/tests/worker/test_executor_attach_server.py new file mode 100644 index 000000000..0144301c7 --- /dev/null +++ b/sdks/python/tests/worker/test_executor_attach_server.py @@ -0,0 +1,231 @@ +"""The same attach assertions, against the real ``taskito-server`` binary. + +Gated on ``TASKITO_SERVER_BIN`` so the default suite needs no Rust build. Its +job is to keep `test_executor_attach.py`'s hand-rolled scheduler honest: that +file proves the executor speaks the protocol it was told to, this one proves the +protocol it was told to is the one the server actually speaks. + +Build the binary with:: + + cargo build -p taskito-server + TASKITO_SERVER_BIN=target/debug/taskito-server uv run pytest tests/worker +""" + +from __future__ import annotations + +import os +import subprocess +import sys +import time +from collections.abc import Iterator +from pathlib import Path + +import pytest + +# tests/worker is not a package, so pytest's rootdir insertion is what makes +# this a plain module import rather than a relative one. +from test_executor_attach import ( + APP_DIR, + APP_PATH, + BOOM, + ECHO, + SLOW, + read_stderr, + spawn_executor, + terminate, + wait_started, +) + +from taskito import Queue + +SERVER_BIN = os.environ.get("TASKITO_SERVER_BIN") + +pytestmark = [ + pytest.mark.skipif( + not SERVER_BIN, + reason="set TASKITO_SERVER_BIN to a built taskito-server to run these", + ), + pytest.mark.skipif( + sys.platform == "win32", + reason="the executor runs tasks on prefork children, which Windows does not support", + ), +] + +SETTLE = 60.0 + + +def free_port() -> int: + import socket + + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as probe: + probe.bind(("127.0.0.1", 0)) + return int(probe.getsockname()[1]) + + +@pytest.fixture +def scheduler(tmp_path: Path) -> Iterator[tuple[int, Path]]: + """Run a real scheduler over a temp SQLite database.""" + db_path = tmp_path / "server.db" + port = free_port() + + env = dict(os.environ) + env["TASKITO_BACKEND"] = "sqlite" + env["TASKITO_DSN"] = str(db_path) + env["TASKITO_LISTEN"] = f"127.0.0.1:{port}" + # Unset, not "off": the dashboard is disabled by having no bind address. + env.pop("TASKITO_DASHBOARD", None) + env.pop("TASKITO_ATTACH_TOKEN", None) + + assert SERVER_BIN is not None + process = subprocess.Popen( + [SERVER_BIN], + env=env, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + ) + try: + wait_for_port(port, process) + yield port, db_path + finally: + terminate(process) + + +def wait_for_port(port: int, process: subprocess.Popen[str], timeout: float = SETTLE) -> None: + """Block until the attach listener accepts connections.""" + import socket + + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + if process.poll() is not None: + raise AssertionError(f"the server exited early: {read_stderr(process)}") + try: + with socket.create_connection(("127.0.0.1", port), timeout=1): + return + except OSError: + time.sleep(0.05) + raise AssertionError(f"the server never bound port {port}") + + +def enqueue(db_path: Path, task_name: str, args: tuple = ()) -> str: + """Enqueue through a Queue over the same database the server reads.""" + queue = Queue(db_path=str(db_path)) + return str(queue.enqueue(task_name, args).id) + + +def wait_for_attach(db_path: Path) -> None: + """Block until an executor is attached and dispatching. + + The scheduler starts on the first attach, and nothing records that in + storage, so the only non-racy proof is a job the executor has to run. A + fixed sleep would be a guess about app import and prefork spawn times on a + loaded runner. + """ + probe = enqueue(db_path, ECHO, ("ready",)) + wait_for_status(db_path, probe, "complete") + + +def wait_for_status(db_path: Path, job_id: str, status: str, timeout: float = SETTLE) -> None: + queue = Queue(db_path=str(db_path)) + deadline = time.monotonic() + timeout + seen = None + while time.monotonic() < deadline: + job = queue.get_job(job_id) + seen = None if job is None else job.status + if seen == status: + return + time.sleep(0.05) + raise AssertionError(f"job {job_id} was {seen}, expected {status}") + + +def test_a_real_scheduler_dispatches_to_an_attached_executor( + scheduler: tuple[int, Path], tmp_path: Path +) -> None: + port, db_path = scheduler + process = spawn_executor(port, db_path) + try: + wait_for_attach(db_path) + job_id = enqueue(db_path, ECHO, ("hello",)) + wait_for_status(db_path, job_id, "complete") + finally: + terminate(process) + + +def test_a_failure_is_retried_by_the_real_scheduler( + scheduler: tuple[int, Path], tmp_path: Path +) -> None: + port, db_path = scheduler + process = spawn_executor(port, db_path) + try: + wait_for_attach(db_path) + job_id = enqueue(db_path, BOOM) + + # `boom` always raises. The error reaching storage is what proves the + # executor's failure crossed the wire and the scheduler applied it. + queue = Queue(db_path=str(db_path)) + deadline = time.monotonic() + SETTLE + while time.monotonic() < deadline: + job = queue.get_job(job_id) + if job is not None and job.error: + assert "deliberate failure" in job.error + return + time.sleep(0.05) + raise AssertionError("the failure never reached storage") + finally: + terminate(process) + + +def test_sigterm_drains_against_a_real_scheduler( + scheduler: tuple[int, Path], tmp_path: Path +) -> None: + import signal + + port, db_path = scheduler + markers = tmp_path / "markers" + process = spawn_executor(port, db_path, markers=markers) + try: + wait_for_attach(db_path) + job_id = enqueue(db_path, SLOW, (600,)) + wait_started(markers, job_id) + + process.send_signal(signal.SIGTERM) + (markers / "release").write_text("1") + + wait_for_status(db_path, job_id, "complete") + assert process.wait(timeout=SETTLE) == 0 + finally: + terminate(process) + + +def test_a_bad_token_is_refused_by_the_real_listener(tmp_path: Path) -> None: + """The security gate, against the listener that actually enforces it.""" + db_path = tmp_path / "server.db" + port = free_port() + + env = dict(os.environ) + env["TASKITO_BACKEND"] = "sqlite" + env["TASKITO_DSN"] = str(db_path) + env["TASKITO_LISTEN"] = f"127.0.0.1:{port}" + env.pop("TASKITO_DASHBOARD", None) + env["TASKITO_ATTACH_TOKEN"] = "correct-token-0123456789" + + assert SERVER_BIN is not None + server = subprocess.Popen( + [SERVER_BIN], env=env, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True + ) + try: + wait_for_port(port, server) + executor = spawn_executor(port, db_path, token="wrong-token-0123456789") + try: + assert executor.wait(timeout=SETTLE) != 0 + assert "token" in read_stderr(executor).lower() + finally: + terminate(executor) + finally: + terminate(server) + + +def test_the_app_dir_is_the_one_under_test() -> None: + """Guards the import above: these tests share the other file's fixtures.""" + assert (APP_DIR / "attach_app.py").exists() + assert APP_PATH == "attach_app:queue"