diff --git a/mooncake-store/benchmarks/ha_rebuild_bench.md b/mooncake-store/benchmarks/ha_rebuild_bench.md new file mode 100644 index 00000000..1cce6b8f --- /dev/null +++ b/mooncake-store/benchmarks/ha_rebuild_bench.md @@ -0,0 +1,155 @@ +# HA Client-Driven Metadata Rebuild — Benchmark Report + +Corresponding change: `feat: add client metadata rebuild for master high availability` +(branch `feat/client-metadata-rebuild-ha`). This report presents measured data for this HA rebuild feature across three dimensions — **scale, rebuild latency, and recovery-window read performance** — for code reviewers to evaluate. + +> One-line takeaway: under a single-client / single-replica, value=100B configuration, after the master is +> killed with `kill -9` and restarted on the same port, the client automatically resends its local +> key→location table, the master completes metadata rebuild and restores full readability, achieving +> **zero recomputation**. Rebuild latency is about 1.5–2.1 seconds within 1M keys (dominated by +> heartbeat-detection latency), and about 12 seconds at 5M keys (resend + rebuild workload starts to +> dominate). After recovery completes, read latency matches the pre-failure level (no long-term degradation). + +--- + +## 1. Test Methodology (how it was measured, why it is trustworthy) + +### 1.1 Test Setup + +This uses a **separate-process** live test (not an in-process mock), which most closely resembles a real deployment: + +- A real `mooncake_master` process (non-HA mode, `-enable_ha=false`). +- A standalone benchmark client (`ha_scale_bench_main`) that mounts a segment, loads N + keys, records a baseline, then enters polling; meanwhile an external script `kill -9`s the master and restarts it on the same port. +- After the client observes "reads fail → all readable again", it emits the rebuild latency and recovery-window statistics. + +Timeline: `start master → load N keys → print READY_FOR_KILL → kill -9 master → +wait 4s (let the client detect the failure) → restart master on the same port → client reconnects and resends its local table → +master rebuilds → probes turn fully green → emit JSON`. + +### 1.2 Key Measurement Design (why it is measured this way) + +- **Keys are loaded with BatchPut** (2000 per batch): individual Puts would be dominated by per-RPC overhead and would not accurately measure fill throughput. +- **Rebuild completion is judged with a "sampling probe"**: **uniformly sample 1000 keys** over `[0, N)`; rebuild is deemed complete + once all of them are readable. Getting all N keys every round would itself take several seconds at the million scale and would pollute the latency measurement. + Uniform sampling probes every region of the key space, balancing "detecting partial rebuild" with "low overhead". +- **Timing uses `steady_clock`, with a 20ms recovery polling interval**: rebuild-latency resolution is about ±20ms. +- Each scale runs on a **fresh master and a fresh port** (clean state, no cross-interference). + +### 1.3 Definitions of the Three Metrics + +| Metric | Field | Definition | +|---|---|---| +| **Rebuild latency** | `rebuild_ms` | Wall-clock time from "probe first sees a failure" (detecting the master is down) to "probe first turns fully green" (rebuild complete) | +| **Scale** | `nkeys` sweep | The same procedure repeated at 1k / 10k / 100k / 1M / 5M keys | +| **Recovery-window read latency** | `outage_get_lat_ms` | Average latency of **successful** Gets within the recovery window, compared against `steady_get_lat_ms` (pre-failure steady state) | +| (auxiliary) | `min_avail_pct` | The minimum probe availability within the recovery window (0 means everything was unreadable at some point) | + +--- + +## 2. Test Environment + +| Item | Value | +|---|---| +| Machine memory | 3.0 TB (~1.6 TB available) | +| CPU | 640 cores | +| Runtime | Container `shenshuwei-xllm` (image `xllm-dev-a3-arm-cann9`), openEuler / ARM64, glibc 2.38 | +| Transport protocol | TCP (`-protocol=tcp`) | +| Master mode | Non-HA (`-enable_ha=false -enable_metric_reporting=false`) | +| Replica count | replica_num=1 (single replica) | +| Value size | 100 bytes/key | +| Segment size | Auto-estimated at ~1200B/object (9298 MB actual for the 5M scale) | +| client_ttl | 10s (master default; affects failure-detection latency, see §4) | + +> Note: the test runs inside the container because the binary is compiled in that container (glibc 2.38); the host +> (glibc 2.34) lacks matching runtime libraries and cannot run it directly. + +--- + +## 3. Results + +### 3.1 Summary Table (single run, not averaged over multiple runs) + +| nkeys | Fill rate (keys/s) | **Rebuild latency (ms)** | Steady-state read latency (ms) | Recovery-window read latency (ms) | Recovery-window min availability | +|--:|--:|--:|--:|--:|--:| +| 1,000 | 230,585 | 1,546 | 0.043 | 0.040 | 0% | +| 10,000 | 246,441 | 1,559 | 0.041 | 0.043 | 0% | +| 100,000 | 232,930 | 1,791 | 0.042 | 0.046 | 0% | +| 1,000,000 | 206,332 | 2,091 | 0.045 | 0.042 | 0% | +| 5,000,000 | 196,708 | **12,068** | 0.038 | 0.042 | 0% | + +Every scale ultimately reports `RESULT=PASS`, the sampling probe recovers 1000/1000, and data is rebuilt with zero recomputation. + +### 3.2 Metrics 1 & 2: How Rebuild Latency Scales + +- **1k → 1M**: rebuild latency goes from 1.5s → 2.1s, a very gentle increase. This range is dominated **not** by rebuild workload, + but by **failure-detection latency** (from the master dying to the client's heartbeat deciding to reconnect takes on the order of seconds, and is nearly independent of key count). +- **1M → 5M**: latency jumps from 2.1s to **12s**, clearly superlinear. Here the resend workload (5M entries at 256 per batch + = about 20k `RebuildMetadata` RPCs) plus the master's per-key rebuild — **the workload itself** — + starts to dominate and overtakes the fixed detection latency. +- Fill throughput slowly drops from ~230k/s to ~200k/s as scale grows (allocator pressure, larger single-segment capacity). + +### 3.3 Metric 3: Recovery-Window Read Latency + +- Steady-state read latency is stable at **~0.04ms** (single-threaded sequential Get, local TCP). +- The latency of successful Gets during recovery is likewise **~0.04ms**, **essentially indistinguishable from steady state** — i.e., once rebuild completes, + readable keys are just as fast to read as before the failure, with no long-term degradation. +- `min_avail_pct=0` indicates that between the master being killed and rebuild completing, there is a window where **all keys are unreadable** + (as expected — the new master is empty and has not yet received the metadata resent by the client). During rebuild the transition is + "all-none → all-present", not a gradual recovery. + +--- + +## 4. Important Limitations and Caveats (required reading for reviewers) + +1. **`rebuild_ms` includes failure-detection latency.** It measures the end-to-end time from "client detects failure → fully green again", + which includes the seconds-level latency for the heartbeat to conclude the master is dead. Therefore the + 1.5–2s seen within 1M keys is **not pure rebuild time**; it is largely the floor set by detection latency. To measure "pure resend + rebuild" latency separately, + a latency-histogram metric on `RebuildMetadata` would need to be added on the master side (there are currently only requests/failures + counters, no latency statistics) — recommended as a separate follow-up change. + +2. **Single-client / single-replica scope.** In this test one client Puts data to its own segment (via the + `RecordLocalReplica` local-bookkeeping path), with replica_num=1. **Not covered**: the cross-client + notify bookkeeping path and multi-replica merge rebuild. The **correctness** of these paths is already covered by unit tests + (tests 4 and 6 in `client_metadata_rebuild_test.cpp`), but their **performance** was not measured at this scale. + +3. **Single run, not averaged over multiple runs.** Each scale is run only once, so there is jitter from heartbeat timing (in an earlier run, + the 100k scale once measured an anomalously low 143ms because it hit a different heartbeat moment). For a formal benchmark, repeating each + scale 3–5 times and taking the median is recommended. + +4. **Recovery-window latency is the latency of single-threaded sequential Gets**, not high-concurrency throughput. Latency/failure rates under concurrent workloads + require a multi-threaded client to measure. + +5. **Per-object memory footprint (a byproduct observation):** each key measured about **~1073 bytes** in the segment + (only 100B is value; the rest is object metadata + allocator minimum-unit/alignment/fragmentation). The segment + capacity for scale tests must be estimated accordingly, otherwise `NO_AVAILABLE_HANDLE` occurs (segment full during the data-loading phase). + +--- + +## 5. Reproduction + +Test code (included with this PR): +- `mooncake-store/tests/ha_scale_bench_main.cpp` — benchmark client +- `mooncake-store/tests/ha_scale_bench.sh` — driver script (start/kill/restart master, collect JSON) + +Run inside the build container: + +```bash +# Build +cd build && make ha_scale_bench_main -j32 + +# Run the full sweep (1k / 10k / 100k / 1M / 5M) +cd /path/to/Mooncake +OUT_DIR=/tmp/ha_scale bash mooncake-store/tests/ha_scale_bench.sh \ + 1000 10000 100000 1000000 5000000 + +# Or a custom single scale (tunable value size / segment size) +VSIZE=100 SEG_MB=10240 bash mooncake-store/tests/ha_scale_bench.sh 5000000 +``` + +Results are written as TSV (`$OUT_DIR/results.tsv`) plus per-scale master/client logs. + +--- + +*Data collected: 2026-07-27. This report and the test scripts were generated with AI assistance; all data comes from real runs. +Before submission, please have a human review every changed line and independently reproduce the results.* diff --git a/mooncake-store/include/allocator.h b/mooncake-store/include/allocator.h index de74e7ff..aeeb8b50 100644 --- a/mooncake-store/include/allocator.h +++ b/mooncake-store/include/allocator.h @@ -213,6 +213,13 @@ class OffsetBufferAllocator std::unique_ptr allocate(size_t size) override; + // HA rebuild: allocate `size` to obtain a legit ownership handle (correct + // accounting + safe deallocation), but point the buffer's data address at + // `real_addr` (the client's actual address where data physically lives). + // The self-chosen allocate address is discarded. + std::unique_ptr AllocateForRebuild(size_t size, + void* real_addr); + void deallocate(AllocatedBuffer* handle) override; size_t capacity() const override { return total_size_; } diff --git a/mooncake-store/include/client_metric.h b/mooncake-store/include/client_metric.h index 1ba40614..82e8d6fc 100644 --- a/mooncake-store/include/client_metric.h +++ b/mooncake-store/include/client_metric.h @@ -663,6 +663,8 @@ struct ClientMetric { MasterClientMetric master_client_metric; TransferOperationMetric transfer_operation_metric; SsdMetric ssd_metric; + ylt::metric::counter_t rebuild_failed_batches; + ylt::metric::counter_t rebuild_retries; /** * @brief Creates a ClientMetric instance based on environment variables diff --git a/mooncake-store/include/client_service.h b/mooncake-store/include/client_service.h index 621b4edd..c081f347 100644 --- a/mooncake-store/include/client_service.h +++ b/mooncake-store/include/client_service.h @@ -23,6 +23,7 @@ #include "transfer_task.h" #include "types.h" #include "replica.h" +#include "rebuild_types.h" #include "master_metric_manager.h" #include "count_min_sketch.h" #include "local_hot_cache.h" @@ -67,6 +68,29 @@ class Client { const UUID& getClientId() const { return client_id_; } const std::string& tenant_id() const { return master_client_.tenant_id(); } + // --- test-only helpers for the notify-reliability backstop --- + // Number of endpoints with parked (failed, awaiting-retry) notifies. + size_t PendingNotifyBucketCountForTest() const { + std::lock_guard lk(pending_notifies_mutex_); + return pending_notifies_.size(); + } + // Park a notify aimed at `ep` carrying `key` (used to simulate a dropped + // send, then verify the background loop re-delivers it). + void ParkNotifyForTest(const std::string& ep, const std::string& key, + const Replica::Descriptor& replica, uint64_t size, + ObjectDataType data_type, + const std::string& group_id, + const std::string& tenant_id) { + KeyReplicaEntry e; + e.key = key; + e.tenant_id = tenant_id; + e.size = size; + e.data_type = data_type; + e.group_id = group_id; + e.replicas = {replica}; + ParkPendingNotify(ep, {std::move(e)}); + } + /** * @brief Creates and initializes a new Client instance * @param local_hostname Local host address (IP:Port) @@ -778,6 +802,50 @@ class Client { std::unordered_map>& slices); ReplicateConfig AttachHostId(const ReplicateConfig& config) const; + // === HA rebuild: client-side helpers (impl in client_service.cpp) === + // Record a replica physically located in this client's own segment. Called + // both when this client Put()s onto its own segment and when an UPSERT + // notify arrives. Internally address-overwrites the stale key at the same + // address. + void RecordLocalReplica(const std::string& key, + const Replica::Descriptor& replica, uint64_t size, + ObjectDataType data_type, + const std::string& group_id, + const std::string& tenant_id); + // Evict the stale key occupying `addr` (it was just reused). Caller must + // hold local_replica_table_mutex_. + void EraseByAddressLocked(uint64_t addr); + // Is `ep` one of THIS client's mounted segments' te_endpoint? (Never + // compare against local_hostname_ -- a segment's te_endpoint = + // getLocalIpAndPort().) + bool IsMyEndpoint(const std::string& ep); + // Tell the segment owner at `ep` that we stored `key` there (full + // metadata). + void NotifyOwnerUpsert(const std::string& ep, const std::string& key, + const Replica::Descriptor& replica, uint64_t size, + ObjectDataType data_type, + const std::string& group_id, + const std::string& tenant_id); + // Batched notify: pack multiple keys landing on the same endpoint into one + // notify (BatchPut high-throughput optimization). + void NotifyOwnerUpsertBatch( + const std::unordered_map>& + by_ep); + // On reconnect, resend the whole local table to the (empty) new master. + tl::expected ResendLocalReplicaTable( + ViewVersionId view_version); + // Background loop: poll getNotifies() and apply UPSERT entries. + void RebuildNotifyLoop(); + // Send one UPSERT notify carrying `entries` to endpoint `ep`. Returns true + // on success. Shared by NotifyOwnerUpsert and the pending re-send path. + bool SendUpsertNotify(const std::string& ep, + const std::vector& entries); + // Park a failed notify for later re-send (reliability backstop). + void ParkPendingNotify(const std::string& ep, + const std::vector& entries); + // Re-send all parked notifies; drop the ones that now succeed. + void FlushPendingNotifies(); + // Client identification const UUID client_id_; @@ -793,6 +861,31 @@ class Client { mutable std::mutex mounted_segments_mutex_; std::unordered_map> mounted_segments_; + // === HA rebuild: local replica table === + // Maps a key physically stored in THIS client's segment -> its replica + // location + rebuild metadata. Filled two ways: (1) this client Put()s and + // a replica lands on its own segment; (2) an UPSERT notify arrives from + // another client. Value is LocalReplicaMeta (single replica, see + // rebuild_types.h): a key's multiple replicas are forced onto different + // segments, so from one client's view a key has at most one replica in its + // own segment. + mutable std::mutex local_replica_table_mutex_; + std::unordered_map local_replica_table_; + // Address reverse index: buffer_address_ -> key, for the same client's + // segments. Core of lazy-delete: when an address is reused, locate and + // evict the stale key entry occupying it (see RecordLocalReplica). Same + // mutex as local_replica_table_. + std::unordered_map addr_index_; + std::atomic rebuild_notify_thread_running_{false}; + std::thread rebuild_notify_thread_; // polls getNotifies() + // Reliability backstop: UPSERT notifies whose send failed (peer flapping / + // not yet up) are parked here keyed by target endpoint, and re-sent by + // RebuildNotifyLoop each tick until they succeed. Guards against silent + // multi-replica loss when a notify is dropped. + mutable std::mutex pending_notifies_mutex_; + std::unordered_map> + pending_notifies_; + // Segments in graceful unmount: readable by remote peers, not allocatable // locally. TE MR remains registered until master confirms removal. std::unordered_map> @@ -843,6 +936,7 @@ class Client { std::atomic last_ping_success_{false}; std::atomic segment_desc_publish_pending_{false}; std::atomic rpc_meta_publish_pending_{false}; + std::atomic rebuild_retry_pending_{false}; ErrorCode SwitchLeader(const ha::MasterView& target_view); void LeaderMonitorThreadMain(); void StorageHeartbeatThreadMain(); diff --git a/mooncake-store/include/master_client.h b/mooncake-store/include/master_client.h index 7b541230..5bc34302 100644 --- a/mooncake-store/include/master_client.h +++ b/mooncake-store/include/master_client.h @@ -17,6 +17,7 @@ #include "segment.h" #include "types.h" #include "rpc_types.h" +#include "rebuild_types.h" #include "master_metric_manager.h" #include "task_manager.h" @@ -348,6 +349,20 @@ class MasterClient { [[nodiscard]] tl::expected ReMountSegment( const std::vector& segments); + /** + * @brief HA rebuild: resend object-level metadata (key -> replica location) + * to the (empty) new master after a restart, so it can rebuild metadata. + */ + [[nodiscard]] tl::expected RebuildMetadata( + std::vector&& entries, ViewVersionId view_version); + + /** + * @brief HA rebuild: 通知(空的新)master "本client已把所有metadata重发完", + * master 收齐所有已报到 client 的此信号后开服务。 + */ + [[nodiscard]] tl::expected SignalRebuildComplete( + ViewVersionId view_version); + /** * @brief Re-mount NoF ssd segments, invoked when the client is the first * time to connect to the master or the client Ping TTL is expired and need diff --git a/mooncake-store/include/master_config.h b/mooncake-store/include/master_config.h index fd893953..e5cd419d 100644 --- a/mooncake-store/include/master_config.h +++ b/mooncake-store/include/master_config.h @@ -425,6 +425,7 @@ class WrappedMasterServiceConfig { double nof_eviction_high_watermark_ratio = DEFAULT_NOF_EVICTION_HIGH_WATERMARK_RATIO; ViewVersionId view_version = 0; + bool initially_serving = true; int64_t client_live_ttl_sec = DEFAULT_CLIENT_LIVE_TTL_SEC; int64_t nof_heartbeat_interval_sec = DEFAULT_NOF_HEARTBEAT_INTERVAL_SEC; uint32_t nof_heartbeat_probe_timeout_ms = diff --git a/mooncake-store/include/master_metric_manager.h b/mooncake-store/include/master_metric_manager.h index 293c6206..930d22cb 100644 --- a/mooncake-store/include/master_metric_manager.h +++ b/mooncake-store/include/master_metric_manager.h @@ -166,6 +166,15 @@ class MasterMetricManager { void inc_unmount_nof_segment_failures(int64_t val = 1); void inc_remount_segment_requests(int64_t val = 1); void inc_remount_segment_failures(int64_t val = 1); + void inc_rebuild_metadata_requests(int64_t val = 1); + void inc_rebuild_metadata_failures(int64_t val = 1); + void set_rebuild_state(int64_t state); + void set_rebuild_expected_clients(int64_t clients); + void set_rebuild_completed_clients(int64_t clients); + void set_rebuild_missing_clients(int64_t clients); + void set_rebuild_duration_ms(int64_t duration_ms); + void inc_rebuild_force_open(int64_t val = 1); + void inc_rebuild_stale_epoch_requests(int64_t val = 1); void inc_remount_nof_segment_requests(int64_t val = 1); void inc_remount_nof_segment_failures(int64_t val = 1); void inc_ping_requests(int64_t val = 1); @@ -590,6 +599,15 @@ class MasterMetricManager { ylt::metric::counter_t unmount_segment_failures_; ylt::metric::counter_t remount_segment_requests_; ylt::metric::counter_t remount_segment_failures_; + ylt::metric::counter_t rebuild_metadata_requests_; + ylt::metric::counter_t rebuild_metadata_failures_; + ylt::metric::gauge_t rebuild_state_; + ylt::metric::gauge_t rebuild_expected_clients_; + ylt::metric::gauge_t rebuild_completed_clients_; + ylt::metric::gauge_t rebuild_missing_clients_; + ylt::metric::gauge_t rebuild_duration_ms_; + ylt::metric::counter_t rebuild_force_open_; + ylt::metric::counter_t rebuild_stale_epoch_requests_; ylt::metric::counter_t mount_nof_segment_requests_; ylt::metric::counter_t mount_nof_segment_failures_; ylt::metric::counter_t unmount_nof_segment_requests_; diff --git a/mooncake-store/include/master_service.h b/mooncake-store/include/master_service.h index ac73ecfa..56099bdf 100644 --- a/mooncake-store/include/master_service.h +++ b/mooncake-store/include/master_service.h @@ -33,6 +33,7 @@ #include "master_config.h" #include "rpc_types.h" #include "replica.h" +#include "rebuild_types.h" #include "ha/ha_types.h" #include "ha/snapshot/object/snapshot_object_store.h" #include "task_manager.h" @@ -165,6 +166,19 @@ class MasterService { auto ReMountSegment(const std::vector& segments, const UUID& client_id) -> tl::expected; + std::unordered_set> getAliveClientsSnapshot() const; + + /** + * @brief HA rebuild: accept object-level metadata (key -> replica location) + * resent by a client after the master restarted empty, and rebuild it into + * metadata_shards_. For an existing key, MERGE the incoming replica(s) + * (multi-replica redundancy recovery) rather than skipping. Idempotent per + * (endpoint,address). + */ + auto RebuildMetadata(const std::vector& entries, + const UUID& client_id) + -> tl::expected; + /** * @brief Re-mount NoF SSD segments, invoked when the client is the first * time to connect to the master or the client Ping TTL is expired and need @@ -810,6 +824,15 @@ class MasterService { private: std::unique_ptr CreateSnapshotCatalogStore(); + // === HA rebuild helpers === + // Convert a serializable Replica::Descriptor back into a holding Replica. + // The hard part is MEMORY type: it needs the owning segment's allocator, + // looked up by the descriptor's transport_endpoint_. Returns nullopt on + // failure (segment not mounted / not OK / unknown type). + std::optional DescriptorToReplica(const Replica::Descriptor& desc); + // ReplicaAlreadyPresent is declared after ObjectMetadata is defined (it + // takes const ObjectMetadata&, a private nested type). See below. + // Restore master state void RestoreState(); void ResetStateAfterFailedRestoreAttempt(); @@ -843,7 +866,6 @@ class MasterService { // Helper to get a snapshot of alive clients (under client_mutex_ shared // lock) - std::unordered_set> getAliveClientsSnapshot() const; void UpdateClientHostId(const UUID& client_id, const std::string& host_id); std::string GetClientHostId(const UUID& client_id) const; @@ -1401,6 +1423,12 @@ class MasterService { bool IsTenantRegistered(const std::string& tenant_id) const; bool TenantHasObjects(const std::string& tenant_id) const; + // HA rebuild: is a replica with the same (endpoint,address) already present + // in meta? Declared here (after ObjectMetadata is defined) because it takes + // const ObjectMetadata&. + bool ReplicaAlreadyPresent(const ObjectMetadata& meta, + const Replica& r) const; + static std::string MakeTenantScopedKey(const std::string& tenant_id, const std::string& key) { const auto normalized_tenant = NormalizeTenantId(tenant_id); diff --git a/mooncake-store/include/rebuild_retry.h b/mooncake-store/include/rebuild_retry.h new file mode 100644 index 00000000..cde0dfe9 --- /dev/null +++ b/mooncake-store/include/rebuild_retry.h @@ -0,0 +1,23 @@ +#pragma once + +#include + +#include + +#include "types.h" + +namespace mooncake { + +struct RebuildRecoveryOps { + std::function()> remount; + std::function()> publish_segment_descriptor; + std::function()> publish_rpc_metadata; + std::function()> resend_metadata; + std::function()> signal_complete; +}; + +tl::expected RunRebuildRecovery( + const RebuildRecoveryOps& ops, int max_attempts, + const std::function& wait_before_retry); + +} // namespace mooncake diff --git a/mooncake-store/include/rebuild_types.h b/mooncake-store/include/rebuild_types.h new file mode 100644 index 00000000..daee2361 --- /dev/null +++ b/mooncake-store/include/rebuild_types.h @@ -0,0 +1,54 @@ +// HA metadata rebuild: shared types for client<->master metadata rebuild. +#pragma once + +#include +#include +#include + +#include "replica.h" +#include "types.h" + +namespace mooncake { + +// Client-side local table value: the single replica physically located in this +// client's own segment, plus the metadata master needs to rebuild the object. +// Single replica (not a vector): different replicas of one key are forced onto +// different segments, so from one client's view a key has at most one replica +// in its own segment. +struct LocalReplicaMeta { + Replica::Descriptor replica; + uint64_t size{0}; + ObjectDataType data_type{ObjectDataType::UNKNOWN}; + std::string group_id; + std::string tenant_id{"default"}; +}; + +// One key's rebuild entry: key + its replica location(s). Descriptor is already +// serializable (YLT_REFL at replica.h:477), so it travels over RPC/notify +// as-is. +struct KeyReplicaEntry { + std::string key; + std::string tenant_id{"default"}; + uint64_t size{0}; + ObjectDataType data_type{ObjectDataType::UNKNOWN}; + std::string group_id; + std::vector replicas; + KeyReplicaEntry() = default; +}; +YLT_REFL(KeyReplicaEntry, key, tenant_id, size, data_type, group_id, replicas); + +// Notify payload: one notify may carry multiple keys (multi-key compatible); +// with single-key it just holds one entry. Under lazy-delete only UPSERT is +// used (a reuse write tells the owner to overwrite the stale key at that addr); +// REMOVE is reserved but never sent. +enum class RebuildNotifyOp : uint8_t { UPSERT = 0, REMOVE = 1 /*reserved*/ }; + +struct RebuildNotify { + std::string sender_client_id; + RebuildNotifyOp op{RebuildNotifyOp::UPSERT}; + std::vector entries; + RebuildNotify() = default; +}; +YLT_REFL(RebuildNotify, sender_client_id, op, entries); + +} // namespace mooncake diff --git a/mooncake-store/include/rpc_service.h b/mooncake-store/include/rpc_service.h index 8e2cfa48..8a18682a 100644 --- a/mooncake-store/include/rpc_service.h +++ b/mooncake-store/include/rpc_service.h @@ -2,6 +2,10 @@ #include +#include +#include +#include +#include #include #include #include @@ -17,6 +21,12 @@ namespace mooncake { +enum class StoreServingState : uint8_t { + REBUILDING = 0, + SERVING = 1, + DEGRADED = 2, +}; + // Forward declaration class HttpMetadataServer; class WrappedMasterService { @@ -144,7 +154,7 @@ class WrappedMasterService { const std::string& str, bool force = false, const std::string& tenant_id = "default"); - long RemoveAll(bool force = false, + tl::expected RemoveAll(bool force = false, const std::string& tenant_id = "default"); std::vector> BatchRemove( @@ -160,6 +170,10 @@ class WrappedMasterService { tl::expected ReMountSegment( const std::vector& segments, const UUID& client_id); + tl::expected RebuildMetadata( + const std::vector& entries, const UUID& client_id, + ViewVersionId view_version); + tl::expected ReMountNoFSegment( const std::vector& segments, const UUID& client_id); @@ -303,8 +317,44 @@ class WrappedMasterService { bool KvEventsEnabled() const; KvEventPublisher::Stats GetKvEventStats() const; + void SetServing(bool on); + bool IsServing() const { + return serving_state_.load(std::memory_order_acquire) != + StoreServingState::REBUILDING; + } + StoreServingState GetServingState() const { + return serving_state_.load(std::memory_order_acquire); + } + ViewVersionId GetViewVersion() const { return view_version_; } + + using ClientSet = + std::unordered_set>; + ClientSet GetAliveClientsSnapshot() const; + void LockRebuildExpectedClients(ClientSet expected_clients); + tl::expected SignalRebuildComplete( + const UUID& client_id, ViewVersionId view_version); + void ForceServingAfterTimeout(); + tl::expected SignalRebuildCompleteRpc( + const UUID& client_id, ViewVersionId view_version); + private: + bool IsCurrentView(ViewVersionId view_version) const { + return view_version == view_version_; + } + void MaybeFinishRebuildLocked(); + void TransitionToLocked(StoreServingState state, const char* reason); + std::string MissingClientsLocked() const; + MasterService master_service_; + const ViewVersionId view_version_; + std::atomic serving_state_; + const std::chrono::steady_clock::time_point rebuild_started_at_; + + std::mutex rebuild_mu_; + bool rebuild_window_locked_{false}; + ClientSet rebuild_expected_clients_; + ClientSet rebuild_done_before_lock_; + ClientSet rebuild_done_clients_; }; void RegisterRpcService(coro_rpc::coro_rpc_server& server, diff --git a/mooncake-store/include/segment.h b/mooncake-store/include/segment.h index 9fc5d2e8..1403be86 100644 --- a/mooncake-store/include/segment.h +++ b/mooncake-store/include/segment.h @@ -239,6 +239,14 @@ class ScopedSegmentAccess { */ void UnmountLocalDiskSegment(const UUID& client_id); + // HA rebuild: find the owning segment's buffer allocator for a replica the + // client reports by (te_endpoint, buffer_address). endpoint alone is + // ambiguous (same host -> shared endpoint, 1:N), so disambiguate by + // requiring buffer_address in [segment.base, base+size). Returns nullptr if + // no OK-status segment matches. + std::shared_ptr FindAllocatorByEndpointAndAddr( + const std::string& te_endpoint, uintptr_t buffer_address) const; + private: SegmentManager* segment_manager_; std::unique_lock lock_; diff --git a/mooncake-store/src/CMakeLists.txt b/mooncake-store/src/CMakeLists.txt index 6b86a248..59268260 100644 --- a/mooncake-store/src/CMakeLists.txt +++ b/mooncake-store/src/CMakeLists.txt @@ -21,6 +21,7 @@ set(MOONCAKE_STORE_SOURCES tenant_quota.cpp tenant_quota_policy_store.cpp rpc_service.cpp + rebuild_retry.cpp master_admin_service.cpp offset_allocator.cpp posix_file.cpp diff --git a/mooncake-store/src/allocator.cpp b/mooncake-store/src/allocator.cpp index 23311b83..8742a516 100644 --- a/mooncake-store/src/allocator.cpp +++ b/mooncake-store/src/allocator.cpp @@ -279,6 +279,51 @@ std::unique_ptr OffsetBufferAllocator::allocate(size_t size) { return allocated_buffer; } +std::unique_ptr OffsetBufferAllocator::AllocateForRebuild( + size_t size, void* real_addr) { + if (!offset_allocator_) { + LOG(ERROR) << "allocator_status=not_initialized"; + return nullptr; + } + std::unique_ptr allocated_buffer = nullptr; + try { + // Allocate to obtain a legit ownership handle (correct accounting + + // safe RAII deallocation). We DISCARD the allocator's self-chosen + // address and instead point the buffer at `real_addr` (the client's + // actual address). + auto allocation_handle = offset_allocator_->allocate(size); + if (!allocation_handle) { + VLOG(1) << "rebuild_allocation_failed size=" << size + << " segment=" << segment_name_ + << " current_size=" << cur_size_; + return nullptr; + } + // Data address = client's real address; ownership handle = the legit + // one just allocated. deallocate() only touches the handle + size, + // never the data address, so this is safe. + allocated_buffer = std::make_unique( + shared_from_this(), real_addr, size, std::move(allocation_handle)); + VLOG(1) << "rebuild_allocation_succeeded size=" << size + << " segment=" << segment_name_ + << " real_address=" << real_addr; + } catch (const std::exception& e) { + LOG(ERROR) << "rebuild_allocation_exception error=" << e.what(); + return nullptr; + } catch (...) { + LOG(ERROR) << "rebuild_allocation_unknown_exception"; + return nullptr; + } + cur_size_.fetch_add(size); + if (replica_type_ == ReplicaType::MEMORY) { + MasterMetricManager::instance().inc_allocated_mem_size(segment_name_, + size); + } else if (replica_type_ == ReplicaType::NOF_SSD) { + MasterMetricManager::instance().inc_allocated_nof_size(segment_name_, + size); + } + return allocated_buffer; +} + void OffsetBufferAllocator::deallocate(AllocatedBuffer* handle) { try { // The OffsetAllocator handles deallocation automatically through RAII diff --git a/mooncake-store/src/client_metric.cpp b/mooncake-store/src/client_metric.cpp index 306d479d..c329fefb 100644 --- a/mooncake-store/src/client_metric.cpp +++ b/mooncake-store/src/client_metric.cpp @@ -83,6 +83,10 @@ ClientMetric::ClientMetric(uint64_t interval_seconds, master_client_metric(labels), transfer_operation_metric(labels), ssd_metric(labels), + rebuild_failed_batches("client_rebuild_failed_batches_total", + "Failed metadata rebuild batches", labels), + rebuild_retries("client_rebuild_retries_total", + "HA metadata rebuild retries", labels), should_stop_metrics_thread_(false), metrics_interval_seconds_(interval_seconds), bandwidth_reporting_enabled_(bandwidth_reporting_enabled), @@ -128,6 +132,8 @@ void ClientMetric::serialize(std::string& str) { } transfer_operation_metric.serialize(str); ssd_metric.serialize(str); + rebuild_failed_batches.serialize(str); + rebuild_retries.serialize(str); } std::string ClientMetric::summary_metrics() { diff --git a/mooncake-store/src/client_service.cpp b/mooncake-store/src/client_service.cpp index c4b74fd0..d1638c52 100644 --- a/mooncake-store/src/client_service.cpp +++ b/mooncake-store/src/client_service.cpp @@ -3,7 +3,9 @@ #include #include "allocator.h" +#include "rebuild_retry.h" #include "segment.h" +#include "utils/base64.h" #include #include @@ -315,6 +317,12 @@ Client::~Client() { storage_heartbeat_thread_.join(); } + // === HA rebuild: stop the notify-receiving loop (set flag then join). + rebuild_notify_thread_running_.store(false); + if (rebuild_notify_thread_.joinable()) { + rebuild_notify_thread_.join(); + } + leader_monitor_running_ = false; if (leader_monitor_thread_.joinable()) { leader_monitor_thread_.join(); @@ -977,6 +985,13 @@ std::optional> Client::Create( LOG(ERROR) << "Failed to initialize local hot cache"; } + // === HA rebuild: start the notify-receiving loop now that transfer_engine_ + // is ready. It polls getNotifies() and applies cross-client UPSERT entries + // into local_replica_table_. Stopped in ~Client. + client->rebuild_notify_thread_running_.store(true); + client->rebuild_notify_thread_ = + std::thread(&Client::RebuildNotifyLoop, client.get()); + return client; } @@ -1635,6 +1650,227 @@ tl::expected Client::Put(const ObjectKey& key, return tl::unexpected(finalize_decision.error); } + // === HA rebuild: account by owner === + // A replica landing on our own segment -> record locally (we are the + // owner). A replica landing on someone else's segment -> notify that owner. + { + uint64_t value_size = 0; + for (auto n : slice_lengths) value_size += n; + const std::string tenant_id = master_client_.tenant_id(); + std::string group_id; + if (config.group_ids && !config.group_ids->empty()) + group_id = config.group_ids->front(); + for (const auto& replica : start_result.value()) { + if (!replica.is_memory_replica()) continue; + const std::string& ep = replica.get_memory_descriptor() + .buffer_descriptor.transport_endpoint_; + if (IsMyEndpoint(ep)) { + RecordLocalReplica(key, replica, value_size, config.data_type, + group_id, tenant_id); + } else { + NotifyOwnerUpsert(ep, key, replica, value_size, + config.data_type, group_id, tenant_id); + } + } + } + + return {}; +} + +// =========================================================================== +// HA rebuild: client-side local replica table helpers +// =========================================================================== + +void Client::EraseByAddressLocked(uint64_t addr) { + // Caller must hold local_replica_table_mutex_. + auto it = addr_index_.find(addr); + if (it != addr_index_.end()) { + local_replica_table_.erase(it->second); + addr_index_.erase(it); + } +} + +void Client::RecordLocalReplica(const std::string& key, + const Replica::Descriptor& replica, + uint64_t size, ObjectDataType data_type, + const std::string& group_id, + const std::string& tenant_id) { + if (!replica.is_memory_replica()) return; // only memory replicas tracked + const uint64_t addr = + replica.get_memory_descriptor().buffer_descriptor.buffer_address_; + std::lock_guard lk(local_replica_table_mutex_); + // Reuse-overwrite: evict whatever stale key currently occupies this + // address. + EraseByAddressLocked(addr); + // If the same key previously sat at a different address, drop that stale + // reverse-index entry too (rare: same key relocated). + auto old = local_replica_table_.find(key); + if (old != local_replica_table_.end()) { + const uint64_t old_addr = old->second.replica.get_memory_descriptor() + .buffer_descriptor.buffer_address_; + if (old_addr != addr) addr_index_.erase(old_addr); + } + local_replica_table_[key] = + LocalReplicaMeta{replica, size, data_type, group_id, tenant_id}; + addr_index_[addr] = key; +} + +bool Client::IsMyEndpoint(const std::string& ep) { + std::lock_guard lk(mounted_segments_mutex_); + for (const auto& [id, seg] : mounted_segments_) { + if (seg.te_endpoint == ep) return true; + } + return false; +} + +// --- notify send ------------------------------------------------------------ +// Tell the segment owner at `ep` "I stored `key` on your segment" with full +// metadata, over the TE control-plane notify channel (not one-sided RDMA). +void Client::NotifyOwnerUpsert(const std::string& ep, const std::string& key, + const Replica::Descriptor& replica, + uint64_t size, ObjectDataType data_type, + const std::string& group_id, + const std::string& tenant_id) { + KeyReplicaEntry e; + e.key = key; + e.tenant_id = tenant_id; + e.size = size; + e.data_type = data_type; + e.group_id = group_id; + e.replicas = {replica}; + std::vector entries{std::move(e)}; + // Reliability: if the send fails (peer flapping / not yet up), park it for + // the background loop to retry, so a dropped notify never silently loses a + // replica. + if (!SendUpsertNotify(ep, entries)) { + ParkPendingNotify(ep, entries); + } +} + +// Build + base64 + send one UPSERT notify. Returns true iff the peer accepted. +bool Client::SendUpsertNotify(const std::string& ep, + const std::vector& entries) { + RebuildNotify n; + n.sender_client_id = UuidToString(client_id_); + n.op = RebuildNotifyOp::UPSERT; + n.entries = entries; + TransferMetadata::NotifyDesc desc; + desc.name = n.sender_client_id; + { + // notify_msg travels as a JSON string field (UTF-8), so binary + // struct_pack output MUST be base64-encoded or it gets corrupted. + auto b = struct_pack::serialize(n); + desc.notify_msg = base64::Encode(std::string(b.begin(), b.end())); + } + int rc = transfer_engine_->sendNotifyByName(ep, desc); + if (rc != 0) { + LOG(WARNING) << "sendNotify UPSERT to " << ep << " rc=" << rc + << " (will retry)"; + return false; + } + return true; +} + +void Client::ParkPendingNotify(const std::string& ep, + const std::vector& entries) { + std::lock_guard lk(pending_notifies_mutex_); + auto& bucket = pending_notifies_[ep]; + bucket.insert(bucket.end(), entries.begin(), entries.end()); +} + +void Client::FlushPendingNotifies() { + // Snapshot + clear under lock, retry outside lock, re-park what still + // fails. + std::unordered_map> to_retry; + { + std::lock_guard lk(pending_notifies_mutex_); + if (pending_notifies_.empty()) return; + to_retry.swap(pending_notifies_); + } + for (auto& [ep, entries] : to_retry) { + if (!SendUpsertNotify(ep, entries)) { + ParkPendingNotify(ep, entries); // still down, keep for next tick + } + } +} + +void Client::NotifyOwnerUpsertBatch( + const std::unordered_map>& + by_ep) { + for (const auto& [ep, entries] : by_ep) { + // Same reliability backstop as the singular path: park on failure. + if (!SendUpsertNotify(ep, entries)) { + ParkPendingNotify(ep, entries); + } + } +} + +// --- notify receive loop ----------------------------------------------------- +// Poll getNotifies(), decode UPSERT entries, apply via RecordLocalReplica +// (which does the address-overwrite that lazy-delete correctness depends on). +void Client::RebuildNotifyLoop() { + while (rebuild_notify_thread_running_.load()) { + // Reliability backstop: retry any notifies whose send previously + // failed. + FlushPendingNotifies(); + std::vector notifies; + int rc = transfer_engine_->getNotifies(notifies); + if (rc == 0) { + for (auto& nd : notifies) { + // Reverse of the send side: base64-decode the JSON-carried + // string back to binary, then struct_pack-deserialize. + std::string bin = base64::Decode(nd.notify_msg); + RebuildNotify n; + auto ec = + struct_pack::deserialize_to(n, bin.data(), bin.size()); + if (ec != struct_pack::errc::ok) continue; + for (auto& e : n.entries) { + if (e.replicas.empty()) continue; + RecordLocalReplica(e.key, e.replicas.front(), e.size, + e.data_type, e.group_id, e.tenant_id); + } + } + } + std::this_thread::sleep_for(std::chrono::milliseconds(50)); + } +} + +// --- reconnect resend -------------------------------------------------------- +// On reconnect, snapshot the local table and resend it (batched) to the empty +// new master via the RebuildMetadata RPC. +tl::expected Client::ResendLocalReplicaTable( + ViewVersionId view_version) { + std::vector snapshot; + { + std::lock_guard lk(local_replica_table_mutex_); + snapshot.reserve(local_replica_table_.size()); + for (const auto& [key, metadata] : local_replica_table_) { + KeyReplicaEntry entry; + entry.key = key; + entry.tenant_id = metadata.tenant_id; + entry.size = metadata.size; + entry.data_type = metadata.data_type; + entry.group_id = metadata.group_id; + entry.replicas = {metadata.replica}; + snapshot.emplace_back(std::move(entry)); + } + } + + constexpr size_t kBatchSize = 256; + for (size_t i = 0; i < snapshot.size(); i += kBatchSize) { + std::vector batch( + snapshot.begin() + i, + snapshot.begin() + std::min(i + kBatchSize, snapshot.size())); + auto result = master_client_.RebuildMetadata(std::move(batch), + view_version); + if (!result) { + if (metrics_) metrics_->rebuild_failed_batches.inc(); + LOG(ERROR) << "RebuildMetadata resend failed: " + << toString(result.error()); + return tl::make_unexpected(result.error()); + } + } + return {}; } @@ -1790,6 +2026,14 @@ class PutOperation { std::vector slices; std::vector> batched_slices; + // === HA rebuild: per-key metadata for local-table accounting === + // PutOperation itself has no size/data_type/group_id/tenant_id; filled in + // StartBatchPut/StartBatchUpsert from config + slice lengths. + uint64_t meta_size{0}; + ObjectDataType meta_data_type{ObjectDataType::UNKNOWN}; + std::string meta_group_id; + std::string meta_tenant_id{"default"}; + // Enhanced state tracking PutOperationState state = PutOperationState::PENDING; tl::expected result; @@ -1927,6 +2171,16 @@ void Client::StartBatchPut(std::vector& ops, // Process individual responses with robust error handling for (size_t i = 0; i < ops.size(); ++i) { ops[i].InitializeRequestedReplicas(config); + // === HA rebuild: fill per-key metadata for local-table accounting === + { + uint64_t sz = 0; + for (const auto& s : ops[i].slices) sz += s.size; + ops[i].meta_size = sz; + ops[i].meta_data_type = config.data_type; + ops[i].meta_tenant_id = master_client_.tenant_id(); + if (config.group_ids && i < config.group_ids->size()) + ops[i].meta_group_id = config.group_ids->at(i); + } if (!start_responses[i]) { ops[i].SetTerminalError(start_responses[i].error(), PutOperationState::MASTER_FAILED, @@ -2308,6 +2562,21 @@ void Client::FinalizeBatchPut(std::vector& ops) { } if (should_succeed[i]) { op.SetSuccess(); + // === HA rebuild: account by owner === + for (const auto& replica : op.replicas) { + if (!replica.is_memory_replica()) continue; + const std::string& ep = + replica.get_memory_descriptor() + .buffer_descriptor.transport_endpoint_; + if (IsMyEndpoint(ep)) + RecordLocalReplica(op.key, replica, op.meta_size, + op.meta_data_type, op.meta_group_id, + op.meta_tenant_id); + else + NotifyOwnerUpsert(ep, op.key, replica, op.meta_size, + op.meta_data_type, op.meta_group_id, + op.meta_tenant_id); + } continue; } op.SetTerminalError(terminal_errors[i], @@ -3713,58 +3982,80 @@ void Client::StorageHeartbeatThreadMain() { int ping_fail_count = 0; auto remount_segment = [this]() { - // This lock must be held until the remount rpc is finished, - // otherwise there will be corner cases, e.g., a segment is - // unmounted successfully first, and then remounted again in - // this thread. - std::lock_guard lock(mounted_segments_mutex_); - std::vector segments; - for (auto it : mounted_segments_) { - auto& segment = it.second; - segments.emplace_back(segment); - } - auto remount_result = master_client_.ReMountSegment(segments); - if (!remount_result) { - ErrorCode err = remount_result.error(); - LOG(ERROR) << "Failed to remount segments: " << err; - } - // Re-publish Transfer Engine segment descriptors to the HTTP - // metadata server. When Master (which hosts the HTTP metadata - // server in the same process) is killed and restarted, all - // in-memory KV entries are lost. ReMountSegment above only - // restores Master-side allocation state; it does NOT write back - // the transport-level segment descriptors. Without this, remote - // peers get HTTP 404 when querying our segment descriptor and - // data transfers fail. - auto metadata = transfer_engine_->getMetadata(); - if (metadata) { - int rc = metadata->updateLocalSegmentDesc(); - if (rc != 0) { - LOG(ERROR) << "Failed to re-publish segment descriptor " - << "to metadata server, rc=" << rc - << ", will retry in next heartbeat cycle"; - segment_desc_publish_pending_.store(true); - } else { - segment_desc_publish_pending_.store(false); + rebuild_retry_pending_.store(true); + constexpr int kMaxAttempts = 3; + constexpr int kBackoffMs[kMaxAttempts] = {100, 500, 1000}; + auto wait_before_retry = [&](int attempt) { + if (metrics_) metrics_->rebuild_retries.inc(); + int remaining_ms = kBackoffMs[attempt]; + while (storage_heartbeat_running_.load() && remaining_ms > 0) { + const int sleep_ms = std::min(remaining_ms, 50); + std::this_thread::sleep_for( + std::chrono::milliseconds(sleep_ms)); + remaining_ms -= sleep_ms; } - // Also re-publish RPC meta entry (mooncake/rpc_meta/). - // Remote peers need this to locate our RDMA RPC port for - // handshake. Like segment descriptors, this entry is lost - // when the HTTP metadata server is cleared on Master restart. - rc = metadata->rePublishRpcMetaEntry(local_hostname_); - if (rc != 0) { - LOG(ERROR) << "Failed to re-publish RPC meta entry " - << "to metadata server, rc=" << rc - << ", will retry in next heartbeat cycle"; - rpc_meta_publish_pending_.store(true); - } else { - rpc_meta_publish_pending_.store(false); + return storage_heartbeat_running_.load(); + }; + + ViewVersionId view_version = 0; + { + std::lock_guard lock(leader_switch_mutex_); + if (current_master_view_.has_value()) { + view_version = current_master_view_->view_version; + } + } + + RebuildRecoveryOps ops; + ops.remount = [&]() -> tl::expected { + std::lock_guard lock(mounted_segments_mutex_); + std::vector segments; + segments.reserve(mounted_segments_.size()); + for (const auto& [id, segment] : mounted_segments_) { + segments.emplace_back(segment); + } + return master_client_.ReMountSegment(segments); + }; + ops.publish_segment_descriptor = [&]() -> tl::expected { + auto metadata = transfer_engine_->getMetadata(); + if (!metadata) { + return tl::make_unexpected(ErrorCode::INTERNAL_ERROR); + } + int rc = metadata->updateLocalSegmentDesc(); + segment_desc_publish_pending_.store(rc != 0); + if (rc != 0) return tl::make_unexpected(ErrorCode::RPC_FAIL); + return {}; + }; + ops.publish_rpc_metadata = [&]() -> tl::expected { + auto metadata = transfer_engine_->getMetadata(); + if (!metadata) { + return tl::make_unexpected(ErrorCode::INTERNAL_ERROR); } + int rc = metadata->rePublishRpcMetaEntry(local_hostname_); + rpc_meta_publish_pending_.store(rc != 0); + if (rc != 0) return tl::make_unexpected(ErrorCode::RPC_FAIL); + return {}; + }; + ops.resend_metadata = [&] { + return ResendLocalReplicaTable(view_version); + }; + ops.signal_complete = [&] { + return master_client_.SignalRebuildComplete(view_version); + }; + auto result = RunRebuildRecovery(ops, kMaxAttempts, wait_before_retry); + if (result) { + rebuild_retry_pending_.store(false); + LOG(INFO) << "[HA-REBUILD] client sent rebuild-complete signal view=" + << view_version; + return; + } + if (result.error() == ErrorCode::INVALID_VERSION) { + rebuild_retry_pending_.store(false); + LOG(WARNING) << "Aborting stale rebuild view=" << view_version; + return; } - // Note: LOCAL_DISK segment remount is NOT done here. - // It is handled by FileStorage::Heartbeat() when it detects - // SEGMENT_NOT_FOUND, which also triggers ScanMeta to - // re-register offloaded object metadata. + LOG(ERROR) << "HA metadata rebuild attempt exhausted: " + << toString(result.error()) + << "; waiting for the next heartbeat/remount trigger"; }; // Use another thread to remount segments to avoid blocking the ping // thread @@ -3785,48 +4076,23 @@ void Client::StorageHeartbeatThreadMain() { ping_fail_count = 0; last_ping_success_.store(true); auto& ping_response = ping_result.value(); + if (!leader_coordinator_) { + std::lock_guard lock(leader_switch_mutex_); + current_master_view_ = ha::MasterView{ + direct_master_address_, ping_response.view_version_id}; + } if (ping_response.client_status == ClientStatus::NEED_REMOUNT && !remount_segment_future.valid()) { // Ensure at most one remount segment thread is running remount_segment_future = std::async(std::launch::async, remount_segment); - } else if (segment_desc_publish_pending_.load() && - !remount_segment_future.valid()) { - // Previous remount succeeded but updateLocalSegmentDesc() - // failed (e.g. transient HTTP error). Retry it directly - // without re-running ReMountSegment. - auto metadata = transfer_engine_->getMetadata(); - if (metadata) { - int rc = metadata->updateLocalSegmentDesc(); - if (rc != 0) { - LOG(ERROR) - << "Retry: failed to re-publish segment " - << "descriptor to metadata server, rc=" << rc; - } else { - LOG(INFO) << "Retry: successfully re-published " - << "segment descriptor to metadata server"; - segment_desc_publish_pending_.store(false); - } - } - } else if (rpc_meta_publish_pending_.load() && + } else if ((segment_desc_publish_pending_.load() || + rpc_meta_publish_pending_.load() || + rebuild_retry_pending_.load()) && !remount_segment_future.valid()) { - // Previous remount succeeded but rePublishRpcMetaEntry() - // failed. Retry it directly. - auto metadata = transfer_engine_->getMetadata(); - if (metadata) { - int rc = metadata->rePublishRpcMetaEntry(local_hostname_); - if (rc != 0) { - LOG(ERROR) - << "Retry: failed to re-publish RPC " - << "meta entry to metadata server, rc=" << rc; - } else { - LOG(INFO) << "Retry: successfully re-published " - << "RPC meta entry to metadata server"; - rpc_meta_publish_pending_.store(false); - } - } + remount_segment_future = + std::async(std::launch::async, remount_segment); } - std::this_thread::sleep_for( std::chrono::milliseconds(success_ping_interval_ms)); continue; diff --git a/mooncake-store/src/ha/leadership/master_service_supervisor.cpp b/mooncake-store/src/ha/leadership/master_service_supervisor.cpp index 4e7093b9..4918123d 100644 --- a/mooncake-store/src/ha/leadership/master_service_supervisor.cpp +++ b/mooncake-store/src/ha/leadership/master_service_supervisor.cpp @@ -10,6 +10,7 @@ #include #include +#include #include #include "ha/leadership/leader_coordinator_factory.h" @@ -17,8 +18,36 @@ #include "ha/standby_controller.h" #include "k8s_lease_helper.h" #include "master_admin_service.h" +#include "master_metric_manager.h" #include "rpc_service.h" +// [HA rebuild gate] ★功能总开关(ld 要求:本功能必须可整体关闭)。 +// true(默认)=启用两阶段重建门控;false=升主后立即服务,作为紧急回退, +// 不封死任何请求；非HA/单master不受影响。启用时使用两阶段重建门控 +// (封死 store → 握手窗口收集 client → 盯传输收齐完成信号 → 开服务)。 +// 这是唯一的对外开关;下面 rebuild_collect_window_ms / force_serve_timeout_ms +// 是可选高级旋钮,仅在本开关=true 时生效,有合理默认值,平时不必设置。 +DEFINE_bool(enable_ha_rebuild_gate, true, + "HA: master feature switch; when true, close store service on " + "promotion and only open after client metadata rebuild completes " + "(two-phase). true (default); false = serve immediately (rollback)"); + +// [HA rebuild gate] 升主后"重建收集窗口"毫秒数。>0 时:新 leader 起服务后先 +// 封死 store 读写(serving_=false),放行重建流量(RebuildMetadata/ReMount), +// 等待本窗口时长让所有故障时已存在的 client 重连+重建完,再开服务(SetServing +// true)。<=0(默认)时:若总开关开启则回退到内置默认 7000ms;仅作高级旋钮。 +// 窗口大小建议 ≈ client 最坏重连耗时(检测3s+选主4s+重连1s≈8s)+余量。 +DEFINE_int32(rebuild_collect_window_ms, 0, + "HA: (advanced, only when enable_ha_rebuild_gate) phase-1 handshake " + "window ms to collect client ReMount before locking N; " + "<=0 => built-in default 7000 when gate enabled"); + +// [HA rebuild 两阶段] 兜底上限:升主后最多封死这么久,即使没收齐所有 client 的 +// 重建完成信号也强制开服务(防某 client 传输中挂了永不发信号导致永久封死)。 +DEFINE_int32(rebuild_force_serve_timeout_ms, 120000, + "HA: hard upper bound; force store service open this long after " + "promotion even if not all clients signaled rebuild-complete"); + namespace mooncake { namespace ha { @@ -159,9 +188,43 @@ void ActivateServingState(MasterAdminServer& admin_server, const std::shared_ptr& service, LeaderLabelReconciler& label_reconciler) { admin_server.SetServiceDelegate(service); + + // [HA rebuild gate] ★总开关:关闭(默认)→ 升主立即服务=原生行为,直接返回, + // 不封死、不开线程,完全不受本功能影响。只有显式打开才走两阶段门控。 + if (!FLAGS_enable_ha_rebuild_gate) { + service->SetServing(true); + admin_server.SetServiceAvailable(true); + SetRuntimeState(admin_server, MasterRuntimeState::kServing); + label_reconciler.SetLeader(true); + return; + } admin_server.SetServiceAvailable(true); - SetRuntimeState(admin_server, MasterRuntimeState::kServing); + SetRuntimeState(admin_server, MasterRuntimeState::kLeaderWarmup); label_reconciler.SetLeader(true); + // [HA rebuild 两阶段](仅在总开关开启时执行)升主后先封死 store,只放行重建流量。 + // 阶段1(握手窗口,规模无关):等 window_ms 让故障时已存在的 client 重连+ + // ReMount 报到,窗口结束锁定 N=已报到client数。阶段2(盯传输,规模相关): + // 等这 N 个 client 各自发 SignalRebuildComplete,收齐→开服务。兜底:force_ms + // 上限超时强制开服务。实现"重建完成前不提供 store 命中"(ld all-or-nothing)。 + int window_ms = FLAGS_rebuild_collect_window_ms; + if (window_ms <= 0) window_ms = 7000; // 开关开启但未设窗口 → 内置默认7s + LOG(INFO) << "[HA-REBUILD-GATE] store CLOSED from construction; " + << "handshake window " << window_ms + << " ms (phase-1: collect client ReMount)"; + std::weak_ptr weak = service; + const int force_ms = FLAGS_rebuild_force_serve_timeout_ms; + // 阶段1线程:等窗口时长 → 锁定 N = 当前活跃(已ReMount)client数。 + std::thread([weak, window_ms] { + std::this_thread::sleep_for(std::chrono::milliseconds(window_ms)); + if (auto s = weak.lock()) { + s->LockRebuildExpectedClients(s->GetAliveClientsSnapshot()); + } + }).detach(); + // 兜底线程:force_ms 后无论如何开服务。 + std::thread([weak, force_ms] { + std::this_thread::sleep_for(std::chrono::milliseconds(force_ms)); + if (auto s = weak.lock()) s->ForceServingAfterTimeout(); + }).detach(); } void DeactivateServingState(MasterAdminServer& admin_server, @@ -381,10 +444,12 @@ int RunSupervisorLoop(const HABackendSpec& spec, // The serving primary handles heartbeats/unmounts, so forward the // metadata cleanup config here like the non-HA path does. + auto wrapped_config = mooncake::WrappedMasterServiceConfig( + config, leadership_session->view.view_version); + wrapped_config.initially_serving = !FLAGS_enable_ha_rebuild_gate; auto wrapped_master_service = std::make_shared( - mooncake::WrappedMasterServiceConfig( - config, leadership_session->view.view_version), - config.http_metadata_server, config.http_metadata_remote_url); + wrapped_config, config.http_metadata_server, + config.http_metadata_remote_url); mooncake::RegisterRpcService(server, *wrapped_master_service); auto serve_preflight = diff --git a/mooncake-store/src/master_client.cpp b/mooncake-store/src/master_client.cpp index e2d9db34..6d626684 100644 --- a/mooncake-store/src/master_client.cpp +++ b/mooncake-store/src/master_client.cpp @@ -157,6 +157,16 @@ struct RpcNameTraits<&WrappedMasterService::ReMountSegment> { static constexpr const char* value = "ReMountSegment"; }; +template <> +struct RpcNameTraits<&WrappedMasterService::RebuildMetadata> { + static constexpr const char* value = "RebuildMetadata"; +}; + +template <> +struct RpcNameTraits<&WrappedMasterService::SignalRebuildCompleteRpc> { + static constexpr const char* value = "SignalRebuildCompleteRpc"; +}; + template <> struct RpcNameTraits<&WrappedMasterService::ReMountNoFSegment> { static constexpr const char* value = "ReMountNoFSegment"; @@ -814,6 +824,29 @@ tl::expected MasterClient::ReMountSegment( return result; } +tl::expected MasterClient::RebuildMetadata( + std::vector&& entries, ViewVersionId view_version) { + ScopedVLogTimer timer(1, "MasterClient::RebuildMetadata"); + timer.LogRequest("entries_num=", entries.size(), + ", client_id=", client_id_); + + auto result = invoke_rpc<&WrappedMasterService::RebuildMetadata, void>( + entries, client_id_, view_version); + timer.LogResponseExpected(result); + return result; +} + +tl::expected MasterClient::SignalRebuildComplete( + ViewVersionId view_version) { + ScopedVLogTimer timer(1, "MasterClient::SignalRebuildComplete"); + timer.LogRequest("client_id=", client_id_); + auto result = + invoke_rpc<&WrappedMasterService::SignalRebuildCompleteRpc, void>( + client_id_, view_version); + timer.LogResponseExpected(result); + return result; +} + tl::expected MasterClient::ReMountNoFSegment( const std::vector& segments) { ScopedVLogTimer timer(1, "MasterClient::ReMountNofSegment"); diff --git a/mooncake-store/src/master_metric_manager.cpp b/mooncake-store/src/master_metric_manager.cpp index 41202ebf..8f01e691 100644 --- a/mooncake-store/src/master_metric_manager.cpp +++ b/mooncake-store/src/master_metric_manager.cpp @@ -126,6 +126,27 @@ MasterMetricManager::MasterMetricManager() remount_segment_failures_( "master_remount_segment_failures_total", "Total number of failed RemountSegment requests"), + rebuild_metadata_requests_( + "master_rebuild_metadata_requests_total", + "Total number of RebuildMetadata requests received"), + rebuild_metadata_failures_( + "master_rebuild_metadata_failures_total", + "Total number of failed RebuildMetadata requests"), + rebuild_state_("master_rebuild_state", + "HA rebuild state: 0 rebuilding, 1 serving, 2 degraded"), + rebuild_expected_clients_("master_rebuild_expected_clients", + "Expected clients in the current rebuild"), + rebuild_completed_clients_("master_rebuild_completed_clients", + "Completed clients in the current rebuild"), + rebuild_missing_clients_("master_rebuild_missing_clients", + "Missing clients in the current rebuild"), + rebuild_duration_ms_("master_rebuild_duration_ms", + "Duration of the current rebuild in milliseconds"), + rebuild_force_open_("master_rebuild_force_open_total", + "Total number of degraded force opens"), + rebuild_stale_epoch_requests_( + "master_rebuild_stale_epoch_requests_total", + "Total number of stale rebuild requests"), mount_nof_segment_requests_( "master_mount_nof_segment_requests_total", "Total number of MountNoFSegment requests received"), @@ -996,6 +1017,33 @@ void MasterMetricManager::inc_remount_segment_requests(int64_t val) { void MasterMetricManager::inc_remount_segment_failures(int64_t val) { remount_segment_failures_.inc(val); } +void MasterMetricManager::inc_rebuild_metadata_requests(int64_t val) { + rebuild_metadata_requests_.inc(val); +} +void MasterMetricManager::inc_rebuild_metadata_failures(int64_t val) { + rebuild_metadata_failures_.inc(val); +} +void MasterMetricManager::set_rebuild_state(int64_t state) { + rebuild_state_.update(state); +} +void MasterMetricManager::set_rebuild_expected_clients(int64_t clients) { + rebuild_expected_clients_.update(clients); +} +void MasterMetricManager::set_rebuild_completed_clients(int64_t clients) { + rebuild_completed_clients_.update(clients); +} +void MasterMetricManager::set_rebuild_missing_clients(int64_t clients) { + rebuild_missing_clients_.update(clients); +} +void MasterMetricManager::set_rebuild_duration_ms(int64_t duration_ms) { + rebuild_duration_ms_.update(duration_ms); +} +void MasterMetricManager::inc_rebuild_force_open(int64_t val) { + rebuild_force_open_.inc(val); +} +void MasterMetricManager::inc_rebuild_stale_epoch_requests(int64_t val) { + rebuild_stale_epoch_requests_.inc(val); +} void MasterMetricManager::inc_remount_nof_segment_requests(int64_t val) { remount_nof_segment_requests_.inc(val); } @@ -1789,6 +1837,11 @@ std::string MasterMetricManager::serialize_metrics() { serialize_metric(key_count_); serialize_metric(soft_pin_key_count_); serialize_metric(active_clients_); + serialize_metric(rebuild_state_); + serialize_metric(rebuild_expected_clients_); + serialize_metric(rebuild_completed_clients_); + serialize_metric(rebuild_missing_clients_); + serialize_metric(rebuild_duration_ms_); // Serialize Histogram serialize_metric(value_size_distribution_); @@ -1819,6 +1872,10 @@ std::string MasterMetricManager::serialize_metrics() { serialize_metric(unmount_segment_failures_); serialize_metric(remount_segment_requests_); serialize_metric(remount_segment_failures_); + serialize_metric(rebuild_metadata_requests_); + serialize_metric(rebuild_metadata_failures_); + serialize_metric(rebuild_force_open_); + serialize_metric(rebuild_stale_epoch_requests_); serialize_metric(mount_nof_segment_requests_); serialize_metric(mount_nof_segment_failures_); serialize_metric(unmount_nof_segment_requests_); diff --git a/mooncake-store/src/master_service.cpp b/mooncake-store/src/master_service.cpp index e899f21a..e5089a45 100644 --- a/mooncake-store/src/master_service.cpp +++ b/mooncake-store/src/master_service.cpp @@ -900,6 +900,147 @@ auto MasterService::ReMountSegment(const std::vector& segments, return {}; } +// =========================================================================== +// HA rebuild: master side +// =========================================================================== + +namespace { +// (endpoint,address) identity of a memory replica; used for merge de-dup. +// get_descriptor() returns Descriptor BY VALUE, so bind it to a named var +// first -- taking a reference into a temporary would dangle. +bool RebuildSameMemoryLocation(const Replica& a, const Replica& b) { + if (!a.is_memory_replica() || !b.is_memory_replica()) return false; + const Replica::Descriptor da_desc = a.get_descriptor(); + const Replica::Descriptor db_desc = b.get_descriptor(); + const auto& da = da_desc.get_memory_descriptor().buffer_descriptor; + const auto& db = db_desc.get_memory_descriptor().buffer_descriptor; + return da.transport_endpoint_ == db.transport_endpoint_ && + da.buffer_address_ == db.buffer_address_; +} +} // namespace + +bool MasterService::ReplicaAlreadyPresent(const ObjectMetadata& meta, + const Replica& r) const { + bool found = false; + meta.VisitReplicas([](const Replica&) { return true; }, + [&](const Replica& existing) { + if (RebuildSameMemoryLocation(existing, r)) + found = true; + }); + return found; +} + +std::optional MasterService::DescriptorToReplica( + const Replica::Descriptor& desc) { + return std::visit( + [&](auto&& d) -> std::optional { + using T = std::decay_t; + if constexpr (std::is_same_v) { + const auto& bd = d.buffer_descriptor; + // Find owning segment's allocator by endpoint + address range. + std::shared_ptr base_alloc; + { + auto seg_access = segment_manager_.getSegmentAccess(); + base_alloc = seg_access.FindAllocatorByEndpointAndAddr( + bd.transport_endpoint_, bd.buffer_address_); + } + if (!base_alloc) { + LOG(WARNING) << "rebuild: no OK segment for endpoint=" + << bd.transport_endpoint_ + << " addr=" << bd.buffer_address_; + return std::nullopt; + } + // Method-1 (allocate-placeholder + rebind addr) is only safe on + // OffsetBufferAllocator (deallocate frees via offset_handle, + // not buffer_ptr). Cachelib would double-free -> refuse for + // now. + auto offset_alloc = + std::dynamic_pointer_cast( + base_alloc); + if (!offset_alloc) { + LOG(WARNING) << "rebuild: segment allocator is not " + "OffsetBufferAllocator; skip key rebuild"; + return std::nullopt; + } + auto buffer = offset_alloc->AllocateForRebuild( + bd.size_, reinterpret_cast(bd.buffer_address_)); + if (!buffer) { + LOG(WARNING) << "rebuild: AllocateForRebuild failed size=" + << bd.size_; + return std::nullopt; + } + return Replica(std::move(buffer), ReplicaStatus::COMPLETE); + } else if constexpr (std::is_same_v) { + return Replica(d.file_path, d.object_size, + ReplicaStatus::COMPLETE); + } else if constexpr (std::is_same_v) { + return Replica(d.client_id, d.object_size, d.transport_endpoint, + ReplicaStatus::COMPLETE); + } + // NoFDescriptor or others: not rebuilt in the first version. + return std::nullopt; + }, + desc.descriptor_variant); +} + +auto MasterService::RebuildMetadata(const std::vector& entries, + const UUID& client_id) + -> tl::expected { + std::shared_lock snap_lock(snapshot_mutex_); + for (const auto& e : entries) { + // (a) Descriptor -> holding Replica. + std::vector replicas; + replicas.reserve(e.replicas.size()); + bool ok = true; + for (const auto& desc : e.replicas) { + auto rep = DescriptorToReplica(desc); + if (!rep) { + ok = false; + break; + } + replicas.emplace_back(std::move(*rep)); + } + if (!ok || replicas.empty()) { + LOG(ERROR) << "rebuild: failed to restore key=" << e.key + << "; rejecting batch so the client retries"; + return tl::make_unexpected(ErrorCode::INTERNAL_ERROR); + } + + // (b) Insert or MERGE (multi-replica redundancy recovery). + const std::string tenant = + e.tenant_id.empty() ? "default" : e.tenant_id; + const ObjectIdentity oid{tenant, e.key}; + MetadataAccessorRW accessor(this, oid); + if (!accessor.Exists()) { + accessor.Create(client_id, e.size, std::move(replicas), + /*enable_soft_pin=*/false, + /*enable_hard_pin=*/false, e.data_type, e.group_id); + } else { + // Another owner already reported this key (replica_num>1): merge + // the incoming replica(s) instead of dropping them, de-duping by + // (endpoint,address). + auto& meta = accessor.Get(); + for (auto& r : replicas) { + if (!ReplicaAlreadyPresent(meta, r)) { + std::vector one; + one.emplace_back(std::move(r)); + meta.AddReplicas(std::move(one)); + } + } + } + + // (c) Mark available: GrantLease (mirrors PutEnd). Replicas rebuilt by + // DescriptorToReplica are already COMPLETE, so only mark_complete the + // ones still PROCESSING (avoids "already marked as complete" warnings). + auto& meta = accessor.Get(); + meta.VisitReplicas([](const Replica& r) { return !r.is_completed(); }, + [](Replica& r) { r.mark_complete(); }); + meta.GrantLease(0, default_kv_soft_pin_ttl_); + SyncCacheTotalAccounting(meta); + } + return {}; +} + auto MasterService::ReMountNoFSegment(const std::vector& segments, const UUID& client_id) -> tl::expected { diff --git a/mooncake-store/src/rebuild_retry.cpp b/mooncake-store/src/rebuild_retry.cpp new file mode 100644 index 00000000..2a615c7d --- /dev/null +++ b/mooncake-store/src/rebuild_retry.cpp @@ -0,0 +1,40 @@ +#include "rebuild_retry.h" + +namespace mooncake { + +tl::expected RunRebuildRecovery( + const RebuildRecoveryOps& ops, int max_attempts, + const std::function& wait_before_retry) { + ErrorCode last_error = ErrorCode::INTERNAL_ERROR; + for (int attempt = 0; attempt < max_attempts; ++attempt) { + auto remount = ops.remount(); + if (!remount) { + last_error = remount.error(); + } else { + auto descriptor = ops.publish_segment_descriptor(); + if (!descriptor) { + last_error = descriptor.error(); + } else { + auto rpc_metadata = ops.publish_rpc_metadata(); + if (!rpc_metadata) { + last_error = rpc_metadata.error(); + } else { + auto metadata = ops.resend_metadata(); + if (!metadata) { + last_error = metadata.error(); + } else { + auto complete = ops.signal_complete(); + if (complete) return {}; + last_error = complete.error(); + } + } + } + } + if (last_error == ErrorCode::INVALID_VERSION || + attempt + 1 == max_attempts || !wait_before_retry(attempt)) { + return tl::make_unexpected(last_error); + } + } + return tl::make_unexpected(last_error); +} +} // namespace mooncake diff --git a/mooncake-store/src/rpc_service.cpp b/mooncake-store/src/rpc_service.cpp index d5de362f..d3939988 100644 --- a/mooncake-store/src/rpc_service.cpp +++ b/mooncake-store/src/rpc_service.cpp @@ -1,5 +1,7 @@ #include "rpc_service.h" +#include + #include #include @@ -16,7 +18,17 @@ WrappedMasterService::WrappedMasterService( const WrappedMasterServiceConfig& config, HttpMetadataServer* http_metadata_server, const std::string& http_metadata_remote_url) - : master_service_(MasterServiceConfig(config)) { + : master_service_(MasterServiceConfig(config)), + view_version_(config.view_version), + serving_state_(config.initially_serving ? StoreServingState::SERVING + : StoreServingState::REBUILDING), + rebuild_started_at_(std::chrono::steady_clock::now()) { + MasterMetricManager::instance().set_rebuild_state( + static_cast(serving_state_.load(std::memory_order_relaxed))); + MasterMetricManager::instance().set_rebuild_expected_clients(0); + MasterMetricManager::instance().set_rebuild_completed_clients(0); + MasterMetricManager::instance().set_rebuild_missing_clients(0); + MasterMetricManager::instance().set_rebuild_duration_ms(0); // Configure metadata cleanup on client timeout. Prefer the co-located // in-process server; otherwise fall back to a separately-deployed HTTP // metadata server derived from the cluster configuration. @@ -29,13 +41,27 @@ WrappedMasterService::WrappedMasterService( WrappedMasterService::~WrappedMasterService() = default; +void WrappedMasterService::SetServing(bool on) { + const auto state = on ? StoreServingState::SERVING + : StoreServingState::REBUILDING; + serving_state_.store(state, std::memory_order_release); + MasterMetricManager::instance().set_rebuild_state( + static_cast(state)); +} + tl::expected WrappedMasterService::CalcCacheStats() { + if (!IsServing()) { + return tl::make_unexpected( + ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS); + } return MasterMetricManager::instance().calculate_cache_stats(); } tl::expected WrappedMasterService::ExistKey( const std::string& key, const std::string& tenant_id) { + if (!IsServing()) + return tl::make_unexpected(ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS); return execute_rpc( "ExistKey", [&] { return master_service_.ExistKey(key, tenant_id); }, [&](auto& timer) { timer.LogRequest("key=", key); }, @@ -45,6 +71,10 @@ tl::expected WrappedMasterService::ExistKey( std::vector> WrappedMasterService::BatchExistKey( const std::vector& keys, const std::string& tenant_id) { + if (!IsServing()) + return std::vector>( + keys.size(), + tl::make_unexpected(ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS)); ScopedVLogTimer timer(1, "BatchExistKey"); const size_t total_keys = keys.size(); timer.LogRequest("keys_count=", total_keys); @@ -80,6 +110,10 @@ tl::expected< std::unordered_map, boost::hash>, ErrorCode> WrappedMasterService::BatchQueryIp(const std::vector& client_ids) { + if (!IsServing()) { + return tl::make_unexpected( + ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS); + } ScopedVLogTimer timer(1, "BatchQueryIp"); const size_t total_client_ids = client_ids.size(); timer.LogRequest("client_ids_count=", total_client_ids); @@ -120,6 +154,10 @@ tl::expected, ErrorCode> WrappedMasterService::BatchReplicaClear( const std::vector& object_keys, const UUID& client_id, const std::string& segment_name) { + if (!IsServing()) { + return tl::make_unexpected( + ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS); + } ScopedVLogTimer timer(1, "BatchReplicaClear"); const size_t total_keys = object_keys.size(); timer.LogRequest("object_keys_count=", total_keys, @@ -159,6 +197,10 @@ tl::expected>, ErrorCode> WrappedMasterService::GetReplicaListByRegex(const std::string& str, const std::string& tenant_id) { + if (!IsServing()) { + return tl::make_unexpected( + ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS); + } return execute_rpc( "GetReplicaListByRegex", [&] { return master_service_.GetReplicaListByRegex(str, tenant_id); }, @@ -176,6 +218,8 @@ WrappedMasterService::GetReplicaListByRegex(const std::string& str, tl::expected WrappedMasterService::GetReplicaList(const std::string& key, const std::string& tenant_id) { + if (!IsServing()) + return tl::make_unexpected(ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS); return execute_rpc( "GetReplicaList", [&] { return master_service_.GetReplicaList(key, tenant_id); }, @@ -189,6 +233,10 @@ WrappedMasterService::GetReplicaList(const std::string& key, std::vector> WrappedMasterService::BatchGetReplicaList(const std::vector& keys, const std::string& tenant_id) { + if (!IsServing()) + return std::vector>( + keys.size(), + tl::make_unexpected(ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS)); ScopedVLogTimer timer(1, "BatchGetReplicaList"); const size_t total_keys = keys.size(); timer.LogRequest("keys_count=", total_keys); @@ -233,12 +281,21 @@ WrappedMasterService::BatchGetReplicaList(const std::vector& keys, std::vector> WrappedMasterService::BatchGetReplicaListForAdmin( const std::vector& keys, const std::string& tenant_id) { + if (!IsServing()) { + return std::vector>( + keys.size(), tl::make_unexpected( + ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS)); + } return master_service_.BatchGetReplicaListForAdmin(keys, tenant_id); } tl::expected WrappedMasterService::GetReplicaListForAdmin(const std::string& key, const std::string& tenant_id) { + if (!IsServing()) { + return tl::make_unexpected( + ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS); + } return execute_rpc( "GetReplicaListForAdmin", [&] { return master_service_.GetReplicaListForAdmin(key, tenant_id); }, @@ -250,6 +307,8 @@ WrappedMasterService::PutStart(const UUID& client_id, const std::string& key, const uint64_t slice_length, const ReplicateConfig& config, const std::string& tenant_id) { + if (!IsServing()) + return tl::make_unexpected(ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS); return execute_rpc( "PutStart", [&] { @@ -267,6 +326,10 @@ WrappedMasterService::PutStart(const UUID& client_id, const std::string& key, tl::expected WrappedMasterService::PutEnd( const UUID& client_id, const std::string& key, ReplicaType replica_type, const std::string& tenant_id) { + if (!IsServing()) { + return tl::make_unexpected( + ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS); + } return execute_rpc( "PutEnd", [&] { @@ -284,6 +347,10 @@ tl::expected WrappedMasterService::PutEnd( tl::expected WrappedMasterService::PutRevoke( const UUID& client_id, const std::string& key, ReplicaType replica_type, const std::string& tenant_id) { + if (!IsServing()) { + return tl::make_unexpected( + ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS); + } return execute_rpc( "PutRevoke", [&] { @@ -304,6 +371,11 @@ WrappedMasterService::BatchPutStart(const UUID& client_id, const std::vector& slice_lengths, const ReplicateConfig& config, const std::string& tenant_id) { + if (!IsServing()) + return std::vector< + tl::expected, ErrorCode>>( + keys.size(), + tl::make_unexpected(ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS)); ScopedVLogTimer timer(1, "BatchPutStart"); const size_t total_keys = keys.size(); timer.LogRequest("client_id=", client_id, ", keys_count=", total_keys); @@ -396,6 +468,11 @@ WrappedMasterService::BatchPutStart(const UUID& client_id, std::vector> WrappedMasterService::BatchPutEnd( const UUID& client_id, const std::vector& keys, ReplicaType replica_type, const std::string& tenant_id) { + if (!IsServing()) { + return std::vector>( + keys.size(), tl::make_unexpected( + ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS)); + } ScopedVLogTimer timer(1, "BatchPutEnd"); const size_t total_keys = keys.size(); timer.LogRequest("client_id=", client_id, ", keys_count=", total_keys); @@ -436,6 +513,11 @@ std::vector> WrappedMasterService::BatchPutEnd( std::vector> WrappedMasterService::BatchPutRevoke( const UUID& client_id, const std::vector& keys, ReplicaType replica_type, const std::string& tenant_id) { + if (!IsServing()) { + return std::vector>( + keys.size(), tl::make_unexpected( + ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS)); + } ScopedVLogTimer timer(1, "BatchPutRevoke"); const size_t total_keys = keys.size(); timer.LogRequest("client_id=", client_id, ", keys_count=", total_keys); @@ -478,6 +560,10 @@ WrappedMasterService::UpsertStart(const UUID& client_id, const std::string& key, const uint64_t slice_length, const ReplicateConfig& config, const std::string& tenant_id) { + if (!IsServing()) { + return tl::make_unexpected( + ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS); + } return execute_rpc( "UpsertStart", [&] { @@ -495,6 +581,10 @@ WrappedMasterService::UpsertStart(const UUID& client_id, const std::string& key, tl::expected WrappedMasterService::UpsertEnd( const UUID& client_id, const std::string& key, ReplicaType replica_type, const std::string& tenant_id) { + if (!IsServing()) { + return tl::make_unexpected( + ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS); + } return execute_rpc( "UpsertEnd", [&] { @@ -512,6 +602,10 @@ tl::expected WrappedMasterService::UpsertEnd( tl::expected WrappedMasterService::UpsertRevoke( const UUID& client_id, const std::string& key, ReplicaType replica_type, const std::string& tenant_id) { + if (!IsServing()) { + return tl::make_unexpected( + ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS); + } return execute_rpc( "UpsertRevoke", [&] { @@ -531,6 +625,11 @@ WrappedMasterService::BatchUpsertStart( const UUID& client_id, const std::vector& keys, const std::vector& slice_lengths, const ReplicateConfig& config, const std::string& tenant_id) { + if (!IsServing()) { + return std::vector, ErrorCode>>( + keys.size(), tl::make_unexpected( + ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS)); + } ScopedVLogTimer timer(1, "BatchUpsertStart"); const size_t total_keys = keys.size(); timer.LogRequest("client_id=", client_id, ", keys_count=", total_keys); @@ -566,6 +665,11 @@ WrappedMasterService::BatchUpsertStart( std::vector> WrappedMasterService::BatchUpsertEnd( const UUID& client_id, const std::vector& keys, const std::string& tenant_id) { + if (!IsServing()) { + return std::vector>( + keys.size(), tl::make_unexpected( + ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS)); + } ScopedVLogTimer timer(1, "BatchUpsertEnd"); const size_t total_keys = keys.size(); timer.LogRequest("client_id=", client_id, ", keys_count=", total_keys); @@ -601,6 +705,11 @@ std::vector> WrappedMasterService::BatchUpsertRevoke(const UUID& client_id, const std::vector& keys, const std::string& tenant_id) { + if (!IsServing()) { + return std::vector>( + keys.size(), tl::make_unexpected( + ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS)); + } ScopedVLogTimer timer(1, "BatchUpsertRevoke"); const size_t total_keys = keys.size(); timer.LogRequest("client_id=", client_id, ", keys_count=", total_keys); @@ -635,6 +744,10 @@ WrappedMasterService::BatchUpsertRevoke(const UUID& client_id, tl::expected WrappedMasterService::Remove( const std::string& key, bool force, const std::string& tenant_id) { + if (!IsServing()) { + return tl::make_unexpected( + ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS); + } return execute_rpc( "Remove", [&] { return master_service_.Remove(key, tenant_id, force); }, [&](auto& timer) { timer.LogRequest("key=", key, ", force=", force); }, @@ -644,6 +757,10 @@ tl::expected WrappedMasterService::Remove( tl::expected WrappedMasterService::RemoveByRegex( const std::string& str, bool force, const std::string& tenant_id) { + if (!IsServing()) { + return tl::make_unexpected( + ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS); + } return execute_rpc( "RemoveByRegex", [&] { return master_service_.RemoveByRegex(str, tenant_id, force); }, @@ -654,7 +771,11 @@ tl::expected WrappedMasterService::RemoveByRegex( [] { MasterMetricManager::instance().inc_remove_by_regex_failures(); }); } -long WrappedMasterService::RemoveAll(bool force, const std::string& tenant_id) { +tl::expected WrappedMasterService::RemoveAll(bool force, const std::string& tenant_id) { + if (!IsServing()) { + return tl::make_unexpected( + ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS); + } ScopedVLogTimer timer(1, "RemoveAll"); timer.LogRequest("action=remove_all_objects, force=", force); MasterMetricManager::instance().inc_remove_all_requests(); @@ -666,6 +787,11 @@ long WrappedMasterService::RemoveAll(bool force, const std::string& tenant_id) { std::vector> WrappedMasterService::BatchRemove( const std::vector& keys, bool force, const std::string& tenant_id) { + if (!IsServing()) { + return std::vector>( + keys.size(), tl::make_unexpected( + ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS)); + } ScopedVLogTimer timer(1, "BatchRemove"); const size_t total_keys = keys.size(); timer.LogRequest("keys_count=", total_keys, ", force=", force); @@ -689,6 +815,10 @@ std::vector> WrappedMasterService::BatchRemove( tl::expected WrappedMasterService::MountSegment( const Segment& segment, const UUID& client_id) { + if (!IsServing()) { + return tl::make_unexpected( + ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS); + } return execute_rpc( "MountSegment", [&] { return master_service_.MountSegment(segment, client_id); }, @@ -703,6 +833,10 @@ tl::expected WrappedMasterService::MountSegment( tl::expected WrappedMasterService::MountNoFSegment( const NoFSegment& segment, const UUID& client_id) { + if (!IsServing()) { + return tl::make_unexpected( + ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS); + } return execute_rpc( "MountNoFSegment", [&] { return master_service_.MountNoFSegment(segment, client_id); }, @@ -733,6 +867,159 @@ tl::expected WrappedMasterService::ReMountSegment( [] { MasterMetricManager::instance().inc_remount_segment_failures(); }); } +tl::expected WrappedMasterService::RebuildMetadata( + const std::vector& entries, const UUID& client_id, + ViewVersionId view_version) { + if (!IsCurrentView(view_version)) { + MasterMetricManager::instance().inc_rebuild_stale_epoch_requests(); + return tl::make_unexpected(ErrorCode::INVALID_VERSION); + } + return execute_rpc( + "RebuildMetadata", + [&] { return master_service_.RebuildMetadata(entries, client_id); }, + [&](auto& timer) { + timer.LogRequest("entries_count=", entries.size(), + ", client_id=", client_id); + }, + [] { MasterMetricManager::instance().inc_rebuild_metadata_requests(); }, + [] { + MasterMetricManager::instance().inc_rebuild_metadata_failures(); + }); +} + +WrappedMasterService::ClientSet +WrappedMasterService::GetAliveClientsSnapshot() const { + return master_service_.getAliveClientsSnapshot(); +} + +std::string WrappedMasterService::MissingClientsLocked() const { + std::ostringstream missing; + bool first = true; + for (const auto& client_id : rebuild_expected_clients_) { + if (rebuild_done_clients_.contains(client_id)) continue; + if (!first) missing << ","; + missing << client_id.first << "-" << client_id.second; + first = false; + } + return missing.str(); +} + +void WrappedMasterService::TransitionToLocked(StoreServingState state, + const char* reason) { + const auto previous = serving_state_.load(std::memory_order_acquire); + if (previous == state || + (previous == StoreServingState::SERVING && + state == StoreServingState::DEGRADED)) { + return; + } + serving_state_.store(state, std::memory_order_release); + MasterMetricManager::instance().set_rebuild_state( + static_cast(state)); + auto& metrics = MasterMetricManager::instance(); + metrics.set_rebuild_duration_ms( + std::chrono::duration_cast( + std::chrono::steady_clock::now() - rebuild_started_at_) + .count()); + metrics.set_rebuild_missing_clients(static_cast( + rebuild_expected_clients_.size() - rebuild_done_clients_.size())); + if (state == StoreServingState::DEGRADED) { + metrics.inc_rebuild_force_open(); + } + LOG(INFO) << "[HA-REBUILD-GATE] view=" << view_version_ + << " store state=" + << (state == StoreServingState::SERVING ? "SERVING" : "DEGRADED") + << " reason=" << reason << " completed=" + << rebuild_done_clients_.size() << "/" + << rebuild_expected_clients_.size() << " missing_clients=[" + << MissingClientsLocked() << "]"; +} + +void WrappedMasterService::MaybeFinishRebuildLocked() { + if (!rebuild_window_locked_ || + rebuild_done_clients_.size() != rebuild_expected_clients_.size()) { + return; + } + TransitionToLocked(StoreServingState::SERVING, "all expected clients rebuilt"); +} + +void WrappedMasterService::LockRebuildExpectedClients( + ClientSet expected_clients) { + std::lock_guard lk(rebuild_mu_); + if (rebuild_window_locked_) return; + rebuild_window_locked_ = true; + rebuild_expected_clients_ = std::move(expected_clients); + for (const auto& client_id : rebuild_done_before_lock_) { + if (rebuild_expected_clients_.contains(client_id)) { + rebuild_done_clients_.insert(client_id); + } + } + rebuild_done_before_lock_.clear(); + MasterMetricManager::instance().set_rebuild_expected_clients( + static_cast(rebuild_expected_clients_.size())); + MasterMetricManager::instance().set_rebuild_completed_clients( + static_cast(rebuild_done_clients_.size())); + MasterMetricManager::instance().set_rebuild_missing_clients( + static_cast(rebuild_expected_clients_.size() - + rebuild_done_clients_.size())); + LOG(INFO) << "[HA-REBUILD-GATE] view=" << view_version_ + << " handshake window closed; expected_clients=" + << rebuild_expected_clients_.size() << " completed=" + << rebuild_done_clients_.size(); + MaybeFinishRebuildLocked(); +} + +tl::expected WrappedMasterService::SignalRebuildComplete( + const UUID& client_id, ViewVersionId view_version) { + if (!IsCurrentView(view_version)) { + MasterMetricManager::instance().inc_rebuild_stale_epoch_requests(); + LOG(WARNING) << "[HA-REBUILD-GATE] stale rebuild-complete view=" + << view_version << " current_view=" << view_version_; + return tl::make_unexpected(ErrorCode::INVALID_VERSION); + } + std::lock_guard lk(rebuild_mu_); + if (!rebuild_window_locked_) { + if (rebuild_done_before_lock_.size() >= 10000 && + !rebuild_done_before_lock_.contains(client_id)) { + LOG(ERROR) << "[HA-REBUILD-GATE] too many pre-lock completion " + "signals; refusing unbounded growth"; + return tl::make_unexpected(ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS); + } + rebuild_done_before_lock_.insert(client_id); + return {}; + } + if (!rebuild_expected_clients_.contains(client_id)) { + LOG(WARNING) << "[HA-REBUILD-GATE] ignoring non-expected client=(" + << client_id.first << "," << client_id.second << ") view=" + << view_version_; + return {}; + } + rebuild_done_clients_.insert(client_id); + MasterMetricManager::instance().set_rebuild_completed_clients( + static_cast(rebuild_done_clients_.size())); + MasterMetricManager::instance().set_rebuild_missing_clients( + static_cast(rebuild_expected_clients_.size() - + rebuild_done_clients_.size())); + LOG(INFO) << "[HA-REBUILD-GATE] view=" << view_version_ + << " rebuild-complete from client=(" << client_id.first << "," + << client_id.second << "); " << rebuild_done_clients_.size() + << "/" << rebuild_expected_clients_.size(); + MaybeFinishRebuildLocked(); + return {}; +} + +void WrappedMasterService::ForceServingAfterTimeout() { + std::lock_guard lk(rebuild_mu_); + if (serving_state_.load(std::memory_order_acquire) == + StoreServingState::REBUILDING) { + TransitionToLocked(StoreServingState::DEGRADED, "rebuild timeout"); + } +} + +tl::expected WrappedMasterService::SignalRebuildCompleteRpc( + const UUID& client_id, ViewVersionId view_version) { + return SignalRebuildComplete(client_id, view_version); +} + tl::expected WrappedMasterService::ReMountNoFSegment( const std::vector& segments, const UUID& client_id) { return execute_rpc( @@ -752,6 +1039,10 @@ tl::expected WrappedMasterService::ReMountNoFSegment( tl::expected WrappedMasterService::UnmountSegment( const UUID& segment_id, const UUID& client_id) { + if (!IsServing()) { + return tl::make_unexpected( + ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS); + } return execute_rpc( "UnmountSegment", [&] { return master_service_.UnmountSegment(segment_id, client_id); }, @@ -765,6 +1056,10 @@ tl::expected WrappedMasterService::UnmountSegment( tl::expected WrappedMasterService::GracefulUnmountSegment( const UUID& segment_id, const UUID& client_id, uint64_t grace_period_ms) { + if (!IsServing()) { + return tl::make_unexpected( + ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS); + } return execute_rpc( "GracefulUnmountSegment", [&] { @@ -788,6 +1083,10 @@ tl::expected WrappedMasterService::GracefulUnmountSegment( tl::expected WrappedMasterService::UnmountNoFSegment( const UUID& segment_id, const UUID& client_id) { + if (!IsServing()) { + return tl::make_unexpected( + ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS); + } return execute_rpc( "UnmountNoFSegment", [&] { @@ -807,6 +1106,10 @@ tl::expected WrappedMasterService::UnmountNoFSegment( tl::expected, ErrorCode> WrappedMasterService::GetAllNoFSegments() { + if (!IsServing()) { + return tl::make_unexpected( + ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS); + } return execute_rpc( "GetAllNoFSegments", [&] { return master_service_.GetAllNoFSegments(); }, @@ -816,6 +1119,10 @@ WrappedMasterService::GetAllNoFSegments() { tl::expected, ErrorCode> WrappedMasterService::GetNoFSegmentsByName(const std::string& segment_name) { + if (!IsServing()) { + return tl::make_unexpected( + ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS); + } return execute_rpc( "GetNoFSegmentsByName", [&] { return master_service_.GetNoFSegmentsByName(segment_name); }, @@ -827,6 +1134,10 @@ tl::expected WrappedMasterService::CopyStart( const UUID& client_id, const std::string& key, const std::string& tenant_id, const std::string& src_segment, const std::vector& tgt_segments) { + if (!IsServing()) { + return tl::make_unexpected( + ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS); + } return execute_rpc( "CopyStart", [&] { @@ -846,6 +1157,10 @@ tl::expected WrappedMasterService::CopyStart( tl::expected WrappedMasterService::CopyEnd( const UUID& client_id, const std::string& key, const std::string& tenant_id) { + if (!IsServing()) { + return tl::make_unexpected( + ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS); + } return execute_rpc( "CopyEnd", [&] { return master_service_.CopyEnd(client_id, key, tenant_id); }, @@ -860,6 +1175,10 @@ tl::expected WrappedMasterService::CopyEnd( tl::expected WrappedMasterService::CopyRevoke( const UUID& client_id, const std::string& key, const std::string& tenant_id) { + if (!IsServing()) { + return tl::make_unexpected( + ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS); + } return execute_rpc( "CopyRevoke", [&] { return master_service_.CopyRevoke(client_id, key, tenant_id); }, @@ -874,6 +1193,10 @@ tl::expected WrappedMasterService::CopyRevoke( tl::expected WrappedMasterService::MoveStart( const UUID& client_id, const std::string& key, const std::string& tenant_id, const std::string& src_segment, const std::string& tgt_segment) { + if (!IsServing()) { + return tl::make_unexpected( + ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS); + } return execute_rpc( "MoveStart", [&] { @@ -893,6 +1216,10 @@ tl::expected WrappedMasterService::MoveStart( tl::expected WrappedMasterService::MoveEnd( const UUID& client_id, const std::string& key, const std::string& tenant_id) { + if (!IsServing()) { + return tl::make_unexpected( + ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS); + } return execute_rpc( "MoveEnd", [&] { return master_service_.MoveEnd(client_id, key, tenant_id); }, @@ -907,6 +1234,10 @@ tl::expected WrappedMasterService::MoveEnd( tl::expected WrappedMasterService::MoveRevoke( const UUID& client_id, const std::string& key, const std::string& tenant_id) { + if (!IsServing()) { + return tl::make_unexpected( + ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS); + } return execute_rpc( "MoveRevoke", [&] { return master_service_.MoveRevoke(client_id, key, tenant_id); }, @@ -921,6 +1252,10 @@ tl::expected WrappedMasterService::MoveRevoke( tl::expected WrappedMasterService::EvictDiskReplica( const UUID& client_id, const std::string& key, const std::string& tenant_id, ReplicaType replica_type) { + if (!IsServing()) { + return tl::make_unexpected( + ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS); + } return execute_rpc( "EvictDiskReplica", [&] { @@ -944,6 +1279,11 @@ std::vector> WrappedMasterService::BatchEvictDiskReplica( const UUID& client_id, const std::vector& keys, const std::string& tenant_id, ReplicaType replica_type) { + if (!IsServing()) { + return std::vector>( + keys.size(), tl::make_unexpected( + ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS)); + } ScopedVLogTimer timer(1, "BatchEvictDiskReplica"); const size_t total_keys = keys.size(); timer.LogRequest("client_id=", client_id, ", keys_count=", total_keys, @@ -976,6 +1316,10 @@ WrappedMasterService::BatchEvictDiskReplica( tl::expected WrappedMasterService::CreateCopyTask( const std::string& key, const std::string& tenant_id, const std::vector& targets) { + if (!IsServing()) { + return tl::make_unexpected( + ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS); + } return execute_rpc( "CreateCopyTask", [&] { return master_service_.CreateCopyTask(key, tenant_id, targets); }, @@ -992,6 +1336,10 @@ tl::expected WrappedMasterService::CreateCopyTask( tl::expected WrappedMasterService::CreateMoveTask( const std::string& key, const std::string& tenant_id, const std::string& source, const std::string& target) { + if (!IsServing()) { + return tl::make_unexpected( + ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS); + } return execute_rpc( "CreateMoveTask", [&] { @@ -1010,6 +1358,10 @@ tl::expected WrappedMasterService::CreateMoveTask( tl::expected WrappedMasterService::QueryTask( const UUID& task_id) { + if (!IsServing()) { + return tl::make_unexpected( + ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS); + } return execute_rpc( "QueryTask", [&] { return master_service_.QueryTask(task_id); }, [&](auto& timer) { timer.LogRequest("task_id=", task_id); }, @@ -1019,6 +1371,10 @@ tl::expected WrappedMasterService::QueryTask( tl::expected, ErrorCode> WrappedMasterService::FetchTasks(const UUID& client_id, size_t batch_size) { + if (!IsServing()) { + return tl::make_unexpected( + ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS); + } return execute_rpc( "FetchTasks", [&] { return master_service_.FetchTasks(client_id, batch_size); }, @@ -1032,6 +1388,10 @@ WrappedMasterService::FetchTasks(const UUID& client_id, size_t batch_size) { tl::expected WrappedMasterService::MarkTaskToComplete( const UUID& client_id, const TaskCompleteRequest& request) { + if (!IsServing()) { + return tl::make_unexpected( + ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS); + } return execute_rpc( "MarkTaskToComplete", [&] { return master_service_.MarkTaskToComplete(client_id, request); }, @@ -1082,6 +1442,10 @@ tl::expected WrappedMasterService::ServiceReady() { tl::expected, ErrorCode> WrappedMasterService::ListTenantQuotaSnapshots() { + if (!IsServing()) { + return tl::make_unexpected( + ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS); + } if (!master_service_.IsTenantQuotaEnabled()) { return tl::make_unexpected(ErrorCode::UNAVAILABLE_IN_CURRENT_MODE); } @@ -1090,6 +1454,10 @@ WrappedMasterService::ListTenantQuotaSnapshots() { tl::expected WrappedMasterService::GetTenantQuotaSnapshot(const std::string& tenant_id) { + if (!IsServing()) { + return tl::make_unexpected( + ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS); + } if (!master_service_.IsTenantQuotaEnabled()) { return tl::make_unexpected(ErrorCode::UNAVAILABLE_IN_CURRENT_MODE); } @@ -1103,6 +1471,10 @@ WrappedMasterService::GetTenantQuotaSnapshot(const std::string& tenant_id) { tl::expected WrappedMasterService::UpsertTenantQuotaPolicy(const std::string& tenant_id, uint64_t requested_quota_bytes) { + if (!IsServing()) { + return tl::make_unexpected( + ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS); + } if (!master_service_.IsTenantQuotaEnabled()) { return tl::make_unexpected(ErrorCode::UNAVAILABLE_IN_CURRENT_MODE); } @@ -1112,6 +1484,10 @@ WrappedMasterService::UpsertTenantQuotaPolicy(const std::string& tenant_id, tl::expected, ErrorCode> WrappedMasterService::DeleteTenantQuotaPolicy(const std::string& tenant_id) { + if (!IsServing()) { + return tl::make_unexpected( + ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS); + } if (!master_service_.IsTenantQuotaEnabled()) { return tl::make_unexpected(ErrorCode::UNAVAILABLE_IN_CURRENT_MODE); } @@ -1120,6 +1496,10 @@ WrappedMasterService::DeleteTenantQuotaPolicy(const std::string& tenant_id) { tl::expected WrappedMasterService::GetTenantQuotaAllocatableCapacityBytes() { + if (!IsServing()) { + return tl::make_unexpected( + ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS); + } if (!master_service_.IsTenantQuotaEnabled()) { return tl::make_unexpected(ErrorCode::UNAVAILABLE_IN_CURRENT_MODE); } @@ -1128,6 +1508,10 @@ WrappedMasterService::GetTenantQuotaAllocatableCapacityBytes() { tl::expected, ErrorCode> WrappedMasterService::GetAllKeysForAdmin() { + if (!IsServing()) { + return tl::make_unexpected( + ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS); + } // Compatibility endpoint: /get_all_keys historically listed only the // default tenant's keys. return master_service_.GetAllKeys("default"); @@ -1135,21 +1519,37 @@ WrappedMasterService::GetAllKeysForAdmin() { tl::expected, ErrorCode> WrappedMasterService::GetAllSegmentsForAdmin() { + if (!IsServing()) { + return tl::make_unexpected( + ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS); + } return master_service_.GetAllSegments(); } tl::expected, ErrorCode> WrappedMasterService::GetSegmentsDetailForAdmin() { + if (!IsServing()) { + return tl::make_unexpected( + ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS); + } return master_service_.GetSegmentsDetail(); } tl::expected, ErrorCode> WrappedMasterService::QuerySegmentForAdmin(const std::string& segment) { + if (!IsServing()) { + return tl::make_unexpected( + ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS); + } return master_service_.QuerySegments(segment); } tl::expected WrappedMasterService::MountLocalDiskSegment( const UUID& client_id, bool enable_offloading) { + if (!IsServing()) { + return tl::make_unexpected( + ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS); + } ScopedVLogTimer timer(1, "MountLocalDiskSegment"); timer.LogRequest("action=mount_local_disk_segment"); LOG(INFO) << "Mount local disk segment with client id is : " << client_id @@ -1164,6 +1564,10 @@ tl::expected WrappedMasterService::MountLocalDiskSegment( tl::expected, ErrorCode> WrappedMasterService::OffloadObjectHeartbeat(const UUID& client_id, bool enable_offloading) { + if (!IsServing()) { + return tl::make_unexpected( + ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS); + } ScopedVLogTimer timer(1, "OffloadObjectHeartbeat"); timer.LogRequest("action=offload_object_heartbeat"); auto result = @@ -1173,6 +1577,10 @@ WrappedMasterService::OffloadObjectHeartbeat(const UUID& client_id, tl::expected WrappedMasterService::ReportSsdCapacity( const UUID& client_id, int64_t ssd_total_capacity_bytes) { + if (!IsServing()) { + return tl::make_unexpected( + ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS); + } ScopedVLogTimer timer(1, "ReportSsdCapacity"); timer.LogRequest("client_id=", client_id, ", ssd_total_capacity_bytes=", ssd_total_capacity_bytes); @@ -1183,6 +1591,10 @@ tl::expected WrappedMasterService::ReportSsdCapacity( tl::expected WrappedMasterService::NotifyOffloadSuccess( const UUID& client_id, const std::vector& tasks, const std::vector& metadatas) { + if (!IsServing()) { + return tl::make_unexpected( + ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS); + } ScopedVLogTimer timer(1, "NotifyOffloadSuccess"); timer.LogRequest("action=notify_offload_success"); @@ -1194,6 +1606,10 @@ tl::expected WrappedMasterService::NotifyOffloadSuccess( tl::expected, ErrorCode> WrappedMasterService::PromotionObjectHeartbeat(const UUID& client_id) { + if (!IsServing()) { + return tl::make_unexpected( + ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS); + } ScopedVLogTimer timer(1, "PromotionObjectHeartbeat"); timer.LogRequest("action=promotion_object_heartbeat"); return master_service_.PromotionObjectHeartbeat(client_id); @@ -1203,6 +1619,10 @@ tl::expected WrappedMasterService::PromotionAllocStart( const UUID& client_id, const std::string& key, const std::string& tenant_id, uint64_t size, const std::vector& preferred_segments) { + if (!IsServing()) { + return tl::make_unexpected( + ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS); + } ScopedVLogTimer timer(1, "PromotionAllocStart"); timer.LogRequest("action=promotion_alloc_start"); auto result = master_service_.PromotionAllocStart(client_id, key, tenant_id, @@ -1214,6 +1634,10 @@ WrappedMasterService::PromotionAllocStart( tl::expected WrappedMasterService::NotifyPromotionSuccess( const UUID& client_id, const std::string& key, const std::string& tenant_id) { + if (!IsServing()) { + return tl::make_unexpected( + ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS); + } ScopedVLogTimer timer(1, "NotifyPromotionSuccess"); timer.LogRequest("action=notify_promotion_success"); auto result = @@ -1225,6 +1649,10 @@ tl::expected WrappedMasterService::NotifyPromotionSuccess( tl::expected WrappedMasterService::NotifyPromotionFailure( const UUID& client_id, const std::string& key, const std::string& tenant_id) { + if (!IsServing()) { + return tl::make_unexpected( + ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS); + } ScopedVLogTimer timer(1, "NotifyPromotionFailure"); timer.LogRequest("action=notify_promotion_failure"); auto result = @@ -1235,26 +1663,46 @@ tl::expected WrappedMasterService::NotifyPromotionFailure( tl::expected WrappedMasterService::CreateDrainJob( const CreateDrainJobRequest& request) { + if (!IsServing()) { + return tl::make_unexpected( + ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS); + } return master_service_.CreateDrainJob(request); } tl::expected WrappedMasterService::QueryDrainJob( const UUID& job_id) { + if (!IsServing()) { + return tl::make_unexpected( + ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS); + } return master_service_.QueryDrainJob(job_id); } tl::expected WrappedMasterService::CancelDrainJob( const UUID& job_id) { + if (!IsServing()) { + return tl::make_unexpected( + ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS); + } return master_service_.CancelDrainJob(job_id); } tl::expected WrappedMasterService::QuerySegmentStatus( const std::string& segment_name) { + if (!IsServing()) { + return tl::make_unexpected( + ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS); + } return master_service_.QuerySegmentStatus(segment_name); } tl::expected WrappedMasterService::QuerySegmentStatusById(const UUID& segment_id) { + if (!IsServing()) { + return tl::make_unexpected( + ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS); + } return master_service_.QuerySegmentStatusById(segment_id); } @@ -1321,6 +1769,11 @@ void RegisterRpcService( &wrapped_master_service); server.register_handler<&mooncake::WrappedMasterService::ReMountSegment>( &wrapped_master_service); + server.register_handler<&mooncake::WrappedMasterService::RebuildMetadata>( + &wrapped_master_service); + server.register_handler< + &mooncake::WrappedMasterService::SignalRebuildCompleteRpc>( + &wrapped_master_service); server.register_handler<&mooncake::WrappedMasterService::ReMountNoFSegment>( &wrapped_master_service); server.register_handler<&mooncake::WrappedMasterService::UnmountSegment>( diff --git a/mooncake-store/src/segment.cpp b/mooncake-store/src/segment.cpp index 7009b896..d6c111e7 100644 --- a/mooncake-store/src/segment.cpp +++ b/mooncake-store/src/segment.cpp @@ -412,6 +412,22 @@ ErrorCode ScopedSegmentAccess::GetClientSegments( return ErrorCode::OK; } +std::shared_ptr +ScopedSegmentAccess::FindAllocatorByEndpointAndAddr( + const std::string& te_endpoint, uintptr_t buffer_address) const { + for (const auto& [id, ms] : segment_manager_->mounted_segments_) { + if (ms.status != SegmentStatus::OK) continue; + if (ms.segment.te_endpoint != te_endpoint) continue; + // Disambiguate 1:N endpoint sharing by address range. + const uintptr_t base = ms.segment.base; + const uintptr_t end = base + ms.segment.size; + if (buffer_address >= base && buffer_address < end) { + return ms.buf_allocator; + } + } + return nullptr; +} + void ScopedSegmentAccess::UnmountLocalDiskSegment(const UUID& client_id) { auto it = segment_manager_->client_local_disk_segment_.find(client_id); if (it != segment_manager_->client_local_disk_segment_.end()) { diff --git a/mooncake-store/tests/CMakeLists.txt b/mooncake-store/tests/CMakeLists.txt index e7e66c9f..32a70c65 100644 --- a/mooncake-store/tests/CMakeLists.txt +++ b/mooncake-store/tests/CMakeLists.txt @@ -62,6 +62,8 @@ if(ENABLE_KV_EVENTS) endif() endif() add_store_test(master_service_test master_service_test.cpp) +add_store_test(ha_rebuild_gate_test ha_rebuild_gate_test.cpp) +add_store_test(rebuild_retry_test rebuild_retry_test.cpp) add_store_test(master_service_tenant_quota_test master_service_tenant_quota_test.cpp) add_store_test(batch_remove_test batch_remove_test.cpp) @@ -115,6 +117,21 @@ add_store_test(master_snapshot_codec_test add_store_test(master_service_test_for_snapshot ha/snapshot/master_service_test_for_snapshot.cpp) add_store_test(non_ha_reconnect_test non_ha_reconnect_test.cpp) +add_store_test(client_metadata_rebuild_test client_metadata_rebuild_test.cpp) +# Live HA recovery test: standalone client (own main, no gtest) connecting to a +# real separate mooncake_master process. Driven by ha_live_test.sh. +add_executable(ha_recovery_live_main ha_recovery_live_main.cpp) +target_include_directories(ha_recovery_live_main PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}) +target_link_libraries(ha_recovery_live_main + PUBLIC mooncake_store transfer_engine cachelib_memory_allocator + ${ETCD_WRAPPER_LIB} glog gflags ibverbs pthread) +# HA scale benchmark: measures rebuild latency, scale, and recovery-window +# latency/availability. Driven by ha_scale_bench.sh. +add_executable(ha_scale_bench_main ha_scale_bench_main.cpp) +target_include_directories(ha_scale_bench_main PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}) +target_link_libraries(ha_scale_bench_main + PUBLIC mooncake_store transfer_engine cachelib_memory_allocator + ${ETCD_WRAPPER_LIB} glog gflags ibverbs pthread) add_store_test(storage_backend_test storage_backend_test.cpp) add_store_test(mutex_test mutex_test.cpp) add_store_test(file_storage_test file_storage_test.cpp) diff --git a/mooncake-store/tests/client_metadata_rebuild_test.cpp b/mooncake-store/tests/client_metadata_rebuild_test.cpp new file mode 100644 index 00000000..5e9c2925 --- /dev/null +++ b/mooncake-store/tests/client_metadata_rebuild_test.cpp @@ -0,0 +1,716 @@ +// Unit tests for the client-driven metadata rebuild HA feature. After the +// master process crashes and restarts empty, each client resends its locally +// tracked key-to-location table so the master rebuilds its metadata with zero +// recomputation. +// +// Tests: +// RebuildObjectMetadataAfterMasterRestart: a single client rebuilds its +// object metadata after a master restart and reads data back with no +// recompute. +// RebuiltMetadataPointsToRealData: rebuilt metadata must point to the real, +// correct data, not to some other key's bytes. +// LazyDelete_RemovedButNotReused_MayRevive: a removed key may revive after +// rebuild if its space was not reused, and its data is still correct. +// CrossClientRebuildViaNotify: client A stores onto client B's segment, so B +// records the replica via notify and rebuilds it after a master restart. +// NotifyRetryBackstopRedeliversDroppedNotify: a failed notify is parked and +// retried by a background thread until the owner receives it. +// RemovedKeySpaceReuseNoStaleMapping: once a removed key's space is reused, +// the old key must not revive and point at the new key's address. +// MultiReplicaMergedOnRebuild: with replica_num=2, the master merges the two +// replicas reported by different owners back into two on rebuild. + +#include +#include + +#include +#include +#include +#include +#include +#include + +#include "allocator.h" +#include "client_service.h" +#include "types.h" +#include "utils.h" // allocate_buffer_allocator_memory, SimpleAllocator +#include "test_server_helpers.h" // InProcMaster, InProcMasterConfigBuilder +#include "default_config.h" + +DEFINE_string(protocol, "tcp", "Transfer protocol: rdma|tcp"); + +namespace mooncake { +namespace testing { + +namespace { + +// Wait until all keys can be read back, not just a single probe key. +// If RebuildMetadata pushes keys one by one instead of atomically, a probe +// key may arrive while keys[N-1] has not, so asserting immediately would +// wrongly report a live key as lost. Requiring every key to Get successfully +// works for both incremental and atomic implementations. Returns false if +// some key is still not rebuilt when the timeout is reached. +bool WaitForAllKeysRebuilt(std::shared_ptr& client, + SimpleAllocator& allocator, + const std::vector& keys, + const std::vector& values, + int max_attempts = 40, int interval_ms = 500) { + for (int attempt = 0; attempt < max_attempts; ++attempt) { + bool all_ok = true; + for (size_t i = 0; i < keys.size(); ++i) { + void* buf = allocator.allocate(values[i].size()); + std::vector slices{Slice{buf, values[i].size()}}; + auto res = client->Get(keys[i], slices); + allocator.deallocate(buf, values[i].size()); + if (!res.has_value()) { + all_ok = false; + break; + } + } + if (all_ok) { + LOG(INFO) << "All " << keys.size() << " keys rebuilt after " + << attempt << " polls"; + return true; + } + std::this_thread::sleep_for(std::chrono::milliseconds(interval_ms)); + } + return false; +} + +tl::expected PutString(std::shared_ptr& client, + SimpleAllocator& allocator, + const std::string& key, + const std::string& value) { + void* buf = allocator.allocate(value.size()); + std::memcpy(buf, value.data(), value.size()); + std::vector slices{Slice{buf, value.size()}}; + ReplicateConfig config; + config.replica_num = 1; + auto res = client->Put(key, slices, config); + allocator.deallocate(buf, value.size()); + return res; +} + +// Get the value and compare it byte for byte against the expected content. +bool GetAndVerify(std::shared_ptr& client, SimpleAllocator& allocator, + const std::string& key, const std::string& expected) { + void* buf = allocator.allocate(expected.size()); + std::vector slices{Slice{buf, expected.size()}}; + auto res = client->Get(key, slices); + bool ok = res.has_value() && slices[0].size == expected.size() && + std::memcmp(slices[0].ptr, expected.data(), expected.size()) == 0; + allocator.deallocate(buf, expected.size()); + return ok; +} + +} // namespace + +class ClientMetadataRebuildTest : public ::testing::Test { + protected: + void SetUp() override { + // In-process non-HA master (auto-selected port). + ASSERT_TRUE(master_.Start(InProcMasterConfigBuilder().build())); + master_address_ = master_.master_address(); + + // Create the client. + local_hostname_ = "127.0.0.1:19100"; + auto client_opt = + Client::Create(local_hostname_, "P2PHANDSHAKE", FLAGS_protocol, + std::nullopt, master_address_); + ASSERT_TRUE(client_opt.has_value()) << "Failed to create client"; + client_ = client_opt.value(); + + // Data buffer allocator, registered as local memory. Put must have the + // buffer registered before it can transfer from it. + allocator_ = std::make_unique(kAllocSize); + auto reg = client_->RegisterLocalMemory( + allocator_->getBase(), kAllocSize, "cpu:0", false, false); + ASSERT_TRUE(reg.has_value()) << "RegisterLocalMemory failed"; + + // Mount a segment where data will land, using the three-argument + // overload that takes a protocol. + seg_ptr_ = allocate_buffer_allocator_memory(kSegmentSize); + ASSERT_NE(seg_ptr_, nullptr); + auto mount = + client_->MountSegment(seg_ptr_, kSegmentSize, FLAGS_protocol); + ASSERT_TRUE(mount.has_value()) << toString(mount.error()); + } + + void TearDown() override { + if (client_ && seg_ptr_) { + client_->UnmountSegment(seg_ptr_, kSegmentSize); + } + master_.Stop(); + } + + // Simulate a master failure by bringing up an empty new master on the same + // port, which is what lets the client reconnect back to it. + void RestartMasterEmpty() { + master_.Stop(); + std::this_thread::sleep_for( + std::chrono::seconds(3)); // let heartbeats fail + ASSERT_TRUE(master_.Start( + InProcMasterConfigBuilder() + .set_rpc_port(master_.rpc_port()) + .set_http_metrics_port(master_.http_metrics_port()) + .build())); + } + + static constexpr size_t kSegmentSize = 128 * 1024 * 1024; // 128MB segment + static constexpr size_t kAllocSize = 64 * 1024 * 1024; // 64MB buffer + + InProcMaster master_; + std::string master_address_; + std::string local_hostname_; + std::shared_ptr client_; + void* seg_ptr_ = nullptr; + std::unique_ptr allocator_; +}; + +// --------------------------------------------------------------------------- +// Test 1 (core): after a master restart the client rebuilds object metadata +// and the data reads back with no recomputation. +// --------------------------------------------------------------------------- +TEST_F(ClientMetadataRebuildTest, RebuildObjectMetadataAfterMasterRestart) { + const int kNumKeys = 50; + std::vector keys, values; + for (int i = 0; i < kNumKeys; ++i) { + keys.push_back("rebuild_key_" + std::to_string(i)); + values.push_back("rebuild_value_" + std::to_string(i)); + } + for (int i = 0; i < kNumKeys; ++i) { + auto r = PutString(client_, *allocator_, keys[i], values[i]); + ASSERT_TRUE(r.has_value()) + << "Put failed " << keys[i] << ": " << toString(r.error()); + } + // Baseline: everything reads back before the restart. + for (int i = 0; i < kNumKeys; ++i) + ASSERT_TRUE(GetAndVerify(client_, *allocator_, keys[i], values[i])) + << "Baseline Get failed " << keys[i]; + + RestartMasterEmpty(); + + // Wait for all keys to rebuild, not just one probe, to avoid a false + // failure under an incremental implementation. + ASSERT_TRUE(WaitForAllKeysRebuilt(client_, *allocator_, keys, values)) + << "Metadata not fully rebuilt in time: new master still returns " + "NOT_FOUND, or the resend/rebuild path is not working"; + + // Core assertion: after rebuild every key reads back with matching content + // (no recomputation). + for (int i = 0; i < kNumKeys; ++i) + EXPECT_TRUE(GetAndVerify(client_, *allocator_, keys[i], values[i])) + << "After rebuild, Get/verify failed " << keys[i]; +} + +// --------------------------------------------------------------------------- +// Test 2 (guard against fake recovery): rebuilt metadata must point to the +// real and correct data, never to another key's bytes. +// --------------------------------------------------------------------------- +TEST_F(ClientMetadataRebuildTest, RebuiltMetadataPointsToRealData) { + std::vector> kv = { + {"distinct_A", std::string(1024, 'A')}, + {"distinct_B", std::string(2048, 'B')}, + {"distinct_C", std::string(512, 'C')}, + {"distinct_D", std::string(4096, 'D')}, + }; + std::vector keys, values; + for (auto& [k, v] : kv) { + keys.push_back(k); + values.push_back(v); + } + + for (auto& [k, v] : kv) { + auto r = PutString(client_, *allocator_, k, v); + ASSERT_TRUE(r.has_value()) << "Put failed " << k; + } + RestartMasterEmpty(); + ASSERT_TRUE(WaitForAllKeysRebuilt(client_, *allocator_, keys, values)) + << "Metadata not fully rebuilt in time"; + for (auto& [k, v] : kv) + EXPECT_TRUE(GetAndVerify(client_, *allocator_, k, v)) + << "Rebuilt metadata for '" << k + << "' points to wrong/corrupt data"; +} + +// --------------------------------------------------------------------------- +// Test 3 (lazy-delete semantics): a key that was Removed but whose space was +// not reused may revive after rebuild, and the revived data is still correct +// because Remove does not wipe memory. This is the expected lazy-delete +// behavior, not a bug. +// --------------------------------------------------------------------------- +// NOTE on the semantic change: an earlier eager-delete version asserted that a +// removed key does not revive. It is now lazy-delete: Remove leaves the client +// local table untouched, so a removed-but-not-reused key revives and still +// points to the correct old data. This test checks that live keys work +// normally and that removed-but-not-reused keys revive with intact data, which +// is acceptable. The "old key does not revive after reuse" property is covered +// by test 5, which is the correctness that must be guaranteed. +TEST_F(ClientMetadataRebuildTest, LazyDelete_RemovedButNotReused_MayRevive) { + const int kNumKeys = 20; + std::vector keys, values; + for (int i = 0; i < kNumKeys; ++i) { + keys.push_back("mix_key_" + std::to_string(i)); + values.push_back("mix_value_" + std::to_string(i)); + auto r = PutString(client_, *allocator_, keys.back(), values.back()); + ASSERT_TRUE(r.has_value()) << "Put failed " << keys.back(); + } + // Remove half of the keys (even indices), with force=true to bypass the + // lease. No new Put happens after the removes, so the space is not reused. + for (int i = 0; i < kNumKeys; i += 2) { + auto r = client_->Remove(keys[i], /*force=*/true); + ASSERT_TRUE(r.has_value()) + << "Remove failed " << keys[i] << ": " << toString(r.error()); + } + RestartMasterEmpty(); + + // Lazy-delete: every key (including the removed ones) may rebuild, so wait + // for all of them. + ASSERT_TRUE(WaitForAllKeysRebuilt(client_, *allocator_, keys, values)) + << "Rebuild did not complete (under lazy-delete a removed-but-not-" + "reused key should also be able to revive)"; + + // Assertion: every key, removed or not, reads back with correct content, + // which is the expected lazy-delete outcome. A revived removed key is + // acceptable; what matters is that the data is intact (Remove does not + // wipe memory). + for (int i = 0; i < kNumKeys; ++i) { + EXPECT_TRUE(GetAndVerify(client_, *allocator_, keys[i], values[i])) + << (i % 2 == 0 ? "removed-not-reused key data should be correct " + "after revive: " + : "live key data should be correct: ") + << keys[i]; + } +} + +// --------------------------------------------------------------------------- +// Test 4 (cross-client / notify path, the core multi-client scenario): +// clientA mounts no segment and clientB does, so A's Put data necessarily +// lands on B's segment because it is the only segment in the global pool. +// Those keys are recorded on B via an A-to-B notify, and after a master +// restart B (where the data physically lives) resends them for rebuild. This +// is the test that best represents a real multi-client scenario and where the +// core value of the feature lies. +// +// Reliable way to force data onto B's segment (following the two-client +// pattern in client_integration_test.cpp, where segment_provider_client_ +// mounts and test_client_ only RegisterLocalMemory): +// clientB mounts the only allocatable segment. +// clientA only RegisterLocalMemory for its local read/write buffer and does +// not MountSegment. +// At PutStart the global pool has only B's segment, so data lands on B +// deterministically and the test is not flaky. +// +// NOTE: this depends on notify send/receive being implemented. Without it the +// data is on B, A's local table has none of these keys, and B was not +// notified, so after the restart nobody resends them and the test fails. That +// failure is exactly the evidence that notify is being exercised. +// NOTE: uses its own fixture that brings up A and B, rather than reusing the +// single-client ClientMetadataRebuildTest. +class ClientCrossNotifyTest : public ::testing::Test { + protected: + void SetUp() override { + ASSERT_TRUE(master_.Start(InProcMasterConfigBuilder().build())); + master_address_ = master_.master_address(); + + // clientB is the segment owner and mounts the only allocatable + // segment. + auto b = Client::Create("127.0.0.1:19201", "P2PHANDSHAKE", + FLAGS_protocol, std::nullopt, master_address_); + ASSERT_TRUE(b.has_value()); + clientB_ = b.value(); + segB_ = allocate_buffer_allocator_memory(kSeg); + ASSERT_NE(segB_, nullptr); + ASSERT_TRUE( + clientB_->MountSegment(segB_, kSeg, FLAGS_protocol).has_value()); + + // clientA is the writer and only registers a local read/write buffer; + // it does not mount a segment. + auto a = Client::Create("127.0.0.1:19202", "P2PHANDSHAKE", + FLAGS_protocol, std::nullopt, master_address_); + ASSERT_TRUE(a.has_value()); + clientA_ = a.value(); + allocA_ = std::make_unique(kAlloc); + ASSERT_TRUE(clientA_ + ->RegisterLocalMemory(allocA_->getBase(), kAlloc, + "cpu:0", false, false) + .has_value()); + // B also needs a local read/write buffer, used when it Gets to verify. + allocB_ = std::make_unique(kAlloc); + ASSERT_TRUE(clientB_ + ->RegisterLocalMemory(allocB_->getBase(), kAlloc, + "cpu:0", false, false) + .has_value()); + } + + void TearDown() override { + if (clientB_ && segB_) clientB_->UnmountSegment(segB_, kSeg); + master_.Stop(); + } + + void RestartMasterEmpty() { + master_.Stop(); + std::this_thread::sleep_for(std::chrono::seconds(3)); + ASSERT_TRUE(master_.Start( + InProcMasterConfigBuilder() + .set_rpc_port(master_.rpc_port()) + .set_http_metrics_port(master_.http_metrics_port()) + .build())); + } + + static constexpr size_t kSeg = 128 * 1024 * 1024; + static constexpr size_t kAlloc = 64 * 1024 * 1024; + InProcMaster master_; + std::string master_address_; + std::shared_ptr clientA_, clientB_; + void* segB_ = nullptr; + std::unique_ptr allocA_, allocB_; +}; + +TEST_F(ClientCrossNotifyTest, CrossClientRebuildViaNotify) { + const int kNumKeys = 30; + std::vector keys, values; + for (int i = 0; i < kNumKeys; ++i) { + keys.push_back("cross_key_" + std::to_string(i)); + values.push_back("cross_val_" + std::to_string(i)); + } + // 1. A Put (data must land on B's segment; the A-to-B notify makes B + // record it). + for (int i = 0; i < kNumKeys; ++i) { + auto r = PutString(clientA_, *allocA_, keys[i], values[i]); + ASSERT_TRUE(r.has_value()) + << "A Put failed " << keys[i] << ": " << toString(r.error()); + } + // (Optional strong assertion) Confirm the data really landed on B's segment + // by comparing the replica endpoint from Query against B (see + // client_integration_test.cpp:419-423). This can be enabled if Client + // exposes GetTransportEndpoint: + // { auto q = clientA_->Query(keys[0]); ASSERT_TRUE(q.has_value()); + // EXPECT_EQ(q.value().replicas[0].get_memory_descriptor() + // .buffer_descriptor.transport_endpoint_, + // clientB_->GetTransportEndpoint()); } + + // 2. Baseline: A can read back (data is on B's segment, so Get asks the + // master for the location and then reads over the transport engine). + for (int i = 0; i < kNumKeys; ++i) + ASSERT_TRUE(GetAndVerify(clientA_, *allocA_, keys[i], values[i])) + << "baseline A Get " << keys[i]; + + // 3. Master crashes and restarts empty. + RestartMasterEmpty(); + + // 4. Wait for rebuild. The key point: data is on B's segment and A's local + // table has none of these keys, so only B (which recorded them via + // notify) can resend them. If notify is not working this times out. + ASSERT_TRUE(WaitForAllKeysRebuilt(clientA_, *allocA_, keys, values)) + << "Cross-client metadata not rebuilt in time: notify recording or B's " + "resend path is not working"; + + // 5. Core assertion: after rebuild both A and B can read with correct + // content. + for (int i = 0; i < kNumKeys; ++i) { + EXPECT_TRUE(GetAndVerify(clientA_, *allocA_, keys[i], values[i])) + << "After rebuild, A Get/verify failed " << keys[i]; + EXPECT_TRUE(GetAndVerify(clientB_, *allocB_, keys[i], values[i])) + << "After rebuild, B Get/verify failed " << keys[i]; + } +} + +// --------------------------------------------------------------------------- +// Test 7 (notify reliability backstop): a failed notify must not be dropped +// silently. A failed notify is parked in a pending queue and retried by a +// background thread until the peer receives it. Without this backstop, one +// lost notify makes the owner miss a replica, and that redundancy is silently +// lost when the master rebuilds. +// --------------------------------------------------------------------------- +// Approach (deterministic and reproducible): A first Puts a key normally so +// the data lands on B's segment, and uses Query to get the real Descriptor +// pointing at B's segment. Then it uses ParkNotifyForTest to park a notify for +// a new key headed to B, simulating "this notify originally failed to send". +// Then: +// 1. Assert the pending bucket count is 1 (it really was parked). +// 2. Wait for the background RebuildNotifyLoop's FlushPendingNotifies to +// redeliver it, so pending drains to zero. +// 3. Restart the master and assert that this new key, which was recorded +// only thanks to redelivery, can also be resent and rebuilt by B. +// If the backstop is missing (a failed send is just dropped), pending never +// drains and the new key never rebuilds, so the test fails. +TEST_F(ClientCrossNotifyTest, NotifyRetryBackstopRedeliversDroppedNotify) { + // 1. A Puts a carrier key normally (lands on B's segment) and gets its + // real Descriptor pointing at B's segment. + const std::string carrier = "carrier_key"; + const std::string carrier_val = std::string(4096, 'C'); + ASSERT_TRUE(PutString(clientA_, *allocA_, carrier, carrier_val).has_value()) + << "carrier Put failed"; + auto q = clientA_->Query(carrier); + ASSERT_TRUE(q.has_value() && !q.value().replicas.empty()) + << "carrier Query failed"; + const Replica::Descriptor& carrier_desc = q.value().replicas.front(); + const std::string owner_ep = carrier_desc.get_memory_descriptor() + .buffer_descriptor.transport_endpoint_; + + // 2. Simulate "the notify to B originally failed" by parking a notify for + // a new key into pending. Reuse the carrier's Descriptor as this new + // key's replica location; the focus is the redelivery path, not address + // authenticity, and the new key goes through B's RecordLocalReplica so B + // can later resend it. + const std::string dropped = "dropped_notify_key"; + clientA_->ParkNotifyForTest(owner_ep, dropped, carrier_desc, + carrier_val.size(), ObjectDataType::UNKNOWN, "", + "default"); + + // 1. It really was parked. + EXPECT_GE(clientA_->PendingNotifyBucketCountForTest(), 1u) + << "A failed notify should be parked in the pending queue (without the " + "backstop it would not be parked)"; + + // 2. Wait for the background thread to redeliver, draining pending to zero. + bool drained = false; + for (int i = 0; i < 40 && !drained; ++i) { + if (clientA_->PendingNotifyBucketCountForTest() == 0) { + drained = true; + break; + } + std::this_thread::sleep_for(std::chrono::milliseconds(200)); + } + EXPECT_TRUE(drained) + << "pending notify did not drain in time via redelivery: the " + "background retry backstop is not working"; + + // Give B's receive thread a moment to record the redelivered notify into + // its local table. + std::this_thread::sleep_for(std::chrono::seconds(1)); + + // 3. Master crashes and restarts empty. + RestartMasterEmpty(); + + // 3. Assert the dropped key, recorded only thanks to redelivery, can be + // resent and rebuilt by B (the master knows it). + bool rebuilt = false; + for (int i = 0; i < 40 && !rebuilt; ++i) { + auto qq = clientA_->Query(dropped); + if (qq.has_value() && !qq.value().replicas.empty()) { + rebuilt = true; + break; + } + std::this_thread::sleep_for(std::chrono::milliseconds(500)); + } + EXPECT_TRUE(rebuilt) + << "the key for the redelivered notify was not rebuilt: without the " + "retry backstop this notify would be lost, the owner would miss the " + "record, and the replica would be silently lost on rebuild"; +} + +// --------------------------------------------------------------------------- +// Test 5 (address-reuse overwrite, the core correctness guarantee of +// lazy-delete): once a removed key's space is reused by a new key, the removed +// key must not revive on rebuild and point at the address now owned by the new +// key, which would be silent data corruption. Under lazy-delete this is the +// property that MUST be guaranteed. The "not-reused may revive" case from test +// 3 is acceptable, but "old key still present after reuse" is never +// acceptable. +// --------------------------------------------------------------------------- +// Principle: key_A is removed (lazy-delete leaves the client local table +// untouched), so its space in the segment returns to the allocator freelist +// and a later Put reuses the same address. Recording must overwrite by +// address: the new key must clear the old key_A entry in the local table that +// points at the same (segment, address). RecordLocalReplica does this via +// EraseByAddressLocked. This case is a single client Putting onto its own +// segment, so it takes RecordLocalReplica's self-overwrite path and sends no +// notify. In the cross-client case (A writes to B's segment) the UPSERT notify +// triggers the same RecordLocalReplica overwrite on B. Without the +// address-overwrite, resend after restart would report key_A pointing at the +// old address, which now holds the new key's data, causing silent corruption. +// This test forces reuse and verifies that key_A does not revive (its old +// entry was overwritten) and that the new key data is entirely correct. +TEST_F(ClientMetadataRebuildTest, RemovedKeySpaceReuseNoStaleMapping) { + const std::string kA = "reuse_victim_A"; + const std::string vA = std::string(4096, 'X'); // reused by same-size key + + // 1. Put key_A and confirm it reads back (it occupies some address in the + // segment). + ASSERT_TRUE(PutString(client_, *allocator_, kA, vA).has_value()) + << "Put key_A failed"; + ASSERT_TRUE(GetAndVerify(client_, *allocator_, kA, vA)) << "baseline key_A"; + + // 2. force-remove key_A (lazy-delete: the client local table is untouched; + // the space returns to the allocator freelist). + ASSERT_TRUE(client_->Remove(kA, /*force=*/true).has_value()) + << "Remove key_A failed"; + + // 3. Put a batch of same-size new keys to force the allocator to reuse the + // address key_A just freed. The UPSERT recording on reuse must overwrite + // key_A's old local-table entry by address. + const int kNumNew = 64; + std::vector new_keys, new_values; + for (int i = 0; i < kNumNew; ++i) { + new_keys.push_back("reuse_new_" + std::to_string(i)); + new_values.push_back( + std::string(4096, static_cast('a' + i % 26))); + ASSERT_TRUE(PutString(client_, *allocator_, new_keys[i], new_values[i]) + .has_value()) + << "Put new key failed " << new_keys[i]; + } + + // 4. Master crashes and restarts empty, then wait for all new keys to + // rebuild. + RestartMasterEmpty(); + ASSERT_TRUE( + WaitForAllKeysRebuilt(client_, *allocator_, new_keys, new_values)) + << "New keys not fully rebuilt in time"; + + // 5a. Core assertion 1: the removed key_A must not revive. + { + void* buf = allocator_->allocate(vA.size()); + std::vector slices{Slice{buf, vA.size()}}; + auto res = client_->Get(kA, slices); + allocator_->deallocate(buf, vA.size()); + EXPECT_FALSE(res.has_value()) + << "removed key_A revived (bug in the delete cleanup logic): if it " + "still points at the address now reused by a new key, that is " + "silent data corruption"; + } + + // 5b. Core assertion 2: all new key data must be entirely correct (not + // polluted by key_A's stale mapping). + for (int i = 0; i < kNumNew; ++i) + EXPECT_TRUE( + GetAndVerify(client_, *allocator_, new_keys[i], new_values[i])) + << "New key data polluted or lost: " << new_keys[i]; +} + +// --------------------------------------------------------------------------- +// Test 6 (multi-replica merge, replica_num=2): the two replicas of one key +// land on different segments (different clients) and are resent by their +// respective owners. On rebuild the master must merge them back into "this key +// has 2 replicas" rather than keeping only one (RebuildMetadata takes the +// merge branch for an already-present key instead of continue-skipping it). +// --------------------------------------------------------------------------- +// NOTE: this is the only test for redundant multi-replica recovery; all other +// tests use replica_num=1 and never reach the merge branch. +// Preconditions: both clients MountSegment so there are two different segments +// for replica_num=2 to spread across, and owner recording, per-owner resend, +// and master merge are all implemented. +// Fixture: both clients mount a segment, unlike test 4 where A mounts none. +class ClientMultiReplicaTest : public ::testing::Test { + protected: + void SetUp() override { + ASSERT_TRUE(master_.Start(InProcMasterConfigBuilder().build())); + master_address_ = master_.master_address(); + for (int i = 0; i < 2; ++i) { + auto c = Client::Create("127.0.0.1:1930" + std::to_string(i + 1), + "P2PHANDSHAKE", FLAGS_protocol, + std::nullopt, master_address_); + ASSERT_TRUE(c.has_value()); + clients_[i] = c.value(); + seg_[i] = allocate_buffer_allocator_memory(kSeg); + ASSERT_NE(seg_[i], nullptr); + ASSERT_TRUE(clients_[i] + ->MountSegment(seg_[i], kSeg, FLAGS_protocol) + .has_value()); + alloc_[i] = std::make_unique(kAlloc); + ASSERT_TRUE(clients_[i] + ->RegisterLocalMemory(alloc_[i]->getBase(), kAlloc, + "cpu:0", false, false) + .has_value()); + } + } + void TearDown() override { + for (int i = 0; i < 2; ++i) + if (clients_[i] && seg_[i]) + clients_[i]->UnmountSegment(seg_[i], kSeg); + master_.Stop(); + } + void RestartMasterEmpty() { + master_.Stop(); + std::this_thread::sleep_for(std::chrono::seconds(3)); + ASSERT_TRUE(master_.Start( + InProcMasterConfigBuilder() + .set_rpc_port(master_.rpc_port()) + .set_http_metrics_port(master_.http_metrics_port()) + .build())); + } + // Return the replica count the master records for this key, taken from + // replicas.size() via Query. + int ReplicaCount(const std::string& key) { + auto q = clients_[0]->Query(key); + return q.has_value() ? static_cast(q.value().replicas.size()) : -1; + } + static constexpr size_t kSeg = 128 * 1024 * 1024; + static constexpr size_t kAlloc = 64 * 1024 * 1024; + InProcMaster master_; + std::string master_address_; + std::shared_ptr clients_[2]; + void* seg_[2] = {nullptr, nullptr}; + std::unique_ptr alloc_[2]; +}; + +TEST_F(ClientMultiReplicaTest, MultiReplicaMergedOnRebuild) { + const int kNumKeys = 20; + std::vector keys, values; + for (int i = 0; i < kNumKeys; ++i) { + keys.push_back("dual_key_" + std::to_string(i)); + values.push_back("dual_val_" + std::to_string(i)); + } + // 1. Put with replica_num=2: two replicas per key, spread across the two + // clients' segments. + for (int i = 0; i < kNumKeys; ++i) { + void* buf = alloc_[0]->allocate(values[i].size()); + std::memcpy(buf, values[i].data(), values[i].size()); + std::vector slices{Slice{buf, values[i].size()}}; + ReplicateConfig cfg; + cfg.replica_num = 2; // key point: 2 replicas + auto r = clients_[0]->Put(keys[i], slices, cfg); + alloc_[0]->deallocate(buf, values[i].size()); + ASSERT_TRUE(r.has_value()) << "Put(replica_num=2) failed " << keys[i] + << ": " << toString(r.error()); + } + + // 2. Baseline: each key should have 2 replicas before the restart. + for (int i = 0; i < kNumKeys; ++i) + ASSERT_EQ(ReplicaCount(keys[i]), 2) + << "baseline: key should have 2 replicas " << keys[i]; + + // 3. Master crashes and restarts empty. + RestartMasterEmpty(); + + // 4. Wait for rebuild: the two owners each report their own replica and the + // master merges them. A replica count of 2 marks completion. + bool merged = false; + for (int attempt = 0; attempt < 40 && !merged; ++attempt) { + merged = true; + for (int i = 0; i < kNumKeys; ++i) + if (ReplicaCount(keys[i]) != 2) { + merged = false; + break; + } + if (!merged) + std::this_thread::sleep_for(std::chrono::milliseconds(500)); + } + + // 5. KEY ASSERTION: after rebuild every key is restored to 2 replicas + // (merge succeeded, redundancy preserved). If the master used continue + // to skip (old code), this would be 1 and the test fails. + for (int i = 0; i < kNumKeys; ++i) + EXPECT_EQ(ReplicaCount(keys[i]), 2) + << "key should have 2 replicas after rebuild (multi-replica " + "merge): " + << keys[i] + << " -- a count of 1 means the master did not merge and lost the " + "second replica (redundancy lost)"; + + // 6. Data is still readable and correct. + for (int i = 0; i < kNumKeys; ++i) { + void* buf = alloc_[0]->allocate(values[i].size()); + std::vector slices{Slice{buf, values[i].size()}}; + auto res = clients_[0]->Get(keys[i], slices); + bool ok = + res.has_value() && + std::memcmp(slices[0].ptr, values[i].data(), values[i].size()) == 0; + alloc_[0]->deallocate(buf, values[i].size()); + EXPECT_TRUE(ok) << "data should be correct after rebuild " << keys[i]; + } +} + +} // namespace testing +} // namespace mooncake diff --git a/mooncake-store/tests/ha_live_test.sh b/mooncake-store/tests/ha_live_test.sh new file mode 100755 index 00000000..75698e47 --- /dev/null +++ b/mooncake-store/tests/ha_live_test.sh @@ -0,0 +1,120 @@ +#!/usr/bin/env bash +# Live HA recovery test (out-of-process): +# Drives ha_recovery_live_main against a REAL, separate mooncake_master +# PROCESS. Verifies that after the master is killed and restarted on the same +# port, the standalone client re-registers its held key->location metadata so +# that every key is readable again with ZERO recompute. +# +# Timeline: +# 1. start mooncake_master (non-HA) on $MASTER_PORT +# 2. start ha_recovery_live_main -> mounts a segment, Puts N keys, prints +# "READY_FOR_KILL", then polls Get in a loop +# 3. once we see READY_FOR_KILL: kill -9 the master (reads start failing) +# 4. restart mooncake_master on the SAME port -> client reconnects, resends +# its local_replica_table_, master rebuilds metadata, reads succeed again +# 5. client prints RESULT=PASS/FAIL; this script propagates that as exit code +set -u + +BUILD_DIR="${BUILD_DIR:-/export/home/shenshuwei.3/Mooncake/build}" +MASTER_BIN="${MASTER_BIN:-$BUILD_DIR/mooncake-store/src/mooncake_master}" +CLIENT_BIN="${CLIENT_BIN:-$BUILD_DIR/mooncake-store/tests/ha_recovery_live_main}" + +MASTER_PORT="${MASTER_PORT:-50055}" +LOCAL_ADDR="${LOCAL_ADDR:-127.0.0.1:19110}" +NKEYS="${NKEYS:-50}" +PROTOCOL="${PROTOCOL:-tcp}" + +WORKDIR="$(mktemp -d /tmp/ha_live_test.XXXXXX)" +MASTER_LOG="$WORKDIR/master.log" +CLIENT_LOG="$WORKDIR/client.log" + +MASTER_PID="" +CLIENT_PID="" + +log() { echo "[ha_live_test] $*"; } + +cleanup() { + [ -n "$CLIENT_PID" ] && kill -9 "$CLIENT_PID" 2>/dev/null + [ -n "$MASTER_PID" ] && kill -9 "$MASTER_PID" 2>/dev/null + wait 2>/dev/null +} +trap cleanup EXIT + +start_master() { + "$MASTER_BIN" -rpc_port="$MASTER_PORT" -enable_ha=false \ + -enable_metric_reporting=false >>"$MASTER_LOG" 2>&1 & + MASTER_PID=$! + log "started master pid=$MASTER_PID port=$MASTER_PORT" +} + +wait_for_line() { # $1=file $2=pattern $3=timeout_sec + local f="$1" pat="$2" t="$3" i=0 + while [ "$i" -lt "$((t * 2))" ]; do + grep -q "$pat" "$f" 2>/dev/null && return 0 + # bail out early if the process we depend on already died + [ -n "$CLIENT_PID" ] && ! kill -0 "$CLIENT_PID" 2>/dev/null && \ + grep -q "$pat" "$f" 2>/dev/null && return 0 + sleep 0.5 + i=$((i + 1)) + done + return 1 +} + +[ -x "$MASTER_BIN" ] || { log "FATAL master bin missing: $MASTER_BIN"; exit 3; } +[ -x "$CLIENT_BIN" ] || { log "FATAL client bin missing: $CLIENT_BIN"; exit 3; } + +log "workdir=$WORKDIR" + +# --- 1. start master --- +start_master +sleep 2 +if ! kill -0 "$MASTER_PID" 2>/dev/null; then + log "FATAL master died on startup; log:"; cat "$MASTER_LOG"; exit 3 +fi + +# --- 2. start client --- +"$CLIENT_BIN" -master="127.0.0.1:$MASTER_PORT" -local="$LOCAL_ADDR" \ + -protocol="$PROTOCOL" -nkeys="$NKEYS" >>"$CLIENT_LOG" 2>&1 & +CLIENT_PID=$! +log "started client pid=$CLIENT_PID" + +# --- 3. wait for baseline + READY_FOR_KILL --- +if ! wait_for_line "$CLIENT_LOG" "READY_FOR_KILL" 60; then + log "FATAL client never reached READY_FOR_KILL; client log:"; cat "$CLIENT_LOG" + exit 3 +fi +log "client is READY_FOR_KILL; baseline done" + +# --- 4. KILL the master hard --- +log "kill -9 master pid=$MASTER_PID" +kill -9 "$MASTER_PID" 2>/dev/null +wait "$MASTER_PID" 2>/dev/null +MASTER_PID="" +# give the client time to observe read failures (saw_down) +sleep 4 + +# --- 5. RESTART master on the same port --- +log "restart master on same port $MASTER_PORT" +start_master +sleep 2 +if ! kill -0 "$MASTER_PID" 2>/dev/null; then + log "FATAL master failed to restart; log:"; cat "$MASTER_LOG"; exit 3 +fi + +# --- 6. wait for client's verdict --- +if ! wait_for_line "$CLIENT_LOG" "RESULT=" 150; then + log "FATAL client never printed RESULT; client log tail:"; tail -30 "$CLIENT_LOG" + exit 3 +fi + +RESULT_LINE="$(grep -m1 "RESULT=" "$CLIENT_LOG")" +log "client verdict: $RESULT_LINE" +log "----- client log tail -----"; tail -20 "$CLIENT_LOG" + +if echo "$RESULT_LINE" | grep -q "RESULT=PASS"; then + log "OVERALL: PASS (out-of-process master kill+restart, metadata rebuilt)" + exit 0 +else + log "OVERALL: FAIL" + exit 1 +fi diff --git a/mooncake-store/tests/ha_rebuild_gate_test.cpp b/mooncake-store/tests/ha_rebuild_gate_test.cpp new file mode 100644 index 00000000..16bd0133 --- /dev/null +++ b/mooncake-store/tests/ha_rebuild_gate_test.cpp @@ -0,0 +1,180 @@ +#include + +#include "rpc_service.h" + +namespace mooncake { +namespace { + +WrappedMasterServiceConfig MakeConfig(bool initially_serving, + ViewVersionId view_version = 42) { + WrappedMasterServiceConfig config; + config.default_kv_lease_ttl = 100; + config.enable_metric_reporting = false; + config.initially_serving = initially_serving; + config.view_version = view_version; + return config; +} + +template +void ExpectUnavailable(const tl::expected& result) { + ASSERT_FALSE(result.has_value()); + EXPECT_EQ(result.error(), ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS); +} + +template +void ExpectUnavailable(const std::vector>& results) { + ASSERT_FALSE(results.empty()); + for (const auto& result : results) ExpectUnavailable(result); +} + +TEST(HaRebuildGateTest, ConstructorHonorsInitialState) { + WrappedMasterService closed(MakeConfig(false)); + EXPECT_FALSE(closed.IsServing()); + EXPECT_EQ(closed.GetServingState(), StoreServingState::REBUILDING); + + WrappedMasterService open(MakeConfig(true)); + EXPECT_TRUE(open.IsServing()); + EXPECT_EQ(open.GetServingState(), StoreServingState::SERVING); +} + +TEST(HaRebuildGateTest, ExactRosterAndEpochControlOpening) { + WrappedMasterService service(MakeConfig(false)); + const UUID first{1, 1}; + const UUID second{2, 2}; + const UUID outsider{3, 3}; + + auto stale = service.SignalRebuildComplete(first, 41); + ASSERT_FALSE(stale.has_value()); + EXPECT_EQ(stale.error(), ErrorCode::INVALID_VERSION); + ASSERT_TRUE(service.SignalRebuildComplete(first, 42).has_value()); + service.LockRebuildExpectedClients({first, second}); + EXPECT_FALSE(service.IsServing()); + + ASSERT_TRUE(service.SignalRebuildComplete(first, 42).has_value()); + ASSERT_TRUE(service.SignalRebuildComplete(outsider, 42).has_value()); + EXPECT_FALSE(service.IsServing()); + + ASSERT_TRUE(service.SignalRebuildComplete(second, 42).has_value()); + EXPECT_TRUE(service.IsServing()); + EXPECT_EQ(service.GetServingState(), StoreServingState::SERVING); +} + +TEST(HaRebuildGateTest, EmptyRosterOpensAndTimeoutIsDegraded) { + WrappedMasterService empty(MakeConfig(false)); + empty.LockRebuildExpectedClients({}); + EXPECT_EQ(empty.GetServingState(), StoreServingState::SERVING); + + WrappedMasterService timed_out(MakeConfig(false)); + timed_out.LockRebuildExpectedClients({UUID{1, 1}}); + timed_out.ForceServingAfterTimeout(); + EXPECT_TRUE(timed_out.IsServing()); + EXPECT_EQ(timed_out.GetServingState(), StoreServingState::DEGRADED); + ASSERT_TRUE(timed_out.SignalRebuildComplete(UUID{1, 1}, 42).has_value()); + EXPECT_EQ(timed_out.GetServingState(), StoreServingState::SERVING); + + WrappedMasterService opened(MakeConfig(false)); + opened.LockRebuildExpectedClients({}); + opened.SetServing(false); + EXPECT_EQ(opened.GetServingState(), StoreServingState::REBUILDING); +} + + +TEST(HaRebuildGateTest, ConcurrentDuplicateDoneOpensOnlyAfterExactRoster) { + WrappedMasterService service(MakeConfig(false)); + const UUID first{1, 1}; + const UUID second{2, 2}; + service.LockRebuildExpectedClients({first, second}); + + std::vector duplicates; + for (int i = 0; i < 16; ++i) { + duplicates.emplace_back([&] { + EXPECT_TRUE(service.SignalRebuildComplete(first, 42).has_value()); + }); + } + for (auto& thread : duplicates) thread.join(); + EXPECT_EQ(service.GetServingState(), StoreServingState::REBUILDING); + + ASSERT_TRUE(service.SignalRebuildComplete(second, 42).has_value()); + EXPECT_EQ(service.GetServingState(), StoreServingState::SERVING); +} + +TEST(HaRebuildGateTest, TimeoutAndFinalDoneConvergeSafely) { + for (int iteration = 0; iteration < 5; ++iteration) { + WrappedMasterService service(MakeConfig(false)); + const UUID client{1, 1}; + service.LockRebuildExpectedClients({client}); + std::thread timeout([&] { service.ForceServingAfterTimeout(); }); + std::thread done([&] { + EXPECT_TRUE(service.SignalRebuildComplete(client, 42).has_value()); + }); + timeout.join(); + done.join(); + EXPECT_EQ(service.GetServingState(), StoreServingState::SERVING); + } +} + +TEST(HaRebuildGateTest, RebuildingRejectsBusinessAndAllowsRecoveryRpc) { + WrappedMasterService service(MakeConfig(false)); + const UUID client{1, 1}; + ReplicateConfig config; + config.replica_num = 1; + + ExpectUnavailable(service.ExistKey("key")); + ExpectUnavailable(service.BatchExistKey({"key"})); + ExpectUnavailable(service.GetReplicaListByRegex(".*")); + ExpectUnavailable(service.GetReplicaList("key")); + ExpectUnavailable(service.BatchGetReplicaList({"key"})); + ExpectUnavailable(service.PutStart(client, "key", 8, config)); + ExpectUnavailable(service.PutEnd(client, "key")); + ExpectUnavailable(service.PutRevoke(client, "key")); + ExpectUnavailable(service.BatchPutStart(client, {"key"}, {8}, config)); + ExpectUnavailable(service.BatchPutEnd(client, {"key"})); + ExpectUnavailable(service.BatchPutRevoke(client, {"key"})); + ExpectUnavailable(service.UpsertStart(client, "key", 8, config)); + ExpectUnavailable(service.UpsertEnd(client, "key")); + ExpectUnavailable(service.UpsertRevoke(client, "key")); + ExpectUnavailable(service.BatchUpsertStart(client, {"key"}, {8}, config)); + ExpectUnavailable(service.BatchUpsertEnd(client, {"key"})); + ExpectUnavailable(service.BatchUpsertRevoke(client, {"key"})); + ExpectUnavailable(service.Remove("key")); + ExpectUnavailable(service.RemoveByRegex(".*")); + ExpectUnavailable(service.RemoveAll()); + ExpectUnavailable(service.BatchRemove({"key"})); + Segment segment; + segment.id = UUID{4, 4}; + segment.name = "blocked_mount"; + segment.size = 4096; + segment.base = 4096; + segment.te_endpoint = "127.0.0.1:1"; + ExpectUnavailable(service.MountSegment(segment, client)); + NoFSegment nof_segment; + ExpectUnavailable(service.MountNoFSegment(nof_segment, client)); + ExpectUnavailable(service.MountLocalDiskSegment(client, true)); + ExpectUnavailable(service.CreateCopyTask("key", "default", {})); + ExpectUnavailable(service.CreateMoveTask("key", "default", "a", "b")); + ExpectUnavailable(service.CopyStart(client, "key", "default", "a", {"b"})); + ExpectUnavailable(service.CopyEnd(client, "key", "default")); + ExpectUnavailable(service.CopyRevoke(client, "key", "default")); + ExpectUnavailable(service.MoveStart(client, "key", "default", "a", "b")); + ExpectUnavailable(service.MoveEnd(client, "key", "default")); + ExpectUnavailable(service.MoveRevoke(client, "key", "default")); + ExpectUnavailable(service.PromotionObjectHeartbeat(client)); + ExpectUnavailable(service.OffloadObjectHeartbeat(client, true)); + + ASSERT_TRUE(service.Ping(client).has_value()); + ASSERT_TRUE(service.ReMountSegment({}, client).has_value()); + auto nof_remount = service.ReMountNoFSegment({}, client); + EXPECT_NE(nof_remount.error(), ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS); + ASSERT_TRUE(service.RebuildMetadata({}, client, 42).has_value()); + ASSERT_TRUE(service.SignalRebuildCompleteRpc(client, 42).has_value()); +} + +TEST(HaRebuildGateTest, StaleMetadataBatchIsRejected) { + WrappedMasterService service(MakeConfig(false)); + auto result = service.RebuildMetadata({}, UUID{1, 1}, 41); + ASSERT_FALSE(result.has_value()); + EXPECT_EQ(result.error(), ErrorCode::INVALID_VERSION); +} + +} // namespace +} // namespace mooncake diff --git a/mooncake-store/tests/ha_recovery_live_main.cpp b/mooncake-store/tests/ha_recovery_live_main.cpp new file mode 100644 index 00000000..a68d31b8 --- /dev/null +++ b/mooncake-store/tests/ha_recovery_live_main.cpp @@ -0,0 +1,135 @@ +// Live HA recovery test: standalone client that connects to a REAL, separate +// mooncake_master PROCESS (not InProcMaster). Driven by an external shell +// script (ha_live_test.sh): +// 1. mount a segment, Put N keys, verify baseline +// 2. print "READY_FOR_KILL", then poll Get until reads fail (master killed) +// and then succeed again (master restarted + metadata rebuilt) +// 3. print RESULT=PASS/FAIL +#include +#include + +#include +#include +#include +#include +#include +#include + +#include "allocator.h" +#include "client_service.h" +#include "types.h" +#include "utils.h" + +DEFINE_string(protocol, "tcp", "transfer protocol"); +DEFINE_string(master, "127.0.0.1:50055", "master rpc ip:port"); +DEFINE_string(metadata, "", "http metadata server url (empty => P2PHANDSHAKE)"); +DEFINE_string(local, "127.0.0.1:19110", "local hostname ip:port"); +DEFINE_int32(nkeys, 50, "number of keys to put"); + +using namespace mooncake; + +static constexpr size_t kSeg = 128ull * 1024 * 1024; +static constexpr size_t kAlloc = 64ull * 1024 * 1024; + +static tl::expected PutStr(std::shared_ptr& c, + SimpleAllocator& a, + const std::string& k, + const std::string& v) { + void* buf = a.allocate(v.size()); + std::memcpy(buf, v.data(), v.size()); + std::vector s{Slice{buf, v.size()}}; + ReplicateConfig cfg; + cfg.replica_num = 1; + auto r = c->Put(k, s, cfg); + a.deallocate(buf, v.size()); + return r; +} + +static bool GetVerify(std::shared_ptr& c, SimpleAllocator& a, + const std::string& k, const std::string& exp) { + void* buf = a.allocate(exp.size()); + std::vector s{Slice{buf, exp.size()}}; + auto r = c->Get(k, s); + bool ok = r.has_value() && s[0].size == exp.size() && + std::memcmp(s[0].ptr, exp.data(), exp.size()) == 0; + a.deallocate(buf, exp.size()); + return ok; +} + +int main(int argc, char** argv) { + gflags::ParseCommandLineFlags(&argc, &argv, true); + google::InitGoogleLogging(argv[0]); + FLAGS_logtostderr = 1; + + std::vector keys, vals; + for (int i = 0; i < FLAGS_nkeys; ++i) { + keys.push_back("live_key_" + std::to_string(i)); + vals.push_back("live_value_payload_" + std::to_string(i) + + std::string(200, 'x')); + } + + const std::string meta = + FLAGS_metadata.empty() ? "P2PHANDSHAKE" : FLAGS_metadata; + auto co = Client::Create(FLAGS_local, meta, FLAGS_protocol, std::nullopt, + FLAGS_master); + if (!co.has_value()) { + LOG(ERROR) << "RESULT=FAIL reason=client_create_failed"; + return 2; + } + auto client = co.value(); + + auto alloc = std::make_unique(kAlloc); + auto reg = client->RegisterLocalMemory(alloc->getBase(), kAlloc, "cpu:0", + false, false); + if (!reg.has_value()) { + LOG(ERROR) << "RESULT=FAIL reason=register_local_memory_failed"; + return 2; + } + void* seg = allocate_buffer_allocator_memory(kSeg); + auto mnt = client->MountSegment(seg, kSeg, FLAGS_protocol); + if (!mnt.has_value()) { + LOG(ERROR) << "RESULT=FAIL reason=mount_failed"; + return 2; + } + + for (int i = 0; i < FLAGS_nkeys; ++i) { + auto r = PutStr(client, *alloc, keys[i], vals[i]); + if (!r.has_value()) { + LOG(ERROR) << "RESULT=FAIL reason=put_failed key=" << keys[i]; + return 2; + } + } + int base_ok = 0; + for (int i = 0; i < FLAGS_nkeys; ++i) + if (GetVerify(client, *alloc, keys[i], vals[i])) ++base_ok; + LOG(INFO) << "BASELINE ok=" << base_ok << "/" << FLAGS_nkeys; + if (base_ok != FLAGS_nkeys) { + LOG(ERROR) << "RESULT=FAIL reason=baseline_incomplete"; + return 2; + } + + LOG(INFO) << "READY_FOR_KILL"; + fflush(stderr); + + bool saw_down = false; + int final_ok = -1; + for (int attempt = 0; attempt < 240; ++attempt) { + int ok = 0; + for (int i = 0; i < FLAGS_nkeys; ++i) + if (GetVerify(client, *alloc, keys[i], vals[i])) ++ok; + if (ok < FLAGS_nkeys) saw_down = true; + if (saw_down && ok == FLAGS_nkeys) { + final_ok = ok; + break; + } + std::this_thread::sleep_for(std::chrono::milliseconds(500)); + } + + if (final_ok == FLAGS_nkeys) { + LOG(INFO) << "RESULT=PASS recovered=" << final_ok << "/" << FLAGS_nkeys + << " zero-recompute-after-master-kill-restart"; + return 0; + } + LOG(ERROR) << "RESULT=FAIL reason=not_recovered saw_down=" << saw_down; + return 1; +} diff --git a/mooncake-store/tests/ha_scale_bench.sh b/mooncake-store/tests/ha_scale_bench.sh new file mode 100755 index 00000000..5a7e094f --- /dev/null +++ b/mooncake-store/tests/ha_scale_bench.sh @@ -0,0 +1,178 @@ +#!/usr/bin/env bash +# HA scale benchmark driver. For each --nkeys in the sweep: +# 1. start a fresh mooncake_master (non-HA) on $MASTER_PORT +# 2. start ha_scale_bench_main -> mounts a segment, BatchPuts N keys, +# records fill throughput + steady Get latency, prints READY_FOR_KILL, +# then polls a sampled probe every --poll_ms +# 3. on READY_FOR_KILL: kill -9 the master (probe reads start failing) +# 4. restart master on the SAME port -> client resends its local table, +# master rebuilds; client measures rebuild latency + outage-window stats +# 5. scrape the JSON_RESULT line into a summary table +# +# Each nkeys value runs on a fresh master (clean state). Results -> $OUT_TSV. +set -u + +BUILD_DIR="${BUILD_DIR:-/export/home/shenshuwei.3/Mooncake/build}" +MASTER_BIN="${MASTER_BIN:-$BUILD_DIR/mooncake-store/src/mooncake_master}" +CLIENT_BIN="${CLIENT_BIN:-$BUILD_DIR/mooncake-store/tests/ha_scale_bench_main}" + +# Per-run ports are derived from these bases + RUN_IDX*10 (see run_one), so +# consecutive runs never reuse a socket still in TIME_WAIT. +BASE_MASTER_PORT="${MASTER_PORT:-50057}" +BASE_METRICS_PORT="${METRICS_PORT:-9013}" +BASE_LOCAL_PORT="${LOCAL_PORT:-19120}" +MASTER_PORT="" # set per-run in run_one +METRICS_PORT="" +LOCAL_ADDR="" +PROTOCOL="${PROTOCOL:-tcp}" +VSIZE="${VSIZE:-100}" +BATCH="${BATCH:-2000}" +PROBE="${PROBE:-1000}" +POLL_MS="${POLL_MS:-20}" +SEG_MB="${SEG_MB:-0}" # 0 => client auto-sizes from nkeys*vsize +ALLOC_MB="${ALLOC_MB:-0}" # 0 => client default (64MB) +DOWN_WAIT="${DOWN_WAIT:-4}" # seconds to let client observe the outage +# nkeys sweep (override by passing args: ha_scale_bench.sh 1000 10000 ...) +NKEYS_SWEEP=("$@") +if [ "${#NKEYS_SWEEP[@]}" -eq 0 ]; then + NKEYS_SWEEP=(1000 10000 100000 1000000) +fi + +OUT_DIR="${OUT_DIR:-$(mktemp -d /tmp/ha_scale_bench.XXXXXX)}" +mkdir -p "$OUT_DIR" +OUT_TSV="$OUT_DIR/results.tsv" +MASTER_PID="" +CLIENT_PID="" + +log() { echo "[ha_scale_bench] $*"; } + +cleanup() { + [ -n "$CLIENT_PID" ] && kill -9 "$CLIENT_PID" 2>/dev/null + [ -n "$MASTER_PID" ] && kill -9 "$MASTER_PID" 2>/dev/null + wait 2>/dev/null +} +trap cleanup EXIT + +start_master() { + # $1=logfile. Uses per-run $MASTER_PORT and $METRICS_PORT (set by run_one) + # so consecutive runs don't collide on a socket still in TIME_WAIT. + "$MASTER_BIN" -rpc_port="$MASTER_PORT" -metrics_port="$METRICS_PORT" \ + -enable_ha=false -enable_metric_reporting=false >>"$1" 2>&1 & + MASTER_PID=$! +} + +wait_for_line() { # $1=file $2=pattern $3=timeout_sec + local f="$1" pat="$2" t="$3" i=0 + while [ "$i" -lt "$((t * 2))" ]; do + grep -q "$pat" "$f" 2>/dev/null && return 0 + sleep 0.5 + i=$((i + 1)) + done + return 1 +} + +[ -x "$MASTER_BIN" ] || { log "FATAL master bin missing: $MASTER_BIN"; exit 3; } +[ -x "$CLIENT_BIN" ] || { log "FATAL client bin missing: $CLIENT_BIN"; exit 3; } + +log "out_dir=$OUT_DIR sweep=${NKEYS_SWEEP[*]}" +printf 'nkeys\tvsize\tfill_keys_per_s\trebuild_ms\trecover_since_ready_ms\tsteady_get_lat_ms\toutage_get_lat_ms\tmin_avail_pct\n' >"$OUT_TSV" + +run_one() { + local NKEYS="$1" + local tag="n${NKEYS}" + local MLOG="$OUT_DIR/master_${tag}.log" + local CLOG="$OUT_DIR/client_${tag}.log" + : >"$MLOG"; : >"$CLOG" + MASTER_PID=""; CLIENT_PID="" + + # Unique ports per run so a socket left in TIME_WAIT by the previous run + # can't stop this run's master from binding. RUN_IDX is bumped by the caller. + MASTER_PORT=$(( BASE_MASTER_PORT + RUN_IDX * 10 )) + METRICS_PORT=$(( BASE_METRICS_PORT + RUN_IDX * 10 )) + local LOCAL_PORT=$(( BASE_LOCAL_PORT + RUN_IDX * 10 )) + LOCAL_ADDR="127.0.0.1:${LOCAL_PORT}" + + log "=== nkeys=$NKEYS : start master (rpc=$MASTER_PORT metrics=$METRICS_PORT local=$LOCAL_ADDR) ===" + start_master "$MLOG" + sleep 2 + if ! kill -0 "$MASTER_PID" 2>/dev/null; then + log "FATAL master died on startup (nkeys=$NKEYS); log:"; tail -20 "$MLOG"; return 1 + fi + + # client: fill + baseline + wait-for-kill + local FILL_TIMEOUT=$(( 120 + NKEYS / 5000 )) # scale fill wait with size + "$CLIENT_BIN" -master="127.0.0.1:$MASTER_PORT" -local="$LOCAL_ADDR" \ + -protocol="$PROTOCOL" -nkeys="$NKEYS" -vsize="$VSIZE" -batch="$BATCH" \ + -probe="$PROBE" -poll_ms="$POLL_MS" -seg_mb="$SEG_MB" -alloc_mb="$ALLOC_MB" \ + >>"$CLOG" 2>&1 & + CLIENT_PID=$! + + if ! wait_for_line "$CLOG" "READY_FOR_KILL" "$FILL_TIMEOUT"; then + log "FATAL client never reached READY_FOR_KILL (nkeys=$NKEYS); tail:"; tail -25 "$CLOG" + kill -9 "$CLIENT_PID" 2>/dev/null; kill -9 "$MASTER_PID" 2>/dev/null + return 1 + fi + grep -m1 "FILL done" "$CLOG" | sed 's/^/[ha_scale_bench] /' + log " nkeys=$NKEYS: baseline done, killing master pid=$MASTER_PID" + + kill -9 "$MASTER_PID" 2>/dev/null + wait "$MASTER_PID" 2>/dev/null + MASTER_PID="" + sleep "$DOWN_WAIT" + + log " nkeys=$NKEYS: restart master on same port" + start_master "$MLOG" + sleep 2 + if ! kill -0 "$MASTER_PID" 2>/dev/null; then + log "FATAL master failed to restart (nkeys=$NKEYS); log:"; tail -20 "$MLOG"; return 1 + fi + + local REC_TIMEOUT=$(( 120 + NKEYS / 2000 )) + if ! wait_for_line "$CLOG" "JSON_RESULT=" "$REC_TIMEOUT"; then + log "FATAL client never printed JSON_RESULT (nkeys=$NKEYS); tail:"; tail -25 "$CLOG" + kill -9 "$CLIENT_PID" 2>/dev/null; kill -9 "$MASTER_PID" 2>/dev/null + return 1 + fi + + local JLINE + JLINE="$(grep -m1 "JSON_RESULT=" "$CLOG" | sed 's/.*JSON_RESULT=//')" + log " nkeys=$NKEYS JSON: $JLINE" + + # parse with python for robustness, append a TSV row + python3 - "$JLINE" >>"$OUT_TSV" <<'PY' +import sys, json +d = json.loads(sys.argv[1]) +print("%d\t%d\t%.0f\t%.1f\t%.1f\t%.3f\t%.3f\t%.2f" % ( + d["nkeys"], d["vsize"], d["fill_keys_per_s"], d["rebuild_ms"], + d["recover_since_ready_ms"], d["steady_get_lat_ms"], + d["outage_get_lat_ms"], d["min_avail_pct"])) +PY + + kill -9 "$CLIENT_PID" 2>/dev/null + wait "$CLIENT_PID" 2>/dev/null + CLIENT_PID="" + # IMPORTANT: kill this run's (restarted) master too, else it lingers holding + # its rpc/metrics ports and the next run that reuses a nearby port fails. + kill -9 "$MASTER_PID" 2>/dev/null + wait "$MASTER_PID" 2>/dev/null + MASTER_PID="" + log " nkeys=$NKEYS: DONE" + return 0 +} + +overall=0 +RUN_IDX=0 +for nk in "${NKEYS_SWEEP[@]}"; do + if ! run_one "$nk"; then + log "nkeys=$nk FAILED" + overall=1 + fi + RUN_IDX=$(( RUN_IDX + 1 )) + sleep 2 +done + +log "================ SUMMARY ================" +column -t -s $'\t' "$OUT_TSV" | sed 's/^/[ha_scale_bench] /' +log "TSV: $OUT_TSV" +log "logs: $OUT_DIR" +exit "$overall" diff --git a/mooncake-store/tests/ha_scale_bench_main.cpp b/mooncake-store/tests/ha_scale_bench_main.cpp new file mode 100644 index 00000000..adfe1bf0 --- /dev/null +++ b/mooncake-store/tests/ha_scale_bench_main.cpp @@ -0,0 +1,289 @@ +// ============================================================================= +// HA scale benchmark: standalone client against a REAL separate mooncake_master +// PROCESS, driven by ha_scale_bench.sh. Extends ha_recovery_live_main.cpp with +// timing + scale + recovery-window latency/failure statistics. +// +// Measures three metrics: +// [1] rebuild latency : t_recovered - t_first_read_fail (and since restart) +// [2] scale : same run repeated over the --nkeys sweep (via .sh) +// [3] recovery-window : per-poll sampled success-rate + Get latency during +// the outage->recovery window, vs steady-state Get lat +// +// Key design choices (why, so a reviewer can trust the numbers): +// - Fill with BatchPut (--batch keys per RPC): filling millions of keys one +// Put at a time is dominated by per-RPC overhead, not the thing we measure. +// - Completion is judged by a UNIFORM SAMPLE of --probe keys, not by Getting +// all N. Getting millions of keys each poll would itself take seconds and +// pollute the latency measurement. Sampling every key-space region detects +// partial rebuild while staying cheap. +// - steady_clock everywhere; recovery poll interval is --poll_ms (default 20) +// so rebuild-latency resolution is +/-poll_ms, not the old 500ms. +// - Emits a machine-readable JSON line (JSON_RESULT={...}) the .sh collects. +// ============================================================================= +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "allocator.h" +#include "client_service.h" +#include "types.h" +#include "utils.h" + +DEFINE_string(protocol, "tcp", "transfer protocol"); +DEFINE_string(master, "127.0.0.1:50055", "master rpc ip:port"); +DEFINE_string(metadata, "", "http metadata server url (empty => P2PHANDSHAKE)"); +DEFINE_string(local, "127.0.0.1:19110", "local hostname ip:port"); +DEFINE_int64(nkeys, 50, "number of keys to put"); +DEFINE_int32(vsize, 100, "value size in bytes per key"); +DEFINE_int32(batch, 2000, "keys per BatchPut RPC while filling"); +DEFINE_int32(probe, 1000, + "number of uniformly-sampled keys used as the " + "rebuild-completion probe (<=nkeys)"); +DEFINE_int32(poll_ms, 20, "recovery poll interval in ms"); +DEFINE_int32(max_recovery_sec, 600, "give up waiting for recovery after this"); +DEFINE_int64(seg_mb, 0, "segment size in MB (0 => auto from nkeys*vsize)"); +DEFINE_int64(alloc_mb, 0, "local buffer size in MB (0 => auto)"); + +using namespace mooncake; +using Clock = std::chrono::steady_clock; + +static double ms_since(Clock::time_point a, Clock::time_point b) { + return std::chrono::duration(b - a).count(); +} + +// Build the value payload for key i deterministically (so Get can verify). +static std::string MakeValue(int64_t i, int vsize) { + std::string v = "v" + std::to_string(i) + "_"; + if (static_cast(v.size()) >= vsize) { + v.resize(vsize); + } else { + v.append(vsize - v.size(), static_cast('a' + (i % 26))); + } + return v; +} + +static std::string MakeKey(int64_t i) { + return "scale_key_" + std::to_string(i); +} + +// Get one key and byte-verify against expected. Returns {ok, latency_ms}. +static std::pair GetVerify(std::shared_ptr& c, + SimpleAllocator& a, + const std::string& k, + const std::string& exp) { + void* buf = a.allocate(exp.size()); + std::vector s{Slice{buf, exp.size()}}; + auto t0 = Clock::now(); + auto r = c->Get(k, s); + auto t1 = Clock::now(); + bool ok = r.has_value() && s[0].size == exp.size() && + std::memcmp(s[0].ptr, exp.data(), exp.size()) == 0; + a.deallocate(buf, exp.size()); + return {ok, ms_since(t0, t1)}; +} + +// Probe the sampled key set once. Returns {num_ok, avg_latency_ms_over_ok}. +static std::pair ProbeOnce(std::shared_ptr& c, + SimpleAllocator& a, + const std::vector& idx, + int vsize) { + int ok = 0; + double sum = 0; + for (int64_t i : idx) { + auto [good, lat] = GetVerify(c, a, MakeKey(i), MakeValue(i, vsize)); + if (good) { + ++ok; + sum += lat; + } + } + return {ok, ok ? sum / ok : 0.0}; +} + +int main(int argc, char** argv) { + gflags::ParseCommandLineFlags(&argc, &argv, true); + google::InitGoogleLogging(argv[0]); + FLAGS_logtostderr = 1; + + const int64_t N = FLAGS_nkeys; + const int vsize = FLAGS_vsize; + const int probe_n = std::min(FLAGS_probe, N); + + // Auto-size segment/buffer. Per-object footprint in the segment is far + // larger than the value payload: object metadata + allocator min-unit / + // alignment / fragmentation. Measured ~1073 B/object at vsize=100, so we + // budget max(vsize+1200, 2*vsize) bytes per key plus a floor and headroom. + int64_t per_obj = std::max(vsize + 1200, vsize * 2); + int64_t data_mb = (N * per_obj) / (1024 * 1024) + 1; + int64_t seg_mb = + FLAGS_seg_mb ? FLAGS_seg_mb : std::max(128, data_mb * 3 / 2); + int64_t alloc_mb = FLAGS_alloc_mb ? FLAGS_alloc_mb : 64; + const size_t kSeg = static_cast(seg_mb) * 1024 * 1024; + const size_t kAlloc = static_cast(alloc_mb) * 1024 * 1024; + + LOG(INFO) << "CONFIG nkeys=" << N << " vsize=" << vsize + << " batch=" << FLAGS_batch << " probe=" << probe_n + << " poll_ms=" << FLAGS_poll_ms << " seg_mb=" << seg_mb + << " alloc_mb=" << alloc_mb; + + const std::string meta = + FLAGS_metadata.empty() ? "P2PHANDSHAKE" : FLAGS_metadata; + auto co = Client::Create(FLAGS_local, meta, FLAGS_protocol, std::nullopt, + FLAGS_master); + if (!co.has_value()) { + LOG(ERROR) << "RESULT=FAIL reason=client_create_failed"; + return 2; + } + auto client = co.value(); + + auto alloc = std::make_unique(kAlloc); + auto reg = client->RegisterLocalMemory(alloc->getBase(), kAlloc, "cpu:0", + false, false); + if (!reg.has_value()) { + LOG(ERROR) << "RESULT=FAIL reason=register_local_memory_failed"; + return 2; + } + void* seg = allocate_buffer_allocator_memory(kSeg); + if (!seg) { + LOG(ERROR) << "RESULT=FAIL reason=segment_alloc_failed seg_mb=" + << seg_mb; + return 2; + } + auto mnt = client->MountSegment(seg, kSeg, FLAGS_protocol); + if (!mnt.has_value()) { + LOG(ERROR) << "RESULT=FAIL reason=mount_failed"; + return 2; + } + + // --- fill N keys via BatchPut --- + // A dedicated fill buffer, reused per batch (separate from the Get path). + auto fill_alloc = std::make_unique(std::max( + kAlloc, static_cast(FLAGS_batch) * vsize + (1 << 20))); + auto t_fill0 = Clock::now(); + int64_t filled = 0; + ReplicateConfig cfg; + cfg.replica_num = 1; + for (int64_t base = 0; base < N; base += FLAGS_batch) { + int64_t cnt = std::min(FLAGS_batch, N - base); + std::vector keys; + std::vector> slices; + std::vector vals; + keys.reserve(cnt); + slices.reserve(cnt); + vals.reserve(cnt); + for (int64_t j = 0; j < cnt; ++j) { + int64_t i = base + j; + keys.push_back(MakeKey(i)); + vals.push_back(MakeValue(i, vsize)); + } + for (int64_t j = 0; j < cnt; ++j) { + void* b = fill_alloc->allocate(vals[j].size()); + std::memcpy(b, vals[j].data(), vals[j].size()); + slices.push_back({Slice{b, vals[j].size()}}); + } + auto rs = client->BatchPut(keys, slices, cfg); + for (auto& s : slices) fill_alloc->deallocate(s[0].ptr, s[0].size); + for (size_t j = 0; j < rs.size(); ++j) { + if (!rs[j].has_value()) { + LOG(ERROR) << "RESULT=FAIL reason=batchput_failed key=" + << keys[j] << " err=" << toString(rs[j].error()); + return 2; + } + } + filled += cnt; + if (base / FLAGS_batch % 50 == 0) + LOG(INFO) << "FILL progress " << filled << "/" << N; + } + double fill_ms = ms_since(t_fill0, Clock::now()); + LOG(INFO) << "FILL done " << filled << "/" << N << " in " << fill_ms + << " ms (" << (filled / (fill_ms / 1000.0)) << " keys/s)"; + + // --- build the uniform probe sample --- + std::vector probe_idx; + probe_idx.reserve(probe_n); + for (int p = 0; p < probe_n; ++p) { + // even spacing across [0, N) + probe_idx.push_back( + static_cast((static_cast(p) + 0.5) * N / probe_n)); + } + + // --- baseline: probe must be fully readable, and record steady latency --- + auto [base_ok, steady_lat] = ProbeOnce(client, *alloc, probe_idx, vsize); + LOG(INFO) << "BASELINE probe_ok=" << base_ok << "/" << probe_n + << " steady_get_lat_ms=" << steady_lat; + if (base_ok != probe_n) { + LOG(ERROR) << "RESULT=FAIL reason=baseline_incomplete"; + return 2; + } + + LOG(INFO) << "READY_FOR_KILL"; + fflush(stderr); + + // --- recovery window: poll the probe, capture timing + per-poll stats --- + auto t_ready = Clock::now(); + bool saw_down = false; + Clock::time_point t_first_fail, t_recovered; + int min_ok = probe_n; // worst observed availability + double outage_lat_sum = 0; // avg latency of successful Gets during outage + int outage_lat_polls = 0; + int polls = 0; + + int max_polls = (FLAGS_max_recovery_sec * 1000) / FLAGS_poll_ms; + for (int attempt = 0; attempt < max_polls; ++attempt) { + auto [ok, lat] = ProbeOnce(client, *alloc, probe_idx, vsize); + ++polls; + if (ok < probe_n) { + if (!saw_down) { + saw_down = true; + t_first_fail = Clock::now(); + } + min_ok = std::min(min_ok, ok); + if (ok > 0) { + outage_lat_sum += lat; + ++outage_lat_polls; + } + } + if (saw_down && ok == probe_n) { + t_recovered = Clock::now(); + break; + } + std::this_thread::sleep_for(std::chrono::milliseconds(FLAGS_poll_ms)); + } + + if (!saw_down || t_recovered.time_since_epoch().count() == 0) { + LOG(ERROR) << "RESULT=FAIL reason=not_recovered saw_down=" << saw_down + << " polls=" << polls; + return 1; + } + + double rebuild_ms = ms_since(t_first_fail, t_recovered); + double since_ready_ms = ms_since(t_ready, t_recovered); + double outage_lat = + outage_lat_polls ? outage_lat_sum / outage_lat_polls : 0.0; + double min_avail_pct = 100.0 * min_ok / probe_n; + + LOG(INFO) << "RESULT=PASS recovered=" << probe_n << "/" << probe_n + << " zero-recompute"; + // Machine-readable line for the shell to scrape. + LOG(INFO) << "JSON_RESULT={" + << "\"nkeys\":" << N << ",\"vsize\":" << vsize + << ",\"probe\":" << probe_n << ",\"fill_ms\":" << fill_ms + << ",\"fill_keys_per_s\":" << (filled / (fill_ms / 1000.0)) + << ",\"rebuild_ms\":" << rebuild_ms + << ",\"recover_since_ready_ms\":" << since_ready_ms + << ",\"steady_get_lat_ms\":" << steady_lat + << ",\"outage_get_lat_ms\":" << outage_lat + << ",\"min_avail_pct\":" << min_avail_pct + << ",\"poll_ms\":" << FLAGS_poll_ms << "}"; + fflush(stderr); + + client->UnmountSegment(seg, kSeg); + return 0; +} diff --git a/mooncake-store/tests/rebuild_retry_test.cpp b/mooncake-store/tests/rebuild_retry_test.cpp new file mode 100644 index 00000000..112290ad --- /dev/null +++ b/mooncake-store/tests/rebuild_retry_test.cpp @@ -0,0 +1,111 @@ +#include + +#include +#include +#include + +#include "rebuild_retry.h" + +namespace mooncake { +namespace { + +struct ScriptedRecovery { + std::array failures{}; + std::array calls{}; + std::vector order; + + tl::expected RunStep(int step) { + ++calls[step]; + order.push_back(step); + if (failures[step]-- > 0) { + return tl::make_unexpected(ErrorCode::RPC_FAIL); + } + return {}; + } + + RebuildRecoveryOps Ops() { + return { + [&] { return RunStep(0); }, [&] { return RunStep(1); }, + [&] { return RunStep(2); }, [&] { return RunStep(3); }, + [&] { return RunStep(4); }, + }; + } +}; + +TEST(RebuildRetryTest, ExecutesAllStepsInOrderOnceOnSuccess) { + ScriptedRecovery scripted; + auto result = RunRebuildRecovery(scripted.Ops(), 3, + [](int) { return true; }); + ASSERT_TRUE(result.has_value()); + EXPECT_EQ(scripted.order, (std::vector{0, 1, 2, 3, 4})); +} + +class RebuildFailureStepTest : public testing::TestWithParam {}; + +TEST_P(RebuildFailureStepTest, NeverSignalsCompleteAfterExhaustedStepFailure) { + ScriptedRecovery scripted; + scripted.failures[GetParam()] = 3; + int waits = 0; + auto result = RunRebuildRecovery(scripted.Ops(), 3, + [&](int) { ++waits; return true; }); + ASSERT_FALSE(result.has_value()); + EXPECT_EQ(scripted.calls[GetParam()], 3); + EXPECT_EQ(scripted.calls[4], GetParam() == 4 ? 3 : 0); + EXPECT_EQ(waits, 2); + for (int step = GetParam() + 1; step < 5; ++step) { + if (step != 4) { + EXPECT_EQ(scripted.calls[step], 0); + } + } +} + +INSTANTIATE_TEST_SUITE_P(AllRecoverySteps, RebuildFailureStepTest, + testing::Values(0, 1, 2, 3, 4)); + +TEST(RebuildRetryTest, TransientBatchFailureRestartsWholeTransaction) { + ScriptedRecovery scripted; + scripted.failures[3] = 1; + auto result = RunRebuildRecovery(scripted.Ops(), 3, + [](int) { return true; }); + ASSERT_TRUE(result.has_value()); + EXPECT_EQ(scripted.calls, (std::array{2, 2, 2, 2, 1})); +} + +TEST(RebuildRetryTest, StaleEpochStopsWithoutRetryOrDone) { + RebuildRecoveryOps ops; + int metadata_calls = 0; + int done_calls = 0; + ops.remount = [] { return tl::expected{}; }; + ops.publish_segment_descriptor = [] { + return tl::expected{}; + }; + ops.publish_rpc_metadata = [] { return tl::expected{}; }; + ops.resend_metadata = [&]() -> tl::expected { + ++metadata_calls; + return tl::make_unexpected(ErrorCode::INVALID_VERSION); + }; + ops.signal_complete = [&]() -> tl::expected { + ++done_calls; + return {}; + }; + + auto result = RunRebuildRecovery(ops, 3, [](int) { return true; }); + ASSERT_FALSE(result.has_value()); + EXPECT_EQ(result.error(), ErrorCode::INVALID_VERSION); + EXPECT_EQ(metadata_calls, 1); + EXPECT_EQ(done_calls, 0); +} + +TEST(RebuildRetryTest, ShutdownInterruptsRetries) { + ScriptedRecovery scripted; + scripted.failures[0] = 3; + int waits = 0; + auto result = RunRebuildRecovery(scripted.Ops(), 3, + [&](int) { ++waits; return false; }); + ASSERT_FALSE(result.has_value()); + EXPECT_EQ(scripted.calls[0], 1); + EXPECT_EQ(waits, 1); +} + +} // namespace +} // namespace mooncake