From 41b1df795425faa7f95f59eb5b3c46f809bd329f Mon Sep 17 00:00:00 2001 From: ykwd Date: Wed, 11 Jun 2025 16:37:05 +0800 Subject: [PATCH] [Store] Add initial support for master high availability failover (#451) * A temp version. Better to continue development after merging the latest main branch * Temp version to merge the latest main branch * Allow optional use HA mode, in default use non-HA mode. Fix a minor metrics bug. * Refactor the etcd_helper * refactor ha_helper * Add some unit tests. Refactor the code * Update cmakelists: build etcd_wrapper in default * Fix ci problems. Compile etcd wrapper only when use_etcd or with_store are set. * Update python config relating to mooncake-store client * make some blocking etcd helper function cancellable. bug fix: add string name of new errors that will be used in tostring. * Refactor etcd related code * Bug fix * Add basic masterviewhelper unit tests * In ci flow, install and start etcd to run HA feature unit test. * Fix a ci bug * Reuse master_server_address parameter and remove enable_ha parameter. * Format the code. Fix a minor bug. * Handle the error case: the coro server may fail to start or return internal error. --- .github/workflows/ci.yml | 10 + mooncake-common/CMakeLists.txt | 4 +- mooncake-common/etcd/etcd_wrapper.go | 311 +++++++++++++++++- mooncake-store/CMakeLists.txt | 5 +- mooncake-store/include/client.h | 22 +- mooncake-store/include/etcd_helper.h | 102 ++++++ mooncake-store/include/ha_helper.h | 111 +++++++ mooncake-store/include/master_client.h | 8 + .../include/master_metric_manager.h | 3 + mooncake-store/include/rpc_service.h | 89 ++++- mooncake-store/include/types.h | 12 + mooncake-store/src/CMakeLists.txt | 8 +- mooncake-store/src/client.cpp | 163 ++++++++- mooncake-store/src/etcd_helper.cpp | 151 +++++++++ mooncake-store/src/ha_helper.cpp | 181 ++++++++++ mooncake-store/src/master.cpp | 103 +++--- mooncake-store/src/master_client.cpp | 28 +- mooncake-store/src/master_metric_manager.cpp | 16 +- mooncake-store/src/types.cpp | 5 + mooncake-store/tests/CMakeLists.txt | 13 +- .../tests/high_availability_test.cpp | 209 ++++++++++++ 21 files changed, 1453 insertions(+), 101 deletions(-) create mode 100644 mooncake-store/include/etcd_helper.h create mode 100644 mooncake-store/include/ha_helper.h create mode 100644 mooncake-store/src/etcd_helper.cpp create mode 100644 mooncake-store/src/ha_helper.cpp create mode 100644 mooncake-store/tests/high_availability_test.cpp diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index bb15f088..e7172855 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -23,6 +23,16 @@ jobs: with: python-version: ${{ matrix.python-version }} + - name: Install and start etcd + run: | + wget https://github.com/etcd-io/etcd/releases/download/v3.6.1/etcd-v3.6.1-linux-amd64.tar.gz + tar xzf etcd-v3.6.1-linux-amd64.tar.gz + sudo mv etcd-v3.6.1-linux-amd64/etcd* /usr/local/bin/ + etcd --advertise-client-urls http://127.0.0.1:2379 --listen-client-urls http://127.0.0.1:2379 & + sleep 3 # Give etcd time to start + etcdctl --endpoints=http://127.0.0.1:2379 endpoint health + shell: bash + - name: Free up disk space run: | sudo rm -rf /usr/share/dotnet diff --git a/mooncake-common/CMakeLists.txt b/mooncake-common/CMakeLists.txt index 6041e8a8..4cfaf77d 100644 --- a/mooncake-common/CMakeLists.txt +++ b/mooncake-common/CMakeLists.txt @@ -1,3 +1,3 @@ -if (USE_ETCD AND NOT USE_ETCD_LEGACY) - add_subdirectory(etcd) +if ((USE_ETCD AND NOT USE_ETCD_LEGACY) OR WITH_STORE) + add_subdirectory(etcd) endif() \ No newline at end of file diff --git a/mooncake-common/etcd/etcd_wrapper.go b/mooncake-common/etcd/etcd_wrapper.go index 6933d6e7..ea00eb57 100644 --- a/mooncake-common/etcd/etcd_wrapper.go +++ b/mooncake-common/etcd/etcd_wrapper.go @@ -2,29 +2,44 @@ package main /* #include +#include +#include */ import "C" import ( "context" + "strings" "sync" "time" clientv3 "go.etcd.io/etcd/client/v3" ) +// Use different etcd client so they are not affected by each other, +// and can be configured separately. var ( + // etcd client for transform engine globalClient *clientv3.Client - mutex sync.Mutex - refCount int + globalMutex sync.Mutex + globalRefCount int + // etcd client for store + storeClient *clientv3.Client + storeMutex sync.Mutex + // keep alive contexts for store + storeKeepAliveCtx = make(map[int64]context.CancelFunc) + storeKeepAliveMutex sync.Mutex + // watch contexts for store + storeWatchCtx = make(map[string]context.CancelFunc) + storeWatchMutex sync.Mutex ) //export NewEtcdClient func NewEtcdClient(endpoints *C.char, errMsg **C.char) int { - mutex.Lock() - defer mutex.Unlock() + globalMutex.Lock() + defer globalMutex.Unlock() if globalClient != nil { - refCount++ + globalRefCount++ return 0 } @@ -40,7 +55,7 @@ func NewEtcdClient(endpoints *C.char, errMsg **C.char) int { } globalClient = cli - refCount++ + globalRefCount++ return 0 } @@ -104,15 +119,291 @@ func EtcdDeleteWrapper(key *C.char, errMsg **C.char) int { //export EtcdCloseWrapper func EtcdCloseWrapper() { - mutex.Lock() - defer mutex.Unlock() + globalMutex.Lock() + defer globalMutex.Unlock() if globalClient != nil { - refCount-- - if refCount == 0 { + globalRefCount-- + if globalRefCount == 0 { globalClient.Close() globalClient = nil } } } +//export NewStoreEtcdClient +func NewStoreEtcdClient(endpoints *C.char, errMsg **C.char) int { + storeMutex.Lock() + defer storeMutex.Unlock() + if storeClient != nil { + *errMsg = C.CString("etcd client can be initialized only once") + return -2 + } + + endpointStr := C.GoString(endpoints) + endpointList := strings.Split(endpointStr, ";") + + // Filter out any empty strings that might result from splitting + var validEndpoints []string + for _, ep := range endpointList { + if ep != "" { + validEndpoints = append(validEndpoints, ep) + } + } + + if len(validEndpoints) == 0 { + *errMsg = C.CString("no valid endpoints provided") + return -1 + } + + cli, err := clientv3.New(clientv3.Config{ + Endpoints: validEndpoints, + DialTimeout: 5 * time.Second, + }) + + if err != nil { + *errMsg = C.CString(err.Error()) + return -1 + } + + storeClient = cli + return 0 +} + +//export EtcdStoreGetWrapper +func EtcdStoreGetWrapper(key *C.char, keySize C.int, value **C.char, + valueSize *C.int, revisionId *int64, errMsg **C.char) int { + if storeClient == nil { + *errMsg = C.CString("etcd client not initialized") + return -1 + } + k := C.GoStringN(key, keySize) + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + resp, err := storeClient.Get(ctx, k) + if err != nil { + *errMsg = C.CString(err.Error()) + return -1 + } + if len(resp.Kvs) == 0 { + *errMsg = C.CString("key not found in etcd") + return -2 + } else { + kv := resp.Kvs[0] + *value = C.CString(string(kv.Value)) + *valueSize = C.int(len(kv.Value)) + *revisionId = kv.CreateRevision + return 0 + } +} + +//export EtcdStoreGrantLeaseWrapper +func EtcdStoreGrantLeaseWrapper(ttl int64, leaseId *int64, errMsg **C.char) int { + if storeClient == nil { + *errMsg = C.CString("etcd client not initialized") + return -1 + } + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + resp, err := storeClient.Grant(ctx, ttl) + if err != nil { + *errMsg = C.CString(err.Error()) + return -1 + } + *leaseId = int64(resp.ID) + return 0 +} + +//export EtcdStoreCreateWithLeaseWrapper +func EtcdStoreCreateWithLeaseWrapper(key *C.char, keySize C.int, value *C.char, valueSize C.int, + leaseId int64, revisionId *int64, errMsg **C.char) int { + if storeClient == nil { + *errMsg = C.CString("etcd client not initialized") + return -1 + } + k := C.GoStringN(key, keySize) + v := C.GoStringN(value, valueSize) + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + // Create a transaction + txn := storeClient.Txn(ctx) + + // Only put the key if it does not exist + resp, err := txn.If(clientv3.Compare(clientv3.CreateRevision(k), "=", 0)). + Then(clientv3.OpPut(k, v, clientv3.WithLease(clientv3.LeaseID(leaseId)))). + Commit() + + if err != nil { + *errMsg = C.CString(err.Error()) + return -1 + } + + // If the key already existed, resp.Succeeded will be false + // If we created the key, resp.Succeeded will be true + if resp.Succeeded { + *revisionId = resp.Header.Revision + return 0; + } else { + *errMsg = C.CString("etcd transaction failed") + return -2 + } +} + +/* +* @brief First cancel the watch context, then delete it from the map. +* Cancel must be called before delete in case this is a new context +* other than the one we want to delete. In that case, that context will +* be deleted before being cancelled and will not be able to be cancelled +* anymore. +*/ +func cancelAndDeleteWatch(k string) int { + storeWatchMutex.Lock() + defer storeWatchMutex.Unlock() + + if cancel, exists := storeWatchCtx[k]; exists { + cancel() + delete(storeWatchCtx, k) + return 0 + } + return -1 +} + +//export EtcdStoreWatchUntilDeletedWrapper +func EtcdStoreWatchUntilDeletedWrapper(key *C.char, keySize C.int, errMsg **C.char) int { + if storeClient == nil { + *errMsg = C.CString("etcd client not initialized") + return -1 + } + k := C.GoStringN(key, keySize) + + // Create a context with cancel function + ctx, cancel := context.WithCancel(context.Background()) + + // Store the cancel function + storeWatchMutex.Lock() + if _, exists := storeWatchCtx[k]; exists { + storeWatchMutex.Unlock() + *errMsg = C.CString("This key is already being watched") + return -1 + } + storeWatchCtx[k] = cancel + storeWatchMutex.Unlock() + + // Make sure to delete from the map before returning + defer cancelAndDeleteWatch(k) + + // Start watching the key + watchChan := storeClient.Watch(ctx, k) + + // Wait for the key to be deleted + for { + select { + case watchResp, ok := <-watchChan: + if !ok { + // Channel closed unexpectedly + *errMsg = C.CString("watch channel closed unexpectedly") + return -1 + } + for _, event := range watchResp.Events { + if event.Type == clientv3.EventTypeDelete { + // Clean up the context when done + return 0 + } + } + case <-ctx.Done(): + // Context was cancelled + *errMsg = C.CString("watch context cancelled") + return -2 + } + } +} + +//export EtcdStoreCancelWatchWrapper +func EtcdStoreCancelWatchWrapper(key *C.char, keySize C.int, errMsg **C.char) int { + k := C.GoStringN(key, keySize) + if cancelAndDeleteWatch(k) == -1 { + *errMsg = C.CString("no watch context found for the given key") + return -1 + } + return 0 +} + +/* +* @brief First cancel the keep alive context, then delete it from the map. +* Cancel must be called before deleting in case this is a new context +* other than the one we want to delete. In that case, that context will +* be deleted before being cancelled and will not be able to be cancelled +* anymore. +*/ +func cancelAndDeleteKeepAlive(leaseId int64) int { + storeKeepAliveMutex.Lock() + defer storeKeepAliveMutex.Unlock() + + if cancel, exists := storeKeepAliveCtx[leaseId]; exists { + cancel() + delete(storeKeepAliveCtx, leaseId) + return 0 + } + return -1 +} + +//export EtcdStoreKeepAliveWrapper +func EtcdStoreKeepAliveWrapper(leaseId int64, errMsg **C.char) int { + if storeClient == nil { + *errMsg = C.CString("etcd client not initialized") + return -1 + } + + // Create a context with cancel function + ctx, cancel := context.WithCancel(context.Background()) + + // Store the cancel function + storeKeepAliveMutex.Lock() + if _, exists := storeKeepAliveCtx[leaseId]; exists { + storeKeepAliveMutex.Unlock() + *errMsg = C.CString("This lease id is already being kept alive") + return -1 + } + storeKeepAliveCtx[leaseId] = cancel + storeKeepAliveMutex.Unlock() + // Make sure to delete from the map before returning + defer cancelAndDeleteKeepAlive(leaseId) + + // Start keep alive + keepAliveChan, err := storeClient.KeepAlive(ctx, clientv3.LeaseID(leaseId)) + if err != nil { + *errMsg = C.CString(err.Error()) + return -1 + } + + // Wait for keep alive responses + for { + select { + case resp, ok := <-keepAliveChan: + if !ok { + *errMsg = C.CString("keep alive channel closed") + return -1 + } + if resp == nil { + *errMsg = C.CString("keep alive response is nil") + return -1 + } + // Keep alive successful, continue + case <-ctx.Done(): + // Context cancelled + *errMsg = C.CString("keep alive context cancelled") + return -2 + } + } +} + +//export EtcdStoreCancelKeepAliveWrapper +func EtcdStoreCancelKeepAliveWrapper(leaseId int64, errMsg **C.char) int { + if cancelAndDeleteKeepAlive(leaseId) == -1 { + *errMsg = C.CString("no keep alive context found for the given lease ID") + return -1 + } + return 0 +} + func main() {} diff --git a/mooncake-store/CMakeLists.txt b/mooncake-store/CMakeLists.txt index a65189ec..ec0c0e9c 100644 --- a/mooncake-store/CMakeLists.txt +++ b/mooncake-store/CMakeLists.txt @@ -1,5 +1,7 @@ project(MooncakeStore) +set(ETCD_WRAPPER_INCLUDE ${CMAKE_CURRENT_BINARY_DIR}/../mooncake-common/etcd/) +set(ETCD_WRAPPER_LIB ${CMAKE_CURRENT_BINARY_DIR}/../mooncake-common/etcd/libetcd_wrapper.so) # Add include directories include_directories( @@ -9,8 +11,9 @@ include_directories( ${CMAKE_CURRENT_SOURCE_DIR}/include/mooncake-store/proto/ ${CMAKE_CURRENT_SOURCE_DIR}/include/ ${CMAKE_CURRENT_SOURCE_DIR}/../mooncake-transfer-engine/include + ${ETCD_WRAPPER_INCLUDE} ) # Add subdirectories add_subdirectory(src) -add_subdirectory(tests) +add_subdirectory(tests) \ No newline at end of file diff --git a/mooncake-store/include/client.h b/mooncake-store/include/client.h index 7e62ab44..74369489 100644 --- a/mooncake-store/include/client.h +++ b/mooncake-store/include/client.h @@ -10,6 +10,7 @@ #include "rpc_service.h" #include "transfer_engine.h" #include "types.h" +#include "ha_helper.h" namespace mooncake { @@ -26,7 +27,9 @@ class Client { * @param metadata_connstring Connection string for metadata service * @param protocol Transfer protocol ("rdma" or "tcp") * @param protocol_args Protocol-specific arguments - * @param master_addr Master server address + * @param master_server_entry The entry of master server (IP:Port of master + * address for non-HA mode, etcd://IP:Port;IP:Port;...;IP:Port for + * HA mode) * @return std::optional containing a shared_ptr to Client if successful, * std::nullopt otherwise */ @@ -34,7 +37,7 @@ class Client { const std::string& local_hostname, const std::string& metadata_connstring, const std::string& protocol, void** protocol_args, - const std::string& master_addr = kDefaultMasterAddress); + const std::string& master_server_entry = kDefaultMasterAddress); /** * @brief Retrieves data for a given key @@ -197,7 +200,7 @@ class Client { /** * @brief Internal helper functions for initialization and data transfer */ - ErrorCode ConnectToMaster(const std::string& master_addr); + ErrorCode ConnectToMaster(const std::string& master_server_entry); ErrorCode InitTransferEngine(const std::string& local_hostname, const std::string& metadata_connstring, const std::string& protocol, @@ -216,13 +219,24 @@ class Client { TransferEngine transfer_engine_; MasterClient master_client_; + // Client local segments + struct Segment{ + void* buffer; + size_t size; + }; // Mutex to protect mounted_segments_ std::mutex mounted_segments_mutex_; - std::unordered_map mounted_segments_; + std::unordered_map mounted_segments_; // Configuration const std::string local_hostname_; const std::string metadata_connstring_; + + // For high availability + MasterViewHelper master_view_helper_; + std::thread ping_thread_; + std::atomic ping_running_{false}; + void PingThreadFunc(int current_version); }; } // namespace mooncake \ No newline at end of file diff --git a/mooncake-store/include/etcd_helper.h b/mooncake-store/include/etcd_helper.h new file mode 100644 index 00000000..e81ce562 --- /dev/null +++ b/mooncake-store/include/etcd_helper.h @@ -0,0 +1,102 @@ +#pragma once + +#include + +#include "libetcd_wrapper.h" +#include "types.h" + +namespace mooncake { + +/* + * @brief A helper class for etcd operations. + * This class is used to handle the requests to the etcd cluster. + * All methods of this class are thread-safe. + */ +class EtcdHelper { + public: + /* + * @brief Connect to the etcd store client. There is a global etcd client in + * libetcd. It is used for all the etcd operations for mooncake-store. This + * function ensures the client is only connected once. + * @param etcd_endpoints: The endpoints of the etcd store client. + * Multiple endpoints are separated by semicolons. + * @return: Error code. + */ + static ErrorCode ConnectToEtcdStoreClient( + const std::string& etcd_endpoints); + + /* + * @brief Get the value of a key from the etcd. + * @param key: The key to get the value of. + * @param key_size: The size of the key in bytes. + * @param value: Output param, the value of the key. + * @param revision_id: Output param, the create revision id of the key. + * @return: Error code. + */ + static ErrorCode Get(const char* key, const size_t key_size, + std::string& value, EtcdRevisionId& revision_id); + + /* + * @brief Create a key-value pair that binds to a given lease. + * @param key: The key to create. + * @param key_size: The size of the key in bytes. + * @param value: The value to create. + * @param value_size: The size of the value in bytes. + * @param lease_id: The lease id to bind to the key. + * @param revision_id: Output param, the create revision id of the key. + * @return: Error code. + */ + static ErrorCode CreateWithLease(const char* key, const size_t key_size, + const char* value, const size_t value_size, + EtcdLeaseId lease_id, + EtcdRevisionId& revision_id); + + /* + * @brief Grant a lease from the etcd. + * @param lease_ttl: The ttl of the lease, in seconds. + * @param lease_id: Output param, the lease id. + * @return: Error code. + */ + static ErrorCode GrantLease(int64_t lease_ttl, EtcdLeaseId& lease_id); + + /* + * @brief Watch a key until it is deleted. This is a blocking function. + * @param key: The key to watch. + * @param key_size: The size of the key in bytes. + * @return: Error code. + */ + static ErrorCode WatchUntilDeleted(const char* key, const size_t key_size); + + /* + * @brief Cancel watching a key + * @param key: The key to cancel watch. + * @param key_size: The size of the key in bytes. + * @return: Error code. + */ + static ErrorCode CancelWatch(const char* key, const size_t key_size); + + /* + * @brief Keep a lease alive. This is a blocking function. + * @param lease_id: The lease id to keep alive. + * @return: Error code. + */ + static ErrorCode KeepAlive(EtcdLeaseId lease_id); + + /* + * @brief Cancel a lease keep alive. Returning error means + * the lease id does not exist or the goroutine to + * keep the lease alive is already closed. + * @param lease_id: The lease id to cancel keep alive. + * @return: Error code. + */ + static ErrorCode CancelKeepAlive(EtcdLeaseId lease_id); + + private: + // Variables that are used to ensure the etcd client + // is only connected once. + static std::string connected_endpoints_; + static std::mutex etcd_mutex_; + static bool etcd_connected_; +}; + +} // namespace mooncake \ No newline at end of file diff --git a/mooncake-store/include/ha_helper.h b/mooncake-store/include/ha_helper.h new file mode 100644 index 00000000..0ced521d --- /dev/null +++ b/mooncake-store/include/ha_helper.h @@ -0,0 +1,111 @@ +#ifndef MOONCAKE_HA_HELPER_H_ +#define MOONCAKE_HA_HELPER_H_ + +#include + +#include +#include +#include +#include +#include + +#include "etcd_helper.h" +#include "rpc_service.h" +#include "types.h" + +namespace mooncake { + +// The key to store the master view in etcd +inline const char* const MASTER_VIEW_KEY = "mooncake-store/master_view"; + +/* + * @brief A helper class for maintain and monitor the master view change. + * The cluster is assumed to have multiple master servers, but only + * one master can be elected as leader to serve client requests. + * Each master view is associated with a unique version id, which + * is incremented monotonically each time the master view is changed. + */ +class MasterViewHelper { + public: + MasterViewHelper(const MasterViewHelper&) = delete; + MasterViewHelper& operator=(const MasterViewHelper&) = delete; + MasterViewHelper() = default; + + /* + * @brief Connect to the etcd cluster. This function should be called at + * first + * @param etcd_endpoints: The endpoints of the etcd store client. + * Multiple endpoints are separated by semicolons. + * @return: Error code. + */ + ErrorCode ConnectToEtcd(const std::string& etcd_endpoints); + + /* + * @brief Elect the master to be the leader. This is a blocking function. + * @param master_address: The ip:port address of the master to be elected. + * @param version: Output param, the version of the new master view. + * @param lease_id: Output param, the lease id of the leader. + */ + void ElectLeader(const std::string& master_address, ViewVersionId& version, + EtcdLeaseId& lease_id); + + /* + * @brief Keep the master to be the leader. This function blocks until the + * master is no longer the leader. + * @param lease_id: The lease id of the leader. + */ + void KeepLeader(EtcdLeaseId lease_id); + + /* + * @brief Get the current master view. + * @param master: Output param, the ip:port address of the master. + * @param version: Output param, the version of the master view. + * @return: Error code. + */ + ErrorCode GetMasterView(std::string& master_address, + ViewVersionId& version); +}; + +/* + * @brief A supervisor class for the master service, only used in HA mode. + * This class will continuously do the following procedures after start: + * 1. Elect local master to be the leader. + * 2. Start the master service when it is elected as leader. + * 3. Stop the master service when it is no longer the leader. + */ +class MasterServiceSupervisor { + public: + MasterServiceSupervisor( + int port, int server_thread_num, bool enable_gc, + bool enable_metric_reporting, int metrics_port, + int64_t default_kv_lease_ttl, double eviction_ratio, + double eviction_high_watermark_ratio, + const std::string& etcd_endpoints = "0.0.0.0:2379", + const std::string& local_hostname = "0.0.0.0:50051"); + int Start(); + ~MasterServiceSupervisor(); + + private: + // Master service parameters + int port_; + int server_thread_num_; + bool enable_gc_; + bool enable_metric_reporting_; + int metrics_port_; + int64_t default_kv_lease_ttl_; + double eviction_ratio_; + double eviction_high_watermark_ratio_; + + // coro_rpc server thread + std::thread server_thread_; + + // ETCD parameters + std::string etcd_endpoints_; + + // Local hostname for leader election + std::string local_hostname_; +}; + +} // namespace mooncake + +#endif // MOONCAKE_HA_HELPER_H_ \ No newline at end of file diff --git a/mooncake-store/include/master_client.h b/mooncake-store/include/master_client.h index 3138f313..77b97c3a 100644 --- a/mooncake-store/include/master_client.h +++ b/mooncake-store/include/master_client.h @@ -148,6 +148,14 @@ class MasterClient { [[nodiscard]] UnmountSegmentResponse UnmountSegment( const std::string& segment_name); + /** + * @brief Pings master to check its availability + * @param No parameters + * @return current master view version + * @return ErrorCode indicating success/failure + */ + [[nodiscard]] PingResponse Ping(); + private: coro_rpc_client client_; }; diff --git a/mooncake-store/include/master_metric_manager.h b/mooncake-store/include/master_metric_manager.h index d16ac6da..14229ad7 100644 --- a/mooncake-store/include/master_metric_manager.h +++ b/mooncake-store/include/master_metric_manager.h @@ -53,6 +53,7 @@ class MasterMetricManager { void inc_mount_segment_failures(int64_t val = 1); void inc_unmount_segment_requests(int64_t val = 1); void inc_unmount_segment_failures(int64_t val = 1); + void inc_ping_requests(int64_t val = 1); // Operation Statistics Getters int64_t get_put_start_requests(); @@ -73,6 +74,7 @@ class MasterMetricManager { int64_t get_mount_segment_failures(); int64_t get_unmount_segment_requests(); int64_t get_unmount_segment_failures(); + int64_t get_ping_requests(); // Eviction Metrics void inc_eviction_success(int64_t key_count, int64_t size); @@ -131,6 +133,7 @@ class MasterMetricManager { ylt::metric::counter_t mount_segment_failures_; ylt::metric::counter_t unmount_segment_requests_; ylt::metric::counter_t unmount_segment_failures_; + ylt::metric::counter_t ping_requests_; // Eviction Metrics ylt::metric::counter_t eviction_success_; diff --git a/mooncake-store/include/rpc_service.h b/mooncake-store/include/rpc_service.h index f72c669d..bef7981b 100644 --- a/mooncake-store/include/rpc_service.h +++ b/mooncake-store/include/rpc_service.h @@ -1,18 +1,21 @@ #pragma once +#include +#include + #include #include #include #include #include #include +#include #include -#include -#include #include "master_metric_manager.h" #include "master_service.h" #include "types.h" #include "utils/scoped_vlog_timer.h" + namespace mooncake { struct ExistKeyResponse { @@ -82,6 +85,12 @@ struct UnmountSegmentResponse { }; YLT_REFL(UnmountSegmentResponse, error_code) +struct PingResponse { + ViewVersionId view_version = 0; + ErrorCode error_code = ErrorCode::OK; +}; +YLT_REFL(PingResponse, view_version, error_code) + constexpr uint64_t kMetricReportIntervalSeconds = 10; class WrappedMasterService { @@ -90,10 +99,14 @@ class WrappedMasterService { bool enable_metric_reporting = true, uint16_t http_port = 9003, double eviction_ratio = DEFAULT_EVICTION_RATIO, - double eviction_low_watermark_ratio = DEFAULT_EVICTION_HIGH_WATERMARK_RATIO) - : master_service_(enable_gc, default_kv_lease_ttl, eviction_ratio, eviction_low_watermark_ratio), + double eviction_high_watermark_ratio = + DEFAULT_EVICTION_HIGH_WATERMARK_RATIO, + ViewVersionId view_version = 0) + : master_service_(enable_gc, default_kv_lease_ttl, eviction_ratio, + eviction_high_watermark_ratio), http_server_(4, http_port), - metric_report_running_(enable_metric_reporting) { + metric_report_running_(enable_metric_reporting), + view_version_(view_version) { // Initialize HTTP server for metrics init_http_server(); @@ -152,8 +165,9 @@ class WrappedMasterService { response = GetReplicaList(std::string(key)); resp.add_header("Content-Type", "text/plain; version=0.0.4"); std::string ss = ""; - for(size_t i = 0; i < response.replica_list.size(); i++) { - for(const auto& handle : response.replica_list[i].buffer_descriptors) { + for (size_t i = 0; i < response.replica_list.size(); i++) { + for (const auto& handle : + response.replica_list[i].buffer_descriptors) { std::string tmp = ""; struct_json::to_json(handle, tmp); ss += tmp; @@ -169,9 +183,9 @@ class WrappedMasterService { [&](coro_http_request& req, coro_http_response& resp) { resp.add_header("Content-Type", "text/plain; version=0.0.4"); std::string ss = ""; - std::vector all_keys; + std::vector all_keys; master_service_.GetAllKeys(all_keys); - for(const auto & key : all_keys) { + for (const auto& key : all_keys) { ss += key; ss += "\n"; } @@ -184,9 +198,9 @@ class WrappedMasterService { [&](coro_http_request& req, coro_http_response& resp) { resp.add_header("Content-Type", "text/plain; version=0.0.4"); std::string ss = ""; - std::vector all_segments; + std::vector all_segments; master_service_.GetAllSegments(all_segments); - for(const auto & segment: all_segments) { + for (const auto& segment : all_segments) { ss += segment; ss += "\n"; } @@ -201,8 +215,8 @@ class WrappedMasterService { resp.add_header("Content-Type", "text/plain; version=0.0.4"); std::string ss = ""; size_t used = 0, capacity = 0; - if(master_service_.QuerySegments(std::string(segment), used, capacity) - == ErrorCode::OK) { + if (master_service_.QuerySegments(std::string(segment), used, + capacity) == ErrorCode::OK) { ss += segment; ss += "\n"; ss += "Used(bytes): "; @@ -480,11 +494,58 @@ class WrappedMasterService { return response; } + PingResponse Ping() { + ScopedVLogTimer timer(1, "Ping"); + timer.LogRequest("action=ping"); + + MasterMetricManager::instance().inc_ping_requests(); + + PingResponse response(view_version_, ErrorCode::OK); + + timer.LogResponseJson(response); + return response; + } + private: MasterService master_service_; std::thread metric_report_thread_; coro_http::coro_http_server http_server_; std::atomic metric_report_running_; + ViewVersionId view_version_; }; -} // namespace mooncake +inline void RegisterRpcService( + coro_rpc::coro_rpc_server& server, + mooncake::WrappedMasterService& wrapped_master_service) { + server.register_handler<&mooncake::WrappedMasterService::ExistKey>( + &wrapped_master_service); + server.register_handler<&mooncake::WrappedMasterService::GetReplicaList>( + &wrapped_master_service); + server + .register_handler<&mooncake::WrappedMasterService::BatchGetReplicaList>( + &wrapped_master_service); + server.register_handler<&mooncake::WrappedMasterService::PutStart>( + &wrapped_master_service); + server.register_handler<&mooncake::WrappedMasterService::PutEnd>( + &wrapped_master_service); + server.register_handler<&mooncake::WrappedMasterService::PutRevoke>( + &wrapped_master_service); + server.register_handler<&mooncake::WrappedMasterService::BatchPutStart>( + &wrapped_master_service); + server.register_handler<&mooncake::WrappedMasterService::BatchPutEnd>( + &wrapped_master_service); + server.register_handler<&mooncake::WrappedMasterService::BatchPutRevoke>( + &wrapped_master_service); + server.register_handler<&mooncake::WrappedMasterService::Remove>( + &wrapped_master_service); + server.register_handler<&mooncake::WrappedMasterService::RemoveAll>( + &wrapped_master_service); + server.register_handler<&mooncake::WrappedMasterService::MountSegment>( + &wrapped_master_service); + server.register_handler<&mooncake::WrappedMasterService::UnmountSegment>( + &wrapped_master_service); + server.register_handler<&mooncake::WrappedMasterService::Ping>( + &wrapped_master_service); +} + +} // namespace mooncake \ No newline at end of file diff --git a/mooncake-store/include/types.h b/mooncake-store/include/types.h index 0e3a4885..2b2b5132 100644 --- a/mooncake-store/include/types.h +++ b/mooncake-store/include/types.h @@ -12,6 +12,7 @@ #include "Slab.h" #include "ylt/struct_json/json_reader.h" #include "ylt/struct_json/json_writer.h" +#include "libetcd_wrapper.h" namespace mooncake { @@ -23,6 +24,7 @@ static constexpr uint64_t DEFAULT_DEFAULT_KV_LEASE_TTL = 200; // in milliseconds static constexpr double DEFAULT_EVICTION_RATIO = 0.1; static constexpr double DEFAULT_EVICTION_HIGH_WATERMARK_RATIO = 1.0; +static constexpr int64_t ETCD_MASTER_VIEW_LEASE_TTL = 5; // in seconds // Forward declarations class BufferAllocator; @@ -39,6 +41,10 @@ using BufHandleList = std::vector>; using ReplicaList = std::unordered_map; using BufferResources = std::map>>; +// Mapping between c++ and go types +using EtcdRevisionId = GoInt64; +using ViewVersionId = EtcdRevisionId; +using EtcdLeaseId = GoInt64; /** * @brief Error codes for various operations in the system @@ -83,6 +89,12 @@ enum class ErrorCode : int32_t { // RPC errors (Range: -900 to -999) RPC_FAIL = -900, ///< RPC operation failed. + + // ETCD errors (Range: -1000 to -1099) + ETCD_OPERATION_ERROR = -1000, ///< etcd operation failed. + ETCD_KEY_NOT_EXIST = -1001, ///< key not found in etcd. + ETCD_TRANSACTION_FAIL = -1002, ///< etcd transaction failed. + ETCD_CTX_CANCELLED = -1003, ///< etcd context cancelled. }; int32_t toInt(ErrorCode errorCode) noexcept; diff --git a/mooncake-store/src/CMakeLists.txt b/mooncake-store/src/CMakeLists.txt index 9d9f6366..08eb445d 100644 --- a/mooncake-store/src/CMakeLists.txt +++ b/mooncake-store/src/CMakeLists.txt @@ -9,13 +9,15 @@ set(MOONCAKE_STORE_SOURCES master_client.cpp utils.cpp master_metric_manager.cpp + etcd_helper.cpp + ha_helper.cpp ) # The cache_allocator library include_directories(${Python3_INCLUDE_DIRS}) add_library(mooncake_store ${MOONCAKE_STORE_SOURCES}) -target_link_libraries(mooncake_store PUBLIC transfer_engine glog::glog gflags::gflags) - +target_link_libraries(mooncake_store PUBLIC transfer_engine ${ETCD_WRAPPER_LIB} glog::glog gflags::gflags) +add_dependencies(mooncake_store build_etcd_wrapper) # Master binary add_executable(mooncake_master master.cpp) @@ -23,6 +25,8 @@ target_link_libraries(mooncake_master PRIVATE mooncake_store cachelib_memory_allocator pthread + ${ETCD_WRAPPER_LIB} ) +add_dependencies(mooncake_master build_etcd_wrapper) install(TARGETS mooncake_master DESTINATION bin) diff --git a/mooncake-store/src/client.cpp b/mooncake-store/src/client.cpp index 5c563760..734a9dff 100644 --- a/mooncake-store/src/client.cpp +++ b/mooncake-store/src/client.cpp @@ -31,11 +31,11 @@ Client::~Client() { // No need for mutex here since the client is being destroyed(protected by // shared_ptr) // Make a copy of mounted_segments_ to avoid modifying while iterating - std::unordered_map segments_to_unmount = + std::unordered_map segments_to_unmount = mounted_segments_; for (auto& entry : segments_to_unmount) { - auto err_code = UnmountSegment(entry.first, entry.second); + auto err_code = UnmountSegment(entry.first, entry.second.buffer); if (err_code != ErrorCode::OK) { LOG(ERROR) << "Failed to unmount segment: " << toString(err_code); } @@ -43,10 +43,14 @@ Client::~Client() { // Clear any remaining segments mounted_segments_.clear(); -} -ErrorCode Client::ConnectToMaster(const std::string& master_addr) { - return master_client_.Connect(master_addr); + // Stop ping thread only after no need to contact master anymore + if (ping_running_) { + ping_running_ = false; + if (ping_thread_.joinable()) { + ping_thread_.join(); + } + } } static bool get_auto_discover() { @@ -105,6 +109,42 @@ static std::vector get_auto_discover_filters(bool auto_discover) { return whitelst_filters; } +ErrorCode Client::ConnectToMaster(const std::string& master_server_entry) { + if (master_server_entry.find("etcd://") == 0) { + std::string etcd_entry = master_server_entry.substr(strlen("etcd://")); + + // Get master address from etcd + auto err = master_view_helper_.ConnectToEtcd(etcd_entry); + if (err != ErrorCode::OK) { + LOG(ERROR) << "Failed to connect to etcd"; + return err; + } + std::string master_address; + ViewVersionId master_version = 0; + err = master_view_helper_.GetMasterView(master_address, master_version); + if (err != ErrorCode::OK) { + LOG(ERROR) << "Failed to get master address"; + return err; + } + + err = master_client_.Connect(master_address); + if (err != ErrorCode::OK) { + LOG(ERROR) << "Failed to connect to master"; + return err; + } + + // Start Ping thread to monitor master view changes and remount segments + // if needed + ping_running_ = true; + ping_thread_ = + std::thread(&Client::PingThreadFunc, this, master_version); + + return ErrorCode::OK; + } else { + return master_client_.Connect(master_server_entry); + } +} + ErrorCode Client::InitTransferEngine(const std::string& local_hostname, const std::string& metadata_connstring, const std::string& protocol, @@ -145,22 +185,18 @@ ErrorCode Client::InitTransferEngine(const std::string& local_hostname, std::optional> Client::Create( const std::string& local_hostname, const std::string& metadata_connstring, const std::string& protocol, void** protocol_args, - const std::string& master_addr) { + const std::string& master_server_entry) { auto client = std::shared_ptr( new Client(local_hostname, metadata_connstring)); - // Connect to master service - ErrorCode err = client->ConnectToMaster(master_addr); + ErrorCode err = client->ConnectToMaster(master_server_entry); if (err != ErrorCode::OK) { - LOG(ERROR) << "Failed to connect to Master"; return std::nullopt; } - LOG(INFO) << "Connect to Master success"; - // 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; @@ -492,7 +528,7 @@ ErrorCode Client::MountSegment(const std::string& segment_name, { std::lock_guard lock(mounted_segments_mutex_); - mounted_segments_[segment_name] = (void*)buffer; + mounted_segments_[segment_name] = {(void*)buffer, size}; } return ErrorCode::OK; } @@ -502,11 +538,11 @@ ErrorCode Client::UnmountSegment(const std::string& segment_name, void* addr) { { std::lock_guard lock(mounted_segments_mutex_); auto it = mounted_segments_.find(segment_name); - if (it == mounted_segments_.end() || it->second != addr) { + if (it == mounted_segments_.end() || it->second.buffer != addr) { LOG(ERROR) << "segment_not_found segment_name=" << segment_name; return ErrorCode::INVALID_PARAMS; } - segment_addr = it->second; + segment_addr = it->second.buffer; // Remove from map first to prevent any further access to this segment mounted_segments_.erase(it); @@ -674,4 +710,101 @@ ErrorCode Client::TransferRead( return TransferData(handles, slices, TransferRequest::READ); } +void Client::PingThreadFunc(int current_version) { + // How many failed pings before getting latest master view from etcd + const int max_ping_fail_count = 3; + // How long to wait for next ping after success + const int success_ping_interval_ms = 1000; + // How long to wait for next ping after failure + const int fail_ping_interval_ms = 1000; + // Increment after a ping failure, reset after a ping success + int ping_fail_count = 0; + // Set to true when there is a view change. + // When set true, will try to remount periodically. + bool need_remount = false; + + auto remount_segment = [this]() { + std::lock_guard lock(mounted_segments_mutex_); + for (auto it : mounted_segments_) { + auto& name = it.first; + auto& segment = it.second; + auto err = + master_client_.MountSegment(name, segment.buffer, segment.size) + .error_code; + // If err is INVALID_PARAMS, it means the segment is already + // mounted, or cannot be mounted with current parameters. Either + // way, there is nothing we can do for this segment. + if (err != ErrorCode::OK && err != ErrorCode::INVALID_PARAMS) { + LOG(ERROR) << "Failed to remount segment " << name << ": " + << toString(err); + return err; + } + } + return ErrorCode::OK; + }; + + while (ping_running_) { + auto ping_result = master_client_.Ping(); + if (ping_result.error_code == ErrorCode::OK) { + ping_fail_count = 0; + if (ping_result.view_version > current_version) { + // There is an unknown view change, we need to update + // local view version and remount segments. + LOG(ERROR) << "Master view version has changed, need to " + "remount segments"; + current_version = ping_result.view_version; + need_remount = true; + } + // Only try to remount if the ping succeeds and need_remount is true + if (need_remount && remount_segment() == ErrorCode::OK) { + LOG(INFO) << "Successfully remounted all segments"; + need_remount = false; + } + std::this_thread::sleep_for( + std::chrono::milliseconds(success_ping_interval_ms)); + continue; + } + + ping_fail_count++; + if (ping_fail_count < max_ping_fail_count) { + LOG(ERROR) << "Failed to ping master"; + std::this_thread::sleep_for( + std::chrono::milliseconds(fail_ping_interval_ms)); + continue; + } + + // Too many ping failures, we need to check if the master view has + // changed + LOG(ERROR) << "Failed to ping master for " << ping_fail_count + << " times, try to get latest master view and reconnect"; + std::string master_address; + ViewVersionId next_version = 0; + auto err = + master_view_helper_.GetMasterView(master_address, next_version); + if (err != ErrorCode::OK) { + LOG(ERROR) << "Failed to get new master view: " << toString(err); + std::this_thread::sleep_for( + std::chrono::milliseconds(fail_ping_interval_ms)); + continue; + } + + err = master_client_.Connect(master_address); + if (err != ErrorCode::OK) { + LOG(ERROR) << "Failed to connect to master " << master_address + << ": " << toString(err); + std::this_thread::sleep_for( + std::chrono::milliseconds(fail_ping_interval_ms)); + continue; + } + + LOG(INFO) << "Reconnected to master " << master_address; + ping_fail_count = 0; + if (next_version > current_version) { + // Master view has changed + current_version = next_version; + need_remount = true; + } + } +} + } // namespace mooncake diff --git a/mooncake-store/src/etcd_helper.cpp b/mooncake-store/src/etcd_helper.cpp new file mode 100644 index 00000000..399a7527 --- /dev/null +++ b/mooncake-store/src/etcd_helper.cpp @@ -0,0 +1,151 @@ +#include "etcd_helper.h" + +#include + +namespace mooncake { + +std::string EtcdHelper::connected_endpoints_ = ""; +std::mutex EtcdHelper::etcd_mutex_; +bool EtcdHelper::etcd_connected_ = false; + +ErrorCode EtcdHelper::ConnectToEtcdStoreClient( + const std::string& etcd_endpoints) { + std::lock_guard lock(etcd_mutex_); + if (etcd_connected_) { + if (connected_endpoints_ != etcd_endpoints) { + LOG(ERROR) << "Etcd client already connected to " + << connected_endpoints_ + << ", while trying to connect to " << etcd_endpoints; + return ErrorCode::INVALID_PARAMS; + } + return ErrorCode::OK; + } else { + char* err_msg = nullptr; + int ret = NewStoreEtcdClient((char*)etcd_endpoints.c_str(), &err_msg); + // ret == -2 means the etcd client has already been initialized + if (ret != 0 && ret != -2) { + LOG(ERROR) << "Failed to initialize etcd client: " << err_msg; + free(err_msg); + err_msg = nullptr; + return ErrorCode::ETCD_OPERATION_ERROR; + } + // Record the connection to avoid future connection attempts + connected_endpoints_ = etcd_endpoints; + etcd_connected_ = true; + return ErrorCode::OK; + } +} + +ErrorCode EtcdHelper::Get(const char* key, const size_t key_size, + std::string& value, EtcdRevisionId& revision_id) { + char* err_msg = nullptr; + char* value_ptr = nullptr; + int value_size = 0; + int ret = EtcdStoreGetWrapper((char*)key, (int)key_size, &value_ptr, + &value_size, &revision_id, &err_msg); + if (ret == -2) { + LOG(ERROR) << "key=" << std::string(key, key_size) + << ", error=" << err_msg; + free(err_msg); + return ErrorCode::ETCD_KEY_NOT_EXIST; + } + if (ret != 0) { + LOG(ERROR) << "key=" << std::string(key, key_size) + << ", error=" << err_msg; + free(err_msg); + return ErrorCode::ETCD_OPERATION_ERROR; + } + value = std::string(value_ptr, value_size); + free(value_ptr); + return ErrorCode::OK; +} + +ErrorCode EtcdHelper::CreateWithLease(const char* key, const size_t key_size, + const char* value, + const size_t value_size, + EtcdLeaseId lease_id, + EtcdRevisionId& revision_id) { + char* err_msg = nullptr; + int ret = EtcdStoreCreateWithLeaseWrapper((char*)key, (int)key_size, + (char*)value, (int)value_size, + lease_id, &revision_id, &err_msg); + if (ret == -2) { + VLOG(1) << "key=" << std::string(key, key_size) + << ", lease_id=" << lease_id << ", error=" << err_msg; + free(err_msg); + return ErrorCode::ETCD_TRANSACTION_FAIL; + } else if (ret != 0) { + LOG(ERROR) << "key=" << std::string(key, key_size) + << ", lease_id=" << lease_id << ", error=" << err_msg; + free(err_msg); + return ErrorCode::ETCD_OPERATION_ERROR; + } else { + return ErrorCode::OK; + } +} + +ErrorCode EtcdHelper::GrantLease(int64_t lease_ttl, EtcdLeaseId& lease_id) { + char* err_msg = nullptr; + if (0 != EtcdStoreGrantLeaseWrapper(lease_ttl, &lease_id, &err_msg)) { + LOG(ERROR) << "lease_ttl=" << lease_ttl << ", error=" << err_msg; + free(err_msg); + return ErrorCode::ETCD_OPERATION_ERROR; + } + return ErrorCode::OK; +} + +ErrorCode EtcdHelper::WatchUntilDeleted(const char* key, + const size_t key_size) { + char* err_msg = nullptr; + int err_code = + EtcdStoreWatchUntilDeletedWrapper((char*)key, (int)key_size, &err_msg); + if (err_code != 0) { + LOG(ERROR) << "key=" << std::string(key, key_size) + << ", error=" << err_msg; + free(err_msg); + if (err_code == -2) { + return ErrorCode::ETCD_CTX_CANCELLED; + } else { + return ErrorCode::ETCD_OPERATION_ERROR; + } + } + return ErrorCode::OK; +} + +ErrorCode EtcdHelper::CancelWatch(const char* key, const size_t key_size) { + char* err_msg = nullptr; + if (0 != EtcdStoreCancelWatchWrapper((char*)key, (int)key_size, &err_msg)) { + LOG(ERROR) << "key=" << std::string(key, key_size) + << ", error=" << err_msg; + free(err_msg); + return ErrorCode::ETCD_OPERATION_ERROR; + } + return ErrorCode::OK; +} + +ErrorCode EtcdHelper::KeepAlive(EtcdLeaseId lease_id) { + char* err_msg = nullptr; + int err_code = EtcdStoreKeepAliveWrapper(lease_id, &err_msg); + if (err_code != 0) { + LOG(ERROR) << "lease_id=" << lease_id << ", error=" << err_msg; + free(err_msg); + if (err_code == -2) { + return ErrorCode::ETCD_CTX_CANCELLED; + } else { + return ErrorCode::ETCD_OPERATION_ERROR; + } + } + return ErrorCode::OK; +} + +ErrorCode EtcdHelper::CancelKeepAlive(EtcdLeaseId lease_id) { + char* err_msg = nullptr; + if (0 != EtcdStoreCancelKeepAliveWrapper(lease_id, &err_msg)) { + LOG(ERROR) << "Failed to cancel keep lease: " << err_msg; + free(err_msg); + return ErrorCode::ETCD_OPERATION_ERROR; + } + return ErrorCode::OK; +} + +} // namespace mooncake \ No newline at end of file diff --git a/mooncake-store/src/ha_helper.cpp b/mooncake-store/src/ha_helper.cpp new file mode 100644 index 00000000..61e58100 --- /dev/null +++ b/mooncake-store/src/ha_helper.cpp @@ -0,0 +1,181 @@ +#include "ha_helper.h" + +namespace mooncake { + +ErrorCode MasterViewHelper::ConnectToEtcd(const std::string& etcd_endpoints) { + return EtcdHelper::ConnectToEtcdStoreClient(etcd_endpoints); +} + +void MasterViewHelper::ElectLeader(const std::string& master_address, + ViewVersionId& version, + EtcdLeaseId& lease_id) { + while (true) { + // Check if there is already a leader + ViewVersionId current_version = 0; + std::string current_master; + auto ret = EtcdHelper::Get(MASTER_VIEW_KEY, strlen(MASTER_VIEW_KEY), + current_master, current_version); + if (ret != ErrorCode::OK && ret != ErrorCode::ETCD_KEY_NOT_EXIST) { + LOG(ERROR) << "Failed to get current leader: " << ret; + std::this_thread::sleep_for(std::chrono::seconds(1)); + continue; + } else if (ret != ErrorCode::ETCD_KEY_NOT_EXIST) { + LOG(INFO) << "CurrentLeader=" << current_master + << ", CurrentVersion=" << current_version; + // In rare cases, the leader may be ourselves, but it does not + // matter. We will watch the key until it's deleted. + LOG(INFO) << "Waiting for leadership change..."; + auto ret = EtcdHelper::WatchUntilDeleted(MASTER_VIEW_KEY, + strlen(MASTER_VIEW_KEY)); + if (ret != ErrorCode::OK) { + LOG(ERROR) << "Etcd error when waiting for leadership change: " + << ret; + std::this_thread::sleep_for(std::chrono::seconds(1)); + continue; + } + // From now, the key is deleted + } else { + LOG(INFO) << "No leader found, trying to elect self as leader"; + } + + // Here, the key is either deleted or not set. We can + // try to elect ourselves as the leader. We vote ourselfves + // as the leader by trying to creating the key in a transaction. + // The one who successfully creates the key is the leader. + ret = EtcdHelper::GrantLease(ETCD_MASTER_VIEW_LEASE_TTL, lease_id); + if (ret != ErrorCode::OK) { + LOG(ERROR) << "Failed to grant lease: " << ret; + std::this_thread::sleep_for(std::chrono::seconds(1)); + continue; + } + + ret = EtcdHelper::CreateWithLease( + MASTER_VIEW_KEY, strlen(MASTER_VIEW_KEY), master_address.c_str(), + master_address.size(), lease_id, version); + if (ret == ErrorCode::ETCD_TRANSACTION_FAIL) { + LOG(INFO) << "Failed to elect self as leader: " << ret; + std::this_thread::sleep_for(std::chrono::seconds(1)); + continue; + } else if (ret != ErrorCode::OK) { + LOG(ERROR) << "Failed to create key with lease: " << ret; + std::this_thread::sleep_for(std::chrono::seconds(1)); + continue; + } else { + LOG(INFO) << "Successfully elected self as leader"; + return; + } + } +} + +void MasterViewHelper::KeepLeader(EtcdLeaseId lease_id) { + EtcdHelper::KeepAlive(lease_id); +} + +ErrorCode MasterViewHelper::GetMasterView(std::string& master_address, + ViewVersionId& version) { + auto err_code = EtcdHelper::Get(MASTER_VIEW_KEY, strlen(MASTER_VIEW_KEY), + master_address, version); + if (err_code != ErrorCode::OK) { + if (err_code == ErrorCode::ETCD_KEY_NOT_EXIST) { + LOG(ERROR) << "No master is available"; + } else { + LOG(ERROR) << "Failed to get master address due to etcd error"; + } + return err_code; + } else { + LOG(INFO) << "Get master address: " << master_address + << ", version: " << version; + return ErrorCode::OK; + } +} + +MasterServiceSupervisor::MasterServiceSupervisor( + int port, int server_thread_num, bool enable_gc, + bool enable_metric_reporting, int metrics_port, + int64_t default_kv_lease_ttl, double eviction_ratio, + double eviction_high_watermark_ratio, const std::string& etcd_endpoints, + const std::string& local_hostname) + : port_(port), + server_thread_num_(server_thread_num), + enable_gc_(enable_gc), + enable_metric_reporting_(enable_metric_reporting), + metrics_port_(metrics_port), + default_kv_lease_ttl_(default_kv_lease_ttl), + eviction_ratio_(eviction_ratio), + eviction_high_watermark_ratio_(eviction_high_watermark_ratio), + etcd_endpoints_(etcd_endpoints), + local_hostname_(local_hostname) {} + +int MasterServiceSupervisor::Start() { + while (true) { + LOG(INFO) << "Init master service..."; + coro_rpc::coro_rpc_server server(server_thread_num_, port_); + LOG(INFO) << "Init leader election helper..."; + MasterViewHelper mv_helper; + if (mv_helper.ConnectToEtcd(etcd_endpoints_) != ErrorCode::OK) { + LOG(ERROR) << "Failed to connect to etcd endpoints: " + << etcd_endpoints_; + return -1; + } + LOG(INFO) << "Trying to elect self as leader..."; + ViewVersionId version = 0; + EtcdLeaseId lease_id = 0; + mv_helper.ElectLeader(local_hostname_, version, lease_id); + + // Start a thread to keep the leader alive + auto keep_leader_thread = + std::thread([&server, &mv_helper, lease_id]() { + mv_helper.KeepLeader(lease_id); + LOG(INFO) << "Trying to stop server..."; + server.stop(); + }); + + // To prevent potential split-brain, wait long enough for the old leader + // to retire. + const int waiting_time = ETCD_MASTER_VIEW_LEASE_TTL; + std::this_thread::sleep_for(std::chrono::seconds(waiting_time)); + + LOG(INFO) << "Starting master service..."; + mooncake::WrappedMasterService wrapped_master_service( + enable_gc_, default_kv_lease_ttl_, enable_metric_reporting_, + metrics_port_, eviction_ratio_, eviction_high_watermark_ratio_, + version); + mooncake::RegisterRpcService(server, wrapped_master_service); + // Metric reporting is now handled by WrappedMasterService. + + async_simple::Future ec = + server.async_start(); // won't block here + if (ec.hasResult()) { + LOG(ERROR) << "Failed to start master service: " + << ec.result().value(); + auto etcd_err = EtcdHelper::CancelKeepAlive(lease_id); + if (etcd_err != ErrorCode::OK) { + LOG(ERROR) << "Failed to cancel keep leader alive: " + << etcd_err; + } + // Even if CancelKeepAlive fails, the keep alive context are closed. + // We can safely join the keep leader thread. + keep_leader_thread.join(); + return -1; + } + // Block until the server is stopped + auto server_err = std::move(ec).get(); + LOG(ERROR) << "Master service stopped: " << server_err; + + // If the server is closed due to internal errors, we need to manually + // stop keep leader alive. + auto etcd_err = EtcdHelper::CancelKeepAlive(lease_id); + // The error here is predicatable, no need to log it as ERROR. + LOG(INFO) << "Cancel keep leader alive: " << etcd_err; + keep_leader_thread.join(); + } + return 0; +} + +MasterServiceSupervisor::~MasterServiceSupervisor() { + if (server_thread_.joinable()) { + server_thread_.join(); + } +} + +} // namespace mooncake \ No newline at end of file diff --git a/mooncake-store/src/master.cpp b/mooncake-store/src/master.cpp index 1afc179e..cf7b3b45 100644 --- a/mooncake-store/src/master.cpp +++ b/mooncake-store/src/master.cpp @@ -5,8 +5,10 @@ #include #include +#include "ha_helper.h" #include "rpc_service.h" #include "types.h" + using namespace coro_rpc; using namespace async_simple; using namespace async_simple::coro; @@ -17,9 +19,12 @@ DEFINE_bool(enable_gc, false, "Enable garbage collection"); DEFINE_bool(enable_metric_reporting, true, "Enable periodic metric reporting"); DEFINE_int32(metrics_port, 9003, "Port for HTTP metrics server to listen on"); DEFINE_uint64(default_kv_lease_ttl, mooncake::DEFAULT_DEFAULT_KV_LEASE_TTL, - "Default lease time for kv objects"); -DEFINE_double(eviction_ratio, mooncake::DEFAULT_EVICTION_RATIO, "Ratio of objects to evict when storage space is full"); -DEFINE_double(eviction_high_watermark_ratio, mooncake::DEFAULT_EVICTION_HIGH_WATERMARK_RATIO, "Ratio of high watermark trigger eviction"); + "Default lease time for kv objects"); +DEFINE_double(eviction_ratio, mooncake::DEFAULT_EVICTION_RATIO, + "Ratio of objects to evict when storage space is full"); +DEFINE_double(eviction_high_watermark_ratio, + mooncake::DEFAULT_EVICTION_HIGH_WATERMARK_RATIO, + "Ratio of high watermark trigger eviction"); DEFINE_validator(eviction_ratio, [](const char* flagname, double value) { if (value < 0.0 || value > 1.0) { LOG(FATAL) << "Eviction ratio must be between 0.0 and 1.0"; @@ -27,18 +32,19 @@ DEFINE_validator(eviction_ratio, [](const char* flagname, double value) { } return true; }); +DEFINE_bool(enable_ha, false, + "Enable high availability, which depends on ETCD"); +DEFINE_string( + etcd_endpoints, "", + "Endpoints of ETCD server, separated by semicolon, required in HA mode"); +DEFINE_string(local_hostname, "", + "Local host address (IP:Port), required in HA mode"); int main(int argc, char* argv[]) { easylog::set_min_severity(easylog::Severity::WARN); // Initialize gflags gflags::ParseCommandLineFlags(&argc, &argv, true); - // init rpc server - coro_rpc_server server( - /*thread=*/std::min( - FLAGS_max_threads, - static_cast(std::thread::hardware_concurrency())), - /*port=*/FLAGS_port); LOG(INFO) << "Master service started on port " << FLAGS_port << ", enable_gc=" << FLAGS_enable_gc << ", max_threads=" << FLAGS_max_threads @@ -46,40 +52,53 @@ int main(int argc, char* argv[]) { << ", metrics_port=" << FLAGS_metrics_port << ", default_kv_lease_ttl=" << FLAGS_default_kv_lease_ttl << ", eviction_ratio=" << FLAGS_eviction_ratio - << ", eviction_high_watermark_ratio=" << FLAGS_eviction_high_watermark_ratio; + << ", eviction_high_watermark_ratio=" + << FLAGS_eviction_high_watermark_ratio + << ", enable_ha=" << FLAGS_enable_ha + << ", etcd_endpoints=" << FLAGS_etcd_endpoints + << ", local_hostname=" << FLAGS_local_hostname; - mooncake::WrappedMasterService wrapped_master_service( - FLAGS_enable_gc, FLAGS_default_kv_lease_ttl, - FLAGS_enable_metric_reporting, FLAGS_metrics_port, - FLAGS_eviction_ratio, FLAGS_eviction_high_watermark_ratio); - server.register_handler<&mooncake::WrappedMasterService::ExistKey>( - &wrapped_master_service); - server.register_handler<&mooncake::WrappedMasterService::GetReplicaList>( - &wrapped_master_service); - server.register_handler<&mooncake::WrappedMasterService::BatchGetReplicaList>( - &wrapped_master_service); - server.register_handler<&mooncake::WrappedMasterService::PutStart>( - &wrapped_master_service); - server.register_handler<&mooncake::WrappedMasterService::PutEnd>( - &wrapped_master_service); - server.register_handler<&mooncake::WrappedMasterService::PutRevoke>( - &wrapped_master_service); - server.register_handler<&mooncake::WrappedMasterService::BatchPutStart>( - &wrapped_master_service); - server.register_handler<&mooncake::WrappedMasterService::BatchPutEnd>( - &wrapped_master_service); - server.register_handler<&mooncake::WrappedMasterService::BatchPutRevoke>( - &wrapped_master_service); - server.register_handler<&mooncake::WrappedMasterService::Remove>( - &wrapped_master_service); - server.register_handler<&mooncake::WrappedMasterService::RemoveAll>( - &wrapped_master_service); - server.register_handler<&mooncake::WrappedMasterService::MountSegment>( - &wrapped_master_service); - server.register_handler<&mooncake::WrappedMasterService::UnmountSegment>( - &wrapped_master_service); + int server_thread_num = + std::min(FLAGS_max_threads, + static_cast(std::thread::hardware_concurrency())); - // Metric reporting is now handled by WrappedMasterService + if (FLAGS_enable_ha && FLAGS_etcd_endpoints.empty()) { + LOG(FATAL) << "Etcd endpoints must be set when enable_ha is true"; + return 1; + } + if (!FLAGS_enable_ha && !FLAGS_etcd_endpoints.empty()) { + LOG(WARNING) + << "Etcd endpoints are set but will not be used in non-HA mode"; + } - return !server.start(); + if (FLAGS_enable_ha && FLAGS_local_hostname.empty()) { + LOG(FATAL) << "Local hostname must be set when enable_ha is true"; + return 1; + } + if (!FLAGS_enable_ha && !FLAGS_local_hostname.empty()) { + LOG(WARNING) + << "Local hostname is set but will not be used in non-HA mode"; + } + + if (FLAGS_enable_ha) { + mooncake::MasterServiceSupervisor supervisor( + FLAGS_port, server_thread_num, FLAGS_enable_gc, + FLAGS_enable_metric_reporting, FLAGS_metrics_port, + FLAGS_default_kv_lease_ttl, FLAGS_eviction_ratio, + FLAGS_eviction_high_watermark_ratio, FLAGS_etcd_endpoints, + FLAGS_local_hostname); + + return supervisor.Start(); + } else { + // version is not used in non-HA mode, just pass a dummy value + mooncake::ViewVersionId version = 0; + coro_rpc::coro_rpc_server server(server_thread_num, FLAGS_port); + mooncake::WrappedMasterService wrapped_master_service( + FLAGS_enable_gc, FLAGS_default_kv_lease_ttl, + FLAGS_enable_metric_reporting, FLAGS_metrics_port, + FLAGS_eviction_ratio, FLAGS_eviction_high_watermark_ratio, version); + + mooncake::RegisterRpcService(server, wrapped_master_service); + return server.start(); + } } diff --git a/mooncake-store/src/master_client.cpp b/mooncake-store/src/master_client.cpp index e4c2ee48..6cd3b83c 100644 --- a/mooncake-store/src/master_client.cpp +++ b/mooncake-store/src/master_client.cpp @@ -27,7 +27,7 @@ ErrorCode MasterClient::Connect(const std::string& master_addr) { auto result = coro::syncAwait(client_.connect(master_addr)); if (result.val() != 0) { LOG(ERROR) << "Failed to connect to master: " << result.message(); - return ErrorCode::INTERNAL_ERROR; + return ErrorCode::RPC_FAIL; } timer.LogResponse("error_code=", ErrorCode::OK); return ErrorCode::OK; @@ -388,4 +388,30 @@ UnmountSegmentResponse MasterClient::UnmountSegment( return result.value(); } +PingResponse MasterClient::Ping() { + ScopedVLogTimer timer(1, "MasterClient::Ping"); + timer.LogRequest("action=ping"); + + auto request_result = + client_.send_request<&WrappedMasterService::Ping>(); + std::optional result = + coro::syncAwait([&]() -> coro::Lazy> { + auto result = co_await co_await request_result; + if (!result) { + LOG(ERROR) << "Failed to ping master: " << result.error().msg; + co_return std::nullopt; + } + co_return result->result(); + }()); + + if (!result) { + auto response = PingResponse{0, ErrorCode::RPC_FAIL}; + timer.LogResponseJson(response); + return response; + } + + timer.LogResponseJson(result.value()); + return result.value(); +} + } // namespace mooncake diff --git a/mooncake-store/src/master_metric_manager.cpp b/mooncake-store/src/master_metric_manager.cpp index 29e90b63..9632a4b4 100644 --- a/mooncake-store/src/master_metric_manager.cpp +++ b/mooncake-store/src/master_metric_manager.cpp @@ -71,6 +71,8 @@ MasterMetricManager::MasterMetricManager() unmount_segment_failures_( "master_unmount_segment_failures_total", "Total number of failed UnmountSegment requests"), + ping_requests_("master_ping_requests_total", + "Total number of ping requests received"), // Initialize Eviction Counters eviction_success_("master_successful_evictions_total", @@ -183,6 +185,9 @@ void MasterMetricManager::inc_unmount_segment_requests(int64_t val) { void MasterMetricManager::inc_unmount_segment_failures(int64_t val) { unmount_segment_failures_.inc(val); } +void MasterMetricManager::inc_ping_requests(int64_t val) { + ping_requests_.inc(val); +} int64_t MasterMetricManager::get_put_start_requests() { return put_start_requests_.value(); @@ -256,6 +261,10 @@ int64_t MasterMetricManager::get_unmount_segment_failures() { return unmount_segment_failures_.value(); } +int64_t MasterMetricManager::get_ping_requests() { + return ping_requests_.value(); +} + // Eviction Metrics void MasterMetricManager::inc_eviction_success(int64_t key_count, int64_t size) { evicted_key_count_.inc(key_count); @@ -325,6 +334,7 @@ std::string MasterMetricManager::serialize_metrics() { serialize_metric(mount_segment_failures_); serialize_metric(unmount_segment_requests_); serialize_metric(unmount_segment_failures_); + serialize_metric(ping_requests_); // Serialize Eviction Counters serialize_metric(eviction_success_); @@ -373,6 +383,7 @@ std::string MasterMetricManager::get_summary_string() { int64_t remove_fails = remove_failures_.value(); int64_t remove_all = remove_all_requests_.value(); int64_t remove_all_fails = remove_all_failures_.value(); + int64_t pings = ping_requests_.value(); // Eviction counters int64_t eviction_success = eviction_success_.value(); @@ -385,7 +396,7 @@ std::string MasterMetricManager::get_summary_string() { << format_bytes(capacity); if (capacity > 0) { ss << " (" << std::fixed << std::setprecision(1) - << (allocated / capacity * 100.0) << "%)"; + << ((double) allocated / (double)capacity * 100.0) << "%)"; } ss << " | Keys: " << keys; @@ -397,7 +408,8 @@ std::string MasterMetricManager::get_summary_string() { ss << "Get=" << get_replicas - get_replica_fails << "/" << get_replicas << ", "; ss << "Exist=" << exist_keys - exist_key_fails << "/" << exist_keys << ", "; ss << "Del=" << removes - remove_fails << "/" << removes << ", "; - ss << "DelAll=" << remove_all - remove_all_fails << "/" << remove_all; + ss << "DelAll=" << remove_all - remove_all_fails << "/" << remove_all << ", "; + ss << "Ping=" << pings; // Eviction summary ss << " | Eviction: " diff --git a/mooncake-store/src/types.cpp b/mooncake-store/src/types.cpp index f27918e2..c8b616f3 100644 --- a/mooncake-store/src/types.cpp +++ b/mooncake-store/src/types.cpp @@ -20,8 +20,13 @@ const std::string& toString(ErrorCode errorCode) noexcept { {ErrorCode::REPLICA_IS_NOT_READY, "REPLICA_IS_NOT_READY"}, {ErrorCode::OBJECT_NOT_FOUND, "OBJECT_NOT_FOUND"}, {ErrorCode::OBJECT_ALREADY_EXISTS, "OBJECT_ALREADY_EXISTS"}, + {ErrorCode::OBJECT_HAS_LEASE, "OBJECT_HAS_LEASE"}, {ErrorCode::TRANSFER_FAIL, "TRANSFER_FAIL"}, {ErrorCode::RPC_FAIL, "RPC_FAIL"}, + {ErrorCode::ETCD_OPERATION_ERROR, "ETCD_OPERATION_ERROR"}, + {ErrorCode::ETCD_KEY_NOT_EXIST, "ETCD_KEY_NOT_EXIST"}, + {ErrorCode::ETCD_TRANSACTION_FAIL, "ETCD_TRANSACTION_FAIL"}, + {ErrorCode::ETCD_CTX_CANCELLED, "ETCD_CTX_CANCELLED"}, }; auto it = errorCodeMap.find(errorCode); diff --git a/mooncake-store/tests/CMakeLists.txt b/mooncake-store/tests/CMakeLists.txt index f3bb6dd3..d5bc90d8 100644 --- a/mooncake-store/tests/CMakeLists.txt +++ b/mooncake-store/tests/CMakeLists.txt @@ -26,16 +26,13 @@ target_link_libraries(client_integration_test PUBLIC add_test(NAME client_integration_test COMMAND client_integration_test) add_executable(master_metrics_test master_metrics_test.cpp) -target_link_libraries(master_metrics_test PUBLIC - mooncake_store - cachelib_memory_allocator - glog - gtest - gtest_main - pthread -) +target_link_libraries(master_metrics_test PUBLIC mooncake_store cachelib_memory_allocator glog gtest gtest_main pthread) add_test(NAME master_metrics_test COMMAND master_metrics_test) +add_executable(high_availability_test high_availability_test.cpp) +target_link_libraries(high_availability_test PUBLIC mooncake_store cachelib_memory_allocator glog gtest gtest_main pthread ${ETCD_WRAPPER_LIB}) +add_test(NAME high_availability_test COMMAND high_availability_test) + add_executable(stress_workload_test stress_workload_test.cpp) target_link_libraries(stress_workload_test PUBLIC mooncake_store diff --git a/mooncake-store/tests/high_availability_test.cpp b/mooncake-store/tests/high_availability_test.cpp new file mode 100644 index 00000000..a908ef30 --- /dev/null +++ b/mooncake-store/tests/high_availability_test.cpp @@ -0,0 +1,209 @@ +#include +#include +#include + +#include +#include + +#include "etcd_helper.h" +#include "ha_helper.h" +#include "types.h" + +namespace mooncake { +namespace testing { + +DEFINE_string(etcd_endpoints, "0.0.0.0:2379", "Etcd endpoints"); +DEFINE_string(etcd_test_key_prefix, "mooncake-store/test/", + "The prefix of the test keys in ETCD"); + +class HighAvailabilityTest : public ::testing::Test { + protected: + static void SetUpTestSuite() { + // Initialize glog + google::InitGoogleLogging("ClientIntegrationTest"); + + // Set VLOG level to 1 for detailed logs + google::SetVLOGLevel("*", 1); + FLAGS_logtostderr = 1; + + // Initialize etcd client + ASSERT_EQ(ErrorCode::OK, + EtcdHelper::ConnectToEtcdStoreClient(FLAGS_etcd_endpoints)); + } + + static void TearDownTestSuite() { google::ShutdownGoogleLogging(); } +}; + +TEST_F(HighAvailabilityTest, EtcdBasicOperations) { + // == Test grant lease, create kv and get kv == + int64_t lease_ttl = 10; + std::vector keys; + std::vector values; + // Ordinary key-value pair + keys.push_back(FLAGS_etcd_test_key_prefix + std::string("test_key1")); + values.push_back("test_value1"); + // Key-value pair with null bytes in the middle + keys.push_back(FLAGS_etcd_test_key_prefix + std::string("test_\0\0key2")); + values.push_back("test_\0\0value2"); + // Key-value pair with null bytes at the end + keys.push_back(FLAGS_etcd_test_key_prefix + std::string("test_key3\0\0")); + values.push_back("test_value3\0\0"); + // Key-value pair with null bytes at the beginning + keys.push_back(FLAGS_etcd_test_key_prefix + std::string("\0\0test_key4")); + values.push_back("\0\0test_value4"); + + for (size_t i = 0; i < keys.size(); i++) { + auto &key = keys[i]; + auto &value = values[i]; + EtcdLeaseId lease_id; + EtcdRevisionId version = 0; + + ASSERT_EQ(ErrorCode::OK, EtcdHelper::GrantLease(lease_ttl, lease_id)); + ASSERT_EQ(ErrorCode::OK, EtcdHelper::CreateWithLease( + key.c_str(), key.size(), value.c_str(), + value.size(), lease_id, version)); + std::string get_value; + EtcdRevisionId get_version; + ASSERT_EQ(ErrorCode::OK, EtcdHelper::Get(key.c_str(), key.size(), + get_value, get_version)); + ASSERT_EQ(value, get_value); + ASSERT_EQ(version, get_version); + } + + // == Test keep alive and cancel keep alive == + lease_ttl = 2; + EtcdLeaseId lease_id; + ASSERT_EQ(ErrorCode::OK, EtcdHelper::GrantLease(lease_ttl, lease_id)); + + std::promise promise; + std::future future = promise.get_future(); + + std::thread keep_alive_thread([&]() { + ErrorCode result = EtcdHelper::KeepAlive(lease_id); + promise.set_value(result); + }); + // Check if keep alive can extend the lease's life time + ASSERT_NE(future.wait_for(std::chrono::seconds(lease_ttl * 3)), + std::future_status::ready); + std::string key = + FLAGS_etcd_test_key_prefix + std::string("keep_alive_key"); + std::string value = "keep_alive_value"; + EtcdRevisionId version = 0; + ASSERT_EQ(ErrorCode::OK, EtcdHelper::CreateWithLease( + key.c_str(), key.size(), value.c_str(), + value.size(), lease_id, version)); + ASSERT_EQ(ErrorCode::OK, + EtcdHelper::Get(key.c_str(), key.size(), value, version)); + + // Test cancel keep alive + ASSERT_EQ(ErrorCode::OK, EtcdHelper::CancelKeepAlive(lease_id)); + ASSERT_EQ(future.wait_for(std::chrono::seconds(1)), + std::future_status::ready); + ASSERT_EQ(future.get(), ErrorCode::ETCD_CTX_CANCELLED); + keep_alive_thread.join(); + + // == Test watch key and cancel watch == + lease_ttl = 2; + ASSERT_EQ(ErrorCode::OK, EtcdHelper::GrantLease(lease_ttl, lease_id)); + std::string watch_key = + FLAGS_etcd_test_key_prefix + std::string("watch_key"); + std::string watch_value = "watch_value"; + EtcdRevisionId watch_version = 0; + ASSERT_EQ(ErrorCode::OK, + EtcdHelper::CreateWithLease( + watch_key.c_str(), watch_key.size(), watch_value.c_str(), + watch_value.size(), lease_id, watch_version)); + + promise = std::promise(); + future = promise.get_future(); + keep_alive_thread = std::thread([&]() { EtcdHelper::KeepAlive(lease_id); }); + std::thread watch_thread([&]() { + ErrorCode result = + EtcdHelper::WatchUntilDeleted(watch_key.c_str(), watch_key.size()); + promise.set_value(result); + }); + // Check the watch thread is blocked if the key is not deleted + ASSERT_NE(future.wait_for(std::chrono::seconds(lease_ttl * 3)), + std::future_status::ready); + // Check the watch thread returns after the key is deleted + ASSERT_EQ(ErrorCode::OK, EtcdHelper::CancelKeepAlive(lease_id)); + ASSERT_EQ(future.wait_for(std::chrono::seconds(lease_ttl * 3)), + std::future_status::ready); + ASSERT_EQ(future.get(), ErrorCode::OK); + watch_thread.join(); + keep_alive_thread.join(); + + // Test cancel watch + lease_ttl = 10; + int64_t watch_wait_time = 2; + ASSERT_EQ(ErrorCode::OK, EtcdHelper::GrantLease(lease_ttl, lease_id)); + watch_key = FLAGS_etcd_test_key_prefix + std::string("watch_key2"); + watch_value = "watch_value2"; + ASSERT_EQ(ErrorCode::OK, + EtcdHelper::CreateWithLease( + watch_key.c_str(), watch_key.size(), watch_value.c_str(), + watch_value.size(), lease_id, watch_version)); + + promise = std::promise(); + future = promise.get_future(); + watch_thread = std::thread([&]() { + ErrorCode result = + EtcdHelper::WatchUntilDeleted(watch_key.c_str(), watch_key.size()); + promise.set_value(result); + }); + // Wait for the watch thread to call WatchUntilDeleted + std::this_thread::sleep_for(std::chrono::seconds(1)); + // Cancel the watch + ASSERT_EQ(ErrorCode::OK, + EtcdHelper::CancelWatch(watch_key.c_str(), watch_key.size())); + ASSERT_EQ(future.wait_for(std::chrono::seconds(watch_wait_time)), + std::future_status::ready); + ASSERT_EQ(future.get(), ErrorCode::ETCD_CTX_CANCELLED); + watch_thread.join(); +} + +TEST_F(HighAvailabilityTest, BasicMasterViewOperations) { + MasterViewHelper mv_helper; + mv_helper.ConnectToEtcd(FLAGS_etcd_endpoints); + std::string master_address = "0.0.0.0:8888"; + ViewVersionId version = 0; + + // Initially, the master view is not set + ASSERT_NE(ErrorCode::OK, mv_helper.GetMasterView(master_address, version)); + + // Elect and keep leader + EtcdLeaseId lease_id = 0; + mv_helper.ElectLeader(master_address, version, lease_id); + std::thread keep_alive_thread([&]() { mv_helper.KeepLeader(lease_id); }); + + // Check the master view is correctly set + std::string get_master_address; + ASSERT_EQ(ErrorCode::OK, + mv_helper.GetMasterView(get_master_address, version)); + ASSERT_EQ(get_master_address, master_address); + + // Check the master view does not change + std::this_thread::sleep_for( + std::chrono::seconds(ETCD_MASTER_VIEW_LEASE_TTL + 2)); + ASSERT_EQ(ErrorCode::OK, + mv_helper.GetMasterView(get_master_address, version)); + ASSERT_EQ(get_master_address, master_address); + + EtcdHelper::CancelKeepAlive(lease_id); + keep_alive_thread.join(); +} + +} // namespace testing + +} // namespace mooncake + +int main(int argc, char **argv) { + // Initialize Google's flags library + gflags::ParseCommandLineFlags(&argc, &argv, true); + + // Initialize Google Test + ::testing::InitGoogleTest(&argc, argv); + + // Run all tests + return RUN_ALL_TESTS(); +}