Skip to content

[ISSUE #7544]🚀Harden Tokio runtime lifecycle: service-tree ownership, shutdown reports, scheduler guardrails, and diagnostics#7545

Merged
rocketmq-rust-bot merged 1 commit into
mainfrom
feat-7544
Jun 24, 2026
Merged

[ISSUE #7544]🚀Harden Tokio runtime lifecycle: service-tree ownership, shutdown reports, scheduler guardrails, and diagnostics#7545
rocketmq-rust-bot merged 1 commit into
mainfrom
feat-7544

Conversation

@mxsm

@mxsm mxsm commented Jun 23, 2026

Copy link
Copy Markdown
Owner

Which Issue(s) This PR Fixes(Closes)

Brief Description

How Did You Test This Change?

Summary by CodeRabbit

  • New Features

    • Enhanced shutdown health reporting across broker, remoting, and proxy components to track graceful termination status.
    • Hierarchical task group support enabling better organization of background operations.
  • Refactor

    • Improved runtime initialization with explicit blocking thread configuration in broker, namesrv, and admin TUI binaries.
    • Streamlined shutdown procedures with comprehensive diagnostic reporting during component lifecycle.

… shutdown reports, scheduler guardrails, and diagnostics
@rocketmq-rust-robot rocketmq-rust-robot added the feature🚀 Suggest an idea for this project. label Jun 23, 2026
@rocketmq-rust-bot

Copy link
Copy Markdown
Collaborator

🔊@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💥.

@coderabbitai

coderabbitai Bot commented Jun 23, 2026

Copy link
Copy Markdown
Contributor

Review Change Stack

Walkthrough

This PR hardens the Tokio runtime lifecycle across the RocketMQ Rust codebase: binary entrypoints are refactored to explicit multi-thread runtime construction, TaskGroup parent-child hierarchy is propagated into schedulers, stats, controllers, remoting clients, proxy, and tieredstore components, RocketMQRuntime/ScheduledTaskManager::new() are deprecated, structured shutdown reports are introduced for remoting and broker components, and the runtime-audit script gains a scheduler-sites category.

Changes

Runtime lifecycle hardening

Layer / File(s) Summary
TaskGroup & diagnostics foundation
rocketmq-runtime/src/task_group.rs, rocketmq-runtime/src/diagnostics.rs, rocketmq-runtime/src/legacy.rs, rocketmq-runtime/src/lib.rs, rocketmq-tieredstore/src/runtime.rs
Adds TaskGroup::child_count, updates RuntimeDiagnosticsSnapshot to use typed TaskGroupId/TaskGroupLifecycleState, marks RocketMQRuntime as #[deprecated], suppresses the re-export warning, and introduces the task_group_with_parent helper in tieredstore.
Explicit Tokio runtime construction in entrypoints
rocketmq-broker/src/bin/broker_bootstrap_server.rs, rocketmq-namesrv/src/bin/namesrv_bootstrap_server.rs, rocketmq-proxy/src/bin/rocketmq-proxy-rust.rs, rocketmq-tools/.../main.rs
Replaces #[rocketmq::main]/#[tokio::main] macro entrypoints with synchronous main that builds a Tokio multi-thread runtime with ENTRYPOINT_MAX_BLOCKING_THREADS, calling an async run() via block_on.
ScheduledTaskManager legacy compatibility and scheduler/executor/service task-group hierarchy
rocketmq/src/schedule.rs, rocketmq/src/schedule/executor.rs, rocketmq/src/schedule/scheduler.rs, rocketmq/src/task/service_task.rs, rocketmq-client/src/factory/mq_client_instance.rs, rocketmq-client/src/producer/transaction_mq_produce_builder.rs, rocketmq/benches/scheduler_lifecycle_bench.rs, rocketmq-store/src/stats/broker_stats_manager.rs
Adds LEGACY_SCHEDULED_TASK_MANAGER_BOUNDARY constants and LegacyScheduledTaskManagerBoundary struct, deprecates ScheduledTaskManager::new() in favour of new_legacy_compatibility(), threads optional parent TaskGroup into ScheduledTaskManager, TaskExecutor, ExecutorPool, TaskScheduler, and ServiceManager; updates all call sites and tests.
Remoting client ServiceContext, shutdown reports, and connection lifecycle
rocketmq-remoting/src/clients/rocketmq_tokio_client.rs, rocketmq-remoting/src/clients/client.rs, rocketmq-remoting/src/clients/connection_pool.rs
Adds optional ServiceContext to RocketmqDefaultClient with new_with_service_context constructors, parents background/worker TaskGroups under it, introduces RemotingClientShutdownReport, implements shutdown_with_report that drains connections, adds ConnectionPool::clear, and wires a dedicated send-shutdown signal into Client::connect/run_send with close_with_report.
Controller, heartbeat, and RPC server task-group hierarchy
rocketmq-controller/src/controller/controller_manager.rs, rocketmq-controller/src/controller/open_raft_controller.rs, rocketmq-controller/src/heartbeat/default_broker_heartbeat_manager.rs, rocketmq-controller/src/metadata/broker.rs, rocketmq-controller/src/rpc/server.rs
Adds parent_task_group field and new_with_task_group constructors to all five controller-layer components; each derives its internal TaskGroup from the parent when provided, falling back to root creation otherwise.
Stats/statistics task-group hierarchy
rocketmq-common/src/common/statistics/statistics_manager.rs, rocketmq-common/src/common/stats/moment_stats_item.rs, rocketmq-common/src/common/stats/moment_stats_item_set.rs, rocketmq-common/src/utils/http_tiny_client.rs
Wires optional parent TaskGroup into StatisticsManager, MomentStatsItem, and MomentStatsItemSet; get_and_create_stats_item selects constructor based on parent presence; adds HttpSyncRuntimeBoundary boundary metadata struct and constants.
Broker runtime shutdown consolidation and outer API reporting
rocketmq-broker/src/broker_runtime.rs, rocketmq-broker/src/out_api/broker_outer_api.rs
Adds broker_outer_api and client_housekeeping fields to BrokerBasicServiceShutdownReport, consolidates their shutdown into shutdown_basic_service_with_report, removes duplicate direct shutdown calls from shutdown(), adds BrokerOuterAPI::shutdown_with_report, and wires ScheduledTaskManager to service context TaskGroup.
Proxy, observability, and tieredstore task-group hierarchy
rocketmq-proxy/src/grpc/service.rs, rocketmq-proxy/src/grpc/server.rs, rocketmq-observability/src/exporter/prometheus.rs, rocketmq-tieredstore/src/dispatcher/tiered_dispatcher.rs, rocketmq-tieredstore/src/service.rs
Adds parent TaskGroup support to ProxyGrpcService, ProxyGrpcServer, Prometheus exporter, DefaultTieredDispatcher, and TieredServiceSet; each conditionally creates a child task group when a parent is provided, with new _with_task_group public constructors and matching tests.
Client runtime boundary metadata
rocketmq-client/src/runtime.rs
Adds CLIENT_SHARED_FALLBACK_RUNTIME_BOUNDARY/CLIENT_SHARED_FALLBACK_RUNTIME_COMPATIBILITY constants and extends ClientSharedFallbackSnapshot with boundary, compatibility, and prefers_injected_runtime fields.
Runtime audit tooling and model tests
scripts/runtime-audit.ps1, scripts/runtime-audit-baseline.json, rocketmq-runtime/tests/runtime_model.rs
Adds scheduler-sites category to the audit script and baseline, introduces path-list exceptions for all new legacy/compat constructors, and extends runtime_model.rs with tests validating entrypoint explicit-runtime construction, RocketMQRuntime deprecation, and richer diagnostics snapshot assertions.

Estimated code review effort

🎯 5 (Critical) | ⏱️ ~120 minutes

Possibly related PRs

  • mxsm/rocketmq-rust#7524: Refactors ClientHousekeepingService to use a TaskGroup for its background loop with shutdown health reporting, which is directly consumed by the broker shutdown consolidation in this PR.
  • mxsm/rocketmq-rust#7533: Introduces the runtime-audit.ps1 and runtime-audit-baseline.json infrastructure that this PR extends with the new scheduler-sites category and additional exception rules.

Suggested labels

refactor♻️, enhancement⚡️, approved

Suggested reviewers

  • rocketmq-rust-bot
  • SpaceXCN
  • TeslaRustor
  • rocketmq-rust-robot

Poem

🐰 A rabbit hops through task-group trees,
Each parent spawns a child with ease.
Old new() is deprecated now—
new_legacy_compat() takes the bow.
Shutdown reports flow leaf to root,
The Tokio runtime stands resolute! 🚀

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 inconclusive)

Check name Status Explanation Resolution
Linked Issues check ❓ Inconclusive The linked issue #7544 contains only a feature title with no description, acceptance criteria, or specific requirements, making it impossible to validate whether code changes meet defined objectives. Provide detailed requirements, acceptance criteria, and specific objectives in issue #7544 to enable proper compliance assessment.
✅ Passed checks (4 passed)
Check name Status Explanation
Description Check ✅ Passed Check skipped - CodeRabbit’s high-level summary is enabled.
Title check ✅ Passed The PR title clearly summarizes the main changes: hardening Tokio runtime lifecycle across service-tree ownership, shutdown reports, scheduler guardrails, and diagnostics, which aligns with the extensive refactoring throughout the codebase.
Out of Scope Changes check ✅ Passed All changes consistently implement task-group parenting, shutdown reporting, legacy-compatibility boundaries, and runtime lifecycle hardening across 31 files without unrelated modifications.
Docstring Coverage ✅ Passed Docstring coverage is 100.00% 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 docstrings
  • Create stacked PR
  • Commit on current branch
🧪 Generate unit tests (beta)
  • Create PR with unit tests
  • Commit unit tests in branch feat-7544

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.

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

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 win

Keep the parent TaskGroup for executor reuse after shutdown.

Line 111 only stores the derived child group. Once shutdown_task_group() takes it, the next execution falls back to new_executor_task_group() and creates a root group, so a parented executor loses service-tree ownership after shutdown_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 a new_with_task_group executor, 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 lift

Propagate parent_task_group into RaftController construction.

Line 447-451 passes the parent into DefaultBrokerHeartbeatManager, but Line 460-463 still builds RaftController through new_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 win

Add a focused test for the parented MomentStatsItem path.

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 contains rocketmq-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 win

Exercise 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 new MomentStatsItem::new_with_task_group branch at Lines 156-161 can regress unnoticed. Extend the test to create/init one item and assert the parent shutdown report includes its rocketmq-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 win

Assert 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 win

Consider 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

📥 Commits

Reviewing files that changed from the base of the PR and between c138ddd and bffd177.

📒 Files selected for processing (40)
  • rocketmq-broker/src/bin/broker_bootstrap_server.rs
  • rocketmq-broker/src/broker_runtime.rs
  • rocketmq-broker/src/out_api/broker_outer_api.rs
  • rocketmq-client/src/factory/mq_client_instance.rs
  • rocketmq-client/src/producer/transaction_mq_produce_builder.rs
  • rocketmq-client/src/runtime.rs
  • rocketmq-common/src/common/statistics/statistics_manager.rs
  • rocketmq-common/src/common/stats/moment_stats_item.rs
  • rocketmq-common/src/common/stats/moment_stats_item_set.rs
  • rocketmq-common/src/utils/http_tiny_client.rs
  • rocketmq-controller/src/controller/controller_manager.rs
  • rocketmq-controller/src/controller/open_raft_controller.rs
  • rocketmq-controller/src/heartbeat/default_broker_heartbeat_manager.rs
  • rocketmq-controller/src/metadata/broker.rs
  • rocketmq-controller/src/rpc/server.rs
  • rocketmq-namesrv/src/bin/namesrv_bootstrap_server.rs
  • rocketmq-observability/src/exporter/prometheus.rs
  • rocketmq-proxy/src/bin/rocketmq-proxy-rust.rs
  • rocketmq-proxy/src/grpc/server.rs
  • rocketmq-proxy/src/grpc/service.rs
  • rocketmq-remoting/src/clients/client.rs
  • rocketmq-remoting/src/clients/connection_pool.rs
  • rocketmq-remoting/src/clients/rocketmq_tokio_client.rs
  • rocketmq-runtime/src/diagnostics.rs
  • rocketmq-runtime/src/legacy.rs
  • rocketmq-runtime/src/lib.rs
  • rocketmq-runtime/src/task_group.rs
  • rocketmq-runtime/tests/runtime_model.rs
  • rocketmq-store/src/stats/broker_stats_manager.rs
  • rocketmq-tieredstore/src/dispatcher/tiered_dispatcher.rs
  • rocketmq-tieredstore/src/runtime.rs
  • rocketmq-tieredstore/src/service.rs
  • rocketmq-tools/rocketmq-admin/rocketmq-admin-tui/src/main.rs
  • rocketmq/benches/scheduler_lifecycle_bench.rs
  • rocketmq/src/schedule.rs
  • rocketmq/src/schedule/executor.rs
  • rocketmq/src/schedule/scheduler.rs
  • rocketmq/src/task/service_task.rs
  • scripts/runtime-audit-baseline.json
  • scripts/runtime-audit.ps1

Comment on lines +674 to 678
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()));

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

🩺 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/clients

Repository: 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.rs

Repository: 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.rs

Repository: 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.rs

Repository: 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.

Comment on lines +1158 to +1166
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:?}"),
)
};

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

🎯 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/clients

Repository: 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.rs

Repository: 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.rs

Repository: 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.rs

Repository: 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.rs

Repository: 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 -20

Repository: 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.rs

Repository: 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.rs

Repository: 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.rs

Repository: 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.rs

Repository: 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 -20

Repository: mxsm/rocketmq-rust

Length of output: 1260


🏁 Script executed:

# Check all timeout constants defined
rg -n "SHUTDOWN_TIMEOUT" rocketmq-broker/src/broker_runtime.rs

Repository: 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.rs

Repository: 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.rs

Repository: 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 -50

Repository: 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.rs

Repository: 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 10

Repository: 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 rs

Repository: mxsm/rocketmq-rust

Length of output: 394


🏁 Script executed:

# Read the ShutdownReport struct definition
cat rocketmq-runtime/src/shutdown_report.rs

Repository: 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.

Comment on lines +475 to +481
let task_group = parent_task_group.map_or_else(
|| {
TaskGroup::root(
"rocketmq-proxy.grpc.housekeeping",
RuntimeHandle::new(tokio::runtime::Handle::current()),
)
},

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

🩺 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/src

Repository: 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.rs

Repository: 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 -30

Repository: 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.rs

Repository: 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' -A5

Repository: 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 -20

Repository: 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 -A5

Repository: 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 -50

Repository: 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 -30

Repository: 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.

Comment on lines +723 to +736
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
}

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

🩺 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

Comment on lines +1018 to +1022
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 });
}

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

🩺 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.

Suggested change
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

codecov Bot commented Jun 23, 2026

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 80.78078% with 192 lines in your changes missing coverage. Please review.
✅ Project coverage is 69.08%. Comparing base (c138ddd) to head (bffd177).

Files with missing lines Patch % Lines
...etmq-remoting/src/clients/rocketmq_tokio_client.rs 79.61% 32 Missing ⚠️
rocketmq-controller/src/rpc/server.rs 0.00% 21 Missing ⚠️
...-controller/src/controller/open_raft_controller.rs 37.93% 18 Missing ⚠️
...ketmq-common/src/common/stats/moment_stats_item.rs 38.09% 13 Missing ⚠️
rocketmq-proxy/src/bin/rocketmq-proxy-rust.rs 0.00% 11 Missing ⚠️
...mq-tieredstore/src/dispatcher/tiered_dispatcher.rs 86.11% 10 Missing ⚠️
rocketmq-broker/src/bin/broker_bootstrap_server.rs 0.00% 9 Missing ⚠️
...cketmq-namesrv/src/bin/namesrv_bootstrap_server.rs 0.00% 9 Missing ⚠️
...common/src/common/statistics/statistics_manager.rs 81.39% 8 Missing ⚠️
...mq-controller/src/controller/controller_manager.rs 60.00% 8 Missing ⚠️
... and 14 more
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.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
  • 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

@rocketmq-rust-bot rocketmq-rust-bot merged commit a099b99 into main Jun 24, 2026
30 of 31 checks passed
@rocketmq-rust-bot rocketmq-rust-bot added approved PR has approved and removed ready to review waiting-review waiting review this PR labels Jun 24, 2026
@mxsm mxsm deleted the feat-7544 branch June 24, 2026 01:00
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

AI review first Ai review pr first approved PR has approved auto merge feature🚀 Suggest an idea for this project.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[Feature🚀] Harden Tokio runtime lifecycle: service-tree ownership, shutdown reports, scheduler guardrails, and diagnostics

3 participants