[yokinoshitayoki]

Description

Implements the Store-side portion of RFC #4255.
K3 prefix lookups may inspect multiple KDA checkpoints while restoring only one. This PR adds GrantLeasePolicy to BatchProbeKey to support selective lease acquisition:

• None: preserves the existing lease-free probe behavior.
• LastHitOnly: preserves per-key existence results and grants leases to the last complete candidate in caller order. The last all-true candidate in the returned results identifies the selected candidate.
  candidate_size defines how many consecutive keys form a candidate. The implementation first probes without granting leases, then rechecks and leases candidates in reverse order using the existing per-shard locking. Recheck results are written back into the returned results. Concurrent eviction may leave partial leases on failed candidates until their TTL expires.
  The policy is exposed through the C++ client/RPC layers and Python binding. Existing existence APIs, single-key probes, soft-pin behavior, and shared-group lease semantics remain unchanged. Inference-engine integration is left to follow-up PRs.

Compatibility: batch-probe RPC signatures change. Master, client daemon, and clients must be upgraded together; mixed-version batch-probe calls are not supported.

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?

CPU/TCP validation on A100 :

• 154 C++ tests passed across Master, policy, group, and client suites.
• 6 Python tests passed through direct-client and production-daemon RPC paths.
• Concurrent selection/removal test passed 20 repetitions, with 64 races per repetition.
• Scoped pre-commit, changed-line formatting, and documentation build passed.

Test results:

• Unit tests pass

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

• AI tools were used.

Codex assisted with implementation, testing, and documentation. And I have reviewed the code and docs.

[rishabhsinha17]

The project needs this mode. #3801 converged on the two stage K3 lookup, and LastHitOnly is the right primitive. Placement still needs confirming. That thread ended with lease policies on BatchExistKey, added as a new RPC so no existing signature changes. #4255 reverses both points and cites a discussion with @Aionw that is not on the issue. Changing BatchProbeKey before any release includes it is defensible. Could @Aionw or @staryxchen confirm on #4255?

Correctness looks good. Each candidate takes its shard locks shared, in ascending index order, and every other serving path locks one shard at a time, so they cannot deadlock. No lease goes out until the candidate passes the readable check in full, so a racing removal cannot yield a partial selection.

Nothing in the master service scenario suites pins the new mask or its lease effect. A BatchProbe spec beside BatchExistence would drive the call. ExpireAt and EvictMemory then confirm eviction removes unselected candidates only, with Object().IsReadable() on the leased keys. I can add the verb and a first case in a follow-up.

Thanks for the careful RFC and tests.

[Aionw]

Thanks for the careful RFC and tests. I checked the LastHitOnly path and the results are correct. I think the interface and the implementation can both be simpler, though, while keeping the current BatchProbeKey signature and return type.

A lease policy should only change which leases are granted, not what the results mean.

Right now LastHitOnly turns the result into a selected-only mask, and 0 then means "not selected" rather than "missing". This causes three problems:

• Callers have to decode the mask and guard against states the server never produces, such as partly true groups or more than one all-true group.
• A None-style reading such as [k for k, r in zip(keys, res) if r == 1] gives a silently wrong answer when applied to a LastHitOnly result.
• The two modes end up as two unrelated code paths behind one if.

I'd suggest that the results mean per-key existence in both modes, with LastHitOnly only adding the lease side effect. The contract then becomes: under LastHitOnly, the last all-true group of candidate_size keys in the result is exactly the leased candidate. The server can guarantee this by writing what it observed at lease time back into the results:

// validate policy first
auto results = BatchExistKeyImpl(keys, tenant, /*grant_lease=*/false);
if (policy.lease_mode == ProbeLeaseMode::None) return results;

const size_t n = policy.candidate_size;
for (size_t end = keys.size(); end > 0; end -= n) {
    const size_t begin = end - n;
    if (!AllTrue(results, begin, end)) continue;
    std::vector<std::string> candidate(keys.begin() + begin,
                                       keys.begin() + end);
    auto leased = BatchExistKeyImpl(candidate, tenant, /*grant_lease=*/true);
    // A key evicted since the probe flips to false, so this group is no
    // longer all-true and the caller will not pick it.
    std::copy(leased.begin(), leased.end(), results.begin() + begin);
    if (AllTrue(leased)) break;
}
return results;

The caller then picks the leased candidate like this:

res = store.batch_probe_key(keys, policy="LastHitOnly", candidate_size=n)
selected = next((c for c in reversed(range(len(keys) // n))
                 if all(r == 1 for r in res[c * n:(c + 1) * n])), None)

This also simplifies the implementation:

• The ordered multi-shard locking can go. A read lease protects each key on its own, so a candidate for which BatchExistKeyImpl(..., /*grant_lease=*/true) returns all true is already fully protected.
  • Trade-off: when an eviction lands between the probe and the lease, a failed candidate can keep a few partial leases. They expire by TTL, just like the extra KV-side leases the RFC already accepts.
• Much less locking. The common path becomes one lease-free scan (one shard lock at a time) plus a lease call on a single candidate. Today every candidate from the end, including obviously incomplete ones, goes through the full locked check: snapshot lock, a std::map, one unique_ptr accessor per shard, and up to candidate_size shard locks held at once.
• One place for validation. Keep only the check in Master, drop the duplicate in store_py.cpp, and log which check failed (unknown mode, zero size, or non-divisible size).

The docs and docstrings also need updating: with this change a LastHitOnly result reads the same way as a None result, so the selected-only mask wording goes away.

[yokinoshitayoki]

@Aionw Thanks for your review, I have updated this PR accordingly, please have a look when you have time.

[staryxchen]

LGTM

[rishabhsinha17]

The scenario verb from my review is up as #4523, written against the BatchProbeKey signature on main with the lease-free eviction case, and the LastHitOnly selection case will follow the result contract once it settles here.

[yokinoshitayoki]

@stmatengss @zxpdemonio Could you have a look at this PR when you have time?

[zxpdemonio]

I reviewed the current head. The implementation, concurrency behavior, RPC propagation, boundary handling, and test coverage look good; I found no code-level blockers.

There is one documentation mismatch to fix before merging: the PR description still says that LastHitOnly returns a “selected-only mask.” The current contract preserves per-key existence results, and the last all-true candidate in the returned results identifies the candidate that received leases. Please update the description accordingly.

[yokinoshitayoki]

@zxpdemonio Thanks for your review, I have updated PR description accordingly.

[rishabhsinha17]

The LastHitOnly scenario from my review is up as #4789, on top of the merged #4523 verb and pinning the per-key contract approved here.