diff --git a/.gitmodules b/.gitmodules index e8ec2dd74..227ca4681 100644 --- a/.gitmodules +++ b/.gitmodules @@ -9,3 +9,6 @@ path = 3rdparty/spdk url = https://github.com/spdk/spdk.git shallow = true +[submodule "3rdparty/redis-plus-plus"] + path = 3rdparty/redis-plus-plus + url = https://github.com/sewenew/redis-plus-plus.git diff --git a/3rdparty/redis-plus-plus b/3rdparty/redis-plus-plus new file mode 160000 index 000000000..a63ac43bf --- /dev/null +++ b/3rdparty/redis-plus-plus @@ -0,0 +1 @@ +Subproject commit a63ac43bf192772910b52e27cd2b42a6098a0071 diff --git a/include/mori/metrics/prometheus_metrics_server.hpp b/include/mori/metrics/prometheus_metrics_server.hpp index e2bc9cac7..ca71633a8 100644 --- a/include/mori/metrics/prometheus_metrics_server.hpp +++ b/include/mori/metrics/prometheus_metrics_server.hpp @@ -114,6 +114,12 @@ class MetricsServer { void observe(std::string_view name, std::string_view help, const std::vector& bounds, double value); + // Record one histogram observation into a LABELED series (e.g. + // {op=...,backend=...}). Same first-write-wins bucket layout as the unlabeled + // observe(); reuses the labeled-histogram storage that observeAggregated uses. + void observe(std::string_view name, std::string_view help, const Labels& labels, + const std::vector& bounds, double value); + // Merge a pre-aggregated histogram batch into a labeled series. // `bucket_counts` is CUMULATIVE (bucket_counts[i] = #obs with value <= // bounds[i]) and is added per-bucket into the stored series; `count` and diff --git a/setup.py b/setup.py index 92500de81..4b438eaf7 100644 --- a/setup.py +++ b/setup.py @@ -381,6 +381,7 @@ def build_extension(self, ext: Extension) -> None: build_tests = os.environ.get("BUILD_TESTS", "OFF") build_umbp = "ON" if build_umbp_enabled else "OFF" build_umbp_spdk = "ON" if build_umbp_spdk_enabled else "OFF" + use_redis_backend = os.environ.get("USE_REDIS_BACKEND", "OFF") build_xla_ffi_ops = os.environ.get("BUILD_XLA_FFI_OPS", "OFF") with_mpi = ( "ON" @@ -411,6 +412,7 @@ def build_extension(self, ext: Extension) -> None: f"-DBUILD_BENCHMARK={build_benchmark}", f"-DBUILD_TESTS={build_tests}", f"-DBUILD_UMBP={build_umbp}", + f"-DUSE_REDIS_BACKEND={use_redis_backend}", f"-DUSE_SPDK={build_umbp_spdk}", f"-DWITH_MPI={with_mpi}", "-DBUILD_TORCH_BOOTSTRAP=OFF", diff --git a/src/application/transport/rdma/providers/ibverbs/ibverbs.cpp b/src/application/transport/rdma/providers/ibverbs/ibverbs.cpp index 3ab1b7bd4..04fc9c599 100644 --- a/src/application/transport/rdma/providers/ibverbs/ibverbs.cpp +++ b/src/application/transport/rdma/providers/ibverbs/ibverbs.cpp @@ -170,25 +170,32 @@ RdmaEndpoint IBVerbsDeviceContext::CreateRdmaEndpoint(const RdmaEndpointConfig& // TODO: we need to add more options in config, include min cqe num for ib_create_cq endpoint.ibvHandle.compCh = config.withCompChannel ? ibv_create_comp_channel(context) : nullptr; if (config.withCompChannel && !endpoint.ibvHandle.compCh) { - MORI_APP_ERROR("ibv_create_comp_channel failed: errno={} ({}); dev={}", errno, strerror(errno), + // Capture errno immediately: any intervening libc call before we read it + // (e.g. isatty() inside the logging sink when stderr is not a tty, which + // sets ENOTTY) would otherwise overwrite the real failure code. + const int err = errno; + MORI_APP_ERROR("ibv_create_comp_channel failed: errno={} ({}); dev={}", err, strerror(err), GetRdmaDevice()->Name()); - throw std::runtime_error("ibv_create_comp_channel failed: " + std::string(strerror(errno))); + throw std::runtime_error("ibv_create_comp_channel failed: " + std::string(strerror(err))); } endpoint.ibvHandle.cq = ibv_create_cq(context, config.maxCqeNum, NULL, endpoint.ibvHandle.compCh, 0); if (!endpoint.ibvHandle.cq) { + // Capture errno before the lock/log path can overwrite it (see the note at + // the comp-channel site). + const int err = errno; size_t cqPoolSize = 0; { std::lock_guard lock(poolMu); cqPoolSize = cqPool.size(); } MORI_APP_ERROR( - "ibv_create_cq failed: errno={} ({}); dev={} max_cqe={} dev_max_cqe={} cqs_in_pool={}", - errno, strerror(errno), GetRdmaDevice()->Name(), config.maxCqeNum, - deviceAttr->orig_attr.max_cqe, cqPoolSize); + "ibv_create_cq failed: errno={} ({}); dev={} max_cqe={} dev_max_cqe={} cqs_in_pool={}", err, + strerror(err), GetRdmaDevice()->Name(), config.maxCqeNum, deviceAttr->orig_attr.max_cqe, + cqPoolSize); if (endpoint.ibvHandle.compCh) ibv_destroy_comp_channel(endpoint.ibvHandle.compCh); - throw std::runtime_error("ibv_create_cq failed: " + std::string(strerror(errno))); + throw std::runtime_error("ibv_create_cq failed: " + std::string(strerror(err))); } // TODO: should also manage the lifecycle of completion channel && srq @@ -216,6 +223,12 @@ RdmaEndpoint IBVerbsDeviceContext::CreateRdmaEndpoint(const RdmaEndpointConfig& .qp_type = IBV_QPT_RC}; endpoint.ibvHandle.qp = ibv_create_qp(pd, &qpAttr); if (!endpoint.ibvHandle.qp) { + // Capture errno immediately. Otherwise the lock/log path below overwrites + // it before we read it: spdlog's console sink calls isatty(stderr), which + // sets errno=ENOTTY when stderr is not a terminal, masking the real error + // (commonly EINVAL when the requested QP capacity — max_send_wr x per-WQE + // size from sge+inline — exceeds the device's per-QP work-queue budget). + const int err = errno; size_t qpPoolSize = 0; { std::lock_guard lock(poolMu); @@ -224,13 +237,13 @@ RdmaEndpoint IBVerbsDeviceContext::CreateRdmaEndpoint(const RdmaEndpointConfig& MORI_APP_ERROR( "ibv_create_qp failed: errno={} ({}); dev={} port={} max_send_wr={} max_recv_wr={} " "max_send_sge={} max_cqe={} dev_caps(max_qp_wr={} max_qp={} max_cqe={}) qps_in_pool={}", - errno, strerror(errno), GetRdmaDevice()->Name(), config.portId, qpAttr.cap.max_send_wr, + err, strerror(err), GetRdmaDevice()->Name(), config.portId, qpAttr.cap.max_send_wr, qpAttr.cap.max_recv_wr, qpAttr.cap.max_send_sge, config.maxCqeNum, deviceAttr->orig_attr.max_qp_wr, deviceAttr->orig_attr.max_qp, deviceAttr->orig_attr.max_cqe, qpPoolSize); ibv_destroy_cq(endpoint.ibvHandle.cq); if (endpoint.ibvHandle.compCh) ibv_destroy_comp_channel(endpoint.ibvHandle.compCh); - throw std::runtime_error("ibv_create_qp failed: " + std::string(strerror(errno))); + throw std::runtime_error("ibv_create_qp failed: " + std::string(strerror(err))); } endpoint.handle.qpn = endpoint.ibvHandle.qp->qp_num; diff --git a/src/metrics/prometheus_metrics_server.cpp b/src/metrics/prometheus_metrics_server.cpp index 116a35227..5fefb5f5f 100644 --- a/src/metrics/prometheus_metrics_server.cpp +++ b/src/metrics/prometheus_metrics_server.cpp @@ -179,6 +179,25 @@ void MetricsServer::observe(std::string_view name, std::string_view help, static std::string FormatLabels( const mori::metrics::MetricsServer::Labels& labels); // forward decl +void MetricsServer::observe(std::string_view name, std::string_view help, const Labels& labels, + const std::vector& bounds, double value) { + auto label_str = FormatLabels(labels); + std::lock_guard lk(mutex_); + auto& family = labeled_histograms_[std::string(name)]; + family.help = help; + auto& s = family.series[label_str]; + if (s.bounds.empty()) { + s.bounds = bounds; + s.bucket_counts.assign(bounds.size(), 0u); + } + // Increment all cumulative buckets whose upper bound >= value. + for (std::size_t i = 0; i < s.bounds.size(); ++i) { + if (value <= s.bounds[i]) s.bucket_counts[i]++; + } + s.count++; + s.sum += value; +} + void MetricsServer::observeAggregated(std::string_view name, std::string_view help, const Labels& labels, const std::vector& bounds, const std::vector& bucket_counts, uint64_t count, diff --git a/src/umbp/CMakeLists.txt b/src/umbp/CMakeLists.txt index 68a47c367..152b8398a 100644 --- a/src/umbp/CMakeLists.txt +++ b/src/umbp/CMakeLists.txt @@ -309,6 +309,7 @@ add_library( ${UMBP_PEER_PROTO_SRCS} ${UMBP_PEER_GRPC_SRCS} distributed/master/in_memory_master_metadata_store.cpp + distributed/master/master_metadata_store_factory.cpp distributed/master/master_server.cpp distributed/master/master_client.cpp distributed/master/rpc_latency_timer.cpp @@ -347,6 +348,53 @@ target_link_libraries( target_compile_features(umbp_common PUBLIC cxx_std_17) +# ---------- Optional Redis / RESP master metadata backend -------------------- +# Makes the master stateless by storing all metadata in an external RESP store +# (Redis / Dragonfly / Valkey). OFF by default so existing builds are +# unaffected; enable with -DUSE_REDIS_BACKEND=ON. The RESP client seam is built +# on hiredis (RespClient isolates the client library from the store). +option(USE_REDIS_BACKEND + "Build the RESP-compatible Redis master metadata backend" OFF) +if(USE_REDIS_BACKEND) + find_path(HIREDIS_INCLUDE_DIR hiredis/hiredis.h REQUIRED) + find_library(HIREDIS_LIB NAMES hiredis REQUIRED) + message( + STATUS + "UMBP Redis backend ENABLED (hiredis lib: ${HIREDIS_LIB}, include: ${HIREDIS_INCLUDE_DIR})" + ) + + # redis-plus-plus (3rdparty submodule) — the cluster/sentinel/pooling-capable + # client the Redis backend standardizes on. Built from source so any build + # image works without an extra system package; only the client library needs + # a submodule (hiredis stays the system lib). Static + PIC (umbp_common is + # linked into the shared pybind module), C++17, tests off. The store code does + # not depend on it yet (that migration is a later step); building it here + # de-risks the dependency (toolchain, -Werror, link) before the seam rebase. + set(_REDISPP_DIR "${CMAKE_SOURCE_DIR}/3rdparty/redis-plus-plus") + if(NOT EXISTS "${_REDISPP_DIR}/CMakeLists.txt") + message( + FATAL_ERROR + "USE_REDIS_BACKEND=ON but 3rdparty/redis-plus-plus is empty; run " + "'git submodule update --init --recursive 3rdparty/redis-plus-plus'") + endif() + set(REDIS_PLUS_PLUS_BUILD_TEST OFF CACHE BOOL "" FORCE) + set(REDIS_PLUS_PLUS_BUILD_SHARED OFF CACHE BOOL "" FORCE) + set(REDIS_PLUS_PLUS_BUILD_STATIC ON CACHE BOOL "" FORCE) + set(REDIS_PLUS_PLUS_CXX_STANDARD 17 CACHE STRING "" FORCE) + add_subdirectory("${_REDISPP_DIR}" + "${CMAKE_BINARY_DIR}/3rdparty/redis-plus-plus") + + target_sources( + umbp_common PRIVATE distributed/master/redis/resp_client.cpp + distributed/master/redis/resp_cluster_client.cpp + distributed/master/redis_master_metadata_store.cpp) + target_include_directories(umbp_common PRIVATE ${HIREDIS_INCLUDE_DIR}) + # redis++_static carries its own (and hiredis') include dirs; link it before + # hiredis so the static archive's hiredis symbols resolve. + target_link_libraries(umbp_common PUBLIC redis++_static ${HIREDIS_LIB}) + target_compile_definitions(umbp_common PUBLIC USE_REDIS_BACKEND) +endif() + if(MORI_UMBP_TESTING) target_compile_definitions(umbp_common PUBLIC MORI_UMBP_TESTING) endif() diff --git a/src/umbp/distributed/master/eviction_manager.cpp b/src/umbp/distributed/master/eviction_manager.cpp index 729fb410d..323f1dc7d 100644 --- a/src/umbp/distributed/master/eviction_manager.cpp +++ b/src/umbp/distributed/master/eviction_manager.cpp @@ -71,7 +71,13 @@ void EvictionManager::EvictionLoop() { [this] { return !running_.load(std::memory_order_relaxed); }); } if (!running_.load(std::memory_order_relaxed)) break; - RunOnce(); + // A metadata-store hiccup (e.g. RESP transport error) must not kill the + // eviction thread; log and retry on the next tick. + try { + RunOnce(); + } catch (const std::exception& e) { + MORI_UMBP_ERROR("[EvictionManager] metadata-store error: {}", e.what()); + } } } diff --git a/src/umbp/distributed/master/master_metadata_store_factory.cpp b/src/umbp/distributed/master/master_metadata_store_factory.cpp new file mode 100644 index 000000000..e65d10bc3 --- /dev/null +++ b/src/umbp/distributed/master/master_metadata_store_factory.cpp @@ -0,0 +1,274 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. +#include "umbp/distributed/master/master_metadata_store_factory.h" + +#include +#include +#include +#include +#include +#include + +#include "mori/utils/mori_log.hpp" +#include "umbp/distributed/master/in_memory_master_metadata_store.h" + +#ifdef USE_REDIS_BACKEND +#include "umbp/distributed/master/redis/resp_cluster_client.h" +#include "umbp/distributed/master/redis_master_metadata_store.h" +#endif + +namespace mori::umbp { + +namespace { + +std::string GetEnvStr(const char* key, const std::string& def) { + const char* v = std::getenv(key); + return (v != nullptr && *v != '\0') ? std::string(v) : def; +} + +int GetEnvInt(const char* key, int def) { + const char* v = std::getenv(key); + if (v == nullptr || *v == '\0') return def; + char* end = nullptr; + const long n = std::strtol(v, &end, 10); + if (end == v || n <= 0) { + MORI_UMBP_WARN("[MetadataStore] ignoring invalid {}='{}'", key, v); + return def; + } + return static_cast(n); +} + +bool GetEnvBool(const char* key, bool def) { + const char* v = std::getenv(key); + if (v == nullptr || *v == '\0') return def; + const std::string s(v); + if (s == "1" || s == "true" || s == "TRUE" || s == "yes" || s == "on") return true; + if (s == "0" || s == "false" || s == "FALSE" || s == "no" || s == "off") return false; + MORI_UMBP_WARN("[MetadataStore] ignoring invalid {}='{}', using default {}", key, v, def); + return def; +} + +// Split a comma-separated list, dropping empty entries. +std::vector SplitCsv(const std::string& s) { + std::vector out; + size_t pos = 0; + while (pos <= s.size()) { + const size_t comma = s.find(',', pos); + const size_t end = (comma == std::string::npos) ? s.size() : comma; + std::string tok = s.substr(pos, end - pos); + if (!tok.empty()) out.push_back(tok); + if (comma == std::string::npos) break; + pos = comma + 1; + } + return out; +} + +#ifdef USE_REDIS_BACKEND +// Best-effort probe: does the single-endpoint server report itself as Dragonfly? +// (`INFO server` carries "dragonfly_version" on Dragonfly; Redis/Valkey do not.) +// Only used to nudge operators to shard the block keyspace so Dragonfly's +// proactor threads are actually used; any error just returns false (the real +// readiness probe below surfaces an outage). +bool ServerLooksLikeDragonfly(const std::string& uri, const std::string& password, + int connect_timeout_ms, int socket_timeout_ms) { + try { + redis::RespClient::Options o; + o.uri = uri; + o.password = password; + o.connect_timeout_ms = connect_timeout_ms; + o.socket_timeout_ms = socket_timeout_ms; + o.pool_size = 1; + redis::RespClient client(std::move(o)); + const redis::RespValue r = client.Command({"INFO", "server"}); + return r.str.find("dragonfly") != std::string::npos; + } catch (const std::exception&) { + return false; + } +} +#endif + +} // namespace + +std::unique_ptr MakeMasterMetadataStore() { + const std::string backend = GetEnvStr("UMBP_METADATA_BACKEND", "inmemory"); + + if (backend == "redis") { +#ifdef USE_REDIS_BACKEND + RedisMasterMetadataStore::Config cfg; + cfg.uri = GetEnvStr("UMBP_REDIS_URI", "tcp://127.0.0.1:6379"); + cfg.namespace_id = GetEnvStr("UMBP_REDIS_NAMESPACE", "default"); + cfg.password = GetEnvStr("UMBP_REDIS_PASSWORD", ""); + cfg.connect_timeout_ms = GetEnvInt("UMBP_REDIS_CONNECT_TIMEOUT_MS", 1000); + cfg.socket_timeout_ms = GetEnvInt("UMBP_REDIS_SOCKET_TIMEOUT_MS", 1000); + // Cap the default pool: a single-threaded Redis serializes every command, + // so a pool sized to a big host's CPU count (e.g. 448 on 224 cores) only + // deepens the queue and pushes tail latency past the socket timeout. 32 is + // a healthy default; override with UMBP_REDIS_POOL_SIZE for sharded/cluster + // deployments. + unsigned hw = std::thread::hardware_concurrency(); + const int default_pool = static_cast(std::min(32u, std::max(4u, hw ? hw * 2u : 8u))); + cfg.pool_size = static_cast(GetEnvInt("UMBP_REDIS_POOL_SIZE", default_pool)); + // Redis Cluster mode (UMBP_REDIS_CLUSTER=1): one redis-plus-plus client + // routes by hash-tag slot with MOVED/ASK + master-failover handled for us. + // Read early because it changes the block-shard default. + const bool cluster = GetEnvBool("UMBP_REDIS_CLUSTER", false); + + // Block-keyspace shards. Default 1 = legacy single-tag layout (no change) for + // single/multi. In cluster mode, if left unset it is auto-sized to ~2x the + // discovered master count (below) so a RouteGet batch spreads across nodes + // without over-splitting into too many per-shard scripts; 16 is only the + // fallback if discovery fails. Override explicitly with UMBP_REDIS_BLOCK_SHARDS. + const char* bs_env = std::getenv("UMBP_REDIS_BLOCK_SHARDS"); + const bool bs_explicit = (bs_env != nullptr && *bs_env != '\0'); + constexpr int kMaxBlockShards = 4096; + int block_shards = GetEnvInt("UMBP_REDIS_BLOCK_SHARDS", cluster ? 16 : 1); + if (block_shards > kMaxBlockShards) { + MORI_UMBP_WARN("[MetadataStore] clamping UMBP_REDIS_BLOCK_SHARDS={} to max {}", block_shards, + kMaxBlockShards); + block_shards = kMaxBlockShards; + } + cfg.block_shards = static_cast(block_shards); + // Multi-endpoint mode: comma-separated Redis URIs, one instance per block + // shard, so their scripts run on independent server threads (the way past a + // single instance's single-thread ceiling). When set, it supersedes the + // single-endpoint block_shards knob. + cfg.shard_uris = SplitCsv(GetEnvStr("UMBP_REDIS_SHARD_URIS", "")); + + // single-endpoint / multi-endpoint / cluster are mutually exclusive. + if (cluster && cfg.shard_uris.size() > 1) { + throw std::runtime_error( + "UMBP_REDIS_CLUSTER=1 and UMBP_REDIS_SHARD_URIS are mutually exclusive; set only one"); + } + + if (cluster) { + cfg.cluster = true; + // Seeds: the UMBP_REDIS_URI comma list (any reachable node bootstraps the + // rest via CLUSTER SLOTS). + cfg.cluster_seeds = SplitCsv(GetEnvStr("UMBP_REDIS_URI", "tcp://127.0.0.1:6379")); + // Balanced placement (unless the operator set UMBP_REDIS_BLOCK_SHARDS): + // read CLUSTER SLOTS and put exactly one block-shard tag on each master + // (a tag whose CRC16 slot that node owns), so a RouteGet batch spreads + // evenly — one EVALSHA per node — instead of piling onto whichever node the + // formulaic tags happen to hash to. Measured to lift cluster throughput + // from ~single-node to ~85% of multi-endpoint. Falls back to formulaic + // shards if discovery / tag search fails (readiness probe surfaces a real + // outage). + if (!bs_explicit) { + try { + redis::RespClusterClient::Options opts; + opts.seeds = cfg.cluster_seeds; + opts.password = cfg.password; + opts.connect_timeout_ms = cfg.connect_timeout_ms; + opts.socket_timeout_ms = cfg.socket_timeout_ms; + opts.pool_size = cfg.pool_size; + const auto ranges = redis::RespClusterClient::DiscoverMasterSlotRanges(opts); + std::vector tags; + tags.reserve(ranges.size()); + for (const auto& node_ranges : ranges) { + std::string tag; + if (redis::FindTagForRanges(cfg.namespace_id, node_ranges, &tag)) tags.push_back(tag); + } + if (!tags.empty() && tags.size() == ranges.size()) { + cfg.cluster_block_tags = tags; + cfg.block_shards = tags.size(); // one balanced tag per master + MORI_UMBP_INFO("[MetadataStore] cluster balanced placement: {} tags, one per master", + tags.size()); + } else { + const std::size_t fallback = std::min( + kMaxBlockShards, std::max(1, 2 * std::max(1, ranges.size()))); + cfg.block_shards = fallback; + MORI_UMBP_WARN( + "[MetadataStore] cluster balanced tag search matched {}/{} masters; " + "using formulaic block_shards={}", + tags.size(), ranges.size(), fallback); + } + } catch (const std::exception& e) { + MORI_UMBP_WARN( + "[MetadataStore] cluster topology discovery failed ({}); using block_shards={}", + e.what(), cfg.block_shards); + } + } + MORI_UMBP_INFO("[MetadataStore] backend=redis namespace={} seeds={} block_shards={} (cluster)", + cfg.namespace_id, cfg.cluster_seeds.size(), cfg.block_shards); + } else if (cfg.shard_uris.size() > 1) { + MORI_UMBP_INFO("[MetadataStore] backend=redis namespace={} endpoints={} (multi-endpoint)", + cfg.namespace_id, cfg.shard_uris.size()); + } else { + // Single-endpoint. Dragonfly is multi-threaded, but with the default + // block_shards=1 every block key shares one hash tag => one proactor + // thread, so the RouteGet hot path does not scale. We do NOT silently bump + // the default (block_shards is fixed for a deployment's lifetime — changing + // it strands existing block keys), but nudge the operator to set it. + if (!bs_explicit && ServerLooksLikeDragonfly(cfg.uri, cfg.password, cfg.connect_timeout_ms, + cfg.socket_timeout_ms)) { + MORI_UMBP_WARN( + "[MetadataStore] Dragonfly at {} with default UMBP_REDIS_BLOCK_SHARDS=1: block reads " + "run on a single proactor thread. Set UMBP_REDIS_BLOCK_SHARDS to ~the Dragonfly " + "--proactor_threads count (e.g. 8) to parallelize the RouteGet hot path. NOTE: this " + "value is fixed for the deployment's lifetime — pick one and keep it (changing it " + "strands existing block keys).", + cfg.uri); + } + MORI_UMBP_INFO("[MetadataStore] backend=redis uri={} namespace={} block_shards={}", cfg.uri, + cfg.namespace_id, cfg.block_shards); + } + + auto store = std::make_unique(cfg); + + // Startup readiness probe. Pinging every configured endpoint turns a + // misconfigured / unreachable store into a clear, immediate failure instead + // of a master that starts "healthy" and then returns UNAVAILABLE on every + // RPC. Gated by UMBP_REDIS_REQUIRED (default true): set 0 to start in a + // degraded state (e.g. when some shards are expected to come up later) and + // rely on the runtime reconnect path. + const std::string where = cfg.shard_uris.size() > 1 + ? ("endpoints=" + std::to_string(cfg.shard_uris.size())) + : cfg.uri; + if (store->Ping()) { + MORI_UMBP_INFO("[MetadataStore] backend=redis readiness probe OK ({})", where); + } else if (GetEnvBool("UMBP_REDIS_REQUIRED", true)) { + throw std::runtime_error( + "backend=redis but the store is unreachable at " + where + + " (readiness probe failed). Fix the store/connection, or set UMBP_REDIS_REQUIRED=0 to " + "start in a degraded state."); + } else { + MORI_UMBP_WARN( + "[MetadataStore] backend=redis store unreachable at {} at startup; starting degraded " + "(UMBP_REDIS_REQUIRED=0) — RPCs return UNAVAILABLE until it recovers.", + where); + } + return store; +#else + throw std::runtime_error( + "UMBP_METADATA_BACKEND=redis but the Redis backend was not compiled in; " + "rebuild with -DUSE_REDIS_BACKEND=ON"); +#endif + } + + if (backend != "inmemory") { + MORI_UMBP_WARN("[MetadataStore] unknown UMBP_METADATA_BACKEND='{}', using inmemory", backend); + } + MORI_UMBP_INFO("[MetadataStore] backend=inmemory"); + return std::make_unique(); +} + +} // namespace mori::umbp diff --git a/src/umbp/distributed/master/master_server.cpp b/src/umbp/distributed/master/master_server.cpp index 2d99ccc05..ee249bc62 100644 --- a/src/umbp/distributed/master/master_server.cpp +++ b/src/umbp/distributed/master/master_server.cpp @@ -37,6 +37,7 @@ #include "umbp/distributed/master/evict_strategy.h" #include "umbp/distributed/master/in_memory_master_metadata_store.h" #include "umbp/distributed/master/master_metadata_store.h" +#include "umbp/distributed/master/master_metadata_store_factory.h" #include "umbp/distributed/master/master_metrics.h" #include "umbp/distributed/routing/router.h" #include "umbp_peer.grpc.pb.h" @@ -203,322 +204,358 @@ class MasterServer::UMBPMasterServiceImpl final : public ::umbp::UMBPMaster::Ser const ClientRegistryConfig& config, mori::metrics::MetricsServer* metrics) : store_(store), router_(router), config_(config), metrics_(metrics) {} + // Run a handler body, converting any exception that escapes the metadata + // store into a clean grpc::UNAVAILABLE instead of letting it propagate out of + // the gRPC handler and terminate the whole master process. The RESP/Redis + // backend signals transport failures (and Phase-2-unimplemented methods) by + // throwing; the in-memory backend never throws, so this is a no-op cost + // there. This is the master's half of the "degrade, don't fabricate" + // contract in design-redis-metadata-store.md §7 — a struggling store drains + // this replica via failed RPCs rather than crashing it. + template + static grpc::Status GuardStore(const char* rpc, Fn&& fn) { + try { + return fn(); + } catch (const std::exception& e) { + MORI_UMBP_ERROR("[Master] {} metadata-store error: {}", rpc, e.what()); + return grpc::Status(grpc::StatusCode::UNAVAILABLE, + std::string(rpc) + ": metadata store unavailable: " + e.what()); + } + } + // -------- Client lifecycle -------- grpc::Status RegisterClient(grpc::ServerContext* /*ctx*/, const ::umbp::RegisterClientRequest* request, ::umbp::RegisterClientResponse* response) override { - std::map caps; - for (const auto& tc : request->tier_capacities()) { - TierCapacity c; - c.total_bytes = tc.total_capacity_bytes(); - c.available_bytes = tc.available_capacity_bytes(); - caps[static_cast(tc.tier())] = c; - } + return GuardStore("RegisterClient", [&]() -> grpc::Status { + std::map caps; + for (const auto& tc : request->tier_capacities()) { + TierCapacity c; + c.total_bytes = tc.total_capacity_bytes(); + c.available_bytes = tc.available_capacity_bytes(); + caps[static_cast(tc.tier())] = c; + } - const auto& engine_desc_str = request->engine_desc(); - - ClientRegistration registration; - registration.node_id = request->node_id(); - registration.node_address = request->node_address(); - registration.tier_capacities = caps; - registration.peer_address = request->peer_address(); - registration.engine_desc_bytes.assign(engine_desc_str.begin(), engine_desc_str.end()); - registration.tags.assign(request->tags().begin(), request->tags().end()); - - // stale_after mirrors ClientRegistry::ExpiryDuration() (heartbeat_ttl × - // max_missed_heartbeats) so a TTL-stale ALIVE record the reaper hasn't yet - // flipped can still be re-registered (hazard #2). - const auto stale_after = config_.heartbeat_ttl * config_.max_missed_heartbeats; - const bool registered = - store_.RegisterClient(registration, std::chrono::system_clock::now(), stale_after); - if (!registered) { - return grpc::Status(grpc::StatusCode::ALREADY_EXISTS, - "node is already alive and cannot be re-registered"); - } + const auto& engine_desc_str = request->engine_desc(); + + ClientRegistration registration; + registration.node_id = request->node_id(); + registration.node_address = request->node_address(); + registration.tier_capacities = caps; + registration.peer_address = request->peer_address(); + registration.engine_desc_bytes.assign(engine_desc_str.begin(), engine_desc_str.end()); + registration.tags.assign(request->tags().begin(), request->tags().end()); + + // stale_after mirrors ClientRegistry::ExpiryDuration() (heartbeat_ttl × + // max_missed_heartbeats) so a TTL-stale ALIVE record the reaper hasn't yet + // flipped can still be re-registered (hazard #2). + const auto stale_after = config_.heartbeat_ttl * config_.max_missed_heartbeats; + const bool registered = + store_.RegisterClient(registration, std::chrono::system_clock::now(), stale_after); + if (!registered) { + return grpc::Status(grpc::StatusCode::ALREADY_EXISTS, + "node is already alive and cannot be re-registered"); + } - UpdateClientCountMetric(); - UpdateClientCapacityMetrics(request->node_id(), caps); + UpdateClientCountMetric(); + UpdateClientCapacityMetrics(request->node_id(), caps); - auto interval_ms = - static_cast(config_.heartbeat_ttl.count() * 1000) / HeartbeatIntervalDivisor(); - response->set_heartbeat_interval_ms(interval_ms); - response->set_ack_seq(0); - return grpc::Status::OK; + auto interval_ms = + static_cast(config_.heartbeat_ttl.count() * 1000) / HeartbeatIntervalDivisor(); + response->set_heartbeat_interval_ms(interval_ms); + response->set_ack_seq(0); + return grpc::Status::OK; + }); } grpc::Status UnregisterClient(grpc::ServerContext* /*ctx*/, const ::umbp::UnregisterClientRequest* request, ::umbp::UnregisterClientResponse* /*response*/) override { - store_.UnregisterClient(request->node_id()); - UpdateClientCountMetric(); - return grpc::Status::OK; + return GuardStore("UnregisterClient", [&]() -> grpc::Status { + store_.UnregisterClient(request->node_id()); + UpdateClientCountMetric(); + return grpc::Status::OK; + }); } // -------- Heartbeat (event-driven) -------- grpc::Status Heartbeat(grpc::ServerContext* /*ctx*/, const ::umbp::HeartbeatRequest* request, ::umbp::HeartbeatResponse* response) override { - std::map caps; - for (const auto& tc : request->tier_capacities()) { - TierCapacity c; - c.total_bytes = tc.total_capacity_bytes(); - c.available_bytes = tc.available_capacity_bytes(); - caps[static_cast(tc.tier())] = c; - } + return GuardStore("Heartbeat", [&]() -> grpc::Status { + std::map caps; + for (const auto& tc : request->tier_capacities()) { + TierCapacity c; + c.total_bytes = tc.total_capacity_bytes(); + c.available_bytes = tc.available_capacity_bytes(); + caps[static_cast(tc.tier())] = c; + } - std::vector bundles; - bundles.reserve(static_cast(request->bundles_size())); - for (const auto& pb : request->bundles()) { - EventBundle bundle; - bundle.seq = pb.seq(); - bundle.events.reserve(static_cast(pb.events_size())); - for (const auto& pe : pb.events()) bundle.events.push_back(FromProtoEvent(pe)); - bundles.push_back(std::move(bundle)); - } + std::vector bundles; + bundles.reserve(static_cast(request->bundles_size())); + for (const auto& pb : request->bundles()) { + EventBundle bundle; + bundle.seq = pb.seq(); + bundle.events.reserve(static_cast(pb.events_size())); + for (const auto& pe : pb.events()) bundle.events.push_back(FromProtoEvent(pe)); + bundles.push_back(std::move(bundle)); + } - // §3e heartbeat adapter: the wire protocol ships EventBundle[] (each with - // its own seq) but IMasterMetadataStore::ApplyHeartbeat takes one seq at a - // time. Translate here — the store sees one seq per call, which is exactly - // what the seq-CAS (hazard #1) requires. - const auto now = std::chrono::system_clock::now(); - uint64_t acked_seq = 0; - bool request_full_sync = false; - ClientStatus client_status = ClientStatus::ALIVE; - - if (request->is_full_sync()) { - // Full sync replaces this node's locations wholesale and re-baselines - // last_applied_seq to delta_seq_baseline. Flatten every bundle's events - // into one ApplyHeartbeat call (ReplaceNodeLocations keeps only ADDs). - std::vector events; - for (const auto& bundle : bundles) { - for (const auto& ev : bundle.events) events.push_back(ev); - } - auto result = store_.ApplyHeartbeat(request->node_id(), request->delta_seq_baseline(), now, - caps, events, /*is_full_sync=*/true); - if (result.status == HeartbeatResult::UNKNOWN) { - client_status = ClientStatus::UNKNOWN; - } else { - acked_seq = result.acked_seq; - } - } else if (bundles.empty()) { - // Keepalive heartbeat: no events, but liveness must still refresh so the - // reaper doesn't expire an idle-but-alive node. seq=0 deterministically - // hits the SEQ_GAP branch, which bumps last_heartbeat + status←ALIVE - // without advancing last_applied_seq; the gap is ignored (no full-sync - // request) because there is nothing to recover. - auto result = store_.ApplyHeartbeat(request->node_id(), /*seq=*/0, now, caps, - /*events=*/{}, /*is_full_sync=*/false); - if (result.status == HeartbeatResult::UNKNOWN) { - client_status = ClientStatus::UNKNOWN; - } else { - acked_seq = result.acked_seq; - } - } else { - // Delta path: apply bundles in ascending-seq order. A real forward gap - // short-circuits the loop and requests a full sync, leaving earlier - // bundles applied (Risk item 2 — application is per-bundle, not atomic - // across the batch). - for (const auto& bundle : bundles) { - auto result = store_.ApplyHeartbeat(request->node_id(), bundle.seq, now, caps, - bundle.events, /*is_full_sync=*/false); + // §3e heartbeat adapter: the wire protocol ships EventBundle[] (each with + // its own seq) but IMasterMetadataStore::ApplyHeartbeat takes one seq at a + // time. Translate here — the store sees one seq per call, which is exactly + // what the seq-CAS (hazard #1) requires. + const auto now = std::chrono::system_clock::now(); + uint64_t acked_seq = 0; + bool request_full_sync = false; + ClientStatus client_status = ClientStatus::ALIVE; + + if (request->is_full_sync()) { + // Full sync replaces this node's locations wholesale and re-baselines + // last_applied_seq to delta_seq_baseline. Flatten every bundle's events + // into one ApplyHeartbeat call (ReplaceNodeLocations keeps only ADDs). + std::vector events; + for (const auto& bundle : bundles) { + for (const auto& ev : bundle.events) events.push_back(ev); + } + auto result = store_.ApplyHeartbeat(request->node_id(), request->delta_seq_baseline(), now, + caps, events, /*is_full_sync=*/true); + if (result.status == HeartbeatResult::UNKNOWN) { + client_status = ClientStatus::UNKNOWN; + } else { + acked_seq = result.acked_seq; + } + } else if (bundles.empty()) { + // Keepalive heartbeat: no events, but liveness must still refresh so the + // reaper doesn't expire an idle-but-alive node. seq=0 deterministically + // hits the SEQ_GAP branch, which bumps last_heartbeat + status←ALIVE + // without advancing last_applied_seq; the gap is ignored (no full-sync + // request) because there is nothing to recover. + auto result = store_.ApplyHeartbeat(request->node_id(), /*seq=*/0, now, caps, + /*events=*/{}, /*is_full_sync=*/false); if (result.status == HeartbeatResult::UNKNOWN) { client_status = ClientStatus::UNKNOWN; - break; + } else { + acked_seq = result.acked_seq; } - if (result.status == HeartbeatResult::SEQ_GAP) { + } else { + // Delta path: apply bundles in ascending-seq order. A real forward gap + // short-circuits the loop and requests a full sync, leaving earlier + // bundles applied (Risk item 2 — application is per-bundle, not atomic + // across the batch). + for (const auto& bundle : bundles) { + auto result = store_.ApplyHeartbeat(request->node_id(), bundle.seq, now, caps, + bundle.events, /*is_full_sync=*/false); + if (result.status == HeartbeatResult::UNKNOWN) { + client_status = ClientStatus::UNKNOWN; + break; + } + if (result.status == HeartbeatResult::SEQ_GAP) { + acked_seq = result.acked_seq; + request_full_sync = true; + break; + } acked_seq = result.acked_seq; - request_full_sync = true; - break; } - acked_seq = result.acked_seq; } - } - response->set_status(static_cast<::umbp::ClientStatus>(client_status)); - response->set_acked_seq(acked_seq); - response->set_request_full_sync(request_full_sync); + response->set_status(static_cast<::umbp::ClientStatus>(client_status)); + response->set_acked_seq(acked_seq); + response->set_request_full_sync(request_full_sync); - UpdateClientCapacityMetrics(request->node_id(), caps); + UpdateClientCapacityMetrics(request->node_id(), caps); - if (metrics_ != nullptr && request->tier_kv_counts_size() > 0) { - mori::metrics::MetricsServer::Labels base = {{"node", request->node_id()}}; - for (const auto& tag : store_.GetClientTags(request->node_id())) { - const auto sep = tag.find('='); - if (sep != std::string::npos) { - base.push_back({tag.substr(0, sep), tag.substr(sep + 1)}); + if (metrics_ != nullptr && request->tier_kv_counts_size() > 0) { + mori::metrics::MetricsServer::Labels base = {{"node", request->node_id()}}; + for (const auto& tag : store_.GetClientTags(request->node_id())) { + const auto sep = tag.find('='); + if (sep != std::string::npos) { + base.push_back({tag.substr(0, sep), tag.substr(sep + 1)}); + } + } + uint64_t total = 0; + for (const auto& tkc : request->tier_kv_counts()) { + total += tkc.count(); + auto labels = base; + labels.push_back({"tier", TierTypeName(static_cast(tkc.tier()))}); + metrics_->setGauge(MORI_UMBP_METRIC_CLIENT_KV_LIVE_COUNT, + MORI_UMBP_METRIC_CLIENT_KV_LIVE_COUNT_HELP, labels, + static_cast(tkc.count())); } + metrics_->setGauge(MORI_UMBP_METRIC_CLIENT_KV_LIVE_COUNT_TOTAL, + MORI_UMBP_METRIC_CLIENT_KV_LIVE_COUNT_TOTAL_HELP, base, + static_cast(total)); } - uint64_t total = 0; - for (const auto& tkc : request->tier_kv_counts()) { - total += tkc.count(); - auto labels = base; - labels.push_back({"tier", TierTypeName(static_cast(tkc.tier()))}); - metrics_->setGauge(MORI_UMBP_METRIC_CLIENT_KV_LIVE_COUNT, - MORI_UMBP_METRIC_CLIENT_KV_LIVE_COUNT_HELP, labels, - static_cast(tkc.count())); - } - metrics_->setGauge(MORI_UMBP_METRIC_CLIENT_KV_LIVE_COUNT_TOTAL, - MORI_UMBP_METRIC_CLIENT_KV_LIVE_COUNT_TOTAL_HELP, base, - static_cast(total)); - } - if (request_full_sync && metrics_ != nullptr) { - metrics_->addCounter("mori_umbp_heartbeat_seq_gap_total", - "Heartbeats rejected due to seq gap (full sync requested)", - {{"node", request->node_id()}}); - } - size_t event_count = 0; - for (const auto& bundle : bundles) event_count += bundle.events.size(); - if (metrics_ != nullptr && event_count > 0) { - metrics_->addCounter("mori_umbp_heartbeat_events_applied_total", - "KvEvents applied to GlobalBlockIndex via heartbeat", - {{"node", request->node_id()}}, static_cast(event_count)); - } - return grpc::Status::OK; + if (request_full_sync && metrics_ != nullptr) { + metrics_->addCounter("mori_umbp_heartbeat_seq_gap_total", + "Heartbeats rejected due to seq gap (full sync requested)", + {{"node", request->node_id()}}); + } + size_t event_count = 0; + for (const auto& bundle : bundles) event_count += bundle.events.size(); + if (metrics_ != nullptr && event_count > 0) { + metrics_->addCounter("mori_umbp_heartbeat_events_applied_total", + "KvEvents applied to GlobalBlockIndex via heartbeat", + {{"node", request->node_id()}}, static_cast(event_count)); + } + return grpc::Status::OK; + }); } // -------- Routing (read-only) -------- grpc::Status RoutePut(grpc::ServerContext* /*ctx*/, const ::umbp::RoutePutRequest* request, ::umbp::RoutePutResponse* response) override { - if (request->key().empty()) { - return grpc::Status(grpc::StatusCode::INVALID_ARGUMENT, "key cannot be empty"); - } - std::unordered_set excludes(request->exclude_nodes().begin(), - request->exclude_nodes().end()); - auto result = - router_.RoutePut(request->key(), request->node_id(), request->block_size(), excludes); - if (!result.has_value()) { - response->set_outcome(::umbp::ROUTE_PUT_OUTCOME_UNAVAILABLE); - return grpc::Status::OK; - } - if (result->outcome == RoutePutOutcome::kAlreadyExists) { - response->set_outcome(::umbp::ROUTE_PUT_OUTCOME_ALREADY_EXISTS); + return GuardStore("RoutePut", [&]() -> grpc::Status { + if (request->key().empty()) { + return grpc::Status(grpc::StatusCode::INVALID_ARGUMENT, "key cannot be empty"); + } + std::unordered_set excludes(request->exclude_nodes().begin(), + request->exclude_nodes().end()); + auto result = + router_.RoutePut(request->key(), request->node_id(), request->block_size(), excludes); + if (!result.has_value()) { + response->set_outcome(::umbp::ROUTE_PUT_OUTCOME_UNAVAILABLE); + return grpc::Status::OK; + } + if (result->outcome == RoutePutOutcome::kAlreadyExists) { + response->set_outcome(::umbp::ROUTE_PUT_OUTCOME_ALREADY_EXISTS); + return grpc::Status::OK; + } + response->set_outcome(::umbp::ROUTE_PUT_OUTCOME_ROUTED); + response->set_node_id(result->node_id); + response->set_tier(static_cast<::umbp::TierType>(result->tier)); + response->set_peer_address(result->peer_address); + + if (metrics_) { + metrics_->addCounter(MORI_UMBP_METRIC_CLIENT_ROUTE_PUT, + MORI_UMBP_METRIC_CLIENT_ROUTE_PUT_HELP, {{"node", result->node_id}}); + } return grpc::Status::OK; - } - response->set_outcome(::umbp::ROUTE_PUT_OUTCOME_ROUTED); - response->set_node_id(result->node_id); - response->set_tier(static_cast<::umbp::TierType>(result->tier)); - response->set_peer_address(result->peer_address); - - if (metrics_) { - metrics_->addCounter(MORI_UMBP_METRIC_CLIENT_ROUTE_PUT, - MORI_UMBP_METRIC_CLIENT_ROUTE_PUT_HELP, {{"node", result->node_id}}); - } - return grpc::Status::OK; + }); } grpc::Status RouteGet(grpc::ServerContext* /*ctx*/, const ::umbp::RouteGetRequest* request, ::umbp::RouteGetResponse* response) override { - if (request->key().empty()) { - return grpc::Status(grpc::StatusCode::INVALID_ARGUMENT, "key cannot be empty"); - } - std::unordered_set excludes(request->exclude_nodes().begin(), - request->exclude_nodes().end()); - auto result = router_.RouteGet(request->key(), request->node_id(), excludes); - if (!result.has_value()) { - response->set_found(false); + return GuardStore("RouteGet", [&]() -> grpc::Status { + if (request->key().empty()) { + return grpc::Status(grpc::StatusCode::INVALID_ARGUMENT, "key cannot be empty"); + } + std::unordered_set excludes(request->exclude_nodes().begin(), + request->exclude_nodes().end()); + auto result = router_.RouteGet(request->key(), request->node_id(), excludes); + if (!result.has_value()) { + response->set_found(false); + return grpc::Status::OK; + } + response->set_found(true); + response->set_node_id(result->location.node_id); + response->set_tier(static_cast<::umbp::TierType>(result->location.tier)); + response->set_size(result->location.size); + response->set_peer_address(result->peer_address); + + if (metrics_) { + metrics_->addCounter(MORI_UMBP_METRIC_CLIENT_ROUTE_GET, + MORI_UMBP_METRIC_CLIENT_ROUTE_GET_HELP, + {{"node", result->location.node_id}}); + } return grpc::Status::OK; - } - response->set_found(true); - response->set_node_id(result->location.node_id); - response->set_tier(static_cast<::umbp::TierType>(result->location.tier)); - response->set_size(result->location.size); - response->set_peer_address(result->peer_address); - - if (metrics_) { - metrics_->addCounter(MORI_UMBP_METRIC_CLIENT_ROUTE_GET, - MORI_UMBP_METRIC_CLIENT_ROUTE_GET_HELP, - {{"node", result->location.node_id}}); - } - return grpc::Status::OK; + }); } grpc::Status BatchRoutePut(grpc::ServerContext* /*ctx*/, const ::umbp::BatchRoutePutRequest* request, ::umbp::BatchRoutePutResponse* response) override { - if (request->keys_size() != request->block_sizes_size()) { - return grpc::Status(grpc::StatusCode::INVALID_ARGUMENT, - "keys and block_sizes must have the same length"); - } - std::vector keys(request->keys().begin(), request->keys().end()); - std::vector block_sizes(request->block_sizes().begin(), request->block_sizes().end()); - std::unordered_set excludes(request->exclude_nodes().begin(), - request->exclude_nodes().end()); - - auto results = router_.BatchRoutePut(keys, request->node_id(), block_sizes, excludes); - for (auto& opt : results) { - auto* entry = response->add_entries(); - if (!opt.has_value()) continue; // default UNAVAILABLE - if (opt->outcome == RoutePutOutcome::kAlreadyExists) { - entry->set_outcome(::umbp::ROUTE_PUT_OUTCOME_ALREADY_EXISTS); - continue; - } - entry->set_outcome(::umbp::ROUTE_PUT_OUTCOME_ROUTED); - entry->set_node_id(opt->node_id); - entry->set_tier(static_cast<::umbp::TierType>(opt->tier)); - entry->set_peer_address(opt->peer_address); - if (metrics_) { - metrics_->addCounter(MORI_UMBP_METRIC_CLIENT_BATCH_ROUTE_PUT, - MORI_UMBP_METRIC_CLIENT_BATCH_ROUTE_PUT_HELP, - {{"node", opt->node_id}}); + return GuardStore("BatchRoutePut", [&]() -> grpc::Status { + if (request->keys_size() != request->block_sizes_size()) { + return grpc::Status(grpc::StatusCode::INVALID_ARGUMENT, + "keys and block_sizes must have the same length"); } - } - return grpc::Status::OK; + std::vector keys(request->keys().begin(), request->keys().end()); + std::vector block_sizes(request->block_sizes().begin(), + request->block_sizes().end()); + std::unordered_set excludes(request->exclude_nodes().begin(), + request->exclude_nodes().end()); + + auto results = router_.BatchRoutePut(keys, request->node_id(), block_sizes, excludes); + for (auto& opt : results) { + auto* entry = response->add_entries(); + if (!opt.has_value()) continue; // default UNAVAILABLE + if (opt->outcome == RoutePutOutcome::kAlreadyExists) { + entry->set_outcome(::umbp::ROUTE_PUT_OUTCOME_ALREADY_EXISTS); + continue; + } + entry->set_outcome(::umbp::ROUTE_PUT_OUTCOME_ROUTED); + entry->set_node_id(opt->node_id); + entry->set_tier(static_cast<::umbp::TierType>(opt->tier)); + entry->set_peer_address(opt->peer_address); + if (metrics_) { + metrics_->addCounter(MORI_UMBP_METRIC_CLIENT_BATCH_ROUTE_PUT, + MORI_UMBP_METRIC_CLIENT_BATCH_ROUTE_PUT_HELP, + {{"node", opt->node_id}}); + } + } + return grpc::Status::OK; + }); } grpc::Status BatchRouteGet(grpc::ServerContext* /*ctx*/, const ::umbp::BatchRouteGetRequest* request, ::umbp::BatchRouteGetResponse* response) override { - std::vector keys(request->keys().begin(), request->keys().end()); - std::unordered_set excludes(request->exclude_nodes().begin(), - request->exclude_nodes().end()); - auto results = router_.BatchRouteGet(keys, request->node_id(), excludes); - // Columnar response: distinct (node_id, peer_address) pairs are emitted - // once into `nodes`; each key carries a 1-based node_ref index (0 = not - // found). node_ref/tier/size are parallel arrays aligned with the request - // keys, so the per-key fields default to 0 for unresolved keys. - response->mutable_node_ref()->Reserve(static_cast(results.size())); - response->mutable_tier()->Reserve(static_cast(results.size())); - response->mutable_size()->Reserve(static_cast(results.size())); - // Maps "node_id\0peer_address" -> 1-based index into response->nodes(). - std::unordered_map node_index; - for (auto& opt : results) { - if (!opt.has_value()) { - response->add_node_ref(0); - response->add_tier(::umbp::TIER_UNKNOWN); - response->add_size(0); - continue; - } - std::string node_key = opt->location.node_id; - node_key.push_back('\0'); - node_key.append(opt->peer_address); - auto [it, inserted] = node_index.try_emplace(node_key, 0); - if (inserted) { - auto* node = response->add_nodes(); - node->set_node_id(opt->location.node_id); - node->set_peer_address(opt->peer_address); - it->second = static_cast(response->nodes_size()); // 1-based - } - response->add_node_ref(it->second); - response->add_tier(static_cast<::umbp::TierType>(opt->location.tier)); - response->add_size(opt->location.size); - if (metrics_) { - metrics_->addCounter(MORI_UMBP_METRIC_CLIENT_BATCH_ROUTE_GET, - MORI_UMBP_METRIC_CLIENT_BATCH_ROUTE_GET_HELP, - {{"node", opt->location.node_id}}); + return GuardStore("BatchRouteGet", [&]() -> grpc::Status { + std::vector keys(request->keys().begin(), request->keys().end()); + std::unordered_set excludes(request->exclude_nodes().begin(), + request->exclude_nodes().end()); + auto results = router_.BatchRouteGet(keys, request->node_id(), excludes); + // Columnar response: distinct (node_id, peer_address) pairs are emitted + // once into `nodes`; each key carries a 1-based node_ref index (0 = not + // found). node_ref/tier/size are parallel arrays aligned with the request + // keys, so the per-key fields default to 0 for unresolved keys. + response->mutable_node_ref()->Reserve(static_cast(results.size())); + response->mutable_tier()->Reserve(static_cast(results.size())); + response->mutable_size()->Reserve(static_cast(results.size())); + // Maps "node_id\0peer_address" -> 1-based index into response->nodes(). + std::unordered_map node_index; + for (auto& opt : results) { + if (!opt.has_value()) { + response->add_node_ref(0); + response->add_tier(::umbp::TIER_UNKNOWN); + response->add_size(0); + continue; + } + std::string node_key = opt->location.node_id; + node_key.push_back('\0'); + node_key.append(opt->peer_address); + auto [it, inserted] = node_index.try_emplace(node_key, 0); + if (inserted) { + auto* node = response->add_nodes(); + node->set_node_id(opt->location.node_id); + node->set_peer_address(opt->peer_address); + it->second = static_cast(response->nodes_size()); // 1-based + } + response->add_node_ref(it->second); + response->add_tier(static_cast<::umbp::TierType>(opt->location.tier)); + response->add_size(opt->location.size); + if (metrics_) { + metrics_->addCounter(MORI_UMBP_METRIC_CLIENT_BATCH_ROUTE_GET, + MORI_UMBP_METRIC_CLIENT_BATCH_ROUTE_GET_HELP, + {{"node", opt->location.node_id}}); + } } - } - return grpc::Status::OK; + return grpc::Status::OK; + }); } grpc::Status BatchLookup(grpc::ServerContext* /*ctx*/, const ::umbp::BatchLookupRequest* request, ::umbp::BatchLookupResponse* response) override { - std::vector keys(request->keys().begin(), request->keys().end()); - auto found = store_.BatchExistsBlock(keys); - for (bool b : found) response->add_found(b); - return grpc::Status::OK; + return GuardStore("BatchLookup", [&]() -> grpc::Status { + std::vector keys(request->keys().begin(), request->keys().end()); + auto found = store_.BatchExistsBlock(keys); + for (bool b : found) response->add_found(b); + return grpc::Status::OK; + }); } // -------- External KV mutation/query -------- @@ -526,204 +563,218 @@ class MasterServer::UMBPMasterServiceImpl final : public ::umbp::UMBPMaster::Ser grpc::Status ReportExternalKvBlocks( grpc::ServerContext* /*ctx*/, const ::umbp::ReportExternalKvBlocksRequest* request, ::umbp::ReportExternalKvBlocksResponse* /*response*/) override { - if (request->node_id().empty()) { - return grpc::Status(grpc::StatusCode::INVALID_ARGUMENT, "node_id must not be empty"); - } - if (request->hashes_size() == 0) { - return grpc::Status(grpc::StatusCode::INVALID_ARGUMENT, "hashes must not be empty"); - } + return GuardStore("ReportExternalKvBlocks", [&]() -> grpc::Status { + if (request->node_id().empty()) { + return grpc::Status(grpc::StatusCode::INVALID_ARGUMENT, "node_id must not be empty"); + } + if (request->hashes_size() == 0) { + return grpc::Status(grpc::StatusCode::INVALID_ARGUMENT, "hashes must not be empty"); + } + + const TierType tier = static_cast(request->tier()); + std::vector hashes(request->hashes().begin(), request->hashes().end()); + // Alive-check + write are fused into one atomic store call (closes the + // TOCTOU gap between the old IsClientAlive check and Register — §2b.5). + const bool applied = store_.RegisterExternalKvIfAlive(request->node_id(), hashes, tier); + if (!applied) { + MORI_UMBP_WARN("[Server] ReportExternalKvBlocks rejected: node not alive: {}", + request->node_id()); + if (metrics_) { + metrics_->addCounter(MORI_UMBP_METRIC_EXT_KV_REPORT_TOTAL, + MORI_UMBP_METRIC_EXT_KV_REPORT_TOTAL_HELP, + {{"node", request->node_id()}, + {"tier", TierTypeName(tier)}, + {"result", "rejected_not_alive"}}); + } + return grpc::Status::OK; + } - const TierType tier = static_cast(request->tier()); - std::vector hashes(request->hashes().begin(), request->hashes().end()); - // Alive-check + write are fused into one atomic store call (closes the - // TOCTOU gap between the old IsClientAlive check and Register — §2b.5). - const bool applied = store_.RegisterExternalKvIfAlive(request->node_id(), hashes, tier); - if (!applied) { - MORI_UMBP_WARN("[Server] ReportExternalKvBlocks rejected: node not alive: {}", - request->node_id()); if (metrics_) { - metrics_->addCounter(MORI_UMBP_METRIC_EXT_KV_REPORT_TOTAL, - MORI_UMBP_METRIC_EXT_KV_REPORT_TOTAL_HELP, - {{"node", request->node_id()}, - {"tier", TierTypeName(tier)}, - {"result", "rejected_not_alive"}}); + const mori::metrics::MetricsServer::Labels labels = {{"node", request->node_id()}, + {"tier", TierTypeName(tier)}}; + metrics_->addCounter(MORI_UMBP_METRIC_EXT_KV_REPORT_BLOCKS_TOTAL, + MORI_UMBP_METRIC_EXT_KV_REPORT_BLOCKS_TOTAL_HELP, labels, + static_cast(hashes.size())); + metrics_->addCounter( + MORI_UMBP_METRIC_EXT_KV_REPORT_TOTAL, MORI_UMBP_METRIC_EXT_KV_REPORT_TOTAL_HELP, + {{"node", request->node_id()}, {"tier", TierTypeName(tier)}, {"result", "ok"}}); } return grpc::Status::OK; - } - - if (metrics_) { - const mori::metrics::MetricsServer::Labels labels = {{"node", request->node_id()}, - {"tier", TierTypeName(tier)}}; - metrics_->addCounter(MORI_UMBP_METRIC_EXT_KV_REPORT_BLOCKS_TOTAL, - MORI_UMBP_METRIC_EXT_KV_REPORT_BLOCKS_TOTAL_HELP, labels, - static_cast(hashes.size())); - metrics_->addCounter( - MORI_UMBP_METRIC_EXT_KV_REPORT_TOTAL, MORI_UMBP_METRIC_EXT_KV_REPORT_TOTAL_HELP, - {{"node", request->node_id()}, {"tier", TierTypeName(tier)}, {"result", "ok"}}); - } - return grpc::Status::OK; + }); } grpc::Status RevokeExternalKvBlocks( grpc::ServerContext* /*ctx*/, const ::umbp::RevokeExternalKvBlocksRequest* request, ::umbp::RevokeExternalKvBlocksResponse* /*response*/) override { - if (request->node_id().empty()) { - return grpc::Status(grpc::StatusCode::INVALID_ARGUMENT, "node_id must not be empty"); - } - if (request->hashes_size() == 0) { - return grpc::Status(grpc::StatusCode::INVALID_ARGUMENT, "hashes must not be empty"); - } + return GuardStore("RevokeExternalKvBlocks", [&]() -> grpc::Status { + if (request->node_id().empty()) { + return grpc::Status(grpc::StatusCode::INVALID_ARGUMENT, "node_id must not be empty"); + } + if (request->hashes_size() == 0) { + return grpc::Status(grpc::StatusCode::INVALID_ARGUMENT, "hashes must not be empty"); + } - const TierType tier = static_cast(request->tier()); - std::vector hashes(request->hashes().begin(), request->hashes().end()); - store_.UnregisterExternalKv(request->node_id(), hashes, tier); - if (metrics_) { - const mori::metrics::MetricsServer::Labels labels = {{"node", request->node_id()}, - {"tier", TierTypeName(tier)}}; - metrics_->addCounter(MORI_UMBP_METRIC_EXT_KV_REVOKE_BLOCKS_TOTAL, - MORI_UMBP_METRIC_EXT_KV_REVOKE_BLOCKS_TOTAL_HELP, labels, - static_cast(hashes.size())); - metrics_->addCounter( - MORI_UMBP_METRIC_EXT_KV_REVOKE_TOTAL, MORI_UMBP_METRIC_EXT_KV_REVOKE_TOTAL_HELP, - {{"node", request->node_id()}, {"tier", TierTypeName(tier)}, {"result", "ok"}}); - } - return grpc::Status::OK; + const TierType tier = static_cast(request->tier()); + std::vector hashes(request->hashes().begin(), request->hashes().end()); + store_.UnregisterExternalKv(request->node_id(), hashes, tier); + if (metrics_) { + const mori::metrics::MetricsServer::Labels labels = {{"node", request->node_id()}, + {"tier", TierTypeName(tier)}}; + metrics_->addCounter(MORI_UMBP_METRIC_EXT_KV_REVOKE_BLOCKS_TOTAL, + MORI_UMBP_METRIC_EXT_KV_REVOKE_BLOCKS_TOTAL_HELP, labels, + static_cast(hashes.size())); + metrics_->addCounter( + MORI_UMBP_METRIC_EXT_KV_REVOKE_TOTAL, MORI_UMBP_METRIC_EXT_KV_REVOKE_TOTAL_HELP, + {{"node", request->node_id()}, {"tier", TierTypeName(tier)}, {"result", "ok"}}); + } + return grpc::Status::OK; + }); } grpc::Status RevokeAllExternalKvBlocksAtTier( grpc::ServerContext* /*ctx*/, const ::umbp::RevokeAllExternalKvBlocksAtTierRequest* request, ::umbp::RevokeAllExternalKvBlocksAtTierResponse* /*response*/) override { - if (request->node_id().empty()) { - return grpc::Status(grpc::StatusCode::INVALID_ARGUMENT, "node_id must not be empty"); - } + return GuardStore("RevokeAllExternalKvBlocksAtTier", [&]() -> grpc::Status { + if (request->node_id().empty()) { + return grpc::Status(grpc::StatusCode::INVALID_ARGUMENT, "node_id must not be empty"); + } - const TierType tier = static_cast(request->tier()); - // Whole-tier wipe. UnregisterExternalKvByTier returns void (the store - // interface does not surface a mutated-block count), so the per-block - // counter is no longer emitted for this admin path; the revoke_total - // counter still records the operation. - store_.UnregisterExternalKvByTier(request->node_id(), tier); - if (metrics_) { - metrics_->addCounter( - MORI_UMBP_METRIC_EXT_KV_REVOKE_TOTAL, MORI_UMBP_METRIC_EXT_KV_REVOKE_TOTAL_HELP, - {{"node", request->node_id()}, {"tier", TierTypeName(tier)}, {"result", "ok"}}); - } - return grpc::Status::OK; + const TierType tier = static_cast(request->tier()); + // Whole-tier wipe. UnregisterExternalKvByTier returns void (the store + // interface does not surface a mutated-block count), so the per-block + // counter is no longer emitted for this admin path; the revoke_total + // counter still records the operation. + store_.UnregisterExternalKvByTier(request->node_id(), tier); + if (metrics_) { + metrics_->addCounter( + MORI_UMBP_METRIC_EXT_KV_REVOKE_TOTAL, MORI_UMBP_METRIC_EXT_KV_REVOKE_TOTAL_HELP, + {{"node", request->node_id()}, {"tier", TierTypeName(tier)}, {"result", "ok"}}); + } + return grpc::Status::OK; + }); } grpc::Status RevokeAllExternalKvBlocksForNode( grpc::ServerContext* /*ctx*/, const ::umbp::RevokeAllExternalKvBlocksForNodeRequest* request, ::umbp::RevokeAllExternalKvBlocksForNodeResponse* /*response*/) override { - if (request->node_id().empty()) { - return grpc::Status(grpc::StatusCode::INVALID_ARGUMENT, "node_id must not be empty"); - } + return GuardStore("RevokeAllExternalKvBlocksForNode", [&]() -> grpc::Status { + if (request->node_id().empty()) { + return grpc::Status(grpc::StatusCode::INVALID_ARGUMENT, "node_id must not be empty"); + } - // All-tier wipe for one node (full-sync recovery path). Returns void, so - // the per-block counter is dropped here as in the per-tier wipe above. - store_.UnregisterExternalKvByNode(request->node_id()); - if (metrics_) { - metrics_->addCounter(MORI_UMBP_METRIC_EXT_KV_REVOKE_TOTAL, - MORI_UMBP_METRIC_EXT_KV_REVOKE_TOTAL_HELP, - {{"node", request->node_id()}, {"tier", "ALL"}, {"result", "ok"}}); - } - return grpc::Status::OK; + // All-tier wipe for one node (full-sync recovery path). Returns void, so + // the per-block counter is dropped here as in the per-tier wipe above. + store_.UnregisterExternalKvByNode(request->node_id()); + if (metrics_) { + metrics_->addCounter(MORI_UMBP_METRIC_EXT_KV_REVOKE_TOTAL, + MORI_UMBP_METRIC_EXT_KV_REVOKE_TOTAL_HELP, + {{"node", request->node_id()}, {"tier", "ALL"}, {"result", "ok"}}); + } + return grpc::Status::OK; + }); } grpc::Status MatchExternalKv(grpc::ServerContext* /*ctx*/, const ::umbp::MatchExternalKvRequest* request, ::umbp::MatchExternalKvResponse* response) override { - std::vector hashes(request->hashes().begin(), request->hashes().end()); - // The store fuses the match with the hit-count increment + last_seen stamp - // (when count_as_hit) into one lock acquisition; the handler no longer - // makes a separate IncrementHits call. `now` crosses the store boundary - // (system_clock) and feeds GarbageCollectHits. - auto matches = - store_.MatchExternalKv(hashes, request->count_as_hit(), std::chrono::system_clock::now()); - - auto peer_map = store_.GetAlivePeerView(); - for (auto& m : matches) { - auto* proto_match = response->add_matches(); - proto_match->set_node_id(m.node_id); - auto peer_it = peer_map.find(m.node_id); - if (peer_it != peer_map.end()) proto_match->set_peer_address(peer_it->second); - for (const auto& [tier, hashes] : m.hashes_by_tier) { - auto* proto_bucket = proto_match->add_hashes_by_tier(); - proto_bucket->set_tier(static_cast<::umbp::TierType>(tier)); - for (const auto& hash : hashes) proto_bucket->add_hashes(hash); + return GuardStore("MatchExternalKv", [&]() -> grpc::Status { + std::vector hashes(request->hashes().begin(), request->hashes().end()); + // The store fuses the match with the hit-count increment + last_seen stamp + // (when count_as_hit) into one lock acquisition; the handler no longer + // makes a separate IncrementHits call. `now` crosses the store boundary + // (system_clock) and feeds GarbageCollectHits. + auto matches = + store_.MatchExternalKv(hashes, request->count_as_hit(), std::chrono::system_clock::now()); + + auto peer_map = store_.GetAlivePeerView(); + for (auto& m : matches) { + auto* proto_match = response->add_matches(); + proto_match->set_node_id(m.node_id); + auto peer_it = peer_map.find(m.node_id); + if (peer_it != peer_map.end()) proto_match->set_peer_address(peer_it->second); + for (const auto& [tier, hashes] : m.hashes_by_tier) { + auto* proto_bucket = proto_match->add_hashes_by_tier(); + proto_bucket->set_tier(static_cast<::umbp::TierType>(tier)); + for (const auto& hash : hashes) proto_bucket->add_hashes(hash); + } } - } - size_t total_matched = 0; - for (const auto& m : matches) total_matched += m.MatchedHashCount(); - if (metrics_) { - metrics_->addCounter(MORI_UMBP_METRIC_EXT_KV_MATCH_TOTAL, - MORI_UMBP_METRIC_EXT_KV_MATCH_TOTAL_HELP); - metrics_->addCounter(MORI_UMBP_METRIC_EXT_KV_MATCH_QUERIED_BLOCKS_TOTAL, - MORI_UMBP_METRIC_EXT_KV_MATCH_QUERIED_BLOCKS_TOTAL_HELP, - static_cast(hashes.size())); - metrics_->addCounter(MORI_UMBP_METRIC_EXT_KV_MATCH_MATCHED_BLOCKS_TOTAL, - MORI_UMBP_METRIC_EXT_KV_MATCH_MATCHED_BLOCKS_TOTAL_HELP, - static_cast(total_matched)); - } - return grpc::Status::OK; + size_t total_matched = 0; + for (const auto& m : matches) total_matched += m.MatchedHashCount(); + if (metrics_) { + metrics_->addCounter(MORI_UMBP_METRIC_EXT_KV_MATCH_TOTAL, + MORI_UMBP_METRIC_EXT_KV_MATCH_TOTAL_HELP); + metrics_->addCounter(MORI_UMBP_METRIC_EXT_KV_MATCH_QUERIED_BLOCKS_TOTAL, + MORI_UMBP_METRIC_EXT_KV_MATCH_QUERIED_BLOCKS_TOTAL_HELP, + static_cast(hashes.size())); + metrics_->addCounter(MORI_UMBP_METRIC_EXT_KV_MATCH_MATCHED_BLOCKS_TOTAL, + MORI_UMBP_METRIC_EXT_KV_MATCH_MATCHED_BLOCKS_TOTAL_HELP, + static_cast(total_matched)); + } + return grpc::Status::OK; + }); } grpc::Status GetExternalKvHitCounts(grpc::ServerContext* /*ctx*/, const ::umbp::GetExternalKvHitCountsRequest* request, ::umbp::GetExternalKvHitCountsResponse* response) override { - const size_t max_batch = static_cast(HitQueryMaxBatch()); - if (static_cast(request->hashes_size()) > max_batch) { - return grpc::Status( - grpc::StatusCode::INVALID_ARGUMENT, - "hashes size exceeds UMBP_HIT_QUERY_MAX_BATCH=" + std::to_string(max_batch)); - } + return GuardStore("GetExternalKvHitCounts", [&]() -> grpc::Status { + const size_t max_batch = static_cast(HitQueryMaxBatch()); + if (static_cast(request->hashes_size()) > max_batch) { + return grpc::Status( + grpc::StatusCode::INVALID_ARGUMENT, + "hashes size exceeds UMBP_HIT_QUERY_MAX_BATCH=" + std::to_string(max_batch)); + } - std::vector hashes(request->hashes().begin(), request->hashes().end()); - auto entries = store_.GetExternalKvHitCounts(hashes); - for (const auto& e : entries) { - auto* entry = response->add_entries(); - entry->set_hash(e.hash); - entry->set_hit_count_total(e.hit_count_total); - } - return grpc::Status::OK; + std::vector hashes(request->hashes().begin(), request->hashes().end()); + auto entries = store_.GetExternalKvHitCounts(hashes); + for (const auto& e : entries) { + auto* entry = response->add_entries(); + entry->set_hash(e.hash); + entry->set_hit_count_total(e.hit_count_total); + } + return grpc::Status::OK; + }); } grpc::Status ReportMetrics(grpc::ServerContext* /*ctx*/, const ::umbp::ReportMetricsRequest* request, ::umbp::ReportMetricsResponse* /*response*/) override { - if (!metrics_) return grpc::Status::OK; + return GuardStore("ReportMetrics", [&]() -> grpc::Status { + if (!metrics_) return grpc::Status::OK; - mori::metrics::MetricsServer::Labels base = {{"node", request->node_id()}}; - for (const auto& tag : store_.GetClientTags(request->node_id())) { - const auto sep = tag.find('='); - if (sep != std::string::npos) { - base.push_back({tag.substr(0, sep), tag.substr(sep + 1)}); + mori::metrics::MetricsServer::Labels base = {{"node", request->node_id()}}; + for (const auto& tag : store_.GetClientTags(request->node_id())) { + const auto sep = tag.find('='); + if (sep != std::string::npos) { + base.push_back({tag.substr(0, sep), tag.substr(sep + 1)}); + } } - } - for (const auto& s : request->metrics()) { - mori::metrics::MetricsServer::Labels labels = base; - for (const auto& l : s.labels()) labels.push_back({l.name(), l.value()}); - switch (s.value_case()) { - case ::umbp::MetricSample::kCounterDelta: - metrics_->addCounter(s.name(), s.help(), labels, - static_cast(s.counter_delta())); - break; - case ::umbp::MetricSample::kGaugeValue: - metrics_->setGauge(s.name(), s.help(), labels, s.gauge_value()); - break; - case ::umbp::MetricSample::kHistogramAggregate: { - const auto& a = s.histogram_aggregate(); - std::vector bounds(a.bounds().begin(), a.bounds().end()); - std::vector counts(a.bucket_counts().begin(), a.bucket_counts().end()); - metrics_->observeAggregated(s.name(), s.help(), labels, bounds, counts, a.count(), - a.sum()); - break; + for (const auto& s : request->metrics()) { + mori::metrics::MetricsServer::Labels labels = base; + for (const auto& l : s.labels()) labels.push_back({l.name(), l.value()}); + switch (s.value_case()) { + case ::umbp::MetricSample::kCounterDelta: + metrics_->addCounter(s.name(), s.help(), labels, + static_cast(s.counter_delta())); + break; + case ::umbp::MetricSample::kGaugeValue: + metrics_->setGauge(s.name(), s.help(), labels, s.gauge_value()); + break; + case ::umbp::MetricSample::kHistogramAggregate: { + const auto& a = s.histogram_aggregate(); + std::vector bounds(a.bounds().begin(), a.bounds().end()); + std::vector counts(a.bucket_counts().begin(), a.bucket_counts().end()); + metrics_->observeAggregated(s.name(), s.help(), labels, bounds, counts, a.count(), + a.sum()); + break; + } + default: + break; } - default: - break; } - } - return grpc::Status::OK; + return grpc::Status::OK; + }); } void SetMetrics(mori::metrics::MetricsServer* metrics) { metrics_ = metrics; } @@ -771,7 +822,7 @@ class MasterServer::UMBPMasterServiceImpl final : public ::umbp::UMBPMaster::Ser // --------------------------------------------------------------------------- MasterServer::MasterServer(MasterServerConfig config) : config_(std::move(config)), - store_(std::make_unique()), + store_(MakeMasterMetadataStore()), router_(*store_, std::move(config_.get_strategy), std::move(config_.put_strategy)), service_(std::make_unique(*store_, router_, config_.registry_config, nullptr)), @@ -791,6 +842,7 @@ void MasterServer::Run() { if (config_.metrics_port > 0) { metrics_server_ = std::make_unique(config_.metrics_port); service_->SetMetrics(metrics_server_.get()); + store_->SetMetricsSink(metrics_server_.get()); metrics_server_->setGauge(MORI_UMBP_METRIC_CLIENT_COUNT, MORI_UMBP_METRIC_CLIENT_COUNT_HELP, 0.0); MORI_UMBP_INFO("[Master] Metrics server listening on port {}", config_.metrics_port); @@ -875,9 +927,15 @@ void MasterServer::HitIndexGcLoop() { // cutoff is a system_clock time_point now (hazard #7) so it's comparable // to the last_seen the store stamps in MatchExternalKv(count_as_hit=true). const auto cutoff = std::chrono::system_clock::now() - ttl; - const size_t dropped = store_->GarbageCollectHits(cutoff); - if (dropped > 0) { - MORI_UMBP_DEBUG("[Master] External KV hit index GC dropped {} entries", dropped); + // A store hiccup (e.g. RESP transport error) must not kill the GC thread; + // swallow, log, and retry on the next tick. + try { + const size_t dropped = store_->GarbageCollectHits(cutoff); + if (dropped > 0) { + MORI_UMBP_DEBUG("[Master] External KV hit index GC dropped {} entries", dropped); + } + } catch (const std::exception& e) { + MORI_UMBP_ERROR("[Master] hit-index GC metadata-store error: {}", e.what()); } } } @@ -916,9 +974,14 @@ void MasterServer::ReaperLoop() { if (!reaper_running_) break; const auto cutoff = std::chrono::system_clock::now() - expiry; - auto expired = store_->ExpireStaleClients(cutoff); - for (const auto& node_id : expired) { - MORI_UMBP_WARN("[Reaper] Expired client: {}", node_id); + // A store hiccup must not kill the reaper thread; log and retry next tick. + try { + auto expired = store_->ExpireStaleClients(cutoff); + for (const auto& node_id : expired) { + MORI_UMBP_WARN("[Reaper] Expired client: {}", node_id); + } + } catch (const std::exception& e) { + MORI_UMBP_ERROR("[Reaper] metadata-store error: {}", e.what()); } } } diff --git a/src/umbp/distributed/master/redis/resp_client.cpp b/src/umbp/distributed/master/redis/resp_client.cpp new file mode 100644 index 000000000..f49be1a6a --- /dev/null +++ b/src/umbp/distributed/master/redis/resp_client.cpp @@ -0,0 +1,414 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. +#include "umbp/distributed/master/redis/resp_client.h" + +#include + +#include +#include +#include + +#include "mori/metrics/prometheus_metrics_server.hpp" + +namespace mori::umbp::redis { + +namespace { + +// Backend label shared by every RESP-client cold-path metric. +const mori::metrics::MetricsServer::Labels kRedisLabels = {{"backend", "redis"}}; + +// Parse "tcp://host:port" (or "host:port") into (host, port). Throws on a +// malformed URI so a misconfigured master fails fast at startup. +void ParseUri(const std::string& uri, std::string* host, int* port) { + std::string rest = uri; + const std::string scheme = "tcp://"; + if (rest.rfind(scheme, 0) == 0) rest = rest.substr(scheme.size()); + // Strip any comma-separated cluster seeds; Phase 1 uses the first only. + const auto comma = rest.find(','); + if (comma != std::string::npos) rest = rest.substr(0, comma); + const auto colon = rest.rfind(':'); + if (colon == std::string::npos) { + *host = rest; + *port = 6379; + return; + } + *host = rest.substr(0, colon); + try { + *port = std::stoi(rest.substr(colon + 1)); + } catch (const std::exception&) { + throw RespError("RespClient: invalid port in URI '" + uri + "'"); + } +} + +} // namespace + +RespClient::Lease::~Lease() { owner_->Release(ctx_, healthy_); } + +RespClient::RespClient(Options options) : options_(std::move(options)) { + ParseUri(options_.uri, &host_, &port_); + if (options_.pool_size == 0) options_.pool_size = 1; + idle_.reserve(options_.pool_size); +} + +RespClient::~RespClient() { + std::lock_guard lk(mu_); + for (redisContext* c : idle_) { + if (c) redisFree(c); + } + idle_.clear(); +} + +void RespClient::CountTransportError() { + if (metrics_ == nullptr) return; + metrics_->addCounter("mori_umbp_redis_transport_errors_total", + "RESP client transport failures (connect/socket/protocol)", kRedisLabels); +} + +void RespClient::CountNoscriptReload() { + if (metrics_ == nullptr) return; + metrics_->addCounter("mori_umbp_redis_noscript_reload_total", + "EVALSHA NOSCRIPT reloads (server evicted the cached script)", kRedisLabels); +} + +redisContext* RespClient::Connect() { + timeval tv{}; + tv.tv_sec = options_.connect_timeout_ms / 1000; + tv.tv_usec = (options_.connect_timeout_ms % 1000) * 1000; + redisContext* ctx = redisConnectWithTimeout(host_.c_str(), port_, tv); + if (ctx == nullptr || ctx->err) { + const std::string msg = + ctx ? std::string(ctx->errstr) : std::string("redisConnectWithTimeout returned null"); + if (ctx) redisFree(ctx); + throw RespError("RespClient: connect to " + host_ + ":" + std::to_string(port_) + + " failed: " + msg); + } + timeval sock{}; + sock.tv_sec = options_.socket_timeout_ms / 1000; + sock.tv_usec = (options_.socket_timeout_ms % 1000) * 1000; + redisSetTimeout(ctx, sock); + + if (!options_.password.empty()) { + redisReply* r = + static_cast(redisCommand(ctx, "AUTH %s", options_.password.c_str())); + const bool bad = (r == nullptr) || (r->type == REDIS_REPLY_ERROR); + if (r) freeReplyObject(r); + if (bad) { + redisFree(ctx); + throw RespError("RespClient: AUTH failed"); + } + } + return ctx; +} + +RespClient::Lease RespClient::Acquire() { + std::unique_lock lk(mu_); + // Pool-wait timing is armed lazily: only when we actually block on an exhausted + // pool (the else branch below). The fast path (idle conn available or room to + // create one) reads no clock and touches no metric. + bool waited = false; + std::chrono::steady_clock::time_point wait_start; + auto observe_wait = [&] { + if (!waited || metrics_ == nullptr) return; + const double sec = + std::chrono::duration(std::chrono::steady_clock::now() - wait_start).count(); + static const std::vector kBounds = {0.0001, 0.00025, 0.0005, 0.001, 0.0025, + 0.005, 0.01, 0.025, 0.05, 0.1}; + metrics_->observe("mori_umbp_redis_pool_wait_seconds", + "Time a caller blocked waiting for a free pooled RESP connection", + kRedisLabels, kBounds, sec); + }; + for (;;) { + if (!idle_.empty()) { + redisContext* c = idle_.back(); + idle_.pop_back(); + observe_wait(); + return Lease(this, c); + } + if (created_ < options_.pool_size) { + ++created_; + lk.unlock(); + redisContext* c = nullptr; + try { + c = Connect(); + } catch (...) { + std::lock_guard relk(mu_); + --created_; + cv_.notify_one(); + throw; + } + observe_wait(); + return Lease(this, c); + } + if (metrics_ != nullptr && !waited) { + waited = true; + wait_start = std::chrono::steady_clock::now(); + metrics_->addCounter("mori_umbp_redis_pool_exhausted_total", + "Times a caller found the RESP connection pool exhausted and blocked", + kRedisLabels); + } + cv_.wait(lk); + } +} + +void RespClient::Release(redisContext* ctx, bool healthy) { + std::lock_guard lk(mu_); + if (healthy && ctx != nullptr && ctx->err == 0) { + idle_.push_back(ctx); + } else { + if (ctx) redisFree(ctx); + if (created_ > 0) --created_; + } + cv_.notify_one(); +} + +RespValue RespClient::Convert(void* reply_ptr) { + RespValue out; + auto* reply = static_cast(reply_ptr); + if (reply == nullptr) { + out.type = RespValue::Type::Nil; + return out; + } + switch (reply->type) { + case REDIS_REPLY_STRING: + out.type = RespValue::Type::String; + out.str.assign(reply->str, reply->len); + break; + case REDIS_REPLY_STATUS: + out.type = RespValue::Type::Status; + out.str.assign(reply->str, reply->len); + break; + case REDIS_REPLY_ERROR: + out.type = RespValue::Type::Error; + out.str.assign(reply->str, reply->len); + break; + case REDIS_REPLY_INTEGER: + out.type = RespValue::Type::Integer; + out.integer = reply->integer; + break; + case REDIS_REPLY_NIL: + out.type = RespValue::Type::Nil; + break; + case REDIS_REPLY_ARRAY: + out.type = RespValue::Type::Array; + out.elements.reserve(reply->elements); + for (size_t i = 0; i < reply->elements; ++i) { + out.elements.push_back(Convert(reply->element[i])); + } + break; + default: + // REDIS_REPLY_DOUBLE / MAP / SET / etc. (RESP3): treat as string/array + // best-effort. Phase 1 uses RESP2, so this is rarely hit. + if (reply->str != nullptr) { + out.type = RespValue::Type::String; + out.str.assign(reply->str, reply->len); + } else { + out.type = RespValue::Type::Nil; + } + break; + } + return out; +} + +RespValue RespClient::RunArgv(redisContext* ctx, const std::vector& args, + bool* broke) { + std::vector argv; + std::vector argvlen; + ToArgv(args, argv, argvlen); + auto* reply = static_cast( + redisCommandArgv(ctx, static_cast(argv.size()), argv.data(), argvlen.data())); + if (reply == nullptr) { + *broke = true; + CountTransportError(); + const std::string err = ctx ? std::string(ctx->errstr) : std::string("null reply"); + throw RespError("RespClient: command failed (transport): " + err); + } + RespValue out = Convert(reply); + freeReplyObject(reply); + return out; +} + +RespValue RespClient::Command(const std::vector& args) { + // Command backs only pure-READ metadata lookups (HGETALL/HGET/SCARD), so a + // transport failure on a stale pooled connection — a server/LB drops an idle + // socket and hiredis trips ctx->err only on the next use — is safe to retry + // once on a fresh connection. This is deliberately NOT done for Eval / + // EvalPipeline: route_get_batch bumps _lease/_lacc/_acnt, so a blind retry of a + // script that may already have run server-side could double-apply a side + // effect. A connect failure (Acquire throws) is a different case and propagates + // immediately. A second transport failure propagates. + for (int attempt = 0; attempt < 2; ++attempt) { + Lease lease = Acquire(); + bool broke = false; + try { + return RunArgv(lease.get(), args, &broke); + } catch (const RespError&) { + if (broke) lease.MarkBroken(); + if (broke && attempt == 0) continue; // stale socket: retry once, fresh conn. + throw; + } + } + throw RespError("RespClient: command retry exhausted"); // unreachable +} + +std::vector RespClient::Pipeline(const std::vector>& commands) { + std::vector replies; + replies.reserve(commands.size()); + Lease lease = Acquire(); + redisContext* ctx = lease.get(); + + for (const auto& cmd : commands) { + std::vector argv; + std::vector argvlen; + ToArgv(cmd, argv, argvlen); + if (redisAppendCommandArgv(ctx, static_cast(argv.size()), argv.data(), argvlen.data()) != + REDIS_OK) { + lease.MarkBroken(); + throw RespError("RespClient: pipeline append failed"); + } + } + for (size_t i = 0; i < commands.size(); ++i) { + void* r = nullptr; + if (redisGetReply(ctx, &r) != REDIS_OK) { + lease.MarkBroken(); + const std::string err = ctx ? std::string(ctx->errstr) : std::string("null"); + throw RespError("RespClient: pipeline read failed: " + err); + } + replies.push_back(Convert(r)); + if (r) freeReplyObject(static_cast(r)); + } + return replies; +} + +std::string RespClient::GetOrLoadSha(redisContext* ctx, const std::string& script, bool* broke) { + return script_cache_.GetOrLoad(script, [&](const std::string& s) { + RespValue r = RunArgv(ctx, {"SCRIPT", "LOAD", s}, broke); + if (r.type != RespValue::Type::String) { + throw RespError("RespClient: SCRIPT LOAD did not return a sha: " + r.str); + } + return r.str; + }); +} + +RespValue RespClient::Eval(const std::string& script, const std::vector& keys, + const std::vector& args) { + Lease lease = Acquire(); + redisContext* ctx = lease.get(); + bool broke = false; + try { + const std::string sha = GetOrLoadSha(ctx, script, &broke); + std::vector cmd = BuildEvalshaArgv(sha, keys, args); + + RespValue r = RunArgv(ctx, cmd, &broke); + if (r.is_error() && r.str.rfind("NOSCRIPT", 0) == 0) { + // Script evicted from the server cache; reload and retry once. + CountNoscriptReload(); + script_cache_.Invalidate(script); + cmd[1] = GetOrLoadSha(ctx, script, &broke); + r = RunArgv(ctx, cmd, &broke); + } + return r; + } catch (...) { + if (broke) lease.MarkBroken(); + throw; + } +} + +std::vector RespClient::PipelineEvalshaBatch( + redisContext* ctx, const std::string& sha, + const std::vector>& keys_per_call, + const std::vector& shared_args, const std::vector& indices, + std::vector* replies, bool* broke) { + // Queue one EVALSHA per index. + for (const std::size_t idx : indices) { + std::vector cmd = BuildEvalshaArgv(sha, keys_per_call[idx], shared_args); + std::vector argv; + std::vector argvlen; + ToArgv(cmd, argv, argvlen); + if (redisAppendCommandArgv(ctx, static_cast(argv.size()), argv.data(), argvlen.data()) != + REDIS_OK) { + *broke = true; + CountTransportError(); + throw RespError("RespClient: EvalPipeline append failed"); + } + } + + // Read replies back in the same order; flag the ones the server didn't have. + std::vector noscript; + for (const std::size_t idx : indices) { + void* r = nullptr; + if (redisGetReply(ctx, &r) != REDIS_OK) { + *broke = true; + CountTransportError(); + const std::string err = ctx ? std::string(ctx->errstr) : std::string("null"); + throw RespError("RespClient: EvalPipeline read failed: " + err); + } + RespValue val = Convert(r); + if (r) freeReplyObject(static_cast(r)); + if (val.is_error() && val.str.rfind("NOSCRIPT", 0) == 0) noscript.push_back(idx); + (*replies)[idx] = std::move(val); + } + return noscript; +} + +std::vector RespClient::EvalPipeline( + const std::string& script, const std::vector>& keys_per_call, + const std::vector& shared_args) { + std::vector replies(keys_per_call.size()); + if (keys_per_call.empty()) return replies; + + Lease lease = Acquire(); + redisContext* ctx = lease.get(); + bool broke = false; + try { + const std::string sha = GetOrLoadSha(ctx, script, &broke); + + std::vector all(keys_per_call.size()); + for (std::size_t i = 0; i < all.size(); ++i) all[i] = i; + std::vector missing = + PipelineEvalshaBatch(ctx, sha, keys_per_call, shared_args, all, &replies, &broke); + + // NOSCRIPT => the server evicted the script from its cache. Reload once and + // retry ONLY the calls that missed, so calls that already executed (and, for + // route_get_batch, already bumped lease/access) are not run a second time. + if (!missing.empty()) { + CountNoscriptReload(); + script_cache_.Invalidate(script); + const std::string sha2 = GetOrLoadSha(ctx, script, &broke); + PipelineEvalshaBatch(ctx, sha2, keys_per_call, shared_args, missing, &replies, &broke); + } + return replies; + } catch (...) { + if (broke) lease.MarkBroken(); + throw; + } +} + +bool RespClient::Ping() { + try { + RespValue r = Command({"PING"}); + return r.type == RespValue::Type::Status || r.type == RespValue::Type::String; + } catch (const std::exception&) { + return false; + } +} + +} // namespace mori::umbp::redis diff --git a/src/umbp/distributed/master/redis/resp_cluster_client.cpp b/src/umbp/distributed/master/redis/resp_cluster_client.cpp new file mode 100644 index 000000000..e2a654ea7 --- /dev/null +++ b/src/umbp/distributed/master/redis/resp_cluster_client.cpp @@ -0,0 +1,334 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. +#include "umbp/distributed/master/redis/resp_cluster_client.h" + +#include +#include + +#include +#include +#include +#include +#include +#include + +#include "mori/metrics/prometheus_metrics_server.hpp" +#include "mori/utils/mori_log.hpp" +#include "umbp/distributed/master/redis/resp_client.h" +#include "umbp/distributed/master/redis/thread_pool.h" + +namespace mori::umbp::redis { + +namespace { + +// Backend label shared by every cluster-client cold-path metric. +const mori::metrics::MetricsServer::Labels kRedisLabels = {{"backend", "redis"}}; + +// Parse "tcp://host:port" (or "host:port") into a redis-plus-plus +// ConnectionOptions host/port. Throws RespError on a malformed seed. +void ParseSeed(const std::string& uri, std::string* host, int* port) { + std::string rest = uri; + const std::string scheme = "tcp://"; + if (rest.rfind(scheme, 0) == 0) rest = rest.substr(scheme.size()); + const auto slash = rest.find('/'); + if (slash != std::string::npos) rest = rest.substr(0, slash); + const auto colon = rest.rfind(':'); + if (colon == std::string::npos) { + *host = rest; + *port = 6379; + return; + } + *host = rest.substr(0, colon); + try { + *port = std::stoi(rest.substr(colon + 1)); + } catch (const std::exception&) { + throw RespError("RespClusterClient: invalid port in seed '" + uri + "'"); + } +} + +// Build a redis-plus-plus ConnectionOptions for one seed from our Options +// (host/port from the seed URI, plus the shared password + timeouts). Shared by +// ConnectCluster and DiscoverMasterSlotRanges so the seed->options mapping lives +// in one place. +sw::redis::ConnectionOptions SeedConnectionOptions(const RespClusterClient::Options& o, + const std::string& seed) { + sw::redis::ConnectionOptions co; + ParseSeed(seed, &co.host, &co.port); + if (!o.password.empty()) co.password = o.password; + co.connect_timeout = std::chrono::milliseconds(o.connect_timeout_ms); + co.socket_timeout = std::chrono::milliseconds(o.socket_timeout_ms); + return co; +} + +// Connect to the cluster by trying each seed until one lets redis-plus-plus +// fetch the topology (its RedisCluster ctor runs CLUSTER SLOTS and throws if the +// seed is down). Shared by the client ctor and DiscoverMasterSlotRanges's +// callers. +std::unique_ptr ConnectCluster(const RespClusterClient::Options& o) { + if (o.seeds.empty()) throw RespError("RespClusterClient: no cluster seeds configured"); + sw::redis::ConnectionPoolOptions pool_opts; + pool_opts.size = o.pool_size == 0 ? 1 : o.pool_size; + std::string last_err; + for (const auto& seed : o.seeds) { + sw::redis::ConnectionOptions co = SeedConnectionOptions(o, seed); + try { + return std::make_unique(co, pool_opts); + } catch (const sw::redis::Error& e) { + last_err = e.what(); + MORI_UMBP_WARN("[RedisCluster] seed {} unreachable, trying next: {}", seed, e.what()); + } + } + throw RespError("RespClusterClient: no reachable cluster seed (last error: " + last_err + ")"); +} + +// Parse a CLUSTER SLOTS reply into per-master slot ranges, ordered by each +// master's lowest slot (stable). Reply shape: array of +// [start, end, [master_ip, master_port, master_id, ...], [replica...], ...]. +std::vector> ParseClusterSlots(const redisReply* reply) { + std::vector> ranges; + std::vector min_slot; + std::unordered_map master_idx; + if (reply == nullptr || reply->type != REDIS_REPLY_ARRAY) return ranges; + + for (std::size_t i = 0; i < reply->elements; ++i) { + const redisReply* e = reply->element[i]; + if (e == nullptr || e->type != REDIS_REPLY_ARRAY || e->elements < 3) continue; + if (e->element[0]->type != REDIS_REPLY_INTEGER || e->element[1]->type != REDIS_REPLY_INTEGER) { + continue; + } + const auto start = static_cast(e->element[0]->integer); + const auto end = static_cast(e->element[1]->integer); + const redisReply* m = e->element[2]; // master node descriptor + if (m == nullptr || m->type != REDIS_REPLY_ARRAY || m->elements < 2) continue; + + // Key a master by its node id (element[2]) when present, else ip:port. + std::string key; + if (m->elements >= 3 && m->element[2] != nullptr && m->element[2]->type == REDIS_REPLY_STRING) { + key.assign(m->element[2]->str, m->element[2]->len); + } else { + const std::string ip = (m->element[0]->type == REDIS_REPLY_STRING) + ? std::string(m->element[0]->str, m->element[0]->len) + : std::string(); + const long long port = + (m->element[1]->type == REDIS_REPLY_INTEGER) ? m->element[1]->integer : 0; + key = ip + ":" + std::to_string(port); + } + + auto it = master_idx.find(key); + if (it == master_idx.end()) { + master_idx.emplace(key, ranges.size()); + ranges.push_back({SlotRange{start, end}}); + min_slot.push_back(start); + } else { + ranges[it->second].push_back(SlotRange{start, end}); + if (start < min_slot[it->second]) min_slot[it->second] = start; + } + } + + // Order masters by their lowest slot so the tag assignment is deterministic. + std::vector order(ranges.size()); + std::iota(order.begin(), order.end(), 0); + std::sort(order.begin(), order.end(), + [&](std::size_t a, std::size_t b) { return min_slot[a] < min_slot[b]; }); + std::vector> ordered; + ordered.reserve(ranges.size()); + for (std::size_t o : order) ordered.push_back(std::move(ranges[o])); + return ordered; +} + +} // namespace + +std::vector> RespClusterClient::DiscoverMasterSlotRanges( + const Options& options) { + if (options.seeds.empty()) throw RespError("RespClusterClient: no cluster seeds configured"); + // CLUSTER SLOTS is a node command (no key), so query a plain connection to the + // first reachable seed rather than routing through the cluster client. + std::string last_err; + for (const auto& seed : options.seeds) { + try { + sw::redis::Redis r(SeedConnectionOptions(options, seed)); + auto reply = r.command("CLUSTER", "SLOTS"); + return ParseClusterSlots(reply.get()); + } catch (const sw::redis::Error& e) { + last_err = e.what(); + } + } + throw RespError("RespClusterClient: CLUSTER SLOTS discovery failed (last error: " + last_err + + ")"); +} + +RespClusterClient::RespClusterClient(Options options) : options_(std::move(options)) { + if (options_.pool_size == 0) options_.pool_size = 1; + cluster_ = ConnectCluster(options_); + + // Workers to fan a multi-slot EvalPipeline out concurrently (one round trip + // per node instead of N sequential). Each task is an independent + // redirection-aware call, so a slow/failing node cannot block the others. + const std::size_t threads = + std::min(64, std::max(4, options_.pool_size)); + pool_ = std::make_unique(threads); + + MORI_UMBP_INFO("[RedisCluster] connected via {} seed(s), pool={}, fanout={}", + options_.seeds.size(), options_.pool_size, threads); +} + +RespClusterClient::~RespClusterClient() = default; + +RespValue RespClusterClient::RunArgvRouted(const std::vector& argv, + const std::string& routing_key) { + // Send the raw argv to the node owning routing_key's slot. redis-plus-plus's + // RedisCluster::command redirection loop resolves the slot, follows MOVED/ASK, + // and refreshes the slot map on a moved slot or a failed-over master. + auto sender = [&argv](sw::redis::Connection& conn, const sw::redis::StringView&) { + std::vector ptrs; + std::vector lens; + ToArgv(argv, ptrs, lens); + conn.send(static_cast(ptrs.size()), ptrs.data(), lens.data()); + }; + + try { + auto reply = cluster_->command(sender, sw::redis::StringView(routing_key)); + return RespClient::Convert(reply.get()); + } catch (const sw::redis::ReplyError& e) { + // A server-side logical error (e.g. a Lua runtime error). Surface it as an + // Error value so callers keep the error-as-value contract. (MOVED/ASK are + // handled inside the redirection loop and do not reach here.) + RespValue v; + v.type = RespValue::Type::Error; + v.str = e.what(); + return v; + } catch (const sw::redis::Error& e) { + // Transport / redirection-exhausted / cluster-down: a real failure. + if (metrics_ != nullptr) { + metrics_->addCounter("mori_umbp_redis_transport_errors_total", + "RESP client transport failures (connect/socket/protocol)", + kRedisLabels); + } + throw RespError(std::string("RespClusterClient: ") + e.what()); + } +} + +std::string RespClusterClient::GetOrLoadSha(const std::string& script) { + return script_cache_.GetOrLoad(script, [&](const std::string& s) { + // SCRIPT LOAD on any node returns the content-addressed SHA (same everywhere). + RespValue r = RunArgvRouted({"SCRIPT", "LOAD", s}, "sha"); + if (r.type != RespValue::Type::String) { + throw RespError("RespClusterClient: SCRIPT LOAD did not return a sha: " + r.str); + } + return r.str; + }); +} + +RespValue RespClusterClient::EvalShaRouted(const std::string& script, + const std::vector& keys, + const std::vector& args) { + const std::string sha = GetOrLoadSha(script); + std::vector argv = BuildEvalshaArgv(sha, keys, args); + + RespValue r = RunArgvRouted(argv, keys.front()); + if (r.is_error() && r.str.rfind("NOSCRIPT", 0) == 0) { + // The node serving this slot doesn't have the script cached (e.g. a promoted + // replica or a resharded node). Load it there (routed by the same key) and + // retry once. EVALSHA has no side effects until it runs, so this is safe. + if (metrics_ != nullptr) { + metrics_->addCounter("mori_umbp_redis_noscript_reload_total", + "EVALSHA NOSCRIPT reloads (server evicted the cached script)", + kRedisLabels); + } + RunArgvRouted({"SCRIPT", "LOAD", script}, keys.front()); + r = RunArgvRouted(argv, keys.front()); + } + return r; +} + +RespValue RespClusterClient::Command(const std::vector& args) { + if (args.empty()) throw RespError("RespClusterClient::Command: empty argv"); + // Store commands put the key at args[1] (HGETALL/HGET/SCARD ...); fall + // back to args[0] for keyless commands. + const std::string& routing_key = args.size() > 1 ? args[1] : args[0]; + return RunArgvRouted(args, routing_key); +} + +RespValue RespClusterClient::Eval(const std::string& script, const std::vector& keys, + const std::vector& args) { + if (keys.empty()) { + throw RespError("RespClusterClient::Eval: a routing key (KEYS[1]) is required in cluster mode"); + } + return EvalShaRouted(script, keys, args); +} + +std::vector RespClusterClient::EvalPipeline( + const std::string& script, const std::vector>& keys_per_call, + const std::vector& shared_args) { + std::vector replies(keys_per_call.size()); + if (keys_per_call.empty()) return replies; + + // One independent, redirection-aware EVAL per group; groups are single-slot + // (the caller groups a batch by shard tag) so each routes cleanly by keys[0]. + auto run_one = [&](std::size_t i) { + const auto& keys = keys_per_call[i]; + if (keys.empty()) { + RespValue v; + v.type = RespValue::Type::Error; + v.str = "RespClusterClient::EvalPipeline: empty KEYS group"; + replies[i] = std::move(v); + return; + } + try { + replies[i] = EvalShaRouted(script, keys, shared_args); + } catch (const std::exception& e) { + // Keep parallel-to-input: a failed group becomes an Error value the caller + // can inspect (mirrors the hiredis EvalPipeline's per-call replies). + RespValue v; + v.type = RespValue::Type::Error; + v.str = e.what(); + replies[i] = std::move(v); + } + }; + + if (keys_per_call.size() == 1 || pool_ == nullptr) { + for (std::size_t i = 0; i < keys_per_call.size(); ++i) run_one(i); + return replies; + } + + // Fan the extra groups out to the pool, run the first inline, then join all. + std::vector> futs; + futs.reserve(keys_per_call.size() - 1); + for (std::size_t i = 1; i < keys_per_call.size(); ++i) { + futs.push_back(pool_->Enqueue([&run_one, i] { run_one(i); })); + } + run_one(0); + for (auto& f : futs) f.get(); + return replies; +} + +bool RespClusterClient::Ping() { + try { + // PING ignores its "key"; routing by an arbitrary tag just picks a node. + RespValue r = RunArgvRouted({"PING"}, "ping"); + return r.type == RespValue::Type::Status || r.type == RespValue::Type::String; + } catch (const std::exception&) { + return false; + } +} + +} // namespace mori::umbp::redis diff --git a/src/umbp/distributed/master/redis_master_metadata_store.cpp b/src/umbp/distributed/master/redis_master_metadata_store.cpp new file mode 100644 index 000000000..1c4425630 --- /dev/null +++ b/src/umbp/distributed/master/redis_master_metadata_store.cpp @@ -0,0 +1,1050 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. +#include "umbp/distributed/master/redis_master_metadata_store.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "mori/metrics/prometheus_metrics_server.hpp" +#include "mori/utils/mori_log.hpp" +#include "umbp/distributed/master/redis/lua_scripts.h" +#include "umbp/distributed/master/redis/resp_cluster_client.h" +#include "umbp/distributed/master/redis/sharded_read.h" + +namespace mori::umbp { + +namespace { + +using redis::RespValue; + +// Records store-op latency into mori_umbp_store_op_latency_seconds{op,backend=redis} +// on scope exit when a metrics sink is attached; a cheap no-op otherwise. Lets a +// dashboard separate backend round-trip cost per store method. +class ScopedStoreOp { + public: + ScopedStoreOp(mori::metrics::MetricsServer* metrics, const char* op) + : metrics_(metrics), op_(op) { + if (metrics_) t0_ = std::chrono::steady_clock::now(); + } + ~ScopedStoreOp() { + if (!metrics_) return; + const double sec = + std::chrono::duration(std::chrono::steady_clock::now() - t0_).count(); + static const std::vector kBounds = {0.0001, 0.00025, 0.0005, 0.001, 0.0025, 0.005, 0.01, + 0.025, 0.05, 0.1, 0.25, 0.5, 1.0}; + metrics_->observe("mori_umbp_store_op_latency_seconds", + "Latency of IMasterMetadataStore operations against the backend", + {{"op", op_}, {"backend", "redis"}}, kBounds, sec); + } + ScopedStoreOp(const ScopedStoreOp&) = delete; + ScopedStoreOp& operator=(const ScopedStoreOp&) = delete; + + private: + mori::metrics::MetricsServer* metrics_; + const char* op_; + std::chrono::steady_clock::time_point t0_; +}; + +int64_t ToEpochMs(std::chrono::system_clock::time_point tp) { + return std::chrono::duration_cast(tp.time_since_epoch()).count(); +} + +int64_t ToMs(std::chrono::system_clock::duration d) { + return std::chrono::duration_cast(d).count(); +} + +std::chrono::system_clock::time_point FromEpochMs(int64_t ms) { + return std::chrono::system_clock::time_point(std::chrono::milliseconds(ms)); +} + +// caps: "tier:total:avail;tier:total:avail;..." +std::string EncodeCaps(const std::map& caps) { + std::string out; + for (const auto& [tier, cap] : caps) { + out += std::to_string(static_cast(tier)); + out += ':'; + out += std::to_string(cap.total_bytes); + out += ':'; + out += std::to_string(cap.available_bytes); + out += ';'; + } + return out; +} + +std::map DecodeCaps(const std::string& blob) { + std::map caps; + size_t pos = 0; + while (pos < blob.size()) { + const size_t semi = blob.find(';', pos); + const size_t end = (semi == std::string::npos) ? blob.size() : semi; + const std::string tok = blob.substr(pos, end - pos); + pos = (semi == std::string::npos) ? blob.size() : semi + 1; + if (tok.empty()) continue; + const size_t c1 = tok.find(':'); + const size_t c2 = (c1 == std::string::npos) ? std::string::npos : tok.find(':', c1 + 1); + if (c1 == std::string::npos || c2 == std::string::npos) continue; + try { + const int tier = std::stoi(tok.substr(0, c1)); + const uint64_t total = std::stoull(tok.substr(c1 + 1, c2 - c1 - 1)); + const uint64_t avail = std::stoull(tok.substr(c2 + 1)); + caps[static_cast(tier)] = TierCapacity{total, avail}; + } catch (const std::exception&) { + // best-effort decode; skip malformed token + } + } + return caps; +} + +std::string JoinTags(const std::vector& tags) { + std::string out; + for (size_t i = 0; i < tags.size(); ++i) { + if (i) out += '\n'; + out += tags[i]; + } + return out; +} + +std::vector SplitTags(const std::string& blob) { + std::vector out; + if (blob.empty()) return out; + size_t pos = 0; + while (pos <= blob.size()) { + const size_t nl = blob.find('\n', pos); + const size_t end = (nl == std::string::npos) ? blob.size() : nl; + out.push_back(blob.substr(pos, end - pos)); + if (nl == std::string::npos) break; + pos = nl + 1; + } + return out; +} + +// Dedup a hash list preserving first-seen order. The external-KV read scripts +// used to dedup with a Lua `seen` table; doing it here lets every touched key be +// declared in KEYS[] (no allow-undeclared-keys) and lets a match/hit result be +// scattered back by first-seen position. +std::vector DedupPreserveOrder(const std::vector& in) { + std::vector out; + out.reserve(in.size()); + std::unordered_set seen; + seen.reserve(in.size() * 2); + for (const auto& s : in) { + if (seen.insert(s).second) out.push_back(s); + } + return out; +} + +// Decode a flat HGETALL reply [f,v,f,v,...] into a field->value map. +std::unordered_map FlatToMap(const RespValue& flat) { + std::unordered_map m; + if (!flat.is_array()) return m; + for (size_t i = 0; i + 1 < flat.elements.size(); i += 2) { + m.emplace(flat.elements[i].str, flat.elements[i + 1].str); + } + return m; +} + +ClientRecord DecodeRecord(const std::string& node_id, + const std::unordered_map& f) { + ClientRecord rec; + rec.node_id = node_id; + auto get = [&](const char* k) -> const std::string* { + auto it = f.find(k); + return it == f.end() ? nullptr : &it->second; + }; + if (auto* v = get("addr")) rec.node_address = *v; + if (auto* v = get("peer")) rec.peer_address = *v; + if (auto* v = get("status")) { + try { + rec.status = static_cast(std::stoi(*v)); + } catch (...) { + rec.status = ClientStatus::UNKNOWN; + } + } + if (auto* v = get("last_hb")) { + try { + rec.last_heartbeat = FromEpochMs(std::stoll(*v)); + } catch (...) { + } + } + if (auto* v = get("reg_at")) { + try { + rec.registered_at = FromEpochMs(std::stoll(*v)); + } catch (...) { + } + } + if (auto* v = get("seq")) { + try { + rec.last_applied_seq = std::stoull(*v); + } catch (...) { + } + } + if (auto* v = get("caps")) rec.tier_capacities = DecodeCaps(*v); + if (auto* v = get("engine")) rec.engine_desc_bytes.assign(v->begin(), v->end()); + if (auto* v = get("tags")) rec.tags = SplitTags(*v); + return rec; +} + +// ShardedBatch / GroupKeysByShard / RunShardedRead moved to +// redis/sharded_read.h so the batch fan-out + fault-tolerance + reply-scatter +// logic can be unit-tested against a fake IRespClient (see test_sharded_read). + +} // namespace + +std::size_t RedisMasterMetadataStore::ResolveNumShards(const Config& config) { + if (config.cluster) + return config.block_shards == 0 ? 1 : config.block_shards; // slot-spread tags + if (config.shard_uris.size() > 1) return config.shard_uris.size(); // one shard per endpoint + return config.block_shards == 0 ? 1 : config.block_shards; // single-endpoint knob +} + +redis::KeySchema RedisMasterMetadataStore::BuildKeySchema(const Config& config) { + // Cluster balanced placement: use the explicit per-master tags when the + // factory supplied them; otherwise fall back to formulaic shard tags. + if (config.cluster && !config.cluster_block_tags.empty()) { + return redis::KeySchema(config.namespace_id, config.cluster_block_tags); + } + return redis::KeySchema(config.namespace_id, ResolveNumShards(config)); +} + +RedisMasterMetadataStore::RedisMasterMetadataStore(const Config& config) + : keys_(BuildKeySchema(config)), + mode_(config.cluster ? Mode::kCluster + : config.shard_uris.size() > 1 ? Mode::kMulti + : Mode::kSingle) { + if (mode_ == Mode::kCluster) { + // One Redis Cluster client routes every shard by slot; the store still uses + // the split control + per-shard-block write path (their keys live in + // different slots), driven by split_writes(). + redis::RespClusterClient::Options opts; + opts.seeds = + config.cluster_seeds.empty() ? std::vector{config.uri} : config.cluster_seeds; + opts.password = config.password; + opts.connect_timeout_ms = config.connect_timeout_ms; + opts.socket_timeout_ms = config.socket_timeout_ms; + opts.pool_size = config.pool_size; + clients_.push_back(std::make_unique(std::move(opts))); + MORI_UMBP_INFO("[RedisStore] cluster namespace={} pool={} seeds={} block_shards={}", + config.namespace_id, config.pool_size, opts.seeds.size(), keys_.NumShards()); + return; + } + + // Endpoints: the shard_uris list when given (multi-endpoint), else the single + // `uri`. clients_[0] is the control instance and also backs block shard 0. + std::vector uris = + config.shard_uris.size() > 1 ? config.shard_uris : std::vector{config.uri}; + + clients_.reserve(uris.size()); + for (const auto& uri : uris) { + redis::RespClient::Options opts; + opts.uri = uri; + opts.password = config.password; + opts.connect_timeout_ms = config.connect_timeout_ms; + opts.socket_timeout_ms = config.socket_timeout_ms; + opts.pool_size = config.pool_size; + clients_.push_back(std::make_unique(std::move(opts))); + } + + if (multi_endpoint()) { + // Workers to fan the per-instance read calls out concurrently. Sized so a + // handful of in-flight RouteGets can each parallelize their N instance calls + // without per-call thread churn; workers block on Redis I/O so a generous + // count is cheap. Capped to keep it bounded. + const std::size_t pool_threads = std::min(64, clients_.size() * 8); + fanout_pool_ = std::make_unique(pool_threads); + + std::string joined; + for (size_t i = 0; i < uris.size(); ++i) joined += (i ? "," : "") + uris[i]; + MORI_UMBP_INFO("[RedisStore] multi-endpoint namespace={} pool={} endpoints={} fanout={} [{}]", + config.namespace_id, config.pool_size, uris.size(), pool_threads, joined); + } else { + MORI_UMBP_INFO("[RedisStore] backend at {} namespace={} pool={} block_shards={}", config.uri, + config.namespace_id, config.pool_size, keys_.NumShards()); + } +} + +// ===================================================================== +// Multi-endpoint write helpers: run per-shard block scripts on each shard's own +// instance. Idempotent (ADD overwrites / REMOVE no-op / full_sync clear+replay) +// so a thrown/retried step is safe; a partial failure surfaces as an exception, +// the peer retries, seq-gaps, and heals via full_sync. +// ===================================================================== + +void RedisMasterMetadataStore::ApplyBlockEventsMulti(const std::string& node_id, + const std::vector& events, + bool is_full_sync, + std::chrono::system_clock::time_point now) { + const std::string node_prefix = "l|" + node_id + "|"; + const std::string now_ms = std::to_string(ToEpochMs(now)); + + // Group events by their shard so each instance gets exactly its own events. + std::vector> events_by_shard(keys_.NumShards()); + for (const auto& ev : events) events_by_shard[keys_.ShardOf(ev.key)].push_back(&ev); + + // full_sync must clear the node on EVERY shard (to drop stale locations), even + // shards with no new ADDs; a delta only touches shards that have events. + for (std::size_t shard = 0; shard < keys_.NumShards(); ++shard) { + const auto& shard_events = events_by_shard[shard]; + if (!is_full_sync && shard_events.empty()) continue; + + std::vector args; + args.reserve(5 + shard_events.size() * 4); + args.push_back(keys_.NodeBlocks(node_id, shard)); + args.push_back(node_prefix); + args.push_back(is_full_sync ? "1" : "0"); + args.push_back(now_ms); + args.push_back(std::to_string(shard_events.size())); + for (const KvEvent* ev : shard_events) { + args.push_back(ev->kind == KvEvent::Kind::ADD ? "0" : "1"); + args.push_back(keys_.Block(ev->key)); + args.push_back(std::to_string(static_cast(ev->tier))); + args.push_back(std::to_string(ev->size)); + } + // KEYS[1] = this shard's reverse-index key (a shard-tag key) so the cluster + // client can route to the shard's slot; the script reads all its keys from + // ARGV (same slot). Harmless in single / multi-endpoint mode (a standalone + // Redis does not slot-check, and the script ignores KEYS). + RespValue r = client_for_shard(shard).Eval(redis::kApplyBlockEventsLua, + {keys_.NodeBlocks(node_id, shard)}, args); + if (r.is_error()) throw std::runtime_error("[RedisStore] ApplyBlockEvents: " + r.str); + } +} + +void RedisMasterMetadataStore::WipeNodeBlocksMulti(const std::string& node_id) { + const std::string node_prefix = "l|" + node_id + "|"; + // Best-effort per shard: a down shard must not block wiping the reachable + // ones. The node is already gone/EXPIRED on the control instance, so any + // locations lingering on an unreachable shard point at a dead node and are + // filtered out of reads by GetAlivePeerView until that shard returns. + for (std::size_t shard = 0; shard < keys_.NumShards(); ++shard) { + try { + RespValue r = client_for_shard(shard).Eval(redis::kWipeNodeBlocksLua, + {keys_.NodeBlocks(node_id, shard)}, + {keys_.NodeBlocks(node_id, shard), node_prefix}); + if (r.is_error()) { + MORI_UMBP_WARN("[RedisStore] WipeNodeBlocks: shard {} script error, skipped: {}", shard, + r.str); + } + } catch (const std::exception& e) { + MORI_UMBP_WARN("[RedisStore] WipeNodeBlocks: shard {} unavailable, skipped: {}", shard, + e.what()); + } + } +} + +void RedisMasterMetadataStore::WipeNodeExtKvMulti(const std::string& node_id) { + // Best-effort per shard: a down shard must not block wiping the reachable ones. + // The node record is already gone/EXPIRED on the control instance, so any extkv + // entry lingering on an unreachable shard points at a dead node and is filtered + // out of matches by GetAlivePeerView until that shard returns. Members of the + // per-(node, shard) reverse index are full extkv keys, so the script HDELs them + // directly. Idempotent. + for (std::size_t shard = 0; shard < keys_.NumShards(); ++shard) { + try { + RespValue r = client_for_shard(shard).Eval(redis::kUnregisterExternalKvByNodeLua, + {keys_.ExtKvNode(node_id, shard)}, {node_id}); + if (r.is_error()) { + MORI_UMBP_WARN("[RedisStore] WipeNodeExtKv: shard {} script error, skipped: {}", shard, + r.str); + } + } catch (const std::exception& e) { + MORI_UMBP_WARN("[RedisStore] WipeNodeExtKv: shard {} unavailable, skipped: {}", shard, + e.what()); + } + } +} + +// ===================================================================== +// Cross-store writes +// ===================================================================== + +bool RedisMasterMetadataStore::RegisterClient(const ClientRegistration& registration, + std::chrono::system_clock::time_point now, + std::chrono::system_clock::duration stale_after) { + ScopedStoreOp _op(metrics_, "RegisterClient"); + const std::string engine(registration.engine_desc_bytes.begin(), + registration.engine_desc_bytes.end()); + RespValue r = control().Eval( + redis::kRegisterClientLua, {keys_.Node(registration.node_id)}, + {keys_.Tag(), registration.node_id, std::to_string(ToEpochMs(now)), + std::to_string(ToMs(stale_after)), registration.node_address, registration.peer_address, + EncodeCaps(registration.tier_capacities), engine, JoinTags(registration.tags)}); + if (r.is_error()) throw std::runtime_error("[RedisStore] RegisterClient: " + r.str); + return r.integer == 1; +} + +void RedisMasterMetadataStore::UnregisterClient(const std::string& node_id) { + ScopedStoreOp _op(metrics_, "UnregisterClient"); + // wipe_extkv=1 wipes the node's external-kv inline in the control/single script + // (num_shards==1: extkv lives on the control slot). When extkv is sharded it is + // wiped per shard afterwards (WipeNodeExtKvMulti); the inline step is skipped. + const std::string wipe_extkv = extkv_sharded() ? "0" : "1"; + if (!split_writes()) { + RespValue r = control().Eval(redis::kUnregisterClientLua, {keys_.Node(node_id)}, + {keys_.Tag(), node_id, wipe_extkv}); + if (r.is_error()) throw std::runtime_error("[RedisStore] UnregisterClient: " + r.str); + if (extkv_sharded()) WipeNodeExtKvMulti(node_id); + return; + } + // Multi-endpoint: control record first (so the router stops routing to it), + // then wipe its block locations on every shard's instance. A lingering + // location for the now-gone node is filtered out by GetAlivePeerView, and the + // wipe is idempotent, so a mid-way failure + retry is safe. + RespValue r = control().Eval(redis::kUnregisterControlLua, {keys_.Node(node_id)}, + {keys_.Tag(), node_id, wipe_extkv}); + if (r.is_error()) throw std::runtime_error("[RedisStore] UnregisterClient(control): " + r.str); + WipeNodeBlocksMulti(node_id); + if (extkv_sharded()) WipeNodeExtKvMulti(node_id); +} + +HeartbeatResult RedisMasterMetadataStore::ApplyHeartbeat( + const std::string& node_id, uint64_t seq, std::chrono::system_clock::time_point now, + const std::map& caps, const std::vector& events, + bool is_full_sync) { + ScopedStoreOp _op(metrics_, "ApplyHeartbeat"); + RespValue r; + if (!split_writes()) { + // Single instance: one atomic script does seq-CAS + record + blocks. + std::vector args; + args.reserve(7 + events.size() * 4); + args.push_back(keys_.Tag()); + args.push_back(node_id); + args.push_back(std::to_string(seq)); + args.push_back(std::to_string(ToEpochMs(now))); + args.push_back(is_full_sync ? "1" : "0"); + args.push_back(EncodeCaps(caps)); + args.push_back(std::to_string(events.size())); + for (const auto& ev : events) { + args.push_back(ev.kind == KvEvent::Kind::ADD ? "0" : "1"); + // Pass the fully composed (sharded) block key so the Lua script never has + // to know the shard layout; it also becomes the reverse-index member. + args.push_back(keys_.Block(ev.key)); + args.push_back(std::to_string(static_cast(ev.tier))); + args.push_back(std::to_string(ev.size)); + } + r = control().Eval(redis::kApplyHeartbeatLua, {keys_.Node(node_id)}, args); + } else { + // Multi-endpoint: control step (seq-CAS + record + alive/peers) only; block + // events are applied per shard afterwards if this heartbeat is APPLIED. + r = control().Eval(redis::kApplyHeartbeatControlLua, {keys_.Node(node_id)}, + {keys_.Tag(), node_id, std::to_string(seq), std::to_string(ToEpochMs(now)), + is_full_sync ? "1" : "0", EncodeCaps(caps)}); + } + + if (r.is_error()) throw std::runtime_error("[RedisStore] ApplyHeartbeat: " + r.str); + if (!r.is_array() || r.elements.size() < 2) { + throw std::runtime_error("[RedisStore] ApplyHeartbeat: malformed reply"); + } + const std::string& status = r.elements[0].str; + uint64_t acked = 0; + try { + acked = std::stoull(r.elements[1].str); + } catch (...) { + } + if (status == "UNKNOWN") return HeartbeatResult{HeartbeatResult::UNKNOWN, 0}; + if (status == "SEQ_GAP") return HeartbeatResult{HeartbeatResult::SEQ_GAP, acked}; + + // APPLIED. In multi-endpoint mode the control record has advanced (so the node + // is already marked ALIVE and won't be reaped); now apply the block events on + // each shard's instance. If a shard is down, don't fail the RPC — return + // SEQ_GAP so the peer full_syncs and self-heals via the existing recovery path + // once the shard is back (block ops are idempotent, so the replay is safe). + if (split_writes()) { + try { + ApplyBlockEventsMulti(node_id, events, is_full_sync, now); + } catch (const std::exception& e) { + MORI_UMBP_WARN( + "[RedisStore] ApplyHeartbeat: block apply degraded for node {} (a shard is " + "unavailable); requesting full_sync to self-heal: {}", + node_id, e.what()); + return HeartbeatResult{HeartbeatResult::SEQ_GAP, acked}; + } + } + return HeartbeatResult{HeartbeatResult::APPLIED, acked}; +} + +std::vector RedisMasterMetadataStore::ExpireStaleClients( + std::chrono::system_clock::time_point cutoff) { + ScopedStoreOp _op(metrics_, "ExpireStaleClients"); + // Bind a reference (not const char*) so the SHA cache's pointer-identity key + // (&script) stays stable — see lua_scripts.h. + const std::string& script = split_writes() ? redis::kExpireControlLua : redis::kExpireStaleLua; + const std::string wipe_extkv = extkv_sharded() ? "0" : "1"; + // KEYS[1] = a control-tag key (nodes:alive) to route + fix the slot in cluster + // mode; the script reads tag from ARGV and touches only control-tag keys. + RespValue r = control().Eval(script, {keys_.NodesAlive()}, + {keys_.Tag(), std::to_string(ToEpochMs(cutoff)), wipe_extkv}); + if (r.is_error()) throw std::runtime_error("[RedisStore] ExpireStaleClients: " + r.str); + std::vector dead; + if (r.is_array()) { + dead.reserve(r.elements.size()); + for (const auto& e : r.elements) dead.push_back(e.str); + } + // Split-write modes: the control step only flipped status + returned the dead + // ids; wipe each dead node's block locations on every shard. + if (split_writes()) { + for (const auto& id : dead) WipeNodeBlocksMulti(id); + } + // Sharded extkv: wipe each dead node's external-kv per shard (the control step + // skipped it when wipe_extkv==0). + if (extkv_sharded()) { + for (const auto& id : dead) WipeNodeExtKvMulti(id); + } + return dead; +} + +// ===================================================================== +// External-KV writes. The extkv/hit keyspace is sharded by ExtKvShardOf(hash), +// so every op fans out one single-slot script per touched shard (grouped like +// the block read hot path). For num_shards == 1 this collapses onto the control +// slot — one atomic script, byte-identical to the legacy layout. See design +// §4/§5 + the PHASE 2 note in lua_scripts.h. +// ===================================================================== + +bool RedisMasterMetadataStore::RegisterExternalKvIfAlive(const std::string& node_id, + const std::vector& hashes, + TierType tier) { + ScopedStoreOp _op(metrics_, "RegisterExternalKvIfAlive"); + if (hashes.empty()) return IsClientAlive(node_id); + const std::string tier_str = std::to_string(static_cast(tier)); + + if (!extkv_sharded()) { + // num_shards == 1: alive-gate + write in one atomic script (KEYS declared, no + // allow-undeclared-keys flag — Dragonfly-single would still run it global, but + // that path is num_shards>1). extkv reverse index + extkv keys all on shard 0 + // (== control slot). + std::vector keys; + keys.reserve(2 + hashes.size()); + keys.push_back(keys_.Node(node_id)); // KEYS[1] alive-check + keys.push_back(keys_.ExtKvNode(node_id, 0)); // KEYS[2] reverse index + for (const auto& h : hashes) keys.push_back(keys_.ExtKv(h)); // KEYS[3..] + RespValue r = control().Eval(redis::kRegisterExternalKvLua, keys, {node_id, tier_str}); + if (r.is_error()) throw std::runtime_error("[RedisStore] RegisterExternalKvIfAlive: " + r.str); + return r.integer == 1; + } + + // Sharded: alive-check on the control instance first, then fan out the write to + // each touched shard. The check + write are not one atomic step across shards + // (the node key and the sharded extkv keys live in different slots), a TOCTOU + // window consistent with the split block-write path — a node that dies between + // the check and the write leaves extkv entries the unregister/expire cascade + // reaps. Idempotent, so a retried shard is safe. + if (!IsClientAlive(node_id)) return false; + + // Group hashes by shard; each group's KEYS = [reverse index] ++ [extkv keys]. + redis::ShardedBatch batch; + std::vector group_of_shard(keys_.NumShards(), -1); + for (const auto& h : hashes) { + const std::size_t shard = keys_.ExtKvShardOf(h); + int& g = group_of_shard[shard]; + if (g < 0) { + g = static_cast(batch.keys_by_shard.size()); + batch.keys_by_shard.emplace_back(1, keys_.ExtKvNode(node_id, shard)); // KEYS[1] + batch.orig_index_by_shard.emplace_back(); + batch.shard_of_group.push_back(shard); + } + batch.keys_by_shard[g].push_back(keys_.ExtKv(h)); + } + redis::RunShardedRead( + batch, redis::kRegisterExternalKvWriteLua, {node_id, tier_str}, "RegisterExternalKvIfAlive", + fanout_pool_.get(), /*tolerate_shard_failures=*/false, metrics_, + [&](size_t g) { return &client_for_shard(batch.shard_of_group[g]); }, + [](size_t, const RespValue&) {}); + return true; +} + +void RedisMasterMetadataStore::UnregisterExternalKv(const std::string& node_id, + const std::vector& hashes, + TierType tier) { + ScopedStoreOp _op(metrics_, "UnregisterExternalKv"); + if (hashes.empty()) return; + const std::string tier_str = std::to_string(static_cast(tier)); + + // Group hashes by shard; each group's KEYS = [reverse index] ++ [extkv keys]. + // num_shards == 1 => one group on the control slot (one atomic script). + redis::ShardedBatch batch; + std::vector group_of_shard(keys_.NumShards(), -1); + for (const auto& h : hashes) { + const std::size_t shard = keys_.ExtKvShardOf(h); + int& g = group_of_shard[shard]; + if (g < 0) { + g = static_cast(batch.keys_by_shard.size()); + batch.keys_by_shard.emplace_back(1, keys_.ExtKvNode(node_id, shard)); + batch.orig_index_by_shard.emplace_back(); + batch.shard_of_group.push_back(shard); + } + batch.keys_by_shard[g].push_back(keys_.ExtKv(h)); + } + redis::RunShardedRead( + batch, redis::kUnregisterExternalKvLua, {node_id, tier_str}, "UnregisterExternalKv", + fanout_pool_.get(), /*tolerate_shard_failures=*/false, metrics_, + [&](size_t g) { return &client_for_shard(batch.shard_of_group[g]); }, + [](size_t, const RespValue&) {}); +} + +void RedisMasterMetadataStore::UnregisterExternalKvByTier(const std::string& node_id, + TierType tier) { + ScopedStoreOp _op(metrics_, "UnregisterExternalKvByTier"); + // Enumerates via the per-(node, shard) reverse index, so it must touch every + // shard. One script per shard, routed to that shard. + const std::string tier_str = std::to_string(static_cast(tier)); + for (std::size_t shard = 0; shard < keys_.NumShards(); ++shard) { + RespValue r = client_for_shard(shard).Eval(redis::kUnregisterExternalKvByTierLua, + {keys_.ExtKvNode(node_id, shard)}, + {node_id, tier_str}); + if (r.is_error()) throw std::runtime_error("[RedisStore] UnregisterExternalKvByTier: " + r.str); + } +} + +void RedisMasterMetadataStore::UnregisterExternalKvByNode(const std::string& node_id) { + ScopedStoreOp _op(metrics_, "UnregisterExternalKvByNode"); + // Same per-shard wipe the client-record cascades use. + WipeNodeExtKvMulti(node_id); +} + +std::size_t RedisMasterMetadataStore::GarbageCollectHits( + std::chrono::system_clock::time_point cutoff) { + ScopedStoreOp _op(metrics_, "GarbageCollectHits"); + // Drop every hit counter whose last_seen < cutoff via each shard's hit index + // (no SCAN — cluster-routable by the index key). One script per shard; runs on + // the slow hit-GC timer, not the hot path. num_shards == 1 => one call. + const std::string cutoff_str = std::to_string(ToEpochMs(cutoff)); + std::size_t dropped = 0; + for (std::size_t shard = 0; shard < keys_.NumShards(); ++shard) { + RespValue r = client_for_shard(shard).Eval(redis::kGarbageCollectHitsLua, + {keys_.HitIndex(shard)}, {cutoff_str}); + if (r.is_error()) throw std::runtime_error("[RedisStore] GarbageCollectHits: " + r.str); + if (r.type == RespValue::Type::Integer) dropped += static_cast(r.integer); + } + return dropped; +} + +// ===================================================================== +// Block reads +// ===================================================================== + +std::vector RedisMasterMetadataStore::LookupBlock(const std::string& key) const { + ScopedStoreOp _op(metrics_, "LookupBlock"); + RespValue r = client_for_shard(keys_.ShardOf(key)).Command({"HGETALL", keys_.Block(key)}); + std::vector out; + if (!r.is_array()) return out; + for (size_t i = 0; i + 1 < r.elements.size(); i += 2) { + const std::string& f = r.elements[i].str; + if (f.rfind("l|", 0) != 0) continue; + const std::string rest = f.substr(2); + const size_t sep = rest.find('|'); + if (sep == std::string::npos) continue; + Location loc; + loc.node_id = rest.substr(0, sep); + try { + loc.tier = static_cast(std::stoi(rest.substr(sep + 1))); + loc.size = std::stoull(r.elements[i + 1].str); + } catch (...) { + continue; + } + out.push_back(std::move(loc)); + } + return out; +} + +std::vector RedisMasterMetadataStore::LookupBlockForRouteGet( + const std::string& key, const std::unordered_set& exclude_nodes, + std::chrono::system_clock::time_point now, std::chrono::system_clock::duration lease_duration) { + auto batch = BatchLookupBlockForRouteGet({key}, exclude_nodes, now, lease_duration); + return batch.empty() ? std::vector{} : std::move(batch.front()); +} + +std::vector> RedisMasterMetadataStore::BatchLookupBlockForRouteGet( + const std::vector& keys, const std::unordered_set& exclude_nodes, + std::chrono::system_clock::time_point now, std::chrono::system_clock::duration lease_duration) { + std::vector> out(keys.size()); + if (keys.empty()) return out; + ScopedStoreOp _op(metrics_, "BatchLookupBlockForRouteGet"); + + std::vector args; + args.reserve(3 + exclude_nodes.size()); + args.push_back(std::to_string(ToEpochMs(now))); + args.push_back(std::to_string(ToMs(lease_duration))); + args.push_back(std::to_string(exclude_nodes.size())); + for (const auto& n : exclude_nodes) args.push_back(n); + + // Fan out one single-slot route_get_batch per shard; groups on the same + // instance share one pipeline, groups on different instances run against + // their own instance. Each key's reply is a flat [node, size, tier, ...] list. + const redis::ShardedBatch batch = redis::GroupKeysByShard(keys_, keys); + redis::RunShardedRead( + batch, redis::kRouteGetBatchLua, args, "BatchLookupBlockForRouteGet", fanout_pool_.get(), + /*tolerate_shard_failures=*/mode_ != Mode::kSingle, metrics_, + [&](size_t group) { return &client_for_shard(batch.shard_of_group[group]); }, + [&](size_t orig_index, const RespValue& locs) { + if (!locs.is_array()) return; + for (size_t j = 0; j + 3 <= locs.elements.size(); j += 3) { + Location loc; + loc.node_id = locs.elements[j].str; + try { + loc.size = std::stoull(locs.elements[j + 1].str); + loc.tier = static_cast(std::stoi(locs.elements[j + 2].str)); + } catch (...) { + continue; + } + out[orig_index].push_back(std::move(loc)); + } + }); + return out; +} + +std::vector RedisMasterMetadataStore::BatchExistsBlock( + const std::vector& keys) const { + std::vector results(keys.size(), false); + if (keys.empty()) return results; + ScopedStoreOp _op(metrics_, "BatchExistsBlock"); + + const redis::ShardedBatch batch = redis::GroupKeysByShard(keys_, keys); + redis::RunShardedRead( + batch, redis::kExistsBatchLua, {}, "BatchExistsBlock", fanout_pool_.get(), + /*tolerate_shard_failures=*/mode_ != Mode::kSingle, metrics_, + [&](size_t group) { return &client_for_shard(batch.shard_of_group[group]); }, + [&](size_t orig_index, const RespValue& has) { results[orig_index] = has.integer != 0; }); + return results; +} + +std::map> +RedisMasterMetadataStore::EnumerateEvictionCandidates( + const std::vector& buckets, EvictionOrder order, size_t max_per_bucket, + std::chrono::system_clock::time_point now) const { + std::map> result; + if (buckets.empty()) return result; + ScopedStoreOp _op(metrics_, "EnumerateEvictionCandidates"); + + // The read hot path is untouched: eviction reuses the per-node block reverse + // index + the _lease/_lacc already maintained on each block hash. Candidates + // are enumerated per shard by kEnumerateEvictionLua and aggregated here; the + // ordering / per-bucket cap is applied in C++ (an eviction tick is seconds, + // not a hot path), mirroring the in-memory backend's scan-and-sort. + + // Distinct nodes to scan, and the wanted (node, tier) pairs the script filters + // on (passed as ARGV pairs, so a node id may contain any byte). + std::vector nodes; + { + std::unordered_set seen; + for (const auto& b : buckets) { + if (seen.insert(b.node_id).second) nodes.push_back(b.node_id); + } + } + std::vector shared_args; + shared_args.reserve(2 + buckets.size() * 2); + shared_args.push_back(std::to_string(ToEpochMs(now))); + shared_args.push_back(std::to_string(buckets.size())); + for (const auto& b : buckets) { + shared_args.push_back(b.node_id); + shared_args.push_back(std::to_string(static_cast(b.tier))); + } + + // Build the per-shard groups. split_writes() (multi-endpoint / cluster) keeps a + // reverse index per (node, shard) on that shard's instance/slot, so one group + // per shard routed to that shard. Single mode keeps one control-tag reverse + // index per node, so one group on the control instance. + redis::ShardedBatch batch; + auto add_group = [&](std::size_t shard, std::vector node_block_keys) { + const std::size_t g = batch.keys_by_shard.size(); + batch.orig_index_by_shard.emplace_back(); + for (std::size_t j = 0; j < node_block_keys.size(); ++j) { + batch.orig_index_by_shard[g].push_back(j); // unused (node is in each tuple) + } + batch.keys_by_shard.push_back(std::move(node_block_keys)); + batch.shard_of_group.push_back(shard); + }; + if (split_writes()) { + for (std::size_t s = 0; s < keys_.NumShards(); ++s) { + std::vector ks; + ks.reserve(nodes.size()); + for (const auto& n : nodes) ks.push_back(keys_.NodeBlocks(n, s)); + add_group(s, std::move(ks)); + } + } else { + std::vector ks; + ks.reserve(nodes.size()); + for (const auto& n : nodes) ks.push_back(keys_.NodeBlocks(n)); + add_group(/*shard=*/0, std::move(ks)); + } + + // Strip ":block:" from a redis block key to recover the user key. + auto user_key_of = [](const std::string& block_key) -> std::string { + const size_t pos = block_key.find(":block:"); + return pos == std::string::npos ? block_key : block_key.substr(pos + 7); + }; + + redis::RunShardedRead( + batch, redis::kEnumerateEvictionLua, shared_args, "EnumerateEvictionCandidates", + fanout_pool_.get(), /*tolerate_shard_failures=*/mode_ != Mode::kSingle, metrics_, + [&](size_t group) { return &client_for_shard(batch.shard_of_group[group]); }, + [&](size_t /*orig_index*/, const RespValue& cands) { + if (!cands.is_array()) return; + // Flat 5-tuples: [block_key, node, tier, size, last_accessed_ms]. + for (size_t j = 0; j + 5 <= cands.elements.size(); j += 5) { + EvictionCandidate c; + c.key = user_key_of(cands.elements[j].str); + c.location.node_id = cands.elements[j + 1].str; + try { + c.location.tier = static_cast(std::stoi(cands.elements[j + 2].str)); + c.location.size = std::stoull(cands.elements[j + 3].str); + c.last_accessed_at = FromEpochMs(std::stoll(cands.elements[j + 4].str)); + } catch (...) { + continue; + } + c.size = c.location.size; + result[NodeTierKey{c.location.node_id, c.location.tier}].push_back(std::move(c)); + } + }); + + // Honor the ordering hint and the per-bucket cap (same policy as in-memory). + const auto older_first = [](const EvictionCandidate& a, const EvictionCandidate& b) { + return a.last_accessed_at < b.last_accessed_at; + }; + for (auto& [ntk, candidates] : result) { + (void)ntk; + const bool cap = max_per_bucket > 0 && candidates.size() > max_per_bucket; + if (order == EvictionOrder::kLeastRecentlyAccessed) { + if (cap) { + std::partial_sort(candidates.begin(), candidates.begin() + max_per_bucket, candidates.end(), + older_first); + candidates.resize(max_per_bucket); + } else { + std::sort(candidates.begin(), candidates.end(), older_first); + } + } else if (cap) { + candidates.resize(max_per_bucket); + } + } + return result; +} + +// ===================================================================== +// Client reads +// ===================================================================== + +std::optional RedisMasterMetadataStore::GetClient(const std::string& node_id) const { + RespValue r = control().Command({"HGETALL", keys_.Node(node_id)}); + if (!r.is_array() || r.elements.empty()) return std::nullopt; + return DecodeRecord(node_id, FlatToMap(r)); +} + +bool RedisMasterMetadataStore::IsClientAlive(const std::string& node_id) const { + RespValue r = control().Command({"HGET", keys_.Node(node_id), "status"}); + return r.type == RespValue::Type::String && r.str == "1"; +} + +std::optional RedisMasterMetadataStore::GetPeerAddress( + const std::string& node_id) const { + RespValue r = control().Command({"HGET", keys_.Node(node_id), "peer"}); + if (r.is_nil()) return std::nullopt; + return r.str; +} + +std::vector RedisMasterMetadataStore::ListAliveClients() const { + ScopedStoreOp _op(metrics_, "ListAliveClients"); + // KEYS[1] = a control-tag key (nodes:alive) to route + fix the slot in cluster + // mode; the script reads tag from ARGV and touches only control-tag keys. + RespValue r = control().Eval(redis::kListAliveLua, {keys_.NodesAlive()}, {keys_.Tag()}); + if (r.is_error()) throw std::runtime_error("[RedisStore] ListAliveClients: " + r.str); + std::vector out; + if (!r.is_array()) return out; + out.reserve(r.elements.size()); + for (const auto& entry : r.elements) { + if (!entry.is_array() || entry.elements.size() < 2) continue; + const std::string& id = entry.elements[0].str; + out.push_back(DecodeRecord(id, FlatToMap(entry.elements[1]))); + } + return out; +} + +std::unordered_map RedisMasterMetadataStore::GetAlivePeerView() const { + ScopedStoreOp _op(metrics_, "GetAlivePeerView"); + RespValue r = control().Command({"HGETALL", keys_.AlivePeers()}); + std::unordered_map view; + if (!r.is_array()) return view; + for (size_t i = 0; i + 1 < r.elements.size(); i += 2) { + view.emplace(r.elements[i].str, r.elements[i + 1].str); + } + return view; +} + +std::size_t RedisMasterMetadataStore::AliveClientCount() const { + RespValue r = control().Command({"SCARD", keys_.NodesAlive()}); + return r.type == RespValue::Type::Integer ? static_cast(r.integer) : 0; +} + +std::vector RedisMasterMetadataStore::GetClientTags(const std::string& node_id) const { + RespValue r = control().Command({"HGET", keys_.Node(node_id), "tags"}); + if (r.type != RespValue::Type::String) return {}; + return SplitTags(r.str); +} + +// ===================================================================== +// External-KV reads. The extkv/hit keyspace is sharded by ExtKvShardOf(hash), so +// a batch fans out one single-slot script per touched shard (grouped + scattered +// by RunShardedRead like the block read hot path). num_shards == 1 collapses to +// one group on the control instance. +// ===================================================================== + +std::vector RedisMasterMetadataStore::MatchExternalKv( + const std::vector& hashes, bool count_as_hit, + std::chrono::system_clock::time_point now) { + ScopedStoreOp _op(metrics_, "MatchExternalKv"); + std::vector result; + if (hashes.empty()) return result; + + // Dedup preserving first-seen order (lets each touched key be declared in + // KEYS[] — no allow-undeclared-keys — and lets a per-hash reply be scattered + // back by first-seen position). + const std::vector uniq = DedupPreserveOrder(hashes); + + // Group unique hashes by their extkv shard. Each group's KEYS starts as the k + // extkv keys (reply is parallel to these, scattered back by orig index); when + // counting hits, the k hit keys + the shard's hit index are appended so the + // script can bump the counter in the same single-slot call. + redis::ShardedBatch batch; + std::vector group_of_shard(keys_.NumShards(), -1); + for (std::size_t i = 0; i < uniq.size(); ++i) { + const std::size_t shard = keys_.ExtKvShardOf(uniq[i]); + int& g = group_of_shard[shard]; + if (g < 0) { + g = static_cast(batch.keys_by_shard.size()); + batch.keys_by_shard.emplace_back(); + batch.orig_index_by_shard.emplace_back(); + batch.shard_of_group.push_back(shard); + } + batch.keys_by_shard[g].push_back(keys_.ExtKv(uniq[i])); + batch.orig_index_by_shard[g].push_back(i); + } + if (count_as_hit) { + for (std::size_t g = 0; g < batch.keys_by_shard.size(); ++g) { + auto& ks = batch.keys_by_shard[g]; + const auto& idx = batch.orig_index_by_shard[g]; + const std::size_t kg = idx.size(); + ks.reserve(2 * kg + 1); + for (std::size_t j = 0; j < kg; ++j) ks.push_back(keys_.Hit(uniq[idx[j]])); + ks.push_back(keys_.HitIndex(batch.shard_of_group[g])); + } + } + + const std::vector shared_args = {count_as_hit ? "1" : "0", + std::to_string(ToEpochMs(now))}; + + // Group by node, decoding each node's tier bitmask (bit == 1<>> acc; + redis::RunShardedRead( + batch, redis::kMatchExternalKvLua, shared_args, "MatchExternalKv", fanout_pool_.get(), + /*tolerate_shard_failures=*/mode_ != Mode::kSingle, metrics_, + [&](size_t g) { return &client_for_shard(batch.shard_of_group[g]); }, + [&](size_t orig_index, const RespValue& flat) { + if (!flat.is_array()) return; + const std::string& hash = uniq[orig_index]; + for (size_t i = 0; i + 1 < flat.elements.size(); i += 2) { + const std::string& node = flat.elements[i].str; + long long mask = 0; + try { + mask = std::stoll(flat.elements[i + 1].str); + } catch (...) { + continue; + } + auto& by_tier = acc[node]; + for (int t = 0; t < 16; ++t) { + if ((mask >> t) & 1) by_tier[static_cast(t)].push_back(hash); + } + } + }); + result.reserve(acc.size()); + for (auto& [node_id, by_tier] : acc) { + NodeMatch m; + m.node_id = node_id; + m.hashes_by_tier = std::move(by_tier); + result.push_back(std::move(m)); + } + return result; +} + +std::vector RedisMasterMetadataStore::GetExternalKvHitCounts( + const std::vector& hashes) const { + std::vector out; + if (hashes.empty()) return out; + const std::vector uniq = DedupPreserveOrder(hashes); + + // Group hit keys by shard; reply per shard is parallel to that shard's hit + // keys, scattered back by orig index. + redis::ShardedBatch batch; + std::vector group_of_shard(keys_.NumShards(), -1); + for (std::size_t i = 0; i < uniq.size(); ++i) { + const std::size_t shard = keys_.ExtKvShardOf(uniq[i]); + int& g = group_of_shard[shard]; + if (g < 0) { + g = static_cast(batch.keys_by_shard.size()); + batch.keys_by_shard.emplace_back(); + batch.orig_index_by_shard.emplace_back(); + batch.shard_of_group.push_back(shard); + } + batch.keys_by_shard[g].push_back(keys_.Hit(uniq[i])); + batch.orig_index_by_shard[g].push_back(i); + } + + std::vector scratch(uniq.size()); + std::vector present(uniq.size(), false); + redis::RunShardedRead( + batch, redis::kGetHitCountsLua, {}, "GetExternalKvHitCounts", fanout_pool_.get(), + /*tolerate_shard_failures=*/mode_ != Mode::kSingle, metrics_, + [&](size_t g) { return &client_for_shard(batch.shard_of_group[g]); }, + [&](size_t orig_index, const RespValue& c) { + if (c.type != RespValue::Type::String) return; + try { + scratch[orig_index].hit_count_total = std::stoull(c.str); + } catch (...) { + return; + } + scratch[orig_index].hash = uniq[orig_index]; + present[orig_index] = true; + }); + out.reserve(uniq.size()); + for (std::size_t i = 0; i < uniq.size(); ++i) { + if (present[i]) out.push_back(std::move(scratch[i])); + } + return out; +} + +std::size_t RedisMasterMetadataStore::GetExternalKvCount(const std::string& node_id) const { + // Sum the per-(node, shard) reverse-index cardinalities. num_shards == 1 => one + // SCARD on the control slot (byte-identical to the legacy single index). + std::size_t total = 0; + for (std::size_t shard = 0; shard < keys_.NumShards(); ++shard) { + RespValue r = client_for_shard(shard).Command({"SCARD", keys_.ExtKvNode(node_id, shard)}); + if (r.type == RespValue::Type::Integer) total += static_cast(r.integer); + } + return total; +} + +} // namespace mori::umbp diff --git a/src/umbp/doc/design-redis-metadata-store.md b/src/umbp/doc/design-redis-metadata-store.md new file mode 100644 index 000000000..31d90f499 --- /dev/null +++ b/src/umbp/doc/design-redis-metadata-store.md @@ -0,0 +1,678 @@ +# UMBP Master — Redis / Dragonfly Metadata Store Design + +**Scope:** A RESP-protocol-compatible backend for +`IMasterMetadataStore` (`src/umbp/include/umbp/distributed/master/master_metadata_store.h`) +that makes the UMBP master **stateless**: all client/block/external-KV/hit +metadata lives in an external RESP store (Redis, Dragonfly, or Valkey), and +master processes only accept RPCs, query the store, and return data. Fault +tolerance is delegated to the store's replication + persistence. + +**Companion doc:** [`design-master-control-plane.md`](./design-master-control-plane.md) +is authoritative for the master's RPC contract, router, eviction, and the +`InMemoryMasterMetadataStore` (the single-master / unit-test backend). This +doc only covers the Redis backend of the same interface. + +**Status:** Phase 1 delivered — RESP client seam (hiredis), key schema, Lua +hot-path scripts, `RedisMasterMetadataStore` (six hot methods + the methods the +running master needs), the `MakeMasterMetadataStore` factory, and a +backend-selectable microbench. Conformance passes on Redis and Dragonfly. See +the go/no-go results in +[`redis-backend-phase1-bench.md`](./redis-backend-phase1-bench.md); Phase 2 +directions are in that report's §5. + +Phase 2 external-KV / hit-count / eviction (#5) is now delivered too: the +external-KV read/write family, `MatchExternalKv` hit accounting, +`GarbageCollectHits`, and `EnumerateEvictionCandidates` are implemented against +real Lua and pass the parameterized conformance suite in all modes +(single / multi-endpoint / Redis Cluster; Dragonfly). The read hot path +(`route_get_batch`) is unchanged, so read throughput does not regress. This +unblocks a real hicache PD (`kv_events_subscriber=true`) on the Redis backend. + +--- + +## 1. Goals and non-goals + +### Goals + +1. **Stateless master.** Any number of master replicas can serve traffic + concurrently; each holds no durable state. A crashed master restarts and + reads the full picture back from the store. This is the split-brain fix + the interface header already calls out: today `GlobalBlockIndex`, + `ClientRegistry`, `ExternalKvBlockIndex`, `ExternalKvHitIndex` are + in-process `unordered_map`s, so a second master replica immediately + diverges. +2. **RESP + Lua common subset.** One implementation, swappable across + Redis / Dragonfly / Valkey by connection config only. We program to the + intersection of what all three support over the RESP wire: strings / + hashes / sets / sorted-sets, `MULTI`-free single-key atomics, pipelines, + and server-side `EVAL` / `EVALSHA` Lua. +3. **Single round trip per RPC.** Every hot method is one network round trip + (a pipeline or a single Lua `EVALSHA`); no N-round-trip loops. The + interface was designed for this (`Batch*` methods, and the + "implementations SHOULD make each read a single backend round-trip" + contract). +4. **Fault tolerance via the store.** AOF/replication + failover in the store + is the durability and HA mechanism. The master adds reconnect/backoff and + fails RPCs cleanly (`UNAVAILABLE`) rather than serving stale local data. + +### Non-goals + +- No change to the master RPC contract, router, eviction policy, or peer. + The only production wiring change is the store construction line in + `master_server.cpp`. +- No KV data-plane involvement. Pages stay on peers; the master (and thus the + store) holds metadata only. +- No SQL backend. The read path does one store round-trip per RouteGet with a + Lua lookup+lease+access; that is fine on Redis and ruinous on an OLTP DB. + +--- + +## 2. Where it plugs in + +`MasterServer` owns a single `std::unique_ptr` and hands +references to the router, service, and eviction manager. Today it is +hard-wired to the in-memory impl: + +```774:777:src/umbp/distributed/master/master_server.cpp + : config_(std::move(config)), + store_(std::make_unique()), + router_(*store_, std::move(config_.get_strategy), std::move(config_.put_strategy)), + service_(std::make_unique(*store_, router_, config_.registry_config, +``` + +The only production change: replace the `make_unique()` +with a factory `MakeMasterMetadataStore(config)` that returns either the +in-memory or the Redis impl based on `UMBP_METADATA_BACKEND`. Everything +downstream (`router_`, `service_`, `eviction_manager_`, reaper, hit-GC loops) +is untouched because they only see `IMasterMetadataStore&`. + +```mermaid +flowchart LR + subgraph masters [Stateless master replicas] + M1[master #1] + M2[master #2] + M3[master #3] + end + Peers[Peers / PoolClients] -->|gRPC RouteGet/Put/Heartbeat| masters + M1 -->|RESP EVALSHA / pipeline| Store + M2 -->|RESP| Store + M3 -->|RESP| Store + subgraph Store [RESP store: Redis / Dragonfly / Valkey] + primary[(primary)] + replica[(replica)] + primary -.->|replication + AOF| replica + end +``` + +### The hot path is narrow + +`router.cpp` uses only these store methods on the RouteGet / RoutePut path: + +```73:93:src/umbp/distributed/routing/router.cpp + auto exists_mask = store_.BatchExistsBlock(keys); + auto candidates = store_.ListAliveClients(); + ... + auto node_to_peer = store_.GetAlivePeerView(); + ... + auto all_locs = store_.BatchLookupBlockForRouteGet( +``` + +Plus the write hot path `ApplyHeartbeat` and the (cold-ish) `RegisterClient`. +These six methods decide the whole backend's performance and are the entire +Phase 1 scope: + +- `BatchExistsBlock` (RoutePut dedup) +- `ListAliveClients` (RoutePut candidate set) +- `GetAlivePeerView` (RouteGet node -> peer_address projection) +- `BatchLookupBlockForRouteGet` (RouteGet lookup + lease + access, batched) +- `ApplyHeartbeat` (index update, seq-CAS) +- `RegisterClient` + +--- + +## 3. RESP compatibility strategy + +We do **not** couple to any redis-server-only feature. The portable substrate +is: + +- Data types: `STRING`, `HASH`, `SET`, `ZSET`. +- Commands: `HSET/HGET/HGETALL/HDEL/HINCRBY`, `SADD/SREM/SMEMBERS/SCARD`, + `ZADD/ZRANGEBYSCORE/ZREM`, `EXISTS`, `DEL`, `EXPIRE`, plus `EVAL/EVALSHA/SCRIPT LOAD`. +- Pipelining for the pure-read fan-outs (`ListAliveClients`, `GetAlivePeerView`). +- Server-side Lua for every multi-key atomic mutation. + +### Atomicity contract (the key tension) + +The interface header requires several methods to be atomic across former store +boundaries (`ApplyHeartbeat`, `UnregisterClient`, `ExpireStaleClients`, +`RegisterExternalKvIfAlive`, and the hit-counting branch of `MatchExternalKv`). +We satisfy this with **one Lua script per such method**, and we make the script +cross-key-atomic by co-locating all keys in **one hash slot** via a shared hash +tag (see [KeySchema](#4-keyschema)). Then: + +| Deployment | Cross-key Lua atomicity | Notes | +| --- | --- | --- | +| Redis single node | Yes (single-threaded) | Baseline; strongest guarantee | +| Redis Cluster | Yes, **iff** all keys share one slot | Hash tag guarantees this; a stray cross-slot key -> `CROSSSLOT` error | +| Valkey | Same as Redis | RESP + Lua semantics match Redis | +| Dragonfly | Yes within a script over same-slot keys | Multi-threaded engine; Lua runs atomically, but we still verify with the conformance race tests, not by assumption | + +Design rule: **all keys touched by one script share the deployment hash tag**, +so every script runs in a single slot and stays atomic on all four +deployments. This is why we do NOT split the interface into three stores; a +single tagged keyspace lets one script mutate node + block + extkv together. + +Determinism: scripts never call `redis.call('TIME')` or `randomkey`; every +timestamp is passed in from the caller as `now_ms` (`system_clock` epoch +milliseconds; see hazard #7 in the interface header). This keeps scripts +replication-safe regardless of the store's replication mode. + +--- + +## 4. KeySchema + +Two hash-tag families, split so block lookups can scale past one slot: + +- **Control tag** `H = {umbp:}` (`` from + `UMBP_REDIS_NAMESPACE`, default `default`): client records, ALIVE set, peer + projection, and the per-node reverse indexes all co-locate here so one Lua + script mutates them atomically. +- **Block-shard tags** `Bs = {umbp::bS}` for `S in [0, N)` where + `N = UMBP_REDIS_BLOCK_SHARDS`: a block key is placed under the shard chosen by + a stable hash (FNV-1a) of the user key. Spreading blocks over `N` slots is + what lets `BatchLookupBlockForRouteGet` / `BatchExistsBlock` run one + single-slot script per shard instead of piling every lookup onto one slot / + one server thread. **`N == 1` collapses `Bs` back onto `H`, so the emitted key + strings are byte-identical to the original single-tag schema.** + +| Purpose | Key | Type | Fields / members | +| --- | --- | --- | --- | +| Client record | `H:node:` | HASH | `addr`, `peer`, `status` (1=ALIVE,2=EXPIRED), `last_hb`, `reg_at`, `seq`, `caps`, `engine`, `tags` | +| Alive membership | `H:nodes:alive` | SET | `` for every ALIVE node | +| Alive peer projection | `H:alive_peers` | HASH | `` -> `peer_address` (ALIVE only) | +| Block locations | `Bs:block:` | HASH | `l\|\|` -> `size`; meta `_lease`, `_lacc`, `_acnt`, `_created` | +| Node -> its block keys | `H:node::blocks` | SET | full `Bs:block:` strings (reverse index for node-scoped wipe) | +| External-KV entry | `H:extkv:` | HASH | `` -> tier bitmask (bit `1< its extkv hashes | `H:extkv:node:` | SET | `` (reverse index) | +| Hit counter | `H:hit:` | HASH | `c` (count), `ls` (last_seen ms) | +| Hit reverse index | `H:hit:index` | SET | every `` with a live counter (drives GC without SCAN) | + +**External-KV / hit placement.** `extkv:*`, `hit:*`, and `hit:index` all live on the +**control tag** `H`, so `RegisterExternalKvIfAlive` / `MatchExternalKv` (with the +hit-count branch) / `Unregister*` / `GarbageCollectHits` are each one single-slot Lua +on the control instance (control slot in cluster) — no cross-instance step. This adds to +the control-plane load noted in §8; sharding it is a separate design if a match-heavy +workload needs it. The tier-set is a bitmask int (bit `1<:blocks`) plus the `_lease`/`_lacc` already on each block hash** — the read hot +path is untouched — enumerated per shard (one keyed Lua per shard, aggregated by +`RunShardedRead`) with the ordering/cap applied in C++, exactly like the in-memory backend's +scan-and-sort. + +Sharding notes: + +- The per-node reverse index stays a single set on the control tag, but its + **members are now full (already-sharded) block-key strings**. The caller + (`KeySchema::Block`) composes the block key and hands it to the write scripts, + and the wipe scripts (full_sync / unregister / expire) drain those members and + `HDEL` them directly — no shard math in Lua. +- **Atomicity when `N > 1`**: each block key is still mutated / read atomically + in one single-slot script, but a batch read is no longer a whole-batch + snapshot (it is one single-slot script per shard). This matches the in-memory + backend, whose block index is likewise key-hashed shards read under per-shard + locks; the interface only promises per-key semantics. +- **Write scripts vs. Redis Cluster**: `apply_heartbeat` / `unregister` / + `expire` touch the control tag AND block keys across shards in one script, + which is valid on a single instance (non-cluster Redis, or Dragonfly with + `allow-undeclared-keys`) but would `CROSSSLOT` on Redis Cluster. Splitting + those writes into a control script + per-shard block scripts (with idempotent + replay) is the documented next step for cluster / multi-endpoint fan-out; the + read hot path is already single-slot per script. + +### 4.1 Multi-endpoint mode (horizontal scaling) + +A single Redis instance is single-threaded, so splitting block lookups into more +scripts on ONE instance does not raise throughput (it lowers per-script cost but +runs more scripts on the same thread — measured flat/slightly-worse). The only +way past that ceiling is to put block shards on **different Redis instances**. + +`UMBP_REDIS_SHARD_URIS` gives one instance per block shard (shard count = number +of URIs). `clients_[0]` (the first URI) is the **control instance**: it holds +all control-plane keys (`node:`, `nodes:alive`, `alive_peers`, `extkv:`) plus +block shard 0. Block shard `s` lives on `clients_[s]` under tag `{umbp::bs}`, +with its own per-node reverse index `{umbp::bs}:node::blocks` co-located +on that instance. + +Because no Lua script can span instances, the cross-store writes are split: + +| Method | Control instance step | Per-shard-instance step | +| --- | --- | --- | +| `ApplyHeartbeat` | `apply_heartbeat_control` (seq-CAS + record + alive/peers) | `apply_block_events` on each shard with events (full_sync clears every shard) | +| `UnregisterClient` | `unregister_control` (record + alive + peers + extkv) | `wipe_node_blocks` on every shard | +| `ExpireStaleClients` | `expire_control` (flip status, return dead ids) | `wipe_node_blocks` per dead node per shard | + +The control step runs first, then the block steps. Block ADD/REMOVE, full_sync +clear+replay, and node-wipe are all **idempotent**, so a failed/retried block +step is safe; a partial failure surfaces as an exception, the peer retries, the +next heartbeat seq-gaps, and `full_sync` heals the node's locations wholesale. +This is the same "atomic per key/shard, not globally atomic" posture the +in-memory backend already documents. + +Reads (`BatchLookupBlockForRouteGet`, `BatchExistsBlock`) group keys by shard +and issue each instance's `EvalPipeline` **concurrently** through a small +persistent worker pool (`redis/thread_pool.h`), so a batch's latency is one +round trip rather than N — measured ~3.3x RouteGet throughput at 4 instances on +a dedicated host (t8: 5.2k -> 17.3k ops/s, p50 1527us -> 462us), better than the +sequential fan-out (13k) and far better than per-call `std::async` (8k, killed by +thread churn). The pool is sized `min(64, endpoints*8)`; single-endpoint mode +uses no pool and runs inline (zero change). Replies are scattered on the calling +thread, so the per-key decode needs no locking. + +**Read fault tolerance.** In multi-endpoint mode, if one shard instance is down +or errors, the batch read does NOT fail: that shard's keys are left at the +caller's default (a miss — empty locations / `exists=false`) and a WARN is +logged, so RouteGet keeps serving keys on the healthy shards. Single-endpoint +mode still propagates the error (a total outage must surface, not masquerade as +misses). + +**Write fault tolerance.** A down shard no longer aborts writes either. The +control step runs first (seq-CAS + record + alive/peers), so the node is already +ALIVE and won't be reaped. If a per-shard block step then fails, `ApplyHeartbeat` +returns `SEQ_GAP` instead of throwing — the master already turns that into a +`full_sync` request, so the peer re-ships a full snapshot and the node self-heals +once the shard is back (block ops are idempotent, so the replay is safe). +`UnregisterClient` / `ExpireStaleClients` wipe best-effort per shard: a down +shard is skipped (its lingering locations point at a gone node and are filtered +out of reads by `GetAlivePeerView`) and cleaned when it returns. + +Single-endpoint mode (no `UMBP_REDIS_SHARD_URIS`) keeps the original single +atomic scripts unchanged. + +Notes: + +- **Timestamps** are `system_clock` epoch milliseconds (int64), never + `steady_clock`. They cross the process boundary and must be meaningful after + restart and across replicas (hazard #7). +- **`caps`** encodes `std::map` as + `tier:total:avail;tier:total:avail;...`. +- **`tags`** is `\n`-joined (tag values contain `=`, e.g. `sgl_role=prefill`, + so `=`/`,` are unsafe separators; `\n` is not used inside tags). +- **`engine`** stores `engine_desc_bytes` raw (RESP is binary-safe). +- The block hash mixes location fields (`l|node|tier`) and meta fields + (`_`-prefixed) in **one key** so a RouteGet touches exactly one key and the + lookup+lease+access is a single-key atomic in Lua. +- **`alive_peers`** exists so `GetAlivePeerView` is one `HGETALL`, and + `nodes:alive` exists so `AliveClientCount` is one `SCARD`. Both are + projections maintained by the same scripts that flip status. + +--- + +## 5. Method -> RESP/Lua mapping + +Legend: **RT** = round trips. **P** = pipeline. **L** = single Lua `EVALSHA`. + +### 5.1 Hot read path + +| Method | Impl | RT | +| --- | --- | --- | +| `BatchExistsBlock` | Lua: for each `H:block:` return `EXISTS && HLEN(location fields)>0` (a key with only meta fields but no locations is "absent") | 1 (L) | +| `ListAliveClients` | `SMEMBERS H:nodes:alive`, then pipelined `HGETALL H:node:` per member; decode records | 2 (P) | +| `GetAlivePeerView` | `HGETALL H:alive_peers` | 1 | +| `BatchLookupBlockForRouteGet` | Lua `route_get_batch` (see [6.1](#61-route_get_batch)) | 1 (L) | +| `AliveClientCount` | `SCARD H:nodes:alive` | 1 | +| `GetPeerAddress` | `HGET H:node: peer` | 1 | +| `IsClientAlive` | `HGET H:node: status` == 1 | 1 | +| `GetClient` | `HGETALL H:node:` + decode | 1 | +| `GetClientTags` | `HGET H:node: tags` | 1 | +| `LookupBlock` | `HGETALL H:block:`, project location fields (no lease/access) | 1 | + +### 5.2 Hot write path + +| Method | Impl | RT | +| --- | --- | --- | +| `ApplyHeartbeat` | Lua `apply_heartbeat` (seq-CAS + record update + block ADD/REMOVE or full replace + reverse index + `nodes:alive`/`alive_peers` + `lru` maintenance) | 1 (L) | +| `RegisterClient` | Lua `register_client` (TTL-stale/EXPIRED revive CAS + write record + `nodes:alive`/`alive_peers`) | 1 (L) | + +### 5.3 Cascading writes and external-KV + +| Method | Impl | RT | +| --- | --- | --- | +| `UnregisterClient` | Lua: `DEL H:node:`, `SREM nodes:alive`, `HDEL alive_peers`, drain `H:node::blocks` deleting each block's node fields + `H:extkv:node:` reverse wipe | 1 (L) | +| `ExpireStaleClients` | Lua: scan `nodes:alive`, for each with `last_hb < cutoff` flip status EXPIRED, `SREM nodes:alive`, `HDEL alive_peers`, wipe its blocks + extkv; return dead node ids | 1 (L) | +| `RegisterExternalKvIfAlive` | Lua: check `HGET H:node: status`==1; if alive, for each hash set tier bit in `H:extkv:` + `SADD H:extkv:node:` | 1 (L) | +| `UnregisterExternalKv` | Lua: clear tier bit for (node,hash) pairs; drop empty entries + reverse-index members | 1 (L) | +| `UnregisterExternalKvByTier` | Lua over `H:extkv:node:` members | 1 (L) | +| `UnregisterExternalKvByNode` | Lua over `H:extkv:node:`, then `DEL` it | 1 (L) | +| `MatchExternalKv` | Lua `match_external_kv`: for each hash read `H:extkv:`, group by node/tier; if `count_as_hit`, `HINCRBY H:hit: c 1` + set `ls=now` for each unique matched hash | 1 (L) | +| `GetExternalKvHitCounts` | pipelined `HGET H:hit: c` | 1 (P) | +| `GetExternalKvCount` | `SCARD H:extkv:node:` | 1 | +| `GarbageCollectHits` | Lua `garbage_collect_hits`: walk `H:hit:index`, drop each `H:hit:` whose `ls < cutoff` (+`SREM` from the index). One keyed Lua, no `SCAN` | 1 (L) | +| `EnumerateEvictionCandidates` | Lua `enumerate_eviction` per shard: walk each requested node's `node::blocks` set, skip leased (`_lease > now`), emit `(key,node,tier,size,_lacc)` for wanted `(node,tier)`; C++ sorts by `_lacc` + caps. Aggregated across shards by `RunShardedRead`. Read hot path unchanged | 1 per shard | + +--- + +## 6. Hot-path Lua scripts (Phase 1) + +Representative pseudo-Lua; the final scripts are Phase 1 deliverables in +`src/umbp/distributed/master/redis/lua_scripts.h`. All scripts receive the hash +tag prefix as `ARGV[1]` so they can compose auxiliary key names +(`lru`, reverse indexes) that stay in the same slot. + +### 6.1 route_get_batch + +Mirrors `BatchLookupBlockForRouteGet`: locations are returned only for keys +with at least one non-excluded location, and only those keys get a lease + +access bump (fully-excluded / absent keys are untouched, matching the in-memory +impl at `in_memory_master_metadata_store.cpp:490-497`). + +```lua +-- KEYS[1..n] = H:block: +-- ARGV[1]=prefix H, ARGV[2]=now_ms, ARGV[3]=lease_ms, +-- ARGV[4]=n_exclude, ARGV[5..]=exclude node ids, +-- then trailing ARGV = the n raw user keys (for lru members) +local out = {} +for i = 1, #KEYS do + local fields = redis.call('HGETALL', KEYS[i]) -- flat [f,v,f,v,...] + local locs, touched = {}, false + for j = 1, #fields, 2 do + local f, v = fields[j], fields[j+1] + if f:sub(1,2) == 'l|' then -- 'l||' + local node, tier = parse_loc_field(f) + if not excluded(node) then + locs[#locs+1] = { node, tonumber(v), tier } + touched = true + end + end + end + if touched then + redis.call('HSET', KEYS[i], '_lease', now_ms + lease_ms, '_lacc', now_ms) + redis.call('HINCRBY', KEYS[i], '_acnt', 1) + for _, loc in ipairs(locs) do -- maintain per-bucket LRU + redis.call('ZADD', prefix..':lru:'..loc[1]..':'..loc[3], now_ms, userkey_i) + end + end + out[i] = locs +end +return out +``` + +### 6.2 apply_heartbeat + +Mirrors `ApplyHeartbeat` (`in_memory_master_metadata_store.cpp:323-372`), +including the three hazards: seq-CAS (#1), status stays ALIVE on `SEQ_GAP` +without advancing seq/caps, full_sync atomic replace (#payload-sizing bound), +delta ADD/REMOVE with reverse-index + LRU maintenance. + +```lua +-- KEYS[1] = H:node: +-- ARGV: prefix, node_id, seq, now_ms, is_full_sync, caps_blob, +-- n_events, then per-event {kind,key,tier,size} +local rec = redis.call('HGETALL', KEYS[1]) +if #rec == 0 then return {'UNKNOWN', 0} end +local last = tonumber(hget(rec,'seq')) +if is_full_sync == 0 and seq ~= last + 1 then + redis.call('HSET', KEYS[1], 'last_hb', now_ms, 'status', 1) + redis.call('SADD', prefix..':nodes:alive', node_id) + return {'SEQ_GAP', last} -- caps/seq NOT advanced +end +redis.call('HSET', KEYS[1], 'last_hb', now_ms, 'status', 1, 'seq', seq, 'caps', caps_blob) +redis.call('SADD', prefix..':nodes:alive', node_id) +redis.call('HSET', prefix..':alive_peers', node_id, hget(rec,'peer')) +if is_full_sync == 1 then + wipe_node_blocks(prefix, node_id) -- drain node::blocks + for each ADD event do add_location(...) end -- REMOVE ignored on full_sync +else + for each event do apply_add_or_remove(...) end -- maintains reverse index + lru +end +return {'APPLIED', seq} +``` + +### 6.3 register_client + +Mirrors `RegisterClient` (`in_memory_master_metadata_store.cpp:263-303`): +reject only an ALIVE, non-stale record; revive EXPIRED or TTL-stale ALIVE. + +```lua +-- KEYS[1]=H:node:; ARGV: prefix, node_id, now_ms, stale_after_ms, +local rec = redis.call('HGETALL', KEYS[1]) +if #rec > 0 then + local status, last_hb = tonumber(hget(rec,'status')), tonumber(hget(rec,'last_hb')) + local stale = (now_ms - last_hb > stale_after_ms) or (status == 2) + if status == 1 and not stale then return 0 end -- reject live re-register +end +redis.call('HSET', KEYS[1], 'status',1,'last_hb',now_ms,'reg_at',now_ms,'seq',0, ) +redis.call('SADD', prefix..':nodes:alive', node_id) +redis.call('HSET', prefix..':alive_peers', node_id, peer) +return 1 +``` + +--- + +## 7. Fault tolerance + +- **Master crash / scale-out.** All durable AND volatile master state (records, + block locations, leases, LRU, hit counts) is in the store. A restarted or new + master reads it back; no warm-up, no per-replica drift. This is the entire + point of moving state behind the interface. +- **Store persistence.** Recommend AOF `appendfsync everysec` + at least one + replica. On Dragonfly, periodic snapshot + replica. The master relies on the + store's own HA (Sentinel / Cluster failover / Dragonfly replication). +- **Connection layer** (`RespClient`): connection pool sized to the gRPC + handler concurrency; per-call socket timeout; exponential-backoff reconnect; + automatic `SCRIPT LOAD` re-registration on `NOSCRIPT`; `EVALSHA` -> `EVAL` + fallback. +- **Degradation.** If the store is unreachable, the affected RPC returns + `grpc::UNAVAILABLE` (never fabricated / stale data). A readiness probe reports + the master unhealthy while the store is down, so load balancers drain it. +- **Multi-master safety.** All mutations go through Lua/CAS (no process-local + locks), so N masters against one store cannot split-brain. Verified by the + conformance race tests (reader vs in-flight full_sync/heartbeat asserts + old-or-new, never torn; never resolves a peer for an unregistered node). + +--- + +## 8. Known caveats + +- **full_sync payload sizing.** A full_sync from a peer with millions of keys is + one Lua script of that size, which blocks the (single-threaded) Redis server + for its whole duration. Per the interface `TODO(payload-sizing)`, full_sync + MUST be atomic, so the mitigation is on the peer: cap full_sync batch size + (proposed 100k events) and fragment larger resyncs. Documented as a peer-side + follow-up; the store enforces atomicity per call. +- **`GarbageCollectHits` / hit-key enumeration.** + Rather than `SCAN MATCH H:hit:*` (which has no key and cannot be routed + per-node in cluster mode), the backend maintains a `H:hit:index` reverse SET + of every hash with a live counter, updated by `MatchExternalKv`'s hit branch. + `GarbageCollectHits` is one keyed Lua that walks that set, drops each counter + whose `ls < cutoff`, and `SREM`s it — single-slot, cluster-routable, and runs + on the slow GC timer, not the hot path. +- **Dragonfly Lua atomicity** is validated by the conformance race suite, not + assumed. If a divergence surfaces, the fallback is to gate the Redis-backend + "multi-master" claim to Redis/Valkey and treat Dragonfly as single-writer. +- **Hot-path cost shift.** Every RouteGet moves from an in-process + `shared_mutex` read (nanoseconds) to a network + Lua round trip + (microseconds-to-milliseconds). Whether this is acceptable is exactly what + Phase 1 measures. +- **The store is rebuildable soft state; disk persistence is optional.** All + metadata here is a projection of what storage nodes report via heartbeats — the + source of truth is the peers, not Redis (Redis only holds "who has which + block", never the data itself). There are three recovery sources, and they + cover *different* failures: (1) **replica failover** for a single-node crash, + (2) **disk persistence (RDB/AOF)** for a *total*-outage restart + backups + + rollback, and (3) **heartbeat rebuild** (`SEQ_GAP -> full_sync`), which an + ordinary datastore does not have. Because (3) always exists, disk persistence + is optional: even a total wipe self-heals within one heartbeat cycle (RouteGet + returns *misses* meanwhile — cache misses, not data loss). Persistence is a + *separate axis* from replication (it is **not** "single-instance only"): keep + replicas if a single-node crash must be transparent, and note persistence only + buys a shorter warm-up window here, which is not worth its cost (next bullet). +- **Read = write amplification; do NOT enable synchronous persistence.** + `route_get_batch` bumps `_lease`/`_lacc`/`_acnt` on every touched key, so every + "read" is a *write* on the server. With `AOF appendfsync=always` each read would + block on a disk fsync and read throughput collapses; with replicas each read + also propagates. Deploy with `appendonly no` / `save ""` (both bench and prod). + This — not the durability argument above — is the operational reason disk + persistence stays off. +- **The control plane does not scale horizontally.** `clients_[0]` (the control + instance, or the node owning the control tag in cluster mode) holds every + control-plane key — `nodes:alive`, `alive_peers`, `extkv`, the per-node reverse + indexes — so every control op (`RegisterClient`, `ApplyHeartbeat`, + `ListAliveClients`, `ExpireStaleClients`) serializes on that one + (single-threaded) node. Block reads/writes shard across instances/nodes; the + control plane does not. For typical fleet sizes this is fine, but a very large + fleet or a heartbeat-heavy workload hits this single-node ceiling *before* the + block-read ceiling. Phase-2 extkv/hit placed on the control tag adds to this + load; sharding the control plane (like blocks) is a separate design if needed. +- **Dragonfly-in-cluster underperformance is largely a client limitation, not + purely topological.** Measured: Dragonfly in cluster mode tops out ~2–5k + routeget ops/s vs ~11–43k non-cluster, while a raw `redis-benchmark` (keyed + EVAL) on the same node sustains ~144k rps — the server is not slow. The main + cause is the client, not the topology: single/multi-endpoint pipelines a batch's + N per-shard EVALSHAs into **one** round trip (`RespClient::EvalPipeline`, via + hiredis), whereas `RespClusterClient::EvalPipeline` issues **N independent, + separately-routed** round trips (concurrent through a thread pool, but no + cross-slot pipelining). The slot -> single-thread pinning only bites under + balanced (one-tag-per-node) placement; with more tags Dragonfly *does* spread a + node's slots across its threads (measured CPU rises), yet the per-batch + round-trip fan-out still caps throughput. The lever that would actually help + (Redis Cluster too) is grouping a batch **by node** and pipelining each node's + EVALSHAs in one round trip — nontrivial, since redis++ does not expose + slot -> node (it needs per-node raw connections keyed off `CLUSTER SLOTS` + + CRC16). Until then: for throughput use non-cluster Dragonfly or multi-endpoint; + use cluster only when cross-node automatic failover (HA) is required, and do + **not** run Dragonfly in cluster mode. + +--- + +## 9. Connection / client library + +**Decision (locked):** the backend standardizes on `redis-plus-plus` (on +`hiredis`) as its client library. It covers RESP2/3, pipelines, Lua +(`EVAL/EVALSHA/SCRIPT LOAD`), connection pooling, **and Redis Cluster failover + +Sentinel**, and connects unchanged to Dragonfly and Valkey. Because the +production direction is **true HA** (replica + automatic failover), one library +that speaks single-node, Sentinel, and Cluster with the same API keeps the +maintained surface minimal — we do not hand-roll slot routing / `MOVED`/`ASK` / +failover. + +- **Build:** vendored as a `3rdparty/redis-plus-plus` submodule (pinned to + `1.3.15`) and built from source (static, `-fPIC`, C++17, tests off) behind + `option(USE_REDIS_BACKEND ... OFF)`, using the **system `hiredis`** (>= 0.14 + satisfies redis++'s >= 0.12.1 minimum). Building from source keeps any build + image working without an extra system package; if/when this ships in the + product image, redis++ can move to a prebuilt package. +- **`IRespClient` seam** (`redis/resp_value.h`): the store depends only on this + small surface (`Command`, `Eval`, `EvalPipeline`, `Ping`). Two implementations + live behind it: + - `RespClient` (`redis/resp_client.h`) — raw `hiredis`, one endpoint; drives + single and multi-endpoint (`UMBP_REDIS_SHARD_URIS`) modes. + - `RespClusterClient` (`redis/resp_cluster_client.h`) — `redis-plus-plus` + `RedisCluster`; drives cluster mode. It routes each command (by the key) and + script (by KEYS[0]) to the owning node and delegates MOVED/ASK redirection, + slot-map refresh, and master-failover reconnect to redis++. + Cluster mode is implemented and verified against a real 3-master+3-replica + cluster (conformance suite + failover + live reshard). Converging single/ + multi-endpoint onto redis++ (and adding a Sentinel mode) is a possible later + cleanup; both client paths already share the store's hot-path logic. + +--- + +## 10. Environment variables + +| Var | Default | Meaning | +| --- | --- | --- | +| `UMBP_METADATA_BACKEND` | `inmemory` | `inmemory` or `redis` | +| `UMBP_REDIS_URI` | (none) | e.g. `tcp://127.0.0.1:6379`; comma list for cluster seeds | +| `UMBP_REDIS_NAMESPACE` | `default` | deployment id used inside the hash tag `{umbp:}` | +| `UMBP_REDIS_CLUSTER` | `0` | `1` selects Redis Cluster mode: a redis-plus-plus `RedisCluster` client routes every command/script by hash-tag slot, with MOVED/ASK and master-failover handled by the client. Mutually exclusive with `UMBP_REDIS_SHARD_URIS`. Seeds come from the `UMBP_REDIS_URI` comma list (any reachable node bootstraps the rest). Block keys spread across nodes via `UMBP_REDIS_BLOCK_SHARDS` tags (default 16 in cluster mode). | +| `UMBP_REDIS_REQUIRED` | `1` | Startup readiness gate. `1` (default) fails master startup if the store is unreachable (the factory pings every endpoint), so a misconfigured store surfaces immediately instead of `UNAVAILABLE` on every RPC. `0` starts degraded and relies on runtime reconnect. | +| `UMBP_REDIS_SHARD_URIS` | (none) | comma-separated Redis URIs for **multi-endpoint** mode: one instance per block shard, so their scripts run on independent server processes/threads. This is the way past a single instance's single-thread ceiling — measured ~2.9x RouteGet throughput at 4 instances on a dedicated host (`M1 t16 ~5.2k -> M4 t16 ~15k ops/s`), with M1 flat at the single-slot ceiling. When set (>1 URI) it supersedes `UMBP_REDIS_BLOCK_SHARDS`. The first URI is the control instance. See §4.1. | +| `UMBP_REDIS_BLOCK_SHARDS` | `1` (16 in cluster) | number of hash-tag shards the block keyspace is spread over. `1` = legacy single-tag layout (byte-identical keys, whole-batch-atomic reads). `>1` spreads block lookups across slots (helps a threaded store / cluster, no gain on one single-threaded Redis). Ignored in multi-endpoint mode (shard count = number of URIs); in cluster mode it is the number of block-shard tags spread across nodes by slot (defaults to 16). Fixed for a deployment's lifetime; clamped to `[1, 4096]`; `<=0` → `1`. See §4. | +| `UMBP_REDIS_POOL_SIZE` | (cpu-derived) | connection pool size | +| `UMBP_REDIS_CONNECT_TIMEOUT_MS` | `1000` | connect timeout | +| `UMBP_REDIS_SOCKET_TIMEOUT_MS` | `1000` | per-command socket timeout | +| `UMBP_REDIS_PASSWORD` | (none) | auth (may also be in the URI) | + +These follow the resolution semantics in +[`runtime-env-vars.md`](./runtime-env-vars.md) (cached at startup; one WARN per +invalid value). + +--- + +## 11. Phasing + +The whole approach lives or dies on hot-path latency against a real store, so +we validate that **before** implementing the full interface. + +### Phase 0 — Design (this doc) + +This document plus a cross-reference from `design-master-control-plane.md`. + +### Phase 1 — Core vertical slice + performance comparison + +- Introduce `redis-plus-plus` behind `USE_REDIS_BACKEND` (default OFF). +- `RespClient`, `KeySchema`, `RedisMasterMetadataStore` implementing only the + six hot-path methods (§2). All other methods `throw + std::logic_error("unimplemented (phase 1)")`. +- `MakeMasterMetadataStore(config)` factory + the one-line + `master_server.cpp` change, gated on `UMBP_METADATA_BACKEND`. +- `tools/redis/docker-compose.yml`: single-node Redis 7 and Dragonfly. +- **Deliverable:** run the existing `bench_kvevent_master_pressure` + (`tests/cpp/umbp/distributed/bench_kvevent_master_pressure.cpp`) with + `UMBP_METADATA_BACKEND` = inmemory / redis / dragonfly and produce + `docs/redis-backend-phase1-bench.md`: p50/p99 of `BatchRouteGet` / + `Heartbeat`, throughput, and read-after-write miss rate. This is the go/no-go + and optimization-direction signal. + +### Phase 2 — Full implementation + fault tolerance + +- Implement the remaining ~30 methods (external-KV, hit GC, eviction ZSET, + reaper `ExpireStaleClients`, cascading unregister). +- Parameterize `test_in_memory_master_metadata_store.cpp` into a backend-agnostic + conformance suite that runs against both in-memory and Redis (integration + label; skipped where no Redis is available). +- Fault tolerance: reconnect/backoff, `NOSCRIPT`/`EVALSHA` fallback, + `UNAVAILABLE` degradation, readiness probe. +- Multi-master: two master processes against one store passing the conformance + race cases; kill-a-replica failover. +- Redis Cluster: assert same-slot for every scripted keyset; handle + `MOVED`/`ASK` in the client. + +--- + +## 12. Testing + +- **Phase 1:** targeted unit tests for the six hot methods against a local Redis + (integration label, skippable) + the bench report. +- **Phase 2:** the parameterized conformance suite (same assertions on in-memory + and Redis), the reader-vs-full_sync race cases, and fault-injection tests + (drop the connection mid-call, `NOSCRIPT`, replica failover). + +--- + +## 13. Decisions (locked) + +1. **Fault-tolerance target:** true HA — replica + automatic failover, so a + single store node failure does not drop service or lose cache hits. This is + why Redis Cluster (and/or Sentinel) is on the roadmap; manual sharding + (`UMBP_REDIS_SHARD_URIS`) remains the proven horizontal-scale mode but has no + replication/failover. +2. **Client library:** `redis-plus-plus` (see §9), vendored as a submodule and + built from source, using the system `hiredis`. +3. **Dependency delivery:** submodule + CMake source build now (reproducible, no + image dependency); promote to the official image later if this ships. +4. **Store deployment:** independent containers via `tools/redis/store.sh` + + `docker-compose.yml` (single / dragonfly / cluster profiles, host networking), + not a redis-server built inside the app container. `run_local_backends.sh` is + the no-docker fallback. +5. **Phase 1 backends:** Redis + Dragonfly. +6. **Atomicity posture:** RESP+Lua common subset with single-instance strong + atomicity and documented cluster/Dragonfly caveats. diff --git a/src/umbp/include/umbp/distributed/master/master_metadata_store.h b/src/umbp/include/umbp/distributed/master/master_metadata_store.h index 7f963fb90..ba770e764 100644 --- a/src/umbp/include/umbp/distributed/master/master_metadata_store.h +++ b/src/umbp/include/umbp/distributed/master/master_metadata_store.h @@ -175,6 +175,10 @@ #include "umbp/distributed/types.h" +namespace mori::metrics { +class MetricsServer; +} + namespace mori::umbp { // ===================================================================== @@ -243,6 +247,12 @@ class IMasterMetadataStore { public: virtual ~IMasterMetadataStore() = default; + // Optional observability hook: attach the master's Prometheus metrics server + // so a backend can export per-op latency (e.g. store->Redis round trips). + // Default no-op — backends that don't emit store metrics (in-memory) ignore + // it. Called once at master startup (after the metrics server is created). + virtual void SetMetricsSink(mori::metrics::MetricsServer* /*metrics*/) {} + // =================================================================== // Cross-store write operations — each call is atomic. // =================================================================== diff --git a/src/umbp/include/umbp/distributed/master/master_metadata_store_factory.h b/src/umbp/include/umbp/distributed/master/master_metadata_store_factory.h new file mode 100644 index 000000000..eb121f415 --- /dev/null +++ b/src/umbp/include/umbp/distributed/master/master_metadata_store_factory.h @@ -0,0 +1,46 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. + +// MakeMasterMetadataStore — selects the master's metadata backend at startup. +// +// Reads UMBP_METADATA_BACKEND ("inmemory" | "redis"; default "inmemory") and +// the UMBP_REDIS_* connection knobs. This is the single production wiring point +// for the Redis backend: MasterServer constructs its store through this factory +// so router / eviction / reaper stay backend-agnostic (they only see +// IMasterMetadataStore&). + +#pragma once + +#include + +#include "umbp/distributed/master/master_metadata_store.h" + +namespace mori::umbp { + +// Constructs the metadata store selected by UMBP_METADATA_BACKEND. Falls back +// to the in-memory backend when the env is unset/"inmemory". Throws +// std::runtime_error if "redis" is requested but the Redis backend was not +// compiled in (USE_REDIS_BACKEND=OFF), so a misconfiguration is loud rather +// than silently serving from a wrong backend. +std::unique_ptr MakeMasterMetadataStore(); + +} // namespace mori::umbp diff --git a/src/umbp/include/umbp/distributed/master/redis/cluster_slots.h b/src/umbp/include/umbp/distributed/master/redis/cluster_slots.h new file mode 100644 index 000000000..08ebf1722 --- /dev/null +++ b/src/umbp/include/umbp/distributed/master/redis/cluster_slots.h @@ -0,0 +1,98 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. + +// Redis Cluster slot math, used ONLY at startup to place one block-shard hash +// tag on each master node (balanced placement). The redis-plus-plus client does +// the real per-command routing; this is just so we can pick tags whose slot +// lands on a chosen node so a RouteGet batch spreads evenly instead of piling +// onto whichever node the default {umbp::bS} tags happen to hash to. + +#pragma once + +#include +#include +#include +#include + +namespace mori::umbp::redis { + +// Number of hash slots in a Redis Cluster (fixed by the protocol). +inline constexpr int kNumSlots = 16384; + +// CRC16-CCITT (XMODEM): poly 0x1021, init 0x0000, no reflection — the exact +// function Redis Cluster uses for CLUSTER KEYSLOT. Bitwise (no 256-entry table) +// because it only runs a few hundred times at startup for tag placement, so the +// table's speed is irrelevant and a hand-copied table would risk a typo. +inline uint16_t Crc16(const char* buf, std::size_t len) { + uint16_t crc = 0; + for (std::size_t i = 0; i < len; ++i) { + crc ^= static_cast(static_cast(buf[i])) << 8; + for (int b = 0; b < 8; ++b) { + crc = (crc & 0x8000) ? static_cast((crc << 1) ^ 0x1021) + : static_cast(crc << 1); + } + } + return crc; +} + +// Slot for a key, honoring the hash-tag rule: if the key contains "{...}" with a +// non-empty body, only that body is hashed (so all keys sharing a tag share a +// slot). Matches Redis' keyHashSlot(). +inline int SlotOfKey(const std::string& key) { + const auto open = key.find('{'); + if (open != std::string::npos) { + const auto close = key.find('}', open + 1); + if (close != std::string::npos && close > open + 1) { + return Crc16(key.data() + open + 1, close - open - 1) % kNumSlots; + } + } + return Crc16(key.data(), key.size()) % kNumSlots; +} + +// A master's owned slot ranges (inclusive), as returned by CLUSTER SLOTS. +using SlotRange = std::pair; + +inline bool SlotInRanges(int slot, const std::vector& ranges) { + for (const auto& r : ranges) { + if (slot >= r.first && slot <= r.second) return true; + } + return false; +} + +// Find a block-shard hash tag "{umbp::b}" whose slot lands in `ranges` +// (i.e. is served by a chosen master). Searches k = 0,1,2,... deterministically +// so the same topology yields the same tag across restarts (stable key +// placement). Returns false if none found within max_tries (each master owns +// ~1/N of the slots, so a hit is expected within a few N tries). +inline bool FindTagForRanges(const std::string& ns, const std::vector& ranges, + std::string* out_tag, int max_tries = 200000) { + for (int k = 0; k < max_tries; ++k) { + std::string tag = "{umbp:" + ns + ":b" + std::to_string(k) + "}"; + if (SlotInRanges(SlotOfKey(tag), ranges)) { + *out_tag = std::move(tag); + return true; + } + } + return false; +} + +} // namespace mori::umbp::redis diff --git a/src/umbp/include/umbp/distributed/master/redis/key_schema.h b/src/umbp/include/umbp/distributed/master/redis/key_schema.h new file mode 100644 index 000000000..bb0609717 --- /dev/null +++ b/src/umbp/include/umbp/distributed/master/redis/key_schema.h @@ -0,0 +1,198 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. + +// KeySchema — every master metadata key and the hash tags they hash under. +// +// Two families of keys, deliberately split across hash tags: +// +// * Control-plane keys (client records, ALIVE set, peer projection, the +// per-node reverse indexes) all share the control tag `{umbp:}`, so +// they co-locate in one slot and a single Lua script can mutate them +// atomically. +// +// * Block-location keys are spread across `num_shards` shard tags +// `{umbp::b}` chosen by hashing the user key. Spreading blocks +// over many slots is what lets the read hot path (RouteGet / Exists) run +// one single-slot script per shard instead of piling every block lookup +// onto one slot / one server thread. On a multi-threaded store (Dragonfly) +// different shards land on different proactor threads and run in parallel. +// +// Backwards compatibility: `num_shards == 1` puts block keys back under the +// control tag, so the emitted key strings are byte-identical to the original +// single-tag schema (a deployment with sharding disabled is unchanged). +// +// The shard tag is derived from a stable hash (FNV-1a), NOT std::hash, so a key +// written by one build always resolves to the same shard after a rebuild. +// `num_shards` is fixed for a deployment's lifetime, exactly like the in-memory +// backend's UMBP_MASTER_INDEX_SHARDS: changing it with live data would strand +// keys under their old shard tag. +// +// See design-redis-metadata-store.md §4 for the full schema table. + +#pragma once + +#include +#include +#include +#include + +#include "umbp/distributed/types.h" + +namespace mori::umbp::redis { + +class KeySchema { + public: + explicit KeySchema(const std::string& ns, std::size_t num_shards = 1) + : control_tag_("{umbp:" + ns + "}"), num_shards_(num_shards == 0 ? 1 : num_shards) { + block_tags_.reserve(num_shards_); + if (num_shards_ == 1) { + // Legacy layout: block keys live under the control tag. + block_tags_.push_back(control_tag_); + } else { + for (std::size_t s = 0; s < num_shards_; ++s) { + block_tags_.push_back("{umbp:" + ns + ":b" + std::to_string(s) + "}"); + } + } + } + + // Explicit block-shard tags (Redis Cluster balanced placement): the caller + // supplies one tag per master, each chosen so its slot is served by a distinct + // node, so a RouteGet batch spreads evenly instead of piling onto whichever + // node the formulaic {umbp::bS} tags happen to hash to. num_shards is the + // number of tags; user keys still distribute across them by StableHash. + KeySchema(const std::string& ns, std::vector block_tags) + : control_tag_("{umbp:" + ns + "}"), + num_shards_(block_tags.empty() ? 1 : block_tags.size()), + block_tags_(std::move(block_tags)) { + if (block_tags_.empty()) block_tags_.push_back(control_tag_); + } + + std::size_t NumShards() const { return num_shards_; } + + // The control-plane hash tag, e.g. "{umbp:default}". Passed to Lua as ARGV[1] + // so scripts can compose auxiliary control-slot key names in the same slot. + const std::string& Tag() const { return control_tag_; } + + // HASH: one client record. + std::string Node(const std::string& node_id) const { return control_tag_ + ":node:" + node_id; } + + // SET: ALIVE node ids. + std::string NodesAlive() const { return control_tag_ + ":nodes:alive"; } + + // HASH: node_id -> peer_address for ALIVE nodes only. + std::string AlivePeers() const { return control_tag_ + ":alive_peers"; } + + // The shard a user key belongs to. + std::size_t ShardOf(const std::string& key) const { + return num_shards_ == 1 ? 0 : (StableHash(key) % num_shards_); + } + + // The hash tag backing shard `shard` (0 <= shard < num_shards). + const std::string& ShardTag(std::size_t shard) const { return block_tags_[shard]; } + + // HASH: block locations ("l||" -> size) plus meta fields. Placed + // in the key's shard slot. + std::string Block(const std::string& key) const { + return ShardTag(ShardOf(key)) + ":block:" + key; + } + + // SET: the block keys a node owns (reverse index for node-scoped wipe). Kept + // on the control tag as one set per node; its members are full (already + // sharded) block-key strings, so the wipe scripts can HDEL them directly + // without recomputing the shard in Lua. Used by the single-endpoint path. + std::string NodeBlocks(const std::string& node_id) const { + return control_tag_ + ":node:" + node_id + ":blocks"; + } + + // SET: the block keys a node owns within one shard, placed under that shard's + // tag so it co-locates (same slot) with the shard's block keys. The + // multi-endpoint path keeps one reverse index per (node, shard) on the shard's + // own instance, so each per-shard wipe/full-sync script is single-slot and + // touches only its own instance. Members are full block-key strings. + std::string NodeBlocks(const std::string& node_id, std::size_t shard) const { + return ShardTag(shard) + ":node:" + node_id + ":blocks"; + } + + // The shard a block-hash (external-kv key) belongs to. Same StableHash + // bucketing as ShardOf so an extkv: and its hit: always co-locate, + // and so a batch of hashes spreads across shards exactly like a block batch. + std::size_t ExtKvShardOf(const std::string& hash) const { return ShardOf(hash); } + + // SET: the external-kv hashes a node registered within one shard (reverse + // index for the node-scoped wipe), placed under that shard's tag so it + // co-locates with the shard's extkv: keys. One set per (node, shard) so + // every extkv mutation/wipe stays single-slot and touches only its own shard, + // exactly like NodeBlocks(node, shard). Members are hash strings. + // For num_shards == 1, shard 0's tag IS the control tag, so this is + // byte-identical to the legacy single control-tag reverse index. + std::string ExtKvNode(const std::string& node_id, std::size_t shard) const { + return ShardTag(shard) + ":extkv:node:" + node_id; + } + + // HASH: one external-kv entry, node_id -> tier bitmask (bit == 1< control tag, legacy + // byte-identical layout) so the external-KV hot path spreads across shards / + // instances / proactor threads instead of piling every match + hit-count write + // onto the single control slot. A single Lua still mutates it atomically within + // its shard; cross-shard extkv ops fan out one single-slot script per shard. + std::string ExtKv(const std::string& hash) const { + return ShardTag(ShardOf(hash)) + ":extkv:" + hash; + } + + // HASH: one per-hash hit counter, fields `c` (count) and `ls` (last_seen ms). + // Co-located with its extkv: (same shard) so a match + hit-count bump is + // one single-slot script per shard. + std::string Hit(const std::string& hash) const { + return ShardTag(ShardOf(hash)) + ":hit:" + hash; + } + + // SET: every hash that currently has a hit counter WITHIN one shard. A + // Redis-specific reverse index (the in-memory backend just iterates its map) so + // GarbageCollectHits is one keyed Lua per shard over this set instead of a SCAN + // — SCAN has no key and cannot be routed per-node in cluster mode. One index + // per shard so GC fans out; num_shards == 1 => control tag (byte-identical). + std::string HitIndex(std::size_t shard) const { return ShardTag(shard) + ":hit:index"; } + + private: + // FNV-1a (32-bit): small, fast, and stable across builds/runs. std::hash is + // avoided on purpose — it is implementation-defined, so its bucketing could + // change under a toolchain upgrade and silently relocate every key's shard. + static std::size_t StableHash(const std::string& key) { + uint32_t hash = 2166136261u; + for (const unsigned char c : key) { + hash ^= c; + hash *= 16777619u; + } + return static_cast(hash); + } + + std::string control_tag_; + std::size_t num_shards_; + std::vector block_tags_; // indexed by shard; size == num_shards_ +}; + +// Location hash-field prefix marker. A field named "l||" holds a +// location's size; any field beginning with this marker is a location, and +// "_"-prefixed fields (_lease/_lacc/_acnt/_created) are per-block metadata. +inline constexpr const char* kLocFieldMarker = "l|"; + +} // namespace mori::umbp::redis diff --git a/src/umbp/include/umbp/distributed/master/redis/lua_scripts.h b/src/umbp/include/umbp/distributed/master/redis/lua_scripts.h new file mode 100644 index 000000000..b814cd9e7 --- /dev/null +++ b/src/umbp/include/umbp/distributed/master/redis/lua_scripts.h @@ -0,0 +1,892 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. + +// Server-side Lua for every multi-key atomic mutation in the Redis metadata +// backend. Each script runs atomically on the store; all keys it touches share +// the deployment hash tag (passed as ARGV[1]) so it stays single-slot and +// cross-key atomic on Redis Cluster / Dragonfly / Valkey too. +// +// Dragonfly multithreading (per-script "--!df" flag, NOT the global launch flag): +// Dragonfly parallelizes scripts across proactor threads ONLY when a script +// declares every key it touches (lock-ahead mode: each shard's script runs on +// its slot's thread, concurrently with others). A script that reaches keys not +// in KEYS[] must run as a GLOBAL transaction, which takes a store-wide lock and +// serializes ALL shards — so launching Dragonfly with +// `--default_lua_flags=allow-undeclared-keys` forces EVERY script (including the +// read hot path) global and kills the sharding win. +// Fix: the two read hot-path scripts (route_get_batch / exists_batch) already +// pass every key via KEYS[], so they need no flag and run per-shard in parallel. +// The write/control scripts below DO derive auxiliary same-slot keys from the +// hash tag, so each carries a first-line "--!df flags=allow-undeclared-keys" +// directive that lets Dragonfly run just that script global. Redis treats the +// line as a plain Lua comment (it only honours a "#!" shebang), so this is a +// no-op there and the single-node / Cluster behaviour is byte-for-byte +// unchanged. Deploy Dragonfly WITHOUT the global flag to get the parallelism. +// +// Determinism: scripts never read the server clock or randomkey; every +// timestamp is passed in by the caller as epoch-milliseconds (hazard #7 in +// master_metadata_store.h), so they are replication-safe. +// +// Phase 1 scope: the eviction LRU ZSET is NOT maintained here (eviction / +// EnumerateEvictionCandidates is Phase 2). The block hash still carries the +// _lease / _lacc / _acnt meta so RouteGet semantics match the in-memory impl. + +#pragma once + +#include + +namespace mori::umbp::redis { + +// Each script is a single, stable global object (not a per-use temporary): the +// SHA cache in RespClient / RespClusterClient keys on the script object's ADDRESS +// (`&script`), not its ~1.5KB text, so the read hot path never rehashes/compares +// the whole Lua body under the cache lock on every EVALSHA. That only works if a +// call site passes a reference to THESE objects (an `inline const std::string`, +// one instance per literal across the program) rather than constructing a fresh +// std::string from a `const char*` each call. Do NOT copy a script into a local +// std::string before Eval() — that would defeat the pointer-identity cache. + +// register_client: +// KEYS[1] = node key +// ARGV = [tag, node_id, now_ms, stale_after_ms, addr, peer, caps, engine, tags] +// Returns 1 if registered/revived, 0 if rejected (ALIVE and not stale). +inline const std::string kRegisterClientLua = R"LUA(--!df flags=allow-undeclared-keys +local nodeKey = KEYS[1] +local tag = ARGV[1] +local nodeId = ARGV[2] +local now = tonumber(ARGV[3]) +local stale = tonumber(ARGV[4]) +if redis.call('EXISTS', nodeKey) == 1 then + local status = tonumber(redis.call('HGET', nodeKey, 'status') or '0') + local lastHb = tonumber(redis.call('HGET', nodeKey, 'last_hb') or '0') + local isStale = ((now - lastHb) > stale) or (status == 2) + if status == 1 and (not isStale) then + return 0 + end +end +redis.call('DEL', nodeKey) +redis.call('HSET', nodeKey, + 'status', 1, 'last_hb', now, 'reg_at', now, 'seq', 0, + 'addr', ARGV[5], 'peer', ARGV[6], 'caps', ARGV[7], + 'engine', ARGV[8], 'tags', ARGV[9]) +redis.call('SADD', tag .. ':nodes:alive', nodeId) +redis.call('HSET', tag .. ':alive_peers', nodeId, ARGV[6]) +return 1 +)LUA"; + +// apply_heartbeat: +// KEYS[1] = node key +// ARGV = [tag, node_id, seq, now_ms, is_full_sync, caps_blob, n_events, +// then per event: kind, block_key, tier, size] +// `block_key` is the FULL (already sharded) block key composed by the caller +// (KeySchema::Block), so this script never has to know the shard layout. The +// node's reverse-index set stores these full block keys as members. +// Returns { status_string, acked_seq_string } +// status_string in { "UNKNOWN", "SEQ_GAP", "APPLIED" }. +inline const std::string kApplyHeartbeatLua = R"LUA(--!df flags=allow-undeclared-keys +local nodeKey = KEYS[1] +local tag = ARGV[1] +local nodeId = ARGV[2] +local seq = tonumber(ARGV[3]) +local now = tonumber(ARGV[4]) +local full = tonumber(ARGV[5]) +local caps = ARGV[6] +local nev = tonumber(ARGV[7]) + +if redis.call('EXISTS', nodeKey) == 0 then + return { 'UNKNOWN', '0' } +end +local last = tonumber(redis.call('HGET', nodeKey, 'seq') or '0') +local aliveSet = tag .. ':nodes:alive' +local peers = tag .. ':alive_peers' +local peer = redis.call('HGET', nodeKey, 'peer') or '' + +if full == 0 and seq ~= last + 1 then + redis.call('HSET', nodeKey, 'last_hb', now, 'status', 1) + redis.call('SADD', aliveSet, nodeId) + redis.call('HSET', peers, nodeId, peer) + return { 'SEQ_GAP', tostring(last) } +end + +redis.call('HSET', nodeKey, 'last_hb', now, 'status', 1, 'seq', seq, 'caps', caps) +redis.call('SADD', aliveSet, nodeId) +redis.call('HSET', peers, nodeId, peer) + +local blocksSet = tag .. ':node:' .. nodeId .. ':blocks' +local nodePfx = 'l|' .. nodeId .. '|' + +local function cleanupEmpty(bk) + local flds = redis.call('HKEYS', bk) + local anyLoc = false + for _, f in ipairs(flds) do + if string.sub(f, 1, 2) == 'l|' then anyLoc = true break end + end + if not anyLoc then redis.call('DEL', bk) end +end + +local function addLoc(blockKey, tier, size) + if redis.call('EXISTS', blockKey) == 0 then + redis.call('HSET', blockKey, '_created', now, '_lacc', now, '_acnt', 0) + end + local field = 'l|' .. nodeId .. '|' .. tier + if redis.call('HEXISTS', blockKey, field) == 0 then + redis.call('HSET', blockKey, field, size) + redis.call('SADD', blocksSet, blockKey) + end +end + +local function removeLoc(blockKey, tier) + local field = 'l|' .. nodeId .. '|' .. tier + if redis.call('HDEL', blockKey, field) == 1 then + local flds = redis.call('HKEYS', blockKey) + local nodeStill = false + for _, f in ipairs(flds) do + if string.sub(f, 1, string.len(nodePfx)) == nodePfx then nodeStill = true break end + end + if not nodeStill then redis.call('SREM', blocksSet, blockKey) end + cleanupEmpty(blockKey) + end +end + +if full == 1 then + local members = redis.call('SMEMBERS', blocksSet) + for _, bk in ipairs(members) do + local flds = redis.call('HKEYS', bk) + for _, f in ipairs(flds) do + if string.sub(f, 1, string.len(nodePfx)) == nodePfx then + redis.call('HDEL', bk, f) + end + end + cleanupEmpty(bk) + end + redis.call('DEL', blocksSet) + for i = 0, nev - 1 do + local base = 8 + i * 4 + if tonumber(ARGV[base]) == 0 then + addLoc(ARGV[base + 1], ARGV[base + 2], ARGV[base + 3]) + end + end +else + for i = 0, nev - 1 do + local base = 8 + i * 4 + if tonumber(ARGV[base]) == 0 then + addLoc(ARGV[base + 1], ARGV[base + 2], ARGV[base + 3]) + else + removeLoc(ARGV[base + 1], ARGV[base + 2]) + end + end +end +return { 'APPLIED', tostring(seq) } +)LUA"; + +// route_get_batch: +// KEYS[1..n] = block keys (all in one shard: the caller groups a batch by +// shard and runs one invocation per shard via RespClient::EvalPipeline, so +// KEYS stays single-slot on Redis Cluster / one proactor on Dragonfly). +// ARGV = [now_ms, lease_ms, n_exclude, exclude_node_1..exclude_node_k] +// Returns an array of n elements; element i is a flat array +// [node, size, tier, node, size, tier, ...] of the surviving locations. +// Only keys with >=1 surviving location get a lease + access bump. +// +// The per-key HGETALL already returns _acnt, so the lease/access bump folds +// into a SINGLE HSET (_lease/_lacc/_acnt) instead of HSET + a separate +// HINCRBY. This is semantics-preserving (_acnt still ends at old+1) and drops +// the per-touched-key write commands from 2 to 1 — the single-slot server is +// redis.call()-count bound, so fewer calls per script is a direct win. +inline const std::string kRouteGetBatchLua = R"LUA( +local now = tonumber(ARGV[1]) +local lease = tonumber(ARGV[2]) +local ne = tonumber(ARGV[3]) +local excl = {} +for i = 1, ne do excl[ARGV[3 + i]] = true end +local out = {} +for i = 1, #KEYS do + local flds = redis.call('HGETALL', KEYS[i]) + local locs = {} + local touched = false + local acnt = 0 + local j = 1 + while j <= #flds do + local f = flds[j] + local v = flds[j + 1] + j = j + 2 + if string.sub(f, 1, 2) == 'l|' then + local rest = string.sub(f, 3) + local sep = string.find(rest, '|', 1, true) + if sep ~= nil then + local node = string.sub(rest, 1, sep - 1) + local tier = string.sub(rest, sep + 1) + if not excl[node] then + locs[#locs + 1] = node + locs[#locs + 1] = v + locs[#locs + 1] = tier + touched = true + end + end + elseif f == '_acnt' then + acnt = tonumber(v) or 0 + end + end + if touched then + redis.call('HSET', KEYS[i], '_lease', now + lease, '_lacc', now, '_acnt', acnt + 1) + end + out[i] = locs +end +return out +)LUA"; + +// exists_batch: +// KEYS[1..n] = block keys (single-slot per invocation, grouped by shard by +// the caller, same as route_get_batch). +// Returns an array of n integers (1 if the key has >=1 location, else 0). +inline const std::string kExistsBatchLua = R"LUA( +local out = {} +for i = 1, #KEYS do + local flds = redis.call('HKEYS', KEYS[i]) + local has = 0 + for _, f in ipairs(flds) do + if string.sub(f, 1, 2) == 'l|' then has = 1 break end + end + out[i] = has +end +return out +)LUA"; + +// list_alive: +// ARGV = [tag] +// Returns an array; each element is { node_id, flat_hgetall_of_node_hash } +// for every ALIVE node. +inline const std::string kListAliveLua = R"LUA(--!df flags=allow-undeclared-keys +local tag = ARGV[1] +local members = redis.call('SMEMBERS', tag .. ':nodes:alive') +local out = {} +for _, id in ipairs(members) do + local nk = tag .. ':node:' .. id + local status = tonumber(redis.call('HGET', nk, 'status') or '0') + if status == 1 then + out[#out + 1] = { id, redis.call('HGETALL', nk) } + end +end +return out +)LUA"; + +// unregister_client: +// KEYS[1] = node key +// ARGV = [tag, node_id, wipe_extkv] +// wipe_extkv == 1 (num_shards == 1): the node's external-kv lives on the +// control tag, so wipe it inline (reverse-index members are full extkv keys). +// wipe_extkv == 0 (num_shards > 1): extkv is sharded off the control slot; the +// store wipes it separately per shard (WipeNodeExtKvMulti) after this runs. +// Returns 1 if the client existed, 0 otherwise. +inline const std::string kUnregisterClientLua = R"LUA(--!df flags=allow-undeclared-keys +local nodeKey = KEYS[1] +local tag = ARGV[1] +local nodeId = ARGV[2] +local wipeExtkv = tonumber(ARGV[3]) +if redis.call('EXISTS', nodeKey) == 0 then return 0 end +redis.call('DEL', nodeKey) +redis.call('SREM', tag .. ':nodes:alive', nodeId) +redis.call('HDEL', tag .. ':alive_peers', nodeId) +local blocksSet = tag .. ':node:' .. nodeId .. ':blocks' +local nodePfx = 'l|' .. nodeId .. '|' +local members = redis.call('SMEMBERS', blocksSet) +for _, bk in ipairs(members) do + local flds = redis.call('HKEYS', bk) + for _, f in ipairs(flds) do + if string.sub(f, 1, string.len(nodePfx)) == nodePfx then redis.call('HDEL', bk, f) end + end + local flds2 = redis.call('HKEYS', bk) + local anyLoc = false + for _, f in ipairs(flds2) do + if string.sub(f, 1, 2) == 'l|' then anyLoc = true break end + end + if not anyLoc then redis.call('DEL', bk) end +end +redis.call('DEL', blocksSet) +-- Wipe this node's external-kv: clear its bit from every extkv: it +-- registered (reverse-index members are full keys) and drop the reverse index. +if wipeExtkv == 1 then + local exNode = tag .. ':extkv:node:' .. nodeId + local ekeys = redis.call('SMEMBERS', exNode) + for _, ekey in ipairs(ekeys) do + redis.call('HDEL', ekey, nodeId) + if redis.call('HLEN', ekey) == 0 then redis.call('DEL', ekey) end + end + redis.call('DEL', exNode) +end +return 1 +)LUA"; + +// expire_stale: +// ARGV = [tag, cutoff_ms, wipe_extkv] +// Flips ALIVE->EXPIRED for nodes whose last_hb < cutoff (keeping the row), +// drops their block locations, and (wipe_extkv == 1, num_shards == 1) their +// external-kv inline; when wipe_extkv == 0 the store wipes sharded extkv per +// shard afterwards. Returns the dead node ids. +inline const std::string kExpireStaleLua = R"LUA(--!df flags=allow-undeclared-keys +local tag = ARGV[1] +local cutoff = tonumber(ARGV[2]) +local wipeExtkv = tonumber(ARGV[3]) +local members = redis.call('SMEMBERS', tag .. ':nodes:alive') +local dead = {} +for _, id in ipairs(members) do + local nk = tag .. ':node:' .. id + local status = tonumber(redis.call('HGET', nk, 'status') or '0') + local lastHb = tonumber(redis.call('HGET', nk, 'last_hb') or '0') + if status == 1 and lastHb < cutoff then + redis.call('HSET', nk, 'status', 2) + redis.call('SREM', tag .. ':nodes:alive', id) + redis.call('HDEL', tag .. ':alive_peers', id) + local blocksSet = tag .. ':node:' .. id .. ':blocks' + local nodePfx = 'l|' .. id .. '|' + local ms = redis.call('SMEMBERS', blocksSet) + for _, bk in ipairs(ms) do + local flds = redis.call('HKEYS', bk) + for _, f in ipairs(flds) do + if string.sub(f, 1, string.len(nodePfx)) == nodePfx then redis.call('HDEL', bk, f) end + end + local flds2 = redis.call('HKEYS', bk) + local anyLoc = false + for _, f in ipairs(flds2) do + if string.sub(f, 1, 2) == 'l|' then anyLoc = true break end + end + if not anyLoc then redis.call('DEL', bk) end + end + redis.call('DEL', blocksSet) + if wipeExtkv == 1 then + local exNode = tag .. ':extkv:node:' .. id + local ekeys = redis.call('SMEMBERS', exNode) + for _, ekey in ipairs(ekeys) do + redis.call('HDEL', ekey, id) + if redis.call('HLEN', ekey) == 0 then redis.call('DEL', ekey) end + end + redis.call('DEL', exNode) + end + dead[#dead + 1] = id + end +end +return dead +)LUA"; + +// ===================================================================== +// Multi-endpoint (sharded across Redis instances) scripts. +// +// When block shards live on DIFFERENT Redis instances than the control keys, no +// single Lua script can span them. The cross-store writes are therefore split +// into a control script (runs on the control instance) plus per-shard block +// scripts (run on each shard's own instance, single-slot). The store runs the +// control step first, then the block step(s); block ADD/REMOVE and full-sync +// clear+replay are idempotent, so a failed/retried block step is safe and a +// partial failure is healed by the peer's next SEQ_GAP -> full_sync. The +// single-endpoint path above (M==1) is unchanged and still fully atomic. +// ===================================================================== + +// apply_heartbeat_control (control instance): +// KEYS[1] = node key +// ARGV = [tag, node_id, seq, now_ms, is_full_sync, caps_blob] +// seq-CAS + record + nodes:alive/alive_peers ONLY (no block work). +// Returns { status_string, acked_seq_string } like apply_heartbeat. +inline const std::string kApplyHeartbeatControlLua = R"LUA(--!df flags=allow-undeclared-keys +local nodeKey = KEYS[1] +local tag = ARGV[1] +local nodeId = ARGV[2] +local seq = tonumber(ARGV[3]) +local now = tonumber(ARGV[4]) +local full = tonumber(ARGV[5]) +local caps = ARGV[6] +if redis.call('EXISTS', nodeKey) == 0 then + return { 'UNKNOWN', '0' } +end +local last = tonumber(redis.call('HGET', nodeKey, 'seq') or '0') +local aliveSet = tag .. ':nodes:alive' +local peers = tag .. ':alive_peers' +local peer = redis.call('HGET', nodeKey, 'peer') or '' +if full == 0 and seq ~= last + 1 then + redis.call('HSET', nodeKey, 'last_hb', now, 'status', 1) + redis.call('SADD', aliveSet, nodeId) + redis.call('HSET', peers, nodeId, peer) + return { 'SEQ_GAP', tostring(last) } +end +redis.call('HSET', nodeKey, 'last_hb', now, 'status', 1, 'seq', seq, 'caps', caps) +redis.call('SADD', aliveSet, nodeId) +redis.call('HSET', peers, nodeId, peer) +return { 'APPLIED', tostring(seq) } +)LUA"; + +// apply_block_events (one shard's instance): +// ARGV = [revidx_key, node_prefix, is_full_sync, now_ms, n_events, +// then per event: kind, block_key, tier, size] +// `revidx_key` is this (node, shard) reverse-index set; `node_prefix` is +// 'l||'. All block keys passed in ARGV are on this instance/slot. +// Idempotent: ADD overwrites, REMOVE of a missing loc is a no-op, full_sync +// clears the node's blocks here then replays the ADDs. Returns 'OK'. +inline const std::string kApplyBlockEventsLua = R"LUA(--!df flags=allow-undeclared-keys +local revidx = ARGV[1] +local nodePfx = ARGV[2] +local full = tonumber(ARGV[3]) +local now = tonumber(ARGV[4]) +local nev = tonumber(ARGV[5]) +local pfxLen = string.len(nodePfx) + +local function cleanupEmpty(bk) + local flds = redis.call('HKEYS', bk) + for _, f in ipairs(flds) do + if string.sub(f, 1, 2) == 'l|' then return end + end + redis.call('DEL', bk) +end + +local function addLoc(bk, tier, size) + if redis.call('EXISTS', bk) == 0 then + redis.call('HSET', bk, '_created', now, '_lacc', now, '_acnt', 0) + end + local field = nodePfx .. tier + if redis.call('HEXISTS', bk, field) == 0 then + redis.call('HSET', bk, field, size) + redis.call('SADD', revidx, bk) + end +end + +local function removeLoc(bk, tier) + local field = nodePfx .. tier + if redis.call('HDEL', bk, field) == 1 then + local flds = redis.call('HKEYS', bk) + local still = false + for _, f in ipairs(flds) do + if string.sub(f, 1, pfxLen) == nodePfx then still = true break end + end + if not still then redis.call('SREM', revidx, bk) end + cleanupEmpty(bk) + end +end + +if full == 1 then + local members = redis.call('SMEMBERS', revidx) + for _, bk in ipairs(members) do + local flds = redis.call('HKEYS', bk) + for _, f in ipairs(flds) do + if string.sub(f, 1, pfxLen) == nodePfx then redis.call('HDEL', bk, f) end + end + cleanupEmpty(bk) + end + redis.call('DEL', revidx) + for i = 0, nev - 1 do + local base = 6 + i * 4 + if tonumber(ARGV[base]) == 0 then addLoc(ARGV[base + 1], ARGV[base + 2], ARGV[base + 3]) end + end +else + for i = 0, nev - 1 do + local base = 6 + i * 4 + if tonumber(ARGV[base]) == 0 then + addLoc(ARGV[base + 1], ARGV[base + 2], ARGV[base + 3]) + else + removeLoc(ARGV[base + 1], ARGV[base + 2]) + end + end +end +return 'OK' +)LUA"; + +// wipe_node_blocks (one shard's instance): +// ARGV = [revidx_key, node_prefix] +// Drains the (node, shard) reverse index, deletes the node's location fields +// from each block, drops now-empty blocks, and deletes the reverse index. +// Idempotent. Returns 1. +inline const std::string kWipeNodeBlocksLua = R"LUA(--!df flags=allow-undeclared-keys +local revidx = ARGV[1] +local nodePfx = ARGV[2] +local pfxLen = string.len(nodePfx) +local members = redis.call('SMEMBERS', revidx) +for _, bk in ipairs(members) do + local flds = redis.call('HKEYS', bk) + for _, f in ipairs(flds) do + if string.sub(f, 1, pfxLen) == nodePfx then redis.call('HDEL', bk, f) end + end + local flds2 = redis.call('HKEYS', bk) + local anyLoc = false + for _, f in ipairs(flds2) do + if string.sub(f, 1, 2) == 'l|' then anyLoc = true break end + end + if not anyLoc then redis.call('DEL', bk) end +end +redis.call('DEL', revidx) +return 1 +)LUA"; + +// unregister_control (control instance): +// KEYS[1] = node key; ARGV = [tag, node_id, wipe_extkv] +// Drops the client record + nodes:alive + alive_peers. External-kv is wiped +// inline only when wipe_extkv == 1 (num_shards == 1: extkv on the control tag); +// for num_shards > 1 the store wipes sharded extkv per shard afterwards. Block +// locations are wiped separately per shard. Returns 1 if it existed. +inline const std::string kUnregisterControlLua = R"LUA(--!df flags=allow-undeclared-keys +local nodeKey = KEYS[1] +local tag = ARGV[1] +local nodeId = ARGV[2] +local wipeExtkv = tonumber(ARGV[3]) +if redis.call('EXISTS', nodeKey) == 0 then return 0 end +redis.call('DEL', nodeKey) +redis.call('SREM', tag .. ':nodes:alive', nodeId) +redis.call('HDEL', tag .. ':alive_peers', nodeId) +if wipeExtkv == 1 then + local exNode = tag .. ':extkv:node:' .. nodeId + local ekeys = redis.call('SMEMBERS', exNode) + for _, ekey in ipairs(ekeys) do + redis.call('HDEL', ekey, nodeId) + if redis.call('HLEN', ekey) == 0 then redis.call('DEL', ekey) end + end + redis.call('DEL', exNode) +end +return 1 +)LUA"; + +// expire_control (control instance): +// ARGV = [tag, cutoff_ms, wipe_extkv] +// Flips ALIVE->EXPIRED for nodes whose last_hb < cutoff, drops them from +// nodes:alive/alive_peers, and (wipe_extkv == 1) their extkv inline; returns +// the dead node ids. Block locations (and sharded extkv when wipe_extkv == 0) +// are wiped separately per shard by the caller. +inline const std::string kExpireControlLua = R"LUA(--!df flags=allow-undeclared-keys +local tag = ARGV[1] +local cutoff = tonumber(ARGV[2]) +local wipeExtkv = tonumber(ARGV[3]) +local members = redis.call('SMEMBERS', tag .. ':nodes:alive') +local dead = {} +for _, id in ipairs(members) do + local nk = tag .. ':node:' .. id + local status = tonumber(redis.call('HGET', nk, 'status') or '0') + local lastHb = tonumber(redis.call('HGET', nk, 'last_hb') or '0') + if status == 1 and lastHb < cutoff then + redis.call('HSET', nk, 'status', 2) + redis.call('SREM', tag .. ':nodes:alive', id) + redis.call('HDEL', tag .. ':alive_peers', id) + if wipeExtkv == 1 then + local exNode = tag .. ':extkv:node:' .. id + local ekeys = redis.call('SMEMBERS', exNode) + for _, ekey in ipairs(ekeys) do + redis.call('HDEL', ekey, id) + if redis.call('HLEN', ekey) == 0 then redis.call('DEL', ekey) end + end + redis.call('DEL', exNode) + end + dead[#dead + 1] = id + end +end +return dead +)LUA"; + +// ===================================================================== +// External-KV / hit-count / eviction scripts (Phase 2, #5). +// +// Placement (design §4, handoff): extkv + hit + the hit reverse-index all live +// under the CONTROL tag {umbp:}, so each is one single-slot Lua on the +// control instance in every mode. The tier-set is a bitmask int per node (bit == +// 1< / hit: are placed in the hash's OWN shard slot +// (KeySchema::ExtKv/Hit use ShardOf(hash)), not the single control slot. This +// lets the external-KV hot path spread across shards / instances / Dragonfly +// proactor threads instead of piling every match + hit-count write onto one +// slot. For num_shards == 1 the shard tag IS the control tag, so the key strings +// are byte-identical to the legacy layout. +// +// The per-(node, shard) reverse index (KeySchema::ExtKvNode(node, shard)) stores +// the FULL extkv: KEY as each member (like NodeBlocks stores full block +// keys), and the per-shard hit index (KeySchema::HitIndex(shard)) stores the +// FULL hit: KEY. That lets every wipe / GC script HDEL/DEL a member +// directly with no in-Lua key composition, and — crucially — lets the hot +// register/unregister write scripts take ONLY [node_id, tier] as shared ARGV +// (all per-hash data is in KEYS[]), so a whole batch fans out one single-slot +// script per shard with identical ARGV via EvalPipeline. +// +// The multi-key mutations that used to hang off the control slot (register, +// unregister, match, hit-count GC, node wipe) are therefore driven by the store +// as one single-slot script PER TOUCHED SHARD, grouped + fanned out by +// RunShardedRead exactly like the block read hot path. Cross-shard atomicity is +// not needed (each hash is independent). num_shards == 1 collapses to one group +// on the control instance — one script, still atomic, byte-identical keys. +// ===================================================================== + +// register_external_kv (single-shard atomic path, num_shards == 1): alive-gate + +// write in one script. Keys declared (no flag; Dragonfly lock-ahead). +// KEYS[1] = node key (alive-check), KEYS[2] = extkv reverse index, +// KEYS[3 .. 2+k] = extkv: per hash. ARGV = [node_id, tier] +// Only if the node is ALIVE (status == 1) does it OR `tier`'s bit into +// extkv:[node] and add the extkv KEY to the reverse index. Idempotent; +// the reverse-index SADD is skipped when the node already holds the hash at +// some tier (mask != 0 ⇒ already a member). Returns 1 if alive, 0 if not. +inline const std::string kRegisterExternalKvLua = R"LUA( +local nodeKey = KEYS[1] +local revidx = KEYS[2] +local nodeId = ARGV[1] +local tier = tonumber(ARGV[2]) +if redis.call('HGET', nodeKey, 'status') ~= '1' then return 0 end +local bit = 2 ^ tier +for i = 3, #KEYS do + local ekey = KEYS[i] + local mask = tonumber(redis.call('HGET', ekey, nodeId) or '0') + if mask == 0 then + redis.call('HSET', ekey, nodeId, bit) + redis.call('SADD', revidx, ekey) + elseif math.floor(mask / bit) % 2 == 0 then + redis.call('HSET', ekey, nodeId, mask + bit) + end +end +return 1 +)LUA"; + +// register_external_kv_write (per-shard, num_shards > 1): the write half of the +// hot register, WITHOUT the alive-gate (the caller checked node status on the +// control instance first). One invocation per touched shard, routed to that +// shard. Keys declared (no flag; Dragonfly runs each shard's script on its own +// proactor thread concurrently). +// KEYS[1] = extkv reverse index for (node, shard), KEYS[2..] = extkv:. +// ARGV = [node_id, tier] (uniform across shards → EvalPipeline-friendly) +inline const std::string kRegisterExternalKvWriteLua = R"LUA( +local revidx = KEYS[1] +local nodeId = ARGV[1] +local tier = tonumber(ARGV[2]) +local bit = 2 ^ tier +for i = 2, #KEYS do + local ekey = KEYS[i] + local mask = tonumber(redis.call('HGET', ekey, nodeId) or '0') + if mask == 0 then + redis.call('HSET', ekey, nodeId, bit) + redis.call('SADD', revidx, ekey) + elseif math.floor(mask / bit) % 2 == 0 then + redis.call('HSET', ekey, nodeId, mask + bit) + end +end +return 1 +)LUA"; + +// unregister_external_kv (per-shard): clears `tier`'s bit for each (node, hash). +// When a hash's bitmask for the node reaches 0 the node field is dropped (and the +// extkv: key + reverse-index membership with it). No liveness check. Keys +// declared (no flag). One invocation per touched shard. +// KEYS[1] = extkv reverse index, KEYS[2..] = extkv:. ARGV = [node_id, tier] +inline const std::string kUnregisterExternalKvLua = R"LUA( +local revidx = KEYS[1] +local nodeId = ARGV[1] +local tier = tonumber(ARGV[2]) +local bit = 2 ^ tier +for i = 2, #KEYS do + local ekey = KEYS[i] + local mask = tonumber(redis.call('HGET', ekey, nodeId) or '0') + if math.floor(mask / bit) % 2 == 1 then + local newmask = mask - bit + if newmask == 0 then + redis.call('HDEL', ekey, nodeId) + if redis.call('HLEN', ekey) == 0 then redis.call('DEL', ekey) end + redis.call('SREM', revidx, ekey) + else + redis.call('HSET', ekey, nodeId, newmask) + end + end +end +return 1 +)LUA"; + +// unregister_external_kv_by_tier (per-shard): clears `tier` from every hash the +// node registered in this shard (admin whole-tier wipe). Enumerates via the +// reverse index whose members are full extkv keys (undeclared but same slot, so +// the allow-undeclared-keys flag only matters on Dragonfly — cold path). +// KEYS[1] = extkv reverse index for (node, shard). ARGV = [node_id, tier] +inline const std::string kUnregisterExternalKvByTierLua = R"LUA(--!df flags=allow-undeclared-keys +local revidx = KEYS[1] +local nodeId = ARGV[1] +local tier = tonumber(ARGV[2]) +local bit = 2 ^ tier +local ekeys = redis.call('SMEMBERS', revidx) +for _, ekey in ipairs(ekeys) do + local mask = tonumber(redis.call('HGET', ekey, nodeId) or '0') + if math.floor(mask / bit) % 2 == 1 then + local newmask = mask - bit + if newmask == 0 then + redis.call('HDEL', ekey, nodeId) + if redis.call('HLEN', ekey) == 0 then redis.call('DEL', ekey) end + redis.call('SREM', revidx, ekey) + else + redis.call('HSET', ekey, nodeId, newmask) + end + end +end +return 1 +)LUA"; + +// unregister_external_kv_by_node / cascade wipe (per-shard): drops every +// external-kv entry for the node in this shard (all tiers) and deletes the +// reverse index. Members are full extkv keys. Idempotent. +// KEYS[1] = extkv reverse index for (node, shard). ARGV = [node_id] +inline const std::string kUnregisterExternalKvByNodeLua = R"LUA(--!df flags=allow-undeclared-keys +local revidx = KEYS[1] +local nodeId = ARGV[1] +local ekeys = redis.call('SMEMBERS', revidx) +for _, ekey in ipairs(ekeys) do + redis.call('HDEL', ekey, nodeId) + if redis.call('HLEN', ekey) == 0 then redis.call('DEL', ekey) end +end +redis.call('DEL', revidx) +return 1 +)LUA"; + +// match_external_kv (per-shard): THE external-KV hot read. Keys are declared in +// KEYS[] so the script touches no undeclared key and needs no +// allow-undeclared-keys flag — on Dragonfly each shard's script runs on its own +// proactor thread (lock-ahead) instead of taking a store-wide GLOBAL lock. +// +// count == 0: KEYS = [extkv:h1 .. extkv:hk] ARGV = [0, now] +// count == 1: KEYS = [extkv:h1 .. extkv:hk, hit:h1 .. hit:hk, hitidx] ARGV = [1, now] +// Returns an array PARALLEL to the k extkv keys: element i = flat +// HGETALL(extkv:) ([] when the hash has no registered node). The caller +// scatters element i back to hash i (RunShardedRead) and decodes each node's +// tier bitmask. +// +// Hit path (count == 1) costs 3 redis.call() per matched hash steady state +// instead of the old 5: HGETALL + HINCRBY + HSET(ls); the SADD into the shard's +// hit index fires ONLY on the counter's first hit (HINCRBY returns 1), and ls is +// stamped with an unconditional HSET (the old read-then-guard HGET is dropped — a +// backward clock skew merely makes an entry look marginally less recent to the +// coarse hit-GC, benign). The server is redis.call()-count bound, so this cut is +// a direct throughput win on top of the cross-shard parallelism. +inline const std::string kMatchExternalKvLua = R"LUA( +local countHit = tonumber(ARGV[1]) +local now = tonumber(ARGV[2]) +local k = #KEYS +if countHit == 1 then k = (k - 1) / 2 end +local out = {} +for i = 1, k do + local flat = redis.call('HGETALL', KEYS[i]) + out[i] = flat + if countHit == 1 and #flat > 0 then + local hkey = KEYS[k + i] + local c = redis.call('HINCRBY', hkey, 'c', 1) + redis.call('HSET', hkey, 'ls', now) + if c == 1 then redis.call('SADD', KEYS[2 * k + 1], hkey) end + end +end +return out +)LUA"; + +// get_external_kv_hit_counts (per-shard): keys declared (no flag). Returns an +// array PARALLEL to KEYS; element i = the count string of hit: (nil if no +// counter). The caller scatters back to hash i. +// KEYS = [hit:h1 .. hit:hk] ARGV = [] +inline const std::string kGetHitCountsLua = R"LUA( +local out = {} +for i = 1, #KEYS do + out[i] = redis.call('HGET', KEYS[i], 'c') +end +return out +)LUA"; + +// garbage_collect_hits (per-shard): drops every hit counter in this shard whose +// ls < cutoff, via the shard's hit index (members are full hit keys; no SCAN). +// Returns the number dropped. Cold path (GC timer). +// KEYS[1] = hit index for this shard. ARGV = [cutoff_ms] +inline const std::string kGarbageCollectHitsLua = R"LUA(--!df flags=allow-undeclared-keys +local idx = KEYS[1] +local cutoff = tonumber(ARGV[1]) +local hkeys = redis.call('SMEMBERS', idx) +local dropped = 0 +for _, hkey in ipairs(hkeys) do + local ls = tonumber(redis.call('HGET', hkey, 'ls') or '0') + if ls < cutoff then + redis.call('DEL', hkey) + redis.call('SREM', idx, hkey) + dropped = dropped + 1 + end +end +return dropped +)LUA"; + +// enumerate_eviction (per shard): +// KEYS[1..k] = per-node block reverse-index sets to scan (control-tag +// node::blocks in single mode, or this shard's shard-tag +// node::blocks in split modes). Members are full block-hash keys. +// ARGV = [now_ms, n_wanted, node1, tier1, node2, tier2, ...] +// For each set, walks its block hashes, skips leased ones (_lease > now), and +// for every location field l|node|tier whose (node,tier) is wanted emits a +// flat 5-tuple [block_key, node, tier, size, last_accessed_ms]. Returns one +// flat array per KEYS entry (parallel to KEYS). Ordering / per-bucket cap are +// applied by the caller in C++ (an eviction tick is seconds, not a hot path). +inline const std::string kEnumerateEvictionLua = R"LUA(--!df flags=allow-undeclared-keys +local now = tonumber(ARGV[1]) +local nw = tonumber(ARGV[2]) +local wanted = {} +for i = 1, nw do + local node = ARGV[2 + i * 2 - 1] + local tier = ARGV[2 + i * 2] + wanted[node] = wanted[node] or {} + wanted[node][tier] = true +end +local out = {} +for ki = 1, #KEYS do + local members = redis.call('SMEMBERS', KEYS[ki]) + local cands = {} + for _, bk in ipairs(members) do + local flat = redis.call('HGETALL', bk) + local lease = 0 + local lacc = 0 + for j = 1, #flat, 2 do + local f = flat[j] + if f == '_lease' then + lease = tonumber(flat[j + 1]) or 0 + elseif f == '_lacc' then + lacc = tonumber(flat[j + 1]) or 0 + end + end + if lease <= now then + for j = 1, #flat, 2 do + local f = flat[j] + if string.sub(f, 1, 2) == 'l|' then + local rest = string.sub(f, 3) + local sep = string.find(rest, '|', 1, true) + if sep ~= nil then + local node = string.sub(rest, 1, sep - 1) + local tier = string.sub(rest, sep + 1) + if wanted[node] and wanted[node][tier] then + cands[#cands + 1] = bk + cands[#cands + 1] = node + cands[#cands + 1] = tier + cands[#cands + 1] = flat[j + 1] + cands[#cands + 1] = tostring(lacc) + end + end + end + end + end + end + out[ki] = cands +end +return out +)LUA"; + +} // namespace mori::umbp::redis diff --git a/src/umbp/include/umbp/distributed/master/redis/resp_client.h b/src/umbp/include/umbp/distributed/master/redis/resp_client.h new file mode 100644 index 000000000..e7ed1ba03 --- /dev/null +++ b/src/umbp/include/umbp/distributed/master/redis/resp_client.h @@ -0,0 +1,184 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. + +// RespClient — a thin, RESP-protocol-compatible client seam for the master's +// Redis metadata backend. +// +// This is deliberately the ONLY file in the store that knows which client +// library is used underneath. Phase 1 implements it on hiredis (the library +// already present in the build image); the store code above it depends only on +// the small RespValue / RespClient surface here, so a raw-hiredis client can be +// swapped for redis-plus-plus (or a cluster client) without touching the store. +// +// It speaks only the portable RESP subset (STRING/HASH/SET/ZSET commands plus +// EVAL/EVALSHA/SCRIPT LOAD), so the same binary connects unchanged to Redis, +// Dragonfly, and Valkey — selection is connection config only. +// +// Threading: a fixed-size connection pool fronts a shared instance. The gRPC +// handler thread pool calls Command/Pipeline/Eval concurrently; each call +// borrows one pooled connection for its duration. + +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "umbp/distributed/master/redis/resp_detail.h" +#include "umbp/distributed/master/redis/resp_value.h" + +// Forward-declare the hiredis context so this header stays library-agnostic to +// its includers (only resp_client.cpp includes ). +struct redisContext; + +namespace mori::umbp::redis { + +// RespError and RespValue live in resp_value.h so the redis-plus-plus cluster +// client can share them without pulling in hiredis. + +class RespClient : public IRespClient { + public: + struct Options { + // e.g. "tcp://127.0.0.1:6379". Only tcp is supported in Phase 1. + std::string uri = "tcp://127.0.0.1:6379"; + std::string password; + int connect_timeout_ms = 1000; + int socket_timeout_ms = 1000; + std::size_t pool_size = 8; + }; + + explicit RespClient(Options options); + ~RespClient(); + + RespClient(const RespClient&) = delete; + RespClient& operator=(const RespClient&) = delete; + + // One command. `args` is the full argv (command name first); values are + // binary-safe. Returns the server's reply (which may be an Error value). + RespValue Command(const std::vector& args) override; + + // Pipeline: append every command, then read all replies in order. One round + // trip for the whole batch. (Not part of IRespClient; single-endpoint only.) + // NOTE: currently has no in-tree caller — retained as the single-endpoint + // general-purpose batch primitive (the hot path uses EvalPipeline instead). If + // it stays unused long-term, consider dropping it. + std::vector Pipeline(const std::vector>& commands); + + // EVALSHA with transparent SCRIPT LOAD + NOSCRIPT fallback. `script` is the + // Lua body; its SHA is loaded once and cached. `keys` then `args` are passed + // to the script as KEYS[] / ARGV[]. + RespValue Eval(const std::string& script, const std::vector& keys, + const std::vector& args) override; + + // Pipeline the SAME script over several KEYS groups in one round trip. Call i + // runs `script` with KEYS = keys_per_call[i] and ARGV = shared_args; the + // returned vector is parallel to keys_per_call. Used by the sharded read path + // to fan one batch out to one single-slot EVALSHA per shard. + // + // On NOSCRIPT the script is reloaded and ONLY the affected calls are retried, + // so a script with side effects (e.g. route_get_batch's lease/access bump) is + // never applied twice for a call that already ran. + std::vector EvalPipeline(const std::string& script, + const std::vector>& keys_per_call, + const std::vector& shared_args) override; + + // Liveness probe (PING). Returns false if no connection can be established. + bool Ping() override; + + void SetMetrics(mori::metrics::MetricsServer* metrics) override { metrics_ = metrics; } + + const Options& options() const { return options_; } + + // Convert a raw redisReply (void* to avoid leaking the hiredis type into this + // header) into a RespValue. Public + static so the redis-plus-plus cluster + // client can reuse the exact same reply decoding (its ReplyUPtr is a + // hiredis redisReply under the hood). + static RespValue Convert(void* reply); + + private: + // RAII borrow of one pooled connection. + class Lease { + public: + Lease(RespClient* owner, redisContext* ctx) : owner_(owner), ctx_(ctx) {} + ~Lease(); + Lease(const Lease&) = delete; + Lease& operator=(const Lease&) = delete; + redisContext* get() const { return ctx_; } + void MarkBroken() { healthy_ = false; } + + private: + RespClient* owner_; + redisContext* ctx_; + bool healthy_ = true; + }; + + redisContext* Connect(); // create + AUTH; throws on failure + Lease Acquire(); // borrow (blocks if pool exhausted) + void Release(redisContext* ctx, bool healthy); + + // Run one argv command on a specific connection; sets *broke on transport + // failure. Returns a RespValue (Error type on server error). + RespValue RunArgv(redisContext* ctx, const std::vector& args, bool* broke); + + // Append one EVALSHA (using `sha`) per index in `indices`, then read the + // replies back in order into (*replies)[idx]. Returns the subset of `indices` + // whose reply was a NOSCRIPT error, so the caller can reload + retry just + // those. Sets *broke on transport failure. + std::vector PipelineEvalshaBatch( + redisContext* ctx, const std::string& sha, + const std::vector>& keys_per_call, + const std::vector& shared_args, const std::vector& indices, + std::vector* replies, bool* broke); + + // SHA for `script`, loaded + cached on first use (deterministic across + // connections, so one cache is correct). + std::string GetOrLoadSha(redisContext* ctx, const std::string& script, bool* broke); + + // Cold-path metric helpers (no-op when metrics_ is null). Kept off the hot + // success path — called only on a transport failure / NOSCRIPT reload. + void CountTransportError(); + void CountNoscriptReload(); + + Options options_; + std::string host_; + int port_ = 6379; + + std::mutex mu_; + std::condition_variable cv_; + std::vector idle_; + std::size_t created_ = 0; + + // Optional cold-path metrics sink; null = no metrics. Set once at startup. + mori::metrics::MetricsServer* metrics_ = nullptr; + + // EVALSHA SHA cache, keyed by script identity (see ScriptCache / lua_scripts.h + // for the pointer-identity contract: never pass a freshly-constructed string). + ScriptCache script_cache_; +}; + +} // namespace mori::umbp::redis diff --git a/src/umbp/include/umbp/distributed/master/redis/resp_cluster_client.h b/src/umbp/include/umbp/distributed/master/redis/resp_cluster_client.h new file mode 100644 index 000000000..80b18b1cd --- /dev/null +++ b/src/umbp/include/umbp/distributed/master/redis/resp_cluster_client.h @@ -0,0 +1,136 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. + +// RespClusterClient — the IRespClient implementation for Redis Cluster, built on +// redis-plus-plus (sw::redis::RedisCluster). +// +// It routes each command / script to the node owning the relevant hash-tag slot +// and delegates MOVED / ASK redirection, slot-map refresh, and master-failover +// reconnect to redis-plus-plus (its RedisCluster::command redirection loop). The +// store programs against IRespClient, so single / multi-endpoint (hiredis +// RespClient) and cluster (this class) share the same hot-path logic. +// +// Every scripted call must carry at least one key (KEYS[0]) so we can route it; +// all keys touched by one script share a hash tag (single slot), which is what +// keeps the script cross-key-atomic and cluster-legal. +// +// redis-plus-plus (and hiredis) are confined to resp_cluster_client.cpp; this +// header forward-declares RedisCluster so includers stay library-agnostic. + +#pragma once + +#include +#include +#include +#include +#include +#include + +#include "umbp/distributed/master/redis/cluster_slots.h" +#include "umbp/distributed/master/redis/resp_detail.h" +#include "umbp/distributed/master/redis/resp_value.h" + +namespace sw::redis { +class RedisCluster; +} + +namespace mori::umbp::redis { + +class ThreadPool; + +class RespClusterClient : public IRespClient { + public: + struct Options { + // Cluster seeds, e.g. {"tcp://127.0.0.1:7000", ...}. Only one reachable seed + // is needed; redis-plus-plus discovers the rest via CLUSTER SLOTS. We try + // seeds in order until one connects. + std::vector seeds; + std::string password; + int connect_timeout_ms = 1000; + int socket_timeout_ms = 1000; + std::size_t pool_size = 8; + }; + + explicit RespClusterClient(Options options); + ~RespClusterClient() override; + + RespClusterClient(const RespClusterClient&) = delete; + RespClusterClient& operator=(const RespClusterClient&) = delete; + + // Connect and return each master's owned slot ranges (from CLUSTER SLOTS), one + // vector per master, ordered by the master's lowest slot (stable across + // restarts). Used for balanced placement: pick one block-shard tag whose slot + // lands on each master so a RouteGet batch spreads evenly. Throws RespError if + // no seed is reachable. + static std::vector> DiscoverMasterSlotRanges(const Options& options); + + // Single command; routed by the key at args[1] (all store commands put the key + // there: HGETALL/HGET/SCARD ...). Returns the reply (Error value on a + // server error), throws RespError on a transport/redirection failure. + RespValue Command(const std::vector& args) override; + + // EVALSHA routed by keys[0]'s slot, with a transparent SCRIPT LOAD + NOSCRIPT + // retry (a newly-promoted replica or resharded node may not have the script + // cached). keys must be non-empty. The SHA is loaded once and cached. + RespValue Eval(const std::string& script, const std::vector& keys, + const std::vector& args) override; + + // Run `script` over several single-slot KEYS groups, each routed by its own + // keys[0]. Groups are issued concurrently through a small worker pool (each is + // an independent redirection-aware call), so a batch spanning several nodes is + // ~one round trip per node rather than N sequential ones. + std::vector EvalPipeline(const std::string& script, + const std::vector>& keys_per_call, + const std::vector& shared_args) override; + + bool Ping() override; + + void SetMetrics(mori::metrics::MetricsServer* metrics) override { metrics_ = metrics; } + + private: + // Send `argv` to the node owning `routing_key`'s slot via redis-plus-plus's + // redirection loop (handles MOVED/ASK/failover), decode the reply. Server + // errors become an Error RespValue; transport failures throw RespError. + RespValue RunArgvRouted(const std::vector& argv, const std::string& routing_key); + + // EVALSHA of `script` routed by keys[0], reloading the script on the target + // node and retrying once on NOSCRIPT. keys must be non-empty. + RespValue EvalShaRouted(const std::string& script, const std::vector& keys, + const std::vector& args); + + // SHA1 of `script`, obtained via SCRIPT LOAD once and cached (the digest is + // content-addressed, so it is the same on every node). + std::string GetOrLoadSha(const std::string& script); + + Options options_; + std::unique_ptr cluster_; + std::unique_ptr pool_; + + // Optional cold-path metrics sink (transport errors, NOSCRIPT reloads); null = + // no metrics. Set once at startup. + mori::metrics::MetricsServer* metrics_ = nullptr; + + // EVALSHA SHA cache, keyed by script identity (see ScriptCache / lua_scripts.h). + ScriptCache script_cache_; +}; + +} // namespace mori::umbp::redis diff --git a/src/umbp/include/umbp/distributed/master/redis/resp_detail.h b/src/umbp/include/umbp/distributed/master/redis/resp_detail.h new file mode 100644 index 000000000..c29ccf385 --- /dev/null +++ b/src/umbp/include/umbp/distributed/master/redis/resp_detail.h @@ -0,0 +1,104 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. + +// Small helpers shared by both RESP client implementations (RespClient on +// hiredis, RespClusterClient on redis-plus-plus), so the SHA cache, the EVALSHA +// argv layout, and the const-char*/len splitting live in ONE place instead of +// being copy-pasted per client. Header-only (all tiny / on the hot path). + +#pragma once + +#include +#include +#include +#include +#include + +namespace mori::umbp::redis { + +// Content-addressed EVALSHA SHA cache, keyed by the script object's ADDRESS (not +// its ~1.5KB text): every Lua script is a single `inline const std::string` +// (lua_scripts.h), so `&script` is stable and unique per script. This keeps the +// read hot path from rehashing/comparing the whole script body under the lock on +// every EVALSHA. Both clients own one of these. +class ScriptCache { + public: + // Return the cached SHA for `script`, or run `loader(script)` once (outside the + // lock) to obtain it, cache it, and return it. `loader` does the SCRIPT LOAD + // and may throw (transport error / bad reply); nothing is cached if it throws. + template + std::string GetOrLoad(const std::string& script, Loader&& loader) { + const void* key = static_cast(&script); + { + std::lock_guard lk(mu_); + auto it = cache_.find(key); + if (it != cache_.end()) return it->second; + } + std::string sha = loader(script); + { + std::lock_guard lk(mu_); + cache_[key] = sha; + } + return sha; + } + + // Drop `script`'s cached SHA (after a NOSCRIPT, so the next call reloads it). + void Invalidate(const std::string& script) { + std::lock_guard lk(mu_); + cache_.erase(static_cast(&script)); + } + + private: + std::mutex mu_; + std::unordered_map cache_; // &script -> sha1 +}; + +// Split `args` into the parallel const-char*/length arrays hiredis' *Argv APIs +// want. `ptrs`/`lens` are cleared and refilled; they must outlive the call that +// consumes them (they alias into `args`). +inline void ToArgv(const std::vector& args, std::vector& ptrs, + std::vector& lens) { + ptrs.clear(); + lens.clear(); + ptrs.reserve(args.size()); + lens.reserve(args.size()); + for (const auto& a : args) { + ptrs.push_back(a.data()); + lens.push_back(a.size()); + } +} + +// Build the argv for one EVALSHA: [EVALSHA, sha, nkeys, keys..., args...]. +inline std::vector BuildEvalshaArgv(const std::string& sha, + const std::vector& keys, + const std::vector& args) { + std::vector cmd; + cmd.reserve(3 + keys.size() + args.size()); + cmd.push_back("EVALSHA"); + cmd.push_back(sha); + cmd.push_back(std::to_string(keys.size())); + for (const auto& k : keys) cmd.push_back(k); + for (const auto& a : args) cmd.push_back(a); + return cmd; +} + +} // namespace mori::umbp::redis diff --git a/src/umbp/include/umbp/distributed/master/redis/resp_value.h b/src/umbp/include/umbp/distributed/master/redis/resp_value.h new file mode 100644 index 000000000..e2982c972 --- /dev/null +++ b/src/umbp/include/umbp/distributed/master/redis/resp_value.h @@ -0,0 +1,113 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. + +// RespValue / RespError / IRespClient — the library-agnostic surface the store +// depends on, split out of resp_client.h so a client that is NOT built on raw +// hiredis (e.g. the redis-plus-plus cluster client) can implement the same +// interface without pulling in hiredis. +// +// IRespClient is the small seam the RedisMasterMetadataStore programs against. +// Two implementations live behind it: +// * RespClient (resp_client.h) — raw hiredis, single endpoint; +// used for single / multi-endpoint modes. +// * RespClusterClient (resp_cluster_client.h) — redis-plus-plus RedisCluster; +// slot routing + MOVED/ASK/failover. +// +// Error model: server-returned errors (e.g. a Lua runtime error, NOSCRIPT) are +// surfaced as a RespValue with type == Error so callers can inspect the message +// (error-as-value). Only transport-level failures throw RespError. + +#pragma once + +#include +#include +#include + +// Forward-declared so the interface can carry an optional metrics sink without +// this lightweight header pulling in the Prometheus server. +namespace mori::metrics { +class MetricsServer; +} + +namespace mori::umbp::redis { + +// Thrown on a transport-level failure (connect/socket error, protocol error). +// Command errors returned BY the server are NOT thrown — they surface as a +// RespValue with type == Error. +class RespError : public std::runtime_error { + public: + explicit RespError(const std::string& what) : std::runtime_error(what) {} +}; + +// Owned, recursive mirror of a redisReply. Owning it (rather than exposing the +// raw redisReply*) keeps the client library out of every includer of this +// header. +struct RespValue { + enum class Type { Nil, Status, Error, Integer, String, Array }; + + Type type = Type::Nil; + long long integer = 0; // Integer + std::string str; // Status / Error / String (binary-safe) + std::vector elements; // Array + + bool is_nil() const { return type == Type::Nil; } + bool is_error() const { return type == Type::Error; } + bool is_array() const { return type == Type::Array; } + bool ok() const { return type != Type::Error; } +}; + +// The small RESP surface the store programs against. Implementations own their +// own connection management (pool, cluster slot map, ...); the store only sees +// these four operations, so single-endpoint, multi-endpoint, and cluster modes +// share the same hot-path logic. +class IRespClient { + public: + virtual ~IRespClient() = default; + + // One command (full argv, command name first; binary-safe). Returns the + // server reply (may be an Error value). + virtual RespValue Command(const std::vector& args) = 0; + + // EVALSHA with transparent SCRIPT LOAD + NOSCRIPT fallback. `keys` then `args` + // become KEYS[] / ARGV[]. `keys` must carry at least one key per touched slot + // so a cluster implementation can route the script. + virtual RespValue Eval(const std::string& script, const std::vector& keys, + const std::vector& args) = 0; + + // Run `script` over several KEYS groups; returned vector is parallel to + // keys_per_call. A cluster implementation may internally regroup calls by node + // (different groups can land on different slots/nodes) and fan them out. + virtual std::vector EvalPipeline( + const std::string& script, const std::vector>& keys_per_call, + const std::vector& shared_args) = 0; + + // Liveness probe. Returns false if the endpoint(s) cannot be reached. + virtual bool Ping() = 0; + + // Attach an optional metrics sink for COLD-PATH counters/histograms only + // (transport errors, NOSCRIPT reloads, connection-pool waits). Default no-op. + // Set once at startup, before serving traffic; the hot success path records + // nothing (every emit is gated on a non-null sink and only on rare paths). + virtual void SetMetrics(mori::metrics::MetricsServer* /*metrics*/) {} +}; + +} // namespace mori::umbp::redis diff --git a/src/umbp/include/umbp/distributed/master/redis/sharded_read.h b/src/umbp/include/umbp/distributed/master/redis/sharded_read.h new file mode 100644 index 000000000..edc7da207 --- /dev/null +++ b/src/umbp/include/umbp/distributed/master/redis/sharded_read.h @@ -0,0 +1,199 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. + +// Sharded batch-read fan-out for the block hot path (RouteGet / Exists). +// +// Split out of redis_master_metadata_store.cpp so the batch bucketing + fan-out + +// fault-tolerance + reply-scatter logic — the trickiest, most order-sensitive +// part of the backend — can be unit-tested directly against a fake IRespClient, +// without a live Redis. The store just supplies the KeySchema, the routing +// (client-per-group), and the per-key decoder. + +#pragma once + +#include +#include +#include +#include +#include +#include + +#include "mori/metrics/prometheus_metrics_server.hpp" +#include "mori/utils/mori_log.hpp" +#include "umbp/distributed/master/redis/key_schema.h" +#include "umbp/distributed/master/redis/resp_value.h" +#include "umbp/distributed/master/redis/thread_pool.h" + +namespace mori::umbp::redis { + +// A batch's block keys bucketed by shard, plus the mapping to scatter each +// shard's reply back to the caller's original key order. Only shards the batch +// actually touches get a group, so a single-shard batch is one group. +struct ShardedBatch { + std::vector> keys_by_shard; // group -> block keys + std::vector> orig_index_by_shard; // group -> caller indices + std::vector shard_of_group; // group -> shard index +}; + +// Bucket `user_keys` by shard, composing the full block key for each and +// remembering each key's original position for the scatter step. +inline ShardedBatch GroupKeysByShard(const KeySchema& schema, + const std::vector& user_keys) { + ShardedBatch batch; + // shard index -> group slot, lazily assigned so we only build touched shards. + std::vector group_of_shard(schema.NumShards(), -1); + for (std::size_t i = 0; i < user_keys.size(); ++i) { + const std::size_t shard = schema.ShardOf(user_keys[i]); + int& group = group_of_shard[shard]; + if (group < 0) { + group = static_cast(batch.keys_by_shard.size()); + batch.keys_by_shard.emplace_back(); + batch.orig_index_by_shard.emplace_back(); + batch.shard_of_group.push_back(shard); + } + batch.keys_by_shard[group].push_back(schema.Block(user_keys[i])); + batch.orig_index_by_shard[group].push_back(i); + } + return batch; +} + +// One instance's slice of a sharded batch: which groups it serves, the block +// keys per group, and (after the call) that instance's replies or an error. +struct ShardCall { + IRespClient* client = nullptr; + std::vector groups; + std::vector> keys; + std::vector replies; + bool ok = true; + std::string error; +}; + +// Run `script` over a sharded batch, routing each shard-group to the client +// returned by client_for_group(group) — one EvalPipeline per distinct client. +// A single-endpoint deployment is one pipelined round trip on the calling +// thread (no pool use). A multi-endpoint one issues each instance's pipeline +// CONCURRENTLY via `pool` (one round trip instead of N — the win on high-RTT +// remote stores), then scatters replies on the calling thread so on_key needs +// no locking. +// +// Fault tolerance (tolerate_shard_failures): when the keyspace is spread across +// multiple nodes/instances (multi-endpoint OR cluster), a single node that is +// down / erroring does NOT fail the whole batch — its keys are simply left +// untouched (the caller's default is a miss: empty locations / exists=false), so +// RouteGet keeps serving keys on the healthy shards. A WARN is logged per failed +// shard. In single-endpoint mode there is nothing to fall back to, so a failure +// propagates (a total outage must surface, not masquerade as misses). +// +// NOTE: this must be decided by the caller (topology), NOT inferred from the +// number of distinct clients — in cluster mode a single cluster client fronts +// every node, so `calls.size()` is always 1 even though the keys are spread +// across nodes and per-node failures SHOULD degrade to misses. +template +void RunShardedRead(const ShardedBatch& batch, const std::string& script, + const std::vector& shared_args, const char* method, + ThreadPool* pool, bool tolerate_shard_failures, + mori::metrics::MetricsServer* metrics, ClientForGroup client_for_group, + Fn&& on_key) { + // Cold path only: count a shard whose keys we degraded to miss (node down / + // script error). No-op when no metrics sink is attached. + auto count_degraded = [&] { + if (metrics == nullptr) return; + metrics->addCounter("mori_umbp_redis_degraded_shard_total", + "Shard reads degraded to miss because a node was down / erroring", + {{"backend", "redis"}, {"method", method}}); + }; + // Bucket group indices by the client that serves them. + std::unordered_map> groups_by_client; + for (std::size_t g = 0; g < batch.keys_by_shard.size(); ++g) { + groups_by_client[client_for_group(g)].push_back(g); + } + + std::vector calls; + calls.reserve(groups_by_client.size()); + for (auto& [client, group_indices] : groups_by_client) { + ShardCall c; + c.client = client; + c.groups = std::move(group_indices); + c.keys.reserve(c.groups.size()); + for (std::size_t g : c.groups) c.keys.push_back(batch.keys_by_shard[g]); + calls.push_back(std::move(c)); + } + + // Capture a transport failure into the ShardCall rather than throwing, so one + // dead instance can be tolerated below instead of aborting the fan-out. + auto do_eval = [&](ShardCall& c) { + try { + c.replies = c.client->EvalPipeline(script, c.keys, shared_args); + } catch (const std::exception& e) { + c.ok = false; + c.error = e.what(); + } + }; + + if (calls.size() <= 1 || pool == nullptr) { + for (auto& c : calls) do_eval(c); // single instance / no pool: inline. + } else { + // Fan the extra instances out to the pool, run the first inline, then join + // every future so no task outlives this scope (do_eval never throws). + std::vector> futs; + futs.reserve(calls.size() - 1); + for (std::size_t i = 1; i < calls.size(); ++i) { + futs.push_back(pool->Enqueue([&do_eval, &calls, i] { do_eval(calls[i]); })); + } + do_eval(calls[0]); + for (auto& f : futs) f.get(); + } + + const bool tolerate = + tolerate_shard_failures; // multi-endpoint/cluster: degrade a down shard to misses. + + // Scatter replies back to caller order (single-threaded; on_key is unlocked). + for (const ShardCall& c : calls) { + if (!c.ok) { + if (!tolerate) + throw std::runtime_error(std::string("[RedisStore] ") + method + ": " + c.error); + MORI_UMBP_WARN("[RedisStore] {}: shard instance unavailable, its keys read as miss: {}", + method, c.error); + count_degraded(); + continue; // leave this shard's keys at the caller's default (miss). + } + for (std::size_t k = 0; k < c.groups.size() && k < c.replies.size(); ++k) { + const RespValue& reply = c.replies[k]; + if (reply.is_error()) { + if (!tolerate) { + throw std::runtime_error(std::string("[RedisStore] ") + method + ": " + reply.str); + } + MORI_UMBP_WARN("[RedisStore] {}: shard script error, its keys read as miss: {}", method, + reply.str); + count_degraded(); + continue; + } + if (!reply.is_array()) continue; + const auto& orig_indices = batch.orig_index_by_shard[c.groups[k]]; + for (std::size_t j = 0; j < orig_indices.size() && j < reply.elements.size(); ++j) { + on_key(orig_indices[j], reply.elements[j]); + } + } + } +} + +} // namespace mori::umbp::redis diff --git a/src/umbp/include/umbp/distributed/master/redis/thread_pool.h b/src/umbp/include/umbp/distributed/master/redis/thread_pool.h new file mode 100644 index 000000000..f36bcba21 --- /dev/null +++ b/src/umbp/include/umbp/distributed/master/redis/thread_pool.h @@ -0,0 +1,109 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. + +// A tiny fixed-size worker pool used to issue a multi-endpoint fan-out's +// per-instance Redis calls concurrently, so a RouteGet's latency is one round +// trip instead of N sequential ones (matters on high-RTT remote stores). It is +// deliberately minimal: persistent workers (no per-call thread churn — that was +// measured to cost more than it saves), a mutex+cv task queue, and futures that +// carry a task's exception back to the caller. +// +// Deadlock-free for the fan-out use case: worker threads only RUN tasks (a +// blocking Redis call) and never wait on other tasks; only the caller waits on +// its futures. An undersized pool merely reduces parallelism (tasks queue), it +// cannot deadlock. + +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace mori::umbp::redis { + +class ThreadPool { + public: + explicit ThreadPool(std::size_t num_threads) { + if (num_threads == 0) num_threads = 1; + workers_.reserve(num_threads); + for (std::size_t i = 0; i < num_threads; ++i) { + workers_.emplace_back([this] { WorkerLoop(); }); + } + } + + ~ThreadPool() { + { + std::lock_guard lk(mu_); + stop_ = true; + } + cv_.notify_all(); + for (auto& w : workers_) { + if (w.joinable()) w.join(); + } + } + + ThreadPool(const ThreadPool&) = delete; + ThreadPool& operator=(const ThreadPool&) = delete; + + std::size_t size() const { return workers_.size(); } + + // Enqueue a task; the returned future becomes ready when it runs and rethrows + // any exception the task threw. + std::future Enqueue(std::function task) { + std::packaged_task pt(std::move(task)); + std::future fut = pt.get_future(); + { + std::lock_guard lk(mu_); + tasks_.push(std::move(pt)); + } + cv_.notify_one(); + return fut; + } + + private: + void WorkerLoop() { + for (;;) { + std::packaged_task task; + { + std::unique_lock lk(mu_); + cv_.wait(lk, [this] { return stop_ || !tasks_.empty(); }); + if (stop_ && tasks_.empty()) return; + task = std::move(tasks_.front()); + tasks_.pop(); + } + task(); + } + } + + std::vector workers_; + std::queue> tasks_; + std::mutex mu_; + std::condition_variable cv_; + bool stop_ = false; +}; + +} // namespace mori::umbp::redis diff --git a/src/umbp/include/umbp/distributed/master/redis_master_metadata_store.h b/src/umbp/include/umbp/distributed/master/redis_master_metadata_store.h new file mode 100644 index 000000000..ec1589d1f --- /dev/null +++ b/src/umbp/include/umbp/distributed/master/redis_master_metadata_store.h @@ -0,0 +1,232 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. + +// RedisMasterMetadataStore — a RESP-protocol-compatible IMasterMetadataStore. +// +// Makes the master stateless: all client/block metadata lives in an external +// RESP store (Redis / Dragonfly / Valkey), so any number of master replicas can +// serve traffic and a crashed master reads the full picture back on restart. +// +// PHASE 1 SCOPE (this file): the six hot-path methods that decide the whole +// backend's performance are implemented against real Lua/pipelines, plus the +// handful of methods the running master calls on background threads +// (UnregisterClient, ExpireStaleClients, GarbageCollectHits) and the cheap +// single-command client reads. The external-KV methods and the eviction +// candidate enumeration are Phase 2 and throw std::logic_error if called; they +// are not exercised by the default hot-path benchmark. See +// design-redis-metadata-store.md. + +#pragma once + +#include +#include +#include +#include + +#include "umbp/distributed/master/master_metadata_store.h" +#include "umbp/distributed/master/redis/key_schema.h" +#include "umbp/distributed/master/redis/resp_client.h" +#include "umbp/distributed/master/redis/thread_pool.h" +#include "umbp/distributed/types.h" + +namespace mori::umbp { + +class RedisMasterMetadataStore : public IMasterMetadataStore { + public: + struct Config { + std::string uri = "tcp://127.0.0.1:6379"; + // Multi-endpoint mode: one Redis instance per entry, block shards spread one + // per instance so their scripts run on independent server threads (the only + // way past a single instance's single-thread ceiling). Empty or size 1 => + // single-endpoint mode using `uri`. clients_[0] is the control instance + // (client records, alive set, peer view, extkv all live there). + std::vector shard_uris; + std::string namespace_id = "default"; + std::string password; + int connect_timeout_ms = 1000; + int socket_timeout_ms = 1000; + std::size_t pool_size = 8; + // Single-endpoint only: number of hash-tag shards the block keyspace is + // spread over. 1 keeps the legacy single-tag layout (byte-identical keys, + // whole-batch-atomic reads). In multi-endpoint mode the shard count is fixed + // to the number of endpoints (one shard per instance). In cluster mode it is + // the number of block shard tags spread (by CRC16 slot) across the cluster's + // nodes. See KeySchema. + std::size_t block_shards = 1; + // Redis Cluster mode: one redis-plus-plus RedisCluster client routes every + // command/script to the node owning its hash-tag slot, with MOVED/ASK and + // master-failover handled by the client. Mutually exclusive with shard_uris. + // cluster_seeds are bootstrap nodes (any reachable one; the rest are + // discovered via CLUSTER SLOTS). + bool cluster = false; + std::vector cluster_seeds; + // Cluster balanced placement: explicit block-shard tags (one per master, + // each on a distinct node's slot), computed by the factory from CLUSTER + // SLOTS. When set (cluster mode), these replace the formulaic {umbp::bS} + // tags so a RouteGet batch spreads evenly across nodes. Empty => formulaic + // tags sized by block_shards. + std::vector cluster_block_tags; + }; + + explicit RedisMasterMetadataStore(const Config& config); + ~RedisMasterMetadataStore() override = default; + + RedisMasterMetadataStore(const RedisMasterMetadataStore&) = delete; + RedisMasterMetadataStore& operator=(const RedisMasterMetadataStore&) = delete; + + // --- Cross-store writes --- + bool RegisterClient(const ClientRegistration& registration, + std::chrono::system_clock::time_point now, + std::chrono::system_clock::duration stale_after) override; + void UnregisterClient(const std::string& node_id) override; + HeartbeatResult ApplyHeartbeat(const std::string& node_id, uint64_t seq, + std::chrono::system_clock::time_point now, + const std::map& caps, + const std::vector& events, bool is_full_sync) override; + std::vector ExpireStaleClients( + std::chrono::system_clock::time_point cutoff) override; + + // --- External-KV writes (Phase 2) --- + bool RegisterExternalKvIfAlive(const std::string& node_id, const std::vector& hashes, + TierType tier) override; + void UnregisterExternalKv(const std::string& node_id, const std::vector& hashes, + TierType tier) override; + void UnregisterExternalKvByTier(const std::string& node_id, TierType tier) override; + void UnregisterExternalKvByNode(const std::string& node_id) override; + std::size_t GarbageCollectHits(std::chrono::system_clock::time_point cutoff) override; + + // --- Block reads --- + std::vector LookupBlock(const std::string& key) const override; + std::vector LookupBlockForRouteGet( + const std::string& key, const std::unordered_set& exclude_nodes, + std::chrono::system_clock::time_point now, + std::chrono::system_clock::duration lease_duration) override; + std::vector> BatchLookupBlockForRouteGet( + const std::vector& keys, const std::unordered_set& exclude_nodes, + std::chrono::system_clock::time_point now, + std::chrono::system_clock::duration lease_duration) override; + std::vector BatchExistsBlock(const std::vector& keys) const override; + std::map> EnumerateEvictionCandidates( + const std::vector& buckets, EvictionOrder order, size_t max_per_bucket, + std::chrono::system_clock::time_point now) const override; + + // --- Client reads --- + std::optional GetClient(const std::string& node_id) const override; + bool IsClientAlive(const std::string& node_id) const override; + std::optional GetPeerAddress(const std::string& node_id) const override; + std::vector ListAliveClients() const override; + std::unordered_map GetAlivePeerView() const override; + std::size_t AliveClientCount() const override; + std::vector GetClientTags(const std::string& node_id) const override; + + // --- External-KV reads (Phase 2) --- + std::vector MatchExternalKv(const std::vector& hashes, bool count_as_hit, + std::chrono::system_clock::time_point now) override; + std::vector GetExternalKvHitCounts( + const std::vector& hashes) const override; + std::size_t GetExternalKvCount(const std::string& node_id) const override; + + // Attach the master's Prometheus server so hot ops export + // mori_umbp_store_op_latency_seconds{op,backend="redis"}. Null (default) = + // no metrics (e.g. the standalone microbench). + void SetMetricsSink(mori::metrics::MetricsServer* metrics) override { + metrics_ = metrics; + // Propagate to the underlying RESP clients so their cold-path counters + // (transport errors, NOSCRIPT reloads, pool waits) export too. Clients are + // built in the ctor, before this is called, so wire the sink in here. + for (auto& c : clients_) c->SetMetrics(metrics); + } + + // Best-effort connectivity probe (PING). Every endpoint must answer. + bool Ping() const { + for (const auto& c : clients_) { + if (!c->Ping()) return false; + } + return true; + } + + private: + // How many block shards to build for a config (one per endpoint in + // multi-endpoint mode, else Config::block_shards). + static std::size_t ResolveNumShards(const Config& config); + + // Build the KeySchema: explicit per-master tags for cluster balanced placement + // (Config::cluster_block_tags), else the formulaic shard tags. + static redis::KeySchema BuildKeySchema(const Config& config); + + // Deployment topology. kSingle = one instance, one atomic script (legacy). + // kMulti = one instance per block shard (UMBP_REDIS_SHARD_URIS). kCluster = + // one Redis Cluster client routing by slot. kMulti and kCluster both use the + // split control-script + per-shard-block-script write path (their control and + // block keys live in different slots), so gate that on split_writes(), NOT on + // the endpoint count (cluster is a single client but still needs the split). + enum class Mode { kSingle, kMulti, kCluster }; + bool split_writes() const { return mode_ != Mode::kSingle; } + + // Whether the external-KV keyspace (extkv: / hit: / hit index / + // per-node reverse index) is spread across more than one shard tag. When true, + // every external-KV op fans out one single-slot script per touched shard + // (grouped by KeySchema::ExtKvShardOf), and the client-record cascades wipe the + // node's sharded extkv separately (WipeNodeExtKvMulti) instead of inline. When + // false (num_shards == 1) the extkv keyspace collapses onto the control slot — + // one atomic script, byte-identical to the legacy layout. This is keyed on the + // shard count, NOT split_writes(): a single Dragonfly instance runs with + // num_shards > 1 (kSingle) precisely so extkv spreads across its proactor + // threads. + bool extkv_sharded() const { return keys_.NumShards() > 1; } + + bool multi_endpoint() const { return clients_.size() > 1; } + redis::IRespClient& control() const { return *clients_[0]; } + std::size_t endpoint_of_shard(std::size_t shard) const { return shard % clients_.size(); } + redis::IRespClient& client_for_shard(std::size_t shard) const { + return *clients_[endpoint_of_shard(shard)]; + } + + // Multi-endpoint write helpers (see .cpp). Each runs the per-shard block + // script on the shard's own instance; idempotent so retries are safe. + void ApplyBlockEventsMulti(const std::string& node_id, const std::vector& events, + bool is_full_sync, std::chrono::system_clock::time_point now); + void WipeNodeBlocksMulti(const std::string& node_id); + + // Drop a node's sharded external-kv on every shard (best-effort per shard, + // mirroring WipeNodeBlocksMulti). Used by the client-record cascades and + // UnregisterExternalKvByNode when extkv_sharded(). One single-slot script per + // shard, routed to that shard's instance/slot. + void WipeNodeExtKvMulti(const std::string& node_id); + + redis::KeySchema keys_; + Mode mode_ = Mode::kSingle; + // clients_[0] is the control instance; clients_[s] backs block shard s. In + // cluster mode there is a single entry (the cluster client) that internally + // routes every shard by slot. Held as IRespClient so the same store logic + // drives the hiredis single-node client (single / multi-endpoint) and the + // redis-plus-plus cluster client. mutable so const read methods can borrow a + // pooled connection. + mutable std::vector> clients_; + // Workers that issue a multi-endpoint read fan-out's per-instance calls + // concurrently (one round trip instead of N). Null in single-endpoint mode. + std::unique_ptr fanout_pool_; + // Optional Prometheus sink for per-op latency; null = no metrics. + mori::metrics::MetricsServer* metrics_ = nullptr; +}; + +} // namespace mori::umbp diff --git a/src/umbp/tools/redis/docker-compose.yml b/src/umbp/tools/redis/docker-compose.yml new file mode 100644 index 000000000..2407f087a --- /dev/null +++ b/src/umbp/tools/redis/docker-compose.yml @@ -0,0 +1,145 @@ +# Independent RESP store containers for the UMBP master Redis-metadata backend. +# +# Run these ON THE HOST, next to the (host-networked) app container. Every +# service uses `network_mode: host`, so the app container reaches them at +# 127.0.0.1: with no port mapping — the same URIs the tests/bench already +# use. This is the productized replacement for building redis-server from source +# inside the app container (see run_local_backends.sh, kept only as the +# no-docker fallback). +# +# Profiles (one topology at a time): +# single -> Redis 7 on 6379 (default; existing tests/bench) +# dragonfly -> Dragonfly on 6380 (multi-threaded RESP; portability) +# cluster -> 3 masters + 3 replicas on 7000-7005 (Redis Cluster: failover/HA) +# +# Use the store.sh wrapper, or directly: +# docker compose -f docker-compose.yml --profile single up -d +# docker compose -f docker-compose.yml --profile cluster up -d # cluster-init forms it +# docker compose -f docker-compose.yml --profile cluster down -v +# +# NOTE: cluster nodes are announced at 127.0.0.1 (single-host verification). For +# a multi-host cluster, set CLUSTER_ANNOUNCE_IP to the host's routable IP so +# nodes advertise a reachable address. + +name: umbp-store + +x-cluster-node: &cluster-node + image: redis:7-alpine + network_mode: host + restart: unless-stopped + profiles: ["cluster"] + +services: + # ---- single Redis (default) -------------------------------------------- + redis: + image: redis:7-alpine + network_mode: host + profiles: ["single"] + restart: unless-stopped + command: + ["redis-server", "--port", "6379", + "--appendonly", "yes", "--appendfsync", "everysec", "--save", ""] + healthcheck: + test: ["CMD", "redis-cli", "-p", "6379", "ping"] + interval: 2s + timeout: 2s + retries: 10 + + # ---- Dragonfly (multi-threaded RESP) ----------------------------------- + dragonfly: + image: docker.dragonflydb.io/dragonflydb/dragonfly:latest + network_mode: host + profiles: ["dragonfly"] + restart: unless-stopped + ulimits: + memlock: -1 + # NO global --default_lua_flags: it would force EVERY Lua script (incl. the + # read hot path) into Dragonfly's global-transaction mode (store-wide lock), + # serializing all proactor threads and erasing the block-sharding win. The + # write/control scripts instead carry a first-line "--!df flags=..." directive + # (a plain Lua comment on Redis), and the read scripts declare all keys via + # KEYS[], so route_get_batch/exists_batch run per-shard in parallel. See + # src/umbp/include/umbp/distributed/master/redis/lua_scripts.h. + command: + ["dragonfly", "--port", "6380", "--logtostderr"] + healthcheck: + test: ["CMD", "redis-cli", "-p", "6380", "ping"] + interval: 2s + timeout: 2s + retries: 10 + + # ---- Redis Cluster: 3 masters + 3 replicas ----------------------------- + redis-c0: + <<: *cluster-node + command: &cluster-cmd + ["redis-server", "--port", "7000", "--cluster-enabled", "yes", + "--cluster-config-file", "/tmp/nodes-7000.conf", + "--cluster-node-timeout", "5000", + "--cluster-announce-ip", "${CLUSTER_ANNOUNCE_IP:-127.0.0.1}", + "--appendonly", "no", "--save", "", "--dir", "/tmp"] + redis-c1: + <<: *cluster-node + command: + ["redis-server", "--port", "7001", "--cluster-enabled", "yes", + "--cluster-config-file", "/tmp/nodes-7001.conf", + "--cluster-node-timeout", "5000", + "--cluster-announce-ip", "${CLUSTER_ANNOUNCE_IP:-127.0.0.1}", + "--appendonly", "no", "--save", "", "--dir", "/tmp"] + redis-c2: + <<: *cluster-node + command: + ["redis-server", "--port", "7002", "--cluster-enabled", "yes", + "--cluster-config-file", "/tmp/nodes-7002.conf", + "--cluster-node-timeout", "5000", + "--cluster-announce-ip", "${CLUSTER_ANNOUNCE_IP:-127.0.0.1}", + "--appendonly", "no", "--save", "", "--dir", "/tmp"] + redis-c3: + <<: *cluster-node + command: + ["redis-server", "--port", "7003", "--cluster-enabled", "yes", + "--cluster-config-file", "/tmp/nodes-7003.conf", + "--cluster-node-timeout", "5000", + "--cluster-announce-ip", "${CLUSTER_ANNOUNCE_IP:-127.0.0.1}", + "--appendonly", "no", "--save", "", "--dir", "/tmp"] + redis-c4: + <<: *cluster-node + command: + ["redis-server", "--port", "7004", "--cluster-enabled", "yes", + "--cluster-config-file", "/tmp/nodes-7004.conf", + "--cluster-node-timeout", "5000", + "--cluster-announce-ip", "${CLUSTER_ANNOUNCE_IP:-127.0.0.1}", + "--appendonly", "no", "--save", "", "--dir", "/tmp"] + redis-c5: + <<: *cluster-node + command: + ["redis-server", "--port", "7005", "--cluster-enabled", "yes", + "--cluster-config-file", "/tmp/nodes-7005.conf", + "--cluster-node-timeout", "5000", + "--cluster-announce-ip", "${CLUSTER_ANNOUNCE_IP:-127.0.0.1}", + "--appendonly", "no", "--save", "", "--dir", "/tmp"] + + # One-shot: wait for all nodes, then form the cluster (idempotent). + cluster-init: + image: redis:7-alpine + network_mode: host + profiles: ["cluster"] + restart: "no" + depends_on: [redis-c0, redis-c1, redis-c2, redis-c3, redis-c4, redis-c5] + environment: + CLUSTER_ANNOUNCE_IP: "${CLUSTER_ANNOUNCE_IP:-127.0.0.1}" + command: + - sh + - -c + - | + set -e + IP="${CLUSTER_ANNOUNCE_IP:-127.0.0.1}" + for p in 7000 7001 7002 7003 7004 7005; do + until redis-cli -h "$$IP" -p "$$p" ping 2>/dev/null | grep -q PONG; do sleep 1; done + done + if redis-cli -h "$$IP" -p 7000 cluster info 2>/dev/null | grep -q "cluster_state:ok"; then + echo "[cluster-init] cluster already formed"; exit 0 + fi + redis-cli --cluster create \ + "$$IP:7000" "$$IP:7001" "$$IP:7002" "$$IP:7003" "$$IP:7004" "$$IP:7005" \ + --cluster-replicas 1 --cluster-yes + redis-cli -h "$$IP" -p 7000 cluster info | grep cluster_state diff --git a/src/umbp/tools/redis/run_local_backends.sh b/src/umbp/tools/redis/run_local_backends.sh new file mode 100755 index 000000000..10ea1adbf --- /dev/null +++ b/src/umbp/tools/redis/run_local_backends.sh @@ -0,0 +1,173 @@ +#!/usr/bin/env bash +# FALLBACK launcher — prefer store.sh (docker compose) for the productized path. +# +# This launches single-node Redis and Dragonfly as local PROCESSES (source-built +# / downloaded binaries, no docker), for environments without host docker access +# (e.g. working purely inside an app container). The productized way to run the +# stores as independent containers is `store.sh {up|down|status} ` +# (docker-compose.yml), which also covers the Redis Cluster topology this script +# does not. store.sh auto-falls-back to this script for single/dragonfly when +# docker is unavailable. +# +# Both speak RESP, so the master connects to either via UMBP_REDIS_URI only. +# +# Redis -> tcp://127.0.0.1:6379 +# Dragonfly -> tcp://127.0.0.1:6380 +# +# Getting the binaries (in order of preference, auto-selected): +# Redis: redis-server on PATH -> prebuilt in RUN_DIR -> apt-get -> +# build the official source with make (works when the apt mirror +# is blocked but github over HTTPS is reachable, e.g. this CI image). +# Dragonfly: prebuilt in RUN_DIR -> download the official release tarball. +# Both are the stock upstream binaries; only how they are obtained varies. +# +# Usage: +# src/umbp/tools/redis/run_local_backends.sh up # install (if needed) + start +# src/umbp/tools/redis/run_local_backends.sh down # stop +# src/umbp/tools/redis/run_local_backends.sh status +set -euo pipefail + +RUN_DIR="${UMBP_REDIS_RUN_DIR:-/tmp/umbp_redis_bench}" +REDIS_PORT="${UMBP_REDIS_PORT:-6379}" +DF_PORT="${UMBP_DRAGONFLY_PORT:-6380}" +DF_VERSION="${UMBP_DRAGONFLY_VERSION:-v1.23.2}" +DF_BIN="${RUN_DIR}/dragonfly" +REDIS_VERSION="${UMBP_REDIS_VERSION:-7.2.5}" +REDIS_SRC="${RUN_DIR}/redis-src" + +# Resolved by ensure_redis() / find_redis_cli(). +REDIS_SERVER="" +REDIS_CLI="" + +mkdir -p "${RUN_DIR}" + +# Locate a redis-cli without triggering an install (used by status/down). +find_redis_cli() { + if command -v redis-cli >/dev/null 2>&1; then + command -v redis-cli + elif [[ -x "${RUN_DIR}/redis-cli" ]]; then + echo "${RUN_DIR}/redis-cli" + fi +} + +build_redis_from_source() { + if ! command -v git >/dev/null 2>&1 || ! command -v make >/dev/null 2>&1; then + echo "[run_local_backends] cannot build redis: need git + make" >&2 + return 1 + fi + echo "[run_local_backends] building official redis ${REDIS_VERSION} from source..." + if [[ ! -d "${REDIS_SRC}" ]]; then + git clone --depth 1 --branch "${REDIS_VERSION}" https://github.com/redis/redis.git "${REDIS_SRC}" + fi + make -C "${REDIS_SRC}" -j"$(nproc)" BUILD_TLS=no USE_SYSTEMD=no MALLOC=libc + cp "${REDIS_SRC}/src/redis-server" "${RUN_DIR}/redis-server" + cp "${REDIS_SRC}/src/redis-cli" "${RUN_DIR}/redis-cli" +} + +ensure_redis() { + # 1) Already installed on PATH. + if command -v redis-server >/dev/null 2>&1 && command -v redis-cli >/dev/null 2>&1; then + REDIS_SERVER="$(command -v redis-server)" + REDIS_CLI="$(command -v redis-cli)" + return 0 + fi + # 2) Previously built/downloaded into RUN_DIR. + if [[ -x "${RUN_DIR}/redis-server" && -x "${RUN_DIR}/redis-cli" ]]; then + REDIS_SERVER="${RUN_DIR}/redis-server" + REDIS_CLI="${RUN_DIR}/redis-cli" + return 0 + fi + # 3) apt fast path (skipped automatically when the mirror is unreachable). + if command -v apt-get >/dev/null 2>&1; then + echo "[run_local_backends] trying apt-get install redis-server..." + if apt-get update -qq >/dev/null 2>&1 && apt-get install -y -qq redis-server >/dev/null 2>&1; then + REDIS_SERVER="$(command -v redis-server)" + REDIS_CLI="$(command -v redis-cli)" + return 0 + fi + echo "[run_local_backends] apt unavailable; falling back to source build" + fi + # 4) Build the official source (apt mirror blocked but github reachable). + build_redis_from_source + REDIS_SERVER="${RUN_DIR}/redis-server" + REDIS_CLI="${RUN_DIR}/redis-cli" +} + +ensure_dragonfly() { + if [[ -x "${DF_BIN}" ]]; then return 0; fi + echo "[run_local_backends] downloading dragonfly ${DF_VERSION}..." + local arch tarball url + arch="$(uname -m)" + case "${arch}" in + x86_64) tarball="dragonfly-x86_64.tar.gz" ;; + aarch64) tarball="dragonfly-aarch64.tar.gz" ;; + *) echo "[run_local_backends] unsupported arch ${arch}"; return 1 ;; + esac + url="https://github.com/dragonflydb/dragonfly/releases/download/${DF_VERSION}/${tarball}" + curl -fsSL "${url}" -o "${RUN_DIR}/${tarball}" + tar -xzf "${RUN_DIR}/${tarball}" -C "${RUN_DIR}" + # The tarball unpacks to dragonfly-; normalize to ${DF_BIN}. + local extracted + extracted="$(find "${RUN_DIR}" -maxdepth 1 -name 'dragonfly-*' -type f | head -1)" + [[ -n "${extracted}" ]] && cp "${extracted}" "${DF_BIN}" + chmod +x "${DF_BIN}" +} + +up() { + ensure_redis + ensure_dragonfly || echo "[run_local_backends] WARN: dragonfly unavailable; redis only" + + if ! "${REDIS_CLI}" -p "${REDIS_PORT}" ping >/dev/null 2>&1; then + echo "[run_local_backends] starting redis-server on ${REDIS_PORT}" + "${REDIS_SERVER}" --port "${REDIS_PORT}" --daemonize yes \ + --appendonly yes --appendfsync everysec \ + --dir "${RUN_DIR}" --logfile "${RUN_DIR}/redis.log" \ + --save "" + fi + + if [[ -x "${DF_BIN}" ]]; then + if ! "${REDIS_CLI}" -p "${DF_PORT}" ping >/dev/null 2>&1; then + echo "[run_local_backends] starting dragonfly on ${DF_PORT}" + # NO global --default_lua_flags: it forces every Lua script (incl. the read + # hot path) into Dragonfly's global-transaction mode (store-wide lock), + # serializing all proactor threads and erasing the block-sharding win. The + # write/control scripts carry a per-script "--!df flags=allow-undeclared-keys" + # directive (a no-op Lua comment on Redis) while the read scripts declare + # all keys via KEYS[], so they run per-shard in parallel. See lua_scripts.h. + "${DF_BIN}" --port "${DF_PORT}" --logtostderr --alsologtostderr=false \ + --dir "${RUN_DIR}" >"${RUN_DIR}/dragonfly.log" 2>&1 & + echo $! >"${RUN_DIR}/dragonfly.pid" + sleep 1 + fi + fi + status +} + +down() { + echo "[run_local_backends] stopping backends" + local cli + cli="$(find_redis_cli)" + if [[ -n "${cli}" ]]; then + "${cli}" -p "${REDIS_PORT}" shutdown nosave >/dev/null 2>&1 || true + fi + if [[ -f "${RUN_DIR}/dragonfly.pid" ]]; then + kill "$(cat "${RUN_DIR}/dragonfly.pid")" >/dev/null 2>&1 || true + rm -f "${RUN_DIR}/dragonfly.pid" + fi +} + +status() { + local cli + cli="$(find_redis_cli)" + echo -n "[run_local_backends] redis(${REDIS_PORT}): " + { [[ -n "${cli}" ]] && "${cli}" -p "${REDIS_PORT}" ping 2>/dev/null; } || echo "DOWN" + echo -n "[run_local_backends] dragonfly(${DF_PORT}): " + { [[ -n "${cli}" ]] && "${cli}" -p "${DF_PORT}" ping 2>/dev/null; } || echo "DOWN" +} + +case "${1:-up}" in + up) up ;; + down) down ;; + status) status ;; + *) echo "usage: $0 {up|down|status}"; exit 2 ;; +esac diff --git a/src/umbp/tools/redis/store.sh b/src/umbp/tools/redis/store.sh new file mode 100755 index 000000000..fc5cf1e5f --- /dev/null +++ b/src/umbp/tools/redis/store.sh @@ -0,0 +1,106 @@ +#!/usr/bin/env bash +# Unified entrypoint to bring RESP stores up/down for the UMBP master Redis +# metadata backend. Run this ON THE HOST, next to the host-networked app +# container; every store uses host networking, so the app container reaches it +# at 127.0.0.1:. +# +# Usage: +# store.sh up [single|dragonfly|cluster] # default: single +# store.sh down [single|dragonfly|cluster] +# store.sh status [single|dragonfly|cluster] +# store.sh seeds [single|dragonfly|cluster] # print the URI(s) to export +# +# Topologies: +# single -> UMBP_REDIS_URI=tcp://127.0.0.1:6379 +# dragonfly -> UMBP_REDIS_URI=tcp://127.0.0.1:6380 +# cluster -> UMBP_REDIS_URI=tcp://127.0.0.1:7000,...,7005 (+ UMBP_REDIS_CLUSTER=1) +# +# Docker (compose) is the productized path. When docker is unavailable (e.g. +# working purely inside an app container without host docker access), single and +# dragonfly fall back to run_local_backends.sh (source-built local processes); +# cluster has no process fallback here — use a host with docker. +set -euo pipefail + +HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +COMPOSE_FILE="${HERE}/docker-compose.yml" +ANNOUNCE_IP="${CLUSTER_ANNOUNCE_IP:-127.0.0.1}" + +ACTION="${1:-}" +TOPO="${2:-single}" + +have_docker() { command -v docker >/dev/null 2>&1 && docker info >/dev/null 2>&1; } + +compose() { CLUSTER_ANNOUNCE_IP="${ANNOUNCE_IP}" docker compose -f "${COMPOSE_FILE}" "$@"; } + +# Throwaway redis-cli on the host network (no host redis-cli needed). +rcli() { docker run --rm --network host redis:7-alpine redis-cli "$@"; } + +cluster_ports() { echo 7000 7001 7002 7003 7004 7005; } + +cluster_seeds() { + local s="" p + for p in $(cluster_ports); do s+="${s:+,}tcp://${ANNOUNCE_IP}:${p}"; done + echo "$s" +} + +print_seeds() { + case "$TOPO" in + single) echo "export UMBP_METADATA_BACKEND=redis UMBP_REDIS_URI=tcp://127.0.0.1:6379" ;; + dragonfly) echo "export UMBP_METADATA_BACKEND=redis UMBP_REDIS_URI=tcp://127.0.0.1:6380" ;; + cluster) echo "export UMBP_METADATA_BACKEND=redis UMBP_REDIS_CLUSTER=1 UMBP_REDIS_URI=$(cluster_seeds)" ;; + *) echo "unknown topology: $TOPO" >&2; exit 2 ;; + esac +} + +wait_cluster_ok() { + echo "[store] waiting for cluster_state:ok ..." + for _ in $(seq 1 60); do + if rcli -h "${ANNOUNCE_IP}" -p 7000 cluster info 2>/dev/null | grep -q "cluster_state:ok"; then + echo "[store] cluster is ready"; return 0 + fi + sleep 1 + done + echo "[store] WARN: cluster not ready after 60s; check 'store.sh status cluster'" >&2 + return 1 +} + +up() { + if have_docker; then + compose --profile "$TOPO" up -d + [ "$TOPO" = "cluster" ] && wait_cluster_ok || true + else + case "$TOPO" in + single|dragonfly) + echo "[store] docker unavailable; falling back to local processes (run_local_backends.sh)" + "${HERE}/run_local_backends.sh" up ;; + *) echo "[store] docker unavailable and no process fallback for '$TOPO'" >&2; exit 1 ;; + esac + fi + echo "[store] ready. To point the master/tests at it:" + print_seeds +} + +down() { + if have_docker; then + compose --profile "$TOPO" down -v --remove-orphans + else + "${HERE}/run_local_backends.sh" down || true + fi +} + +status() { + if have_docker; then + compose --profile "$TOPO" ps + [ "$TOPO" = "cluster" ] && rcli -h "${ANNOUNCE_IP}" -p 7000 cluster info 2>/dev/null | grep -E "cluster_state|cluster_known_nodes|cluster_size" || true + else + "${HERE}/run_local_backends.sh" status || true + fi +} + +case "$ACTION" in + up) up ;; + down) down ;; + status) status ;; + seeds) print_seeds ;; + *) echo "usage: $0 {up|down|status|seeds} [single|dragonfly|cluster]" >&2; exit 2 ;; +esac diff --git a/tests/cpp/umbp/distributed/CMakeLists.txt b/tests/cpp/umbp/distributed/CMakeLists.txt index 6e8f3da3e..4344bbf17 100644 --- a/tests/cpp/umbp/distributed/CMakeLists.txt +++ b/tests/cpp/umbp/distributed/CMakeLists.txt @@ -365,3 +365,55 @@ target_link_libraries( bench_umbp_kvevent_master_pressure PRIVATE umbp_core umbp_common ${_TEST_PROTOBUF_LIB} ${_TEST_GRPCPP_LIB} gtest_main) + +# bench_master_metadata_store — store-level microbench isolating the hot-path +# IMasterMetadataStore methods (RouteGet / Exists / Heartbeat) across backends +# (UMBP_METADATA_BACKEND). No gRPC / RDMA; links umbp_common only. Standalone +# executable, not a CTest target. +add_executable(bench_umbp_master_metadata_store bench_master_metadata_store.cpp) +target_link_libraries(bench_umbp_master_metadata_store PRIVATE umbp_common) +target_compile_features(bench_umbp_master_metadata_store PRIVATE cxx_std_17) + +# bench_master_metadata_store_extkv — sibling microbench isolating the +# EXTERNAL-KV hot-path IMasterMetadataStore methods (MatchExternalKv / +# RegisterExternalKvIfAlive / UnregisterExternalKv / GetExternalKvHitCounts) +# across backends (UMBP_METADATA_BACKEND). No gRPC / RDMA; links umbp_common +# only. Standalone executable, not a CTest target. +add_executable(bench_umbp_master_metadata_store_extkv + bench_master_metadata_store_extkv.cpp) +target_link_libraries(bench_umbp_master_metadata_store_extkv PRIVATE umbp_common) +target_compile_features(bench_umbp_master_metadata_store_extkv PRIVATE cxx_std_17) + +# test_redis_master_metadata_store — Phase 1 targeted tests for the RESP/Redis +# backend. Built only when the Redis backend is compiled in; the test itself +# skips at runtime when no RESP store is reachable at UMBP_REDIS_URI. +if(USE_REDIS_BACKEND) + add_executable(test_redis_master_metadata_store + test_redis_master_metadata_store.cpp) + target_link_libraries(test_redis_master_metadata_store + PRIVATE umbp_common GTest::gtest_main) + target_compile_features(test_redis_master_metadata_store PRIVATE cxx_std_17) + gtest_discover_tests(test_redis_master_metadata_store) + + # test_sharded_read — pure unit tests for the block-read fan-out + # (redis/sharded_read.h) against a fake IRespClient. No live store needed, so + # it always runs (unlike test_redis_master_metadata_store, which skips without + # a reachable backend). + add_executable(test_sharded_read test_sharded_read.cpp) + target_link_libraries(test_sharded_read PRIVATE umbp_common GTest::gtest_main) + target_compile_features(test_sharded_read PRIVATE cxx_std_17) + gtest_discover_tests(test_sharded_read) + + # redis-plus-plus link/build smoke test. Proves the submodule compiles with + # this toolchain and links (static redis++ + system hiredis) before any store + # code depends on it; at runtime it connects to UMBP_REDIS_URI if reachable, + # else skips. redis++_static carries its own include dirs; hiredis must be + # linked explicitly (the system hiredis ships no CMake config, so redis++ does + # not propagate it for the static lib). + add_executable(test_redispp_smoke test_redispp_smoke.cpp) + target_link_libraries(test_redispp_smoke + PRIVATE redis++_static ${HIREDIS_LIB} GTest::gtest_main) + target_include_directories(test_redispp_smoke PRIVATE ${HIREDIS_INCLUDE_DIR}) + target_compile_features(test_redispp_smoke PRIVATE cxx_std_17) + gtest_discover_tests(test_redispp_smoke) +endif() diff --git a/tests/cpp/umbp/distributed/bench_kvevent_master_pressure.cpp b/tests/cpp/umbp/distributed/bench_kvevent_master_pressure.cpp index c84234126..fa17f6153 100644 --- a/tests/cpp/umbp/distributed/bench_kvevent_master_pressure.cpp +++ b/tests/cpp/umbp/distributed/bench_kvevent_master_pressure.cpp @@ -70,6 +70,8 @@ #include #include #include +#include +#include #include #include #include @@ -282,6 +284,17 @@ struct BenchOpts { bool randomize = false; // mix broadcast/rotate per round uint64_t seed = 1234; bool with_external_kv = false; + // MatchExternalKv hit-count mode. The real hot path (route/prefix lookup) + // counts hits — the increment + last_seen stamp is part of the cost. Default + // true so --with-external-kv measures the true hot-path Match; set false to + // isolate the pure read. + bool ext_kv_count_as_hit = true; + // Isolate the external-KV path at the gRPC+backend layer: skip BatchPut / + // BatchLookup / BatchGet (and thus all RDMA) and exercise ONLY + // ReportExternalKvBlocks (producer) + MatchExternalKv (consumer). Implies + // --with-external-kv. Lets the proxy compare backends on the external-KV hot + // path without the block data plane in the way. + bool ext_kv_only = false; size_t key_space = 0; // 0 = unique keys per round; >0 = recycle per producer int metrics_port = 0; // 0 = no Prometheus scrape double keep_master_secs = 0; // keep master alive at the end for Grafana @@ -289,6 +302,21 @@ struct BenchOpts { // Put placement (injected directly into the master, NOT via env). std::string put_algo = "most_available"; // most_available | random std::string put_affinity = "local"; // none | same | local + // Multi-process / multi-host: if --external-master is set, do NOT spawn an + // in-process MasterServer; connect this process's clients to a standalone + // umbp_master shared by every process/host. node_id_prefix / node_address let + // each process advertise a globally-unique node id + its own reachable IP. + std::string external_master = ""; // "host:port"; empty => in-process master + std::string node_id_prefix = "node-"; // per-process unique prefix; client idx appended + std::string node_address = "127.0.0.1"; + // Cross-process end-of-run barrier (fetch/both only): every client process + // reaches it before any tears down, so peer RDMA endpoints stay alive until all + // readers are done -- eliminating the teardown reset that otherwise aborts a + // straggler. Coordinated via files in a shared directory (shared FS for + // multi-host). Empty dir or size<=1 disables it (falls back to BENCH_DRAIN_MS). + std::string barrier_dir = ""; + size_t barrier_size = 0; // total client processes across all hosts + std::string barrier_tag = ""; // unique per run; isolates the barrier subdir }; void Usage() { @@ -309,11 +337,20 @@ void Usage() { " --randomize (mix broadcast/rotate per round)\n" " --seed N (default 1234)\n" " --with-external-kv (default off)\n" + " --ext-kv-count-as-hit 0|1 (default 1; MatchExternalKv hit-count write)\n" + " --ext-kv-only (only report+match, no BatchPut/Get/RDMA; implies --with-external-kv)\n" " --key-space N (0=unique per round; >0 recycle)\n" " --metrics-port P (0=off; >0 enable master Prometheus + scrape)\n" " --keep-master-secs F (default 0)\n" " --put-algo most_available|random (default most_available)\n" " --put-affinity none|same|local (default local)\n" + " --external-master host:port (connect to a shared standalone master;\n" + " empty => spawn an in-process master)\n" + " --node-id-prefix S (default node-; client index appended)\n" + " --node-address IP (default 127.0.0.1; this process's reachable IP)\n" + " --barrier-dir DIR (cross-process end-of-run barrier dir; shared FS for multi-host)\n" + " --barrier-size N (total client processes to wait for; <=1 disables)\n" + " --barrier-tag S (unique per run; isolates the barrier subdir)\n" "Heartbeat interval is env-driven (UMBP_HEARTBEAT_TTL_SEC,\n" "UMBP_HEARTBEAT_INTERVAL_DIVISOR); launch one process per scenario.\n"); } @@ -376,6 +413,11 @@ bool ParseArgs(int argc, char** argv, BenchOpts* o) { o->seed = std::strtoull(need("--seed"), nullptr, 10); } else if (a == "--with-external-kv") { o->with_external_kv = true; + } else if (a == "--ext-kv-count-as-hit") { + o->ext_kv_count_as_hit = std::strtol(need("--ext-kv-count-as-hit"), nullptr, 10) != 0; + } else if (a == "--ext-kv-only") { + o->ext_kv_only = true; + o->with_external_kv = true; } else if (a == "--key-space") { o->key_space = std::strtoull(need("--key-space"), nullptr, 10); } else if (a == "--metrics-port") { @@ -388,6 +430,18 @@ bool ParseArgs(int argc, char** argv, BenchOpts* o) { o->put_algo = need("--put-algo"); } else if (a == "--put-affinity") { o->put_affinity = need("--put-affinity"); + } else if (a == "--external-master") { + o->external_master = need("--external-master"); + } else if (a == "--node-id-prefix") { + o->node_id_prefix = need("--node-id-prefix"); + } else if (a == "--node-address") { + o->node_address = need("--node-address"); + } else if (a == "--barrier-dir") { + o->barrier_dir = need("--barrier-dir"); + } else if (a == "--barrier-size") { + o->barrier_size = std::strtoull(need("--barrier-size"), nullptr, 10); + } else if (a == "--barrier-tag") { + o->barrier_tag = need("--barrier-tag"); } else if (a == "-h" || a == "--help") { Usage(); std::exit(0); @@ -460,6 +514,7 @@ struct Metrics { std::atomic hit{0}; std::atomic miss{0}; std::atomic rpc_error_keys{0}; + std::atomic fetch_errors{0}; // BatchGet transport failures (peer down); excluded from hit/miss std::atomic ext_report_calls{0}; std::atomic ext_match_calls{0}; std::atomic ext_match_hits{0}; @@ -543,19 +598,23 @@ void Worker(size_t id, size_t round_begin, size_t round_end, bool measure, Ctx c if (IsProducer(pat, id)) { std::vector keys(batch); for (size_t b = 0; b < batch; ++b) keys[b] = MakeKey(id, r, b, o); - const auto t0 = Clock::now(); - auto res = cli->BatchPut(keys, srcs, sizes); - const auto t1 = Clock::now(); - if (measure) { - put_ms.push_back(DurMs(t1 - t0)); - ctx.m->put_calls.fetch_add(1, std::memory_order_relaxed); - ctx.m->put_keys.fetch_add(batch, std::memory_order_relaxed); - size_t fails = 0; - for (bool ok : res) - if (!ok) ++fails; - if (fails) ctx.m->put_fail_keys.fetch_add(fails, std::memory_order_relaxed); + // --ext-kv-only skips the block data plane (BatchPut + RDMA) so the proxy + // measures only the external-KV control path (report + match) end to end. + if (!o.ext_kv_only) { + const auto t0 = Clock::now(); + auto res = cli->BatchPut(keys, srcs, sizes); + const auto t1 = Clock::now(); + if (measure) { + put_ms.push_back(DurMs(t1 - t0)); + ctx.m->put_calls.fetch_add(1, std::memory_order_relaxed); + ctx.m->put_keys.fetch_add(batch, std::memory_order_relaxed); + size_t fails = 0; + for (bool ok : res) + if (!ok) ++fails; + if (fails) ctx.m->put_fail_keys.fetch_add(fails, std::memory_order_relaxed); + } + if (o.mode == "flush") cli->Master().FlushHeartbeat(); } - if (o.mode == "flush") cli->Master().FlushHeartbeat(); if (o.with_external_kv) { const bool ok = cli->ReportExternalKvBlocks(keys, TierType::DRAM); if (measure && ok) ctx.m->ext_report_calls.fetch_add(1, std::memory_order_relaxed); @@ -580,7 +639,7 @@ void Worker(size_t id, size_t round_begin, size_t round_end, bool measure, Ctx c std::vector rkeys(batch); for (size_t b = 0; b < batch; ++b) rkeys[b] = MakeKey(producer, static_cast(rr), b, o); - if (o.get_mode == GetMode::kExists || o.get_mode == GetMode::kBoth) { + if (!o.ext_kv_only && (o.get_mode == GetMode::kExists || o.get_mode == GetMode::kBoth)) { std::vector found; const auto t0 = Clock::now(); grpc::Status st = cli->Master().BatchLookup(rkeys, &found); @@ -601,29 +660,50 @@ void Worker(size_t id, size_t round_begin, size_t round_end, bool measure, Ctx c } } } - if (o.get_mode == GetMode::kFetch || o.get_mode == GetMode::kBoth) { + if (!o.ext_kv_only && (o.get_mode == GetMode::kFetch || o.get_mode == GetMode::kBoth)) { const auto t0 = Clock::now(); - auto bg = cli->BatchGet(rkeys, dsts, sizes); + // A faster process can tear down its peer service at end-of-run while this + // one is still fetching; the RDMA control-plane read then throws a transport + // error on this worker thread. Record it as a distinct transport-error -- + // NOT a miss, and with no latency sample -- so hit/miss + p50 stay clean, and + // keep going so a straggler does not abort the whole (multi-process) run. The + // cross-process barrier normally makes this never fire. + std::vector bg; + bool bg_ok = true; + try { + bg = cli->BatchGet(rkeys, dsts, sizes); + } catch (const std::exception& e) { + bg_ok = false; + static std::atomic warned{false}; + if (!warned.exchange(true)) + std::fprintf(stderr, "warning: BatchGet transport error (peer down?): %s\n", e.what()); + } catch (...) { + bg_ok = false; + } const auto t1 = Clock::now(); if (measure) { - fetch_ms.push_back(DurMs(t1 - t0)); - ctx.m->fetch_calls.fetch_add(1, std::memory_order_relaxed); - if (o.get_mode == GetMode::kFetch) { - // fetch-only classification is coarse: BatchGet cannot split RPC error from - // not-found. BatchRouteGet RPC errors are visible separately in - // mori_umbp_master_client_rpc_errors_total{rpc="BatchRouteGet"}. - ctx.m->get_keys.fetch_add(batch, std::memory_order_relaxed); - size_t h = 0; - for (bool f : bg) - if (f) ++h; - ctx.m->hit.fetch_add(h, std::memory_order_relaxed); - ctx.m->miss.fetch_add(batch - h, std::memory_order_relaxed); + if (!bg_ok) { + ctx.m->fetch_errors.fetch_add(1, std::memory_order_relaxed); + } else { + fetch_ms.push_back(DurMs(t1 - t0)); + ctx.m->fetch_calls.fetch_add(1, std::memory_order_relaxed); + if (o.get_mode == GetMode::kFetch) { + // fetch-only classification is coarse: BatchGet cannot split RPC error + // from not-found. BatchRouteGet RPC errors are visible separately in + // mori_umbp_master_client_rpc_errors_total{rpc="BatchRouteGet"}. + ctx.m->get_keys.fetch_add(batch, std::memory_order_relaxed); + size_t h = 0; + for (bool f : bg) + if (f) ++h; + ctx.m->hit.fetch_add(h, std::memory_order_relaxed); + ctx.m->miss.fetch_add(batch - h, std::memory_order_relaxed); + } } } } if (o.with_external_kv) { std::vector matches; - const bool ok = cli->MatchExternalKv(rkeys, &matches, /*count_as_hit=*/false); + const bool ok = cli->MatchExternalKv(rkeys, &matches, o.ext_kv_count_as_hit); if (measure && ok) { ctx.m->ext_match_calls.fetch_add(1, std::memory_order_relaxed); size_t matched = 0; @@ -645,6 +725,38 @@ void RunPhase(Ctx ctx, size_t begin, size_t end, bool measure) { for (auto& t : threads) t.join(); } +// Cross-process end-of-measurement barrier over a shared directory: each process +// announces itself with a file, then waits until `size` files exist (or timeout). +// Ensures no process tears down its peer service (RDMA endpoints) while another is +// still fetching -- the fetch/both teardown race that otherwise resets peers. +bool FileBarrierWait(const std::string& dir, const std::string& member, size_t size, + double timeout_s) { + namespace fs = std::filesystem; + std::error_code ec; + fs::create_directories(dir, ec); + { + std::ofstream f(fs::path(dir) / member); + f << "1"; + } + const auto deadline = + Clock::now() + + std::chrono::duration_cast(std::chrono::duration(timeout_s)); + for (;;) { + size_t n = 0; + for (auto it = fs::directory_iterator(dir, ec); + !ec && it != fs::directory_iterator(); it.increment(ec)) { + if (ec) break; + ++n; + } + if (n >= size) return true; + if (Clock::now() >= deadline) { + std::fprintf(stderr, "warning: barrier timeout (%zu/%zu present); proceeding\n", n, size); + return false; + } + std::this_thread::sleep_for(std::chrono::milliseconds(50)); + } +} + constexpr const char* kRpcLatencyMetric = "mori_umbp_master_client_rpc_latency_seconds"; } // namespace @@ -654,25 +766,36 @@ int main(int argc, char** argv) { if (!ParseArgs(argc, argv, &o)) return 2; const size_t N = o.clients; - // ---- master (put strategy injected directly; NOT via env) ---- - MasterServerConfig mcfg; - mcfg.listen_address = "0.0.0.0:0"; - mcfg.metrics_port = o.metrics_port; - mcfg.registry_config = ClientRegistryConfig::FromEnvironment(); - mcfg.put_strategy = std::make_unique(ParseAlgo(o.put_algo), - ParseAffinity(o.put_affinity)); - mcfg.route_put_algo = o.put_algo; - mcfg.route_put_affinity = o.put_affinity; - auto master = std::make_unique(std::move(mcfg)); - std::thread server_thread([&] { master->Run(); }); - for (int i = 0; i < 500 && master->GetBoundPort() == 0; ++i) { - std::this_thread::sleep_for(std::chrono::milliseconds(10)); - } - if (master->GetBoundPort() == 0) { - std::fprintf(stderr, "master failed to start\n"); - return 2; + // ---- master: spawn in-process (default), unless --external-master points at + // a standalone umbp_master shared by many processes/hosts. ---- + std::unique_ptr master; + std::thread server_thread; + std::string master_addr; + if (o.external_master.empty()) { + MasterServerConfig mcfg; + mcfg.listen_address = "0.0.0.0:0"; + mcfg.metrics_port = o.metrics_port; + mcfg.registry_config = ClientRegistryConfig::FromEnvironment(); + mcfg.put_strategy = std::make_unique( + ParseAlgo(o.put_algo), ParseAffinity(o.put_affinity)); + mcfg.route_put_algo = o.put_algo; + mcfg.route_put_affinity = o.put_affinity; + master = std::make_unique(std::move(mcfg)); + server_thread = std::thread([&] { master->Run(); }); + for (int i = 0; i < 500 && master->GetBoundPort() == 0; ++i) { + std::this_thread::sleep_for(std::chrono::milliseconds(10)); + } + if (master->GetBoundPort() == 0) { + std::fprintf(stderr, "master failed to start\n"); + return 2; + } + master_addr = "localhost:" + std::to_string(master->GetBoundPort()); + } else { + // Standalone master owns put strategy / index shards via its own env; this + // process only drives clients. --metrics-port should be 0 here (scrape the + // shared master's own endpoint externally, e.g. via parse_master_hist.py). + master_addr = o.external_master; } - const std::string master_addr = "localhost:" + std::to_string(master->GetBoundPort()); // ---- buffer sizing ---- const size_t total_rounds = o.warmup_rounds + o.rounds; @@ -694,10 +817,13 @@ int main(int argc, char** argv) { if (need_dst) dst_bufs[id].assign(io_bytes, 0); PoolClientConfig cfg; - cfg.master_config.node_id = "node-" + std::to_string(id); - cfg.master_config.node_address = "127.0.0.1"; + cfg.master_config.node_id = o.node_id_prefix + std::to_string(id); + cfg.master_config.node_address = o.node_address; cfg.master_config.master_address = master_addr; - cfg.io_engine.host = "0.0.0.0"; + // Advertise the routable node address for the RDMA engine so cross-machine + // peers can reach it. "0.0.0.0" makes the published EngineDesc non-routable + // (resolves to loopback), which only works when all peers are on one host. + cfg.io_engine.host = o.node_address; cfg.io_engine.port = 0; cfg.peer_service_port = NextPeerServicePort(); cfg.dram_page_size = o.page_bytes; @@ -761,8 +887,31 @@ int main(int argc, char** argv) { const double rpc_err_rate = get_total > 0 ? double(m.rpc_error_keys.load()) / double(get_total) : 0.0; + // ---- end-of-measurement sync so teardown never overlaps peer fetches ---- + // fetch/both read block data over RDMA from peer processes; if a faster process + // tears down first, a straggler's read resets. A cross-process barrier (when the + // harness supplies --barrier-*) is deterministic; otherwise fall back to a fixed + // drain (heuristic: safe only while the finish-time spread < BENCH_DRAIN_MS). + if (o.get_mode != GetMode::kExists) { + if (!o.barrier_dir.empty() && o.barrier_size > 1) { + const std::string tag = o.barrier_tag.empty() ? "barrier" : o.barrier_tag; + const std::string member = o.node_id_prefix + std::to_string(getpid()); + FileBarrierWait(o.barrier_dir + "/" + tag, member, o.barrier_size, /*timeout_s=*/60.0); + } else { + const char* drain_env = std::getenv("BENCH_DRAIN_MS"); + const long drain_ms = drain_env ? std::strtol(drain_env, nullptr, 10) : 3000; + if (drain_ms > 0) std::this_thread::sleep_for(std::chrono::milliseconds(drain_ms)); + } + } + // ---- shutdown clients (flushes their buffered metrics to the master) ---- for (size_t id = 0; id < N; ++id) clients[id]->Shutdown(); + if (m.fetch_errors.load() > 0) { + std::fprintf(stderr, + "NOTE: %llu fetch transport-errors excluded from hit/miss/latency " + "(peer teardown; raise BENCH_DRAIN_MS or use --barrier-*)\n", + static_cast(m.fetch_errors.load())); + } // ---- final Prometheus snapshot + delta ---- std::map rpc_delta; @@ -815,7 +964,7 @@ int main(int argc, char** argv) { std::this_thread::sleep_for( std::chrono::milliseconds(static_cast(o.keep_master_secs * 1000.0))); } - master->Shutdown(); + if (master) master->Shutdown(); if (server_thread.joinable()) server_thread.join(); return 0; } diff --git a/tests/cpp/umbp/distributed/bench_master_metadata_store.cpp b/tests/cpp/umbp/distributed/bench_master_metadata_store.cpp new file mode 100644 index 000000000..89becb47c --- /dev/null +++ b/tests/cpp/umbp/distributed/bench_master_metadata_store.cpp @@ -0,0 +1,275 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. + +// Store-level microbenchmark for IMasterMetadataStore — the Phase 1 go/no-go +// signal for the Redis metadata backend. It isolates exactly the hot-path +// methods the design calls out (BatchLookupBlockForRouteGet, BatchExistsBlock, +// ApplyHeartbeat) and measures per-operation latency + throughput, so the cost +// of moving master metadata from an in-process shared_mutex read (nanoseconds) +// to a network + Lua round trip (microseconds) is measured directly, with no +// RDMA data plane or gRPC layer in the way. +// +// Backend is selected by the factory via UMBP_METADATA_BACKEND / UMBP_REDIS_URI, +// exactly as the master picks its store: +// UMBP_METADATA_BACKEND=inmemory ./bench_master_metadata_store ... +// UMBP_METADATA_BACKEND=redis UMBP_REDIS_URI=tcp://127.0.0.1:6379 ./bench... ... +// +// One process runs one workload against one backend; launch it per scenario. + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "umbp/distributed/master/master_metadata_store_factory.h" +#include "umbp/distributed/types.h" + +using namespace mori::umbp; +using Clock = std::chrono::steady_clock; + +namespace { + +struct Opts { + std::string workload = "routeget"; // routeget | exists | heartbeat | mixed + int threads = 8; + double seconds = 5.0; + int keys = 50000; + int batch = 32; + double warmup_seconds = 1.0; +}; + +double DurUs(Clock::duration d) { return std::chrono::duration(d).count(); } + +double Percentile(std::vector& v, double p) { + if (v.empty()) return 0.0; + std::sort(v.begin(), v.end()); + const double idx = p * (static_cast(v.size()) - 1.0); + const size_t lo = static_cast(std::floor(idx)); + const size_t hi = static_cast(std::ceil(idx)); + if (lo == hi) return v[lo]; + const double frac = idx - static_cast(lo); + return v[lo] * (1.0 - frac) + v[hi] * frac; +} + +ClientRegistration MakeReg(int n) { + ClientRegistration r; + r.node_id = "node-" + std::to_string(n); + r.node_address = "127.0.0.1"; + r.peer_address = "127.0.0.1:" + std::to_string(17000 + n); + r.tier_capacities[TierType::DRAM] = TierCapacity{1ULL << 34, 1ULL << 34}; + r.tags = {"role=bench"}; + return r; +} + +std::string SeededKey(int i) { return "k/" + std::to_string(i); } + +void Usage() { + std::fprintf(stderr, + "Usage: bench_master_metadata_store [options]\n" + " --workload routeget|exists|heartbeat|mixed (default routeget)\n" + " --threads N (default 8)\n" + " --seconds F (default 5)\n" + " --warmup-seconds F (default 1)\n" + " --keys N (default 50000)\n" + " --batch N (default 32)\n" + "Backend is chosen via UMBP_METADATA_BACKEND / UMBP_REDIS_URI.\n"); +} + +} // namespace + +int main(int argc, char** argv) { + Opts o; + for (int i = 1; i < argc; ++i) { + std::string a = argv[i]; + auto need = [&]() -> const char* { + if (i + 1 >= argc) { + Usage(); + std::exit(2); + } + return argv[++i]; + }; + if (a == "--workload") + o.workload = need(); + else if (a == "--threads") + o.threads = std::atoi(need()); + else if (a == "--seconds") + o.seconds = std::atof(need()); + else if (a == "--warmup-seconds") + o.warmup_seconds = std::atof(need()); + else if (a == "--keys") + o.keys = std::atoi(need()); + else if (a == "--batch") + o.batch = std::atoi(need()); + else if (a == "-h" || a == "--help") { + Usage(); + return 0; + } else { + std::fprintf(stderr, "unknown arg %s\n", a.c_str()); + Usage(); + return 2; + } + } + if (o.threads < 1) o.threads = 1; + + const char* backend = std::getenv("UMBP_METADATA_BACKEND"); + std::string backend_name = backend ? backend : "inmemory"; + // Disambiguate Redis vs Dragonfly vs Valkey (all use UMBP_METADATA_BACKEND= + // redis): append the target URI so the CSV artifact is self-describing. + if (backend_name == "redis") { + const char* uri = std::getenv("UMBP_REDIS_URI"); + backend_name += "[" + std::string(uri ? uri : "tcp://127.0.0.1:6379") + "]"; + } + + std::unique_ptr store; + try { + store = MakeMasterMetadataStore(); + } catch (const std::exception& e) { + std::fprintf(stderr, "failed to build store: %s\n", e.what()); + return 2; + } + + // One node per thread so the heartbeat workload keeps a private monotonic seq + // (concurrent heartbeats to one node would seq-gap by design). + const int nodes = o.threads; + const auto now = std::chrono::system_clock::now(); + for (int n = 0; n < nodes; ++n) { + store->RegisterClient(MakeReg(n), now, std::chrono::seconds(120)); + } + + // Seed `keys` block locations, distributed round-robin across nodes, via + // delta heartbeats (chunked so no single Lua script is enormous). Track the + // next seq per node for the heartbeat workload. + std::vector next_seq(nodes, 1); + { + constexpr int kChunk = 256; + std::vector> pending(nodes); + auto flush = [&](int n) { + if (pending[n].empty()) return; + store->ApplyHeartbeat(MakeReg(n).node_id, next_seq[n]++, now, MakeReg(n).tier_capacities, + pending[n], /*is_full_sync=*/false); + pending[n].clear(); + }; + for (int i = 0; i < o.keys; ++i) { + const int n = i % nodes; + pending[n].push_back(KvEvent{KvEvent::Kind::ADD, SeededKey(i), TierType::DRAM, 4096}); + if (static_cast(pending[n].size()) >= kChunk) flush(n); + } + for (int n = 0; n < nodes; ++n) flush(n); + } + + std::atomic measuring{false}; + std::atomic stop{false}; + + auto worker = [&](int tid, std::vector* lat, uint64_t* ops) { + std::mt19937_64 rng(0x9E3779B97F4A7C15ULL ^ (tid + 1)); + std::uniform_int_distribution key_dist(0, std::max(0, o.keys - 1)); + std::uniform_int_distribution wl_dist(0, 1); + const std::string node = MakeReg(tid % nodes).node_id; + const auto caps = MakeReg(tid % nodes).tier_capacities; + uint64_t hb_seq = next_seq[tid % nodes]; + uint64_t hb_key = 0; + std::vector local; + local.reserve(1 << 20); + uint64_t local_ops = 0; + + std::vector keys(o.batch); + while (!stop.load(std::memory_order_relaxed)) { + std::string wl = o.workload; + if (wl == "mixed") wl = wl_dist(rng) ? "routeget" : "heartbeat"; + + const auto t0 = Clock::now(); + if (wl == "routeget") { + for (int b = 0; b < o.batch; ++b) keys[b] = SeededKey(key_dist(rng)); + auto r = store->BatchLookupBlockForRouteGet(keys, {}, std::chrono::system_clock::now(), + std::chrono::seconds(10)); + (void)r; + } else if (wl == "exists") { + for (int b = 0; b < o.batch; ++b) keys[b] = SeededKey(key_dist(rng)); + auto r = store->BatchExistsBlock(keys); + (void)r; + } else { // heartbeat + std::vector events; + events.reserve(o.batch); + for (int b = 0; b < o.batch; ++b) { + events.push_back(KvEvent{KvEvent::Kind::ADD, + "hb/" + std::to_string(tid) + "/" + std::to_string(hb_key++), + TierType::DRAM, 4096}); + } + auto r = store->ApplyHeartbeat(node, hb_seq++, std::chrono::system_clock::now(), caps, + events, /*is_full_sync=*/false); + (void)r; + } + const auto t1 = Clock::now(); + if (measuring.load(std::memory_order_relaxed)) { + local.push_back(DurUs(t1 - t0)); + ++local_ops; + } + } + *lat = std::move(local); + *ops = local_ops; + }; + + std::vector> lats(o.threads); + std::vector ops(o.threads, 0); + std::vector threads; + threads.reserve(o.threads); + for (int t = 0; t < o.threads; ++t) { + threads.emplace_back(worker, t, &lats[t], &ops[t]); + } + + std::this_thread::sleep_for(std::chrono::duration(o.warmup_seconds)); + measuring.store(true, std::memory_order_relaxed); + const auto t_start = Clock::now(); + std::this_thread::sleep_for(std::chrono::duration(o.seconds)); + measuring.store(false, std::memory_order_relaxed); + const auto t_end = Clock::now(); + stop.store(true, std::memory_order_relaxed); + for (auto& th : threads) th.join(); + + std::vector all; + uint64_t total_ops = 0; + for (int t = 0; t < o.threads; ++t) { + all.insert(all.end(), lats[t].begin(), lats[t].end()); + total_ops += ops[t]; + } + const double wall_s = std::chrono::duration(t_end - t_start).count(); + const double ops_per_s = wall_s > 0 ? total_ops / wall_s : 0.0; + const double keys_per_s = ops_per_s * o.batch; + + std::printf( + "backend,workload,threads,batch,keys,wall_s,ops,ops_per_s,keys_per_s," + "lat_us_p50,lat_us_p95,lat_us_p99,lat_us_max\n"); + std::printf("%s,%s,%d,%d,%d,%.3f,%llu,%.0f,%.0f,%.2f,%.2f,%.2f,%.2f\n", backend_name.c_str(), + o.workload.c_str(), o.threads, o.batch, o.keys, wall_s, + static_cast(total_ops), ops_per_s, keys_per_s, + Percentile(all, 0.50), Percentile(all, 0.95), Percentile(all, 0.99), + all.empty() ? 0.0 : *std::max_element(all.begin(), all.end())); + std::fflush(stdout); + return 0; +} diff --git a/tests/cpp/umbp/distributed/bench_master_metadata_store_extkv.cpp b/tests/cpp/umbp/distributed/bench_master_metadata_store_extkv.cpp new file mode 100644 index 000000000..a30220e64 --- /dev/null +++ b/tests/cpp/umbp/distributed/bench_master_metadata_store_extkv.cpp @@ -0,0 +1,337 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. + +// Store-level microbenchmark for the EXTERNAL-KV hot path of +// IMasterMetadataStore. Sibling of bench_master_metadata_store.cpp (which +// covers RouteGet / Exists / Heartbeat); this one isolates the external-KV +// methods that SGLang's KV-events integration drives, so the per-op cost of +// each external-KV interface can be compared directly across backends +// (in-memory / single Redis / sharded Redis / Redis Cluster / Dragonfly) with +// no gRPC or RDMA layer in the way. +// +// Why these methods (SGLang usage, see sglang umbp_store.py + mori +// master_server.cpp): +// - MatchExternalKv (count_as_hit=true) is THE hot read: the prefix-lookup / +// route path asks "which live nodes hold these block hashes?", and each +// match atomically bumps the per-hash hit counter + last_seen. This is the +// external-KV analogue of BatchLookupBlockForRouteGet and the single +// interface most able to move end-to-end latency. +// - RegisterExternalKvIfAlive backs report_external_kv_blocks, fired on every +// BlockStored KV-cache event — the dominant external-KV WRITE. +// - UnregisterExternalKv backs revoke_external_kv_blocks (BlockRemoved). +// - GetExternalKvHitCounts backs the eviction / admin hit-count read. +// +// Design note that this bench is built to expose: external-KV + hit state all +// hang off ONE control hash tag ({umbp:}), i.e. a single Redis slot / +// instance. UMBP_REDIS_SHARD_URIS / UMBP_REDIS_BLOCK_SHARDS shard the BLOCK +// index, NOT this path — so sharded-Redis and Cluster are not expected to scale +// the external-KV hot path, while a single multi-threaded Dragonfly instance +// can. Run this bench across topologies to confirm. +// +// Backend is selected by the factory via UMBP_METADATA_BACKEND / UMBP_REDIS_URI, +// exactly as the master picks its store: +// UMBP_METADATA_BACKEND=inmemory ./bench...extkv --workload match ... +// UMBP_METADATA_BACKEND=redis UMBP_REDIS_URI=tcp://127.0.0.1:6379 ./bench...extkv ... +// +// One process runs one workload against one backend; launch it per scenario. + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "umbp/distributed/master/master_metadata_store_factory.h" +#include "umbp/distributed/types.h" + +using namespace mori::umbp; +using Clock = std::chrono::steady_clock; + +namespace { + +struct Opts { + // match | match_nohit | report | revoke | hitcounts | mixed + std::string workload = "match"; + int threads = 8; + int nodes = 8; // registered clients the hashes are spread across + double seconds = 5.0; + int keys = 50000; // external-kv hash keyspace (seeded across nodes) + int batch = 32; // hashes per op (real prefix match is tens) + double hit_ratio = 1.0; // fraction of queried hashes that actually exist + double warmup_seconds = 1.0; +}; + +double DurUs(Clock::duration d) { return std::chrono::duration(d).count(); } + +double Percentile(std::vector& v, double p) { + if (v.empty()) return 0.0; + std::sort(v.begin(), v.end()); + const double idx = p * (static_cast(v.size()) - 1.0); + const size_t lo = static_cast(std::floor(idx)); + const size_t hi = static_cast(std::ceil(idx)); + if (lo == hi) return v[lo]; + const double frac = idx - static_cast(lo); + return v[lo] * (1.0 - frac) + v[hi] * frac; +} + +ClientRegistration MakeReg(int n) { + ClientRegistration r; + r.node_id = "node-" + std::to_string(n); + r.node_address = "127.0.0.1"; + r.peer_address = "127.0.0.1:" + std::to_string(17000 + n); + r.tier_capacities[TierType::DRAM] = TierCapacity{1ULL << 34, 1ULL << 34}; + r.tags = {"role=bench"}; + return r; +} + +// Seeded external-kv hash present in the store (registered during setup). +std::string SeededHash(int i) { return "h/" + std::to_string(i); } +// Hash guaranteed absent from the store (index beyond the seeded keyspace) — +// used to synthesize misses for --hit-ratio < 1. +std::string MissHash(int i, int keys) { return "h/" + std::to_string(keys + i); } + +void Usage() { + std::fprintf(stderr, + "Usage: bench_master_metadata_store_extkv [options]\n" + " --workload match|match_nohit|report|revoke|hitcounts|mixed (default match)\n" + " match MatchExternalKv(count_as_hit=true) -- hot read (+hit-count write)\n" + " match_nohit MatchExternalKv(count_as_hit=false) -- pure read (isolates hit cost)\n" + " report RegisterExternalKvIfAlive -- hot write (BlockStored)\n" + " revoke UnregisterExternalKv -- write (BlockRemoved)\n" + " hitcounts GetExternalKvHitCounts -- eviction/admin read\n" + " mixed 1:1 match(count_as_hit) : report -- read/write blend\n" + " --threads N (default 8)\n" + " --nodes N registered clients hashes spread across (default 8)\n" + " --seconds F (default 5)\n" + " --warmup-seconds F (default 1)\n" + " --keys N external-kv hash keyspace (default 50000)\n" + " --batch N hashes per op (default 32)\n" + " --hit-ratio F fraction of queried hashes that exist (default 1.0)\n" + "Backend is chosen via UMBP_METADATA_BACKEND / UMBP_REDIS_URI.\n"); +} + +} // namespace + +int main(int argc, char** argv) { + Opts o; + for (int i = 1; i < argc; ++i) { + std::string a = argv[i]; + auto need = [&]() -> const char* { + if (i + 1 >= argc) { + Usage(); + std::exit(2); + } + return argv[++i]; + }; + if (a == "--workload") + o.workload = need(); + else if (a == "--threads") + o.threads = std::atoi(need()); + else if (a == "--nodes") + o.nodes = std::atoi(need()); + else if (a == "--seconds") + o.seconds = std::atof(need()); + else if (a == "--warmup-seconds") + o.warmup_seconds = std::atof(need()); + else if (a == "--keys") + o.keys = std::atoi(need()); + else if (a == "--batch") + o.batch = std::atoi(need()); + else if (a == "--hit-ratio") + o.hit_ratio = std::atof(need()); + else if (a == "-h" || a == "--help") { + Usage(); + return 0; + } else { + std::fprintf(stderr, "unknown arg %s\n", a.c_str()); + Usage(); + return 2; + } + } + if (o.threads < 1) o.threads = 1; + if (o.nodes < 1) o.nodes = 1; + if (o.hit_ratio < 0.0) o.hit_ratio = 0.0; + if (o.hit_ratio > 1.0) o.hit_ratio = 1.0; + + const char* backend = std::getenv("UMBP_METADATA_BACKEND"); + std::string backend_name = backend ? backend : "inmemory"; + // Disambiguate Redis vs Dragonfly vs Valkey / single vs sharded vs cluster + // (all use UMBP_METADATA_BACKEND=redis): append the target so the CSV is + // self-describing. + if (backend_name == "redis") { + const char* uri = std::getenv("UMBP_REDIS_URI"); + const char* shards = std::getenv("UMBP_REDIS_SHARD_URIS"); + const char* cluster = std::getenv("UMBP_REDIS_CLUSTER"); + std::string tag = uri ? uri : "tcp://127.0.0.1:6379"; + if (cluster && std::string(cluster) != "0") tag = "cluster:" + tag; + if (shards && shards[0]) tag = "sharded:" + std::string(shards); + // The CSV is comma-separated; a shard-URI list contains commas — replace + // them so the artifact stays parseable. + for (char& c : tag) + if (c == ',') c = ';'; + backend_name += "[" + tag + "]"; + } + + std::unique_ptr store; + try { + store = MakeMasterMetadataStore(); + } catch (const std::exception& e) { + std::fprintf(stderr, "failed to build store: %s\n", e.what()); + return 2; + } + + const int nodes = o.nodes; + const auto now = std::chrono::system_clock::now(); + for (int n = 0; n < nodes; ++n) { + store->RegisterClient(MakeReg(n), now, std::chrono::seconds(120)); + } + + // Seed `keys` external-kv hashes, round-robin across nodes, at DRAM tier, in + // chunks so no single Lua script is enormous. Mirrors report_external_kv_blocks + // fan-in from many nodes. + { + constexpr int kChunk = 512; + std::vector> pending(nodes); + auto flush = [&](int n) { + if (pending[n].empty()) return; + store->RegisterExternalKvIfAlive(MakeReg(n).node_id, pending[n], TierType::DRAM); + pending[n].clear(); + }; + for (int i = 0; i < o.keys; ++i) { + const int n = i % nodes; + pending[n].push_back(SeededHash(i)); + if (static_cast(pending[n].size()) >= kChunk) flush(n); + } + for (int n = 0; n < nodes; ++n) flush(n); + } + + std::atomic measuring{false}; + std::atomic stop{false}; + + auto worker = [&](int tid, std::vector* lat, uint64_t* ops) { + std::mt19937_64 rng(0x9E3779B97F4A7C15ULL ^ (tid + 1)); + std::uniform_int_distribution key_dist(0, std::max(0, o.keys - 1)); + std::uniform_real_distribution coin(0.0, 1.0); + const std::string node = MakeReg(tid % nodes).node_id; + // report/revoke draw from this thread's own node keyspace slice so writes + // stay attributable to one node (as a real peer's events would be). + std::uniform_int_distribution node_key_dist(0, std::max(0, o.keys / nodes - 1)); + uint64_t report_seq = 0; // grows the write keyspace (fresh BlockStored churn) + + std::vector local; + local.reserve(1 << 20); + uint64_t local_ops = 0; + bool mix_toggle = false; + + std::vector hashes(o.batch); + while (!stop.load(std::memory_order_relaxed)) { + std::string wl = o.workload; + if (wl == "mixed") { + wl = mix_toggle ? "report" : "match"; + mix_toggle = !mix_toggle; + } + + const auto t0 = Clock::now(); + if (wl == "match" || wl == "match_nohit") { + for (int b = 0; b < o.batch; ++b) { + if (coin(rng) < o.hit_ratio) + hashes[b] = SeededHash(key_dist(rng)); + else + hashes[b] = MissHash(key_dist(rng), o.keys); + } + auto r = store->MatchExternalKv(hashes, /*count_as_hit=*/wl == "match", + std::chrono::system_clock::now()); + (void)r; + } else if (wl == "hitcounts") { + for (int b = 0; b < o.batch; ++b) hashes[b] = SeededHash(key_dist(rng)); + auto r = store->GetExternalKvHitCounts(hashes); + (void)r; + } else if (wl == "report") { + // Fresh hashes per op in this node's slice => models BlockStored churn. + for (int b = 0; b < o.batch; ++b) { + const int base = (tid % nodes) + nodes * node_key_dist(rng); + hashes[b] = "hw/" + std::to_string(tid) + "/" + std::to_string(report_seq++) + "/" + + std::to_string(base); + } + auto r = store->RegisterExternalKvIfAlive(node, hashes, TierType::DRAM); + (void)r; + } else { // revoke + for (int b = 0; b < o.batch; ++b) { + const int idx = (tid % nodes) + nodes * node_key_dist(rng); + hashes[b] = SeededHash(idx); + } + store->UnregisterExternalKv(node, hashes, TierType::DRAM); + } + const auto t1 = Clock::now(); + if (measuring.load(std::memory_order_relaxed)) { + local.push_back(DurUs(t1 - t0)); + ++local_ops; + } + } + *lat = std::move(local); + *ops = local_ops; + }; + + std::vector> lats(o.threads); + std::vector ops(o.threads, 0); + std::vector threads; + threads.reserve(o.threads); + for (int t = 0; t < o.threads; ++t) { + threads.emplace_back(worker, t, &lats[t], &ops[t]); + } + + std::this_thread::sleep_for(std::chrono::duration(o.warmup_seconds)); + measuring.store(true, std::memory_order_relaxed); + const auto t_start = Clock::now(); + std::this_thread::sleep_for(std::chrono::duration(o.seconds)); + measuring.store(false, std::memory_order_relaxed); + const auto t_end = Clock::now(); + stop.store(true, std::memory_order_relaxed); + for (auto& th : threads) th.join(); + + std::vector all; + uint64_t total_ops = 0; + for (int t = 0; t < o.threads; ++t) { + all.insert(all.end(), lats[t].begin(), lats[t].end()); + total_ops += ops[t]; + } + const double wall_s = std::chrono::duration(t_end - t_start).count(); + const double ops_per_s = wall_s > 0 ? total_ops / wall_s : 0.0; + const double keys_per_s = ops_per_s * o.batch; + + std::printf( + "backend,workload,threads,nodes,batch,keys,hit_ratio,wall_s,ops,ops_per_s,keys_per_s," + "lat_us_p50,lat_us_p95,lat_us_p99,lat_us_max\n"); + std::printf("%s,%s,%d,%d,%d,%d,%.2f,%.3f,%llu,%.0f,%.0f,%.2f,%.2f,%.2f,%.2f\n", + backend_name.c_str(), o.workload.c_str(), o.threads, o.nodes, o.batch, o.keys, + o.hit_ratio, wall_s, static_cast(total_ops), ops_per_s, keys_per_s, + Percentile(all, 0.50), Percentile(all, 0.95), Percentile(all, 0.99), + all.empty() ? 0.0 : *std::max_element(all.begin(), all.end())); + std::fflush(stdout); + return 0; +} diff --git a/tests/cpp/umbp/distributed/parse_master_hist.py b/tests/cpp/umbp/distributed/parse_master_hist.py new file mode 100755 index 000000000..732472775 --- /dev/null +++ b/tests/cpp/umbp/distributed/parse_master_hist.py @@ -0,0 +1,76 @@ +#!/usr/bin/env python3 +"""Fetch a UMBP master's Prometheus /metrics and print per-RPC p50/p95/p99. + +Aggregates mori_umbp_master_client_rpc_latency_seconds_bucket across ALL node +labels (the histogram is per-client), then linear-interpolates the quantile +within the matched bucket (standard Prometheus histogram_quantile). +""" +import sys +import re +import urllib.request + +METRIC = "mori_umbp_master_client_rpc_latency_seconds_bucket" +COUNT = "mori_umbp_master_client_rpc_latency_seconds_count" + +BUCKET_RE = re.compile(r'^%s\{([^}]*)\}\s+([0-9.eE+]+)\s*$' % re.escape(METRIC)) +COUNT_RE = re.compile(r'^%s\{([^}]*)\}\s+([0-9.eE+]+)\s*$' % re.escape(COUNT)) + + +def labels(s): + d = {} + for m in re.finditer(r'(\w+)="([^"]*)"', s): + d[m.group(1)] = m.group(2) + return d + + +def main(): + url = sys.argv[1] if len(sys.argv) > 1 else "http://127.0.0.1:9091/metrics" + body = urllib.request.urlopen(url, timeout=15).read().decode("utf-8", "replace") + + # rpc -> {le(float) -> summed cumulative count}; rpc -> total count + buckets = {} + counts = {} + for line in body.splitlines(): + m = BUCKET_RE.match(line) + if m: + lab = labels(m.group(1)) + rpc = lab.get("rpc", "?") + le = lab.get("le") + le = float("inf") if le == "+Inf" else float(le) + buckets.setdefault(rpc, {}) + buckets[rpc][le] = buckets[rpc].get(le, 0.0) + float(m.group(2)) + continue + m = COUNT_RE.match(line) + if m: + lab = labels(m.group(1)) + rpc = lab.get("rpc", "?") + counts[rpc] = counts.get(rpc, 0.0) + float(m.group(2)) + + + def quantile(edges, q): + total = edges[-1][1] # +Inf cumulative = total + if total <= 0: + return 0.0 + rank = q * total + prev_edge, prev_cum = 0.0, 0.0 + for le, cum in edges: + if cum >= rank: + if le == float("inf"): + return prev_edge * 1000.0 # cannot interpolate the open bucket + # linear interpolation within (prev_edge, le] + span = cum - prev_cum + frac = (rank - prev_cum) / span if span > 0 else 0.0 + return (prev_edge + (le - prev_edge) * frac) * 1000.0 + prev_edge, prev_cum = le, cum + return prev_edge * 1000.0 + + print("rpc,count,p50_ms,p95_ms,p99_ms") + for rpc in sorted(buckets): + edges = sorted(buckets[rpc].items()) + cnt = counts.get(rpc, edges[-1][1]) + print("%s,%.0f,%.3f,%.3f,%.3f" % ( + rpc, cnt, quantile(edges, 0.50), quantile(edges, 0.95), quantile(edges, 0.99))) + + +if __name__ == "__main__": + main() diff --git a/tests/cpp/umbp/distributed/run_extkv_store_bench.sh b/tests/cpp/umbp/distributed/run_extkv_store_bench.sh new file mode 100755 index 000000000..c4e3f10a8 --- /dev/null +++ b/tests/cpp/umbp/distributed/run_extkv_store_bench.sh @@ -0,0 +1,131 @@ +#!/usr/bin/env bash +# External-KV store-level microbench sweep for the UMBP master metadata backend. +# +# Runs bench_umbp_master_metadata_store_extkv (isolated IMasterMetadataStore +# calls, no gRPC / RDMA) across the external-KV hot-path workloads and every +# backend topology, and prints a side-by-side comparison. This is the cleanest +# apples-to-apples signal for "which backend serves the external-KV hot path +# best" — the analogue of the guide's §3 store microbench, for external KV. +# +# It exists to answer, per interface: +# - MatchExternalKv (count_as_hit=true) -- hot read (+ hit-count write) +# - MatchExternalKv (count_as_hit=false) -- pure read (isolates the hit cost) +# - RegisterExternalKvIfAlive -- hot write (BlockStored) +# - UnregisterExternalKv -- write (BlockRemoved) +# - GetExternalKvHitCounts -- eviction/admin read +# and to expose the design fact that external-KV/hit state lives on ONE control +# hash tag ({umbp:}) — a single Redis slot/instance — so SHARDED and CLUSTER +# are NOT expected to scale this path, while DRAGONFLY (multi-threaded single +# instance) might. The sweep makes that visible. +# +# Env knobs (all optional): +# BACKENDS space list from: inmemory single sharded cluster dragonfly +# (default "inmemory single sharded cluster dragonfly") +# WORKLOADS space list from: match match_nohit report revoke hitcounts mixed +# (default "match match_nohit report revoke mixed") +# THREADS SECONDS WARMUP KEYS BATCH NODES HIT_RATIO +# (defaults 8 5 1 50000 32 8 1.0) +# SINGLE_URI tcp://host:port (default tcp://127.0.0.1:6379) +# DRAGONFLY_URI (default tcp://127.0.0.1:6380) +# DRAGONFLY_BLOCK_SHARDS (default 8; matches Dragonfly threads) +# SHARD_URIS comma list (default 6390..6393 on localhost) +# CLUSTER_URI seed (default tcp://127.0.0.1:7000) +# MORI_BUILD_DIR build dir (default /build) +# REDIS_CLI redis-cli path (default PATH, else /tmp/umbp_redis_bench/redis-cli) +# OUT output dir (default ./extkv_store_out) +set -uo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +REPO_ROOT="$(cd "${SCRIPT_DIR}/../../../.." && pwd)" +BUILD_DIR="${MORI_BUILD_DIR:-${REPO_ROOT}/build}" +BIN="${BUILD_DIR}/tests/cpp/umbp/distributed/bench_umbp_master_metadata_store_extkv" + +BACKENDS="${BACKENDS:-inmemory single sharded cluster dragonfly}" +WORKLOADS="${WORKLOADS:-match match_nohit report revoke mixed}" +THREADS="${THREADS:-8}"; SECONDS_="${SECONDS:-5}"; WARMUP="${WARMUP:-1}" +KEYS="${KEYS:-50000}"; BATCH="${BATCH:-32}"; NODES="${NODES:-8}"; HIT_RATIO="${HIT_RATIO:-1.0}" + +SINGLE_URI="${SINGLE_URI:-tcp://127.0.0.1:6379}" +DRAGONFLY_URI="${DRAGONFLY_URI:-tcp://127.0.0.1:6380}" +DRAGONFLY_BLOCK_SHARDS="${DRAGONFLY_BLOCK_SHARDS:-8}" +SHARD_URIS="${SHARD_URIS:-tcp://127.0.0.1:6390,tcp://127.0.0.1:6391,tcp://127.0.0.1:6392,tcp://127.0.0.1:6393}" +CLUSTER_URI="${CLUSTER_URI:-tcp://127.0.0.1:7000}" + +REDIS_CLI="${REDIS_CLI:-$(command -v redis-cli || echo /tmp/umbp_redis_bench/redis-cli)}" +OUT="${OUT:-./extkv_store_out}" + +if [[ ! -x "$BIN" ]]; then + echo "ERROR: bench not built: $BIN" + echo " build: USE_REDIS_BACKEND=ON BUILD_UMBP=ON BUILD_TESTS=ON pip3 install -e . --no-build-isolation -v" + echo " or: cmake --build $BUILD_DIR --target bench_umbp_master_metadata_store_extkv -j" + exit 2 +fi +mkdir -p "$OUT" +CSV="${OUT}/extkv_store_results.csv" +echo "topology,backend,workload,threads,nodes,batch,keys,hit_ratio,wall_s,ops,ops_per_s,keys_per_s,lat_us_p50,lat_us_p95,lat_us_p99,lat_us_max" > "$CSV" + +host_port() { local u="${1#tcp://}"; echo "${u%%:*} ${u##*:}"; } +flush() { # flush every port that this backend uses, so hit-counts start clean + local ports="$*" + for hp in $ports; do + read -r h p <<<"$(host_port "$hp")" + "$REDIS_CLI" -h "$h" -p "$p" FLUSHALL >/dev/null 2>&1 + "$REDIS_CLI" -h "$h" -p "$p" CONFIG RESETSTAT >/dev/null 2>&1 + done +} + +# Emit the env prefix + a FLUSH target list for a given backend name. +setup_backend() { # $1 backend -> prints "ENV|||FLUSH_PORTS" or "SKIP" + local be="$1" ns="ek_$(date +%s)_$$_${RANDOM}" + case "$be" in + inmemory) echo "UMBP_METADATA_BACKEND=inmemory|||" ;; + single) + read -r h p <<<"$(host_port "$SINGLE_URI")" + "$REDIS_CLI" -h "$h" -p "$p" ping >/dev/null 2>&1 || { echo SKIP; return; } + echo "UMBP_METADATA_BACKEND=redis UMBP_REDIS_URI=${SINGLE_URI} UMBP_REDIS_NAMESPACE=${ns}|||${SINGLE_URI}" ;; + dragonfly) + read -r h p <<<"$(host_port "$DRAGONFLY_URI")" + "$REDIS_CLI" -h "$h" -p "$p" ping >/dev/null 2>&1 || { echo SKIP; return; } + echo "UMBP_METADATA_BACKEND=redis UMBP_REDIS_URI=${DRAGONFLY_URI} UMBP_REDIS_BLOCK_SHARDS=${DRAGONFLY_BLOCK_SHARDS} UMBP_REDIS_NAMESPACE=${ns}|||${DRAGONFLY_URI}" ;; + sharded) + local first="${SHARD_URIS%%,*}"; read -r h p <<<"$(host_port "$first")" + "$REDIS_CLI" -h "$h" -p "$p" ping >/dev/null 2>&1 || { echo SKIP; return; } + echo "UMBP_METADATA_BACKEND=redis UMBP_REDIS_SHARD_URIS=${SHARD_URIS} UMBP_REDIS_NAMESPACE=${ns}|||${SHARD_URIS//,/ }" ;; + cluster) + read -r h p <<<"$(host_port "$CLUSTER_URI")" + "$REDIS_CLI" -h "$h" -p "$p" ping >/dev/null 2>&1 || { echo SKIP; return; } + # Cluster: unique namespace isolates runs; no cluster-wide FLUSHALL. + echo "UMBP_METADATA_BACKEND=redis UMBP_REDIS_CLUSTER=1 UMBP_REDIS_URI=${CLUSTER_URI} UMBP_REDIS_NAMESPACE=${ns}|||" ;; + *) echo SKIP ;; + esac +} + +echo "=== extkv store sweep: threads=${THREADS} nodes=${NODES} batch=${BATCH} keys=${KEYS} hit_ratio=${HIT_RATIO} secs=${SECONDS_} ===" +for be in $BACKENDS; do + spec="$(setup_backend "$be")" + if [[ "$spec" == "SKIP" ]]; then echo "-- $be: unreachable, skipped"; continue; fi + ENV="${spec%%|||*}"; FLUSH_PORTS="${spec##*|||}" + for wl in $WORKLOADS; do + [[ -n "$FLUSH_PORTS" ]] && flush $FLUSH_PORTS + line="$(env $ENV "$BIN" --workload "$wl" --threads "$THREADS" --nodes "$NODES" \ + --seconds "$SECONDS_" --warmup-seconds "$WARMUP" --keys "$KEYS" \ + --batch "$BATCH" --hit-ratio "$HIT_RATIO" 2>/dev/null | tail -n +2)" + echo "${be},${line}" >> "$CSV" + echo " ${be}/${wl}: ${line}" + done +done + +echo +echo "=== comparison (ops/s p50us p95us p99us) ===" +awk -F, ' + { be=$1; wl=$3; ops=$11; p50=$13; p95=$14; p99=$15; + val[be","wl]=sprintf("%8.0f %8.1f %8.1f %8.1f", ops, p50, p95, p99); + bes[be]=1; wls[wl]=1 } + END{ + printf "%-14s %-12s %8s %8s %8s %8s\n","backend","workload","ops/s","p50us","p95us","p99us"; + n=split("inmemory single sharded cluster dragonfly", order, " "); + for(i=1;i<=n;i++){ be=order[i]; if(!(be in bes)) continue; + for(wl in wls){ if((be","wl) in val) printf "%-14s %-12s %s\n", be, wl, val[be","wl]; } } + }' "$CSV" | sort -k2,2 -k1,1 +echo +echo "results CSV -> $CSV" diff --git a/tests/cpp/umbp/distributed/run_mp_redis_bench.sh b/tests/cpp/umbp/distributed/run_mp_redis_bench.sh new file mode 100755 index 000000000..e64aeedb9 --- /dev/null +++ b/tests/cpp/umbp/distributed/run_mp_redis_bench.sh @@ -0,0 +1,143 @@ +#!/usr/bin/env bash +# Multi-PROCESS control-plane pressure harness for the UMBP master metadata +# backend (Redis or in-memory). +# +# One standalone `umbp_master` is shared by PROCS OS processes x CLIENTS clients, +# all driving BatchRouteGet / BatchRoutePut / Heartbeat through the full +# Router/gRPC path (requires bench_kvevent_master_pressure built with +# --external-master, see that file). This reproduces the real +# "N-processes-per-machine + multi-machine" shape that the store microbench +# (1 process, N threads) and the in-process kvevent bench (1 process, N clients) +# cannot: each process has its own connection pool / gRPC channel / heartbeat. +# +# It is the primary *repeatable* benchmark for judging Redis-backend +# optimizations: scale PROCS/CLIENTS until the single-slot ceiling shows, then +# compare master per-RPC p50/p95/p99 + redis evalsha usec/call before vs after. +# +# Env knobs (all optional): +# BACKEND redis | inmemory (default redis) +# REDIS_URI tcp://host:6379 (default tcp://127.0.0.1:6379) +# PROCS OS processes (== workers/machine) (default 8) +# CLIENTS clients per process (default 2) +# ROUNDS WARMUP BATCH GAP GETMODE KEYSPACE (defaults 300 20 32 0 both 4096) +# GETMODE: exists=BatchLookup(BatchExistsBlock, no RDMA); +# fetch/both=BatchRouteGet(route_get_batch)+RDMA fetch. +# MORI_BUILD_DIR path to mori build dir (default /build) +# REDIS_CLI redis-cli path (default: PATH, else /tmp/umbp_redis_bench/redis-cli) +# OUT output dir (default ./mp_redis_out) +# PORT METRICS standalone master ports (default 15560 9092) +# +# For a multi-machine run: start ONE master (this script on the master host with +# PROCS/CLIENTS as desired), then run extra client-only processes on other hosts +# pointing bench_kvevent_master_pressure --external-master : +# --node-address . +set -uo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +REPO_ROOT="$(cd "${SCRIPT_DIR}/../../../../.." && pwd)" +BUILD_DIR="${MORI_BUILD_DIR:-${REPO_ROOT}/build}" +BIN="${BUILD_DIR}/tests/cpp/umbp/distributed/bench_umbp_kvevent_master_pressure" +MASTER_BIN="${BUILD_DIR}/src/umbp/umbp_master" +PARSE="${SCRIPT_DIR}/parse_master_hist.py" + +BACKEND="${BACKEND:-redis}" +REDIS_URI="${REDIS_URI:-tcp://127.0.0.1:6379}" +# Redis deployment: REDIS_CLUSTER=1 => cluster (REDIS_URI is the comma seed list); +# SHARD_URIS set => multi-endpoint (one instance per block shard); else single. +REDIS_CLUSTER="${REDIS_CLUSTER:-0}" +SHARD_URIS="${SHARD_URIS:-}" +PROCS="${PROCS:-8}"; CLIENTS="${CLIENTS:-2}" +ROUNDS="${ROUNDS:-300}"; WARMUP="${WARMUP:-20}"; BATCH="${BATCH:-32}" +GAP="${GAP:-0}"; GETMODE="${GETMODE:-both}"; KEYSPACE="${KEYSPACE:-4096}" +# External-KV knobs (default off). EXT_KV=1 also drives ReportExternalKvBlocks + +# MatchExternalKv through the full Router/gRPC path. EXT_KV_ONLY=1 exercises ONLY +# those two (skips BatchPut/Get + RDMA) so the backend comparison is on the extkv +# hot path alone. EXT_KV_COUNT_AS_HIT=1 (default) makes MatchExternalKv perform +# the real per-hash hit-count write, as the production route/prefix path does. +EXT_KV="${EXT_KV:-0}"; EXT_KV_ONLY="${EXT_KV_ONLY:-0}"; EXT_KV_COUNT_AS_HIT="${EXT_KV_COUNT_AS_HIT:-1}" +PORT="${PORT:-15560}"; METRICS="${METRICS:-9092}" +OUT="${OUT:-./mp_redis_out}" + +REDIS_CLI="${REDIS_CLI:-$(command -v redis-cli || echo /tmp/umbp_redis_bench/redis-cli)}" +RHOST="${REDIS_URI#tcp://}"; RHOST="${RHOST%%:*}" +RPORT="${REDIS_URI##*:}" + +if [[ ! -x "$BIN" ]]; then echo "ERROR: bench not built: $BIN (build with USE_REDIS_BACKEND=ON BUILD_TESTS=ON)"; exit 2; fi +mkdir -p "$OUT"; ulimit -n 1048576 2>/dev/null || true +export MORI_IO_QP_MAX_SEND_WR="${MORI_IO_QP_MAX_SEND_WR:-1024}" +HOSTIP="$(hostname -i 2>/dev/null | awk '{print $1}')"; HOSTIP="${HOSTIP:-127.0.0.1}" +# External-KV client args (appended to every client process launch). +EXTKV_ARGS="" +EXTKV_TAG="" +if [[ "$EXT_KV_ONLY" == "1" ]]; then + EXTKV_ARGS="--ext-kv-only --ext-kv-count-as-hit ${EXT_KV_COUNT_AS_HIT}" + EXTKV_TAG="_extkvonly" +elif [[ "$EXT_KV" == "1" ]]; then + EXTKV_ARGS="--with-external-kv --ext-kv-count-as-hit ${EXT_KV_COUNT_AS_HIT}" + EXTKV_TAG="_extkv" +fi +LABEL="${BACKEND}_p${PROCS}c${CLIENTS}_gap${GAP}_${GETMODE}${EXTKV_TAG}" +echo "=== ${LABEL}: total_clients=$((PROCS*CLIENTS)) backend=${BACKEND} redis=${REDIS_URI} ===" + +# ---- start standalone master ---- +pkill -f "umbp_master 0.0.0.0:${PORT}" 2>/dev/null; sleep 1 +ENV="UMBP_METADATA_BACKEND=${BACKEND}" +if [[ "$BACKEND" == "redis" ]]; then + REDIS_EXTRA="UMBP_REDIS_NAMESPACE=mp_${LABEL}_$(date +%s) UMBP_REDIS_POOL_SIZE=${UMBP_REDIS_POOL_SIZE:-32} UMBP_REDIS_CONNECT_TIMEOUT_MS=1000 UMBP_REDIS_SOCKET_TIMEOUT_MS=1000" + if [[ "$REDIS_CLUSTER" == "1" ]]; then + # Cluster: unique namespace isolates runs, so no cluster-wide FLUSHALL needed. + "$REDIS_CLI" -h "$RHOST" -p "$RPORT" ping >/dev/null 2>&1 || { echo "ERROR: cannot reach cluster seed $RHOST:$RPORT"; exit 2; } + ENV="${ENV} UMBP_REDIS_CLUSTER=1 UMBP_REDIS_URI=${REDIS_URI} ${REDIS_EXTRA}" + elif [[ -n "$SHARD_URIS" ]]; then + ENV="${ENV} UMBP_REDIS_SHARD_URIS=${SHARD_URIS} ${REDIS_EXTRA}" + else + "$REDIS_CLI" -h "$RHOST" -p "$RPORT" FLUSHALL >/dev/null 2>&1 || { echo "ERROR: cannot reach redis $REDIS_URI"; exit 2; } + "$REDIS_CLI" -h "$RHOST" -p "$RPORT" CONFIG RESETSTAT >/dev/null 2>&1 + ENV="${ENV} UMBP_REDIS_URI=${REDIS_URI} ${REDIS_EXTRA}" + fi +fi +env $ENV UMBP_ROUTE_PUT_NODE_AFFINITY=local "$MASTER_BIN" "0.0.0.0:${PORT}" "$METRICS" > "${OUT}/master_${LABEL}.log" 2>&1 & +MPID=$! +for i in $(seq 1 30); do curl -sf "http://127.0.0.1:${METRICS}/metrics" >/dev/null 2>&1 && break; sleep 1; done + +# ---- cross-process end-of-run barrier (fetch/both) ---- +# All PROCS processes reach it before any tears down, so a straggler's RDMA read +# never hits a peer that already shut down (the transport-reset crash). For a +# MULTI-HOST run, point every launcher at the SAME shared-FS BARRIER_DIR and set +# --barrier-size to the TOTAL process count across hosts. +BARRIER_DIR="${BARRIER_DIR:-${OUT}/barrier}"; rm -rf "$BARRIER_DIR" 2>/dev/null; mkdir -p "$BARRIER_DIR" +BARRIER_SIZE="${BARRIER_SIZE:-$PROCS}" + +# ---- launch PROCS client processes (each CLIENTS clients) ---- +pids=() +for p in $(seq 0 $((PROCS-1))); do + "$BIN" --external-master "${HOSTIP}:${PORT}" \ + --node-id-prefix "p${p}-" --node-address "$HOSTIP" \ + --clients "$CLIENTS" --rounds "$ROUNDS" --warmup-rounds "$WARMUP" --batch "$BATCH" \ + --key-space "$KEYSPACE" --read-lag-rounds 1 --pattern rotate --get-mode "$GETMODE" \ + --gap-ms "$GAP" --mode baseline --put-affinity local --metrics-port 0 $EXTKV_ARGS \ + --barrier-dir "$BARRIER_DIR" --barrier-size "$BARRIER_SIZE" --barrier-tag "run" \ + > "${OUT}/${LABEL}_p${p}.csv" 2>&1 & + pids+=($!) +done +for pid in "${pids[@]}"; do wait "$pid" 2>/dev/null || true; done +sleep 2 # let clients' final ReportMetrics flush land at the master + +# ---- results: master per-RPC hist + redis cmdstats + aggregate qps ---- +SUM="${OUT}/${LABEL}_summary.txt" +{ + echo "# ${LABEL} procs=${PROCS} clients/proc=${CLIENTS} batch=${BATCH} gap=${GAP}ms get=${GETMODE}" + echo "## master per-RPC latency (ms)" + python3 "$PARSE" "http://127.0.0.1:${METRICS}/metrics" 2>/dev/null + if [[ "$BACKEND" == "redis" ]]; then + echo "## redis commandstats" + "$REDIS_CLI" -h "$RHOST" -p "$RPORT" INFO commandstats 2>/dev/null | grep -E "evalsha|cmdstat_eval:" + echo "## redis cpu"; "$REDIS_CLI" -h "$RHOST" -p "$RPORT" INFO cpu 2>/dev/null | grep used_cpu + fi + echo "## aggregate put/get qps (sum over processes)" + awk -F, 'FNR>1 && $1!="mode"{p+=$14; g+=$15} END{printf "put_qps_total=%.0f get_qps_total=%.0f\n", p, g}' \ + "${OUT}/${LABEL}"_p*.csv 2>/dev/null +} | tee "$SUM" + +kill "$MPID" 2>/dev/null; wait "$MPID" 2>/dev/null || true +echo "DONE -> $SUM" diff --git a/tests/cpp/umbp/distributed/test_redis_master_metadata_store.cpp b/tests/cpp/umbp/distributed/test_redis_master_metadata_store.cpp new file mode 100644 index 000000000..245e3d59d --- /dev/null +++ b/tests/cpp/umbp/distributed/test_redis_master_metadata_store.cpp @@ -0,0 +1,968 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. + +// Phase 1 targeted tests for RedisMasterMetadataStore: the six hot-path methods +// plus the register/heartbeat/unregister/expire semantics that must match the +// in-memory backend. The store tests run against a live RESP store (Redis / +// Dragonfly / Valkey) at UMBP_REDIS_URI (default tcp://127.0.0.1:6379) and skip +// cleanly when none is reachable, so BUILD_TESTS on a host without Redis does +// not fail. Every store test is parameterized over block_shards (1 = legacy +// single-tag layout, 16 = sharded) so both the whole-batch and the per-shard +// fan-out read paths are exercised with identical assertions. +// +// The KeySchema tests are pure (no store) and always run — they guard the shard +// mapping and the "single shard == legacy key strings" invariant even where no +// Redis is available. + +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "mori/metrics/prometheus_metrics_server.hpp" +#include "umbp/distributed/master/redis/cluster_slots.h" +#include "umbp/distributed/master/redis/key_schema.h" +#include "umbp/distributed/master/redis/resp_client.h" +#include "umbp/distributed/master/redis/resp_cluster_client.h" +#include "umbp/distributed/master/redis_master_metadata_store.h" + +namespace mori::umbp { +namespace { + +using namespace std::chrono_literals; + +// ===================================================================== +// KeySchema — pure unit tests (no live store needed). +// ===================================================================== + +TEST(KeySchemaTest, SingleShardIsLegacyLayout) { + redis::KeySchema schema("ns", 1); + EXPECT_EQ(schema.NumShards(), 1u); + EXPECT_EQ(schema.Tag(), "{umbp:ns}"); + EXPECT_EQ(schema.ShardOf("abc"), 0u); + // Byte-identical to the original single-tag schema. + EXPECT_EQ(schema.Block("abc"), "{umbp:ns}:block:abc"); + EXPECT_EQ(schema.Node("n1"), "{umbp:ns}:node:n1"); + EXPECT_EQ(schema.NodeBlocks("n1"), "{umbp:ns}:node:n1:blocks"); +} + +TEST(KeySchemaTest, ZeroShardsClampsToOne) { + redis::KeySchema schema("ns", 0); + EXPECT_EQ(schema.NumShards(), 1u); + EXPECT_EQ(schema.Block("x"), "{umbp:ns}:block:x"); +} + +TEST(KeySchemaTest, MultiShardIsDeterministicAndInRange) { + constexpr std::size_t kShards = 16; + redis::KeySchema schema("ns", kShards); + EXPECT_EQ(schema.NumShards(), kShards); + for (int i = 0; i < 1000; ++i) { + const std::string key = "key-" + std::to_string(i); + const std::size_t shard = schema.ShardOf(key); + EXPECT_LT(shard, kShards); + EXPECT_EQ(shard, schema.ShardOf(key)) << "shard mapping must be deterministic"; + EXPECT_EQ(schema.Block(key), "{umbp:ns:b" + std::to_string(shard) + "}:block:" + key); + } +} + +TEST(KeySchemaTest, ShardsAreReasonablySpread) { + constexpr std::size_t kShards = 16; + redis::KeySchema schema("ns", kShards); + std::vector counts(kShards, 0); + for (int i = 0; i < 4096; ++i) counts[schema.ShardOf("k/" + std::to_string(i))]++; + for (std::size_t s = 0; s < kShards; ++s) EXPECT_GT(counts[s], 0) << "shard " << s << " is empty"; +} + +// ===================================================================== +// ClusterSlots — pure unit tests for the CRC16 / hash-tag slot math used to +// place one block-shard tag per master (balanced placement). Reference slot +// values verified against `redis-cli CLUSTER KEYSLOT`. +// ===================================================================== + +TEST(ClusterSlotsTest, MatchesRedisKeyslot) { + EXPECT_EQ(redis::SlotOfKey("foo"), 12182); + EXPECT_EQ(redis::SlotOfKey("bar"), 5061); + // 0x31C3 — the canonical CRC16/XMODEM check value for "123456789". + EXPECT_EQ(redis::SlotOfKey("123456789"), 12739); + EXPECT_EQ(redis::SlotOfKey("{umbp:default:b0}"), 14816); +} + +TEST(ClusterSlotsTest, HonorsHashTag) { + // Only the {...} body is hashed, so a block key routes by its shard tag. + EXPECT_EQ(redis::SlotOfKey("{foo}bar"), redis::SlotOfKey("foo")); + EXPECT_EQ(redis::SlotOfKey("prefix{foo}suffix"), redis::SlotOfKey("foo")); + EXPECT_EQ(redis::SlotOfKey("{umbp:default:b0}:block:k/5"), redis::SlotOfKey("{umbp:default:b0}")); + // Empty braces are not a tag: the whole key is hashed. + EXPECT_EQ(redis::SlotOfKey("{}foo"), 9500); + EXPECT_NE(redis::SlotOfKey("{}foo"), redis::SlotOfKey("foo")); +} + +TEST(ClusterSlotsTest, FindTagLandsInRanges) { + const std::vector ranges = {{5000, 5500}, {10000, 10200}}; + std::string tag; + ASSERT_TRUE(redis::FindTagForRanges("ns", ranges, &tag)); + EXPECT_TRUE(redis::SlotInRanges(redis::SlotOfKey(tag), ranges)); + EXPECT_EQ(tag.rfind("{umbp:ns:b", 0), 0u); // expected tag shape +} + +TEST(ClusterSlotsTest, FindTagPerDisjointRangeIsDistinct) { + // Model "one tag per master": two disjoint ranges must yield two tags, each + // in its own range (so different masters get different, correctly-placed tags). + const std::vector a = {{0, 5460}}; + const std::vector b = {{10923, 16383}}; + std::string ta, tb; + ASSERT_TRUE(redis::FindTagForRanges("ns", a, &ta)); + ASSERT_TRUE(redis::FindTagForRanges("ns", b, &tb)); + EXPECT_NE(ta, tb); + EXPECT_TRUE(redis::SlotInRanges(redis::SlotOfKey(ta), a)); + EXPECT_TRUE(redis::SlotInRanges(redis::SlotOfKey(tb), b)); +} + +// ===================================================================== +// Store tests — parameterized over block_shards, require a live RESP store. +// ===================================================================== + +std::string RedisUri() { + const char* v = std::getenv("UMBP_REDIS_URI"); + return (v != nullptr && *v != '\0') ? std::string(v) : std::string("tcp://127.0.0.1:6379"); +} + +// Unique namespace per process so parallel runs / leftover keys never collide. +std::string UniqueNamespace() { + return "test_" + std::to_string(::getpid()) + "_" + + std::to_string(std::chrono::steady_clock::now().time_since_epoch().count() & 0xffffff); +} + +// Cluster seeds for the "cluster" StoreMode, from UMBP_REDIS_CLUSTER_SEEDS +// (comma-separated tcp:// URIs). Empty => cluster tests skip. +std::vector ClusterSeeds() { + std::vector out; + const char* v = std::getenv("UMBP_REDIS_CLUSTER_SEEDS"); + if (v == nullptr || *v == '\0') return out; + std::string s(v); + size_t pos = 0; + while (pos <= s.size()) { + const size_t comma = s.find(',', pos); + const size_t end = (comma == std::string::npos) ? s.size() : comma; + std::string tok = s.substr(pos, end - pos); + if (!tok.empty()) out.push_back(tok); + if (comma == std::string::npos) break; + pos = comma + 1; + } + return out; +} + +// A store mode to run every assertion under: single-endpoint with N block +// shards, or multi-endpoint with E endpoints (one shard each). The multi mode +// points all logical endpoints at the same physical Redis (distinct hash tags +// keep the keyspaces disjoint), so it exercises the split-write / per-endpoint +// fan-out code paths without needing several Redis processes. +struct StoreMode { + std::size_t block_shards; // single-endpoint shard count (endpoints == 1) + std::size_t endpoints; // >1 => multi-endpoint mode (block_shards ignored) + const char* name; + bool cluster = false; // true => Redis Cluster mode (needs UMBP_REDIS_CLUSTER_SEEDS) + bool balanced = false; // cluster only: compute one balanced tag per master + // (exercises the factory's balanced-placement path) +}; + +class RedisStoreTest : public ::testing::TestWithParam { + protected: + void SetUp() override { + const StoreMode& m = GetParam(); + ns_ = UniqueNamespace(); + RedisMasterMetadataStore::Config cfg; + cfg.namespace_id = ns_; + + if (m.cluster) { + const std::vector seeds = ClusterSeeds(); + if (seeds.empty()) { + GTEST_SKIP() << "no UMBP_REDIS_CLUSTER_SEEDS set; skipping cluster mode"; + } + redis::RespClusterClient::Options popts; + popts.seeds = seeds; + popts.pool_size = 2; + try { + probe_ = std::make_unique(popts); + } catch (const std::exception& e) { + GTEST_SKIP() << "no Redis Cluster reachable (" << e.what() << "); skipping"; + } + if (!probe_->Ping()) GTEST_SKIP() << "Redis Cluster not reachable; skipping"; + cfg.cluster = true; + cfg.cluster_seeds = seeds; + if (m.balanced) { + // Balanced placement: same computation the factory does — one tag per + // master, each on a slot that master owns. + std::vector> ranges; + try { + ranges = redis::RespClusterClient::DiscoverMasterSlotRanges(popts); + } catch (const std::exception& e) { + GTEST_SKIP() << "CLUSTER SLOTS discovery failed (" << e.what() << "); skipping"; + } + for (const auto& node_ranges : ranges) { + std::string tag; + if (redis::FindTagForRanges(ns_, node_ranges, &tag)) block_tags_.push_back(tag); + } + if (block_tags_.empty() || block_tags_.size() != ranges.size()) { + GTEST_SKIP() << "balanced tag search incomplete; skipping"; + } + cfg.cluster_block_tags = block_tags_; + cfg.block_shards = block_tags_.size(); + } else { + cfg.block_shards = m.block_shards; + } + } else { + redis::RespClient::Options opts; + opts.uri = RedisUri(); + opts.pool_size = 2; + probe_ = std::make_unique(opts); + if (!probe_->Ping()) { + GTEST_SKIP() << "no RESP store reachable at " << RedisUri() + << " (set UMBP_REDIS_URI); skipping"; + } + cfg.uri = RedisUri(); + if (m.endpoints > 1) { + cfg.shard_uris.assign(m.endpoints, RedisUri()); + } else { + cfg.block_shards = m.block_shards; + } + } + store_ = std::make_unique(cfg); + now_ = std::chrono::system_clock::now(); + } + + // Number of block shards the store built (must match KeySchema for MetaField). + std::size_t NumShards() const { + const StoreMode& m = GetParam(); + if (m.cluster) return m.block_shards; + return m.endpoints > 1 ? m.endpoints : m.block_shards; + } + + ClientRegistration MakeReg(const std::string& id) { + ClientRegistration r; + r.node_id = id; + r.node_address = id + ".addr"; + r.peer_address = id + ".peer:1234"; + r.tier_capacities[TierType::DRAM] = TierCapacity{1u << 30, 1u << 30}; + r.tags = {"sgl_role=prefill", "zone=a"}; + return r; + } + + // Raw HGET of a block-hash meta field via the probe client, so tests can + // assert lease/access bookkeeping the store API does not expose. Returns -1 + // when the field (or the whole block key) is absent. Uses the balanced tags + // when the store was built with them, else the formulaic schema. + long long MetaField(const std::string& user_key, const std::string& field) const { + redis::KeySchema schema = block_tags_.empty() ? redis::KeySchema(ns_, NumShards()) + : redis::KeySchema(ns_, block_tags_); + redis::RespValue r = probe_->Command({"HGET", schema.Block(user_key), field}); + if (r.type != redis::RespValue::Type::String) return -1; + try { + return std::stoll(r.str); + } catch (...) { + return -1; + } + } + + std::string ns_; + // Non-empty only in cluster-balanced mode: the per-master tags the store was + // built with, so MetaField reconstructs the same block keys. + std::vector block_tags_; + std::unique_ptr probe_; + std::unique_ptr store_; + std::chrono::system_clock::time_point now_; +}; + +INSTANTIATE_TEST_SUITE_P( + BlockShards, RedisStoreTest, + ::testing::Values(StoreMode{1, 1, "shards1"}, StoreMode{16, 1, "shards16"}, + StoreMode{1, 3, "endpoints3"}, + // Redis Cluster: 16 block-shard tags spread across nodes by + // slot. Skips unless UMBP_REDIS_CLUSTER_SEEDS is set. + StoreMode{16, 1, "cluster", true}, + // Redis Cluster with balanced placement: one tag per master + // (the factory's default path). block_shards is derived. + StoreMode{0, 1, "clusterbalanced", true, true}), + [](const ::testing::TestParamInfo& info) { return std::string(info.param.name); }); + +TEST_P(RedisStoreTest, RegisterMakesClientAlive) { + EXPECT_TRUE(store_->RegisterClient(MakeReg("n1"), now_, 30s)); + EXPECT_TRUE(store_->IsClientAlive("n1")); + EXPECT_EQ(store_->AliveClientCount(), 1u); + + auto peers = store_->GetAlivePeerView(); + ASSERT_EQ(peers.count("n1"), 1u); + EXPECT_EQ(peers["n1"], "n1.peer:1234"); + + auto rec = store_->GetClient("n1"); + ASSERT_TRUE(rec.has_value()); + EXPECT_EQ(rec->node_address, "n1.addr"); + EXPECT_EQ(rec->status, ClientStatus::ALIVE); + ASSERT_EQ(rec->tier_capacities.count(TierType::DRAM), 1u); + EXPECT_EQ(rec->tier_capacities[TierType::DRAM].total_bytes, 1u << 30); + EXPECT_EQ(store_->GetClientTags("n1").size(), 2u); +} + +TEST_P(RedisStoreTest, RegisterRejectsAliveDuplicateButAllowsStale) { + EXPECT_TRUE(store_->RegisterClient(MakeReg("n1"), now_, 30s)); + // Alive and not stale -> rejected. + EXPECT_FALSE(store_->RegisterClient(MakeReg("n1"), now_, 30s)); + // Stale (stale_after = 0) -> allowed to replace. + EXPECT_TRUE(store_->RegisterClient(MakeReg("n1"), now_ + 1s, 0s)); +} + +TEST_P(RedisStoreTest, ListAliveClientsReturnsRecords) { + ASSERT_TRUE(store_->RegisterClient(MakeReg("n1"), now_, 30s)); + ASSERT_TRUE(store_->RegisterClient(MakeReg("n2"), now_, 30s)); + auto alive = store_->ListAliveClients(); + EXPECT_EQ(alive.size(), 2u); +} + +TEST_P(RedisStoreTest, HeartbeatAddThenRouteGet) { + ASSERT_TRUE(store_->RegisterClient(MakeReg("n1"), now_, 30s)); + + std::vector events = { + KvEvent{KvEvent::Kind::ADD, "k1", TierType::DRAM, 4096}, + KvEvent{KvEvent::Kind::ADD, "k2", TierType::HBM, 8192}, + }; + auto res = store_->ApplyHeartbeat("n1", 1, now_, MakeReg("n1").tier_capacities, events, false); + EXPECT_EQ(res.status, HeartbeatResult::APPLIED); + EXPECT_EQ(res.acked_seq, 1u); + + auto exists = store_->BatchExistsBlock({"k1", "k2", "missing"}); + ASSERT_EQ(exists.size(), 3u); + EXPECT_TRUE(exists[0]); + EXPECT_TRUE(exists[1]); + EXPECT_FALSE(exists[2]); + + auto locs = store_->BatchLookupBlockForRouteGet({"k1", "missing"}, {}, now_, 10s); + ASSERT_EQ(locs.size(), 2u); + ASSERT_EQ(locs[0].size(), 1u); + EXPECT_EQ(locs[0][0].node_id, "n1"); + EXPECT_EQ(locs[0][0].tier, TierType::DRAM); + EXPECT_EQ(locs[0][0].size, 4096u); + EXPECT_TRUE(locs[1].empty()); +} + +TEST_P(RedisStoreTest, RouteGetExcludeFiltersNode) { + ASSERT_TRUE(store_->RegisterClient(MakeReg("n1"), now_, 30s)); + ASSERT_TRUE(store_->RegisterClient(MakeReg("n2"), now_, 30s)); + ASSERT_EQ( + store_ + ->ApplyHeartbeat("n1", 1, now_, {}, {{KvEvent::Kind::ADD, "k", TierType::DRAM, 1}}, false) + .status, + HeartbeatResult::APPLIED); + ASSERT_EQ( + store_ + ->ApplyHeartbeat("n2", 1, now_, {}, {{KvEvent::Kind::ADD, "k", TierType::DRAM, 1}}, false) + .status, + HeartbeatResult::APPLIED); + + auto locs = store_->BatchLookupBlockForRouteGet({"k"}, {"n1"}, now_, 10s); + ASSERT_EQ(locs.size(), 1u); + ASSERT_EQ(locs[0].size(), 1u); + EXPECT_EQ(locs[0][0].node_id, "n2"); +} + +TEST_P(RedisStoreTest, RouteGetBumpsAccessOnlyForTouchedKeys) { + ASSERT_TRUE(store_->RegisterClient(MakeReg("n1"), now_, 30s)); + ASSERT_EQ(store_ + ->ApplyHeartbeat("n1", 1, now_, {}, + {{KvEvent::Kind::ADD, "hit", TierType::DRAM, 7}}, false) + .status, + HeartbeatResult::APPLIED); + + // Freshly added key starts at _acnt = 0 and no lease. + EXPECT_EQ(MetaField("hit", "_acnt"), 0); + EXPECT_EQ(MetaField("hit", "_lease"), -1); + + // One RouteGet over a hit + a miss: only the hit is leased / access-bumped; + // the miss must not be created. + auto locs = store_->BatchLookupBlockForRouteGet({"hit", "miss"}, {}, now_, 10s); + ASSERT_EQ(locs.size(), 2u); + ASSERT_EQ(locs[0].size(), 1u); + EXPECT_TRUE(locs[1].empty()); + + EXPECT_EQ(MetaField("hit", "_acnt"), 1); // touched -> bumped exactly once + EXPECT_GT(MetaField("hit", "_lease"), 0); // touched -> lease granted + EXPECT_EQ(MetaField("miss", "_acnt"), -1); // absent -> still absent (no phantom key) + + // A second RouteGet bumps again by exactly one (guards the NOSCRIPT + // retry-only-failed path from double-applying the write). + store_->BatchLookupBlockForRouteGet({"hit"}, {}, now_, 10s); + EXPECT_EQ(MetaField("hit", "_acnt"), 2); +} + +TEST_P(RedisStoreTest, HeartbeatSeqGapAndUnknown) { + EXPECT_EQ(store_->ApplyHeartbeat("ghost", 1, now_, {}, {}, false).status, + HeartbeatResult::UNKNOWN); + + ASSERT_TRUE(store_->RegisterClient(MakeReg("n1"), now_, 30s)); + // Expect seq 1 next; sending 3 is a gap. + auto gap = store_->ApplyHeartbeat("n1", 3, now_, {}, {}, false); + EXPECT_EQ(gap.status, HeartbeatResult::SEQ_GAP); + EXPECT_EQ(gap.acked_seq, 0u); +} + +TEST_P(RedisStoreTest, HeartbeatRemoveDropsLocation) { + ASSERT_TRUE(store_->RegisterClient(MakeReg("n1"), now_, 30s)); + ASSERT_EQ( + store_ + ->ApplyHeartbeat("n1", 1, now_, {}, {{KvEvent::Kind::ADD, "k", TierType::DRAM, 1}}, false) + .status, + HeartbeatResult::APPLIED); + ASSERT_TRUE(store_->BatchExistsBlock({"k"})[0]); + + ASSERT_EQ(store_ + ->ApplyHeartbeat("n1", 2, now_, {}, + {{KvEvent::Kind::REMOVE, "k", TierType::DRAM, 0}}, false) + .status, + HeartbeatResult::APPLIED); + EXPECT_FALSE(store_->BatchExistsBlock({"k"})[0]); +} + +TEST_P(RedisStoreTest, FullSyncReplacesLocations) { + ASSERT_TRUE(store_->RegisterClient(MakeReg("n1"), now_, 30s)); + ASSERT_EQ(store_ + ->ApplyHeartbeat("n1", 1, now_, {}, + {{KvEvent::Kind::ADD, "old", TierType::DRAM, 1}}, false) + .status, + HeartbeatResult::APPLIED); + // Full sync with a different key set replaces wholesale. + ASSERT_EQ(store_ + ->ApplyHeartbeat("n1", 5, now_, {}, + {{KvEvent::Kind::ADD, "new", TierType::DRAM, 2}}, true) + .status, + HeartbeatResult::APPLIED); + EXPECT_FALSE(store_->BatchExistsBlock({"old"})[0]); + EXPECT_TRUE(store_->BatchExistsBlock({"new"})[0]); +} + +TEST_P(RedisStoreTest, UnregisterWipesClientAndBlocks) { + ASSERT_TRUE(store_->RegisterClient(MakeReg("n1"), now_, 30s)); + ASSERT_EQ( + store_ + ->ApplyHeartbeat("n1", 1, now_, {}, {{KvEvent::Kind::ADD, "k", TierType::DRAM, 1}}, false) + .status, + HeartbeatResult::APPLIED); + + store_->UnregisterClient("n1"); + EXPECT_FALSE(store_->IsClientAlive("n1")); + EXPECT_EQ(store_->AliveClientCount(), 0u); + EXPECT_FALSE(store_->BatchExistsBlock({"k"})[0]); +} + +TEST_P(RedisStoreTest, ExpireStaleFlipsAndWipesBlocks) { + ASSERT_TRUE(store_->RegisterClient(MakeReg("n1"), now_, 30s)); + ASSERT_EQ( + store_ + ->ApplyHeartbeat("n1", 1, now_, {}, {{KvEvent::Kind::ADD, "k", TierType::DRAM, 1}}, false) + .status, + HeartbeatResult::APPLIED); + + // cutoff in the future -> n1's last_hb (==now_) is older -> expired. + auto dead = store_->ExpireStaleClients(now_ + 10s); + ASSERT_EQ(dead.size(), 1u); + EXPECT_EQ(dead[0], "n1"); + EXPECT_FALSE(store_->IsClientAlive("n1")); + EXPECT_FALSE(store_->BatchExistsBlock({"k"})[0]); + // EXPIRED record is kept (hazard #3). + auto rec = store_->GetClient("n1"); + ASSERT_TRUE(rec.has_value()); + EXPECT_EQ(rec->status, ClientStatus::EXPIRED); +} + +// --------------------------------------------------------------------- +// Cross-shard coverage: with block_shards=16 the keys below spread over +// multiple shards, so these exercise the per-shard fan-out + scatter-merge and +// the multi-shard wipe paths (a no-op difference when block_shards=1). +// --------------------------------------------------------------------- + +TEST_P(RedisStoreTest, CrossShardBatchPreservesOrderAndMapping) { + ASSERT_TRUE(store_->RegisterClient(MakeReg("n1"), now_, 30s)); + + constexpr int kKeys = 20; + std::vector events; + events.reserve(kKeys); + for (int i = 0; i < kKeys; ++i) { + // size encodes identity (100 + i) so the merge can be verified per key. + events.push_back({KvEvent::Kind::ADD, "key-" + std::to_string(i), TierType::DRAM, + static_cast(100 + i)}); + } + ASSERT_EQ(store_->ApplyHeartbeat("n1", 1, now_, {}, events, false).status, + HeartbeatResult::APPLIED); + + // Scrambled query order, interleaved with a missing key. + const std::vector query = {"key-5", "missing", "key-0", "key-19", "key-12"}; + auto locs = store_->BatchLookupBlockForRouteGet(query, {}, now_, 10s); + ASSERT_EQ(locs.size(), query.size()); + ASSERT_EQ(locs[0].size(), 1u); + EXPECT_EQ(locs[0][0].size, 105u); + EXPECT_TRUE(locs[1].empty()); + ASSERT_EQ(locs[2].size(), 1u); + EXPECT_EQ(locs[2][0].size, 100u); + ASSERT_EQ(locs[3].size(), 1u); + EXPECT_EQ(locs[3][0].size, 119u); + ASSERT_EQ(locs[4].size(), 1u); + EXPECT_EQ(locs[4][0].size, 112u); + + // BatchExistsBlock in the same scrambled order maps back correctly too. + auto exists = store_->BatchExistsBlock(query); + ASSERT_EQ(exists.size(), query.size()); + EXPECT_TRUE(exists[0]); + EXPECT_FALSE(exists[1]); + EXPECT_TRUE(exists[2]); + EXPECT_TRUE(exists[3]); + EXPECT_TRUE(exists[4]); +} + +TEST_P(RedisStoreTest, FullSyncAcrossShardsWipesAllOldLocations) { + ASSERT_TRUE(store_->RegisterClient(MakeReg("n1"), now_, 30s)); + + constexpr int kKeys = 12; + std::vector old_events; + std::vector new_events; + std::vector old_keys; + std::vector new_keys; + for (int i = 0; i < kKeys; ++i) { + old_keys.push_back("old-" + std::to_string(i)); + new_keys.push_back("new-" + std::to_string(i)); + old_events.push_back({KvEvent::Kind::ADD, old_keys.back(), TierType::DRAM, 1}); + new_events.push_back({KvEvent::Kind::ADD, new_keys.back(), TierType::DRAM, 2}); + } + ASSERT_EQ(store_->ApplyHeartbeat("n1", 1, now_, {}, old_events, false).status, + HeartbeatResult::APPLIED); + ASSERT_EQ(store_->ApplyHeartbeat("n1", 9, now_, {}, new_events, true).status, + HeartbeatResult::APPLIED); + + for (bool e : store_->BatchExistsBlock(old_keys)) EXPECT_FALSE(e); + for (bool e : store_->BatchExistsBlock(new_keys)) EXPECT_TRUE(e); +} + +TEST_P(RedisStoreTest, UnregisterWipesBlocksAcrossShards) { + ASSERT_TRUE(store_->RegisterClient(MakeReg("n1"), now_, 30s)); + + constexpr int kKeys = 12; + std::vector events; + std::vector keys; + for (int i = 0; i < kKeys; ++i) { + keys.push_back("u-" + std::to_string(i)); + events.push_back({KvEvent::Kind::ADD, keys.back(), TierType::DRAM, 1}); + } + ASSERT_EQ(store_->ApplyHeartbeat("n1", 1, now_, {}, events, false).status, + HeartbeatResult::APPLIED); + ASSERT_TRUE(store_->BatchExistsBlock({keys.front()})[0]); + + store_->UnregisterClient("n1"); + for (bool e : store_->BatchExistsBlock(keys)) EXPECT_FALSE(e); +} + +// ===================================================================== +// External-KV: register / match / unregister, tier bitmask, reverse-index count. +// Asserts the master_metadata_store.h contract (not in-memory internals); runs +// in every StoreMode (extkv/hit live on the control tag in all modes). +// ===================================================================== + +TEST_P(RedisStoreTest, ExternalKvAliveGateAndMatch) { + // Unknown / not-alive node is rejected and writes nothing. + EXPECT_FALSE(store_->RegisterExternalKvIfAlive("ghost", {"h1"}, TierType::HBM)); + EXPECT_TRUE(store_->MatchExternalKv({"h1"}, false, now_).empty()); + + ASSERT_TRUE(store_->RegisterClient(MakeReg("n1"), now_, 30s)); + EXPECT_TRUE(store_->RegisterExternalKvIfAlive("n1", {"h1", "h2"}, TierType::HBM)); + + auto m = store_->MatchExternalKv({"h1", "h2"}, false, now_); + ASSERT_EQ(m.size(), 1u); + EXPECT_EQ(m[0].node_id, "n1"); + ASSERT_EQ(m[0].hashes_by_tier.count(TierType::HBM), 1u); + EXPECT_EQ(m[0].hashes_by_tier[TierType::HBM].size(), 2u); + EXPECT_EQ(store_->GetExternalKvCount("n1"), 2u); +} + +TEST_P(RedisStoreTest, ExternalKvUnregisterByTierAndHash) { + ASSERT_TRUE(store_->RegisterClient(MakeReg("n1"), now_, 30s)); + ASSERT_TRUE(store_->RegisterExternalKvIfAlive("n1", {"h1"}, TierType::HBM)); + ASSERT_TRUE(store_->RegisterExternalKvIfAlive("n1", {"h1"}, TierType::DRAM)); + + // Drop only HBM; DRAM remains, so the (node,hash) entry survives. + store_->UnregisterExternalKv("n1", {"h1"}, TierType::HBM); + auto m = store_->MatchExternalKv({"h1"}, false, now_); + ASSERT_EQ(m.size(), 1u); + EXPECT_EQ(m[0].hashes_by_tier.count(TierType::HBM), 0u); + EXPECT_EQ(m[0].hashes_by_tier.count(TierType::DRAM), 1u); + EXPECT_EQ(store_->GetExternalKvCount("n1"), 1u); + + // Whole-tier wipe of DRAM removes the last tier → entry gone. + store_->UnregisterExternalKvByTier("n1", TierType::DRAM); + EXPECT_TRUE(store_->MatchExternalKv({"h1"}, false, now_).empty()); + EXPECT_EQ(store_->GetExternalKvCount("n1"), 0u); +} + +TEST_P(RedisStoreTest, ExternalKvMatchedHashCountAcrossTiers) { + ASSERT_TRUE(store_->RegisterClient(MakeReg("n1"), now_, 30s)); + ASSERT_TRUE(store_->RegisterExternalKvIfAlive("n1", {"h1"}, TierType::HBM)); + ASSERT_TRUE(store_->RegisterExternalKvIfAlive("n1", {"h1"}, TierType::DRAM)); + + auto m = store_->MatchExternalKv({"h1"}, false, now_); + ASSERT_EQ(m.size(), 1u); + EXPECT_EQ(m[0].hashes_by_tier.size(), 2u); // one hash, two tier buckets + EXPECT_EQ(m[0].MatchedHashCount(), 1u); // deduped to one unique hash +} + +TEST_P(RedisStoreTest, ExternalKvUnregisterByNodeWipesAllTiers) { + ASSERT_TRUE(store_->RegisterClient(MakeReg("n1"), now_, 30s)); + ASSERT_TRUE(store_->RegisterExternalKvIfAlive("n1", {"h1", "h2"}, TierType::HBM)); + ASSERT_TRUE(store_->RegisterExternalKvIfAlive("n1", {"h1"}, TierType::DRAM)); + ASSERT_EQ(store_->GetExternalKvCount("n1"), 2u); + + store_->UnregisterExternalKvByNode("n1"); + EXPECT_EQ(store_->GetExternalKvCount("n1"), 0u); + EXPECT_TRUE(store_->MatchExternalKv({"h1", "h2"}, false, now_).empty()); + // Must NOT have touched the client record. + EXPECT_TRUE(store_->IsClientAlive("n1")); +} + +TEST_P(RedisStoreTest, ExternalKvHitCounting) { + ASSERT_TRUE(store_->RegisterClient(MakeReg("n1"), now_, 30s)); + ASSERT_TRUE(store_->RegisterExternalKvIfAlive("n1", {"h1", "h2"}, TierType::HBM)); + + // Pure read leaves the hit map untouched. + store_->MatchExternalKv({"h1", "h2"}, /*count_as_hit=*/false, now_); + EXPECT_TRUE(store_->GetExternalKvHitCounts({"h1", "h2"}).empty()); + + // count_as_hit accumulates across calls; increments once per unique hash. + store_->MatchExternalKv({"h1", "h2"}, /*count_as_hit=*/true, now_); + store_->MatchExternalKv({"h1"}, /*count_as_hit=*/true, now_ + 1s); + + auto counts = store_->GetExternalKvHitCounts({"h1", "h2"}); + std::map by_hash; + for (const auto& e : counts) by_hash[e.hash] = e.hit_count_total; + EXPECT_EQ(by_hash["h1"], 2u); + EXPECT_EQ(by_hash["h2"], 1u); + + // GetExternalKvHitCounts dedupes and skips hashes with no count. + auto c2 = store_->GetExternalKvHitCounts({"missing", "h1", "h1"}); + ASSERT_EQ(c2.size(), 1u); + EXPECT_EQ(c2[0].hash, "h1"); + EXPECT_EQ(c2[0].hit_count_total, 2u); +} + +TEST_P(RedisStoreTest, GarbageCollectHitsByLastSeen) { + ASSERT_TRUE(store_->RegisterClient(MakeReg("n1"), now_, 30s)); + ASSERT_TRUE(store_->RegisterExternalKvIfAlive("n1", {"old", "fresh"}, TierType::HBM)); + store_->MatchExternalKv({"old"}, true, now_); + store_->MatchExternalKv({"fresh"}, true, now_ + 100s); + + // Drop entries last seen before now_+50s → only "old" goes. + EXPECT_EQ(store_->GarbageCollectHits(now_ + 50s), 1u); + auto counts = store_->GetExternalKvHitCounts({"old", "fresh"}); + ASSERT_EQ(counts.size(), 1u); + EXPECT_EQ(counts[0].hash, "fresh"); +} + +// Cascade fix: UnregisterClient / ExpireStaleClients must drop the node's +// external-kv entries, not just the reverse index (previously orphaned). +TEST_P(RedisStoreTest, UnregisterClientCascadesExternalKv) { + ASSERT_TRUE(store_->RegisterClient(MakeReg("n1"), now_, 30s)); + ASSERT_TRUE(store_->RegisterExternalKvIfAlive("n1", {"h1", "h2"}, TierType::HBM)); + ASSERT_EQ(store_->GetExternalKvCount("n1"), 2u); + + store_->UnregisterClient("n1"); + EXPECT_EQ(store_->GetExternalKvCount("n1"), 0u); + EXPECT_TRUE(store_->MatchExternalKv({"h1", "h2"}, false, now_).empty()); +} + +TEST_P(RedisStoreTest, ExpireStaleCascadesExternalKv) { + ASSERT_TRUE(store_->RegisterClient(MakeReg("n1"), now_, 30s)); + ASSERT_TRUE(store_->RegisterExternalKvIfAlive("n1", {"h1"}, TierType::HBM)); + + auto dead = store_->ExpireStaleClients(now_ + 10s); + ASSERT_EQ(dead.size(), 1u); + EXPECT_EQ(store_->GetExternalKvCount("n1"), 0u); + EXPECT_TRUE(store_->MatchExternalKv({"h1"}, false, now_).empty()); +} + +// ===================================================================== +// EnumerateEvictionCandidates: reverse-index scan (route_get_batch untouched). +// ===================================================================== + +namespace { +std::vector Buckets(const std::string& node, TierType tier) { + return {NodeTierKey{node, tier}}; +} +} // namespace + +TEST_P(RedisStoreTest, EvictionLruOrderAndCap) { + ASSERT_TRUE(store_->RegisterClient(MakeReg("n1"), now_, 30s)); + // Each ADD stamps _lacc = the heartbeat's now, so three ADDs at increasing + // times give LRU order k_old < k_mid < k_new. + ASSERT_EQ(store_ + ->ApplyHeartbeat("n1", 1, now_, {}, + {{KvEvent::Kind::ADD, "k_old", TierType::HBM, 100}}, false) + .status, + HeartbeatResult::APPLIED); + ASSERT_EQ(store_ + ->ApplyHeartbeat("n1", 2, now_ + 1s, {}, + {{KvEvent::Kind::ADD, "k_mid", TierType::HBM, 100}}, false) + .status, + HeartbeatResult::APPLIED); + ASSERT_EQ(store_ + ->ApplyHeartbeat("n1", 3, now_ + 2s, {}, + {{KvEvent::Kind::ADD, "k_new", TierType::HBM, 100}}, false) + .status, + HeartbeatResult::APPLIED); + + auto cands = store_->EnumerateEvictionCandidates(Buckets("n1", TierType::HBM), + EvictionOrder::kLeastRecentlyAccessed, + /*max_per_bucket=*/2, now_ + 10s); + ASSERT_EQ(cands.size(), 1u); + auto& bucket = cands.at(NodeTierKey{"n1", TierType::HBM}); + ASSERT_EQ(bucket.size(), 2u); + EXPECT_EQ(bucket[0].key, "k_old"); // oldest first + EXPECT_EQ(bucket[1].key, "k_mid"); + EXPECT_EQ(bucket[0].location.node_id, "n1"); + EXPECT_EQ(bucket[0].location.tier, TierType::HBM); + EXPECT_EQ(bucket[0].size, 100u); +} + +TEST_P(RedisStoreTest, EvictionSkipsLeased) { + ASSERT_TRUE(store_->RegisterClient(MakeReg("n1"), now_, 30s)); + ASSERT_EQ(store_ + ->ApplyHeartbeat("n1", 1, now_, {}, + {{KvEvent::Kind::ADD, "k1", TierType::HBM, 100}}, false) + .status, + HeartbeatResult::APPLIED); + // Lease k1 far past the enumeration time. + store_->BatchLookupBlockForRouteGet({"k1"}, {}, now_, 1h); + EXPECT_TRUE(store_ + ->EnumerateEvictionCandidates(Buckets("n1", TierType::HBM), + EvictionOrder::kLeastRecentlyAccessed, 0, now_ + 1s) + .empty()); +} + +TEST_P(RedisStoreTest, EvictionOnlyRequestedBuckets) { + ASSERT_TRUE(store_->RegisterClient(MakeReg("n1"), now_, 30s)); + ASSERT_EQ(store_ + ->ApplyHeartbeat("n1", 1, now_, {}, + {{KvEvent::Kind::ADD, "kh", TierType::HBM, 10}, + {KvEvent::Kind::ADD, "kd", TierType::DRAM, 10}}, + false) + .status, + HeartbeatResult::APPLIED); + auto cands = store_->EnumerateEvictionCandidates( + Buckets("n1", TierType::HBM), EvictionOrder::kLeastRecentlyAccessed, 0, now_ + 1s); + ASSERT_EQ(cands.size(), 1u); + EXPECT_EQ(cands.begin()->first.tier, TierType::HBM); + EXPECT_EQ(cands.begin()->second.size(), 1u); + EXPECT_EQ(cands.begin()->second[0].key, "kh"); +} + +// Tie timestamps (a batch RouteGet stamps one `now` across keys) and cross-shard +// aggregation: all candidates must appear, none dropped. +TEST_P(RedisStoreTest, EvictionTieTimestampsAndCrossShard) { + ASSERT_TRUE(store_->RegisterClient(MakeReg("n1"), now_, 30s)); + constexpr int kKeys = 30; + std::vector adds; + for (int i = 0; i < kKeys; ++i) { + adds.push_back({KvEvent::Kind::ADD, "ek-" + std::to_string(i), TierType::HBM, 10}); + } + ASSERT_EQ(store_->ApplyHeartbeat("n1", 1, now_, {}, adds, false).status, + HeartbeatResult::APPLIED); + + auto cands = store_->EnumerateEvictionCandidates(Buckets("n1", TierType::HBM), + EvictionOrder::kLeastRecentlyAccessed, + /*max_per_bucket=*/0, now_ + 10s); + ASSERT_EQ(cands.size(), 1u); + EXPECT_EQ(cands.at(NodeTierKey{"n1", TierType::HBM}).size(), static_cast(kKeys)); +} + +// Fault tolerance: with one shard instance down, reads for keys on the healthy +// shards still resolve and keys on the dead shard degrade to a miss (no throw). +// Control + shard 0 live; shard 1 points at a closed port. +TEST(RedisFaultToleranceTest, DownShardDegradesToMissNotError) { + redis::RespClient::Options opts; + opts.uri = RedisUri(); + opts.pool_size = 2; + redis::RespClient probe(opts); + if (!probe.Ping()) { + GTEST_SKIP() << "no RESP store reachable at " << RedisUri() << "; skipping"; + } + + const std::string ns = UniqueNamespace(); + // Two shards: pick a key on each (shard 0 = live endpoint 0, shard 1 = dead). + redis::KeySchema schema(ns, 2); + std::string live_key, dead_key; + for (int i = 0; (live_key.empty() || dead_key.empty()) && i < 100000; ++i) { + const std::string k = "ft-" + std::to_string(i); + if (schema.ShardOf(k) == 0 && live_key.empty()) live_key = k; + if (schema.ShardOf(k) == 1 && dead_key.empty()) dead_key = k; + } + ASSERT_FALSE(live_key.empty()); + ASSERT_FALSE(dead_key.empty()); + + RedisMasterMetadataStore::Config cfg; + cfg.namespace_id = ns; + cfg.connect_timeout_ms = 300; // fail fast on the dead endpoint + cfg.socket_timeout_ms = 300; + cfg.shard_uris = {RedisUri(), "tcp://127.0.0.1:6399"}; // 6399: nothing listening + RedisMasterMetadataStore store(cfg); + + const auto now = std::chrono::system_clock::now(); + ClientRegistration reg; + reg.node_id = "ftn"; + reg.node_address = "a"; + reg.peer_address = "p:1"; + reg.tier_capacities[TierType::DRAM] = TierCapacity{1u << 20, 1u << 20}; + ASSERT_TRUE(store.RegisterClient(reg, now, 30s)); // control endpoint is live + + // Seed only the live shard (a delta touching just shard 0's instance). + ASSERT_EQ(store + .ApplyHeartbeat("ftn", 1, now, {}, + {{KvEvent::Kind::ADD, live_key, TierType::DRAM, 7}}, false) + .status, + HeartbeatResult::APPLIED); + + // Batch spanning the live and the dead shard must NOT throw; live key resolves, + // dead-shard key reads as a miss. + auto exists = store.BatchExistsBlock({live_key, dead_key}); + ASSERT_EQ(exists.size(), 2u); + EXPECT_TRUE(exists[0]); + EXPECT_FALSE(exists[1]); + + auto locs = store.BatchLookupBlockForRouteGet({live_key, dead_key}, {}, now, 10s); + ASSERT_EQ(locs.size(), 2u); + ASSERT_EQ(locs[0].size(), 1u); + EXPECT_EQ(locs[0][0].size, 7u); + EXPECT_TRUE(locs[1].empty()); +} + +// Write fault tolerance: a heartbeat whose events span a down shard does NOT +// error the RPC — it returns SEQ_GAP (so the peer full_syncs and self-heals) and +// the node stays ALIVE. Events on the live shard are still applied. +TEST(RedisWriteFaultToleranceTest, HeartbeatToDownShardSelfHealsNotError) { + redis::RespClient::Options opts; + opts.uri = RedisUri(); + opts.pool_size = 2; + redis::RespClient probe(opts); + if (!probe.Ping()) { + GTEST_SKIP() << "no RESP store reachable at " << RedisUri() << "; skipping"; + } + + const std::string ns = UniqueNamespace(); + redis::KeySchema schema(ns, 2); + std::string live_key, dead_key; + for (int i = 0; (live_key.empty() || dead_key.empty()) && i < 100000; ++i) { + const std::string k = "wft-" + std::to_string(i); + if (schema.ShardOf(k) == 0 && live_key.empty()) live_key = k; + if (schema.ShardOf(k) == 1 && dead_key.empty()) dead_key = k; + } + ASSERT_FALSE(live_key.empty()); + ASSERT_FALSE(dead_key.empty()); + + RedisMasterMetadataStore::Config cfg; + cfg.namespace_id = ns; + cfg.connect_timeout_ms = 300; + cfg.socket_timeout_ms = 300; + cfg.shard_uris = {RedisUri(), "tcp://127.0.0.1:6399"}; // shard 1 down + RedisMasterMetadataStore store(cfg); + + const auto now = std::chrono::system_clock::now(); + ClientRegistration reg; + reg.node_id = "wn"; + reg.node_address = "a"; + reg.peer_address = "p:1"; + reg.tier_capacities[TierType::DRAM] = TierCapacity{1u << 20, 1u << 20}; + ASSERT_TRUE(store.RegisterClient(reg, now, 30s)); // control instance is live + + auto res = store.ApplyHeartbeat("wn", 1, now, {}, + {{KvEvent::Kind::ADD, live_key, TierType::DRAM, 1}, + {KvEvent::Kind::ADD, dead_key, TierType::DRAM, 1}}, + false); + // Down shard -> self-heal signal, not an exception. + EXPECT_EQ(res.status, HeartbeatResult::SEQ_GAP); + EXPECT_TRUE(store.IsClientAlive("wn")); // control committed -> node alive + EXPECT_TRUE(store.BatchExistsBlock({live_key})[0]); // live shard still got its event +} + +// A metrics sink attached to the store exports per-op latency as +// mori_umbp_store_op_latency_seconds{op,backend="redis"}. Drives the store +// directly (no gRPC/GPU) and scrapes the Prometheus endpoint over HTTP. +TEST(RedisStoreMetricsTest, EmitsStoreOpLatencyHistogram) { + redis::RespClient::Options opts; + opts.uri = RedisUri(); + opts.pool_size = 2; + redis::RespClient probe(opts); + if (!probe.Ping()) { + GTEST_SKIP() << "no RESP store reachable at " << RedisUri() << "; skipping"; + } + + const int port = 19099; + std::unique_ptr server; + try { + server = std::make_unique(port); + } catch (const std::exception& e) { + GTEST_SKIP() << "could not start metrics server on " << port << ": " << e.what(); + } + + RedisMasterMetadataStore::Config cfg; + cfg.namespace_id = UniqueNamespace(); + RedisMasterMetadataStore store(cfg); + store.SetMetricsSink(server.get()); + + const auto now = std::chrono::system_clock::now(); + ClientRegistration reg; + reg.node_id = "mn"; + reg.node_address = "a"; + reg.peer_address = "p:1"; + reg.tier_capacities[TierType::DRAM] = TierCapacity{1u << 20, 1u << 20}; + ASSERT_TRUE(store.RegisterClient(reg, now, 30s)); + ASSERT_EQ( + store.ApplyHeartbeat("mn", 1, now, {}, {{KvEvent::Kind::ADD, "mk", TierType::DRAM, 4}}, false) + .status, + HeartbeatResult::APPLIED); + store.BatchLookupBlockForRouteGet({"mk"}, {}, now, 10s); + store.BatchExistsBlock({"mk"}); + + // Scrape /metrics via curl (present in the build image). + const std::string cmd = "curl -s http://127.0.0.1:" + std::to_string(port) + "/metrics"; + std::string body; + FILE* f = popen(cmd.c_str(), "r"); + ASSERT_NE(f, nullptr); + char buf[8192]; + size_t n; + while ((n = fread(buf, 1, sizeof(buf), f)) > 0) body.append(buf, n); + pclose(f); + + EXPECT_NE(body.find("mori_umbp_store_op_latency_seconds"), std::string::npos) << body; + EXPECT_NE(body.find("backend=\"redis\""), std::string::npos); + EXPECT_NE(body.find("op=\"BatchLookupBlockForRouteGet\""), std::string::npos); + EXPECT_NE(body.find("op=\"ApplyHeartbeat\""), std::string::npos); +} + +} // namespace +} // namespace mori::umbp diff --git a/tests/cpp/umbp/distributed/test_redispp_smoke.cpp b/tests/cpp/umbp/distributed/test_redispp_smoke.cpp new file mode 100644 index 000000000..3bbb5fe6a --- /dev/null +++ b/tests/cpp/umbp/distributed/test_redispp_smoke.cpp @@ -0,0 +1,88 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. + +// redis-plus-plus build/link smoke test. +// +// This does NOT test the UMBP store — it only proves that the redis++ submodule +// compiles with this toolchain and links (static redis++ + system hiredis) +// before any store code depends on it. It is the de-risk gate for adopting +// redis++ as the Redis-backend client library (cluster / sentinel / pooling): +// if this builds and links, the later seam migration rests on a proven base. +// +// Runtime behaviour: connects to UMBP_REDIS_URI (default tcp://127.0.0.1:6379) +// and, if a cluster seed is given via UMBP_REDIS_CLUSTER_SEEDS, to that cluster; +// both skip cleanly when nothing is reachable, so BUILD_TESTS on a host without +// a store never fails. + +#include + +#include +#include + +#include + +namespace { + +std::string EnvOr(const char* key, const std::string& def) { + const char* v = std::getenv(key); + return (v != nullptr && *v != '\0') ? std::string(v) : def; +} + +// Compile check: constructing ConnectionOptions exercises the redis++ headers. +// Actual linkage of the static archive is proven by the ping tests below, which +// reference symbols defined in redis++'s .cpp files (Redis::ping / +// RedisCluster) regardless of whether they run or skip at runtime. +TEST(RedisPlusPlusSmoke, HeadersCompile) { + sw::redis::ConnectionOptions opts; + opts.host = "127.0.0.1"; + opts.port = 6379; + EXPECT_EQ(opts.port, 6379); +} + +TEST(RedisPlusPlusSmoke, PingSingleIfReachable) { + const std::string uri = EnvOr("UMBP_REDIS_URI", "tcp://127.0.0.1:6379"); + try { + sw::redis::Redis redis(uri); + const std::string pong = redis.ping(); + EXPECT_EQ(pong, "PONG"); + } catch (const sw::redis::Error& e) { + GTEST_SKIP() << "no RESP store reachable at " << uri << " (" << e.what() << "); skipping"; + } +} + +TEST(RedisPlusPlusSmoke, PingClusterIfSeedGiven) { + const char* seeds = std::getenv("UMBP_REDIS_CLUSTER_SEEDS"); + if (seeds == nullptr || *seeds == '\0') { + GTEST_SKIP() << "UMBP_REDIS_CLUSTER_SEEDS not set; skipping cluster smoke"; + } + try { + // RedisCluster fetches CLUSTER SLOTS on construction, proving cluster-path + // symbols link and a real cluster is reachable. + sw::redis::RedisCluster cluster{std::string(seeds)}; + const std::string pong = cluster.redis("smoke", false).ping(); + EXPECT_EQ(pong, "PONG"); + } catch (const sw::redis::Error& e) { + GTEST_SKIP() << "no Redis Cluster reachable at " << seeds << " (" << e.what() << "); skipping"; + } +} + +} // namespace diff --git a/tests/cpp/umbp/distributed/test_sharded_read.cpp b/tests/cpp/umbp/distributed/test_sharded_read.cpp new file mode 100644 index 000000000..98b660912 --- /dev/null +++ b/tests/cpp/umbp/distributed/test_sharded_read.cpp @@ -0,0 +1,225 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. + +// Deterministic unit tests for the block-read fan-out (redis/sharded_read.h) +// against a fake IRespClient — no live Redis. Covers the order-sensitive, +// fault-tolerance-critical parts that the live integration test can only reach +// with a running store: reply scatter back to caller order, multi-instance +// fan-out, and the tolerate-shard-failures degrade-to-miss behaviour. + +#include + +#include +#include +#include +#include + +#include "umbp/distributed/master/redis/key_schema.h" +#include "umbp/distributed/master/redis/resp_value.h" +#include "umbp/distributed/master/redis/sharded_read.h" +#include "umbp/distributed/master/redis/thread_pool.h" + +namespace mori::umbp::redis { +namespace { + +// A configurable IRespClient stand-in. Only EvalPipeline is exercised by +// RunShardedRead; the rest satisfy the interface. +class FakeRespClient : public IRespClient { + public: + // Per-group reply builder; default echoes each key as a String element, so a + // group's reply is an Array parallel to its keys (what route_get_batch shapes). + std::function&)> reply_fn; + bool throw_transport = false; // simulate a dead / unreachable instance + int eval_pipeline_calls = 0; + + RespValue Command(const std::vector&) override { return {}; } + RespValue Eval(const std::string&, const std::vector&, + const std::vector&) override { + return {}; + } + std::vector EvalPipeline(const std::string&, + const std::vector>& keys_per_call, + const std::vector&) override { + ++eval_pipeline_calls; + if (throw_transport) throw std::runtime_error("fake: transport down"); + std::vector out; + out.reserve(keys_per_call.size()); + for (const auto& keys : keys_per_call) { + out.push_back(reply_fn ? reply_fn(keys) : EchoArray(keys)); + } + return out; + } + bool Ping() override { return true; } + + static RespValue EchoArray(const std::vector& keys) { + RespValue arr; + arr.type = RespValue::Type::Array; + for (const auto& k : keys) { + RespValue e; + e.type = RespValue::Type::String; + e.str = k; + arr.elements.push_back(std::move(e)); + } + return arr; + } + static RespValue ErrorReply(const std::string& msg) { + RespValue v; + v.type = RespValue::Type::Error; + v.str = msg; + return v; + } +}; + +// Build a ShardedBatch by hand for full control over grouping (independent of the +// key hash), mirroring what GroupKeysByShard produces. +ShardedBatch MakeBatch(std::vector> keys_by_shard, + std::vector> orig_index_by_shard, + std::vector shard_of_group) { + ShardedBatch b; + b.keys_by_shard = std::move(keys_by_shard); + b.orig_index_by_shard = std::move(orig_index_by_shard); + b.shard_of_group = std::move(shard_of_group); + return b; +} + +// Collect on_key(orig_index, String reply) into a result vector for assertions. +std::vector RunAndCollect(const ShardedBatch& batch, std::size_t n_keys, + ThreadPool* pool, bool tolerate, + std::function route) { + std::vector out(n_keys); + RunShardedRead( + batch, "script", /*shared_args=*/{}, "TestOp", pool, tolerate, /*metrics=*/nullptr, + [&](std::size_t group) { return route(group); }, + [&](std::size_t orig, const RespValue& v) { out[orig] = v.str; }); + return out; +} + +// ---- GroupKeysByShard -------------------------------------------------------- + +TEST(GroupKeysByShard, BucketsPreserveEveryOriginalIndexOncePerShard) { + KeySchema schema("ns", /*num_shards=*/4); + std::vector user_keys = {"a", "b", "c", "d", "e", "f", "g", "h"}; + const ShardedBatch batch = GroupKeysByShard(schema, user_keys); + + ASSERT_EQ(batch.keys_by_shard.size(), batch.orig_index_by_shard.size()); + ASSERT_EQ(batch.keys_by_shard.size(), batch.shard_of_group.size()); + + std::set seen; + for (std::size_t g = 0; g < batch.keys_by_shard.size(); ++g) { + ASSERT_EQ(batch.keys_by_shard[g].size(), batch.orig_index_by_shard[g].size()); + for (std::size_t j = 0; j < batch.keys_by_shard[g].size(); ++j) { + const std::size_t orig = batch.orig_index_by_shard[g][j]; + // Every original index appears exactly once across all groups. + EXPECT_TRUE(seen.insert(orig).second) << "duplicate orig index " << orig; + // The composed block key matches KeySchema, and its shard matches the group. + EXPECT_EQ(batch.keys_by_shard[g][j], schema.Block(user_keys[orig])); + EXPECT_EQ(schema.ShardOf(user_keys[orig]), batch.shard_of_group[g]); + } + } + EXPECT_EQ(seen.size(), user_keys.size()); +} + +TEST(GroupKeysByShard, SingleShardYieldsOneGroup) { + KeySchema schema("ns", /*num_shards=*/1); + const ShardedBatch batch = GroupKeysByShard(schema, {"a", "b", "c"}); + ASSERT_EQ(batch.keys_by_shard.size(), 1u); + EXPECT_EQ(batch.keys_by_shard[0].size(), 3u); + EXPECT_EQ(batch.shard_of_group[0], 0u); +} + +// ---- RunShardedRead: scatter ------------------------------------------------- + +TEST(RunShardedRead, SingleClientScattersRepliesToOriginalOrder) { + // Two groups, interleaved original indices, one client (inline path, no pool). + const ShardedBatch batch = MakeBatch(/*keys=*/{{"blk0", "blk2"}, {"blk1", "blk3"}}, + /*orig=*/{{0, 2}, {1, 3}}, /*shard=*/{0, 1}); + FakeRespClient client; // echoes each key + const auto out = RunAndCollect(batch, 4, /*pool=*/nullptr, /*tolerate=*/false, + [&](std::size_t) { return &client; }); + EXPECT_EQ(out, (std::vector{"blk0", "blk1", "blk2", "blk3"})); + EXPECT_EQ(client.eval_pipeline_calls, 1); // one client => one pipelined call +} + +TEST(RunShardedRead, MultiClientFanOutPreservesOrder) { + // Groups routed to two different clients; a real pool drives the fan-out. + const ShardedBatch batch = + MakeBatch({{"blk0", "blk3"}, {"blk1"}, {"blk2"}}, {{0, 3}, {1}, {2}}, {0, 1, 2}); + FakeRespClient a, b; + ThreadPool pool(4); + // group 0 -> a, groups 1,2 -> b. + const auto out = RunAndCollect(batch, 4, &pool, /*tolerate=*/false, + [&](std::size_t g) -> IRespClient* { return g == 0 ? &a : &b; }); + EXPECT_EQ(out, (std::vector{"blk0", "blk1", "blk2", "blk3"})); + EXPECT_EQ(a.eval_pipeline_calls, 1); + EXPECT_EQ(b.eval_pipeline_calls, 1); // b's two groups share one pipeline +} + +// ---- RunShardedRead: fault tolerance ---------------------------------------- + +TEST(RunShardedRead, TolerateDegradesDeadInstanceToMiss) { + const ShardedBatch batch = MakeBatch({{"blk0"}, {"blk1"}}, {{0}, {1}}, {0, 1}); + FakeRespClient a; // healthy + FakeRespClient b; // dead + b.throw_transport = true; + ThreadPool pool(2); + const auto out = RunAndCollect(batch, 2, &pool, /*tolerate=*/true, + [&](std::size_t g) -> IRespClient* { return g == 0 ? &a : &b; }); + EXPECT_EQ(out[0], "blk0"); // healthy shard delivered + EXPECT_EQ(out[1], ""); // dead shard left at default (miss), no throw +} + +TEST(RunShardedRead, NoToleratePropagatesDeadInstance) { + const ShardedBatch batch = MakeBatch({{"blk0"}}, {{0}}, {0}); + FakeRespClient dead; + dead.throw_transport = true; + EXPECT_THROW(RunAndCollect(batch, 1, /*pool=*/nullptr, /*tolerate=*/false, + [&](std::size_t) { return &dead; }), + std::runtime_error); +} + +TEST(RunShardedRead, TolerateSkipsScriptErrorGroupButDeliversRest) { + const ShardedBatch batch = MakeBatch({{"blk0"}, {"blk1"}}, {{0}, {1}}, {0, 1}); + FakeRespClient client; + // Group whose key is "blk1" returns a server-side error reply; others echo. + client.reply_fn = [](const std::vector& keys) { + if (!keys.empty() && keys[0] == "blk1") return FakeRespClient::ErrorReply("boom"); + return FakeRespClient::EchoArray(keys); + }; + const auto out = RunAndCollect(batch, 2, /*pool=*/nullptr, /*tolerate=*/true, + [&](std::size_t) { return &client; }); + EXPECT_EQ(out[0], "blk0"); // healthy group delivered + EXPECT_EQ(out[1], ""); // errored group skipped (miss) +} + +TEST(RunShardedRead, NoToleratePropagatesScriptError) { + const ShardedBatch batch = MakeBatch({{"blk0"}}, {{0}}, {0}); + FakeRespClient client; + client.reply_fn = [](const std::vector&) { + return FakeRespClient::ErrorReply("boom"); + }; + EXPECT_THROW(RunAndCollect(batch, 1, /*pool=*/nullptr, /*tolerate=*/false, + [&](std::size_t) { return &client; }), + std::runtime_error); +} + +} // namespace +} // namespace mori::umbp::redis