[Store] Drain in-flight RPCs and keep pools alive across client teardown

[he-yufeng]

Description

Fixes #3909: tearing down a read-side client after consecutive RPC timeouts segfaults in asio's epoll_reactor. Two lifetime violations, both verified by reading current main against ylt's documented contract (only send_request is thread-safe; close/destruction must not race in-flight RPC):

1. In-flight requests. MasterClient/DummyClient/ClientRequester destructors just dropped the client pool while co_await pool->send_request(...) coroutines could still be suspended. A shared RpcDrainGuard (mooncake-common) stops admitting new calls once draining and makes the destructor wait for in-flight ones, bounded at 30s with a loud error rather than silently proceeding.

2. Background reconnect coroutines. ylt's client_pool reconnect loop (client_pool.hpp) holds references into pool storage and resumes across sleeps, so no amount of request draining makes pool destruction safe mid-retry. RpcClientPool now hands out pools from a process-wide registry keyed by address instead of freeing them on teardown or address flap. Pools are one-per-distinct-address per process and deliberately process-lifetime.

Module

• Mooncake Store (mooncake-store)
• Common (mooncake-common)

Type of Change

• Bug fix

How Has This Been Tested?

Test commands:

cmake --build build-linux --target rpc_drain_guard_test rpc_timeout_test master_service_test master_service_ha_test client_integration_test

Test results:

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

Linux container (Ubuntu 22.04): new rpc_drain_guard_test (3 tests: drain waits for in-flight and refuses new calls, empty drain returns immediately, timeout reports false) passes. rpc_timeout_test gains TeardownDrainsInFlightTimeoutRequests: 16 concurrent ExistKey calls in flight against a black-hole master at client teardown, 10 rounds, green with the fix. master_service_test 69/69 and master_service_ha_test 95/95 pass (one earlier run flaked once on an unrelated timing test under container load; two clean runs followed).

Honest limitation: the reporter's crash is timing-dependent at 200-node scale, and my in-process harness did not crash pre-fix in this shape, so I am not claiming a local red-to-green crash repro. The fix is verified by construction against ylt's documented thread-safety contract plus the guard semantics tests above.

Checklist

• I have performed a self-review of my own code
• I have formatted my code using ./scripts/code_format.sh (clang-format plus pre-commit on every touched file; all hooks pass)
• I have run pre-commit on the files changed in this PR and all hooks pass
• I have updated the documentation (not applicable)
• I have added tests to prove my changes are effective
• For changes >500 LOC: I have filed an RFC issue (N/A: +322/-35 incl. tests)

AI Assistance Disclosure

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

Prepared with AI assistance (Kimi Code): the mechanism chain (drain + reconnect-storage) was traced on current main by reading, the fix and tests were drafted by the agent, and I reviewed the complete diff before submitting. All verification runs above were actually executed, and the no-deterministic-crash-repro limitation is stated as observed.

Update (2026-09-11, head f9844e4)

Reviewer questions closed:

• ClientRequester's offload pool collection is no longer freed at teardown at all (detail::KeepClientPoolsAlive): the ylt reconnect loop references pool storage whether or not a user call is in flight, so freeing was the original UAF class on every teardown, not just on drain timeout. The 30s drain now covers response correctness only; on timeout, in-flight calls lose their answers, but no late resume touches freed storage. Verified against the offload RPC call sites (invoke_rpc<&RealClient::batch_get_offload_object> and release_offload_buffer), which both enter through the guarded invoke_rpc.
• The pool registry stores the first configuration's salient knobs (max_connection, connect/request timeouts) and warns loudly on a later config mismatch for the same address instead of silently keeping first-wins.
• New tests: RegistryKeepsFirstConfigForSameAddress, KeepClientPoolsAliveRetainsCollection. rpc_client_io_context_test 5/5, rpc_drain_guard_test 3/3, rpc_timeout_test 6/6, all in the Linux container.

[he-yufeng]

Follow-up on head dbc39a0: the CTest lane caught RpcClientIoContextPoolTest.ReplacesPoolWhenTargetChanges, which pinned the old pool-freeing contract (EXPECT_TRUE(old_pool.expired()) after release). That contract is exactly what this PR changes by design — pools are process-lifetime now because a freed pool's background reconnect coroutines still reference its storage. The test now asserts the new contract (same address same pool, different address different pool, released pool stays alive). Green locally in the container.

[he-yufeng]

The CTest leg failed on CloseDmabufExport.ClosesLiveFdAndClearsIt ("fd 8 should be closed", actual false) in dmabuf_export_test, a transfer-engine binary that does not link the code this PR touches (the RPC pool registry lives in mooncake-store's client teardown path). The assertion checks an fd right after close() in a shared CI process, so an fd-number reuse under parallel load fits the observed "Actual: false" shape. Could a maintainer re-run the job? If it recurs on rerun I will dig into the test's fd bookkeeping instead.

[Icedcoco]

Thanks for putting this together. The underlying issue is important, and the drain guard is a reasonable direction. I had a few questions about the remaining lifecycle and resource semantics before we can consider merging:

1. Drain timeout and pool release

   In the destructors, the pool is still released after drain_for(30s) returns false. Could this still leave a suspended RPC accessing the pool after teardown, preserving the original UAF risk in the timeout or cancellation-disabled cases?

   Would it be safer to keep the pool alive until all in-flight calls have completed, or otherwise explicitly cancel and join those calls before releasing the client state?

2. Process-wide pool registry

   The new registry keeps one pool permanently for every address ever used by the process. Could this grow without bound in deployments where peer addresses or ports are dynamic?

   Also, since the first pool created for an address determines its configuration, how should we handle two users selecting different timeout, protocol, or IO-context settings for the same address? Is the process-wide “first configuration wins” behavior intentional and guaranteed by the surrounding APIs?

3. ClientRequester coverage

   The guard is added to ClientRequester::invoke_rpc(), but could we add teardown tests covering all relevant offload paths, including batch_get_offload_object() and release_offload_buffer()? This would help confirm that no public path can outlive client_pools_.

The fix is valuable, but I would prefer to resolve these points or document the intended trade-offs explicitly before giving approval.

[he-yufeng]

@Icedcoco Good catches — all three points were real gaps, and the new head (will note the SHA after push) closes them:

1. Drain timeout / pool release. You were right, and it was worse than the drain-failure case: the ylt pool's background reconnect loop references pool storage whether or not a user call is in flight, so freeing client_pools_ at teardown was the original UAF class on every teardown, drain or not. ~ClientRequester no longer frees the collection at all — it is parked process-wide (detail::KeepClientPoolsAlive, collections are one per RealClient, so few in practice). The 30s drain now only covers response correctness: if it times out, in-flight calls lose their answers, but no late resume ever touches freed storage. That matches how the MasterClient read path already works (the process-wide registry), so both RPC planes now share one lifetime rule.

2. Registry config. Also fair. The registry now stores the first pool's salient knobs (max_connection, connect/request timeouts) and warns loudly on a later mismatch for the same address instead of silently keeping the first config. I deliberately kept first-wins rather than keying by config fingerprint: two pool configurations for one address in one process is a configuration smell, and a loud warning surfaces it instead of multiplying pools. If you would rather fingerprint the key, it is a three-line change from here.

3. Offload path coverage. Verified end to end: both ClientRequester::batch_get_offload_object and release_offload_buffer route through ClientRequester::invoke_rpc, which takes the RpcDrainGuard::ScopedCall on entry (real_client.cpp around the invoke_rpc<&RealClient::batch_get_offload_object> and release_offload_buffer call sites). So the RPC side of both paths is covered by the same guard. The byte plane (TransferEngine openSegment/transfer requests) is a different subsystem with no pool-backed coroutines and is out of this PR's scope by design.

New tests pin 1 and 2: RegistryKeepsFirstConfigForSameAddress (same address, differing config, same pool object back) and KeepClientPoolsAliveRetainsCollection (collection outlives its owner). The rpc_timeout and drain-guard suites stay green.

[he-yufeng]

@Icedcoco Merged current main into the branch (424c49e). The conflicts with your HA connection-policy split (#3743) and the RpcProtocolConfig restructure (#3976) are resolved by keeping both sides, with one placement change worth calling out: the drain guard moved from invoke_rpc down into invoke_rpc_with_client_pool, so the foreground, HA control, and HA probe paths all drain through one choke point.

The merge also caught a real interaction between this PR and yours, and I want to flag it explicitly because it changes the registry behavior you reviewed. HaControlPolicyPreservesForegroundRetryPolicy failed on the merged tree: the HA probe to a refused address took 4299ms where the test expects a fast fail under 750ms. The cause was my shared-pool registry keying pools by address alone. The test's first (foreground) Connect creates the resilient pool for that address; the later HA probe on the same address then got handed the resilient pool, retry budget and all. The config-mismatch warning fired correctly, but "keep the first configuration" is wrong when two policies legitimately coexist on one address, which is exactly what your HA split does.

So the registry key now carries the behavioral knobs (connection count, retry count, reconnect wait, connect and request timeouts, socket flavor) instead of the bare address. Identical configurations still share one pool, which is the case the keep-alive exists for, and different policies get their own pools. The mismatch warning is gone because a mismatch can no longer alias two policies; RegistryKeepsFirstConfigForSameAddress became RegistrySharesPoolOnlyForIdenticalConfig to pin that.

Verified on the merged head in the dev container: rpc_timeout_test 8/8 (both your HA tests and the drain loop), rpc_client_io_context_test 5/5, rpc_drain_guard_test 3/3, and scripts/code_format.sh --check clean.

[Icedcoco]

Thanks for the follow-up and for addressing the earlier lifecycle and HA-policy concerns. I rechecked the current head and confirmed that the latest CI and relevant tests are green.

I have two remaining questions about teardown behavior:

1. Lifetime of RpcDrainGuard after the 30-second timeout

   If drain_for(30s) times out, the client object and its rpc_drain_ member can still be destroyed while an existing ScopedCall is running. When that RPC eventually returns, ScopedCall::~ScopedCall() calls guard_.leave(), which accesses the guard’s counters and synchronization objects.

   Could this leave a use-after-free in the guard itself, even though the underlying pool is kept alive? Would it be possible for ScopedCall to own a shared state object, or otherwise keep the drain state alive until all in-flight calls have exited?

2. DummyClient teardown RPCs after draining starts

   DummyClient::~DummyClient() calls drain_for() before tearDownAll(). Since drain_for() permanently sets stopping_, the cleanup RPC issued by tearDownAll() through unregister_shm() will be rejected by RpcDrainGuard.

   Is this intentional? If remote shared-memory cleanup is still expected during teardown, should it run before draining, or have a separate teardown path? If the cleanup is deliberately best effort after draining, could we document that behavior and add a focused test?

The registry and HA policy changes look substantially better now, and I am no longer treating CI as an issue. These questions are specifically about the remaining timeout and teardown semantics.

[he-yufeng]

@Icedcoco Both were real bugs, thanks for pushing on the teardown semantics. Fixed in 0279884.

1. Guard lifetime after a timed-out drain

You were right, the pool surviving was only half the problem. ScopedCall held a reference to the guard, so once drain_for(30s) timed out and the owner finished tearing down, the late ~ScopedCall() -> guard_.leave() ran on freed memory. The comment even admitted the pool side of the risk while the guard side sat there unprotected.

The counters now live in a shared State held by shared_ptr from both the guard and each ScopedCall. A call that outlives its guard still finds valid counters, and drain_for semantics are unchanged otherwise.

I gave the regression test teeth by running it under ASAN:

• Pre-fix: ScopedCallSurvivesGuard dies with heap-use-after-free in __atomic_base::fetch_sub, freed stack pointing at the guard destruction. Exactly your scenario.
• Post-fix: clean, and the full suite passes (5/5, including a second new test that pins counter consistency across a timed-out drain followed by a late leave).

2. DummyClient teardown RPCs after draining starts

Not intentional, and worse than you described. unregister_shm() goes through invoke_rpc, so the old order rejected the unmap unconditionally. On top of that the ping thread is only joined inside tearDownAll(), so during the 30s drain the rejected pings flipped it into the reconnection loop, which re-registers shm over UDS while the client is trying to tear down. The pre-drain destructor was just tearDownAll(), so my ordering managed to break cleanup and add spurious re-registrations at the same time.

~DummyClient now runs tearDownAll() first while the guard is still open (unmap goes through, ping thread joins), then drain_for(30s) as the last step before members die. MasterClient and ClientRequester destructors only drain and issue no RPCs afterwards, so they are unaffected by the same trap.

Verification: rpc_drain_guard_test 5/5, rpc_timeout_test 8/8, rpc_client_io_context_test 5/5, store objects rebuild clean, pre-commit format gate green.