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_service.h b/mooncake-store/include/client_service.h index 621b4edd..c273f646 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,49 @@ 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. + void ResendLocalReplicaTable(); + // 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 +860,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> diff --git a/mooncake-store/include/master_client.h b/mooncake-store/include/master_client.h index 7b541230..10d246eb 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,13 @@ 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); + /** * @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_metric_manager.h b/mooncake-store/include/master_metric_manager.h index 293c6206..7d9cd86d 100644 --- a/mooncake-store/include/master_metric_manager.h +++ b/mooncake-store/include/master_metric_manager.h @@ -166,6 +166,8 @@ 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 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 +592,8 @@ 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::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..4d11e263 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,17 @@ class MasterService { auto ReMountSegment(const std::vector& segments, const UUID& client_id) -> tl::expected; + /** + * @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 +822,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(); @@ -1401,6 +1422,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_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..08598eb4 100644 --- a/mooncake-store/include/rpc_service.h +++ b/mooncake-store/include/rpc_service.h @@ -160,6 +160,9 @@ class WrappedMasterService { tl::expected ReMountSegment( const std::vector& segments, const UUID& client_id); + tl::expected RebuildMetadata( + const std::vector& entries, const UUID& client_id); + tl::expected ReMountNoFSegment( const std::vector& segments, const UUID& client_id); 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/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_service.cpp b/mooncake-store/src/client_service.cpp index c4b74fd0..065d2a94 100644 --- a/mooncake-store/src/client_service.cpp +++ b/mooncake-store/src/client_service.cpp @@ -4,6 +4,7 @@ #include "allocator.h" #include "segment.h" +#include "utils/base64.h" #include #include @@ -315,6 +316,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 +984,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,9 +1649,223 @@ 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. +void Client::ResendLocalReplicaTable() { + std::vector snapshot; + { + std::lock_guard lk(local_replica_table_mutex_); + snapshot.reserve(local_replica_table_.size()); + for (auto& [k, m] : local_replica_table_) { + KeyReplicaEntry e; + e.key = k; + e.tenant_id = m.tenant_id; + e.size = m.size; + e.data_type = m.data_type; + e.group_id = m.group_id; + e.replicas = {m.replica}; + snapshot.emplace_back(std::move(e)); + } + } + if (snapshot.empty()) return; + const size_t kBatch = 256; + for (size_t i = 0; i < snapshot.size(); i += kBatch) { + std::vector batch( + snapshot.begin() + i, + snapshot.begin() + std::min(i + kBatch, snapshot.size())); + auto r = master_client_.RebuildMetadata(std::move(batch)); + if (!r) + LOG(ERROR) << "RebuildMetadata resend failed: " + << toString(r.error()); + } +} + tl::expected Client::Upsert(const ObjectKey& key, std::vector& slices, const ReplicateConfig& config) { @@ -1790,6 +2018,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 +2163,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 +2554,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 +3974,71 @@ 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); + { + // 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); } - // 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); + auto remount_result = master_client_.ReMountSegment(segments); + if (!remount_result) { + ErrorCode err = remount_result.error(); + LOG(ERROR) << "Failed to remount segments: " << err; } - } - // 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. + // 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); + } + // 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); + } + } + // 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. + } // release mounted_segments_mutex_ before the (potentially many) + // rebuild RPCs + + // === HA rebuild: after segments are re-mounted (and descriptors + // re-published above), resend object-level metadata so the empty new + // master rebuilds key->location. "Segment before key" is satisfied + // because ReMountSegment ran above. Done OUTSIDE + // mounted_segments_mutex_ so the N batched RebuildMetadata RPCs don't + // block Put/Get that need that lock (ResendLocalReplicaTable takes only + // local_replica_table_mutex_). + ResendLocalReplicaTable(); }; // Use another thread to remount segments to avoid blocking the ping // thread diff --git a/mooncake-store/src/master_client.cpp b/mooncake-store/src/master_client.cpp index e2d9db34..b53656bd 100644 --- a/mooncake-store/src/master_client.cpp +++ b/mooncake-store/src/master_client.cpp @@ -157,6 +157,11 @@ 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::ReMountNoFSegment> { static constexpr const char* value = "ReMountNoFSegment"; @@ -814,6 +819,18 @@ tl::expected MasterClient::ReMountSegment( return result; } +tl::expected MasterClient::RebuildMetadata( + std::vector&& entries) { + 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_); + 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..9d5a4356 100644 --- a/mooncake-store/src/master_metric_manager.cpp +++ b/mooncake-store/src/master_metric_manager.cpp @@ -126,6 +126,12 @@ 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"), mount_nof_segment_requests_( "master_mount_nof_segment_requests_total", "Total number of MountNoFSegment requests received"), @@ -996,6 +1002,12 @@ 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::inc_remount_nof_segment_requests(int64_t val) { remount_nof_segment_requests_.inc(val); } diff --git a/mooncake-store/src/master_service.cpp b/mooncake-store/src/master_service.cpp index e899f21a..a1efb4ca 100644 --- a/mooncake-store/src/master_service.cpp +++ b/mooncake-store/src/master_service.cpp @@ -900,6 +900,143 @@ 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()) continue; // skip this key, keep the rest + + // (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/rpc_service.cpp b/mooncake-store/src/rpc_service.cpp index d5de362f..851119b0 100644 --- a/mooncake-store/src/rpc_service.cpp +++ b/mooncake-store/src/rpc_service.cpp @@ -733,6 +733,21 @@ tl::expected WrappedMasterService::ReMountSegment( [] { MasterMetricManager::instance().inc_remount_segment_failures(); }); } +tl::expected WrappedMasterService::RebuildMetadata( + const std::vector& entries, const UUID& client_id) { + 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(); + }); +} + tl::expected WrappedMasterService::ReMountNoFSegment( const std::vector& segments, const UUID& client_id) { return execute_rpc( @@ -1321,6 +1336,8 @@ 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::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..ddff9547 100644 --- a/mooncake-store/tests/CMakeLists.txt +++ b/mooncake-store/tests/CMakeLists.txt @@ -115,6 +115,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_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; +}