queue: Add custom task queue support

[lcwangchao]

Summary

This PR adds support for using a user-provided task queue in YATP.

• Introduces queue::TaskQueue<T> and queue::PopResult<T> so a queue can report whether work is ready, delayed, or absent.
• Adds Builder::build_custom_future_pool and wires custom queues into the existing injector/local-queue plumbing.
• Extends worker parking to understand PopResult::Pending { retry_at }: a worker may sleep until the retry deadline, while a later spawn can still wake it immediately.
• Adds task::future::yield_to_scheduler() for future tasks that need to force the next poll through the scheduler queue instead of relying on the normal preemption hint.
• Drains custom queues on shutdown and adds coverage for pending deadlines, wake races, shutdown drain, metrics, custom pool construction, and forced scheduler yield.

Custom Queue API

A custom queue implements queue::TaskQueue<T>:

pub trait TaskQueue<T>: Send + Sync + 'static {
    fn push(&self, task_cell: T);

    fn pop(&self) -> PopResult<T>;

    fn drain(&self);

    fn has_ready_task(&self) -> bool;
}

The most important contract is pop():

• PopResult::Ready(pop) means a task can run immediately.
• PopResult::Pending { retry_at } means the queue contains tasks, but none can run yet. The worker will retry no later than retry_at, and can still be woken earlier by a later push.
• PopResult::Empty means the queue contains no tasks.

Important: if the queue still owns delayed or throttled tasks, it should return Pending, not Empty. Returning Empty can let workers park as if no work exists, which may delay those tasks for a long time.

Example

Suppose a custom queue delays tasks and keeps the earliest runnable task at the top of a heap. For brevity, Scheduled<T> is ordered by ready_at, and the heap stores Reverse<Scheduled<T>> so the earliest time is popped first.

struct Scheduled<T> {
    task_cell: T,
    schedule_time: Instant,
    ready_at: Instant,
}

struct DelayQueue<T> {
    tasks: Mutex<BinaryHeap<Reverse<Scheduled<T>>>>,
    delay: Duration,
}

push records when the task entered the queue and when it can first run:

fn push(&self, task_cell: T) {
    let now = Instant::now();
    self.tasks.lock().unwrap().push(Reverse(Scheduled {
        task_cell,
        schedule_time: now,
        ready_at: now + self.delay,
    }));
}

pop checks the earliest scheduled task. If it is ready, the worker can run it now. If it is not ready, the queue returns the earliest retry deadline instead of pretending to be empty.

fn pop(&self) -> PopResult<T> {
    let mut tasks = self.tasks.lock().unwrap();

    let Some(Reverse(next)) = tasks.peek() else {
        return PopResult::Empty;
    };

    if Instant::now() < next.ready_at {
        return PopResult::Pending {
            retry_at: next.ready_at,
        };
    }

    let Reverse(task) = tasks.pop().unwrap();
    PopResult::Ready(Pop {
        task_cell: task.task_cell,
        schedule_time: task.schedule_time,
        from_local: false,
    })
}

The remaining trait methods can stay simple for this queue: drain() clears the heap, and has_ready_task() checks whether tasks.peek() is already ready without removing it.

A custom future pool can then be built by passing the queue as Arc<dyn TaskQueue<future::TaskCell>>:

let queue: Arc<dyn TaskQueue<future::TaskCell>> = Arc::new(DelayQueue::new(delay));

let mut builder = Builder::new("custom-queue");
builder.max_thread_count(4);
let pool = builder.build_custom_future_pool(queue);

pool.spawn(async move {
    // This future will be inserted into DelayQueue first, then run after the delay.
});

yield_to_scheduler

YATP future tasks already have task::future::reschedule(), which gives the scheduler a chance to run other ready work. However, reschedule() still follows the normal preemption check: if the scheduler does not report ready work, the current worker may keep polling the same task for locality.

task::future::yield_to_scheduler() is stricter. It always returns the current future task to the scheduler queue before polling it again, skipping the has_ready_task() / preemption hint. This is useful when the queue itself should make the next scheduling decision, for example when a custom queue uses delayed readiness, rate limiting, or another policy that should be applied between task slices.

pool.spawn(async move {
    process_one_batch().await;

    // Force the next poll to go through the queue again, even if the worker
    // would otherwise continue polling this task immediately.
    yatp::task::future::yield_to_scheduler().await;

    process_next_batch().await;
});

Tests

• cargo fmt --check
• cargo clippy --tests -- -D clippy::all
• cargo clippy --tests --features failpoints -- -D clippy::all
• cargo test --tests
• cargo test --tests --features failpoints
• cargo test yield_to_scheduler

Summary by CodeRabbit

• New Features
  • Added support for custom task queues, including typed queue selection and building custom future pools from a provided queue.
  • Introduced a unified pop-result API (ready/pending/empty) across queue and worker handling.
• Bug Fixes
  • Improved worker scheduling to honor the earliest retry deadline and refine park/validate retry behavior.
  • Improved shutdown handling by draining local work more reliably.
  • Hardened steal-sizing logic to avoid divide-by-zero.
• Tests
  • Expanded coverage for custom queues, timing/retry behavior, shutdown boundaries, and updated tests for the new pop-result API.

[coderabbitai]

Note

Reviews paused

It looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the reviews.auto_review.auto_pause_after_reviewed_commits setting.

Use the following commands to manage reviews:

• @coderabbitai resume to resume automatic reviews.
• @coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

• ▶️ Resume reviews
• 🔍 Trigger review

📝 Walkthrough

Walkthrough

The PR adds typed custom queue selection, a new custom task-queue abstraction, and PopResult-based queue semantics. Pool builder, spawn, worker, and task runtime paths are updated to pass typed queues and retry deadlines through execution.

Changes

Custom queue and PopResult rollout

Layer / File(s)	Summary
Queue contracts src/queue.rs, src/queue/custom.rs	Adds PopResult, generic QueueType<T>, the custom TaskQueue trait, and the custom queue config and builder types.
Custom queue wiring src/queue.rs, src/queue/custom.rs, src/pool/builder.rs, src/pool/spawn.rs	Routes injector and local queue dispatch through the custom variant and updates builder APIs to use typed queue selectors for custom queues.
Worker retry and sleep src/pool/spawn.rs, src/pool/worker.rs, src/task/future.rs	Changes local popping and worker sleep to carry retry deadlines through timeout-aware parking, updates shutdown draining, and switches future reschedule control to action-based state.
Custom queue and worker tests src/queue/custom.rs, src/pool/worker.rs	Adds custom queue mock types and worker test scaffolding, then exercises pending, retry, wake, timeout, shutdown, and scheduler-yield behavior.
Compatibility updates src/queue/multilevel.rs, src/queue/priority.rs, src/task/callback.rs, src/task/future.rs	Adjusts multilevel steal sizing and updates existing queue and task tests to the PopResult API.

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~60 minutes

Possibly related PRs

• tikv/yatp#90: Also changes src/pool/spawn.rs worker parking and validation flow around retry handling.

Poem

I hopped through queues by moonlit gleam,
Where Ready, Pending, Empty stream.
The worker woke and took its turn,
While custom burrows softly churn. 🐇

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title accurately summarizes the main change: adding support for custom task queues.
Docstring Coverage	✅ Passed	Docstring coverage is 84.62% which is sufficient. The required threshold is 80.00%.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

Actionable comments posted: 1

🧹 Nitpick comments (1)

src/queue/custom.rs (1)

42-47: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick win

Clarify or guard repeated drain calls for shared custom queues.

Builder::build gives every local handle the same queue, and each handle forwards drain() directly. In multi-worker shutdown, custom implementations may see multiple/concurrent drain calls; either document that TaskQueue::drain must be idempotent/thread-safe, or guard the adapter so the shared queue is drained once.

Also applies to: 100-110, 176-179

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@src/queue/custom.rs` around lines 42 - 47, `TaskQueue::drain` can be called
multiple times or concurrently because `Builder::build` shares the same queue
across local handles and each handle forwards `drain()` directly. Update the
`TaskQueue`/adapter path to either clearly document that `TaskQueue::drain` must
be idempotent and thread-safe, or add a guard so the shared queue is drained
only once; focus on the `TaskQueue` trait and the local handle/shutdown
forwarding code that invokes `drain()`.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@src/task/future.rs`:
- Line 302: The thread-local initializer in LOCAL uses const block syntax that
requires Rust 1.79, so either add a rust-version declaration in Cargo.toml to
match that minimum toolchain or refactor the LOCAL initialization in future.rs
to avoid const { ... } if older compiler support must be preserved.

---

Nitpick comments:
In `@src/queue/custom.rs`:
- Around line 42-47: `TaskQueue::drain` can be called multiple times or
concurrently because `Builder::build` shares the same queue across local handles
and each handle forwards `drain()` directly. Update the `TaskQueue`/adapter path
to either clearly document that `TaskQueue::drain` must be idempotent and
thread-safe, or add a guard so the shared queue is drained only once; focus on
the `TaskQueue` trait and the local handle/shutdown forwarding code that invokes
`drain()`.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: abce53b9-4521-4d5f-811b-edad27e464d2

📥 Commits

Reviewing files that changed from the base of the PR and between f3acdd2 and da431e7.

📒 Files selected for processing (9)

• src/pool/builder.rs
• src/pool/spawn.rs
• src/pool/worker.rs
• src/queue.rs
• src/queue/custom.rs
• src/queue/multilevel.rs
• src/queue/priority.rs
• src/task/callback.rs
• src/task/future.rs

[lcwangchao]

/retest

[lcwangchao]

@coderabbitai review it again

[coderabbitai]

@lcwangchao I'll review the PR again with a fresh pass.

✅ Action performed

Full review finished.

[jiadebin]

Review Finding:
P1: src/pool/worker.rs:27 / src/pool/spawn.rs:348: timed Pending wakeups can bypass core_thread_count after scale-down. pop_or_sleep() only checks id > core_thread_count while validating the park. When a custom queue parks workers on Pending { retry_at }, those workers wake by timeout and pop_or_sleep() returns None; the outer worker loop then immediately starts the spin path and calls self.local.pop() without any scale-down guard. Repro: custom future queue with 4 delayed tasks, max_thread_count=4, core_thread_count=4, then scale_workers(1) before the retry deadline. At the deadline all 4 workers executed concurrently (max_in_flight=4, expected 1). This breaks the pool’s scale-down contract and is especially risky for the custom-queue/rate-limit use case this PR enables. I’d gate the spin pop for workers above the current core count, or make the timeout path re-enter the scaled-down park state instead of returning to the unguarded spin path.

[lcwangchao]

jiadebin

Thanks for the detailed report. Fixed in a3f5419 by handling the TimedOut + scaled-down worker path inside pop_or_sleep.

When a worker above the current core_thread_count wakes by a pending retry timeout, it now wakes one core worker to re-check the queue and then re-parks without carrying the old timeout, instead of returning to the outer worker loop and entering the spin-pop path.

I chose this approach to keep the existing scale-down behavior for already-running workers unchanged, while fixing the new timed-Pending wakeup path introduced by custom queues. I also added failpoint tests covering:

• a scaled-down worker re-parks after a pending timeout;
• the pending deadline is handed off to a core worker, and the scaled-down worker does not execute the delayed task.