diff --git a/mooncake-store/include/client.h b/mooncake-store/include/client.h index 725c890e5..00239d63b 100644 --- a/mooncake-store/include/client.h +++ b/mooncake-store/include/client.h @@ -364,4 +364,34 @@ class Client { UUID client_id_; }; +/** + * @brief Fluent builder for configuring a mooncake::Client instance. + * + * Provides readable, type-safe setters with sensible defaults so callers can + * specify only the options they need while reusing existing Client::Create + * logic under the hood. + */ +class MooncakeStoreBuilder { + public: + MooncakeStoreBuilder& WithLocalHostname(std::string local_hostname); + MooncakeStoreBuilder& WithMetadataConnectionString( + std::string metadata_connstring); + MooncakeStoreBuilder& WithProtocol(std::string protocol); + MooncakeStoreBuilder& WithProtocolArgs(std::string protocol_args); + MooncakeStoreBuilder& WithMasterEndpoint(std::string master_server_entry); + MooncakeStoreBuilder& WithExistingTransferEngine( + std::shared_ptr transfer_engine); + + [[nodiscard]] tl::expected, std::string> Build() + const; + + private: + std::optional local_hostname_; + std::optional metadata_connstring_; + std::string protocol_ = "tcp"; + std::optional protocol_args_; + std::string master_server_entry_ = kDefaultMasterAddress; + std::shared_ptr transfer_engine_ = nullptr; +}; + } // namespace mooncake diff --git a/mooncake-store/src/client.cpp b/mooncake-store/src/client.cpp index cf5ccbabf..9a02c7ac7 100644 --- a/mooncake-store/src/client.cpp +++ b/mooncake-store/src/client.cpp @@ -7,6 +7,7 @@ #include #include #include +#include #include #include #include @@ -406,6 +407,80 @@ std::optional> Client::Create( return client; } +MooncakeStoreBuilder& MooncakeStoreBuilder::WithLocalHostname( + std::string local_hostname) { + local_hostname_ = std::move(local_hostname); + return *this; +} + +MooncakeStoreBuilder& MooncakeStoreBuilder::WithMetadataConnectionString( + std::string metadata_connstring) { + metadata_connstring_ = std::move(metadata_connstring); + return *this; +} + +MooncakeStoreBuilder& MooncakeStoreBuilder::WithProtocol(std::string protocol) { + protocol_ = std::move(protocol); + return *this; +} + +MooncakeStoreBuilder& MooncakeStoreBuilder::WithProtocolArgs( + std::string protocol_args) { + // Can add some other engine arguments + protocol_args_ = std::move(protocol_args); + return *this; +} + +MooncakeStoreBuilder& MooncakeStoreBuilder::WithMasterEndpoint( + std::string master_server_entry) { + master_server_entry_ = std::move(master_server_entry); + return *this; +} + +MooncakeStoreBuilder& MooncakeStoreBuilder::WithExistingTransferEngine( + std::shared_ptr transfer_engine) { + transfer_engine_ = std::move(transfer_engine); + return *this; +} + +tl::expected, std::string> MooncakeStoreBuilder::Build() + const { + std::vector missing; + if (!local_hostname_) { + missing.emplace_back("local_hostname"); + } + if (!metadata_connstring_) { + missing.emplace_back("metadata_connstring"); + } + + if (!missing.empty()) { + std::string joined; + joined.reserve(missing.size() * 16); + for (size_t i = 0; i < missing.size(); ++i) { + if (i != 0) { + joined.append(", "); + } + joined.append(missing[0]); + for (size_t i = 1; i < missing.size(); ++i) { + joined.append(", "); + joined.append(missing[i]); + } + } + auto error_msg = + "MooncakeStoreBuilder missing required fields: " + joined; + LOG(ERROR) << error_msg; + return tl::make_unexpected(std::move(error_msg)); + } + auto client = + Client::Create(*local_hostname_, *metadata_connstring_, protocol_, + protocol_args_, master_server_entry_, transfer_engine_); + if (!client) { + std::string error_msg = "Client creation failed"; + return tl::make_unexpected(std::move(error_msg)); + } + return *client; +} + tl::expected Client::Get(const std::string& object_key, std::vector& slices) { auto query_result = Query(object_key); diff --git a/mooncake-store/src/pybind_client.cpp b/mooncake-store/src/pybind_client.cpp index 996d4aa75..93e4eb6a3 100644 --- a/mooncake-store/src/pybind_client.cpp +++ b/mooncake-store/src/pybind_client.cpp @@ -177,7 +177,6 @@ tl::expected PyClient::setup_internal( std::optional device_name = (rdma_devices.empty() ? std::nullopt : std::make_optional(rdma_devices)); - auto client_opt = mooncake::Client::Create( this->local_hostname, metadata_server, protocol, device_name, master_server_addr, transfer_engine); diff --git a/mooncake-store/tests/client_integration_test.cpp b/mooncake-store/tests/client_integration_test.cpp index e42d50a3b..15be768d1 100644 --- a/mooncake-store/tests/client_integration_test.cpp +++ b/mooncake-store/tests/client_integration_test.cpp @@ -22,6 +22,68 @@ DEFINE_uint64(default_kv_lease_ttl, mooncake::DEFAULT_DEFAULT_KV_LEASE_TTL, namespace mooncake { namespace testing { +class MooncakeStoreBuilderTest : public ::testing::Test { + protected: + void SetUp() override { + google::InitGoogleLogging("MooncakeStoreBuilderTest"); + FLAGS_logtostderr = true; + } + + void TearDown() override { google::ShutdownGoogleLogging(); } +}; + +TEST_F(MooncakeStoreBuilderTest, MissingAllRequiredFields) { + MooncakeStoreBuilder builder; + + auto result = builder.Build(); + + ASSERT_FALSE(result.has_value()); + const std::string& error = result.error(); + EXPECT_NE(error.find("local_hostname"), std::string::npos); + EXPECT_NE(error.find("metadata_connstring"), std::string::npos); +} + +TEST_F(MooncakeStoreBuilderTest, MissingLocalHostnameOnly) { + MooncakeStoreBuilder builder; + builder.WithMetadataConnectionString("metastore:1234"); + + auto result = builder.Build(); + + ASSERT_FALSE(result.has_value()); + const std::string& error = result.error(); + EXPECT_NE(error.find("local_hostname"), std::string::npos); + EXPECT_EQ(error.find("metadata_connstring"), std::string::npos); +} + +TEST_F(MooncakeStoreBuilderTest, MissingMetadataConnectionStringOnly) { + MooncakeStoreBuilder builder; + builder.WithLocalHostname("localhost:1234"); + + auto result = builder.Build(); + + ASSERT_FALSE(result.has_value()); + const std::string& error = result.error(); + EXPECT_NE(error.find("metadata_connstring"), std::string::npos); + EXPECT_EQ(error.find("local_hostname"), std::string::npos); +} + +TEST_F(MooncakeStoreBuilderTest, ProvidedAllRequiredFields) { + MooncakeStoreBuilder builder; + builder.WithLocalHostname("localhost:1234"); + builder.WithMetadataConnectionString("metadata:5678"); + + auto result = builder.Build(); + + if (result.has_value()) { + SUCCEED(); + } else { + const std::string& error = result.error(); + EXPECT_EQ(error.find("missing required fields"), std::string::npos); + EXPECT_EQ(error.find("local_hostname"), std::string::npos); + EXPECT_EQ(error.find("metadata_connstring"), std::string::npos); + } +} + class ClientIntegrationTest : public ::testing::Test { protected: static std::shared_ptr CreateClient(const std::string& host_name) {