[Store] Revalidate lease-expired reads by replica identity

docs/source/design/store/mooncake-store.md

@@ -428,7 +428,7 @@ To avoid data conflicts, a per-object lease is granted whenever an `ExistKey` re
 
 For grouped objects, a successful `ExistKey` or `GetReplicaList` refreshes the lease for the current members of the group, so recently accessed members are less likely to be separated by memory eviction.
 
-However, if the lease expires before a `Get` operation finishes reading the data, the operation will be considered failed, and no data will be returned, in order to prevent potential data corruption.
+If the lease expires before a `Get` or `BatchGet` finishes reading the data, the client asks the master whether the replica it read is still there, still readable, and not rewritten since the query (same replica ID and write version). If so, the data is returned. Otherwise, the operation fails with `LEASE_EXPIRED` and no data is returned, in order to prevent potential data corruption.
 
 The default lease TTL is 10 seconds and is configurable via a startup parameter of `master_service`.
 

mooncake-store/include/client_service.h

@@ -52,16 +52,21 @@ class QueryResult {
     /** @brief Tenant scope the master resolved the query to, when the master
      * reports it; readers forward it to disk-replica verify requests */
     const std::optional<std::string> resolved_tenant_id;
+    /** @brief Object write version at query time, when the master reports
+     * it; lets a read that outlived its lease be revalidated */
+    const std::optional<uint64_t> write_version;
 
     QueryResult(
         std::vector<Replica::Descriptor>&& replicas_param,
         std::chrono::steady_clock::time_point lease_timeout_param,
         std::optional<uint64_t> object_checksum_param = std::nullopt,
-        std::optional<std::string> resolved_tenant_id_param = std::nullopt)
+        std::optional<std::string> resolved_tenant_id_param = std::nullopt,
+        std::optional<uint64_t> write_version_param = std::nullopt)
         : replicas(std::move(replicas_param)),
           lease_timeout(lease_timeout_param),
           object_checksum(object_checksum_param),
-          resolved_tenant_id(std::move(resolved_tenant_id_param)) {}
+          resolved_tenant_id(std::move(resolved_tenant_id_param)),
+          write_version(write_version_param) {}
 
     bool IsLeaseExpired() const {
         return std::chrono::steady_clock::now() >= lease_timeout;
@@ -1018,6 +1023,25 @@ class Client {
         const std::vector<Replica::Descriptor>& replica_list,
         Replica::Descriptor& replica) const;
 
+    /**
+     * @brief For reads that completed after their lease expired: whether the
+     * master confirms each key's memory replica is still the one queried.
+     * A key that cannot be checked (old master, other replica type, RPC
+     * error) is reported as not live.
+     */
+    std::vector<bool> CheckLeaseExpiredReads(
+        const std::vector<std::string>& keys,
+        const std::vector<const QueryResult*>& query_results);
+
+    /**
+     * @brief Sets LEASE_EXPIRED on each successful read whose lease expired
+     * and that CheckLeaseExpiredReads does not confirm.
+     */
+    void FailLeaseExpiredReads(
+        const std::vector<std::string>& keys,
+        const std::vector<QueryResult>& query_results,
+        std::vector<tl::expected<void, ErrorCode>>& results);
+
     /**
      * @brief Batch put helper methods for structured approach
      */

mooncake-store/include/master_client.h

@@ -231,6 +231,11 @@ class MasterClient {
     BatchGetReplicaList(const std::vector<std::string>& object_keys,
                         const std::string& tenant_id);
 
+    [[nodiscard]] std::vector<tl::expected<bool, ErrorCode>>
+    BatchCheckReplicasLive(const std::vector<std::string>& object_keys,
+                           const std::vector<ReplicaReadToken>& tokens,
+                           const std::string& tenant_id);
+
     /**
      * @brief Starts a put operation
      * @param key Object key

mooncake-store/include/master_service.h

@@ -433,6 +433,16 @@ class MasterService {
     BatchGetReplicaListForAdmin(const std::vector<std::string>& keys,
                                 const TenantId& tenant_id);
 
+    /**
+     * @brief For reads whose lease expired mid-transfer: true when the token's
+     * memory replica is still readable and the object has not been rewritten
+     * since the query. A live master never reuses a replica ID, so the buffer
+     * was not freed in between. Grants no lease.
+     */
+    std::vector<tl::expected<bool, ErrorCode>> BatchCheckReplicasLive(
+        const std::vector<std::string>& keys,
+        const std::vector<ReplicaReadToken>& tokens, const TenantId& tenant_id);
+
     /**
      * @brief Start a put operation for an object
      * @param[out] replica_list Vector to store replica information for the

mooncake-store/include/object_metadata.h

@@ -99,8 +99,14 @@ struct ObjectMetadata {
 
     // Updated by UpsertStart (Case B) to reflect the new writer.
     UUID client_id;
-    // Updated by UpsertStart (Case B) to reset the discard timeout.
+    // Updated by UpsertStart (Case B) to reset the discard timeout. Case B is
+    // the only path that rewrites a memory replica's bytes under its existing
+    // ID, so this also serves as the object's write version.
     std::chrono::system_clock::time_point put_start_time;
+
+    uint64_t WriteVersion() const {
+        return static_cast<uint64_t>(put_start_time.time_since_epoch().count());
+    }
     const ObjectDataType data_type{ObjectDataType::UNKNOWN};
     const std::string group_id;
     const TenantId tenant_id;

mooncake-store/include/pyclient.h

@@ -608,6 +608,9 @@ inline CachedQueryResultResponse to_cached_query_result_response(
     if (query_result->resolved_tenant_id) {
         value.resolved_tenant_id = *query_result->resolved_tenant_id;
     }
+    if (query_result->write_version) {
+        value.write_version = *query_result->write_version;
+    }
     return CachedQueryResultResponse(std::move(value));
 }
 
@@ -623,7 +626,8 @@ inline tl::expected<QueryResult, ErrorCode> from_cached_query_result_response(
                                          cached_result.value.replicas.end()),
         now + std::chrono::milliseconds(cached_result.value.lease_ttl_ms),
         cached_result.value.object_checksum,
-        cached_result.value.resolved_tenant_id);
+        cached_result.value.resolved_tenant_id,
+        cached_result.value.write_version);
 }
 
 inline PyClient::QueryResultCache build_query_result_cache_from_cached_results(

mooncake-store/include/rpc_service.h

@@ -80,6 +80,11 @@ class WrappedMasterService {
     BatchGetReplicaList(const std::vector<std::string>& keys,
                         const std::string& tenant_id = "default");
 
+    std::vector<tl::expected<bool, ErrorCode>> BatchCheckReplicasLive(
+        const std::vector<std::string>& keys,
+        const std::vector<ReplicaReadToken>& tokens,
+        const std::string& tenant_id = "default");
+
     // Read-only admin variants: no lease grants, no promotion, no metric
     // updates.
     std::vector<tl::expected<GetReplicaListResponse, ErrorCode>>

mooncake-store/include/rpc_types.h

@@ -84,6 +84,10 @@ struct GetReplicaListResponse {
     // the exact scope the write path used. Empty when the master predates
     // this field.
     struct_pack::compatible<std::string, 1> resolved_tenant_id;
+    // Identifies the object write the replicas belong to, so a reader whose
+    // lease expired mid-transfer can ask BatchCheckReplicasLive whether the
+    // bytes it read were left untouched. Empty when the master predates it.
+    struct_pack::compatible<uint64_t, 2> write_version;
 
     GetReplicaListResponse() : lease_ttl_ms(0) {}
     GetReplicaListResponse(
@@ -95,7 +99,15 @@ struct GetReplicaListResponse {
           object_checksum(object_checksum_param) {}
 };
 YLT_REFL(GetReplicaListResponse, replicas, lease_ttl_ms, object_checksum,
-         resolved_tenant_id);
+         resolved_tenant_id, write_version);
+
+// The replica a read transferred from and the write version it was queried
+// at.
+struct ReplicaReadToken {
+    ReplicaID replica_id{0};
+    uint64_t write_version{0};
+};
+YLT_REFL(ReplicaReadToken, replica_id, write_version);
 
 struct CachedQueryResultResponse {
     bool success;

mooncake-store/src/client_service.cpp

@@ -1391,7 +1391,8 @@ tl::expected<QueryResult, ErrorCode> Client::Query(
     return tl::expected<QueryResult, ErrorCode>(
         tl::in_place, std::move(result.value().replicas),
         start_time + std::chrono::milliseconds(result.value().lease_ttl_ms),
-        result.value().object_checksum, result.value().resolved_tenant_id);
+        result.value().object_checksum, result.value().resolved_tenant_id,
+        result.value().write_version);
 }
 
 std::vector<tl::expected<QueryResult, ErrorCode>> Client::BatchQuery(
@@ -1426,7 +1427,8 @@ std::vector<tl::expected<QueryResult, ErrorCode>> Client::BatchQuery(
                 start_time +
                     std::chrono::milliseconds(response[i].value().lease_ttl_ms),
                 response[i].value().object_checksum,
-                response[i].value().resolved_tenant_id);
+                response[i].value().resolved_tenant_id,
+                response[i].value().write_version);
         } else {
             results.emplace_back(tl::unexpected(response[i].error()));
         }
@@ -1585,7 +1587,8 @@ tl::expected<void, ErrorCode> Client::Get(const std::string& object_key,
         ProcessSlicesAsync(object_key, slices, replica);
     }
 
-    if (query_result.IsLeaseExpired()) {
+    if (query_result.IsLeaseExpired() &&
+        !CheckLeaseExpiredReads({object_key}, {&query_result})[0]) {
         LOG(WARNING) << "lease_expired_before_data_transfer_completed key="
                      << object_key;
         return tl::unexpected(ErrorCode::LEASE_EXPIRED);
@@ -1631,7 +1634,8 @@ tl::expected<void, ErrorCode> Client::Get(const std::string& object_key,
         LOG(ERROR) << "transfer_read_range_failed key=" << object_key;
         return tl::unexpected(err);
     }
-    if (query_result.IsLeaseExpired()) {
+    if (query_result.IsLeaseExpired() &&
+        !CheckLeaseExpiredReads({object_key}, {&query_result})[0]) {
         LOG(WARNING) << "lease_expired_before_data_transfer_completed key="
                      << object_key;
         return tl::unexpected(ErrorCode::LEASE_EXPIRED);
@@ -1800,14 +1804,7 @@ std::vector<tl::expected<void, ErrorCode>> Client::BatchGetWhenPreferSameNode(
     // A successful transfer is only valid while its read lease is live.
     // Check once all segment batches have completed, as in the regular
     // BatchGet path, and preserve errors from validation or transfer.
-    auto now = std::chrono::steady_clock::now();
-    for (size_t i = 0; i < object_keys.size(); ++i) {
-        if (results[i].has_value() && query_results[i].IsLeaseExpired(now)) {
-            LOG(WARNING) << "lease_expired_before_data_transfer_completed key="
-                         << object_keys[i];
-            results[i] = tl::unexpected(ErrorCode::LEASE_EXPIRED);
-        }
-    }
+    FailLeaseExpiredReads(object_keys, query_results, results);
     return results;
 }
 
@@ -2050,17 +2047,7 @@ std::vector<tl::expected<void, ErrorCode>> Client::BatchGet(
         }
     }
 
-    // As lease expired is a rare case, we check all the results with the same
-    // time_point to avoid too many syscalls
-    std::chrono::steady_clock::time_point now =
-        std::chrono::steady_clock::now();
-    for (size_t i = 0; i < object_keys.size(); ++i) {
-        if (results[i].has_value() && query_results[i].IsLeaseExpired(now)) {
-            LOG(WARNING) << "lease_expired_before_data_transfer_completed key="
-                         << object_keys[i];
-            results[i] = tl::unexpected(ErrorCode::LEASE_EXPIRED);
-        }
-    }
+    FailLeaseExpiredReads(object_keys, query_results, results);
 
     auto us_batch_get = std::chrono::duration_cast<std::chrono::microseconds>(
                             std::chrono::steady_clock::now() - t0_batch_get)
@@ -5682,6 +5669,74 @@ ErrorCode Client::FindFirstCompleteReplica(
     return ErrorCode::INVALID_REPLICA;
 }
 
+std::vector<bool> Client::CheckLeaseExpiredReads(
+    const std::vector<std::string>& keys,
+    const std::vector<const QueryResult*>& query_results) {
+    struct Checks {
+        std::vector<std::string> keys;
+        std::vector<ReplicaReadToken> tokens;
+        std::vector<size_t> indexes;
+    };
+    // One RPC per tenant scope; the keys of one BatchQuery share a scope.
+    std::unordered_map<std::string, Checks> by_tenant;
+    for (size_t i = 0; i < keys.size(); ++i) {
+        const QueryResult& query = *query_results[i];
+        // Same pick as the read paths, so the token names the replica read.
+        Replica::Descriptor replica;
+        if (!query.write_version || !query.resolved_tenant_id ||
+            FindFirstCompleteReplica(query.replicas, replica) !=
+                ErrorCode::OK ||
+            !replica.is_memory_replica()) {
+            continue;
+        }
+        auto& checks = by_tenant[*query.resolved_tenant_id];
+        checks.keys.push_back(keys[i]);
+        checks.tokens.push_back({replica.id, *query.write_version});
+        checks.indexes.push_back(i);
+    }
+
+    std::vector<bool> live(keys.size(), false);
+    for (const auto& [tenant_id, checks] : by_tenant) {
+        auto results = master_client_.BatchCheckReplicasLive(
+            checks.keys, checks.tokens, tenant_id);
+        for (size_t j = 0; j < results.size() && j < checks.indexes.size();
+             ++j) {
+            live[checks.indexes[j]] = results[j].value_or(false);
+        }
+    }
+    return live;
+}
+
+void Client::FailLeaseExpiredReads(
+    const std::vector<std::string>& keys,
+    const std::vector<QueryResult>& query_results,
+    std::vector<tl::expected<void, ErrorCode>>& results) {
+    // Lease expiry is rare, so read the clock once for the whole batch.
+    auto now = std::chrono::steady_clock::now();
+    std::vector<size_t> expired;
+    std::vector<std::string> expired_keys;
+    std::vector<const QueryResult*> expired_queries;
+    for (size_t i = 0; i < keys.size(); ++i) {
+        if (results[i].has_value() && query_results[i].IsLeaseExpired(now)) {
+            expired.push_back(i);
+            expired_keys.push_back(keys[i]);
+            expired_queries.push_back(&query_results[i]);
+        }
+    }
+    if (expired.empty()) {
+        return;
+    }
+
+    const auto live = CheckLeaseExpiredReads(expired_keys, expired_queries);
+    for (size_t j = 0; j < expired.size(); ++j) {
+        if (!live[j]) {
+            LOG(WARNING) << "lease_expired_before_data_transfer_completed key="
+                         << keys[expired[j]];
+            results[expired[j]] = tl::unexpected(ErrorCode::LEASE_EXPIRED);
+        }
+    }
+}
+
 tl::expected<Replica::Descriptor, ErrorCode> Client::GetPreferredReplica(
     const std::vector<Replica::Descriptor>& replica_list) {
     if (replica_list.empty()) {

mooncake-store/src/master_client.cpp

@@ -72,6 +72,11 @@ struct RpcNameTraits<&WrappedMasterService::BatchGetReplicaList> {
     static constexpr const char* value = "BatchGetReplicaList";
 };
 
+template <>
+struct RpcNameTraits<&WrappedMasterService::BatchCheckReplicasLive> {
+    static constexpr const char* value = "BatchCheckReplicasLive";
+};
+
 template <>
 struct RpcNameTraits<&WrappedMasterService::PutStart> {
     static constexpr const char* value = "PutStart";
@@ -663,6 +668,20 @@ MasterClient::BatchGetReplicaList(const std::vector<std::string>& object_keys,
     return result;
 }
 
+std::vector<tl::expected<bool, ErrorCode>> MasterClient::BatchCheckReplicasLive(
+    const std::vector<std::string>& object_keys,
+    const std::vector<ReplicaReadToken>& tokens, const std::string& tenant_id) {
+    ScopedVLogTimer timer(1, "MasterClient::BatchCheckReplicasLive");
+    timer.LogRequest("keys_count=", object_keys.size(),
+                     ", tenant_id=", tenant_id);
+
+    auto result =
+        invoke_batch_rpc<&WrappedMasterService::BatchCheckReplicasLive, bool>(
+            object_keys.size(), object_keys, tokens, tenant_id);
+    timer.LogResponse("result=", result.size(), " keys");
+    return result;
+}
+
 tl::expected<std::vector<Replica::Descriptor>, ErrorCode>
 MasterClient::PutStart(const std::string& key,
                        const std::vector<size_t>& slice_lengths,

mooncake-store/src/master_service.cpp

@@ -4645,9 +4645,11 @@ auto MasterService::ReadReplicaList(const TenantId& tenant_id,
         }
     }
 
-    return GetReplicaListResponse(std::move(replica_list),
-                                  default_kv_lease_ttl_,
-                                  metadata.GetPart(0).checksum());
+    GetReplicaListResponse response(std::move(replica_list),
+                                    default_kv_lease_ttl_,
+                                    metadata.GetPart(0).checksum());
+    response.write_version = metadata.WriteVersion();
+    return response;
 }
 
 void MasterService::ReadFollowUps::Run(MasterService& service) const {
@@ -4746,6 +4748,36 @@ MasterService::BatchGetReplicaListForAdmin(const std::vector<std::string>& keys,
     return results;
 }
 
+std::vector<tl::expected<bool, ErrorCode>>
+MasterService::BatchCheckReplicasLive(
+    const std::vector<std::string>& keys,
+    const std::vector<ReplicaReadToken>& tokens, const TenantId& tenant_id) {
+    if (tokens.size() != keys.size()) {
+        return std::vector<tl::expected<bool, ErrorCode>>(
+            keys.size(), tl::make_unexpected(ErrorCode::INVALID_PARAMS));
+    }
+    const TenantId& normalized_tenant = ResolveRequestTenantId(tenant_id);
+    std::vector<tl::expected<bool, ErrorCode>> results(keys.size(), false);
+    std::shared_lock<std::shared_mutex> shared_lock(master_state_mutex_);
+    for (size_t i = 0; i < keys.size(); ++i) {
+        MetadataAccessorRO accessor(this, normalized_tenant, keys[i]);
+        if (!accessor.Exists()) {
+            continue;
+        }
+        const auto& metadata = accessor.Get();
+        const auto& token = tokens[i];
+        // Disk replicas are excluded: they are read by key, and a LOCAL_DISK
+        // replica can be re-registered under its existing ID.
+        results[i] = metadata.WriteVersion() == token.write_version &&
+                     metadata.HasReplica([&](const Replica& replica) {
+                         return replica.id() == token.replica_id &&
+                                replica.is_memory_replica() &&
+                                IsReplicaReadable(replica);
+                     });
+    }
+    return results;
+}
+
 auto MasterService::AllocateReplicas(const std::string& key,
                                      uint64_t value_length,
                                      const ReplicateConfig& config,