refactor!: add TaskTracker for structured task lifecycle management - #458
refactor!: add TaskTracker for structured task lifecycle management#458jlucaso1 wants to merge 4 commits into
Conversation
|
Caution Review failedPull request was closed or merged during review 📝 WalkthroughWalkthroughAdds a TaskTracker type and makes background tasks return/produce AbortHandles; many fire-and-forget runtime spawns are replaced with TaskTracker-managed spawns, the client gains separate Changes
Sequence DiagramsequenceDiagram
participant App as Application
participant Client as Client
participant Tracker as TaskTracker
participant Runtime as Runtime
participant Task as Background Task
rect rgba(100,150,200,0.5)
Note over App,Runtime: Old Flow — fire-and-forget
App->>Runtime: spawn(future).detach()
Runtime->>Task: run task independently
end
rect rgba(150,100,200,0.5)
Note over App,Tracker: New Flow — TaskTracker-managed
App->>Tracker: spawn(&runtime, future)
Tracker->>Runtime: spawn(future)
Runtime-->>Tracker: AbortHandle
Tracker->>Tracker: track(handle)
Runtime->>Task: start task
App->>Tracker: abort_all() on teardown
Tracker->>Task: abort via handles
end
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 3✅ Passed checks (3 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 |
|
@claude review |
|
Claude finished @jlucaso1's task in 1m 38s —— View job Code Review: TaskTracker for structured task lifecycle management
Overall this is a clean, well-scoped refactor. The two-tier tracker model (connection vs client lifetime) is the right design, the Handle vec growth in
|
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: a9c9ed22c6
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| rt: &dyn Runtime, | ||
| future: Pin<Box<dyn Future<Output = ()> + Send + 'static>>, | ||
| ) { | ||
| self.track(rt.spawn(future)); |
There was a problem hiding this comment.
TaskTracker::spawn stores every returned AbortHandle, but nothing removes handles when the task finishes, so completed tasks remain in handles until disconnect/drop. In this commit many high-frequency paths now use connection_tasks.spawn (including per-stanza processing), so a long-lived connection will accumulate an unbounded number of finished JoinHandles, increasing memory usage over time and making abort_all() progressively more expensive. Please remove tracked handles on task completion (or avoid tracking short-lived per-message tasks).
Useful? React with 👍 / 👎.
Bencher Report![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 4
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@src/bot.rs`:
- Around line 653-655: The background saver started by
persistence_manager.run_background_saver (bg_saver_handle) is currently tracked
by client_tasks and will be aborted on Client drop (TaskTracker::Drop →
abort_all), risking loss of dirty state; modify the shutdown flow to perform an
explicit flush and graceful shutdown: add an async shutdown/stop method on Bot
or Client that calls persistence_manager.flush().await and then awaits the
bg_saver_handle to finish (or signals it to exit cleanly) before dropping
TaskTracker, and update places that previously dropped Client to call this
shutdown method; alternatively, if you cannot add an async Drop, document the
limitation and provide a public async shutdown API that users must call to
guarantee persistence.
In `@src/history_sync.rs`:
- Around line 183-188: The parse task started with connection_tasks.spawn
(Box::pin(async move { blocking_fut.await; })) is tied to the connection
lifecycle and can be aborted by cleanup_connection_state() via
connection_tasks.abort_all(), causing lost work after an acknowledgement; move
the parsing work out of the connection-scoped task manager or persist resume
state before sending the receipt. Concretely, either (a) submit blocking_fut to
a global/non-connection-scoped executor/runtime (use the global runtime or a
dedicated parser_tasks manager) instead of connection_tasks.spawn so it survives
disconnects, or (b) persist a parse job record (including blob id and progress)
before sending the receipt and have cleanup_connection_state()/connection
startup reconcile and resume parsing on reconnect; update calls around
connection_tasks.spawn, cleanup_connection_state, and where the receipt is sent
to reflect the chosen approach.
In `@src/pair.rs`:
- Around line 294-299: The spawned task using client.connection_tasks.spawn
currently runs client_for_unified.send_unified_session() which may be aborted by
the 515 stream error; change this so send_unified_session() is guaranteed to
complete — either await client_for_unified.send_unified_session().await directly
here (instead of spawning), or move the spawn into a long-lived task registry
(e.g., client.client_tasks) so the send_unified_session() future isn't cancelled
by connection_tasks shutdown; update the call sites around
client.connection_tasks.spawn, client.runtime and
client_for_unified.send_unified_session() accordingly to ensure the future is
tracked to completion.
In `@wacore/src/runtime.rs`:
- Around line 115-167: TaskTracker currently accumulates AbortHandle entries
forever; update it to prune handles for tasks that have already completed to
avoid unbounded growth. Modify TaskTracker::track to prune stale handles before
pushing the new AbortHandle (e.g., call self.handles.lock()... and Vec::retain
to keep only active handles), also do the same trim in len() (and optionally at
the start of abort_all()) so reported length excludes completed tasks; reference
the TaskTracker struct and the methods track, len, and abort_all when making
these changes (use AbortHandle's appropriate status/API such as is_aborted() or
equivalent to detect completed handles).
🪄 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: Organization UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: 7c164903-1b26-427f-9f8c-f8ac0203f764
📒 Files selected for processing (13)
src/bot.rssrc/client.rssrc/handlers/ib.rssrc/handlers/notification.rssrc/history_sync.rssrc/keepalive.rssrc/message.rssrc/pair.rssrc/pdo.rssrc/receipt.rssrc/store/persistence_manager.rswacore/src/runtime.rswacore/src/store/persistence.rs
| let bg_saver_handle = persistence_manager | ||
| .clone() | ||
| .run_background_saver(runtime.clone(), std::time::Duration::from_secs(30)); |
There was a problem hiding this comment.
Potential data loss: background saver aborted without final flush on Client drop.
Tracking the background saver via client_tasks means it will be aborted when Client is dropped (via TaskTracker::Drop → abort_all()). If there's dirty state that hasn't been flushed yet, it will be lost.
The persistence module comment (lines 136-140 in wacore/src/store/persistence.rs) acknowledges this limitation but doesn't address it. Consider either:
- Adding an explicit
flush()call beforeClientdrop (requires adding a shutdown method) - Calling
persistence_manager.flush().awaitinBot::dropor similar cleanup path - Documenting this as a known limitation with guidance for users who need guaranteed persistence
This is a behavioral change from the previous fire-and-forget approach where the saver could complete naturally.
Also applies to: 690-691
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@src/bot.rs` around lines 653 - 655, The background saver started by
persistence_manager.run_background_saver (bg_saver_handle) is currently tracked
by client_tasks and will be aborted on Client drop (TaskTracker::Drop →
abort_all), risking loss of dirty state; modify the shutdown flow to perform an
explicit flush and graceful shutdown: add an async shutdown/stop method on Bot
or Client that calls persistence_manager.flush().await and then awaits the
bg_saver_handle to finish (or signals it to exit cleanly) before dropping
TaskTracker, and update places that previously dropped Client to call this
shutdown method; alternatively, if you cannot add an async Drop, document the
limitation and provide a public async shutdown API that users must call to
guarantee persistence.
| /// Tracks spawned tasks and aborts them all on demand or when dropped. | ||
| /// | ||
| /// Two intended scopes: | ||
| /// - **Connection-scoped**: aborted on disconnect via `abort_all()`. | ||
| /// - **Client-scoped**: aborted when the owning struct is dropped. | ||
| pub struct TaskTracker { | ||
| handles: std::sync::Mutex<Vec<AbortHandle>>, | ||
| } | ||
|
|
||
| impl TaskTracker { | ||
| pub fn new() -> Self { | ||
| Self { | ||
| handles: std::sync::Mutex::new(Vec::new()), | ||
| } | ||
| } | ||
|
|
||
| /// Store an existing `AbortHandle` so it will be aborted on `abort_all()` / drop. | ||
| pub fn track(&self, handle: AbortHandle) { | ||
| self.handles | ||
| .lock() | ||
| .unwrap_or_else(|p| p.into_inner()) | ||
| .push(handle); | ||
| } | ||
|
|
||
| /// Abort every tracked task and clear the list. | ||
| pub fn abort_all(&self) { | ||
| let handles = std::mem::take(&mut *self.handles.lock().unwrap_or_else(|p| p.into_inner())); | ||
| for handle in handles { | ||
| handle.abort(); | ||
| } | ||
| } | ||
|
|
||
| /// Number of tracked handles (includes already-completed tasks). | ||
| pub fn len(&self) -> usize { | ||
| self.handles.lock().unwrap_or_else(|p| p.into_inner()).len() | ||
| } | ||
|
|
||
| pub fn is_empty(&self) -> bool { | ||
| self.len() == 0 | ||
| } | ||
| } | ||
|
|
||
| impl Default for TaskTracker { | ||
| fn default() -> Self { | ||
| Self::new() | ||
| } | ||
| } | ||
|
|
||
| impl Drop for TaskTracker { | ||
| fn drop(&mut self) { | ||
| self.abort_all(); | ||
| } | ||
| } |
There was a problem hiding this comment.
🧹 Nitpick | 🔵 Trivial
Consider pruning completed handles to prevent unbounded memory growth.
TaskTracker accumulates handles indefinitely since there's no mechanism to remove handles for tasks that have already completed. The len() documentation acknowledges this: "includes already-completed tasks."
For connection_tasks this is fine since abort_all() is called on disconnect. However, for long-lived client_tasks, if many short-lived tasks are tracked over hours/days, the Vec<AbortHandle> could grow unbounded.
Given the current usage pattern (only 3 client-scoped tasks per PR summary), this is unlikely to be a practical issue, but worth documenting or addressing in a future iteration.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@wacore/src/runtime.rs` around lines 115 - 167, TaskTracker currently
accumulates AbortHandle entries forever; update it to prune handles for tasks
that have already completed to avoid unbounded growth. Modify TaskTracker::track
to prune stale handles before pushing the new AbortHandle (e.g., call
self.handles.lock()... and Vec::retain to keep only active handles), also do the
same trim in len() (and optionally at the start of abort_all()) so reported
length excludes completed tasks; reference the TaskTracker struct and the
methods track, len, and abort_all when making these changes (use AbortHandle's
appropriate status/API such as is_aborted() or equivalent to detect completed
handles).
There was a problem hiding this comment.
Actionable comments posted: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (6)
src/receipt.rs (1)
70-84:⚠️ Potential issue | 🟠 MajorDo not make retry handling abortable until it is cancellation-safe.
connection_tasks.abort_all()can now cut this future at any.await, buthandle_retry_receipt()insrc/retry.rswrites retry-dedupe state and mutates Signal/session persistence in multiple steps before it finishes. That means a disconnect can record a retry as “handled” without ever completing the resend or the cleanup. Keep this path detached for now, or refactorhandle_retry_receipt()into an idempotent/transactional flow before tracking it.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/receipt.rs` around lines 70 - 84, The retry handling task is currently spawned via connection_tasks.spawn and thus can be aborted by connection_tasks.abort_all(), which may cut off handle_retry_receipt (in src/retry.rs) mid-mutation and leave retry/session state inconsistent; instead run the retry work in a detached/non-abortable task (e.g. spawn on the runtime directly rather than using connection_tasks.spawn) so handle_retry_receipt(&receipt, &node_clone) runs to completion, or alternatively refactor handle_retry_receipt into an idempotent/transactional flow before reintroducing abortable execution; update the spawn call that references connection_tasks.spawn and client_clone.handle_retry_receipt to use a detached runtime spawn (or equivalent non-abortable mechanism) and/or make handle_retry_receipt cancellation-safe.src/pdo.rs (1)
441-460:⚠️ Potential issue | 🟠 MajorCancellation now leaves stale PDO dedupe entries behind.
This future is now abortable on disconnect, but
send_pdo_placeholder_resend_request()inserts intopdo_pending_requestsbefore it reaches the send/error-cleanup path. If teardown cancels after that insert, the stale cache key suppresses the next PDO attempt for the same message on reconnect. Either make that pending-entry lifecycle cancellation-safe or keep this request detached until it is on the wire.src/pair.rs (1)
67-105:⚠️ Potential issue | 🟡 MinorClear stale QR cancellation state on disconnect.
With the QR rotator now living in
connection_tasks,cleanup_connection_state()aborts it before transport teardown. That dropsstop_rx, butpairing_cancellation_txstays populated here until some laterpair-successor newpair-devicehappens, so pairing state can survive a reconnect with no live task behind it. Clear the option during connection cleanup, or before installing a new sender, to keep reconnect state consistent.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/pair.rs` around lines 67 - 105, cleanup_connection_state() currently aborts the QR rotator task but leaves client.pairing_cancellation_tx populated, allowing stale pairing state to persist across reconnects; fix by clearing that Option by locking client.pairing_cancellation_tx and setting it to None during cleanup_connection_state() (and/or just before installing a new sender where stop_tx is assigned) so there is never a hanging sender handle left behind—update the code paths that install stop_tx (where *client.pairing_cancellation_tx.lock().await = Some(stop_tx) is called) and the cleanup_connection_state() function to ensure they explicitly set the Option to None when aborting or replacing the rotator.src/message.rs (2)
437-450:⚠️ Potential issue | 🟠 MajorDon't connection-track per-
enccustom handlers.This branch runs once per incoming
encnode. Putting it inconnection_tasksreintroduces the same unbounded per-stanza handle growth the PR already had to back out elsewhere, andabort_all()can now drop already-received custom messages mid-handler during reconnect. These handlers should stay detached, or move to a tracker that removes completed tasks.Suggested direction
- self.connection_tasks.spawn( - &*self.runtime, - Box::pin(async move { - if let Err(e) = handler_clone - .handle(client_clone, &enc_node_clone, &info_arc) - .await - { - log::warn!( - "Custom handler for enc type '{}' failed: {e:?}", - enc_type_owned - ); - } - }), - ); + self.runtime + .spawn(Box::pin(async move { + if let Err(e) = handler_clone + .handle(client_clone, &enc_node_clone, &info_arc) + .await + { + log::warn!( + "Custom handler for enc type '{}' failed: {e:?}", + enc_type_owned + ); + } + })) + .detach();🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/message.rs` around lines 437 - 450, The current code spawns per-enc custom handler tasks into self.connection_tasks (via connection_tasks.spawn with runtime and Box::pin async) which causes unbounded per-stanza task tracking and allows abort_all() to cancel handlers; change this to run handlers detached instead of tracking them: replace the use of self.connection_tasks.spawn(...) that wraps handler_clone.handle(client_clone, &enc_node_clone, &info_arc).await with a detached spawn (e.g., tokio::spawn) or a dedicated tracker that removes handles when they complete, so custom handlers are not kept in connection_tasks and are not aborted by abort_all(); keep references to handler_clone, client_clone, enc_node_clone, and info_arc as before.
195-250:⚠️ Potential issue | 🟠 MajorKeep retry receipts out of
connection_tasks.Line 201 increments the retry counter before any send occurs. Once this best-effort receipt is tracked,
cleanup_connection_state().abort_all()can cancel it after the count bump but beforesend_retry_receipt(), so reconnects burn retry numbers without ever transmitting a retry. This is also a high-volume error path, so retaining oneAbortHandleper failed message will growconnection_tasksacross long recovery sessions.Suggested direction
- self.connection_tasks.spawn(&*self.runtime, Box::pin(async move { + self.runtime.spawn(Box::pin(async move { let cache_key = client .make_retry_cache_key(&info.source.chat, &info.id, &info.source.sender) .await; @@ if retry_count == 1 { client.spawn_pdo_request(&info); } - })); + })) + .detach();🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/message.rs` around lines 195 - 250, The retry-count bump and send logic should not be scheduled on connection_tasks because cleanup_connection_state().abort_all() can cancel the task after increment_retry_count() but before send_retry_receipt(), burning retry counters and leaking AbortHandles; move the entire async block currently spawned via self.connection_tasks.spawn(...) to a non-abortable executor (e.g. spawn on the runtime directly with self.runtime.spawn or tokio::spawn) or otherwise run increment/send on a short-lived non-connection task so increment_retry_count(), send_retry_receipt(), and the conditional spawn_pdo_request()/spawn_pdo_request_with_options() execute even if connection_tasks are aborted; preserve existing calls to increment_retry_count, send_retry_receipt, spawn_pdo_request, spawn_pdo_request_with_options, and the HIGH_RETRY_COUNT_THRESHOLD logic when relocating the block.src/client.rs (1)
1900-1946: 🧹 Nitpick | 🔵 TrivialNote on TaskTracker memory management.
The
TaskTrackeraccumulates handles for completed tasks untilabort_all()is called. This is acceptable forconnection_taskssinceabort_all()clears the Vec on disconnect. However, for very long-lived connections with frequent spawns, handles would accumulate.The current spawn sites in
handle_successare bounded per-connection (background init, app state sync), and push name updates are infrequent. The intentionaldetach()for per-stanza processing (line 1311) correctly avoids unbounded growth for high-frequency spawns.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/client.rs` around lines 1900 - 1946, The TaskTracker used by connection_tasks can accumulate completed task handles for long-lived connections; make high-frequency/fire-and-forget spawns call detach() instead of storing their handle in connection_tasks and reserve connection_tasks only for bounded-per-connection work (background init, app state sync) that you already call abort_all() for on disconnect; update the spawn sites in handle_success (replace storing handles for per-stanza or frequent tasks with detach()) and, if there are other frequent spawn sites, likewise switch them to detach() or call a pruning method after spawn so completed handles are removed.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@wacore/src/runtime.rs`:
- Around line 132-137: track() currently only appends AbortHandle to
self.handles and never removes completed tasks, causing unbounded growth; change
track() (and the analogous code around the handles usage at the other location
referenced) to prune completed handles before/after pushing a new one: lock
self.handles, call .retain(...) to keep only active handles (for example, if you
store JoinHandle use .is_finished(), or if you must use AbortHandle wrap it so
you can detect completion or remove the handle when the spawned task finishes
via a cleanup callback), then push the new handle; ensure the same pruning logic
is applied where handles are modified so dead handles are not accumulated.
---
Outside diff comments:
In `@src/client.rs`:
- Around line 1900-1946: The TaskTracker used by connection_tasks can accumulate
completed task handles for long-lived connections; make
high-frequency/fire-and-forget spawns call detach() instead of storing their
handle in connection_tasks and reserve connection_tasks only for
bounded-per-connection work (background init, app state sync) that you already
call abort_all() for on disconnect; update the spawn sites in handle_success
(replace storing handles for per-stanza or frequent tasks with detach()) and, if
there are other frequent spawn sites, likewise switch them to detach() or call a
pruning method after spawn so completed handles are removed.
In `@src/message.rs`:
- Around line 437-450: The current code spawns per-enc custom handler tasks into
self.connection_tasks (via connection_tasks.spawn with runtime and Box::pin
async) which causes unbounded per-stanza task tracking and allows abort_all() to
cancel handlers; change this to run handlers detached instead of tracking them:
replace the use of self.connection_tasks.spawn(...) that wraps
handler_clone.handle(client_clone, &enc_node_clone, &info_arc).await with a
detached spawn (e.g., tokio::spawn) or a dedicated tracker that removes handles
when they complete, so custom handlers are not kept in connection_tasks and are
not aborted by abort_all(); keep references to handler_clone, client_clone,
enc_node_clone, and info_arc as before.
- Around line 195-250: The retry-count bump and send logic should not be
scheduled on connection_tasks because cleanup_connection_state().abort_all() can
cancel the task after increment_retry_count() but before send_retry_receipt(),
burning retry counters and leaking AbortHandles; move the entire async block
currently spawned via self.connection_tasks.spawn(...) to a non-abortable
executor (e.g. spawn on the runtime directly with self.runtime.spawn or
tokio::spawn) or otherwise run increment/send on a short-lived non-connection
task so increment_retry_count(), send_retry_receipt(), and the conditional
spawn_pdo_request()/spawn_pdo_request_with_options() execute even if
connection_tasks are aborted; preserve existing calls to increment_retry_count,
send_retry_receipt, spawn_pdo_request, spawn_pdo_request_with_options, and the
HIGH_RETRY_COUNT_THRESHOLD logic when relocating the block.
In `@src/pair.rs`:
- Around line 67-105: cleanup_connection_state() currently aborts the QR rotator
task but leaves client.pairing_cancellation_tx populated, allowing stale pairing
state to persist across reconnects; fix by clearing that Option by locking
client.pairing_cancellation_tx and setting it to None during
cleanup_connection_state() (and/or just before installing a new sender where
stop_tx is assigned) so there is never a hanging sender handle left
behind—update the code paths that install stop_tx (where
*client.pairing_cancellation_tx.lock().await = Some(stop_tx) is called) and the
cleanup_connection_state() function to ensure they explicitly set the Option to
None when aborting or replacing the rotator.
In `@src/receipt.rs`:
- Around line 70-84: The retry handling task is currently spawned via
connection_tasks.spawn and thus can be aborted by connection_tasks.abort_all(),
which may cut off handle_retry_receipt (in src/retry.rs) mid-mutation and leave
retry/session state inconsistent; instead run the retry work in a
detached/non-abortable task (e.g. spawn on the runtime directly rather than
using connection_tasks.spawn) so handle_retry_receipt(&receipt, &node_clone)
runs to completion, or alternatively refactor handle_retry_receipt into an
idempotent/transactional flow before reintroducing abortable execution; update
the spawn call that references connection_tasks.spawn and
client_clone.handle_retry_receipt to use a detached runtime spawn (or equivalent
non-abortable mechanism) and/or make handle_retry_receipt cancellation-safe.
🪄 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: Organization UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: 401b99cc-5225-43d6-9aaf-1382b172d293
📒 Files selected for processing (13)
src/bot.rssrc/client.rssrc/handlers/ib.rssrc/handlers/notification.rssrc/history_sync.rssrc/keepalive.rssrc/message.rssrc/pair.rssrc/pdo.rssrc/receipt.rssrc/store/persistence_manager.rswacore/src/runtime.rswacore/src/store/persistence.rs
| pub fn track(&self, handle: AbortHandle) { | ||
| self.handles | ||
| .lock() | ||
| .unwrap_or_else(|p| p.into_inner()) | ||
| .push(handle); | ||
| } |
There was a problem hiding this comment.
Completed tasks accumulate forever in TaskTracker.
track() only appends, and nothing removes a handle when its task finishes. With the new call sites in src/receipt.rs, src/pdo.rs, src/history_sync.rs, and the periodic cleanup in src/keepalive.rs, a long-lived connection will keep every dead handle until disconnect/drop. This needs completion pruning, or the tracker should be limited to truly long-lived loops.
Also applies to: 147-149
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@wacore/src/runtime.rs` around lines 132 - 137, track() currently only appends
AbortHandle to self.handles and never removes completed tasks, causing unbounded
growth; change track() (and the analogous code around the handles usage at the
other location referenced) to prune completed handles before/after pushing a new
one: lock self.handles, call .retain(...) to keep only active handles (for
example, if you store JoinHandle use .is_finished(), or if you must use
AbortHandle wrap it so you can detect completion or remove the handle when the
spawned task finishes via a cleanup callback), then push the new handle; ensure
the same pruning logic is applied where handles are modified so dead handles are
not accumulated.
| pub fn abort_all(&self) { | ||
| let handles = std::mem::take(&mut *self.handles.lock().unwrap_or_else(|p| p.into_inner())); | ||
| for handle in handles { | ||
| handle.abort(); | ||
| } | ||
| } |
There was a problem hiding this comment.
Close the tracker before aborting the current batch.
Line 141 only drains the current Vec. A concurrent track()/spawn() right after that lands in the new empty Vec and survives the same disconnect that was supposed to cancel all connection-scoped work. For connection_tasks, that reintroduces the stale-task-on-reconnect race this PR is trying to remove. Add a closed/generation state so post-abort_all() handles are rejected or aborted until the next connection is initialized.
|
@claude review |
|
Claude finished @jlucaso1's task in 2m 16s —— View job Code Review: TaskTracker for structured task lifecycle management
Good refactor overall. The three-commit evolution tells a clear story: initial implementation, then backing out per-stanza tracking (unbounded growth), then removing forceful Key observation for other reviewersMost of the CodeRabbit/Codex comments about cancellation-safety ( What the TaskTracker actually provides nowPer commit
Remaining concerns1. Unbounded handle growth in Since For a long-lived client processing many messages, this could grow to tens of thousands of entries. Each The second commit (
These are all fire-and-forget tasks that don't benefit from abort tracking since 2. With 3.
What looks good
VerdictThe core design is solid. The main actionable item is reverting the high-frequency fire-and-forget paths back to |
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@src/client.rs`:
- Around line 732-737: The spawned long-lived task captures a strong Arc<Client>
(via cleanup_arc) which prevents Client's Drop from ever running and leaks the
loop; change the task to capture a Weak<Client> (downgrading the Arc before
spawn) or extract only the minimal non-Client state required by
device_registry_cleanup_loop and move that into the closure so it does not hold
the Client; alternatively add an explicit shutdown path that calls
client_tasks.abort() (or otherwise stops TaskTracker) before the final
Arc<Client> is dropped so the closure cannot keep the Client alive; update
references around arc.client_tasks.spawn, arc.runtime, and
cleanup_arc.device_registry_cleanup_loop() accordingly.
- Around line 1114-1118: connection_tasks is never reset across reconnects so
TaskTracker accumulates AbortHandles and completed counts, leaking handles and
stale work into new sessions; in cleanup_connection_state() replace or swap the
existing connection_tasks TaskTracker with a fresh TaskTracker for each new
connection (or, if you prefer non-aborting teardown, implement a prune step that
removes completed handles and aborts or marks remaining ones as stale) and
update memory_diagnostics().connection_tasks to reflect the new tracker; ensure
all connection-scoped spawns reference the new tracker and still check
connection_generation and is_shutting_down() so no work from the old tracker
runs in the new session.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
| // Tasks detect staleness cooperatively via connection_generation, | ||
| // is_shutting_down(), and this notifier. Forceful abort is intentionally | ||
| // NOT used here — tasks like the post-login sequence must survive 515 | ||
| // reconnect cycles to dispatch Connected. TaskTracker::Drop handles | ||
| // final cleanup when the Client is dropped. |
There was a problem hiding this comment.
connection_tasks now survives every reconnect, which leaks handles and can leak stale work into the next session.
cleanup_connection_state() no longer aborts or replaces the tracker, but TaskTracker only accumulates AbortHandles and len() counts completed ones too. So the tracker grows monotonically across reconnects, memory_diagnostics().connection_tasks becomes misleading, and any connection-scoped spawn that misses a generation/shutdown guard can still run on the new connection. Please reset/swap the tracker per connection, or add a non-aborting prune step once teardown begins.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@src/client.rs` around lines 1114 - 1118, connection_tasks is never reset
across reconnects so TaskTracker accumulates AbortHandles and completed counts,
leaking handles and stale work into new sessions; in cleanup_connection_state()
replace or swap the existing connection_tasks TaskTracker with a fresh
TaskTracker for each new connection (or, if you prefer non-aborting teardown,
implement a prune step that removes completed handles and aborts or marks
remaining ones as stale) and update memory_diagnostics().connection_tasks to
reflect the new tracker; ensure all connection-scoped spawns reference the new
tracker and still check connection_generation and is_shutting_down() so no work
from the old tracker runs in the new session.
jlucaso1
changed the title
Replace ad-hoc .detach() calls with tracked task spawning so connection- scoped tasks are aborted on disconnect and client-scoped tasks are aborted when Client drops. - Add TaskTracker to wacore::runtime (Vec<AbortHandle> + abort_all()) - Add #[must_use] to AbortHandle for compile-time enforcement - Add connection_tasks / client_tasks fields to Client - Abort connection_tasks in cleanup_connection_state() - Invalidate message_queues on disconnect (bug fix: stale workers) - Migrate 21 detach sites to tracked; keep 14 as intentional fire-and-forget - run_background_saver now returns AbortHandle instead of detaching
…ncy) Per-stanza processing is fire-and-forget — it runs for every non-inline stanza and would accumulate unbounded handles in connection_tasks during long connections. The message loop itself exits on disconnect, so individual stanza tasks don't need abort tracking.
Tasks already detect staleness cooperatively via connection_generation, is_shutting_down(), and shutdown_notifier. Forceful abort via abort_all() killed the post-login task before it could dispatch Connected — the mock server sends 515 immediately after auth on post-pairing reconnect. TaskTracker now serves three purposes: - #[must_use] compile-time enforcement at every spawn site - Diagnostics (connection_tasks.len() / client_tasks.len()) - Client drop cleanup (TaskTracker::Drop aborts everything)
Prevents unbounded AbortHandle accumulation across reconnect cycles. By the time connect() runs, the previous connection's cooperative tasks have already received shutdown_notifier and the generation was bumped.
8264b32 to
8536c78
Compare
|
Superseded by a minimal version — the TaskTracker approach was overengineered. The cooperative cancellation via connection_generation/shutdown_notifier already handles task lifecycle correctly. |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 8536c78e93
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| self.connection_tasks.spawn( | ||
| &*self.runtime, | ||
| Box::pin(async move { |
There was a problem hiding this comment.
handle_incoming_message now schedules custom EncHandler::handle work on connection_tasks, but connect() forcibly cancels that tracker via connection_tasks.abort_all() at the beginning of each reconnect cycle. Because this branch immediately continues after spawning the handler, a reconnect during handler execution can silently drop custom message processing with no retry path. This is user-visible for any handler that performs awaited work (DB/network) and runs longer than the reconnect gap.
Useful? React with 👍 / 👎.
1 participant
Summary
TaskTrackertowacore::runtime— a lightweight container that storesAbortHandles and aborts them all on demand or when dropped#[must_use]toAbortHandleso the compiler warns when a spawned task's handle is silently dropped (forces explicit.detach()or tracking)connection_tasks/client_tasksfields toClient:connect()call to prevent unbounded handle accumulation across reconnects. Tasks exit cooperatively viaconnection_generation,is_shutting_down(), andshutdown_notifier— no forceful abort on disconnect (tasks like the post-login sequence must survive 515 reconnect cycles)Clientdrops — covers LID-PN cache warm-up, device registry cleanup, persistence background savermessage_queueson disconnect (bug fix: stale per-chat workers survived reconnects)run_background_savernow returnsAbortHandleinstead of detaching internally.detach()calls remain as intentional fire-and-forget (event dispatch callbacks, best-effort acks/receipts, transport disconnect teardown, background DB writes, sync worker with Weak ref, per-stanza processing)Breaking changes
PersistenceManager::run_background_saver()now returnsAbortHandleinstead of(). Callers must either track the handle (e.g., viaclient_tasks.track()) or call.detach()to preserve the old fire-and-forget behavior.Test plan
cargo clippy --all --tests— zero warningscargo test --all --exclude e2e-tests— all 1090 tests pass#[must_use]produces a warning whenruntime.spawn()result is unused.detach()calls are all intentional fire-and-forgetconnection_tasks.abort_all()is called at start ofconnect()(not duringcleanup_connection_state) to avoid aborting the post-login task during 515 reconnect cyclesmessage_queues.invalidate_all()closes channels so per-chat workers exitconnection_generation/shutdown_notifierhandles stalenessSummary by CodeRabbit
New Features
Improvements