fix: harden txn client cancellation and cleanup

[XuPeng-SH]

What type of PR is this?

• API-change
• BUG
• Improvement
• Documentation
• Feature
• Test and CI
• Code Refactoring

Which issue(s) this PR fixes:

issue #25206

What this PR does / why we need it:

This PR hardens the txn client lifecycle for cancellation, client shutdown, and CN-failure recovery, fixing multiple issues that could cause goroutine leaks, hangs, or process crashes in distributed scenarios.

Changes

1. Replace sync.Cond with channel-based signaling (client.go)

• sync.Cond.Wait() cannot respond to context cancellation or client close, causing goroutines to hang indefinitely.
• Replaced with pausedC and closedC channels; all wait paths now use select to listen on ctx.Done() and closedC, ensuring every exit condition wakes blocked callers.
• Added closed flag with proper cleanup on Close() to unblock paused New() calls.

2. Decouple abort signal from timestamp (client.go)

• Old abortC carried time.Time with buffer=1; new abort events were dropped when the channel was full, causing the worker to operate on stale timestamps and miss transactions that should be aborted.
• Changed abortC to a wakeup-only chan struct{}; the latest abort observation is stored in atomic.Int64 with CAS to ensure monotonic advancement. The worker reads latestAbortAt on each wakeup, never missing an observation.

3. Convert panics to error returns (operator.go)

• checkResponseTxnStatusForReadWrite, checkResponseTxnStatusForCommit, checkResponseTxnStatusForRollback, and checkTxnError previously used panic on malformed TN responses, which crashes the entire CN process.
• All panic calls replaced with return moerr.NewInternalError(...), allowing callers to handle errors gracefully.

4. Add waiter cleanup on context cancellation (timestamp_waiter.go)

• When a GetTimestamp caller was canceled, the waiter remained in the linked list until an unrelated future timestamp notification happened to clean it up — a permanent leak if no further notifications arrived.
• Added removeWaiter to drop the waiter from the list immediately on cancellation and unref it.
• Added index field for O(1) swap-delete removal.

5. Context-aware openTxn and waitMarkAllActiveAbortedLocked

• openTxn now accepts context.Context and checks ctx.Err() at all wait points, preventing unbounded blocking.
• waitMarkAllActiveAbortedLocked now propagates context cancellation and client close signals.

6. Safe type assertion for ResumeLockService

• ResumeLockService type assertion changed from bare cast (panics on mismatch) to ok-pattern assertion.

Tests Added

• TestOpenTxnReturnsWhenPausedContextCanceled — verifies paused New() returns on context timeout.
• TestCloseUnblocksPausedNew — verifies Close() unblocks a paused New() call.
• TestOpenTxnReturnsWhenAbortMarkingContextCanceled — verifies openTxn returns on context cancel during abort marking.
• TestMarkAllActiveTxnAbortedRetainsLatestObservation — verifies CAS monotonic advancement and wakeup signal coalescing.
• TestGetTimestampCanceledWaitIsRemoved — verifies canceled waiters are removed from the list.
• TestCommitAndRollbackDoNotProceedAfterRunningSQLTimeout — verifies commit/rollback abort when running SQL times out.
• TestMalformedTNResponsesReturnErrors — verifies malformed responses return errors instead of panicking.

[qodo-code-review]

Qodo reviews are paused for this user.

Troubleshooting steps vary by plan Learn more →

On a Teams plan?
Reviews resume once this user has a paid seat and their Git account is linked in Qodo.
Link Git account →

Using GitHub Enterprise Server, GitLab Self-Managed, or Bitbucket Data Center?
These require an Enterprise plan - Contact us
Contact us →

[XuPeng-SH]

Self-review blocker before merge:

running SQL timeout/cancel is now sealed through scheduleRollbackAfterRunningSQL(...), but the public terminal entrypoints still do not honor that seal on later calls.

Current flow on latest head:

• Commit / Rollback / WriteAndCommit call sealAndWaitRunningSQLWithSQL(...)
• if waiting fails, scheduleRollbackAfterRunningSQL(...) sets terminalAction = rollback-after-running-sql and queues one detached rollback owner
• checkStatus(...) already treats terminalAction != none as closed
• but later public Commit / Rollback calls never re-check that terminal state before proceeding

So a second terminal call on the same operator can still run after the deferred rollback has been scheduled. In the worst case, an intended abort path can be bypassed by a later Commit once the running SQL finally drains.

Suggested fix:

1. reject subsequent public Commit / Rollback once terminalAction is set (return txn-closed / stored wait error), and
2. keep a separate internal rollback path for the deferred-cleanup callback so only that owner goroutine can finish the transaction.

Please also add a regression test for:

• first Commit/Rollback/WriteAndCommit timing out on running SQL and scheduling deferred rollback;
• a second public terminal call arriving before/after drain;
• assertion that the second call is rejected and cannot bypass the deferred rollback owner.

[XuPeng-SH]

Blocking issue:

The deferred-rollback seal is still not honored by later public terminal calls.

Current latest-head flow:

• Commit / Rollback / WriteAndCommit call sealAndWaitRunningSQLWithSQL(...)
• on wait failure, scheduleRollbackAfterRunningSQL(...) sets terminalAction = rollback-after-running-sql and transfers cleanup ownership to a deferred rollback callback
• checkStatus(...) already treats terminalAction != none as closed
• but the public terminal entrypoints themselves never check terminalAction before continuing

So a second public Commit / Rollback can still enter the terminal path after cleanup ownership has already been handed to the deferred rollback. The existing new tests cover 'one deferred rollback owner' and 'new SQL is rejected', but they still do not cover 'a second public terminal call is rejected immediately'.

I verified this with a scratch txn/client test on the latest head: after a first canceled Commit schedules deferred rollback and the running SQL drains into a blocked rollback workspace, a second public Commit is not rejected at entry. It proceeds into the terminal path and only returns later with a joined txn-closed/context-deadline error. That means the seal is still not being enforced at the public API boundary.

Suggestion: once terminalAction is set, make public Commit / Rollback / WriteAndCommit fail fast before sealAndWaitRunningSQLWithSQL waits or doWrite runs. Keep a private/internal rollback path for the deferred cleanup owner so only that owner goroutine can finish the transaction. Please add a regression that covers a second public terminal call both before and after the running SQL drains.

[XuPeng-SH]

Addressed the latest terminal-ownership blocker in a873e1e and rebased the PR directly onto main@67898b60fe.

Root cause and fix:

• The deferred rollback seal was only represented inside the operator mutex, while public Commit, Rollback, and WriteAndCommit had no unique terminal owner.
• Added a lock-free idle -> active -> sealed terminal-call state machine. Competing public terminal calls now fail immediately with ErrTxnClosed, including while the owner is blocked in workspace or sender code.
• Split the private rollback cleanup implementation from the public Rollback ownership gate. Deferred running-SQL cleanup and workspace-commit failure cleanup retain their single owner and still cover workspace rollback, TN rollback, and lock release.
• Kept the terminal gate sealed through RestartTxn claim and admission, then reopened it only after the new generation run-SQL gate is open. The lock-free rejection path deliberately avoids reading txnID while restart can replace it.

Regression coverage:

• all 3 timeout owners x all 3 later public terminal calls, both before SQL drain and while deferred rollback is blocked after drain;
• concurrent public terminal calls while Commit is blocked in sender RPC;
• exactly one deferred rollback owner;
• restart rejected during deferred cleanup and terminal calls rejected during restart admission.

Validation after rebase:

• focused lifecycle tests: go test -race -count=50
• go test -count=1 ./pkg/txn/client
• go test -race -count=1 ./pkg/txn/client
• full tests for all 9 changed packages: bootstrap, cdc, frontend, frontend/test, incrservice, sql/compile, txn/client, txn/storage/memorystorage, vm/engine
• go vet for all changed packages
• go build ./cmd/mo-service
• RunSQLTracker benchmarks: 0 allocs on enter/exit fast paths

[XuPeng-SH]

Latest head looks good to me.

The previous blocker around deferred running-SQL rollback ownership is now closed: public Commit / Rollback / WriteAndCommit calls are gated by claimTerminalCall, scheduleRollbackAfterRunningSQL seals that gate before handing ownership to the deferred rollback path, and the new regressions cover both before-drain and after-drain second terminal calls.

I also re-ran the latest txn/client targeted terminal-call tests plus the full pkg/txn/client package tests on this head, and I don't see a remaining high-confidence correctness blocker.

[LeftHandCold]

Requesting changes for one confirmed generation-boundary race on the latest head. The normal pkg/txn/client suite and its existing race suite pass, but a targeted Commit/Exit/Restart stress test reports the race described inline. Please make the last-token exit the final old-generation access and add that interleaving as a regression test.

[LeftHandCold]

[P1] Keep the old RunSQL generation alive until Exit finishes

exitToken removes the last tracker entry and broadcasts the terminal waiter, but this method still reads tc.reset.runSqlCounter afterward. Commit can therefore wake, close, and return; RestartTxn then sees no active token in sealForRestart and initReset overwrites the counter while this old Exit is still reading it. I reproduced this on head 67bad915 with a Commit/Exit/Restart stress test under -race: the write is in initReset at operator.go:692 and the conflicting read is here at operator.go:2043.

Please make token removal/drain notification the final old-generation operation (or move all reset-state updates before it / return the drain state while holding the tracker lock), and add a regression that restarts immediately after Commit is released by the exiting SQL goroutine.

[XuPeng-SH]

Fixed in 62659ab by closing the generation boundary rather than synchronizing the reported read in isolation.

Root cause: activeTokens was the restart authority, but runSqlCounter/runningSQL remained a second lifecycle state that Exit accessed after publishing the final-token drain. Also, nextID was reset on restart, allowing token ABA across generations.

Changes:

• moved run-SQL activity and diagnostic counts into runSQLTracker under its existing mutex;
• made tracker exit the final operation in ExitRunSqlWithToken, with no operator reset-state access after drain publication;
• kept token IDs monotonic for the operator lifetime, reject on ID-space exhaustion, and ignore stale/duplicate exits;
• added regressions for Commit -> final Exit -> immediate Restart without waiting for the Exit caller, plus stale old-token vs new-generation token isolation.

Validated with focused -race -count=50, full pkg/txn/client race tests, vet, mo-service build, and RunSQL tracker benchmarks (0 allocations on enter/exit).

[aptend]

Request changes: one P1 generation-boundary data race remains on the latest head 67bad915.

ExitRunSqlWithToken removes the final tracker token and wakes the terminal waiter in exitToken(token), but then still reads tc.reset.runSqlCounter and may write runningSQL. Once the final token is removed, Commit can finish and RestartTxn can pass sealForRestart; initReset then replaces runSqlCounter while the old-generation Exit is still reading it.

I independently reproduced this with a Commit → last-token Exit → immediate RestartTxn stress test under go test -race. The detector reports:

• write: txnOperator.initReset, operator.go:692
• read: txnOperator.ExitRunSqlWithToken, operator.go:2043

Please make token removal/drain publication the final old-generation operation: complete all counter/runningSQL updates before publishing the tracker as drained, or return the necessary drain state under the tracker lock without touching resettable state afterward. Please also add this exact Commit/Exit/Restart interleaving as a race regression test.

I also audited the cancellation/Close wait paths, timestamp waiter ownership, max-active queue cleanup, deferred workspace/TN/lock rollback, terminal-call ownership, downstream API propagation, and accumulation cleanup. I did not find another high-confidence leak/hang/OOM blocker on this head.

[gouhongshen]

Review summary

除已有 review threads 外，当前 HEAD 仍有以下未覆盖问题；第一项需要在合并前修复。

[P1] Restart admission failure leaks client ownership

RestartTxn calls initForRestart before doCreateTxn (pkg/txn/client/client.go:438-446), and initForRestart seals terminalCall (pkg/txn/client/operator.go:643-712). However, doCreateTxn may already admit the operator into activeTxns or append it to waitActiveTxns (pkg/txn/client/client.go:819-841) before snapshot acquisition or the max-active wait fails. Both failure paths then call public op.Rollback (client.go:514-542), whose first operation is claimTerminalCall (operator.go:1087); because the restart operator is already sealed, rollback returns ErrTxnClosed and the active/queued ownership is never cleaned up.

A canceled snapshot wait therefore leaks the active transaction and user count. A canceled max-active wait leaves an abandoned queue entry that can later be activated. The operator is neither closed nor restartable. The existing TestRestartTxnAdmissionFailureLeavesOperatorClosed only cancels before admission and does not cover these post-admission failures.

Please give post-admission failure a private cleanup owner that removes active/queued ownership and closes the operator regardless of the public terminal gate, or defer sealing terminal calls until admission succeeds. Add regression tests for both canceled snapshot wait and canceled max-active wait.

[P2] Malformed TN responses can poison non-terminal operator state

doSend publishes the response timestamp/status before response validation (pkg/txn/client/operator.go:1587-1598). The new invalid-status path returns an error (operator.go:1702), but public Read and non-commit Write do not perform terminal cleanup or restore the published state. A malformed status such as 99 can therefore be stored in tc.mu.txn.Status while the operator remains active and owned by the client.

The current TestMalformedTNResponsesReturnErrors only tests helper functions and does not exercise the public Read/Write lifecycle. Please validate the response before publishing state, define the intended cleanup/rollback behavior for protocol-invalid responses, and add end-to-end tests.

[P2] Exported interface change is source-breaking

pkg/txn/client.TxnOperator changes EnterRunSqlWithTokenAndSQL from (..., string) uint64 to (..., string) (uint64, error) (pkg/txn/client/types.go:201-204). Updating in-tree implementations and mocks does not preserve compatibility for external implementers of this exported interface. The PR description currently does not call out an API change. Please explicitly acknowledge/version this break, or make the new behavior additive through a new interface/method.

Additional scope / implementation notes

• The 11 .agents/.claude files add 715 lines of unrelated skill/documentation changes to a transaction-client lifecycle PR; please split or drop them.
• pkg/txn/storage/memorystorage/storage_txn_client.go:147-152 now returns a fixed token (1) and makes Exit a no-op. This does not implement the new tracker semantics (cancellation tracking, token uniqueness, or sealed-state rejection), and it is used by the memory engine/test engine. Please either implement the real lifecycle contract or make this limitation explicit and keep it out of the production interface path.

Local verification: go test ./pkg/txn/client -count=1 and go test -race ./pkg/txn/client -count=1 pass.

[XuPeng-SH]

Addressed the latest review systematically in f2d1fde, after rebasing the full branch onto main@cb94d94954.

Root contracts and fixes:

1. Admission ownership is now closed on every post-open failure.

   • Once openTxn transfers an operator into the client active/queued graph, doCreateTxn uses a private abort transition rather than public Rollback.
   • This cleanup is independent of the restart terminal-call seal, emits ClosedEvent exactly once, and removes either active ownership/users or the queued admission entry.
   • Cleanup uses a non-canceled context so cancellation cannot suppress ownership release.
   • Added end-to-end RestartTxn regressions for canceled snapshot acquisition and canceled max-active admission; both assert active count, users, queue state, closed generation, and subsequent operator reuse.

2. TN protocol state is validate-then-publish.

   • doSend no longer mutates local txn state.
   • handleError validates the entire response batch first and only then publishes the final CommitTS/Status under the operator mutex.
   • A late non-terminal response also cannot rewrite an already closed operator.
   • Protocol-invalid Read/non-commit Write returns an operation error while leaving the transaction Active and unchanged; terminal ownership remains with the caller, which can explicitly Rollback. Public API tests cover both methods and subsequent rollback.

3. The exported interface change is now additive and source-compatible.

   • Restored the original TxnOperator.EnterRunSqlWithTokenAndSQL(... ) uint64 signature.
   • Added the optional RunSQLAdmissionOperator capability plus client.TryEnterRunSqlWithTokenAndSQL helper for error-aware admission.
   • MatrixOne production call sites use the helper; legacy external implementations continue to compile and receive the original behavior. Added an explicit legacy-fallback compatibility test.

4. Memory storage no longer uses fixed token 1/no-op Exit.

   • It now has unique operator-lifetime tokens, tracked cancel funcs, idempotent Exit, terminal seal/rejection, cancellation propagation, and drain-before-Commit/Rollback.
   • Added a concurrent lifecycle test covering two tokens, cancellation, rejection after seal, duplicate Exit, and terminal drain.

Scope note: the skill/documentation changes remain in this branch because the branch owner explicitly requested that the generalized macOS/CGo and lifecycle-review guidance be updated here. They are platform/PR-independent methods, not PR-specific workarounds.

Post-rebase verification:

• focused new lifecycle/protocol/API/storage tests under -race -count=100
• full go test -race for pkg/txn/client and pkg/txn/storage/memorystorage
• full ordinary tests for all 8 affected packages
• go vet for all 8 affected packages
• make build
• RunSQL tracker benchmark: 0 allocs; serial 49-51 ns/op, parallel 212-222 ns/op

The branch was pushed with --force-with-lease after the rebase.

[XuPeng-SH]

Latest head still looks good to me.

I re-checked the new delta on top of the previous approved state. The added admission / lifecycle fixes look coherent: RestartTxn failure paths now close ownership explicitly, the terminal-call gate and deferred-rollback owner still hold, and response txn state is only published after response validation passes.

I also ran the latest targeted restart / terminal-call regression tests and the full pkg/txn/client package tests on this head, and I don't see a remaining high-confidence correctness blocker.