perf(scheduler): batch dequeue to claim N jobs per round-trip

crates/taskito-core/src/scheduler/mod.rs

@@ -33,6 +33,10 @@ pub struct SchedulerConfig {
     pub cleanup_interval: u32,
     /// TTL for job results in milliseconds. None means no auto-cleanup.
     pub result_ttl_ms: Option<i64>,
+    /// Maximum number of jobs claimed per dispatch round. `1` (the default)
+    /// preserves the original one-job-per-round-trip behavior; values above
+    /// `1` enable batch claiming for higher throughput.
+    pub batch_size: usize,
 }
 
 impl Default for SchedulerConfig {
@@ -44,6 +48,7 @@ impl Default for SchedulerConfig {
             periodic_check_interval: 60,
             cleanup_interval: 1200,
             result_ttl_ms: None,
+            batch_size: 1,
         }
     }
 }
@@ -225,7 +230,12 @@ impl Scheduler {
     /// Returns true if any work was done (job dispatched or periodic task enqueued),
     /// which resets the adaptive poll interval.
     fn tick(&self, job_tx: &tokio::sync::mpsc::Sender<Job>, counters: &mut TickCounters) -> bool {
-        let dispatched = match self.try_dispatch(job_tx) {
+        let dispatch_result = if self.config.batch_size > 1 {
+            self.try_dispatch_batch(job_tx)
+        } else {
+            self.try_dispatch(job_tx)
+        };
+        let dispatched = match dispatch_result {
             Ok(d) => d,
             Err(e) => {
                 error!("scheduler error: {e}");

crates/taskito-core/src/scheduler/poller.rs

@@ -43,6 +43,66 @@ impl Scheduler {
             None => return Ok(false),
         };
 
+        self.gate_and_dispatch(job, now, job_tx)
+    }
+
+    /// Batch variant of [`try_dispatch`]. Claims up to `batch_size` jobs in a
+    /// single round-trip, then runs the exact same per-job gate/claim/
+    /// concurrency/dispatch logic the single path uses. The pre-sized batch is
+    /// advisory only — the hard per-task and per-queue caps are still enforced
+    /// per job after the claim, rolling back any job that exceeds a limit.
+    pub(super) fn try_dispatch_batch(
+        &self,
+        job_tx: &tokio::sync::mpsc::Sender<Job>,
+    ) -> Result<bool> {
+        let now = now_millis();
+
+        let active_queues = self.active_queues();
+        if active_queues.is_empty() {
+            return Ok(false);
+        }
+
+        // Size the batch to the worker pool's free capacity so we never claim
+        // more jobs than the channel can immediately accept. `try_send` still
+        // guards each hand-off, but pre-sizing avoids needless claim/rollback
+        // churn. Always claim at least one.
+        let budget = self.config.batch_size.min(job_tx.capacity().max(1));
+
+        let jobs = self.storage.dequeue_batch_from(
+            &active_queues,
+            now,
+            self.namespace.as_deref(),
+            budget,
+        )?;
+        if jobs.is_empty() {
+            return Ok(false);
+        }
+
+        // Every job in `jobs` is already claimed (Running). Isolate per-job
+        // failures so one error doesn't strand the rest of the batch in
+        // Running until the reaper times them out.
+        let mut dispatched_any = false;
+        for job in jobs {
+            match self.gate_and_dispatch(job, now, job_tx) {
+                Ok(true) => dispatched_any = true,
+                Ok(false) => {}
+                Err(e) => log::warn!("batch dispatch failed for a claimed job: {e}"),
+            }
+        }
+        Ok(dispatched_any)
+    }
+
+    /// Run the post-dequeue pipeline for a single already-claimed (Running)
+    /// job: soft pre-claim gates, exactly-once claim, hard concurrency caps,
+    /// and hand-off to the worker pool. Returns `Ok(true)` if the job was
+    /// dispatched, `Ok(false)` if it was gated/rolled back. Shared by both the
+    /// single and batch dispatch paths so limit enforcement never drifts.
+    fn gate_and_dispatch(
+        &self,
+        job: Job,
+        now: i64,
+        job_tx: &tokio::sync::mpsc::Sender<Job>,
+    ) -> Result<bool> {
         // Pre-claim soft gates: rate limits and circuit breaker.
         //
         // These don't need to be atomic with the claim — if two schedulers
@@ -60,7 +120,7 @@ impl Scheduler {
 
         // Post-claim hard gate: concurrency caps must be checked AFTER the
         // claim so two schedulers cannot both pass the cap. Status was
-        // already transitioned to `Running` by `dequeue_from`, so the
+        // already transitioned to `Running` by the dequeue, so the
         // running-count includes this job — use strict `>` to allow exactly
         // `max_concurrent` jobs.
         //

crates/taskito-core/src/storage/diesel_common/jobs.rs

@@ -363,6 +363,128 @@ macro_rules! impl_diesel_job_ops {
                 Ok(None)
             }
 
+            /// Atomically claim up to `max` ready jobs from a single queue in
+            /// one transaction. Generalizes `dequeue` to the batch case: scans
+            /// a bounded candidate set and claims eligible rows until the
+            /// budget is met or candidates are exhausted.
+            ///
+            /// Each claim uses an `UPDATE ... WHERE status = Pending` guarded
+            /// by the affected-row count: if another worker already moved the
+            /// row out of `Pending`, the update affects zero rows and the
+            /// candidate is skipped — avoiding a double-claim race.
+            pub fn dequeue_batch(
+                &self,
+                queue_name: &str,
+                now: i64,
+                namespace: Option<&str>,
+                max: usize,
+            ) -> Result<Vec<Job>> {
+                if max == 0 {
+                    return Ok(Vec::new());
+                }
+
+                // Scan more candidates than `max` so dependency/expiry skips
+                // still leave enough eligible rows to fill the batch, bounded
+                // to keep the loaded set small.
+                let scan_limit = (max.saturating_mul(4)).min(400) as i64;
+
+                let mut conn = self.conn()?;
+
+                conn.transaction(|conn| {
+                    let mut query = jobs::table
+                        .filter(jobs::queue.eq(queue_name))
+                        .filter(jobs::status.eq(JobStatus::Pending as i32))
+                        .filter(jobs::scheduled_at.le(now))
+                        .order((jobs::priority.desc(), jobs::scheduled_at.asc()))
+                        .limit(scan_limit)
+                        .into_boxed();
+
+                    if let Some(ns) = namespace {
+                        query = query.filter(jobs::namespace.eq(ns));
+                    } else {
+                        query = query.filter(jobs::namespace.is_null());
+                    }
+
+                    let candidates: Vec<JobRow> = query.select(JobRow::as_select()).load(conn)?;
+
+                    let mut claimed: Vec<Job> = Vec::with_capacity(max.min(candidates.len()));
+
+                    for row in candidates {
+                        if claimed.len() == max {
+                            break;
+                        }
+
+                        // Skip expired jobs (and mark them cancelled).
+                        if let Some(expires_at) = row.expires_at {
+                            if now > expires_at {
+                                diesel::update(jobs::table)
+                                    .filter(jobs::id.eq(&row.id))
+                                    .set((
+                                        jobs::status.eq(JobStatus::Cancelled as i32),
+                                        jobs::completed_at.eq(now),
+                                        jobs::error.eq("expired before execution"),
+                                    ))
+                                    .execute(conn)?;
+                                continue;
+                            }
+                        }
+
+                        // Common case: jobs with no dependencies skip the
+                        // job_dependencies lookup entirely.
+                        if row.has_deps && !Self::deps_satisfied(conn, &row.id)? {
+                            continue;
+                        }
+
+                        // Claim guarded by the affected-row count: if another
+                        // worker already moved this row out of `Pending`, the
+                        // update touches zero rows — skip it rather than
+                        // claiming a job we don't own.
+                        let affected = diesel::update(jobs::table)
+                            .filter(jobs::id.eq(&row.id))
+                            .filter(jobs::status.eq(JobStatus::Pending as i32))
+                            .set((
+                                jobs::status.eq(JobStatus::Running as i32),
+                                jobs::started_at.eq(now),
+                            ))
+                            .execute(conn)?;
+
+                        if affected == 0 {
+                            continue;
+                        }
+
+                        let updated: JobRow = jobs::table
+                            .find(&row.id)
+                            .select(JobRow::as_select())
+                            .first(conn)?;
+
+                        claimed.push(Job::from(updated));
+                    }
+
+                    Ok(claimed)
+                })
+            }
+
+            /// Claim up to `max` ready jobs across the given queues, checking
+            /// each in order until the budget is exhausted.
+            pub fn dequeue_batch_from(
+                &self,
+                queues: &[String],
+                now: i64,
+                namespace: Option<&str>,
+                max: usize,
+            ) -> Result<Vec<Job>> {
+                let mut claimed: Vec<Job> = Vec::new();
+                for queue_name in queues {
+                    if claimed.len() >= max {
+                        break;
+                    }
+                    let remaining = max - claimed.len();
+                    let mut batch = self.dequeue_batch(queue_name, now, namespace, remaining)?;
+                    claimed.append(&mut batch);
+                }
+                Ok(claimed)
+            }
+
             /// Mark a job as complete with the given result.
             pub fn complete(&self, id: &str, result_bytes: Option<Vec<u8>>) -> Result<()> {
                 let now = now_millis();

crates/taskito-core/src/storage/mod.rs

@@ -117,6 +117,24 @@ macro_rules! impl_storage {
             ) -> $crate::error::Result<Option<$crate::job::Job>> {
                 self.dequeue_from(queues, now, namespace)
             }
+            fn dequeue_batch(
+                &self,
+                queue_name: &str,
+                now: i64,
+                namespace: Option<&str>,
+                max: usize,
+            ) -> $crate::error::Result<Vec<$crate::job::Job>> {
+                self.dequeue_batch(queue_name, now, namespace, max)
+            }
+            fn dequeue_batch_from(
+                &self,
+                queues: &[String],
+                now: i64,
+                namespace: Option<&str>,
+                max: usize,
+            ) -> $crate::error::Result<Vec<$crate::job::Job>> {
+                self.dequeue_batch_from(queues, now, namespace, max)
+            }
             fn complete(
                 &self,
                 id: &str,
@@ -581,6 +599,24 @@ impl Storage for StorageBackend {
     ) -> Result<Option<Job>> {
         delegate!(self, dequeue_from, queues, now, namespace)
     }
+    fn dequeue_batch(
+        &self,
+        queue_name: &str,
+        now: i64,
+        namespace: Option<&str>,
+        max: usize,
+    ) -> Result<Vec<Job>> {
+        delegate!(self, dequeue_batch, queue_name, now, namespace, max)
+    }
+    fn dequeue_batch_from(
+        &self,
+        queues: &[String],
+        now: i64,
+        namespace: Option<&str>,
+        max: usize,
+    ) -> Result<Vec<Job>> {
+        delegate!(self, dequeue_batch_from, queues, now, namespace, max)
+    }
     fn complete(&self, id: &str, result_bytes: Option<Vec<u8>>) -> Result<()> {
         delegate!(self, complete, id, result_bytes)
     }