diff --git a/crates/taskito-java/src/convert.rs b/crates/taskito-java/src/convert.rs index 6d6e4888..f0c7416d 100644 --- a/crates/taskito-java/src/convert.rs +++ b/crates/taskito-java/src/convert.rs @@ -319,9 +319,9 @@ pub struct SubscriptionSpec { pub timeout_ms: Option, } -/// A task's retry-backoff curve. Fields left unset fall back to the core's -/// `RetryPolicy` defaults; the per-job retry budget travels via `maxRetries` on -/// enqueue, so it is deliberately absent here. +/// A task's policy: retry curve, throttling, and concurrency caps. Fields left +/// unset fall back to the core's defaults. The per-job retry *budget* travels +/// via `maxRetries` on enqueue, so it is deliberately absent here. #[derive(Deserialize)] #[serde(rename_all = "camelCase")] pub struct TaskRetryConfig { @@ -336,6 +336,16 @@ pub struct TaskRetryConfig { pub circuit_breaker_cooldown_ms: Option, pub circuit_breaker_half_open_probes: Option, pub circuit_breaker_half_open_success_rate: Option, + /// Rate-limit spec like `"100/m"`, `"50/s"`, `"3600/h"`. + pub rate_limit: Option, + /// Cap on how fast this task may *retry*, across all of its jobs. Same spec + /// as `rate_limit`; once spent, failures dead-letter instead of retrying. + pub retry_budget: Option, + /// Cap on concurrently-running jobs of this task, across the cluster. + pub max_concurrent: Option, + /// Cap on this task's share of one worker's dispatch slots. In-process, + /// unlike `max_concurrent`, which is cluster-wide and costs a DB read. + pub max_in_flight_per_task: Option, } /// Filter accepted by `NativeQueue.listJobs`. diff --git a/crates/taskito-java/src/worker.rs b/crates/taskito-java/src/worker.rs index d5f8d69b..9c36045b 100644 --- a/crates/taskito-java/src/worker.rs +++ b/crates/taskito-java/src/worker.rs @@ -13,6 +13,7 @@ use jni::objects::{GlobalRef, JByteArray, JClass, JObject, JString, JValue}; use jni::sys::jlong; use jni::JNIEnv; use taskito_core::resilience::circuit_breaker::CircuitBreakerConfig; +use taskito_core::resilience::rate_limiter::RateLimitConfig; use taskito_core::resilience::retry::RetryPolicy; use taskito_core::scheduler::{ResultOutcome, TaskConfig}; use taskito_core::worker::WorkerDispatcher; @@ -91,6 +92,12 @@ fn start_worker( #[cfg(feature = "mesh")] let mesh_worker_id = worker_id.clone(); + // Validate every task policy before writing any persistent state: this is + // the last fallible step that has no side effects, and returning Err past + // the writes below would leave a worker row and its subscriptions behind + // with no handle to run the lifecycle loop that cleans them up. + let task_policies = build_task_policies(options.task_configs)?; + // Create the live worker row before its ephemeral subscriptions exist: // the reaper only spares owned rows whose owner is registered, so this // ordering (plus the core's registration grace window) keeps a concurrent @@ -105,7 +112,9 @@ fn start_worker( // Claim execution under this worker's id so dead-worker recovery can // attribute orphaned jobs (matches the worker id registered above). scheduler.set_claim_owner(worker_id.clone()); - register_task_policies(&mut scheduler, options.task_configs); + for (name, policy) in task_policies { + scheduler.register_task(name, policy); + } let scheduler = Arc::new(scheduler); let shutdown = scheduler.shutdown_handle(); let heartbeat_stop = Arc::new(Notify::new()); @@ -156,7 +165,9 @@ fn start_worker( // result sender has dropped (dispatcher done). let drain_scheduler = scheduler; let drain_callbacks = callbacks; - runtime.spawn_blocking(move || drain_results(result_rx, drain_scheduler, drain_callbacks)); + runtime.spawn_blocking(move || { + drain_results(result_rx, drain_scheduler, drain_callbacks, capacity) + }); // Lifecycle loop: heartbeat until stopped, then unregister. spawn_lifecycle( @@ -234,11 +245,21 @@ fn register_subscriptions( Ok(()) } -/// Register each task's retry-backoff curve with the scheduler. Only the curve -/// is set here — the core resolves the retry budget from the job's `max_retries` -/// — so unset fields keep the core [`RetryPolicy`] defaults. -fn register_task_policies(scheduler: &mut Scheduler, configs: Option>) { - let Some(configs) = configs else { return }; +/// Build each task's policy — retry curve, throttling, concurrency caps — from +/// its wire config. Unset fields keep the core's defaults; the per-job retry +/// budget still resolves from the job's `max_retries`. +/// +/// Fails on a malformed rate spec rather than registering the task without it: +/// silently dropping a throttle the caller asked for is worse than not starting. +/// Pure, so the caller can validate before writing any persistent worker state +/// and a rejected config leaves nothing behind. +fn build_task_policies( + configs: Option>, +) -> Result, crate::error::BindingError> { + let Some(configs) = configs else { + return Ok(Vec::new()); + }; + let mut built = Vec::with_capacity(configs.len()); for config in configs { let mut retry_policy = RetryPolicy::default(); if let Some(custom) = config.custom_delays_ms { @@ -271,27 +292,52 @@ fn register_task_policies(scheduler: &mut Scheduler, configs: Option, +) -> Result, crate::error::BindingError> { + match spec { + Some(s) => RateLimitConfig::parse(s).map(Some).ok_or_else(|| { + crate::error::BindingError::new(format!( + "invalid {field} '{s}' on task '{task}' (expected e.g. '100/m')" + )) + }), + None => Ok(None), } } -/// Drain results: apply each to storage and invoke `WorkerBridge.onOutcome`. +/// Drain results: finalize each batch in one transaction, then invoke +/// `WorkerBridge.onOutcome` per job. +/// +/// `max_batch` bounds one drain. Unlike the other shells, this result channel is +/// unbounded (so the dispatcher never blocks a runtime worker), so nothing else +/// would cap how much a single drain swallows. fn drain_results( result_rx: crossbeam_channel::Receiver, scheduler: Arc, callbacks: GlobalRef, + max_batch: usize, ) { let vm = jvm::vm(); let mut env = match vm.attach_current_thread() { @@ -301,25 +347,40 @@ fn drain_results( return; } }; - while let Ok(result) = result_rx.recv() { - let outcome = match scheduler.handle_result(result) { - Ok(outcome) => outcome, - Err(e) => { - log::error!("[taskito-java] result handling error: {e}"); - continue; + while let Ok(first) = result_rx.recv() { + // Finalize everything already queued in one transaction rather than one + // per wake. + let mut batch = vec![first]; + while batch.len() < max_batch { + match result_rx.try_recv() { + Ok(more) => batch.push(more), + Err(_) => break, } - }; - // Each outcome allocates several JNI locals on this long-lived attached - // env; scope them in a frame so a busy worker can't exhaust the - // local-reference table over the loop's lifetime. - let framed = env.with_local_frame(16, |env| { - if let Err(e) = call_on_outcome(env, &callbacks, &outcome) { - log::error!("[taskito-java] onOutcome failed: {e}"); + } + + // One outcome per input, in order, so each job is reported exactly once. + for outcome in scheduler.handle_results(batch) { + let outcome = match outcome { + Ok(outcome) => outcome, + Err(e) => { + log::error!("[taskito-java] result handling error: {e}"); + continue; + } + }; + // Each outcome allocates several JNI locals on this long-lived attached + // env; scope them in a frame so a busy worker can't exhaust the + // local-reference table over the loop's lifetime. The frame stays per + // outcome — sized for one `call_on_outcome`, it would overflow if it + // wrapped the whole batch. + let framed = env.with_local_frame(16, |env| { + if let Err(e) = call_on_outcome(env, &callbacks, &outcome) { + log::error!("[taskito-java] onOutcome failed: {e}"); + } + Ok::<(), jni::errors::Error>(()) + }); + if let Err(e) = framed { + log::error!("[taskito-java] drain local frame failed: {e}"); } - Ok::<(), jni::errors::Error>(()) - }); - if let Err(e) = framed { - log::error!("[taskito-java] drain local frame failed: {e}"); } } } diff --git a/crates/taskito-node/src/config.rs b/crates/taskito-node/src/config.rs index a185fafb..02806571 100644 --- a/crates/taskito-node/src/config.rs +++ b/crates/taskito-node/src/config.rs @@ -109,8 +109,14 @@ pub struct TaskConfigInput { pub retry_base_delay_ms: Option, pub retry_max_delay_ms: Option, pub max_concurrent: Option, + /// Cap on this task's share of one worker's dispatch slots, so a slow task + /// cannot occupy the whole pool. In-process, unlike `max_concurrent`. + pub max_in_flight_per_task: Option, /// Rate-limit spec like `"100/m"`, `"50/s"`, `"3600/h"`. pub rate_limit: Option, + /// Cap on how fast this task may *retry*, across all of its jobs. Same spec + /// as `rate_limit`; once spent, failures dead-letter instead of retrying. + pub retry_budget: Option, pub circuit_breaker: Option, } @@ -160,6 +166,10 @@ pub struct MeshWorkerConfig { pub struct WorkerOptions { pub queues: Option>, pub channel_capacity: Option, + /// Jobs this worker runs at once. Unset leaves in-flight work unbounded, + /// which is the historical behaviour — a worker then claims jobs it cannot + /// run yet, stranding them Running and starving peers on the same database. + pub concurrency: Option, /// Jobs claimed per scheduler poll (default 1). pub batch_size: Option, pub task_configs: Option>, diff --git a/crates/taskito-node/src/convert/task_config.rs b/crates/taskito-node/src/convert/task_config.rs index 9bc63837..e54b67fe 100644 --- a/crates/taskito-node/src/convert/task_config.rs +++ b/crates/taskito-node/src/convert/task_config.rs @@ -15,13 +15,25 @@ const DEFAULT_MAX_RETRIES: i32 = 3; const DEFAULT_HALF_OPEN_PROBES: i32 = 5; const DEFAULT_HALF_OPEN_SUCCESS_RATE: f64 = 0.8; -/// Parse an optional rate-limit spec, failing fast on a malformed value rather -/// than silently disabling throttling (a misconfigured limit is a config error). -fn parse_rate_limit(spec: Option<&str>) -> Result> { +/// Parse an optional rate spec, failing fast on a malformed value rather than +/// silently disabling the cap (a misconfigured limit is a config error). +/// +/// Names the option and what carries it — several options share this grammar, +/// and both tasks and queues configure one, so the error cannot say which +/// config is wrong without both. `scope` is the kind ("task", "queue"), `name` +/// its name. +fn parse_rate_spec( + field: &str, + scope: &str, + name: &str, + spec: Option<&str>, +) -> Result> { match spec { - Some(s) => RateLimitConfig::parse(s) - .map(Some) - .ok_or_else(|| invalid_arg(format!("invalid rateLimit '{s}' (expected e.g. '100/m')"))), + Some(s) => RateLimitConfig::parse(s).map(Some).ok_or_else(|| { + invalid_arg(format!( + "invalid {field} '{s}' on {scope} '{name}' (expected e.g. '100/m')" + )) + }), None => Ok(None), } } @@ -35,17 +47,26 @@ pub fn task_config(input: &TaskConfigInput) -> Result { max_delay_ms: input.retry_max_delay_ms.unwrap_or(DEFAULT_RETRY_MAX_MS), custom_delays_ms: None, }, - rate_limit: parse_rate_limit(input.rate_limit.as_deref())?, + rate_limit: parse_rate_spec( + "rateLimit", + "task", + &input.name, + input.rate_limit.as_deref(), + )?, circuit_breaker: input .circuit_breaker .as_ref() .map(circuit_breaker_config) .transpose()?, - // Not surfaced on this SDK yet; retries stay bounded per job. - retry_budget: None, + // Same "100/m" grammar as rate_limit, so one parser serves both. + retry_budget: parse_rate_spec( + "retryBudget", + "task", + &input.name, + input.retry_budget.as_deref(), + )?, max_concurrent: input.max_concurrent, - // Not surfaced on this SDK yet; the whole pool stays available. - max_in_flight_per_task: None, + max_in_flight_per_task: input.max_in_flight_per_task.map(|n| n.max(1) as usize), }) } @@ -88,7 +109,12 @@ fn circuit_breaker_config(input: &CircuitBreakerInput) -> Result Result { Ok(QueueConfig { - rate_limit: parse_rate_limit(input.rate_limit.as_deref())?, + rate_limit: parse_rate_spec( + "rateLimit", + "queue", + &input.name, + input.rate_limit.as_deref(), + )?, max_concurrent: input.max_concurrent, }) } diff --git a/crates/taskito-node/src/dispatcher.rs b/crates/taskito-node/src/dispatcher.rs index 42f8b2c1..39aec9b7 100644 --- a/crates/taskito-node/src/dispatcher.rs +++ b/crates/taskito-node/src/dispatcher.rs @@ -5,6 +5,7 @@ //! `Promise`, and reports a [`JobResult`] back to the scheduler. This is //! the Node mirror of the Python shell's worker pool. +use std::sync::Arc; use std::time::{Duration, Instant}; use crossbeam_channel::Sender; @@ -14,7 +15,7 @@ use taskito_core::job::Job; use taskito_core::scheduler::JobResult; use taskito_core::worker::WorkerDispatcher; use taskito_core::{Storage, StorageBackend}; -use tokio::sync::oneshot; +use tokio::sync::{oneshot, Semaphore}; use crate::convert::JsTaskInvocation; @@ -25,11 +26,31 @@ type TaskCallback = ThreadsafeFunction; pub struct NodeDispatcher { callback: TaskCallback, storage: StorageBackend, + /// Caps jobs running at once. Without it the loop spawns every job it is + /// handed and immediately takes the next, so nothing bounds concurrency. + /// + /// The scheduler's `max_in_flight` bounds what this worker *claims*; this + /// bounds what it *runs*, and is the only bound on the mesh path, where + /// stolen jobs never pass through this scheduler's in-flight accounting. + concurrency: Arc, } impl NodeDispatcher { - pub fn new(callback: TaskCallback, storage: StorageBackend) -> Self { - Self { callback, storage } + /// `concurrency` of `None` leaves execution unbounded, matching the + /// behaviour of a worker that never set the option. + pub fn new( + callback: TaskCallback, + storage: StorageBackend, + concurrency: Option, + ) -> Self { + let permits = concurrency + .map(|c| c.max(1)) + .unwrap_or(Semaphore::MAX_PERMITS); + Self { + callback, + storage, + concurrency: Arc::new(Semaphore::new(permits)), + } } } @@ -41,13 +62,22 @@ impl WorkerDispatcher for NodeDispatcher { result_tx: Sender, ) { // Run jobs concurrently — each invocation is independent and the JS - // side may be async. The scheduler bounds in-flight work via the - // channel capacity and per-task/queue concurrency gates. + // side may be async — but take a permit first. Spawning unconditionally + // drains the job channel as fast as the scheduler fills it, so the + // channel applies no backpressure and in-flight work is unbounded: the + // worker claims jobs it cannot run, stranding them Running and starving + // peers on the same database. Blocking here is what bounds it; the + // permit rides with the job and is released when the task finishes. while let Some(job) = job_rx.recv().await { + let permit = match self.concurrency.clone().acquire_owned().await { + Ok(p) => p, + Err(_) => break, + }; let callback = self.callback.clone(); let storage = self.storage.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; // A full bounded channel parks the sender — do it on the diff --git a/crates/taskito-node/src/worker.rs b/crates/taskito-node/src/worker.rs index ec672c15..ca713fdc 100644 --- a/crates/taskito-node/src/worker.rs +++ b/crates/taskito-node/src/worker.rs @@ -86,10 +86,20 @@ pub fn start_worker( .map(|c| (c as usize).max(1)) .unwrap_or(DEFAULT_CHANNEL_CAPACITY); + // Jobs this worker runs at once. Separate from `channel_capacity`, which + // buffers hand-offs and advertises the worker's capacity — matching Java, + // where the two are distinct options. + let concurrency = options.concurrency.map(|c| (c as usize).max(1)); + let mut config = SchedulerConfig::default(); if let Some(batch) = options.batch_size { config.batch_size = batch.max(1) as usize; } + // Bound in-flight work to what this worker will actually run, so it never + // claims more and strands the surplus Running while peers sharing the + // database skip it. Left unbounded when unset, so an existing worker does + // not silently acquire a cap it never asked for. + config.max_in_flight = concurrency; // The dispatcher reads cancel flags, and the lifecycle loop registers/heartbeats // — both need their own storage handle before `storage` moves into the scheduler. @@ -178,7 +188,7 @@ pub fn start_worker( } // Dispatcher loop: execute each job in JS, report results on `result_tx`. - let dispatcher = NodeDispatcher::new(callback, dispatcher_storage); + let dispatcher = NodeDispatcher::new(callback, dispatcher_storage, concurrency); spawn(async move { dispatcher.run(job_rx, result_tx).await; }); @@ -188,23 +198,50 @@ pub fn start_worker( // sender has dropped (i.e. after the dispatcher and all in-flight jobs end). let scheduler_results = scheduler; spawn_blocking(move || { - while let Ok(result) = result_rx.recv() { - // A panicking result must not kill the drain loop — a dead loop - // silently drops every later outcome. + while let Ok(first) = result_rx.recv() { + // Finalize everything already queued in one batched transaction + // rather than one per wake. + // + // Capped explicitly: the channel bounds how many results can sit in + // it at once, not how many this loop drains — senders refill slots + // while `try_recv` runs, so an unbounded drain could swallow a whole + // backlog into one Vec and stall finalization behind it. + let mut batch = vec![first]; + while batch.len() < capacity { + match result_rx.try_recv() { + Ok(more) => batch.push(more), + Err(_) => break, + } + } + + // A panicking batch must not kill the drain loop — a dead loop + // silently drops every later outcome. The batch widens a panic's + // reach from one outcome to N, which is acceptable: the realistic + // failure is a batch write error, and `handle_results` already + // falls back to finalizing those jobs one at a time. let handled = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { - scheduler_results.handle_result(result) + scheduler_results.handle_results(batch) })); match handled { // Surface each outcome to JS so the shell can emit events and run - // middleware (the events layer). - Ok(Ok(outcome)) => { - outcome_callback.call( - outcome_to_js(&outcome), - ThreadsafeFunctionCallMode::NonBlocking, - ); + // middleware (the events layer). `handle_results` returns one + // outcome per input, in order, so each job is reported once. + Ok(outcomes) => { + for outcome in outcomes { + match outcome { + Ok(outcome) => { + outcome_callback.call( + outcome_to_js(&outcome), + ThreadsafeFunctionCallMode::NonBlocking, + ); + } + Err(err) => { + log::error!("[taskito-node] result handling error: {err}") + } + } + } } - Ok(Err(err)) => log::error!("[taskito-node] result handling error: {err}"), - Err(_) => log::error!("[taskito-node] result handling panicked; outcome dropped"), + Err(_) => log::error!("[taskito-node] result handling panicked; outcomes dropped"), } } }); 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 dad14bda..d4be4596 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 @@ -198,6 +198,22 @@ private String options(ExecutableElement method) { if (boolValue(mirror, "idempotent", false)) { chain.append(".idempotent(true)"); } + String rateLimit = stringValue(mirror, "rateLimit", ""); + if (!rateLimit.isEmpty()) { + chain.append(".rateLimit(\"").append(rateLimit).append("\")"); + } + String retryBudget = stringValue(mirror, "retryBudget", ""); + if (!retryBudget.isEmpty()) { + chain.append(".retryBudget(\"").append(retryBudget).append("\")"); + } + long maxConcurrent = longValue(mirror, "maxConcurrent", 0); + if (maxConcurrent > 0) { + chain.append(".maxConcurrent(").append(maxConcurrent).append(")"); + } + long maxInFlightPerTask = longValue(mirror, "maxInFlightPerTask", 0); + if (maxInFlightPerTask > 0) { + chain.append(".maxInFlightPerTask(").append(maxInFlightPerTask).append(")"); + } appendCircuitBreaker(chain, mirror); return chain.toString(); } diff --git a/sdks/java/src/main/java/org/byteveda/taskito/annotation/TaskHandler.java b/sdks/java/src/main/java/org/byteveda/taskito/annotation/TaskHandler.java index 0dafa40d..115d46cf 100644 --- a/sdks/java/src/main/java/org/byteveda/taskito/annotation/TaskHandler.java +++ b/sdks/java/src/main/java/org/byteveda/taskito/annotation/TaskHandler.java @@ -36,6 +36,32 @@ /** Auto-derive an idempotency {@code uniqueKey} from the payload on every enqueue. */ boolean idempotent() default false; + /** + * Rate-limit spec like {@code "100/m"} ({@code s}, {@code m} and {@code h} + * suffixes); empty (default) leaves the task unthrottled. A malformed spec + * fails the worker's start rather than running unthrottled. + */ + String rateLimit() default ""; + + /** + * Cap on how fast this task may retry, across all of its jobs — a + * spec like {@code "100/m"}; empty (default) leaves retries uncapped. Once + * spent, failures dead-letter instead of retrying. + */ + String retryBudget() default ""; + + /** + * Cap on concurrently-running jobs of this task across the cluster; 0 + * (default) leaves it uncapped. + */ + int maxConcurrent() default 0; + + /** + * Cap on this task's share of one worker's dispatch slots, so a slow task + * cannot occupy the whole pool; 0 (default) lets it use the whole pool. + */ + int maxInFlightPerTask() default 0; + /** Circuit-breaker failure threshold; 0 (default) leaves the breaker off. */ int circuitBreakerThreshold() default 0; diff --git a/sdks/java/src/main/java/org/byteveda/taskito/task/Task.java b/sdks/java/src/main/java/org/byteveda/taskito/task/Task.java index 34ac14a3..b429ea87 100644 --- a/sdks/java/src/main/java/org/byteveda/taskito/task/Task.java +++ b/sdks/java/src/main/java/org/byteveda/taskito/task/Task.java @@ -22,6 +22,10 @@ public final class Task { private final List codecs; private final boolean idempotent; private final CircuitBreakerConfig circuitBreaker; + private final String rateLimit; + private final String retryBudget; + private final Integer maxConcurrent; + private final Integer maxInFlightPerTask; private Task( String name, @@ -30,7 +34,11 @@ private Task( RetryPolicy retryPolicy, List codecs, boolean idempotent, - CircuitBreakerConfig circuitBreaker) { + CircuitBreakerConfig circuitBreaker, + String rateLimit, + String retryBudget, + Integer maxConcurrent, + Integer maxInFlightPerTask) { this.name = Objects.requireNonNull(name, "task name must not be null"); if (name.trim().isEmpty()) { throw new IllegalArgumentException("task name must not be blank"); @@ -41,21 +49,48 @@ private Task( this.codecs = List.copyOf(codecs); this.idempotent = idempotent; this.circuitBreaker = circuitBreaker; + this.rateLimit = rateLimit; + this.retryBudget = retryBudget; + this.maxConcurrent = maxConcurrent; + this.maxInFlightPerTask = maxInFlightPerTask; } /** A task whose payload deserializes to {@code payloadType}. */ public static Task of(String name, Class payloadType) { - return new Task<>(name, payloadType, EnqueueOptions.none(), null, List.of(), false, null); + return new Task<>( + name, payloadType, EnqueueOptions.none(), null, List.of(), false, null, null, null, null, null); } /** A task whose payload deserializes to a generic type, e.g. {@code new TypeReference>(){}}. */ public static Task of(String name, TypeReference payloadType) { - return new Task<>(name, payloadType.getType(), EnqueueOptions.none(), null, List.of(), false, null); + return new Task<>( + name, + payloadType.getType(), + EnqueueOptions.none(), + null, + List.of(), + false, + null, + null, + null, + null, + null); } /** A copy of this task with the given default options. */ public Task withOptions(EnqueueOptions options) { - return new Task<>(name, payloadType, options, retryPolicy, codecs, idempotent, circuitBreaker); + return new Task<>( + name, + payloadType, + options, + retryPolicy, + codecs, + idempotent, + circuitBreaker, + rateLimit, + retryBudget, + maxConcurrent, + maxInFlightPerTask); } /** @@ -64,7 +99,18 @@ public Task withOptions(EnqueueOptions options) { * from {@link #maxRetries}. */ public Task retryPolicy(RetryPolicy retryPolicy) { - return new Task<>(name, payloadType, options, retryPolicy, codecs, idempotent, circuitBreaker); + return new Task<>( + name, + payloadType, + options, + retryPolicy, + codecs, + idempotent, + circuitBreaker, + rateLimit, + retryBudget, + maxConcurrent, + maxInFlightPerTask); } /** @@ -74,7 +120,18 @@ public Task retryPolicy(RetryPolicy retryPolicy) { * {@code Taskito.builder().codec(name, codec)} on producers and workers. */ public Task codecs(String... codecs) { - return new Task<>(name, payloadType, options, retryPolicy, Arrays.asList(codecs), idempotent, circuitBreaker); + return new Task<>( + name, + payloadType, + options, + retryPolicy, + Arrays.asList(codecs), + idempotent, + circuitBreaker, + rateLimit, + retryBudget, + maxConcurrent, + maxInFlightPerTask); } /** @@ -83,7 +140,18 @@ public Task codecs(String... codecs) { * A per-enqueue {@link EnqueueOptions.Builder#idempotent(boolean)} overrides this default. */ public Task idempotent(boolean idempotent) { - return new Task<>(name, payloadType, options, retryPolicy, codecs, idempotent, circuitBreaker); + return new Task<>( + name, + payloadType, + options, + retryPolicy, + codecs, + idempotent, + circuitBreaker, + rateLimit, + retryBudget, + maxConcurrent, + maxInFlightPerTask); } /** @@ -92,7 +160,111 @@ public Task idempotent(boolean idempotent) { * it recovers. */ public Task circuitBreaker(CircuitBreakerConfig circuitBreaker) { - return new Task<>(name, payloadType, options, retryPolicy, codecs, idempotent, circuitBreaker); + return new Task<>( + name, + payloadType, + options, + retryPolicy, + codecs, + idempotent, + circuitBreaker, + rateLimit, + retryBudget, + maxConcurrent, + maxInFlightPerTask); + } + + /** + * A copy of this task throttled to {@code rateLimit}, a spec like {@code "100/m"} + * ({@code s}, {@code m} and {@code h} suffixes). The worker registers it on + * {@code start()} and rejects a malformed spec rather than running unthrottled. + */ + public Task rateLimit(String rateLimit) { + return new Task<>( + name, + payloadType, + options, + retryPolicy, + codecs, + idempotent, + circuitBreaker, + rateLimit, + retryBudget, + maxConcurrent, + maxInFlightPerTask); + } + + /** + * A copy of this task allowed at most {@code maxConcurrent} jobs running at once + * across the cluster. The scheduler counts running jobs before dispatch, so this + * costs a database read. {@code null} or {@code 0} means no cap — matching the + * annotation's sentinel, and because a literal cap of zero would stop the task + * from ever dispatching. + */ + public Task maxConcurrent(Integer maxConcurrent) { + maxConcurrent = uncappedIfZero(maxConcurrent); + return new Task<>( + name, + payloadType, + options, + retryPolicy, + codecs, + idempotent, + circuitBreaker, + rateLimit, + retryBudget, + maxConcurrent, + maxInFlightPerTask); + } + + /** + * A copy of this task whose retries are capped at {@code retryBudget}, + * a spec like {@code "100/m"} — across all of its jobs, not per job. Once spent, + * failures dead-letter instead of retrying, so a broken dependency cannot become + * a retry storm. Distinct from {@link #maxRetries}, which bounds one job rather + * than the rate. + */ + public Task retryBudget(String retryBudget) { + return new Task<>( + name, + payloadType, + options, + retryPolicy, + codecs, + idempotent, + circuitBreaker, + rateLimit, + retryBudget, + maxConcurrent, + maxInFlightPerTask); + } + + /** + * A copy of this task allowed at most {@code maxInFlightPerTask} of one worker's + * dispatch slots, so a slow task cannot occupy the whole pool and starve the + * others. In-process and free, unlike {@link #maxConcurrent}, which is + * cluster-wide and costs a database read. {@code null} or {@code 0} lets it use + * the whole pool, matching the annotation's sentinel. + */ + public Task maxInFlightPerTask(Integer maxInFlightPerTask) { + maxInFlightPerTask = uncappedIfZero(maxInFlightPerTask); + return new Task<>( + name, + payloadType, + options, + retryPolicy, + codecs, + idempotent, + circuitBreaker, + rateLimit, + retryBudget, + maxConcurrent, + maxInFlightPerTask); + } + + /** A cap of zero is the annotation's "unset" sentinel, never a literal zero. */ + private static Integer uncappedIfZero(Integer cap) { + return cap != null && cap == 0 ? null : cap; } public Task queue(String queue) { @@ -160,4 +332,24 @@ public boolean idempotent() { public CircuitBreakerConfig circuitBreaker() { return circuitBreaker; } + + /** This task's rate-limit spec (e.g. {@code "100/m"}), or {@code null} when unthrottled. */ + public String rateLimit() { + return rateLimit; + } + + /** This task's retry-rate cap (e.g. {@code "100/m"}), or {@code null} when uncapped. */ + public String retryBudget() { + return retryBudget; + } + + /** Cap on this task's concurrently-running jobs, or {@code null} when uncapped. */ + public Integer maxConcurrent() { + return maxConcurrent; + } + + /** Cap on this task's share of one worker's dispatch slots, or {@code null} when uncapped. */ + public Integer maxInFlightPerTask() { + return maxInFlightPerTask; + } } diff --git a/sdks/java/src/main/java/org/byteveda/taskito/worker/Worker.java b/sdks/java/src/main/java/org/byteveda/taskito/worker/Worker.java index 26da11b9..771757ed 100644 --- a/sdks/java/src/main/java/org/byteveda/taskito/worker/Worker.java +++ b/sdks/java/src/main/java/org/byteveda/taskito/worker/Worker.java @@ -8,11 +8,9 @@ import java.util.EnumMap; import java.util.HashMap; import java.util.LinkedHashMap; -import java.util.LinkedHashSet; import java.util.List; import java.util.Map; import java.util.Optional; -import java.util.Set; import java.util.concurrent.CountDownLatch; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; @@ -182,8 +180,14 @@ public static final class Builder { private final ResourceRuntime resources; private final Map codecs; private final Map handlers = new HashMap<>(); - private final Map taskPolicies = new HashMap<>(); - private final Map taskCircuitBreakers = new HashMap<>(); + /** + * Tasks carrying policy the worker must register, keyed by name. Holding the + * task itself rather than a map per knob keeps {@link #encodeTaskConfigs} + * reading from one source: a new knob is then encoded in one place instead of + * needing its own map, its own capture, and its own arm of a name union. + */ + private final Map> taskPolicies = new LinkedHashMap<>(); + private final Map>> listeners = new EnumMap<>(EventName.class); private List subscriptions = List.of(); private List queues; @@ -236,18 +240,23 @@ public Builder register(Handler handler) { return this; } - /** Remember a task's retry-backoff curve and circuit breaker so {@code start()} registers them. */ + /** Remember a task's policy — retries, breaker, throttling, caps — so {@code start()} registers it. */ private void capturePolicy(Task task) { - RetryPolicy policy = task.retryPolicy(); - if (policy != null) { - taskPolicies.put(task.name(), policy); - } - CircuitBreakerConfig breaker = task.circuitBreaker(); - if (breaker != null) { - taskCircuitBreakers.put(task.name(), breaker); + if (hasPolicy(task)) { + taskPolicies.put(task.name(), task); } } + /** Whether a task sets anything the scheduler needs registering. */ + private static boolean hasPolicy(Task task) { + return task.retryPolicy() != null + || task.circuitBreaker() != null + || task.rateLimit() != null + || task.retryBudget() != null + || task.maxConcurrent() != null + || task.maxInFlightPerTask() != null; + } + /** Register every handler in a {@link HandlerRegistry} (e.g. a generated {@code XxxTasks.handlers}). */ public Builder register(HandlerRegistry registry) { registry.handlers().forEach(this::register); @@ -432,7 +441,7 @@ private String encodeOptions() { } else if (concurrency > 0) { options.put("concurrency", concurrency); } - if (!taskPolicies.isEmpty() || !taskCircuitBreakers.isEmpty()) { + if (!taskPolicies.isEmpty()) { options.put("taskConfigs", encodeTaskConfigs()); } if (mesh != null) { @@ -481,40 +490,57 @@ private List> encodeSubscriptions() { * wire shape the binding reads — one entry per task, merging both sources by name. */ private List> encodeTaskConfigs() { - Set names = new LinkedHashSet<>(taskPolicies.keySet()); - names.addAll(taskCircuitBreakers.keySet()); - List> configs = new ArrayList<>(names.size()); - for (String name : names) { + List> configs = new ArrayList<>(taskPolicies.size()); + for (Task task : taskPolicies.values()) { Map config = new LinkedHashMap<>(); - config.put("name", name); - RetryPolicy policy = taskPolicies.get(name); - if (policy != null) { - if (policy.baseDelay() != null) { - config.put("baseDelayMs", policy.baseDelay().toMillis()); - } - if (policy.maxDelay() != null) { - config.put("maxDelayMs", policy.maxDelay().toMillis()); - } - if (!policy.customDelays().isEmpty()) { - List delaysMs = - new ArrayList<>(policy.customDelays().size()); - policy.customDelays().forEach(delay -> delaysMs.add(delay.toMillis())); - config.put("customDelaysMs", delaysMs); - } + config.put("name", task.name()); + encodeRetryPolicy(task.retryPolicy(), config); + encodeCircuitBreaker(task.circuitBreaker(), config); + if (task.rateLimit() != null) { + config.put("rateLimit", task.rateLimit()); + } + if (task.retryBudget() != null) { + config.put("retryBudget", task.retryBudget()); } - CircuitBreakerConfig breaker = taskCircuitBreakers.get(name); - if (breaker != null) { - config.put("circuitBreakerThreshold", breaker.threshold()); - config.put("circuitBreakerWindowMs", breaker.window().toMillis()); - config.put("circuitBreakerCooldownMs", breaker.cooldown().toMillis()); - config.put("circuitBreakerHalfOpenProbes", breaker.halfOpenProbes()); - config.put("circuitBreakerHalfOpenSuccessRate", breaker.halfOpenSuccessRate()); + if (task.maxConcurrent() != null) { + config.put("maxConcurrent", task.maxConcurrent()); + } + if (task.maxInFlightPerTask() != null) { + config.put("maxInFlightPerTask", task.maxInFlightPerTask()); } configs.add(config); } return configs; } + private static void encodeRetryPolicy(RetryPolicy policy, Map config) { + if (policy == null) { + return; + } + if (policy.baseDelay() != null) { + config.put("baseDelayMs", policy.baseDelay().toMillis()); + } + if (policy.maxDelay() != null) { + config.put("maxDelayMs", policy.maxDelay().toMillis()); + } + if (!policy.customDelays().isEmpty()) { + List delaysMs = new ArrayList<>(policy.customDelays().size()); + policy.customDelays().forEach(delay -> delaysMs.add(delay.toMillis())); + config.put("customDelaysMs", delaysMs); + } + } + + private static void encodeCircuitBreaker(CircuitBreakerConfig breaker, Map config) { + if (breaker == null) { + return; + } + config.put("circuitBreakerThreshold", breaker.threshold()); + config.put("circuitBreakerWindowMs", breaker.window().toMillis()); + config.put("circuitBreakerCooldownMs", breaker.cooldown().toMillis()); + config.put("circuitBreakerHalfOpenProbes", breaker.halfOpenProbes()); + config.put("circuitBreakerHalfOpenSuccessRate", breaker.halfOpenSuccessRate()); + } + @SuppressWarnings("unchecked") private static TaskFunction cast(TaskFunction handler) { return (TaskFunction) handler; diff --git a/sdks/java/src/test/java/org/byteveda/taskito/core/ResultBatchingTest.java b/sdks/java/src/test/java/org/byteveda/taskito/core/ResultBatchingTest.java new file mode 100644 index 00000000..eadcfb0a --- /dev/null +++ b/sdks/java/src/test/java/org/byteveda/taskito/core/ResultBatchingTest.java @@ -0,0 +1,76 @@ +package org.byteveda.taskito.core; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import java.nio.file.Path; +import java.util.List; +import java.util.Set; +import java.util.concurrent.ConcurrentLinkedQueue; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; +import java.util.stream.Collectors; +import org.byteveda.taskito.Taskito; +import org.byteveda.taskito.events.EventName; +import org.byteveda.taskito.task.Task; +import org.byteveda.taskito.worker.Worker; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.Timeout; +import org.junit.jupiter.api.io.TempDir; + +/** Results finalized as a batch still surface exactly one outcome per job. */ +class ResultBatchingTest { + + @Test + @Timeout(30) + void mixedBatchReportsEveryJobOnce(@TempDir Path dir) throws Exception { + // The drain loop finalizes everything already queued in one transaction. + // Successes batch into a single write while failures stay on the + // per-result path, so a mixed burst exercises both halves: every job must + // surface exactly one outcome, none lost and none double-counted. + int each = 6; + Task ok = Task.of("ok", String.class); + Task boom = Task.of("boom", String.class).maxRetries(0); + + ConcurrentLinkedQueue succeeded = new ConcurrentLinkedQueue<>(); + ConcurrentLinkedQueue dead = new ConcurrentLinkedQueue<>(); + CountDownLatch finished = new CountDownLatch(each * 2); + + try (Taskito queue = Taskito.builder() + .backend("sqlite") + .url(dir.resolve("t.db").toString()) + .open()) { + for (int i = 0; i < each; i++) { + queue.enqueue(ok, String.valueOf(i)); + queue.enqueue(boom, String.valueOf(i)); + } + + try (Worker worker = queue.worker() + .concurrency(8) + .batchSize(each * 2) + .handle(ok, (String p) -> p) + .handle(boom, (String p) -> { + throw new IllegalStateException("expected"); + }) + .on(EventName.SUCCESS, event -> { + succeeded.add(event.jobId); + finished.countDown(); + }) + .on(EventName.DEAD, event -> { + dead.add(event.jobId); + finished.countDown(); + }) + .start()) { + assertTrue(finished.await(20, TimeUnit.SECONDS), "every job should report an outcome"); + } + } + + List succeededIds = List.copyOf(succeeded); + List deadIds = List.copyOf(dead); + assertEquals(each, succeededIds.size(), "one success outcome per successful job"); + assertEquals(each, deadIds.size(), "one dead outcome per failing job"); + assertEquals(each, succeededIds.stream().collect(Collectors.toSet()).size(), "no job may be reported twice"); + Set overlap = succeededIds.stream().filter(deadIds::contains).collect(Collectors.toSet()); + assertTrue(overlap.isEmpty(), "a job cannot both succeed and dead-letter: " + overlap); + } +} diff --git a/sdks/java/src/test/java/org/byteveda/taskito/core/TaskPolicyConfigTest.java b/sdks/java/src/test/java/org/byteveda/taskito/core/TaskPolicyConfigTest.java new file mode 100644 index 00000000..a5ccf067 --- /dev/null +++ b/sdks/java/src/test/java/org/byteveda/taskito/core/TaskPolicyConfigTest.java @@ -0,0 +1,187 @@ +package org.byteveda.taskito.core; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import java.nio.file.Path; +import java.time.Duration; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; +import org.byteveda.taskito.Taskito; +import org.byteveda.taskito.events.EventName; +import org.byteveda.taskito.task.Task; +import org.byteveda.taskito.worker.Worker; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.Timeout; +import org.junit.jupiter.api.io.TempDir; + +/** Per-task throttling and concurrency caps reach the scheduler. */ +class TaskPolicyConfigTest { + + /** A cap of one means jobs of this task never overlap, however many the worker runs. */ + @Test + @Timeout(30) + void maxConcurrentSerializesOneTask(@TempDir Path dir) throws Exception { + Task solo = Task.of("solo", String.class).maxConcurrent(1); + AtomicInteger running = new AtomicInteger(); + AtomicInteger peak = new AtomicInteger(); + CountDownLatch done = new CountDownLatch(4); + + try (Taskito queue = Taskito.builder() + .backend("sqlite") + .url(dir.resolve("t.db").toString()) + .open()) { + for (int i = 0; i < 4; i++) { + queue.enqueue(solo, String.valueOf(i)); + } + + try (Worker worker = queue.worker() + .concurrency(4) + .handle(solo, (String p) -> { + peak.accumulateAndGet(running.incrementAndGet(), Math::max); + try { + Thread.sleep(150); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } + running.decrementAndGet(); + return null; + }) + .on(EventName.SUCCESS, event -> done.countDown()) + .start()) { + assertTrue(done.await(20, TimeUnit.SECONDS), "all four jobs should finish"); + } + } + + assertEquals(1, peak.get(), "maxConcurrent=1 must never let two jobs of the task overlap"); + } + + /** The in-process bulkhead holds a task to its share of one worker's slots. */ + @Test + @Timeout(30) + void maxInFlightPerTaskBoundsOneTask(@TempDir Path dir) throws Exception { + // Same visible effect as maxConcurrent=1 here, but enforced in-process + // from the worker's in-flight map rather than a running-count query. + Task solo = Task.of("solo", String.class).maxInFlightPerTask(1); + AtomicInteger running = new AtomicInteger(); + AtomicInteger peak = new AtomicInteger(); + CountDownLatch done = new CountDownLatch(4); + + try (Taskito queue = Taskito.builder() + .backend("sqlite") + .url(dir.resolve("t.db").toString()) + .open()) { + for (int i = 0; i < 4; i++) { + queue.enqueue(solo, String.valueOf(i)); + } + + try (Worker worker = queue.worker() + .concurrency(4) + .batchSize(4) + .handle(solo, (String p) -> { + peak.accumulateAndGet(running.incrementAndGet(), Math::max); + try { + Thread.sleep(150); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } + running.decrementAndGet(); + return null; + }) + .on(EventName.SUCCESS, event -> done.countDown()) + .start()) { + assertTrue(done.await(20, TimeUnit.SECONDS), "all four jobs should finish"); + } + } + + assertEquals(1, peak.get(), "maxInFlightPerTask=1 must never let two jobs of the task overlap"); + } + + /** Throttling drags a burst out past the window it would otherwise finish in. */ + @Test + @Timeout(30) + void rateLimitThrottlesDispatch(@TempDir Path dir) throws Exception { + // The bucket starts full, so the spec's count is also its burst: "60/m" + // would let 60 jobs through at once and prove nothing. "1/s" holds one + // token, so of three jobs the first runs at once and the rest wait for a + // refill — roughly a second each. + Task limited = Task.of("limited", String.class).rateLimit("1/s"); + CountDownLatch done = new CountDownLatch(3); + + try (Taskito queue = Taskito.builder() + .backend("sqlite") + .url(dir.resolve("t.db").toString()) + .open()) { + for (int i = 0; i < 3; i++) { + queue.enqueue(limited, String.valueOf(i)); + } + + long start = System.nanoTime(); + try (Worker worker = queue.worker() + .handle(limited, (String p) -> null) + .on(EventName.SUCCESS, event -> done.countDown()) + .start()) { + assertTrue(done.await(20, TimeUnit.SECONDS), "all three jobs should finish"); + } + Duration elapsed = Duration.ofNanos(System.nanoTime() - start); + assertTrue( + elapsed.toMillis() >= 1200, + "a throttled burst must not finish instantly, took " + elapsed.toMillis() + "ms"); + } + } + + /** The budget caps retries across jobs, so a storm dead-letters rather than retrying forever. */ + @Test + @Timeout(30) + void retryBudgetDeadLettersOnceSpent(@TempDir Path dir) throws Exception { + // maxRetries(5) would allow plenty of retries per job; one budget token + // across all of them means the rest dead-letter instead. + Task flaky = Task.of("flaky", String.class).retryBudget("1/m").maxRetries(5); + CountDownLatch dead = new CountDownLatch(1); + + try (Taskito queue = Taskito.builder() + .backend("sqlite") + .url(dir.resolve("t.db").toString()) + .open()) { + for (int i = 0; i < 3; i++) { + queue.enqueue(flaky, String.valueOf(i)); + } + + try (Worker worker = queue.worker() + .concurrency(4) + .handle(flaky, (String p) -> { + throw new IllegalStateException("dependency down"); + }) + .on(EventName.DEAD, event -> dead.countDown()) + .start()) { + assertTrue(dead.await(20, TimeUnit.SECONDS), "an over-budget job should dead-letter"); + } + + long budgetKilled = queue.listDead(10, 0).stream() + .filter(entry -> "retry_budget_exhausted".equals(entry.metadata)) + .count(); + assertTrue(budgetKilled > 0, "a job should be dead-lettered by the budget, not by retry exhaustion"); + } + } + + /** A typo must fail the start, not silently run the task unthrottled. */ + @Test + @Timeout(30) + void malformedRateLimitFailsTheStart(@TempDir Path dir) { + Task bad = Task.of("bad", String.class).rateLimit("100/mm"); + + try (Taskito queue = Taskito.builder() + .backend("sqlite") + .url(dir.resolve("t.db").toString()) + .open()) { + RuntimeException error = assertThrows( + RuntimeException.class, + () -> queue.worker().handle(bad, (String p) -> null).start()); + assertTrue( + error.getMessage().contains("rateLimit"), + "the error should name the offending option, got: " + error.getMessage()); + } + } +} diff --git a/sdks/node/src/types.ts b/sdks/node/src/types.ts index eb3e2cdc..80db1fb6 100644 --- a/sdks/node/src/types.ts +++ b/sdks/node/src/types.ts @@ -98,10 +98,26 @@ export interface TaskOptions { retryBackoff?: { baseMs?: number; maxMs?: number }; /** Per-job timeout default (ms); enforced by the worker. */ timeoutMs?: number; - /** Cap on concurrently-running jobs of this task. */ + /** Cap on concurrently-running jobs of this task, across the cluster. */ maxConcurrent?: number; + /** + * Cap on this task's share of a single worker's dispatch slots, so one slow + * task cannot occupy the whole pool and starve the others. In-process and + * free, unlike {@link TaskOptions.maxConcurrent}, which is cluster-wide and + * costs a database read. + */ + maxInFlightPerTask?: number; /** Rate-limit spec like `"100/m"`, `"50/s"`, `"3600/h"`. */ rateLimit?: RateLimit; + /** + * Cap on how fast this task may **retry**, across all of its jobs — same spec + * as {@link TaskOptions.rateLimit}. Once spent, failures dead-letter instead + * of retrying, so a broken dependency cannot become a retry storm. Distinct + * from {@link TaskOptions.maxRetries}, which bounds one job rather than the + * rate, and from {@link TaskOptions.circuitBreaker}, which trips on hard + * failure rather than aggregate retry rate. + */ + retryBudget?: RateLimit; /** Trip the task's circuit breaker after repeated failures. */ circuitBreaker?: CircuitBreakerInput; /** @@ -160,8 +176,18 @@ export interface RegisteredTask { export interface WorkerRunOptions { /** Queues to consume (default `["default"]`). */ queues?: string[]; - /** In-flight channel capacity (default 128). */ + /** + * Buffer size for job hand-offs, and the capacity this worker advertises + * (default 128). This does **not** bound how many jobs run at once — see + * {@link WorkerRunOptions.concurrency}. + */ channelCapacity?: number; + /** + * Jobs this worker runs at once. Unset leaves it unbounded, so the worker can + * claim more than it can run, stranding the surplus `running` while peers + * sharing the database skip it. Set this on any worker sharing a database. + */ + concurrency?: number; /** Jobs claimed per scheduler poll (default 1). */ batchSize?: number; /** diff --git a/sdks/node/src/worker.ts b/sdks/node/src/worker.ts index 7b78c165..586882dd 100644 --- a/sdks/node/src/worker.ts +++ b/sdks/node/src/worker.ts @@ -21,7 +21,7 @@ import type { import { type ResourceRuntime, runWithResolver } from "./resources"; import { deserializeCall, type PayloadCodec, type Serializer } from "./serializers"; import { encodeTaskError } from "./task-error"; -import type { QueueLimits, RegisteredTask, WorkerRunOptions } from "./types"; +import type { QueueLimits, RegisteredTask, TaskOptions, WorkerRunOptions } from "./types"; import { createLogger } from "./utils"; import type { WorkflowTracker } from "./workflows"; import { CACHE_TASK } from "./workflows/cache"; @@ -225,6 +225,7 @@ export class Worker { const nativeOptions: NativeWorkerOptions = { queues: run?.queues, channelCapacity: run?.channelCapacity, + concurrency: run?.concurrency, batchSize: run?.batchSize, taskConfigs: applyTaskOverrides( buildTaskConfigs(tasks), @@ -312,30 +313,43 @@ export class Worker { function buildTaskConfigs(tasks: ReadonlyMap): TaskConfigInput[] { const configs: TaskConfigInput[] = []; for (const [name, task] of tasks) { - const options = task.options; - if ( - !options || - (options.maxRetries === undefined && - options.retryBackoff === undefined && - options.maxConcurrent === undefined && - options.rateLimit === undefined && - options.circuitBreaker === undefined) - ) { + if (!task.options) { continue; } - configs.push({ - name, - maxRetries: options.maxRetries, - retryBaseDelayMs: options.retryBackoff?.baseMs, - retryMaxDelayMs: options.retryBackoff?.maxMs, - maxConcurrent: options.maxConcurrent, - rateLimit: options.rateLimit, - circuitBreaker: options.circuitBreaker, - }); + const config = toTaskConfig(name, task.options); + if (setsSomething(config)) { + configs.push(config); + } } return configs; } +function toTaskConfig(name: string, options: TaskOptions): TaskConfigInput { + return { + name, + maxRetries: options.maxRetries, + retryBaseDelayMs: options.retryBackoff?.baseMs, + retryMaxDelayMs: options.retryBackoff?.maxMs, + maxConcurrent: options.maxConcurrent, + maxInFlightPerTask: options.maxInFlightPerTask, + rateLimit: options.rateLimit, + retryBudget: options.retryBudget, + circuitBreaker: options.circuitBreaker, + }; +} + +/** + * Whether a task set any policy worth registering. + * + * Derived from the built config rather than a hand-listed set of option names: + * a list silently drops any option missing from it, so a task setting only the + * new option would never reach the scheduler — with no error, and invisible to + * type-checking. `name` is always present, so it can't stand in for a setting. + */ +function setsSomething({ name: _name, ...policy }: TaskConfigInput): boolean { + return Object.values(policy).some((value) => value !== undefined); +} + function buildQueueConfigs(limits: ReadonlyMap): QueueConfigInput[] { return [...limits].map(([name, limit]) => ({ name, diff --git a/sdks/node/test/worker/backpressure.test.ts b/sdks/node/test/worker/backpressure.test.ts new file mode 100644 index 00000000..f198e9dc --- /dev/null +++ b/sdks/node/test/worker/backpressure.test.ts @@ -0,0 +1,65 @@ +import { mkdtempSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { afterEach, expect, it } from "vitest"; +import { Queue, type Worker } from "../../src/index"; + +let worker: Worker | undefined; + +afterEach(() => { + worker?.stop(); + worker = undefined; +}); + +function newQueue(): Queue { + return new Queue({ dbPath: join(mkdtempSync(join(tmpdir(), "taskito-bp-")), "q.db") }); +} + +async function waitFor(predicate: () => boolean, timeoutMs = 14000): Promise { + const deadline = Date.now() + timeoutMs; + while (Date.now() < deadline) { + if (predicate()) { + return true; + } + await new Promise((resolve) => setTimeout(resolve, 20)); + } + return false; +} + +it("never runs more jobs at once than concurrency", async () => { + // The dispatcher used to spawn every job it was handed and immediately take + // the next, so the job channel drained as fast as the scheduler filled it and + // nothing bounded in-flight work — the worker claimed jobs it could not run, + // stranding them Running and starving peers on the same database. + // + // batchSize is the whole backlog on purpose. Several bounds overlap here — + // poll cadence (~50ms), task duration, channel capacity, the in-flight cap — + // and the test only proves something if the one under test is the tightest. + // Claiming all 12 in a single poll removes cadence from the picture: unfixed, + // all 12 dispatch at once and peak is 12; fixed, the in-flight clamp holds the + // claim to 3 regardless of how fast the poller runs. + const queue = newQueue(); + const total = 12; + const concurrency = 3; + let running = 0; + let peak = 0; + let done = 0; + + queue.task("slow", async () => { + running += 1; + peak = Math.max(peak, running); + await new Promise((resolve) => setTimeout(resolve, 400)); + running -= 1; + done += 1; + }); + + for (let i = 0; i < total; i += 1) { + await queue.enqueue("slow"); + } + + worker = queue.runWorker({ queues: ["default"], concurrency, batchSize: total }); + + expect(await waitFor(() => done >= total)).toBe(true); + expect(peak).toBeLessThanOrEqual(concurrency); + expect(peak).toBeGreaterThan(1); +}); diff --git a/sdks/node/test/worker/rateScope.test.ts b/sdks/node/test/worker/rateScope.test.ts new file mode 100644 index 00000000..912ae464 --- /dev/null +++ b/sdks/node/test/worker/rateScope.test.ts @@ -0,0 +1,42 @@ +import { mkdtempSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { expect, it } from "vitest"; +import { Queue } from "../../src/index"; + +function newQueue(): Queue { + return new Queue({ dbPath: join(mkdtempSync(join(tmpdir(), "taskito-sc-")), "q.db") }); +} + +/** + * The message a failed `runWorker` threw. + * + * `toThrow(string)` only checks containment, and the contract here is that + * exactly one scope is named — a message mentioning both would satisfy a + * substring match. So compare the whole message instead. + */ +function startError(queue: Queue): string { + try { + queue.runWorker({ queues: ["default"] }); + } catch (error) { + return (error as Error).message; + } + throw new Error("runWorker should have rejected the malformed rate limit"); +} + +// Tasks and queues configure a rate limit through one parser, so the error has +// to name which kind is wrong. +it("names the task, and only the task, for a malformed task rate limit", () => { + const queue = newQueue(); + queue.task("t", () => {}, { rateLimit: "100/mm" as never }); + expect(startError(queue)).toBe("invalid rateLimit '100/mm' on task 't' (expected e.g. '100/m')"); +}); + +it("names the queue, and only the queue, for a malformed queue rate limit", () => { + const queue = newQueue(); + queue.task("t", () => {}); + queue.configureQueue("default", { rateLimit: "100/mm" as never }); + expect(startError(queue)).toBe( + "invalid rateLimit '100/mm' on queue 'default' (expected e.g. '100/m')", + ); +}); diff --git a/sdks/node/test/worker/resultBatching.test.ts b/sdks/node/test/worker/resultBatching.test.ts new file mode 100644 index 00000000..bc509b00 --- /dev/null +++ b/sdks/node/test/worker/resultBatching.test.ts @@ -0,0 +1,68 @@ +import { mkdtempSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { afterEach, expect, it } from "vitest"; +import { Queue, type Worker } from "../../src/index"; + +let worker: Worker | undefined; + +afterEach(() => { + worker?.stop(); + worker = undefined; +}); + +function newQueue(): Queue { + return new Queue({ dbPath: join(mkdtempSync(join(tmpdir(), "taskito-rb-")), "q.db") }); +} + +async function waitFor(predicate: () => boolean, timeoutMs = 14000): Promise { + const deadline = Date.now() + timeoutMs; + while (Date.now() < deadline) { + if (predicate()) { + return true; + } + await new Promise((resolve) => setTimeout(resolve, 20)); + } + return false; +} + +it("reports every job exactly once when results are drained as a batch", async () => { + // The drain loop finalizes everything already queued in one transaction, so a + // burst of results is handled as a batch rather than one per wake. Successes + // batch into a single write while failures stay on the per-result path, so a + // mixed burst exercises both halves: each job must still surface exactly one + // outcome, none lost and none double-counted. + const queue = newQueue(); + const each = 6; + const completed: string[] = []; + const dead: string[] = []; + + queue.on("job.completed", (event) => completed.push(event.jobId)); + queue.on("job.dead", (event) => dead.push(event.jobId)); + + queue.task("ok", (n: number) => n); + queue.task( + "boom", + () => { + throw new Error("expected"); + }, + { maxRetries: 0 }, + ); + + for (let i = 0; i < each; i += 1) { + await queue.enqueue("ok", [i]); + await queue.enqueue("boom"); + } + + worker = queue.runWorker({ queues: ["default"], concurrency: 8, batchSize: each * 2 }); + + expect(await waitFor(() => completed.length >= each && dead.length >= each)).toBe(true); + + expect(completed).toHaveLength(each); + expect(dead).toHaveLength(each); + expect(new Set(completed).size).toBe(each); + expect(new Set(dead).size).toBe(each); + // Uniqueness within each list would still miss a job reported once as both. + const overlap = completed.filter((id) => dead.includes(id)); + expect(overlap).toEqual([]); +}); diff --git a/sdks/node/test/worker/taskConfig.test.ts b/sdks/node/test/worker/taskConfig.test.ts new file mode 100644 index 00000000..42a53feb --- /dev/null +++ b/sdks/node/test/worker/taskConfig.test.ts @@ -0,0 +1,110 @@ +import { mkdtempSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { afterEach, expect, it } from "vitest"; +import { Queue, type Worker } from "../../src/index"; + +let worker: Worker | undefined; + +afterEach(() => { + worker?.stop(); + worker = undefined; +}); + +function newQueue(): Queue { + return new Queue({ dbPath: join(mkdtempSync(join(tmpdir(), "taskito-tc-")), "q.db") }); +} + +/** Poll an async predicate. `waitFor` takes a sync one — a promise is always truthy there. */ +async function waitForAsync( + predicate: () => Promise, + timeoutMs = 10000, +): Promise { + const deadline = Date.now() + timeoutMs; + while (Date.now() < deadline) { + if (await predicate()) { + return true; + } + await new Promise((resolve) => setTimeout(resolve, 20)); + } + return false; +} + +async function waitFor(predicate: () => boolean, timeoutMs = 10000): Promise { + const deadline = Date.now() + timeoutMs; + while (Date.now() < deadline) { + if (predicate()) { + return true; + } + await new Promise((resolve) => setTimeout(resolve, 20)); + } + return false; +} + +it("registers a task that sets only maxInFlightPerTask", async () => { + // The config collector used to gate on a hand-listed set of option names, so + // a task setting only an option missing from that list was dropped before it + // reached the scheduler — silently, and invisibly to type-checking. A job + // still runs either way, so assert the option reaches native by giving it a + // value the scheduler must honour: one slot means strictly serial execution. + const queue = newQueue(); + const total = 4; + let running = 0; + let peak = 0; + let done = 0; + + queue.task( + "solo", + async () => { + running += 1; + peak = Math.max(peak, running); + await new Promise((resolve) => setTimeout(resolve, 150)); + running -= 1; + done += 1; + }, + { maxInFlightPerTask: 1 }, + ); + + for (let i = 0; i < total; i += 1) { + await queue.enqueue("solo"); + } + + worker = queue.runWorker({ queues: ["default"], concurrency: 4, batchSize: total }); + + expect(await waitFor(() => done >= total)).toBe(true); + expect(peak).toBe(1); +}); + +it("rejects a malformed rateLimit rather than silently disabling it", async () => { + const queue = newQueue(); + queue.task("bad", async () => {}, { rateLimit: "100/mm" as never }); + expect(() => queue.runWorker({ queues: ["default"] })).toThrow(/rateLimit/); +}); + +it("registers a task that sets only retryBudget", async () => { + // Nothing was added to the config gate for this option — it derives from the + // built config — so this asserts a task setting only retryBudget still reaches + // the scheduler. One token means the second failure dead-letters rather than + // retrying, even though maxRetries would allow more. + const queue = newQueue(); + + queue.task( + "flaky", + async () => { + throw new Error("dependency down"); + }, + { retryBudget: "1/m", maxRetries: 5 }, + ); + + for (let i = 0; i < 3; i += 1) { + await queue.enqueue("flaky"); + } + + worker = queue.runWorker({ queues: ["default"], concurrency: 4, batchSize: 3 }); + + const budgetKilled = async (): Promise => { + const dead = await queue.deadLetters(10); + return dead.filter((entry) => entry.metadata === "retry_budget_exhausted").length; + }; + expect(await waitForAsync(async () => (await budgetKilled()) > 0)).toBe(true); +});