feat: task scheduler #2406
feat: task scheduler #2406
Conversation
ryanolson
requested review from
PeaBrane,
grahamking,
keivenchang,
kthui and
nnshah1
WalkthroughAdds Tokio tokio-util "rt" feature in Cargo.toml. Introduces utils::tasks module with submodules critical and tracker. Moves CriticalTask* implementation from utils/task.rs to utils/tasks/critical.rs, keeping public API via re-export. Adds CriticalTaskExecutionHandle with monitor, cancellation, join, and detach behavior; exposes new constructors and status methods. Changes
Sequence Diagram(s)sequenceDiagram
participant Caller
participant Runtime
participant MainTask
participant Monitor
participant ParentToken
participant Oneshoot as OneShot(Result)
Caller->>Runtime: new()/new_with_runtime(task_fn, parent_token, desc)
Runtime->>MainTask: spawn task_fn(graceful_token)
Runtime->>Monitor: spawn monitor(MainTask, parent_token, OneShot)
alt MainTask Ok(())
MainTask-->>Monitor: Ok
Monitor-->>Oneshoot: send Ok(())
else MainTask Err or Panic
MainTask-->>Monitor: Err/Panic
Monitor->>ParentToken: cancel()
Monitor-->>Oneshoot: send Err(...)
end
Caller->>Oneshoot: join()
Oneshoot-->>Caller: Result<()>
Caller->>Caller: detach() optional
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes Possibly related PRs
Poem
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. 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
SupportNeed help? Create a ticket on our support page for assistance with any issues or questions. CodeRabbit Commands (Invoked using PR comments)
Other keywords and placeholders
Documentation and Community
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (2)
lib/runtime/src/utils/tasks/critical.rs (1)
170-174: Doc nit: clarify subject of is_finished()"Server's background event loop" looks like a leftover reference.
- /// Check if the task awaiting on the [Server]s background event loop has finished. + /// Check if the critical task (monitored by this handle) has finished.lib/runtime/src/utils/task.rs (1)
4-4: Optional: Soft-deprecate this legacy path to nudge callers to utils::tasks::criticalKeep the re-export for compatibility, but mark it deprecated to encourage gradual migration.
- pub use super::tasks::critical::*; +#[deprecated( + since = "0.4.0", + note = "Moved to utils::tasks::critical; this re-export remains for compatibility." +)] +pub use super::tasks::critical::*;
📜 Review details
Configuration used: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
⛔ Files ignored due to path filters (2)
Cargo.lockis excluded by!**/*.locklib/bindings/python/Cargo.lockis excluded by!**/*.lock
📒 Files selected for processing (5)
Cargo.toml(1 hunks)lib/runtime/src/utils.rs(1 hunks)lib/runtime/src/utils/task.rs(1 hunks)lib/runtime/src/utils/tasks.rs(1 hunks)lib/runtime/src/utils/tasks/critical.rs(1 hunks)
🧰 Additional context used
🧬 Code Graph Analysis (1)
lib/runtime/src/utils/tasks/critical.rs (3)
lib/runtime/src/utils/tasks/tracker.rs (4)
tokio(2571-2571)tokio(2833-2833)tokio(2878-2878)tokio(2910-2910)lib/runtime/src/pipeline/nodes/sources/base.rs (1)
oneshot(73-73)lib/runtime/src/runtime.rs (1)
handle(123-128)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (5)
- GitHub Check: Mirror Repository to GitLab
- GitHub Check: Build and Test - dynamo
- GitHub Check: pre-merge-rust (lib/runtime/examples)
- GitHub Check: pre-merge-rust (lib/bindings/python)
- GitHub Check: pre-merge-rust (.)
🔇 Additional comments (5)
Cargo.toml (1)
62-62: No feature gate needed for CancellationTokenAccording to the tokio-util v0.7 source (see https://docs.rs/tokio-util/latest/source/src/lib.rs), the
syncmodule (and thusCancellationToken) is exposed unconditionally—there is no separate"sync"Cargo feature to enable. Your current dependency line already provides it:tokio-util = { version = "0.7", features = ["codec", "net", "rt"] }No change is required to use
tokio_util::sync::CancellationToken.Likely an incorrect or invalid review comment.
lib/runtime/src/utils/tasks/critical.rs (2)
49-60: Constructor ergonomics LGTMnew() delegating to new_with_runtime via Handle::try_current() is clean and fail-fast outside a runtime. API is clear.
104-160: Monitor task correctly propagates failures and panics to parent tokenThe monitor awaits the main task, cancels parent on error/panic, and forwards the result through oneshot. Good separation of concerns and observability via tracing.
lib/runtime/src/utils/tasks.rs (1)
4-5: Module layout LGTMPublicly exposing critical and tracker modules under utils::tasks is clear and future-proof.
lib/runtime/src/utils.rs (1)
22-22: Export of utils::tasks is correct and non-breakingKeeps existing modules and adds tasks; pairs well with the re-export in utils::task for compatibility.
I got rid of the prom format/output method entirely since it's handled by the registry. we only every inc/dec/set counters/gauges |
The idea is the root can be a prometheus root or a local root.
All are |
|
@keivenchang - we can remove these: as active = issued - completed and queued. we could just have a started counter, so then we could get queues from active - started. |
removed the active gauge in favor of only monotonic counters. |
|
Approving without review, as requested. 7k lines. What can you do? |
Signed-off-by: Ryan Olson <ryanolson@users.noreply.github.com>
|
Just looked through the metrics part. The refactor looks more understandable now. Let me regurgitate this as I type it, so that I can remember it (and for future reference). Basically:
Thank you for simplifying it (reducing the number of Metrics) and making it more readable. |
Signed-off-by: Ryan Olson <ryanolson@users.noreply.github.com> Signed-off-by: Hannah Zhang <hannahz@nvidia.com>
Signed-off-by: Ryan Olson <ryanolson@users.noreply.github.com>
I needed something similar for kvbm, so I made it general.
Summary by CodeRabbit