Skip to content

export more schedule info via Extras#86

Open
zyguan wants to merge 2 commits intotikv:masterfrom
zyguan:dev/observability-1
Open

export more schedule info via Extras#86
zyguan wants to merge 2 commits intotikv:masterfrom
zyguan:dev/observability-1

Conversation

@zyguan
Copy link
Contributor

@zyguan zyguan commented Feb 10, 2026

Add two new enums TaskSource and AcquireState to represent where a task was popped from and how it was acquired by WorkerThread. Gather and expose these information through Extras to help us identify the characteristics of tasks with high scheduling delays.

Summary by CodeRabbit

Release Notes

  • New Features

    • Tasks now track their source (local queue, global queue, or other workers) and acquisition state (immediate, after spinning, or after parking).
  • Improvements

    • Enhanced shutdown behavior with comprehensive queue draining.
  • Tests

    • Extended test coverage for source and state tracking.

Signed-off-by: zyguan <zhongyangguan@gmail.com>
@coderabbitai
Copy link

coderabbitai bot commented Feb 10, 2026

Warning

Rate limit exceeded

@zyguan has exceeded the limit for the number of commits that can be reviewed per hour. Please wait 12 minutes and 50 seconds before requesting another review.

⌛ How to resolve this issue?

After the wait time has elapsed, a review can be triggered using the @coderabbitai review command as a PR comment. Alternatively, push new commits to this PR.

We recommend that you space out your commits to avoid hitting the rate limit.

🚦 How do rate limits work?

CodeRabbit enforces hourly rate limits for each developer per organization.

Our paid plans have higher rate limits than the trial, open-source and free plans. In all cases, we re-allow further reviews after a brief timeout.

Please see our FAQ for further information.

📝 Walkthrough

Walkthrough

The PR introduces task acquisition tracking by adding TaskSource and AcquireState enums to describe task origins and acquisition methods. The pop() method signature is updated to return acquisition state alongside task data, and Pop<T> replaces a boolean flag with the new TaskSource enum. Acquisition metadata is propagated through the Extras structure.

Changes

Cohort / File(s) Summary
Queue Infrastructure
src/queue/extras.rs, src/queue.rs
Introduces TaskSource (LocalQueue, GlobalQueue, OtherWorker) and AcquireState (Immediate, AfterSpin, AfterPark) enums. Extends Extras with task_source and acquire_state fields. Replaces Pop<T>.from_local: bool with Pop<T>.task_source: TaskSource.
Queue Implementations
src/queue/multilevel.rs, src/queue/priority.rs, src/queue/single_level.rs
Updates queue implementations to propagate TaskSource variants instead of boolean flags. Adjusts Pop construction across all pop paths (local, injector steal, worker steal) with appropriate TaskSource values and includes test coverage for source tracking.
Worker Loop
src/pool/worker.rs
Updates pop() method to accept retry_after_park parameter and return (Pop<T>, AcquireState) tuple. Adds state progression logic during spinning and parking. Propagates acquired state and source into task extras before handling, enabling downstream observers to track task acquisition context.

Sequence Diagram

sequenceDiagram
    participant Worker as Worker
    participant LocalQueue as Local Queue
    participant Injector as Global Injector
    participant OtherWorker as Other Worker
    participant TaskExtras as Task Extras

    Worker->>Worker: pop(retry_after_park=false)<br/>state=Immediate
    Worker->>LocalQueue: Try local pop<br/>(LocalQueue source)
    alt Local pop succeeds
        LocalQueue-->>Worker: (Pop, TaskSource::LocalQueue)
        Worker->>TaskExtras: Set task_source & acquire_state
        Note over TaskExtras: source=LocalQueue<br/>state=Immediate
    else Local pop fails & spin
        Worker->>Worker: Advance state<br/>to AfterSpin
        Worker->>Injector: Steal from global<br/>(GlobalQueue source)
        alt Global steal succeeds
            Injector-->>Worker: (Pop, TaskSource::GlobalQueue)
            Worker->>TaskExtras: Set task_source & acquire_state
            Note over TaskExtras: source=GlobalQueue<br/>state=AfterSpin
        else Exhausted, park
            Worker->>Worker: Park thread
            Worker->>Worker: Resume & set state<br/>to AfterPark
            Worker->>OtherWorker: Steal from worker<br/>(OtherWorker source)
            OtherWorker-->>Worker: (Pop, TaskSource::OtherWorker)
            Worker->>TaskExtras: Set task_source & acquire_state
            Note over TaskExtras: source=OtherWorker<br/>state=AfterPark
        end
    end
Loading

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~45 minutes

Poem

🐰 A task's journey, now clearly seen,
From local queue, or injector's gleam,
Immediate or spun, or parked then reborn,
Each acquisition marked, no context forlorn!
With extras enriched, we hop on with glee. 🌟

🚥 Pre-merge checks | ✅ 3
✅ Passed checks (3 passed)
Check name Status Explanation
Description Check ✅ Passed Check skipped - CodeRabbit’s high-level summary is enabled.
Title check ✅ Passed The title 'export more schedule info via Extras' accurately summarizes the main objective: introducing new enums (TaskSource and AcquireState) and exposing them through the Extras struct to provide additional scheduling metadata.
Docstring Coverage ✅ Passed Docstring coverage is 81.25% which is sufficient. The required threshold is 80.00%.

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

✨ Finishing touches
🧪 Generate unit tests (beta)
  • Create PR with unit tests
  • Post copyable unit tests in a comment

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

Comment @coderabbitai help to get the list of available commands and usage tips.

Signed-off-by: zyguan <zhongyangguan@gmail.com>
Copy link

Copilot AI left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

This PR expands task scheduling observability by introducing TaskSource (where a task was popped from) and AcquireState (how a worker acquired a task), then propagating these signals through queue Pop results and into Extras for inspection by runners.

Changes:

  • Add TaskSource and AcquireState enums and expose them via queue::Extras accessors.
  • Replace Pop::from_local with Pop::task_source across queue implementations and add tests validating task source attribution.
  • Track acquisition state in WorkerThread and persist both acquisition state and task source into Extras before handing tasks to the runner.

Reviewed changes

Copilot reviewed 7 out of 7 changed files in this pull request and generated 3 comments.

Show a summary per file
File Description
src/queue/extras.rs Introduces TaskSource/AcquireState and stores them in Extras with public getters.
src/queue.rs Re-exports new enums and changes Pop to carry task_source instead of from_local.
src/queue/single_level.rs Sets task_source on pops from local/global/steal paths; adds tests.
src/queue/multilevel.rs Sets task_source on pops from local/global/steal paths; adds tests.
src/queue/priority.rs Sets task_source for priority queue pops (always global); adds test.
src/pool/worker.rs Computes AcquireState during pop/park flow and writes AcquireState + TaskSource into Extras; adds tests.
src/pool/spawn.rs Suppresses dead_code warnings for compile-time Sync/Send assertion traits.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

/// The source of the task, indicating where the task comes from.
pub task_source: TaskSource,
}

Copy link

Copilot AI Feb 10, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pop is a public struct; replacing the public from_local: bool field with task_source: TaskSource is a breaking API change for downstream users. If backward compatibility is needed, consider keeping from_local (possibly deprecated) or adding a from_local() accessor/compat shim while introducing task_source.

Suggested change
impl<T> Pop<T> {
/// Returns whether this task was popped from a local queue.
///
/// This provides a compatibility shim for the former `from_local: bool`
/// field, using the newer `task_source` information instead.
#[deprecated(
since = "0.0.0",
note = "use `task_source` instead; this accessor is kept for backward compatibility"
)]
pub fn from_local(&self) -> bool {
self.task_source.is_local()
}
}

Copilot uses AI. Check for mistakes.
Comment on lines +30 to +34
let mut state = if retry_after_park {
AcquireState::AfterPark
} else {
AcquireState::Immediate
};
Copy link

Copilot AI Feb 10, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

When retry_after_park is true, state is initialized to AfterPark, so if a task is popped immediately in this call it will still be labeled AfterPark even though no spinning/parking happened during this acquisition. Consider deriving AcquireState based on what happened in the current pop() call (or clarify the semantics in docs), rather than carrying state across iterations via retry_after_park.

Copilot uses AI. Check for mistakes.
Comment on lines +42 to +44
if state == AcquireState::Immediate {
state = AcquireState::AfterSpin;
}
Copy link

Copilot AI Feb 10, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

With the current transition if state == Immediate { state = AfterSpin }, AfterSpin can never be produced when retry_after_park is true (since state starts as AfterPark). If you want to capture “acquired after spinning” even after a previous park/spurious wake, update the state machine accordingly (e.g., track spinning separately or allow transitioning from AfterPark to AfterSpin when spinning occurs).

Copilot uses AI. Check for mistakes.
@cfzjywxk cfzjywxk requested review from cfzjywxk and you06 February 14, 2026 02:00
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants