From fb0b776d2a4db28c3968166f42b181ed4b3193b4 Mon Sep 17 00:00:00 2001 From: Pratyush Sharma <56130065+pratyush618@users.noreply.github.com> Date: Sat, 18 Jul 2026 20:56:06 +0530 Subject: [PATCH] feat: add native worker with task registry and dispatcher Rust consumers previously had no execution path: every dispatcher lived in a binding. Worker::spawn wires scheduler, dispatcher, result drain, and heartbeat/reap into one call. --- crates/taskito-core/examples/hello.rs | 83 +++++ crates/taskito-core/src/lib.rs | 4 +- crates/taskito-core/src/scheduler/mod.rs | 14 + crates/taskito-core/src/worker/dispatcher.rs | 121 +++++++ .../src/{worker.rs => worker/mod.rs} | 14 +- crates/taskito-core/src/worker/registry.rs | 111 ++++++ crates/taskito-core/src/worker/runner.rs | 331 ++++++++++++++++++ crates/taskito-core/src/worker/tests.rs | 169 +++++++++ 8 files changed, 845 insertions(+), 2 deletions(-) create mode 100644 crates/taskito-core/examples/hello.rs create mode 100644 crates/taskito-core/src/worker/dispatcher.rs rename crates/taskito-core/src/{worker.rs => worker/mod.rs} (71%) create mode 100644 crates/taskito-core/src/worker/registry.rs create mode 100644 crates/taskito-core/src/worker/runner.rs create mode 100644 crates/taskito-core/src/worker/tests.rs diff --git a/crates/taskito-core/examples/hello.rs b/crates/taskito-core/examples/hello.rs new file mode 100644 index 00000000..ada396c0 --- /dev/null +++ b/crates/taskito-core/examples/hello.rs @@ -0,0 +1,83 @@ +//! Minimal end-to-end run: enqueue two jobs, execute them on a native worker. +//! +//! ```sh +//! cargo run --example hello -p taskito-core +//! ``` + +use std::time::{Duration, Instant}; + +use taskito_core::{ + now_millis, Job, JobStatus, NewJob, SqliteStorage, Storage, StorageBackend, TaskError, Worker, +}; + +fn new_job(task_name: &str, payload: &[u8]) -> NewJob { + NewJob { + queue: "default".to_string(), + task_name: task_name.to_string(), + payload: payload.to_vec(), + priority: 0, + scheduled_at: now_millis(), + max_retries: 3, + timeout_ms: 30_000, + unique_key: None, + metadata: None, + notes: None, + depends_on: vec![], + expires_at: None, + result_ttl_ms: None, + namespace: None, + } +} + +fn wait_until_completed(storage: &StorageBackend, job_id: &str) -> Job { + let deadline = Instant::now() + Duration::from_secs(10); + loop { + let job = storage + .get_job(job_id) + .expect("get_job") + .expect("job exists"); + if job.status == JobStatus::Complete { + return job; + } + assert!(Instant::now() < deadline, "job did not complete in time"); + std::thread::sleep(Duration::from_millis(20)); + } +} + +fn main() { + let storage = StorageBackend::Sqlite(SqliteStorage::new("hello_taskito.db").expect("storage")); + + let handle = Worker::new(storage.clone()) + .num_workers(2) + .register("greet", |job: &Job| { + let name = String::from_utf8(job.payload.clone()) + .map_err(|invalid| TaskError::fatal(format!("payload not utf-8: {invalid}")))?; + println!("hello, {name}!"); + Ok(Some(format!("greeted {name}").into_bytes())) + }) + .register_async("greet_async", |job: Job| async move { + let name = String::from_utf8(job.payload).unwrap_or_else(|_| "?".to_string()); + println!("hello (async), {name}!"); + Ok(None) + }) + .spawn() + .expect("worker spawn"); + + let sync_job = storage + .enqueue(new_job("greet", b"world")) + .expect("enqueue"); + let async_job = storage + .enqueue(new_job("greet_async", b"tokio")) + .expect("enqueue"); + + let done = wait_until_completed(&storage, &sync_job.id); + println!( + "sync result: {}", + String::from_utf8_lossy(done.result.as_deref().unwrap_or_default()) + ); + wait_until_completed(&storage, &async_job.id); + + handle.shutdown().expect("clean shutdown"); + println!("both jobs completed."); + let _ = std::fs::remove_file("hello_taskito.db"); +} diff --git a/crates/taskito-core/src/lib.rs b/crates/taskito-core/src/lib.rs index e7b0dfa1..dee2cab5 100644 --- a/crates/taskito-core/src/lib.rs +++ b/crates/taskito-core/src/lib.rs @@ -32,4 +32,6 @@ pub use storage::sqlite::SqliteStorage; pub use storage::Storage; pub use storage::StorageBackend; pub use storage::{DeadJob, QueueStats, SubscriptionBacklogStats}; -pub use worker::WorkerDispatcher; +pub use worker::{ + NativeDispatcher, TaskError, TaskRegistry, TaskResult, Worker, WorkerDispatcher, WorkerHandle, +}; diff --git a/crates/taskito-core/src/scheduler/mod.rs b/crates/taskito-core/src/scheduler/mod.rs index fa61bc11..be71a933 100644 --- a/crates/taskito-core/src/scheduler/mod.rs +++ b/crates/taskito-core/src/scheduler/mod.rs @@ -199,6 +199,20 @@ pub struct TaskConfig { pub max_in_flight_per_task: Option, } +impl Default for TaskConfig { + /// Default retry policy, no rate limit, breaker, retry budget, or caps. + fn default() -> Self { + Self { + retry_policy: RetryPolicy::default(), + rate_limit: None, + circuit_breaker: None, + retry_budget: None, + max_concurrent: None, + max_in_flight_per_task: None, + } + } +} + /// In-flight bookkeeping behind the dispatch caps: which jobs are out, and how many /// of each task. Per-task counts are maintained alongside rather than derived so the /// per-task gate stays O(1) — the poller consults it on every dispatch, and the map diff --git a/crates/taskito-core/src/worker/dispatcher.rs b/crates/taskito-core/src/worker/dispatcher.rs new file mode 100644 index 00000000..06213c33 --- /dev/null +++ b/crates/taskito-core/src/worker/dispatcher.rs @@ -0,0 +1,121 @@ +//! Reference [`WorkerDispatcher`]: executes registered Rust handlers. +//! +//! Mirrors the language bindings' pools — bounded concurrency via a semaphore, +//! blocking handlers on `spawn_blocking`, async handlers on the runtime — with +//! a [`TaskRegistry`] in place of a foreign-object registry. Task timeouts are +//! detected server-side by the scheduler's stale-job reap; cancellation relies +//! on the storage cancel flag, so `notify_cancel` is a no-op here. + +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::Arc; +use std::time::Instant; + +use async_trait::async_trait; +use crossbeam_channel::Sender; +use tokio::sync::Semaphore; + +use super::registry::{TaskHandler, TaskRegistry, TaskResult}; +use super::WorkerDispatcher; +use crate::job::Job; +use crate::scheduler::JobResult; + +/// Dispatches dequeued jobs to handlers from a [`TaskRegistry`], at most +/// `num_workers` concurrently. +pub struct NativeDispatcher { + registry: Arc, + num_workers: usize, + shutdown: AtomicBool, +} + +impl NativeDispatcher { + pub fn new(registry: TaskRegistry, num_workers: usize) -> Self { + Self { + registry: Arc::new(registry), + num_workers: num_workers.max(1), + shutdown: AtomicBool::new(false), + } + } +} + +/// Build the [`JobResult`] for a finished handler invocation. +fn job_result(job: &Job, outcome: TaskResult, started: Instant) -> JobResult { + let wall_time_ns: i64 = started.elapsed().as_nanos().try_into().unwrap_or(i64::MAX); + match outcome { + Ok(result) => JobResult::Success { + job_id: job.id.clone(), + result, + task_name: job.task_name.clone(), + wall_time_ns, + }, + Err(task_error) => JobResult::Failure { + job_id: job.id.clone(), + error: task_error.message, + retry_count: job.retry_count, + max_retries: job.max_retries, + task_name: job.task_name.clone(), + wall_time_ns, + should_retry: task_error.retryable, + // Timeouts are synthesized server-side by the stale-job reap. + timed_out: false, + }, + } +} + +#[async_trait] +impl WorkerDispatcher for NativeDispatcher { + async fn run( + &self, + mut job_rx: tokio::sync::mpsc::Receiver, + result_tx: Sender, + ) { + let semaphore = Arc::new(Semaphore::new(self.num_workers)); + + while let Some(job) = job_rx.recv().await { + if self.shutdown.load(Ordering::Relaxed) { + break; + } + let permit = match semaphore.clone().acquire_owned().await { + Ok(permit) => permit, + Err(_) => break, // Semaphore closed. + }; + + let handler = match self.registry.get(&job.task_name) { + Some(handler) => handler.clone(), + None => { + // Fatal, not retryable: no amount of retrying makes an + // unregistered task runnable on this worker. + let error = super::registry::TaskError::fatal(format!( + "task not registered: {}", + job.task_name + )); + let _ = result_tx.send(job_result(&job, Err(error), Instant::now())); + continue; + } + }; + + let tx = result_tx.clone(); + match handler { + TaskHandler::Sync(run) => { + tokio::task::spawn_blocking(move || { + let _permit = permit; // Hold the slot until done. + let started = Instant::now(); + let outcome = run(&job); + let _ = tx.send(job_result(&job, outcome, started)); + }); + } + TaskHandler::Async(make_future) => { + tokio::spawn(async move { + let _permit = permit; + let started = Instant::now(); + let outcome = make_future(job.clone()).await; + let _ = tx.send(job_result(&job, outcome, started)); + }); + } + } + } + } + + fn shutdown(&self) { + self.shutdown.store(true, Ordering::Relaxed); + } +} diff --git a/crates/taskito-core/src/worker.rs b/crates/taskito-core/src/worker/mod.rs similarity index 71% rename from crates/taskito-core/src/worker.rs rename to crates/taskito-core/src/worker/mod.rs index e9b451d6..922eef60 100644 --- a/crates/taskito-core/src/worker.rs +++ b/crates/taskito-core/src/worker/mod.rs @@ -1,3 +1,11 @@ +pub mod dispatcher; +pub mod registry; +pub mod runner; + +pub use dispatcher::NativeDispatcher; +pub use registry::{TaskError, TaskHandler, TaskRegistry, TaskResult}; +pub use runner::{Worker, WorkerHandle}; + use async_trait::async_trait; use crossbeam_channel::Sender; @@ -5,7 +13,8 @@ use crate::job::Job; use crate::scheduler::JobResult; /// Abstraction for worker pool implementations. -/// Core defines the interface; language-specific crates provide implementations. +/// Core defines the interface; the [`NativeDispatcher`] runs registered Rust +/// handlers, and the language bindings provide their own pools. #[async_trait] pub trait WorkerDispatcher: Send + Sync { /// Start dispatching jobs from the receiver. Runs until channel closes. @@ -23,3 +32,6 @@ pub trait WorkerDispatcher: Send + Sync { /// storage. fn notify_cancel(&self, _job_id: &str) {} } + +#[cfg(test)] +mod tests; diff --git a/crates/taskito-core/src/worker/registry.rs b/crates/taskito-core/src/worker/registry.rs new file mode 100644 index 00000000..75a49d65 --- /dev/null +++ b/crates/taskito-core/src/worker/registry.rs @@ -0,0 +1,111 @@ +//! Task registry: maps a job's `task_name` to the Rust handler that runs it. +//! +//! The language bindings keep their own registries (a Python dict, a JS map); +//! this one exists so a pure-Rust consumer has a first-class way to say +//! "`send_email` means this function". + +use std::collections::HashMap; +use std::future::Future; +use std::pin::Pin; +use std::sync::Arc; + +use crate::job::Job; + +/// What a handler produced: optional result bytes (stored on the job and +/// readable via `get_job`), or a [`TaskError`]. +pub type TaskResult = std::result::Result>, TaskError>; + +/// A boxed future returned by async handlers. +pub type TaskFuture = Pin + Send + 'static>>; + +/// A failed task execution. `retryable` feeds the scheduler's retry decision: +/// a retryable failure follows the job's retry policy; a fatal one goes +/// straight to the dead-letter queue. +#[derive(Debug, Clone)] +pub struct TaskError { + pub message: String, + pub retryable: bool, +} + +impl TaskError { + /// A failure the scheduler may retry (subject to the job's retry policy). + pub fn retryable(message: impl Into) -> Self { + TaskError { + message: message.into(), + retryable: true, + } + } + + /// A failure that must not be retried — the job dead-letters immediately. + pub fn fatal(message: impl Into) -> Self { + TaskError { + message: message.into(), + retryable: false, + } + } +} + +impl std::fmt::Display for TaskError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str(&self.message) + } +} + +impl std::error::Error for TaskError {} + +/// A registered handler: blocking (run on a blocking thread) or async +/// (spawned on the worker's tokio runtime). +#[derive(Clone)] +pub enum TaskHandler { + Sync(Arc TaskResult + Send + Sync>), + Async(Arc TaskFuture + Send + Sync>), +} + +/// Name → handler map consulted by the dispatcher for every dequeued job. +#[derive(Clone, Default)] +pub struct TaskRegistry { + handlers: HashMap, +} + +impl TaskRegistry { + pub fn new() -> Self { + Self::default() + } + + /// Register a blocking handler for `task_name`. It runs on a dedicated + /// blocking thread, so it may block freely (I/O, CPU work). + pub fn register( + &mut self, + task_name: impl Into, + handler: impl Fn(&Job) -> TaskResult + Send + Sync + 'static, + ) { + self.handlers + .insert(task_name.into(), TaskHandler::Sync(Arc::new(handler))); + } + + /// Register an async handler for `task_name`. The future is spawned on + /// the worker's runtime, so it must not block the executor. + pub fn register_async(&mut self, task_name: impl Into, handler: F) + where + F: Fn(Job) -> Fut + Send + Sync + 'static, + Fut: Future + Send + 'static, + { + self.handlers.insert( + task_name.into(), + TaskHandler::Async(Arc::new(move |job| Box::pin(handler(job)))), + ); + } + + pub fn get(&self, task_name: &str) -> Option<&TaskHandler> { + self.handlers.get(task_name) + } + + /// Names of every registered task. + pub fn task_names(&self) -> impl Iterator { + self.handlers.keys().map(String::as_str) + } + + pub fn is_empty(&self) -> bool { + self.handlers.is_empty() + } +} diff --git a/crates/taskito-core/src/worker/runner.rs b/crates/taskito-core/src/worker/runner.rs new file mode 100644 index 00000000..6a0623b2 --- /dev/null +++ b/crates/taskito-core/src/worker/runner.rs @@ -0,0 +1,331 @@ +//! Turn-key worker: wires a [`Scheduler`], a [`NativeDispatcher`], the result +//! drain loop, and the heartbeat/reap cadence into one `Worker::spawn()` call — +//! the zero-to-executed-task path for a Rust consumer. +//! +//! Mirrors the orchestration the language bindings hand-roll: a dedicated +//! tokio runtime drives `Scheduler::run` and `WorkerDispatcher::run`, a drain +//! thread feeds `JobResult`s back into `Scheduler::handle_results`, and a +//! heartbeat thread keeps the worker registered and runs the elected reaps. + +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::mpsc as std_mpsc; +use std::sync::Arc; +use std::thread; +use std::time::Duration; + +use crate::error::Result; +use crate::scheduler::{QueueConfig, ResultOutcome, Scheduler, SchedulerConfig, TaskConfig}; +use crate::storage::{ + reap_dead_workers_if_leader, sweep_ephemeral_subscriptions, Storage, StorageBackend, +}; + +use super::dispatcher::NativeDispatcher; +use super::registry::{TaskRegistry, TaskResult}; +use super::WorkerDispatcher; + +const HEARTBEAT_INTERVAL: Duration = Duration::from_secs(5); +const DRAIN_POLL: Duration = Duration::from_millis(100); + +/// Callback invoked with each processed [`ResultOutcome`] (success, retry, +/// dead-letter, cancel) — the native analogue of a binding's middleware hooks. +pub type OutcomeCallback = Arc; + +/// Builder for a running worker. Register handlers, then [`Worker::spawn`]. +pub struct Worker { + storage: StorageBackend, + registry: TaskRegistry, + queues: Vec, + num_workers: usize, + namespace: Option, + scheduler_config: SchedulerConfig, + task_configs: Vec<(String, TaskConfig)>, + queue_configs: Vec<(String, QueueConfig)>, + worker_id: Option, + on_outcome: Option, +} + +impl Worker { + pub fn new(storage: StorageBackend) -> Self { + Self { + storage, + registry: TaskRegistry::new(), + queues: vec!["default".to_string()], + num_workers: 4, + namespace: None, + scheduler_config: SchedulerConfig::default(), + task_configs: Vec::new(), + queue_configs: Vec::new(), + worker_id: None, + on_outcome: None, + } + } + + /// Queues this worker consumes (default: `["default"]`). + pub fn queues(mut self, queues: impl IntoIterator>) -> Self { + self.queues = queues.into_iter().map(Into::into).collect(); + self + } + + /// Maximum concurrently executing tasks (default: 4). + pub fn num_workers(mut self, num_workers: usize) -> Self { + self.num_workers = num_workers.max(1); + self + } + + pub fn namespace(mut self, namespace: impl Into) -> Self { + self.namespace = Some(namespace.into()); + self + } + + /// Override the scheduler configuration. `max_in_flight` is capped to the + /// worker pool size at spawn when unset. + pub fn scheduler_config(mut self, config: SchedulerConfig) -> Self { + self.scheduler_config = config; + self + } + + /// Per-task resilience policy (retry, rate limit, circuit breaker). + pub fn task_config(mut self, task_name: impl Into, config: TaskConfig) -> Self { + self.task_configs.push((task_name.into(), config)); + self + } + + /// Per-queue policy (rate limit, concurrency). + pub fn queue_config(mut self, queue_name: impl Into, config: QueueConfig) -> Self { + self.queue_configs.push((queue_name.into(), config)); + self + } + + /// Explicit worker id (default: `rust-worker-`). + pub fn worker_id(mut self, worker_id: impl Into) -> Self { + self.worker_id = Some(worker_id.into()); + self + } + + /// Observe every processed outcome (the middleware-hook analogue). + pub fn on_outcome(mut self, callback: impl Fn(&ResultOutcome) + Send + Sync + 'static) -> Self { + self.on_outcome = Some(Arc::new(callback)); + self + } + + /// Register a blocking handler. See [`TaskRegistry::register`]. + pub fn register( + mut self, + task_name: impl Into, + handler: impl Fn(&crate::job::Job) -> TaskResult + Send + Sync + 'static, + ) -> Self { + self.registry.register(task_name, handler); + self + } + + /// Register an async handler. See [`TaskRegistry::register_async`]. + pub fn register_async(mut self, task_name: impl Into, handler: F) -> Self + where + F: Fn(crate::job::Job) -> Fut + Send + Sync + 'static, + Fut: std::future::Future + Send + 'static, + { + self.registry.register_async(task_name, handler); + self + } + + /// Register this worker and start the scheduler, dispatcher, result-drain + /// loop, and heartbeat. Returns a handle; call [`WorkerHandle::shutdown`] + /// to drain and stop. + pub fn spawn(self) -> Result { + let Worker { + storage, + registry, + queues, + num_workers, + namespace, + mut scheduler_config, + task_configs, + queue_configs, + worker_id, + on_outcome, + } = self; + + let worker_id = + worker_id.unwrap_or_else(|| format!("rust-worker-{}", uuid::Uuid::now_v7())); + + storage.register_worker( + &worker_id, + &queues.join(","), + None, + None, + None, + num_workers as i32, + None, + Some(std::process::id() as i32), + Some("native"), + )?; + + // Bound dispatch to the pool size so this scheduler never claims more + // than its workers can run; also makes `in_flight_settled` meaningful. + if scheduler_config.max_in_flight.is_none() { + scheduler_config.max_in_flight = Some(num_workers); + } + + let mut scheduler = Scheduler::new(storage.clone(), queues, scheduler_config, namespace); + scheduler.set_claim_owner(worker_id.clone()); + for (task_name, config) in task_configs { + scheduler.register_task(task_name, config); + } + for (queue_name, config) in queue_configs { + scheduler.register_queue_config(queue_name, config); + } + let scheduler = Arc::new(scheduler); + let shutdown = scheduler.shutdown_handle(); + + let (job_tx, job_rx) = tokio::sync::mpsc::channel(num_workers * 2); + let (result_tx, result_rx) = crossbeam_channel::bounded(num_workers * 2); + + let dispatcher = Arc::new(NativeDispatcher::new(registry, num_workers)); + let runtime_done = Arc::new(AtomicBool::new(false)); + + // Runtime thread: scheduler dispatch + task execution. `result_tx` + // moves in, so the drain side disconnects once execution is done. + let runtime_thread = { + let scheduler = scheduler.clone(); + let dispatcher = dispatcher.clone(); + let runtime_done = runtime_done.clone(); + thread::Builder::new() + .name(format!("{worker_id}-runtime")) + .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 { + let scheduler_task = tokio::spawn({ + let scheduler = scheduler.clone(); + async move { scheduler.run(job_tx).await } + }); + let dispatch_task = + tokio::spawn(async move { dispatcher.run(job_rx, result_tx).await }); + let _ = tokio::join!(scheduler_task, dispatch_task); + }); + runtime_done.store(true, Ordering::Release); + }) + .map_err(spawn_error)? + }; + + // Drain thread: feed results back into the scheduler and surface + // outcomes. Exits when execution has finished and every dispatched + // job's result has been handled. + let drain_thread = { + let scheduler = scheduler.clone(); + let runtime_done = runtime_done.clone(); + thread::Builder::new() + .name(format!("{worker_id}-drain")) + .spawn(move || loop { + match result_rx.recv_timeout(DRAIN_POLL) { + Ok(first) => { + let mut batch = vec![first]; + while let Ok(more) = result_rx.try_recv() { + batch.push(more); + } + for handled in scheduler.handle_results(batch) { + match handled { + Ok(outcome) => { + if let Some(callback) = &on_outcome { + callback(&outcome); + } + } + Err(handling_error) => { + log::error!("result handling failed: {handling_error}"); + } + } + } + } + Err(crossbeam_channel::RecvTimeoutError::Timeout) => { + if runtime_done.load(Ordering::Acquire) && scheduler.in_flight_settled() + { + break; + } + } + Err(crossbeam_channel::RecvTimeoutError::Disconnected) => break, + } + }) + .map_err(spawn_error)? + }; + + // Heartbeat thread: liveness + the elected cluster reaps. The stop + // sender doubles as the stop signal — dropping it ends the loop. + let (stop_tx, stop_rx) = std_mpsc::channel::<()>(); + let heartbeat_thread = { + let storage = storage.clone(); + let worker_id = worker_id.clone(); + thread::Builder::new() + .name(format!("{worker_id}-heartbeat")) + .spawn(move || { + while let Err(std_mpsc::RecvTimeoutError::Timeout) = + stop_rx.recv_timeout(HEARTBEAT_INTERVAL) + { + if let Err(heartbeat_error) = storage.heartbeat(&worker_id, None) { + log::warn!("worker heartbeat failed: {heartbeat_error}"); + } + reap_dead_workers_if_leader(&storage, &worker_id); + if let Err(sweep_error) = + sweep_ephemeral_subscriptions(&storage, Some(&worker_id)) + { + log::warn!("ephemeral subscription reap failed: {sweep_error}"); + } + } + }) + .map_err(spawn_error)? + }; + + Ok(WorkerHandle { + worker_id, + storage, + shutdown, + dispatcher, + stop_tx: Some(stop_tx), + threads: vec![runtime_thread, drain_thread, heartbeat_thread], + }) + } +} + +fn spawn_error(io_error: std::io::Error) -> crate::error::QueueError { + crate::error::QueueError::Worker(format!("failed to spawn worker thread: {io_error}")) +} + +/// Handle to a running [`Worker`]. Dropping it without calling +/// [`WorkerHandle::shutdown`] leaves the worker running detached. +pub struct WorkerHandle { + worker_id: String, + storage: StorageBackend, + shutdown: Arc, + dispatcher: Arc, + stop_tx: Option>, + threads: Vec>, +} + +impl WorkerHandle { + pub fn worker_id(&self) -> &str { + &self.worker_id + } + + /// Stop dispatching, drain in-flight work, stop the heartbeat, and + /// unregister the worker. Blocks until every thread has exited. + pub fn shutdown(mut self) -> Result<()> { + self.shutdown.notify_one(); + self.dispatcher.shutdown(); + // Dropping the stop sender ends the heartbeat loop on its next wake. + self.stop_tx.take(); + for thread in self.threads.drain(..) { + if thread.join().is_err() { + log::error!("worker thread panicked during shutdown"); + } + } + self.storage.unregister_worker(&self.worker_id)?; + // This worker's own ephemeral subscriptions are now dead-owned; reap + // immediately instead of waiting for a peer's heartbeat tick. + if let Err(sweep_error) = sweep_ephemeral_subscriptions(&self.storage, None) { + log::warn!("shutdown subscription sweep failed: {sweep_error}"); + } + Ok(()) + } +} diff --git a/crates/taskito-core/src/worker/tests.rs b/crates/taskito-core/src/worker/tests.rs new file mode 100644 index 00000000..86c10a04 --- /dev/null +++ b/crates/taskito-core/src/worker/tests.rs @@ -0,0 +1,169 @@ +//! End-to-end tests for the native worker: enqueue → dispatch → execute → +//! result handling, against an in-memory SQLite backend. + +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}; + +fn test_backend() -> StorageBackend { + StorageBackend::Sqlite(SqliteStorage::in_memory().expect("in-memory sqlite")) +} + +fn make_job(task_name: &str, payload: &[u8], max_retries: i32) -> NewJob { + NewJob { + queue: "default".to_string(), + task_name: task_name.to_string(), + payload: payload.to_vec(), + priority: 0, + scheduled_at: now_millis(), + max_retries, + timeout_ms: 30_000, + unique_key: None, + metadata: None, + notes: None, + depends_on: vec![], + expires_at: None, + result_ttl_ms: None, + namespace: None, + } +} + +fn wait_for_status(storage: &StorageBackend, job_id: &str, wanted: JobStatus) -> Job { + let deadline = Instant::now() + Duration::from_secs(10); + loop { + if let Some(job) = storage.get_job(job_id).expect("get_job") { + if job.status == wanted { + return job; + } + } + assert!( + Instant::now() < deadline, + "job {job_id} never reached {wanted:?}" + ); + std::thread::sleep(Duration::from_millis(10)); + } +} + +fn wait_for_dead(storage: &StorageBackend, job_id: &str) { + let deadline = Instant::now() + Duration::from_secs(10); + loop { + let dead = storage.list_dead(50, 0).expect("list_dead"); + if dead.iter().any(|d| d.original_job_id == job_id) { + return; + } + assert!( + Instant::now() < deadline, + "job {job_id} never dead-lettered" + ); + std::thread::sleep(Duration::from_millis(10)); + } +} + +#[test] +fn sync_handler_executes_and_stores_result() { + let storage = test_backend(); + let handle = Worker::new(storage.clone()) + .num_workers(2) + .register("echo", |job: &Job| Ok(Some(job.payload.clone()))) + .spawn() + .expect("spawn"); + + let job = storage.enqueue(make_job("echo", b"ping", 3)).unwrap(); + let done = wait_for_status(&storage, &job.id, JobStatus::Complete); + assert_eq!(done.result.as_deref(), Some(&b"ping"[..])); + + handle.shutdown().expect("shutdown"); +} + +#[test] +fn async_handler_executes() { + let storage = test_backend(); + let handle = Worker::new(storage.clone()) + .register_async("sleepy", |_job: Job| async move { + tokio::time::sleep(Duration::from_millis(5)).await; + Ok(Some(b"woke".to_vec())) + }) + .spawn() + .expect("spawn"); + + let job = storage.enqueue(make_job("sleepy", b"", 3)).unwrap(); + let done = wait_for_status(&storage, &job.id, JobStatus::Complete); + assert_eq!(done.result.as_deref(), Some(&b"woke"[..])); + + handle.shutdown().expect("shutdown"); +} + +#[test] +fn retryable_failure_retries_then_dead_letters() { + let storage = test_backend(); + let attempts = Arc::new(AtomicU32::new(0)); + let seen = attempts.clone(); + let handle = Worker::new(storage.clone()) + .task_config( + "flaky", + TaskConfig { + retry_policy: RetryPolicy { + max_retries: 2, + base_delay_ms: 10, + max_delay_ms: 20, + custom_delays_ms: None, + }, + ..TaskConfig::default() + }, + ) + .register("flaky", move |_job: &Job| { + seen.fetch_add(1, Ordering::SeqCst); + Err(TaskError::retryable("boom")) + }) + .spawn() + .expect("spawn"); + + let job = storage.enqueue(make_job("flaky", b"", 2)).unwrap(); + wait_for_dead(&storage, &job.id); + // First attempt + 2 retries. + assert_eq!(attempts.load(Ordering::SeqCst), 3); + + handle.shutdown().expect("shutdown"); +} + +#[test] +fn fatal_failure_skips_retries() { + let storage = test_backend(); + let attempts = Arc::new(AtomicU32::new(0)); + let seen = attempts.clone(); + let handle = Worker::new(storage.clone()) + .register("doomed", move |_job: &Job| { + seen.fetch_add(1, Ordering::SeqCst); + Err(TaskError::fatal("unrecoverable")) + }) + .spawn() + .expect("spawn"); + + let job = storage.enqueue(make_job("doomed", b"", 5)).unwrap(); + wait_for_dead(&storage, &job.id); + assert_eq!(attempts.load(Ordering::SeqCst), 1, "fatal must not retry"); + + handle.shutdown().expect("shutdown"); +} + +#[test] +fn unregistered_task_dead_letters_without_retry() { + let storage = test_backend(); + let handle = Worker::new(storage.clone()) + .register("known", |_job: &Job| Ok(None)) + .spawn() + .expect("spawn"); + + let job = storage.enqueue(make_job("unknown", b"", 5)).unwrap(); + wait_for_dead(&storage, &job.id); + + handle.shutdown().expect("shutdown"); +}