[Store] Defer invalid handle cleanup from segment unmount RPC#2877
[Store] Defer invalid handle cleanup from segment unmount RPC#2877Aionw wants to merge 10 commits into
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a SegmentLifetime tracking mechanism to immediately invalidate and hide replicas of unmounted segments from query paths, deferring the physical metadata cleanup asynchronously. The review feedback focuses on optimizing performance by avoiding heap allocations in the default constructor of SegmentLifetime on the data path, and improving safety by adding defensive bounds checks in AllocatorManager to prevent potential out-of-bounds crashes if registrations and allocators get out of sync.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
There was a problem hiding this comment.
Overall this is a solid approach to the unmount-scan cost — separating immediate visibility from physical cleanup via the per-registration SegmentLifetime flag is the right shape. Things I checked and that hold up:
- All descriptor-returning paths are covered and mutually consistent.
GetReplicaList/BatchGetReplicaList/regex/the*ForAdminvariants go throughGetVisibleReplicaDescriptors;ExistKey/BatchExistKey/GetAllKeysthroughHasVisibleCompletedReplica. I checked these two can't disagree during the async window:is_completed_and_available()gates onhas_invalid_mem/nof_handle()→!buffer->isAllocatorValid(), the same predicateget_available_descriptor()uses. So Exist won't say "present" while Get says "not ready". - Per-registration lifetimes correctly disambiguate same-name segments and CXL client segments sharing one underlying allocator, since each
addAllocatormints a freshSegmentAllocator/lifetime. getDescriptorIfAvailablere-checksisAllocatorValid()after building the descriptor — good guard against the flag flipping mid-build.PrepareGracefulUnmount/SetSegmentStatusByNamepassinginvalidate=falsekeeps draining replicas readable, matching the design.
One thing I'd like you to confirm before it merges: snapshot-restored replica buffers don't get the lifetime bound. Serializer<AllocatedBuffer>::deserialize reconnects allocator_ to the remounted buf_allocator, but the buffer's segment_lifetime_ stays default-constructed (always-available), and the registration created in SegmentSerializer::Deserialize (segment.cpp ~944) is a separate lifetime. bindSegmentLifetime only runs in SegmentAllocator::allocate(), so a restored buffer is never bound. After a master restart, a restored replica's unmount-invalidation therefore falls back to allocator_.expired() rather than the flag.
For the common allocator-per-segment case this still invalidates at unmount (the allocator is destroyed once the registration and buf_allocator are released), so no functional regression there. But it (a) breaks the "immediate flag" invariant for restored buffers, and (b) for a shared underlying allocator that outlives a single unmount — CXL client segments, the case this PR explicitly calls out — the restored replica stays isAllocatorValid() == true and remains visible after unmount, i.e. a stale descriptor.
Could you confirm the restore → unmount path (particularly CXL client segments) is covered? The small defensive fix would be to rebind the registration's lifetime onto the restored buffers where the registration is created in SegmentSerializer::Deserialize, so the invariant holds regardless of how a buffer came to exist. A regression test that snapshots → restores → unmounts a memory segment and asserts the object is no longer visible would lock it in.
|
CI failure is expected to be fixed in latest commit. |
|
Codecov Report❌ Patch coverage is 📢 Thoughts on this report? Let us know! |
|
Since this is a critical change to the master service, I also invited @yokinoshitayoki for a double check. |
|
Scope update: I removed the snapshot-specific fixes from this PR, including the pre-fork unavailable-replica sweep, restored lifetime/registration rebinding, LOCAL_DISK-only client restore tracking, and the related snapshot round-trip tests. The reason is to keep #2877 focused on asynchronous invalid-handle cleanup and generation-safe live segment lifetime management. Snapshot consistency and restore behavior have separate correctness and compatibility considerations, so they will be addressed in a dedicated follow-up PR instead of expanding the scope here. The only changes retained under snapshot-related paths are minimal compatibility updates: existing unmount tests now wait for the asynchronous cleanup worker before fixture persistence, and the catalog provider uses the unified replica availability API. |
| * @return ErrorCode::OK if exists | ||
| */ | ||
| auto GetAllKeys(const std::string& tenant_id) | ||
| auto GetAllKeys(const std::string& tenant_id, bool filter_invalid = true) |
There was a problem hiding this comment.
It seems that filter_invalid is currently exposed only on MasterService. The admin HTTP handler ignores the request, and GetAllKeysForAdmin() always uses the default value, so users cannot request filter_invalid=false.
5 participants
Description
MasterService::UnmountSegmentcurrently callsClearInvalidHandles()synchronously on the RPC path. The cleanup walks every metadata shard and examines every key, so its cost scales with the total number of keys managed by the master rather than only the keys stored on the segment being unmounted.This became visible in a stress scenario with four master RPC threads, ten store clients, approximately eight million keys in total, and continuous batch put/get/exists traffic at a 1:1:1 ratio. Unmounting eight store services concurrently caused multiple RPC workers to perform full metadata scans, during which unrelated clients observed connection timeouts. Profiling showed that invalid-handle scanning and metadata destruction dominated the unmount period.
This PR separates immediate replica invalidation from physical metadata removal:
SegmentLifetime.PrepareUnmountSegmentinvalidates the registration immediately, while draining and graceful unmount only remove it from the allocation pool and keep existing replicas readable until final unmount.InvalidHandleCleanupworker instead of scanning every metadata shard synchronously.GetAllKeyshides metadata without an available completed replica by default and retainsfilter_invalid=falsefor callers that need the physical metadata view.The existing
AllocatorManager::getAllocators()interface remains unchanged. The new registration view is used only by allocation and segment lifecycle paths that need lifetime tracking.The behavior change is intentionally limited: replicas become invisible immediately after unmount preparation, while their physical metadata is removed asynchronously. As a result, physical metadata counters may be eventually consistent for a short period. If an object still has another available replica, reads continue using only the surviving replica. The existing synchronous cleanup in the client monitor remains unchanged.
Module
mooncake-transfer-engine)mooncake-store)mooncake-ep)mooncake-pg)mooncake-integration)mooncake-p2p-store)mooncake-wheel)mooncake-common)mooncake-rl)Type of Change
How Has This Been Tested?
The affected test targets were built from the current upstream
main. Tests cover lifetime propagation, same-name segment isolation, CXL shared-allocator behavior, immediate query visibility after unmount, partial-replica survival, asynchronous worker shutdown, admin key filtering, and mounted-segment serialization compatibility.Test commands:
cmake --build build --target allocation_strategy_test segment_test master_service_test serializer_test -j8 ./build/mooncake-store/tests/allocation_strategy_test ./build/mooncake-store/tests/segment_test ./build/mooncake-store/tests/master_service_test --gtest_filter='MasterServiceTest.UnmountSegmentHidesUnavailableReplicas:MasterServiceTest.ReadableAfterPartialUnmountWithReplication:MasterServiceTest.GetAllKeysFiltersUnavailableMetadataByDefault' ./build/mooncake-store/tests/serializer_testTest results:
Results:
Checklist
./scripts/code_format.shpre-commit run --all-filesand all hooks passAI Assistance Disclosure
OpenAI Codex assisted with adapting the change to the latest upstream code, reviewing the implementation, running the listed tests, and drafting this pull request description. The submitter reviewed the resulting code and test results.