[losewayy]

Description

Fix the use-after-free reported in #4556: Slice::markSuccess()/markFailed() publish the terminal slice counters and then keep dereferencing task/batch state inside finishIfComplete(), while a racing freeBatchID() could already delete the BatchDesc and the TransferTask objects it owns (ASan: heap-use-after-free at Slice::markSuccess in tcp_write_visibility_test). The active_completion_callbacks quiescence counter previously existed only under USE_EVENT_DRIVEN_COMPLETION (off by default) and was incremented inside finishIfComplete after the first task-> dereference — a check-then-enter gap in every build.

What this changes:

• Pin the BatchDesc with an RAII guard at the top of markSuccess/markFailed/sealTaskSubmission, before any terminal state is published; the pin drops when the callback unwinds.
• Make BatchDesc::active_completion_callbacks unconditional so non-EDC builds get the same protection.
• Transport::freeBatchID/MultiTransport::tryFreeBatchID return BatchBusy while pins are held; engine callers already convert that into BatchCleanupDeferred handled by the existing deferred-cleanup thread — no new machinery.
• TcpTransport::submitTransfer now sets task.batch_id like every sibling transport, so direct-Transport submissions get the same pin protection.
• Remove a redundant slice->status store after markFailed in HCCL (it landed after the pin was released).
• Deferred-cleanup loop no longer retries permanent (non-busy) errors forever.
• Tests: the callback-quiescence test now runs in the default build, and freeBatchID assertions accept the now-legal BatchCleanupDeferred outcome.

The intermittent ECONNRESET at the WRITE-ack read is consistent with this race in the shared-process loopback (freed slab reuse corrupting sibling state); if it recurs without an ASan signature it should be tracked as a separate TCP flake.

Credit to @koriyoshi2041 for the lifetime-analysis in the issue that pinpointed the retention boundary.

Module

• Transfer Engine (mooncake-transfer-engine)
• Mooncake Store (mooncake-store)
• Mooncake Conductor (mooncake-conductor)
• Reshard (mooncake-reshard)
• Mooncake EP (mooncake-ep)
• Mooncake PG (mooncake-pg)
• Integration (mooncake-integration)
• P2P Store (mooncake-p2p-store)
• Python Wheel (mooncake-wheel)
• Common (mooncake-common)
• Mooncake RL (mooncake-rl)
• CI/CD
• Docs
• Other

Type of Change

• Bug fix
• New feature
• Refactor
• Breaking change
• Documentation update
• Performance improvement
• Other

How Has This Been Tested?

Test commands:

# default (non-EDC) build, ASan and non-ASan
cmake -S mooncake-transfer-engine -B mooncake-transfer-engine/build-asan -DBUILD_UNIT_TESTS=ON -DUSE_ASAN=ON
cmake --build mooncake-transfer-engine/build-asan --target transport_uint_test tcp_write_visibility_test
./transport_uint_test && ./tcp_write_visibility_test

# EDC variant
cmake -S mooncake-transfer-engine -B build-edc -DBUILD_UNIT_TESTS=ON -DUSE_EVENT_DRIVEN_COMPLETION=ON
cmake --build build-edc --target transport_uint_test tcp_write_visibility_test
./transport_uint_test

# format gate
bash scripts/code_format.sh --check --changed-lines

Test results:

• Unit tests pass — transport_uint_test 44/44 (default build, non-ASan and ASan), 46/46 with USE_EVENT_DRIVEN_COMPLETION=ON
• Integration tests pass — tcp_write_visibility_test 53/53 under ASan (the unpatched code reproduced the heap-use-after-free on the first run); multi_transport_batch_test 5/5, worker_pool_rail_state_test 20/20
• Manual testing done — the deferred-cleanup path was observed live in EDC runs ("cleanup has been deferred", then a clean free once the pin drained)

Checklist

• I have performed a self-review of my own code
• I have formatted my code using ./scripts/code_format.sh
• I have run pre-commit on the files changed in this PR and all hooks pass
• I have updated the documentation (if applicable)
• I have added tests to prove my changes are effective
• For changes >500 LOC: I have filed an RFC issue

AI Assistance Disclosure

• No AI tools were used
• AI tools were used (specify below)

AI assistance disclosure: this fix was implemented with AI assistance; the reproduction, the test results, and every changed line were verified by the human submitter.

[chatgpt-codex-connector]

Codex Review Summary

This comment shows the latest Codex review activity on this pull request.

Review	Status	Commit	Review trigger
📝 Code Review	⚠️ Failed 2026-10-11T03:32:39.692035Z	ac67919	PR opened

ℹ️ 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" or "@codex security review".

Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings.

[koriyoshi2041]

Reviewed exact head 50414ff4 against the #4556 failure ordering. The lifetime acquisition now precedes every terminal publication (markSuccess, markFailed, and submission sealing), while both direct and multi-transport free paths require the pin count to drain. That closes the check-then-enter window without making TCP a special case.

The hosted CTest job is green. The two generic wheel reds are not in this diff: both runners were shut down with exit 137 while building the unrelated PG PyTorch extension; the CUDA/EFA/ROCm and transfer-engine lanes passed. I did not find a separate correctness gap in this head.

[alogfans]

LGTM