[ISSUE #7544]🚀Harden Tokio runtime lifecycle: service-tree ownership, shutdown reports, scheduler guardrails, and diagnostics#7545
Conversation
… shutdown reports, scheduler guardrails, and diagnostics
|
🔊@mxsm 🚀Thanks for your contribution🎉! 💡CodeRabbit(AI) will review your code first🔥! Note 🚨The code review suggestions from CodeRabbit are to be used as a reference only, and the PR submitter can decide whether to make changes based on their own judgment. Ultimately, the project management personnel will conduct the final code review💥. |
WalkthroughThis PR hardens the Tokio runtime lifecycle across the RocketMQ Rust codebase: binary entrypoints are refactored to explicit multi-thread runtime construction, ChangesRuntime lifecycle hardening
Estimated code review effort🎯 5 (Critical) | ⏱️ ~120 minutes Possibly related PRs
Suggested labels
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 inconclusive)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
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. Comment |
There was a problem hiding this comment.
Actionable comments posted: 5
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
rocketmq/src/schedule/executor.rs (1)
110-123: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winKeep the parent
TaskGroupfor executor reuse after shutdown.Line 111 only stores the derived child group. Once
shutdown_task_group()takes it, the next execution falls back tonew_executor_task_group()and creates a root group, so a parented executor loses service-tree ownership aftershutdown_all()or scheduler stop/restart.Proposed direction
pub struct TaskExecutor { task_group: Arc<RwLock<Option<TaskGroup>>>, + parent_task_group: Option<TaskGroup>, // ... } impl TaskExecutor { pub fn new(config: ExecutorConfig) -> Self { Self::new_with_optional_task_group(config, None) } pub fn new_with_task_group(config: ExecutorConfig, parent_task_group: TaskGroup) -> Self { - Self::new_with_optional_task_group(config, Some(parent_task_group.child("rocketmq.task-executor"))) + Self::new_with_optional_task_group(config, Some(parent_task_group)) } - fn new_with_optional_task_group(config: ExecutorConfig, task_group: Option<TaskGroup>) -> Self { + fn new_with_optional_task_group(config: ExecutorConfig, parent_task_group: Option<TaskGroup>) -> Self { let semaphore = Arc::new(Semaphore::new(config.max_concurrent_tasks)); + let task_group = parent_task_group + .as_ref() + .map(|parent_task_group| parent_task_group.child("rocketmq.task-executor")); Self { config, semaphore, running_tasks: Arc::new(RwLock::new(HashMap::new())), metrics: Arc::new(RwLock::new(ExecutionMetrics::default())), executions: Arc::new(RwLock::new(HashMap::new())), task_group: Arc::new(RwLock::new(task_group)), + parent_task_group, last_task_group_shutdown_report: Arc::new(RwLock::new(None)), } } async fn task_group(&self, operation: &'static str) -> Result<TaskGroup, SchedulerError> { let mut task_group = self.task_group.write().await; if let Some(task_group) = task_group.as_ref() { return Ok(task_group.clone()); } - let new_group = new_executor_task_group(operation)?; + let new_group = if let Some(parent_task_group) = self.parent_task_group.as_ref() { + parent_task_group.child("rocketmq.task-executor") + } else { + new_executor_task_group(operation)? + }; *task_group = Some(new_group.clone()); Ok(new_group) } }Add a regression test that calls
shutdown_all()on anew_with_task_groupexecutor, executes another task, and asserts the recreated task group still has the service task group as parent.As per coding guidelines, “For behavior changes, add or update focused tests that would fail without the change.”
Also applies to: 279-287, 309-315
🤖 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 `@rocketmq/src/schedule/executor.rs` around lines 110 - 123, The executor loses its parent TaskGroup relationship after shutdown because only the derived child group is stored. In the `new_with_optional_task_group` method, store the parent TaskGroup separately (when provided) in addition to the task_group field. When the task_group needs to be recreated (such as after `shutdown_task_group()` is called and before the next execution), use the stored parent TaskGroup to create a new child with the same parent relationship instead of falling back to `new_executor_task_group()`. Additionally, add a regression test that verifies after calling `shutdown_all()` on a `new_with_task_group` executor, executing another task results in a task group that still maintains the original parent relationship.Source: Coding guidelines
rocketmq-controller/src/controller/controller_manager.rs (1)
438-463: 🩺 Stability & Availability | 🟠 Major | 🏗️ Heavy liftPropagate
parent_task_groupinto RaftController construction.Line 447-451 passes the parent into
DefaultBrokerHeartbeatManager, but Line 460-463 still buildsRaftControllerthroughnew_open_raft_with_heartbeat(...)without parent context. This leaves OpenRaft-owned tasks outside the intended service tree, weakening shutdown/cancellation/report diagnostics hierarchy.🔧 Suggested fix
- let raft_arc = ArcMut::new(RaftController::new_open_raft_with_heartbeat( - config.clone(), - heartbeat_manager.clone(), - )); + let raft_arc = ArcMut::new(match parent_task_group.as_ref() { + Some(parent_task_group) => RaftController::new_open_raft_with_heartbeat_and_task_group( + config.clone(), + heartbeat_manager.clone(), + parent_task_group.clone(), + ), + None => RaftController::new_open_raft_with_heartbeat( + config.clone(), + heartbeat_manager.clone(), + ), + });🤖 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 `@rocketmq-controller/src/controller/controller_manager.rs` around lines 438 - 463, The RaftController construction in the new_with_optional_task_group method does not pass the parent_task_group parameter like the DefaultBrokerHeartbeatManager does. Modify the RaftController creation at the new_open_raft_with_heartbeat call to conditionally pass the parent_task_group, following the same pattern used above for DefaultBrokerHeartbeatManager where you check if parent_task_group exists via a match statement and pass it accordingly when available.
🧹 Nitpick comments (4)
rocketmq-common/src/common/stats/moment_stats_item.rs (1)
48-90: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd a focused test for the parented
MomentStatsItempath.
new_with_task_group(...).init()now has externally visible shutdown-report behavior, but the supplied tests do not exercise this constructor/init branch. Please add a small Tokio test that initializes a parented item, shuts it down, and asserts the parent report containsrocketmq-common.moment-stats.<name>.<key>. As per coding guidelines, “For behavior changes, add or update focused tests that would fail without the change.”🤖 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 `@rocketmq-common/src/common/stats/moment_stats_item.rs` around lines 48 - 90, The parented MomentStatsItem initialization path created by new_with_task_group followed by init is not covered by existing tests. Add a new Tokio test that creates a TaskGroup parent, initializes a MomentStatsItem using new_with_task_group with that parent, calls init on the created item wrapped in Arc, shuts down the item, and then verifies that the parent task group's shutdown report contains the expected group name format rocketmq-common.moment-stats followed by the stats_name and stats_key to confirm the parented initialization and shutdown-report behavior works correctly.Source: Coding guidelines
rocketmq-common/src/common/stats/moment_stats_item_set.rs (1)
156-163: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winExercise the parented item branch in the new set test.
The new test proves the stats-set task group is parented, but it never calls
get_and_create_stats_item, so the newMomentStatsItem::new_with_task_groupbranch at Lines 156-161 can regress unnoticed. Extend the test to create/init one item and assert the parent shutdown report includes itsrocketmq-common.moment-stats.<name>.<key>child. As per coding guidelines, “For behavior changes, add or update focused tests that would fail without the change.”Also applies to: 213-232
🤖 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 `@rocketmq-common/src/common/stats/moment_stats_item_set.rs` around lines 156 - 163, The test for parented stats-set does not exercise the MomentStatsItem::new_with_task_group code path because it never calls get_and_create_stats_item. Extend the test to call get_and_create_stats_item to create a stats item within the parented stats-set, then add an assertion that verifies the parent shutdown report includes the created child item with the expected naming format rocketmq-common.moment-stats.<name>.<key>. This ensures the new_with_task_group branch is tested and cannot regress unnoticed. Apply the same changes to the other test location mentioned at lines 213-232.Source: Coding guidelines
rocketmq-broker/src/broker_runtime.rs (1)
5020-5029: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAssert the parent service shutdown report is healthy.
The child-name assertion can still pass if the parent shutdown report timed out or is unhealthy. Add the health assertion before checking children.
As per coding guidelines, “For behavior changes, add or update focused tests that would fail without the change.”
Proposed test strengthening
let service_report = service.task_group().shutdown(Duration::from_secs(1)).await; + assert!(service_report.is_healthy(), "{}", service_report.to_json()); assert!( service_report .children🤖 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 `@rocketmq-broker/src/broker_runtime.rs` around lines 5020 - 5029, The assertion in the shutdown test for service.task_group() only checks for a specific child in the shutdown report but does not verify that the parent service_report itself is healthy. Add a health assertion before the existing children check to ensure the service_report indicates a successful/healthy shutdown status. This prevents the test from passing when the parent shutdown report has timed out or indicates an unhealthy state while still containing the expected child. Verify the health status using the appropriate method or property on the service_report object before proceeding with the children iteration assertion.Source: Coding guidelines
rocketmq-broker/src/out_api/broker_outer_api.rs (1)
357-360: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winConsider adding documentation for the new public method.
This new public API method would benefit from a doc comment explaining the timeout parameter and the shutdown report return value. While the existing
shutdown()method also lacks documentation, adding it here would help users understand the timeout behavior and what information the report contains.📝 Example documentation
+ /// Shutdown the broker outer API with timeout and report. + /// + /// # Arguments + /// * `timeout` - Maximum duration to wait for graceful shutdown + /// + /// # Returns + /// A shutdown report containing connection closure details and timing information pub async fn shutdown_with_report(&mut self, timeout: std::time::Duration) -> RemotingClientShutdownReport { self.remoting_client.shutdown_with_report(timeout).await }🤖 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 `@rocketmq-broker/src/out_api/broker_outer_api.rs` around lines 357 - 360, Add a documentation comment above the shutdown_with_report method to explain its purpose, parameters, and return value. Document the timeout parameter to clarify its behavior and explain what information the RemotingClientShutdownReport return value provides to help users understand the shutdown behavior and what data they can expect from the report.
🤖 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 `@rocketmq-broker/src/broker_runtime.rs`:
- Around line 674-678: The BrokerOuterAPI instantiation on line 678 is not
propagating the service_context for proper task lifecycle management, unlike the
ScheduledTaskManager on line 674 which correctly uses it. Add a
new_with_service_context constructor method to BrokerOuterAPI (or modify the
existing new method) that accepts a service context parameter and passes it to
the underlying remoting client to enable proper task group parenting. Then
update the BrokerOuterAPI::new call to use this new constructor, passing the
service_context.clone() from the conditional block so that remoting client tasks
are properly managed under the broker shutdown tree.
- Around line 1158-1166: The broker_outer_api shutdown report handling does not
check for timeout status before mapping to unhealthy state, which loses timeout
diagnostics. Refactor the logic for broker_outer_api_report to follow the same
pattern used for client_housekeeping via from_shutdown_report(): first check if
any internal reports (background, workers, connections) within the
RemotingClientShutdownReport timed out and use the timed_out variant if needed,
then check is_healthy() for the completed variant, and finally fall back to the
unhealthy variant. This ensures timeouts are properly captured in the shutdown
report rather than losing that diagnostic information.
In `@rocketmq-proxy/src/grpc/service.rs`:
- Around line 475-481: The no-parent branch in the map_or_else closure uses
tokio::runtime::Handle::current() which panics if there's no active Tokio
runtime, but should use the defensive Handle::try_current() pattern for
consistency with other housekeeping paths in the codebase (gRPC server,
Prometheus exporter, tiered store). Replace Handle::current() with
Handle::try_current() in the RuntimeHandle::new call within the TaskGroup::root
creation, and handle the error case by either returning an error or constructing
an inactive/unhealthy report when the runtime is unavailable.
In `@rocketmq-remoting/src/clients/rocketmq_tokio_client.rs`:
- Around line 723-736: The connect path in the time::timeout block (where
Client::connect_with_service_context and Client::connect are called) does not
check the shutdown_token before or after establishing the connection, allowing
connections to complete and be inserted into connection_tables even after
shutdown_with_report has drained it. Add a shutdown_token check before
initiating the connect operation and after it completes but before returning the
client result to prevent late connection inserts. Additionally, create a
regression test that races shutdown_with_report execution against an in-flight
create_client operation to verify that connections initiated during shutdown do
not get inserted into connection_tables after the shutdown report is built.
- Around line 1018-1022: The connection shutdown loop is applying the full
timeout sequentially to each client, causing total shutdown time to scale
linearly with the number of connections. Replace the sequential iteration in the
loop (where client.close_with_report(timeout) is called for each client one at a
time) with concurrent execution by collecting all the close futures and then
awaiting them concurrently using futures operations like join_all or similar
concurrent combinators. Additionally, add a focused regression test that
verifies multiple connections are shutdown concurrently and confirms the total
shutdown time is bounded by a single timeout value rather than being multiplied
by the connection count.
---
Outside diff comments:
In `@rocketmq-controller/src/controller/controller_manager.rs`:
- Around line 438-463: The RaftController construction in the
new_with_optional_task_group method does not pass the parent_task_group
parameter like the DefaultBrokerHeartbeatManager does. Modify the RaftController
creation at the new_open_raft_with_heartbeat call to conditionally pass the
parent_task_group, following the same pattern used above for
DefaultBrokerHeartbeatManager where you check if parent_task_group exists via a
match statement and pass it accordingly when available.
In `@rocketmq/src/schedule/executor.rs`:
- Around line 110-123: The executor loses its parent TaskGroup relationship
after shutdown because only the derived child group is stored. In the
`new_with_optional_task_group` method, store the parent TaskGroup separately
(when provided) in addition to the task_group field. When the task_group needs
to be recreated (such as after `shutdown_task_group()` is called and before the
next execution), use the stored parent TaskGroup to create a new child with the
same parent relationship instead of falling back to `new_executor_task_group()`.
Additionally, add a regression test that verifies after calling `shutdown_all()`
on a `new_with_task_group` executor, executing another task results in a task
group that still maintains the original parent relationship.
---
Nitpick comments:
In `@rocketmq-broker/src/broker_runtime.rs`:
- Around line 5020-5029: The assertion in the shutdown test for
service.task_group() only checks for a specific child in the shutdown report but
does not verify that the parent service_report itself is healthy. Add a health
assertion before the existing children check to ensure the service_report
indicates a successful/healthy shutdown status. This prevents the test from
passing when the parent shutdown report has timed out or indicates an unhealthy
state while still containing the expected child. Verify the health status using
the appropriate method or property on the service_report object before
proceeding with the children iteration assertion.
In `@rocketmq-broker/src/out_api/broker_outer_api.rs`:
- Around line 357-360: Add a documentation comment above the
shutdown_with_report method to explain its purpose, parameters, and return
value. Document the timeout parameter to clarify its behavior and explain what
information the RemotingClientShutdownReport return value provides to help users
understand the shutdown behavior and what data they can expect from the report.
In `@rocketmq-common/src/common/stats/moment_stats_item_set.rs`:
- Around line 156-163: The test for parented stats-set does not exercise the
MomentStatsItem::new_with_task_group code path because it never calls
get_and_create_stats_item. Extend the test to call get_and_create_stats_item to
create a stats item within the parented stats-set, then add an assertion that
verifies the parent shutdown report includes the created child item with the
expected naming format rocketmq-common.moment-stats.<name>.<key>. This ensures
the new_with_task_group branch is tested and cannot regress unnoticed. Apply the
same changes to the other test location mentioned at lines 213-232.
In `@rocketmq-common/src/common/stats/moment_stats_item.rs`:
- Around line 48-90: The parented MomentStatsItem initialization path created by
new_with_task_group followed by init is not covered by existing tests. Add a new
Tokio test that creates a TaskGroup parent, initializes a MomentStatsItem using
new_with_task_group with that parent, calls init on the created item wrapped in
Arc, shuts down the item, and then verifies that the parent task group's
shutdown report contains the expected group name format
rocketmq-common.moment-stats followed by the stats_name and stats_key to confirm
the parented initialization and shutdown-report behavior works correctly.
🪄 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: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: 1f1bb956-2b3a-4ef9-8937-4337a4730f17
📒 Files selected for processing (40)
rocketmq-broker/src/bin/broker_bootstrap_server.rsrocketmq-broker/src/broker_runtime.rsrocketmq-broker/src/out_api/broker_outer_api.rsrocketmq-client/src/factory/mq_client_instance.rsrocketmq-client/src/producer/transaction_mq_produce_builder.rsrocketmq-client/src/runtime.rsrocketmq-common/src/common/statistics/statistics_manager.rsrocketmq-common/src/common/stats/moment_stats_item.rsrocketmq-common/src/common/stats/moment_stats_item_set.rsrocketmq-common/src/utils/http_tiny_client.rsrocketmq-controller/src/controller/controller_manager.rsrocketmq-controller/src/controller/open_raft_controller.rsrocketmq-controller/src/heartbeat/default_broker_heartbeat_manager.rsrocketmq-controller/src/metadata/broker.rsrocketmq-controller/src/rpc/server.rsrocketmq-namesrv/src/bin/namesrv_bootstrap_server.rsrocketmq-observability/src/exporter/prometheus.rsrocketmq-proxy/src/bin/rocketmq-proxy-rust.rsrocketmq-proxy/src/grpc/server.rsrocketmq-proxy/src/grpc/service.rsrocketmq-remoting/src/clients/client.rsrocketmq-remoting/src/clients/connection_pool.rsrocketmq-remoting/src/clients/rocketmq_tokio_client.rsrocketmq-runtime/src/diagnostics.rsrocketmq-runtime/src/legacy.rsrocketmq-runtime/src/lib.rsrocketmq-runtime/src/task_group.rsrocketmq-runtime/tests/runtime_model.rsrocketmq-store/src/stats/broker_stats_manager.rsrocketmq-tieredstore/src/dispatcher/tiered_dispatcher.rsrocketmq-tieredstore/src/runtime.rsrocketmq-tieredstore/src/service.rsrocketmq-tools/rocketmq-admin/rocketmq-admin-tui/src/main.rsrocketmq/benches/scheduler_lifecycle_bench.rsrocketmq/src/schedule.rsrocketmq/src/schedule/executor.rsrocketmq/src/schedule/scheduler.rsrocketmq/src/task/service_task.rsscripts/runtime-audit-baseline.jsonscripts/runtime-audit.ps1
| let scheduled_task_manager = service_context | ||
| .as_ref() | ||
| .map(|context| ScheduledTaskManager::new_with_task_group(context.task_group().clone())) | ||
| .unwrap_or_else(ScheduledTaskManager::new_legacy_compatibility); | ||
| let broker_outer_api = BrokerOuterAPI::new(Arc::new(TokioClientConfig::default())); |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Inspect whether BrokerOuterAPI can propagate ServiceContext into its remoting client.
rg -n "struct BrokerOuterAPI|impl BrokerOuterAPI|new_with_service_context|RocketmqDefaultClient::new|RocketmqDefaultClient::new_with_service_context" \
rocketmq-broker/src/out_api rocketmq-remoting/src/clientsRepository: mxsm/rocketmq-rust
Length of output: 1934
🏁 Script executed:
#!/bin/bash
# Inspect BrokerOuterAPI full implementation and constructors
ast-grep outline rocketmq-broker/src/out_api/broker_outer_api.rsRepository: mxsm/rocketmq-rust
Length of output: 2774
🏁 Script executed:
#!/bin/bash
# Get the full BrokerOuterAPI struct and impl blocks to see all constructors
sed -n '127,200p' rocketmq-broker/src/out_api/broker_outer_api.rsRepository: mxsm/rocketmq-rust
Length of output: 2954
🏁 Script executed:
#!/bin/bash
# Check the RocketmqDefaultClient::new_with_service_context signature
sed -n '249,270p' rocketmq-remoting/src/clients/rocketmq_tokio_client.rsRepository: mxsm/rocketmq-rust
Length of output: 990
Add service-context support to BrokerOuterAPI constructor to enable proper lifecycle management of remoting client tasks under the broker shutdown tree.
The ScheduledTaskManager on line 674 correctly uses service_context for task group parenting, but the adjacent BrokerOuterAPI::new(...) on line 678 does not. RocketmqDefaultClient supports service-context aware construction via new_with_service_context, but BrokerOuterAPI currently lacks a constructor that propagates it. Add a new_with_service_context method (or modify existing constructors) to accept and pass the service context to the underlying remoting client.
🤖 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 `@rocketmq-broker/src/broker_runtime.rs` around lines 674 - 678, The
BrokerOuterAPI instantiation on line 678 is not propagating the service_context
for proper task lifecycle management, unlike the ScheduledTaskManager on line
674 which correctly uses it. Add a new_with_service_context constructor method
to BrokerOuterAPI (or modify the existing new method) that accepts a service
context parameter and passes it to the underlying remoting client to enable
proper task group parenting. Then update the BrokerOuterAPI::new call to use
this new constructor, passing the service_context.clone() from the conditional
block so that remoting client tasks are properly managed under the broker
shutdown tree.
| shutdown_report.broker_outer_api = if broker_outer_api_report.is_healthy() { | ||
| BrokerShutdownComponentReport::completed("broker_outer_api", started.elapsed()) | ||
| } else { | ||
| BrokerShutdownComponentReport::unhealthy( | ||
| "broker_outer_api", | ||
| started.elapsed(), | ||
| format!("{broker_outer_api_report:?}"), | ||
| ) | ||
| }; |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Verify the timeout API exposed by broker outer API / remoting client shutdown reports.
rg -n "shutdown_with_report|struct .*ShutdownReport|impl .*ShutdownReport|fn .*timed|timed_out|is_healthy" \
rocketmq-broker/src/out_api rocketmq-remoting/src/clientsRepository: mxsm/rocketmq-rust
Length of output: 4907
🏁 Script executed:
# 1. Look at RemotingClientShutdownReport structure and its fields
rg -n "pub struct RemotingClientShutdownReport" -A 10 rocketmq-remoting/src/clients/rocketmq_tokio_client.rs
# 2. Look at BrokerShutdownComponentReport definition and methods
rg -n "struct BrokerShutdownComponentReport|impl.*BrokerShutdownComponentReport" rocketmq-broker/src/broker_runtime.rsRepository: mxsm/rocketmq-rust
Length of output: 855
🏁 Script executed:
# 3. Get the full context around lines 1158-1166 in broker_runtime.rs
sed -n '1140,1180p' rocketmq-broker/src/broker_runtime.rsRepository: mxsm/rocketmq-rust
Length of output: 1945
🏁 Script executed:
# 4. Search for timed_out method/field definition in BrokerShutdownComponentReport
rg -n "fn timed_out|timed_out:" rocketmq-broker/src/broker_runtime.rsRepository: mxsm/rocketmq-rust
Length of output: 486
🏁 Script executed:
# Look at ShutdownReport structure to understand timeout tracking
rg -n "struct ShutdownReport|impl.*ShutdownReport" -A 5 rocketmq-remoting/src/clients/rocketmq_tokio_client.rs | head -40
# Look at from_shutdown_report implementation
rg -n "fn from_shutdown_report" -A 20 rocketmq-broker/src/broker_runtime.rsRepository: mxsm/rocketmq-rust
Length of output: 1277
🏁 Script executed:
# Search for timeout detection logic
rg -n "timed_out|timeout" rocketmq-remoting/src/clients/rocketmq_tokio_client.rs | grep -i "report\|shutdown" | head -20Repository: mxsm/rocketmq-rust
Length of output: 490
🏁 Script executed:
# Find the shutdown_report_has_timed_out function definition
rg -n "fn shutdown_report_has_timed_out" -A 10 rocketmq-broker/src/broker_runtime.rsRepository: mxsm/rocketmq-rust
Length of output: 723
🏁 Script executed:
# Look at ShutdownReport struct definition in remoting
rg -n "pub struct ShutdownReport" -A 15 rocketmq-remoting/src/clients/rocketmq_tokio_client.rsRepository: mxsm/rocketmq-rust
Length of output: 156
🏁 Script executed:
# Find ShutdownReport struct definition
rg -n "pub struct ShutdownReport" -B 2 -A 20 rocketmq-remoting/src/Repository: mxsm/rocketmq-rust
Length of output: 156
🏁 Script executed:
# Check if RemotingClientShutdownReport contains ShutdownReport fields
rg -n "pub struct RemotingClientShutdownReport" -A 6 rocketmq-remoting/src/clients/rocketmq_tokio_client.rsRepository: mxsm/rocketmq-rust
Length of output: 416
🏁 Script executed:
# Check how shutdown timeout is used in message_store to understand the pattern
sed -n '1140,1158p' rocketmq-broker/src/broker_runtime.rsRepository: mxsm/rocketmq-rust
Length of output: 994
🏁 Script executed:
# Look at shutdown_with_report for message_store to see the pattern
sed -n '1128,1150p' rocketmq-broker/src/broker_runtime.rs
# Search for how message_store timeout is detected
rg -n "MESSAGE_STORE_SHUTDOWN_TIMEOUT\|started.elapsed()" rocketmq-broker/src/broker_runtime.rs | head -20Repository: mxsm/rocketmq-rust
Length of output: 1260
🏁 Script executed:
# Check all timeout constants defined
rg -n "SHUTDOWN_TIMEOUT" rocketmq-broker/src/broker_runtime.rsRepository: mxsm/rocketmq-rust
Length of output: 476
🏁 Script executed:
# Verify the complete from_shutdown_report implementation and how it handles timeouts
sed -n '522,535p' rocketmq-broker/src/broker_runtime.rs
# Look at client_housekeeping shutdown code to see the correct pattern
sed -n '1167,1179p' rocketmq-broker/src/broker_runtime.rsRepository: mxsm/rocketmq-rust
Length of output: 1288
🏁 Script executed:
# Check if RemotingClientShutdownReport.is_healthy() loses timeout info
# by looking at how ShutdownReport is used
rg -n "pub struct ShutdownReport\|pub struct RemotingClientConnectionShutdownReport" -A 10 rocketmq-remoting/src/clients/rocketmq_tokio_client.rsRepository: mxsm/rocketmq-rust
Length of output: 156
🏁 Script executed:
# Get a complete picture of how RemotingClientShutdownReport's internal ShutdownReports work
sed -n '203,230p' rocketmq-remoting/src/clients/rocketmq_tokio_client.rs
# Look at the shutdown implementation to see if it sets timed_out on ShutdownReport
rg -n "fn shutdown\(" -A 30 rocketmq-remoting/src/clients/rocketmq_tokio_client.rs | head -50Repository: mxsm/rocketmq-rust
Length of output: 2540
🏁 Script executed:
# Check if there's a method on RemotingClientShutdownReport to detect timeout
rg -n "impl RemotingClientShutdownReport" -A 30 rocketmq-remoting/src/clients/rocketmq_tokio_client.rs
# Check the full structure of ShutdownReport
fd ShutdownReport rocketmq-remoting/src/ --type f
rg -n "pub struct ShutdownReport\|pub fn timed_out\|pub timed_out:" rocketmq-remoting/src/Repository: mxsm/rocketmq-rust
Length of output: 1732
🏁 Script executed:
# Search for how to detect timeout in RemotingClientShutdownReport
# Look for any helper function that checks timeouts in RemotingClientShutdownReport
rg -n "RemotingClientShutdownReport\|shutdown_report_has_timed_out" rocketmq-broker/src/broker_runtime.rs | head -20
# Check the exact location where broker_outer_api timeout could be detected
sed -n '1152,1168p' rocketmq-broker/src/broker_runtime.rsRepository: mxsm/rocketmq-rust
Length of output: 814
🏁 Script executed:
# Look for ShutdownReport definition to understand timed_out field
rg -n "pub struct ShutdownReport" rocketmq-remoting/src/ -A 10Repository: mxsm/rocketmq-rust
Length of output: 156
🏁 Script executed:
# Find where ShutdownReport is defined
fd -e rs -type f | xargs rg -l "pub struct ShutdownReport"
# Get the definition
rg -n "pub struct ShutdownReport" -A 15 --type rsRepository: mxsm/rocketmq-rust
Length of output: 394
🏁 Script executed:
# Read the ShutdownReport struct definition
cat rocketmq-runtime/src/shutdown_report.rsRepository: mxsm/rocketmq-rust
Length of output: 4337
Check for timeout in broker_outer_api shutdown report before mapping to unhealthy.
Line 1156 passes BROKER_OUTER_API_SHUTDOWN_TIMEOUT, but the current code at line 1161 maps all non-healthy reports to unhealthy() without checking if a timeout actually occurred. This loses timeout diagnostics since timed_out_component_names() only includes components with the timed_out flag set.
The RemotingClientShutdownReport contains internal ShutdownReport instances (background, workers, connections) that track timeouts. Follow the pattern used for client_housekeeping via from_shutdown_report(): check if any internal reports timed out first, then check health, then default to unhealthy.
🤖 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 `@rocketmq-broker/src/broker_runtime.rs` around lines 1158 - 1166, The
broker_outer_api shutdown report handling does not check for timeout status
before mapping to unhealthy state, which loses timeout diagnostics. Refactor the
logic for broker_outer_api_report to follow the same pattern used for
client_housekeeping via from_shutdown_report(): first check if any internal
reports (background, workers, connections) within the
RemotingClientShutdownReport timed out and use the timed_out variant if needed,
then check is_healthy() for the completed variant, and finally fall back to the
unhealthy variant. This ensures timeouts are properly captured in the shutdown
report rather than losing that diagnostic information.
| let task_group = parent_task_group.map_or_else( | ||
| || { | ||
| TaskGroup::root( | ||
| "rocketmq-proxy.grpc.housekeeping", | ||
| RuntimeHandle::new(tokio::runtime::Handle::current()), | ||
| ) | ||
| }, |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Verify Tokio runtime handle usage in the proxy housekeeping path and compare with nearby fallible patterns.
rg -n -C3 'Handle::(current|try_current)\(' rocketmq-proxy/src/grpc rocketmq-observability/src/exporter rocketmq-tieredstore/srcRepository: mxsm/rocketmq-rust
Length of output: 2911
🏁 Script executed:
#!/bin/bash
# Get context around the housekeeping path in service.rs
sed -n '440,490p' rocketmq-proxy/src/grpc/service.rsRepository: mxsm/rocketmq-rust
Length of output: 2037
🏁 Script executed:
#!/bin/bash
# Check the ProxyHousekeepingRunReport type and its fields
rg -n 'struct ProxyHousekeepingRunReport|enum ProxyHousekeepingRunReport' rocketmq-proxy/src
# Also check how TaskGroup::root is used in context
rg -A5 'ProxyHousekeepingRunReport' rocketmq-proxy/src/grpc/service.rs | head -30Repository: mxsm/rocketmq-rust
Length of output: 1227
🏁 Script executed:
#!/bin/bash
# Check the ShutdownReport type to see if it can represent error states
rg -n 'struct ShutdownReport|enum ShutdownReport' rocketmq-proxy/src --type rs -A10
# Also check imports to understand ShutdownReport origin
rg -n 'use.*ShutdownReport' rocketmq-proxy/src/grpc/service.rsRepository: mxsm/rocketmq-rust
Length of output: 224
🏁 Script executed:
#!/bin/bash
# Search for ShutdownReport definition (it's from rocketmq_runtime)
find . -name "*.rs" -path "*/rocketmq_runtime/*" | xargs rg -n 'struct ShutdownReport|enum ShutdownReport' -A10
# Also check if TaskGroup::root could fail or handle panic differently
rg -n 'impl.*TaskGroup' --type-list | head -5
rg -n 'fn root' $(find . -name "*.rs" -path "*/rocketmq_runtime/*" 2>/dev/null | head -5)Repository: mxsm/rocketmq-rust
Length of output: 1324
🏁 Script executed:
#!/bin/bash
# Search for TaskGroup::root implementation
find . -name "*.rs" -path "*/rocketmq_runtime/*" | xargs rg -n 'fn root' -B2 -A8
# Also check the TaskGroup struct and impl
find . -name "*.rs" -path "*/rocketmq_runtime/*" | xargs rg -n 'struct TaskGroup|impl TaskGroup' -A5Repository: mxsm/rocketmq-rust
Length of output: 3523
🏁 Script executed:
#!/bin/bash
# Check if RuntimeHandle::new can fail or if Handle::current panics
rg -n 'impl RuntimeHandle|fn new' rocketmq-runtime/src/runtime.rs -A5 | head -30
# Also check if there are any existing error handling patterns in ProxyGrpcService
rg -n 'ProxyError|RocketMQError' rocketmq-proxy/src/grpc/service.rs | head -20Repository: mxsm/rocketmq-rust
Length of output: 1020
🏁 Script executed:
#!/bin/bash
# Find RuntimeHandle definition
find . -name "*.rs" -path "*/rocketmq_runtime/*" | xargs rg -n 'struct RuntimeHandle|impl RuntimeHandle' -A8 | head -50
# Also check if Handle::current() is actually documented to panic
rg -n 'RuntimeHandle::new' rocketmq-proxy/src/grpc/service.rs -B5 -A5Repository: mxsm/rocketmq-rust
Length of output: 1364
🏁 Script executed:
#!/bin/bash
# Check where run_housekeeping_until_with_report is called from
rg -n 'run_housekeeping_until_with_report' rocketmq-proxy/src --type rs
# Also check if there are tests that might call this without a runtime
rg -n 'run_housekeeping' rocketmq-proxy/src --type rs -B3 -A3 | head -50Repository: mxsm/rocketmq-rust
Length of output: 210
🏁 Script executed:
#!/bin/bash
# Check where run_housekeeping_until_with_report is called from
rg -n 'run_housekeeping_until_with_report' rocketmq-proxy/src
# Also check for tests
rg -n '#\[test\]|#\[tokio::test\]' rocketmq-proxy/src/grpc/service.rs -A10 | head -30Repository: mxsm/rocketmq-rust
Length of output: 2049
Use Handle::try_current() instead of Handle::current() in the standalone housekeeping path.
The no-parent branch calls tokio::runtime::Handle::current(), which panics if invoked without a current Tokio runtime. Other housekeeping paths in the codebase (gRPC server, Prometheus exporter, tiered store) use the fallible Handle::try_current() pattern for consistency. Since run_housekeeping_until_with_report is a public API with no runtime guarantee, it should follow the same defensive approach: either return an error or an inactive/unhealthy report when the runtime is unavailable.
🤖 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 `@rocketmq-proxy/src/grpc/service.rs` around lines 475 - 481, The no-parent
branch in the map_or_else closure uses tokio::runtime::Handle::current() which
panics if there's no active Tokio runtime, but should use the defensive
Handle::try_current() pattern for consistency with other housekeeping paths in
the codebase (gRPC server, Prometheus exporter, tiered store). Replace
Handle::current() with Handle::try_current() in the RuntimeHandle::new call
within the TaskGroup::root creation, and handle the error case by either
returning an error or constructing an inactive/unhealthy report when the runtime
is unavailable.
| let service_context = self.service_context.clone(); | ||
| let connect_result = time::timeout(duration, async move { | ||
| if let Some(service_context) = service_context.as_ref() { | ||
| Client::connect_with_service_context( | ||
| service_context, | ||
| addr_inner, | ||
| self.cmd_handler.clone(), | ||
| self.tx.as_ref(), | ||
| tls_config, | ||
| ) | ||
| .await | ||
| } else { | ||
| Client::connect(addr_inner, self.cmd_handler.clone(), self.tx.as_ref(), tls_config).await | ||
| } |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Block late connection inserts after shutdown starts.
shutdown_with_report() cancels and drains the current connection_tables, but cloned clients can still be inside create_client() and this connect path does not check shutdown_token before or after Client::connect*. Because clones share the same connection_tables, a connect that completes after Line 1010 can insert a live client after the report is built.
Possible guard points
async fn create_client(&self, addr: &CheetahString, duration: Duration) -> Option<Client<PR>> {
+ if self.shutdown_token.is_cancelled() {
+ return None;
+ }
+
if let Some(ref pool) = self.connection_pool {
if let Some(pooled_conn) = pool.get(addr) {
if pooled_conn.is_healthy() {
debug!("Reusing pooled connection to {}", addr);
return Some(pooled_conn.client().clone());
@@
match connect_result {
Ok(Ok(new_client)) => {
+ if self.shutdown_token.is_cancelled() {
+ let _ = new_client.close_with_report(duration).await;
+ return None;
+ }
+
// Connection successful - record success in circuit breaker
breaker.record_success();Please add a regression that races shutdown with an in-flight connect. As per coding guidelines, “For bug fixes, add or update a regression test when practical.”
Also applies to: 991-1016
🤖 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 `@rocketmq-remoting/src/clients/rocketmq_tokio_client.rs` around lines 723 -
736, The connect path in the time::timeout block (where
Client::connect_with_service_context and Client::connect are called) does not
check the shutdown_token before or after establishing the connection, allowing
connections to complete and be inserted into connection_tables even after
shutdown_with_report has drained it. Add a shutdown_token check before
initiating the connect operation and after it completes but before returning the
client result to prevent late connection inserts. Additionally, create a
regression test that races shutdown_with_report execution against an in-flight
create_client operation to verify that connections initiated during shutdown do
not get inserted into connection_tables after the shutdown report is built.
Source: Coding guidelines
| let mut connections = Vec::with_capacity(clients.len()); | ||
| for (addr, client) in clients { | ||
| let report = client.close_with_report(timeout).await; | ||
| connections.push(RemotingClientConnectionShutdownReport { addr, report }); | ||
| } |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Close drained connections concurrently to keep shutdown bounded.
Line 1020 applies the full timeout to each connection serially, so shutdown can take connections.len() * timeout after task-group shutdown. Drain the already-collected clients concurrently so one slow peer does not multiply broker/proxy shutdown time.
Proposed concurrent drain
- let mut connections = Vec::with_capacity(clients.len());
- for (addr, client) in clients {
- let report = client.close_with_report(timeout).await;
- connections.push(RemotingClientConnectionShutdownReport { addr, report });
- }
+ let connections = futures::future::join_all(clients.into_iter().map(|(addr, client)| async move {
+ let report = client.close_with_report(timeout).await;
+ RemotingClientConnectionShutdownReport { addr, report }
+ }))
+ .await;Please add a focused multi-connection shutdown regression. As per coding guidelines, “For behavior changes, add or update focused tests that would fail without the change.”
📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| let mut connections = Vec::with_capacity(clients.len()); | |
| for (addr, client) in clients { | |
| let report = client.close_with_report(timeout).await; | |
| connections.push(RemotingClientConnectionShutdownReport { addr, report }); | |
| } | |
| let connections = futures::future::join_all(clients.into_iter().map(|(addr, client)| async move { | |
| let report = client.close_with_report(timeout).await; | |
| RemotingClientConnectionShutdownReport { addr, report } | |
| })) | |
| .await; |
🤖 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 `@rocketmq-remoting/src/clients/rocketmq_tokio_client.rs` around lines 1018 -
1022, The connection shutdown loop is applying the full timeout sequentially to
each client, causing total shutdown time to scale linearly with the number of
connections. Replace the sequential iteration in the loop (where
client.close_with_report(timeout) is called for each client one at a time) with
concurrent execution by collecting all the close futures and then awaiting them
concurrently using futures operations like join_all or similar concurrent
combinators. Additionally, add a focused regression test that verifies multiple
connections are shutdown concurrently and confirms the total shutdown time is
bounded by a single timeout value rather than being multiplied by the connection
count.
Source: Coding guidelines
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #7545 +/- ##
==========================================
+ Coverage 69.04% 69.08% +0.04%
==========================================
Files 1234 1234
Lines 266631 267468 +837
==========================================
+ Hits 184095 184784 +689
- Misses 82536 82684 +148 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
Which Issue(s) This PR Fixes(Closes)
Brief Description
How Did You Test This Change?
Summary by CodeRabbit
New Features
Refactor