feat(store): add batch exist support for master

xiaguan · xiaguan · commit 84ce245e3b81 · 2025-06-23T17:52:41.000+08:00
diff --git a/mooncake-integration/store/store_py.cpp b/mooncake-integration/store/store_py.cpp
@@ -494,6 +494,48 @@ int DistributedObjectStore::isExist(const std::string &key) {
     return toInt(err);                                 // Error
 }
 
+std::vector<int> DistributedObjectStore::batchIsExist(
+    const std::vector<std::string> &keys) {
+    std::vector<int> results;
+
+    if (!client_) {
+        LOG(ERROR) << "Client is not initialized";
+        results.resize(keys.size(), -1);  // Fill with error codes
+        return results;
+    }
+
+    if (keys.empty()) {
+        LOG(WARNING) << "Empty keys vector provided to batchIsExist";
+        return results;  // Return empty vector
+    }
+
+    std::vector<ErrorCode> exist_results;
+    ErrorCode batch_err = client_->BatchIsExist(keys, exist_results);
+
+    results.resize(keys.size());
+
+    if (batch_err != ErrorCode::OK) {
+        LOG(ERROR) << "BatchIsExist operation failed with error: "
+                   << toString(batch_err);
+        // Fill all results with error code
+        std::fill(results.begin(), results.end(), toInt(batch_err));
+        return results;
+    }
+
+    // Convert ErrorCode results to int results
+    for (size_t i = 0; i < keys.size(); ++i) {
+        if (exist_results[i] == ErrorCode::OK) {
+            results[i] = 1;  // Exists
+        } else if (exist_results[i] == ErrorCode::OBJECT_NOT_FOUND) {
+            results[i] = 0;  // Does not exist
+        } else {
+            results[i] = toInt(exist_results[i]);  // Error
+        }
+    }
+
+    return results;
+}
+
 int64_t DistributedObjectStore::getSize(const std::string &key) {
     if (!client_) {
         LOG(ERROR) << "Client is not initialized";
@@ -775,6 +817,10 @@ PYBIND11_MODULE(store, m) {
              py::call_guard<py::gil_scoped_release>())
         .def("is_exist", &DistributedObjectStore::isExist,
              py::call_guard<py::gil_scoped_release>())
+        .def("batch_is_exist", &DistributedObjectStore::batchIsExist,
+             py::call_guard<py::gil_scoped_release>(), py::arg("keys"),
+             "Check if multiple objects exist. Returns list of results: 1 if "
+             "exists, 0 if not exists, -1 if error")
         .def("close", &DistributedObjectStore::tearDownAll)
         .def("get_size", &DistributedObjectStore::getSize,
              py::call_guard<py::gil_scoped_release>())
diff --git a/mooncake-integration/store/store_py.h b/mooncake-integration/store/store_py.h
@@ -164,6 +164,14 @@ class DistributedObjectStore {
      */
     int isExist(const std::string &key);
 
+    /**
+     * @brief Check if multiple objects exist
+     * @param keys Vector of keys to check
+     * @return Vector of existence results: 1 if exists, 0 if not exists, -1 if
+     * error
+     */
+    std::vector<int> batchIsExist(const std::vector<std::string> &keys);
+
     /**
      * @brief Get the size of an object
      * @param key Key of the object
diff --git a/mooncake-store/include/client.h b/mooncake-store/include/client.h
@@ -1,18 +1,18 @@
 #pragma once
 
+#include <boost/functional/hash.hpp>
 #include <memory>
 #include <mutex>
 #include <optional>
 #include <string>
 #include <vector>
-#include <boost/functional/hash.hpp>
 
+#include "ha_helper.h"
 #include "master_client.h"
 #include "rpc_service.h"
 #include "transfer_engine.h"
 #include "transfer_task.h"
 #include "types.h"
-#include "ha_helper.h"
 
 namespace mooncake {
 
@@ -190,6 +190,15 @@ class Client {
      */
     ErrorCode IsExist(const std::string& key);
 
+    /**
+     * @brief Checks if multiple objects exist
+     * @param keys Vector of keys to check
+     * @param exist_results Output vector of existence results for each key
+     * @return ErrorCode indicating success/failure of the batch operation
+     */
+    ErrorCode BatchIsExist(const std::vector<std::string>& keys,
+                           std::vector<ErrorCode>& exist_results);
+
    private:
     /**
      * @brief Private constructor to enforce creation through Create() method
diff --git a/mooncake-store/include/master_client.h b/mooncake-store/include/master_client.h
@@ -38,8 +38,15 @@ class MasterClient {
      * @param object_key Key to query
      * @return ErrorCode indicating exist or not
      */
-    [[nodiscard]] ExistKeyResponse ExistKey(
-        const std::string& object_key);
+    [[nodiscard]] ExistKeyResponse ExistKey(const std::string& object_key);
+
+    /**
+     * @brief Checks if multiple objects exist
+     * @param object_keys Vector of keys to query
+     * @return BatchExistReponse containing existence status for each key
+     */
+    [[nodiscard]] BatchExistReponse BatchExistKey(
+        const std::vector<std::string>& object_keys);
 
     /**
      * @brief Gets object metadata without transferring data
@@ -136,8 +143,8 @@ class MasterClient {
      * @param client_id The uuid of the client
      * @return ErrorCode indicating success/failure
      */
-    [[nodiscard]] MountSegmentResponse MountSegment(
-        const Segment& segment, const UUID& client_id);
+    [[nodiscard]] MountSegmentResponse MountSegment(const Segment& segment,
+                                                    const UUID& client_id);
 
     /**
      * @brief Re-mount segments, invoked when the client is the first time to
@@ -157,8 +164,8 @@ class MasterClient {
      * @param client_id The uuid of the client
      * @return ErrorCode indicating success/failure
      */
-    [[nodiscard]] UnmountSegmentResponse UnmountSegment(
-        const UUID& segment_id, const UUID& client_id);
+    [[nodiscard]] UnmountSegmentResponse UnmountSegment(const UUID& segment_id,
+                                                        const UUID& client_id);
 
     /**
      * @brief Pings master to check its availability
diff --git a/mooncake-store/include/master_service.h b/mooncake-store/include/master_service.h
@@ -105,6 +105,8 @@ class MasterService {
      */
     ErrorCode ExistKey(const std::string& key);
 
+    std::vector<ErrorCode> BatchExistKey(const std::vector<std::string>& keys);
+
     /**
      * @brief Fetch all keys
      * @return ErrorCode::OK if exists
diff --git a/mooncake-store/include/rpc_service.h b/mooncake-store/include/rpc_service.h
@@ -97,6 +97,11 @@ struct PingResponse {
 };
 YLT_REFL(PingResponse, view_version, error_code)
 
+struct BatchExistReponse {
+    std::vector<ErrorCode> exist_responses;
+};
+YLT_REFL(BatchExistReponse, exist_responses)
+
 constexpr uint64_t kMetricReportIntervalSeconds = 10;
 
 class WrappedMasterService {
@@ -274,6 +279,15 @@ class WrappedMasterService {
         return response;
     }
 
+    BatchExistReponse BatchExistKey(const std::vector<std::string>& keys) {
+        ScopedVLogTimer timer(1, "BatchExistKey");
+        timer.LogRequest("keys_count=", keys.size());
+
+        BatchExistReponse response{master_service_.BatchExistKey(keys)};
+        timer.LogResponseJson(response);
+        return response;
+    }
+
     GetReplicaListResponse GetReplicaList(const std::string& key) {
         ScopedVLogTimer timer(1, "GetReplicaList");
         timer.LogRequest("key=", key);
@@ -462,7 +476,8 @@ class WrappedMasterService {
         return response;
     }
 
-    MountSegmentResponse MountSegment(const Segment& segment, const UUID& client_id) {
+    MountSegmentResponse MountSegment(const Segment& segment,
+                                      const UUID& client_id) {
         ScopedVLogTimer timer(1, "MountSegment");
         timer.LogRequest("base=", segment.base, ", size=", segment.size,
                          ", segment_name=", segment.name, ", id=", segment.id);
@@ -504,15 +519,17 @@ class WrappedMasterService {
         return response;
     }
 
-    UnmountSegmentResponse UnmountSegment(const UUID& segment_id, const UUID& client_id) {
+    UnmountSegmentResponse UnmountSegment(const UUID& segment_id,
+                                          const UUID& client_id) {
         ScopedVLogTimer timer(1, "UnmountSegment");
         timer.LogRequest("segment_id=", segment_id);
 
         // Increment request metric
         MasterMetricManager::instance().inc_unmount_segment_requests();
 
         UnmountSegmentResponse response;
-        response.error_code = master_service_.UnmountSegment(segment_id, client_id);
+        response.error_code =
+            master_service_.UnmountSegment(segment_id, client_id);
 
         // Track failures if needed
         if (response.error_code != ErrorCode::OK) {
@@ -581,6 +598,8 @@ inline void RegisterRpcService(
         &wrapped_master_service);
     server.register_handler<&mooncake::WrappedMasterService::Ping>(
         &wrapped_master_service);
+    server.register_handler<&mooncake::WrappedMasterService::BatchExistKey>(
+        &wrapped_master_service);
 }
 
 }  // namespace mooncake
diff --git a/mooncake-store/src/client.cpp b/mooncake-store/src/client.cpp
@@ -43,8 +43,8 @@ Client::~Client() {
     }
 
     for (auto& segment : segments_to_unmount) {
-        auto err_code = UnmountSegment(reinterpret_cast<void*>(segment.base),
-                                       segment.size);
+        auto err_code =
+            UnmountSegment(reinterpret_cast<void*>(segment.base), segment.size);
         if (err_code != ErrorCode::OK) {
             LOG(ERROR) << "Failed to unmount segment: " << toString(err_code);
         }
@@ -148,8 +148,7 @@ ErrorCode Client::ConnectToMaster(const std::string& master_server_entry) {
         // Start Ping thread to monitor master view changes and remount segments
         // if needed
         ping_running_ = true;
-        ping_thread_ =
-            std::thread(&Client::PingThreadFunc, this);
+        ping_thread_ = std::thread(&Client::PingThreadFunc, this);
 
         return ErrorCode::OK;
     } else {
@@ -212,7 +211,7 @@ std::optional<std::shared_ptr<Client>> Client::Create(
 
     // Initialize transfer engine
     err = client->InitTransferEngine(local_hostname, metadata_connstring,
-                                    protocol, protocol_args);
+                                     protocol, protocol_args);
     if (err != ErrorCode::OK) {
         LOG(ERROR) << "Failed to initialize transfer engine";
         return std::nullopt;
@@ -581,8 +580,7 @@ ErrorCode Client::MountSegment(const void* buffer, size_t size) {
     Segment segment(generate_uuid(), local_hostname_,
                     reinterpret_cast<uintptr_t>(buffer), size);
 
-    ErrorCode err =
-        master_client_.MountSegment(segment, client_id_).error_code;
+    ErrorCode err = master_client_.MountSegment(segment, client_id_).error_code;
     if (err != ErrorCode::OK) {
         LOG(ERROR) << "mount_segment_to_master_failed base=" << buffer
                    << " size=" << size << ", error=" << err;
@@ -660,6 +658,42 @@ ErrorCode Client::IsExist(const std::string& key) {
     return response.error_code;
 }
 
+ErrorCode Client::BatchIsExist(const std::vector<std::string>& keys,
+                               std::vector<ErrorCode>& exist_results) {
+    // Check for duplicate keys
+    std::unordered_set<std::string> seen;
+    for (const auto& key : keys) {
+        if (!seen.insert(key).second) {
+            LOG(ERROR) << "Duplicate key not supported for Batch API, key: "
+                       << key;
+            return ErrorCode::INVALID_PARAMS;
+        }
+    }
+
+    auto response = master_client_.BatchExistKey(keys);
+
+    // Resize the output vector to match the number of keys
+    exist_results.resize(keys.size());
+
+    // Check if we got the expected number of responses
+    if (response.exist_responses.size() != keys.size()) {
+        LOG(ERROR) << "BatchExistKey response size mismatch. Expected: "
+                   << keys.size()
+                   << ", Got: " << response.exist_responses.size();
+        // Fill with RPC_FAIL error codes
+        std::fill(exist_results.begin(), exist_results.end(),
+                  ErrorCode::RPC_FAIL);
+        return ErrorCode::RPC_FAIL;
+    }
+
+    // Copy the error codes from the response
+    for (size_t i = 0; i < keys.size(); ++i) {
+        exist_results[i] = response.exist_responses[i];
+    }
+
+    return ErrorCode::OK;
+}
+
 ErrorCode Client::TransferData(
     const std::vector<AllocatedBuffer::Descriptor>& handles,
     std::vector<Slice>& slices, TransferRequest::OpCode op_code) {
diff --git a/mooncake-store/src/master_client.cpp b/mooncake-store/src/master_client.cpp
@@ -39,8 +39,8 @@ ExistKeyResponse MasterClient::ExistKey(const std::string& object_key) {
 
     auto request_result =
         client_.send_request<&WrappedMasterService::ExistKey>(object_key);
-    std::optional<ExistKeyResponse> result = coro::syncAwait(
-        [&]() -> coro::Lazy<std::optional<ExistKeyResponse>> {
+    std::optional<ExistKeyResponse> result =
+        coro::syncAwait([&]() -> coro::Lazy<std::optional<ExistKeyResponse>> {
             auto result = co_await co_await request_result;
             if (!result) {
                 LOG(ERROR) << "Failed to check key existence: "
@@ -60,6 +60,38 @@ ExistKeyResponse MasterClient::ExistKey(const std::string& object_key) {
     return result.value();
 }
 
+BatchExistReponse MasterClient::BatchExistKey(
+    const std::vector<std::string>& object_keys) {
+    ScopedVLogTimer timer(1, "MasterClient::BatchExistKey");
+    timer.LogRequest("keys_count=", object_keys.size());
+
+    auto request_result =
+        client_.send_request<&WrappedMasterService::BatchExistKey>(object_keys);
+    std::optional<BatchExistReponse> result =
+        coro::syncAwait([&]() -> coro::Lazy<std::optional<BatchExistReponse>> {
+            auto result = co_await co_await request_result;
+            if (!result) {
+                LOG(ERROR) << "Failed to check batch key existence: "
+                           << result.error().msg;
+                co_return std::nullopt;
+            }
+            co_return result->result();
+        }());
+
+    if (!result) {
+        BatchExistReponse response;
+        response.exist_responses.resize(object_keys.size());
+        for (auto& exist_response : response.exist_responses) {
+            exist_response = ErrorCode::RPC_FAIL;
+        }
+        timer.LogResponseJson(response);
+        return response;
+    }
+
+    timer.LogResponseJson(result.value());
+    return result.value();
+}
+
 GetReplicaListResponse MasterClient::GetReplicaList(
     const std::string& object_key) {
     ScopedVLogTimer timer(1, "MasterClient::GetReplicaList");
@@ -315,7 +347,7 @@ RemoveAllResponse MasterClient::RemoveAll() {
             auto result = co_await co_await request_result;
             if (!result) {
                 LOG(ERROR) << "Failed to remove all objects: "
-                          << result.error().msg;
+                           << result.error().msg;
                 co_return std::nullopt;
             }
             co_return result->result();
@@ -363,7 +395,8 @@ MountSegmentResponse MasterClient::MountSegment(const Segment& segment,
 ReMountSegmentResponse MasterClient::ReMountSegment(
     const std::vector<Segment>& segments, const UUID& client_id) {
     ScopedVLogTimer timer(1, "MasterClient::ReMountSegment");
-    timer.LogRequest("segments_num=", segments.size(), ", client_id=", client_id);
+    timer.LogRequest("segments_num=", segments.size(),
+                     ", client_id=", client_id);
 
     std::optional<ReMountSegmentResponse> result =
         syncAwait([&]() -> coro::Lazy<std::optional<ReMountSegmentResponse>> {
@@ -431,7 +464,8 @@ PingResponse MasterClient::Ping(const UUID& client_id) {
         }());
 
     if (!result) {
-        auto response = PingResponse{0, ClientStatus::UNDEFINED, ErrorCode::RPC_FAIL};
+        auto response =
+            PingResponse{0, ClientStatus::UNDEFINED, ErrorCode::RPC_FAIL};
         timer.LogResponseJson(response);
         return response;
     }
diff --git a/mooncake-store/src/master_service.cpp b/mooncake-store/src/master_service.cpp
diff --git a/mooncake-store/tests/client_integration_test.cpp b/mooncake-store/tests/client_integration_test.cpp
diff --git a/mooncake-store/tests/master_service_test.cpp b/mooncake-store/tests/master_service_test.cpp
diff --git a/mooncake-wheel/tests/test_distributed_object_store.py b/mooncake-wheel/tests/test_distributed_object_store.py
