[TENT] Fix TaskInfo dangling pointers, data races, stage buffer leaks

[RuiqingFeng]

Description

Fix several critical staging runtime bugs in ProxyManager and TransferEngineImpl:

• Reserve task_list capacity and reject batch capacity overflow to prevent dangling TaskInfo* pointers.
• Read staging_status via __atomic_load to eliminate data races during task status polling.
• Protect stage_buffers_ concurrent map mutations with a recursive mutex.
• Scope StageBufferCache to individual tasks and drain in-flight batches on progress errors.
• Replace stage buffer allocation LOG(FATAL) process crash calls with recoverable error statuses.

Module

• Transfer Engine (mooncake-transfer-engine)
• Mooncake Store (mooncake-store)
• 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?

Added 4 new unit test cases in mooncake-transfer-engine/tent/tests/progress_worker_test.cpp:

• ProxyManager.ConcurrentFirstAllocationIsBounded
• ProxyManager.ReclaimsStageBuffersBetweenWorkerTasks
• ProxyManager.DrainsInflightBatchesAfterProgressError
• TransferEngineImpl.DirectSubmitRejectsBatchCapacityOverflow

Test commands:

cmake -S . -B build-tent \
  -DUSE_TENT=ON \
  -DUSE_CUDA=OFF \
  -DWITH_STORE=OFF \
  -DWITH_STORE_RUST=OFF \
  -DBUILD_UNIT_TESTS=ON

cmake --build build-tent --target tent_progress_worker_test -j

./build-tent/mooncake-transfer-engine/tent/tests/tent_progress_worker_test \
  --gtest_filter='ProxyManager.*:TransferEngineImpl.DirectSubmitRejectsBatchCapacityOverflow'

Test results:

• Unit tests pass
• Integration tests pass (if applicable)
• Manual testing done (describe below)

Ran unit tests on TENT Linux build environment; all 4 newly added unit tests passed.

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 run --all-files 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)

Used Cursor Agent for root cause analysis of TENT staging memory leaks & data races, patch verification, and drafting unit test coverage.