feat(sdks): close Node and Java per-task policy and batching parity

crates/taskito-java/src/convert.rs

@@ -319,9 +319,9 @@ pub struct SubscriptionSpec {
     pub timeout_ms: Option<i64>,
 }
 
-/// 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<i64>,
     pub circuit_breaker_half_open_probes: Option<i32>,
     pub circuit_breaker_half_open_success_rate: Option<f64>,
+    /// Rate-limit spec like `"100/m"`, `"50/s"`, `"3600/h"`.
+    pub rate_limit: Option<String>,
+    /// 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<String>,
+    /// Cap on concurrently-running jobs of this task, across the cluster.
+    pub max_concurrent: Option<i32>,
+    /// 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<i32>,
 }
 
 /// Filter accepted by `NativeQueue.listJobs`.

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<Vec<TaskRetryConfig>>) {
-    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<Vec<TaskRetryConfig>>,
+) -> Result<Vec<(String, TaskConfig)>, 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<Vec<TaskRet
                         .circuit_breaker_half_open_success_rate
                         .unwrap_or(0.8),
                 });
-        scheduler.register_task(
+        let rate_limit = parse_rate_spec("rateLimit", &config.name, config.rate_limit.as_deref())?;
+        let retry_budget =
+            parse_rate_spec("retryBudget", &config.name, config.retry_budget.as_deref())?;
+        built.push((
             config.name,
             TaskConfig {
                 retry_policy,
-                rate_limit: None,
+                rate_limit,
                 circuit_breaker,
-                // Not surfaced on this SDK yet; retries stay bounded per job.
-                retry_budget: None,
-                max_concurrent: None,
-                // Not surfaced on this SDK yet; the whole pool stays available.
-                max_in_flight_per_task: None,
+                retry_budget,
+                max_concurrent: config.max_concurrent,
+                max_in_flight_per_task: config.max_in_flight_per_task.map(|n| n.max(1) as usize),
             },
-        );
+        ));
+    }
+    Ok(built)
+}
+
+/// Parse an optional rate spec, naming the offending task and option so a typo
+/// is actionable. Several options share this `"100/m"` grammar.
+fn parse_rate_spec(
+    field: &str,
+    task: &str,
+    spec: Option<&str>,
+) -> Result<Option<RateLimitConfig>, 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<taskito_core::scheduler::JobResult>,
     scheduler: Arc<Scheduler>,
     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}");
         }
     }
 }

crates/taskito-node/src/config.rs

@@ -109,8 +109,14 @@ pub struct TaskConfigInput {
     pub retry_base_delay_ms: Option<i64>,
     pub retry_max_delay_ms: Option<i64>,
     pub max_concurrent: Option<i32>,
+    /// 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<i32>,
     /// Rate-limit spec like `"100/m"`, `"50/s"`, `"3600/h"`.
     pub rate_limit: Option<String>,
+    /// 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<String>,
     pub circuit_breaker: Option<CircuitBreakerInput>,
 }
 
@@ -160,6 +166,10 @@ pub struct MeshWorkerConfig {
 pub struct WorkerOptions {
     pub queues: Option<Vec<String>>,
     pub channel_capacity: Option<u32>,
+    /// 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<u32>,
     /// Jobs claimed per scheduler poll (default 1).
     pub batch_size: Option<u32>,
     pub task_configs: Option<Vec<TaskConfigInput>>,

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<Option<RateLimitConfig>> {
+/// 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<Option<RateLimitConfig>> {
     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<TaskConfig> {
             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<CircuitBreakerC
 /// Build a [`QueueConfig`] (rate limit, concurrency cap) from JS input.
 pub fn queue_config(input: &QueueConfigInput) -> Result<QueueConfig> {
     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,
     })
 }

crates/taskito-node/src/dispatcher.rs

@@ -5,6 +5,7 @@
 //! `Promise<Buffer>`, 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<JsTaskInvocation, ErrorStrategy::Fatal>;
 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<Semaphore>,
 }
 
 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<usize>,
+    ) -> 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<JobResult>,
     ) {
         // 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