diff --git a/mooncake-store/benchmarks/stress_cluster_bench.cpp b/mooncake-store/benchmarks/stress_cluster_bench.cpp index 299b560182..b4f0ed674f 100644 --- a/mooncake-store/benchmarks/stress_cluster_bench.cpp +++ b/mooncake-store/benchmarks/stress_cluster_bench.cpp @@ -21,7 +21,9 @@ #include "gflags/gflags.h" #include "glog/logging.h" #include "mooncake_logging.h" +#include "dummy_client.h" #include "real_client.h" +#include "shm_helper.h" #include #include @@ -173,6 +175,27 @@ static std::vector DiscoverSegmentsFromMaster( return segments; } + +// Resolve the effective list of peer RPC addresses for client_rpc_bench. +// --peer_rpc_addrs (multi-peer, comma-separated) takes precedence over +// --peer_rpc_addr (single-peer fallback). +std::vector ParsePeerRpcAddrs() { + std::vector addrs; + if (!FLAGS_peer_rpc_addrs.empty()) { + std::istringstream iss(FLAGS_peer_rpc_addrs); + std::string a; + while (std::getline(iss, a, ',')) { + size_t s = a.find_first_not_of(" \t"); + size_t e = a.find_last_not_of(" \t"); + if (s != std::string::npos && e != std::string::npos) { + addrs.push_back(a.substr(s, e - s + 1)); + } + } + } else if (!FLAGS_peer_rpc_addr.empty()) { + addrs.push_back(FLAGS_peer_rpc_addr); + } + return addrs; +} } // namespace DEFINE_string(local_hostname, "localhost", @@ -189,9 +212,65 @@ DEFINE_string(ssd_offload_path, "", "SSD offload directory path"); DEFINE_string(scenario, "local_memory", "Benchmark scenario: local_memory, remote_memory, local_disk, " - "remote_disk, segment_write, segment_read"); + "remote_disk, segment_write, segment_read, client_rpc_bench, " + "list_segments"); +// Peer RPC server address for client_rpc_bench. Must be the +// "local_rpc_addr" (ip:port) of the RealClient that hosts the +// offload_rpc_server_ registered with batch_get_offload_object / +// release_offload_buffer. The peer must be started with +// --enable_ssd_offload=true so the offload_rpc_server_ is up. +DEFINE_string(peer_rpc_addr, "", + "[client_rpc_bench] Address (ip:port) of the peer RealClient's " + "offload_rpc_server_, e.g. \"10.0.0.2:17888\". The peer is " + "expected to have --enable_ssd_offload=true. Single-peer " + "fallback; superseded by --peer_rpc_addrs when that flag is " + "set."); +// Comma-separated list of peer RPC addresses (ip:port). When non-empty, +// takes precedence over --peer_rpc_addr. In each round of the timed +// loop the bench issues one batch_get_offload_object RPC to each alive +// peer in this list, and reports per-peer statistics in addition to the +// aggregate. Each peer must be started with --enable_ssd_offload=true. +DEFINE_string(peer_rpc_addrs, "", + "[client_rpc_bench] Comma-separated list of peer RealClient " + "offload_rpc_server_ addresses (ip:port), e.g. " + "\"10.0.0.2:17888,10.0.0.3:17888,...\". When non-empty, takes " + "precedence over --peer_rpc_addr. Each peer is expected to " + "have --enable_ssd_offload=true. Per-peer statistics are " + "reported alongside an AGGREGATE summary."); +// Optional: if set, we will call batch_get_offload_object with this single +// key and --ssd_value_size. The peer's FileStorage must already hold this +// key in its SSD directory (use scenario=local_disk on the peer first, or +// any other writer that has offloaded to the peer's SSD). When unset, the +// RPC is issued with empty keys (no SSD I/O on the peer); this still +// measures a real client-to-client RPC round-trip, but the response +// contains an empty pointers vector. +DEFINE_string(ssd_key, "", + "[client_rpc_bench] Key that the peer has in its SSD. When " + "set, the bench actually reads from the peer's SSD; the " + "response will carry a real pointer (memory address) to the " + "data buffer on the peer. Leave empty to send an empty " + "request (no SSD I/O)."); +DEFINE_int64(ssd_value_size, 4096, + "[client_rpc_bench] Size in bytes of the SSD-backed object to " + "request from the peer. Only used when --ssd_key is set."); DEFINE_string(role, "writer", "Node role: writer (prefill data) or reader (benchmark reads)"); +DEFINE_string(client_type, "real", + "Underlying client implementation: real (RealClient) or dummy " + "(DummyClient that connects to a remote RealClient via RPC + " + "IPC). The dummy client requires a real client to be running " + "with its RPC server reachable at --dummy_server_address and " + "its IPC server listening on --dummy_ipc_socket_path."); +DEFINE_string(dummy_server_address, "127.0.0.1:12345", + "[dummy client] RealClient RPC server address (IP:port) that " + "the DummyClient connects to."); +DEFINE_string(dummy_ipc_socket_path, "/tmp/mooncake_dummy.sock", + "[dummy client] Abstract-namespace Unix socket path used by " + "DummyClient to register SHM with the real client. Should " + "match the --ipc_socket_path of the real client."); +DEFINE_uint64(dummy_mem_pool_size, 0, + "[dummy client] Memory pool size in bytes allocated inside " + "DummyClient. 0 = reuse --global_segment_size."); DEFINE_uint64(value_size, 4 * MB, "Size of each value in bytes"); DEFINE_uint64(num_keys, 100, "Number of keys to write/read"); DEFINE_uint64(batch_size, 32, "Batch size for put/get operations"); @@ -233,7 +312,7 @@ inline double NanosToMs(int64_t ns) { inline double NanosToSec(int64_t ns) { return static_cast(ns) / 1e9; } struct ThreadResult { - std::vector latencies_ns; // per-query latency + std::vector latencies_ns; size_t total_bytes = 0; size_t total_keys = 0; // number of keys processed size_t total_queries = 0; // number of API calls (get_into / batch_get_into) @@ -378,73 +457,183 @@ class BenchmarkStats { class StressBenchmark { public: StressBenchmark() - : client_(mooncake::RealClient::create()), + : client_(nullptr), + primary_dummy_client_(nullptr), + main_buffer_client_(nullptr), buffer_(nullptr), buffer_size_(0) {} ~StressBenchmark() { - // Early return if already cleaned up - if (!client_) { + // Early return if nothing was set up (covers both real and dummy + // modes - in dummy mode client_ stays null but primary_dummy_client_ + // and dummy_clients_ may be populated). + if (!client_ && !primary_dummy_client_ && dummy_clients_.empty()) { return; } - // Unregister and free thread buffers (these are allocated by this - // class) - for (auto& tb : thread_buffers_) { - if (tb.ptr) { + // Unregister and free per-thread buffers. For dummy mode each + // thread owns its own DummyClient and its own SHM, so we must + // unregister the buffer with the same client that registered it + // (otherwise the real client side will not find the matching + // registration). For real mode all threads share client_, so + // unregister is a no-op for repeated calls. + for (size_t t = 0; t < thread_buffers_.size(); ++t) { + auto& tb = thread_buffers_[t]; + if (!tb.ptr) continue; + std::shared_ptr thread_client; + if (is_dummy_ && t < dummy_clients_.size() && + dummy_clients_[t]) { + thread_client = dummy_clients_[t]; + } else { + thread_client = client_; + } + if (thread_client) { try { - client_->unregister_buffer(tb.ptr); + thread_client->unregister_buffer(tb.ptr); } catch (...) { LOG(WARNING) - << "Failed to unregister thread buffer, ignoring"; + << "Failed to unregister thread " << t + << " buffer, ignoring"; } - numa_free(tb.ptr, tb.size); - tb.ptr = nullptr; } + FreeBuffer(tb.ptr, tb.size); + tb.ptr = nullptr; } thread_buffers_.clear(); - // Unregister and free main buffer (allocated by this class) - if (buffer_) { + // Unregister and free the main buffer (allocated by this class). + if (buffer_ && main_buffer_client_) { try { - client_->unregister_buffer(buffer_); + main_buffer_client_->unregister_buffer(buffer_); } catch (...) { LOG(WARNING) << "Failed to unregister main buffer, ignoring"; } - numa_free(buffer_, buffer_size_); - buffer_ = nullptr; } - client_ = nullptr; + FreeBuffer(buffer_, buffer_size_); + buffer_ = nullptr; + + // Tear down per-thread DummyClients (one per reader thread). Each + // call to DummyClient::tearDownAll() also stops the ping thread and + // closes the IPC / RPC channels belonging to that client. + for (auto& dc : dummy_clients_) { + if (!dc) continue; + try { + dc->tearDownAll(); + } catch (...) { + LOG(WARNING) << "Failed to tearDownAll per-thread dummy " + "client, ignoring"; + } + } + dummy_clients_.clear(); + + // Tear down the primary DummyClient (writer's client). + if (primary_dummy_client_) { + try { + primary_dummy_client_->tearDownAll(); + } catch (...) { + LOG(WARNING) << "Failed to tearDownAll primary dummy " + "client, ignoring"; + } + primary_dummy_client_.reset(); + } + + // Tear down the RealClient (no-op in dummy mode). + if (client_) { + try { + client_->tearDownAll(); + } catch (...) { + LOG(WARNING) << "Failed to tearDownAll real client, ignoring"; + } + client_ = nullptr; + } + main_buffer_client_ = nullptr; } int Setup() { - int ret = client_->setup_real( - FLAGS_local_hostname, FLAGS_metadata_server, - FLAGS_global_segment_size, FLAGS_local_buffer_size, FLAGS_protocol, - FLAGS_device_name, FLAGS_master_server, nullptr, "", - FLAGS_enable_ssd_offload, FLAGS_ssd_offload_path); - if (ret != 0) { - LOG(ERROR) << "RealClient setup_real failed, ret=" << ret; - return ret; + if (FLAGS_client_type == "dummy") { + is_dummy_ = true; + // Build the primary DummyClient that owns the main buffer. The + // single-threaded writer, warmup, and verify paths use this + // client and operate on buffer_. Per-thread DummyClients (one + // per reader thread) are created later in AllocateThreadBuffers. + // local_buffer_size is set to 0 so the primary client does not + // pre-allocate an extra SHM segment - we will allocate buffer_ + // ourselves via ShmHelper and register it explicitly. + primary_dummy_client_ = + CreateDummyClient(/*local_buffer_size=*/0); + if (!primary_dummy_client_) { + return -1; + } + main_buffer_client_ = primary_dummy_client_; + client_ = nullptr; + LOG(INFO) << "DummyClient (primary, for writer) setup succeeded " + << "(server=" << FLAGS_dummy_server_address + << ", ipc=" << FLAGS_dummy_ipc_socket_path << ")"; + } else if (FLAGS_client_type == "real") { + is_dummy_ = false; + auto real = mooncake::RealClient::create(); + int ret = real->setup_real( + FLAGS_local_hostname, FLAGS_metadata_server, + FLAGS_global_segment_size, FLAGS_local_buffer_size, + FLAGS_protocol, FLAGS_device_name, FLAGS_master_server, nullptr, + "", FLAGS_enable_ssd_offload, FLAGS_ssd_offload_path); + if (ret != 0) { + LOG(ERROR) << "RealClient setup_real failed, ret=" << ret; + return ret; + } + client_ = real; + main_buffer_client_ = real; + primary_dummy_client_ = nullptr; + LOG(INFO) << "RealClient setup succeeded" + << (FLAGS_enable_ssd_offload + ? " (SSD offload enabled)" + : ""); + } else { + LOG(ERROR) << "Unknown --client_type: " << FLAGS_client_type + << " (expected 'real' or 'dummy')"; + return -1; } - LOG(INFO) << "RealClient setup succeeded" - << (FLAGS_enable_ssd_offload ? " (SSD offload enabled)" : ""); buffer_size_ = FLAGS_batch_size * FLAGS_value_size; - buffer_ = reinterpret_cast(numa_alloc_local(buffer_size_)); + buffer_ = AllocateBuffer(buffer_size_); if (!buffer_) { LOG(ERROR) << "Failed to allocate buffer of " << buffer_size_ - << " bytes"; + << " bytes (" + << (is_dummy_ ? "ShmHelper" : "numa") << ")"; + // Tear down the primary dummy client immediately: it has an + // open RPC/IPC connection to the real client and a background + // ping thread; leaving it around would leak those resources + // (the destructor would still clean up, but we want the + // failure to be self-contained). + if (primary_dummy_client_) { + primary_dummy_client_->tearDownAll(); + primary_dummy_client_.reset(); + } + main_buffer_client_ = nullptr; return -1; } std::memset(buffer_, 0, buffer_size_); - ret = client_->register_buffer(buffer_, buffer_size_); + // Register the main buffer with the client that owns it: the + // primary dummy client in dummy mode, or the shared real client + // in real mode. If registration fails, free the buffer and the + // primary dummy client immediately so we don't leave a half-open + // RPC/IPC connection to the real client. + int ret = + main_buffer_client_->register_buffer(buffer_, buffer_size_); if (ret != 0) { - LOG(ERROR) << "register_buffer failed, ret=" << ret; + LOG(ERROR) << "register_buffer (main) failed, ret=" << ret; + FreeBuffer(buffer_, buffer_size_); + buffer_ = nullptr; + if (primary_dummy_client_) { + primary_dummy_client_->tearDownAll(); + primary_dummy_client_.reset(); + } + main_buffer_client_ = nullptr; return ret; } - LOG(INFO) << "Registered buffer of " << buffer_size_ / MB << " MB"; + LOG(INFO) << "Registered main buffer of " << buffer_size_ / MB + << " MB"; return 0; } @@ -465,7 +654,7 @@ class StressBenchmark { FillBuffer(i); auto t0 = Clock::now(); - int ret = client_->put_from(key, buffer_, FLAGS_value_size, config); + int ret = main_buffer_client_->put_from(key, buffer_, FLAGS_value_size, config); auto t1 = Clock::now(); if (ret != 0) { @@ -565,7 +754,7 @@ class StressBenchmark { for (size_t i = 0; i < FLAGS_num_keys; ++i) { std::string key = MakeKey(i); FillBuffer(i); - int ret = client_->put_from(key, buffer_, FLAGS_value_size, config); + int ret = main_buffer_client_->put_from(key, buffer_, FLAGS_value_size, config); if (ret != 0) { LOG(ERROR) << "put_from failed for key=" << key; return ret; @@ -631,7 +820,7 @@ class StressBenchmark { for (size_t i = 0; i < FLAGS_num_keys; ++i) { std::string key = MakeKey(i); FillBuffer(i); - int ret = client_->put_from(key, buffer_, FLAGS_value_size, config); + int ret = main_buffer_client_->put_from(key, buffer_, FLAGS_value_size, config); if (ret != 0) { LOG(ERROR) << "put_from failed for key=" << key; return ret; @@ -742,7 +931,7 @@ class StressBenchmark { FillBuffer(i); auto t0 = Clock::now(); - int ret = client_->put_from(key, buffer_, FLAGS_value_size, + int ret = main_buffer_client_->put_from(key, buffer_, FLAGS_value_size, configs[s]); auto t1 = Clock::now(); @@ -812,8 +1001,8 @@ class StressBenchmark { if (buf_ret != 0) return buf_ret; std::vector all_keys; - for (size_t i = 0; i < FLAGS_num_keys; ++i) { - for (size_t s = 0; s < read_segments.size(); ++s) { + for (size_t s = 0; s < read_segments.size(); ++s) { + for (size_t i = 0; i < FLAGS_num_keys; ++i) { all_keys.push_back(MakeSegmentKey(read_segments[s], i)); } } @@ -825,7 +1014,7 @@ class StressBenchmark { LOG(INFO) << "Warmup: reading " << warmup_end << " keys..."; for (size_t i = 0; i < warmup_end; ++i) { int64_t ret = - client_->get_into(all_keys[i], buffer_, FLAGS_value_size); + main_buffer_client_->get_into(all_keys[i], buffer_, FLAGS_value_size); if (ret < 0) { LOG(WARNING) << "Warmup get_into failed for key=" << all_keys[i] @@ -910,9 +1099,18 @@ class StressBenchmark { if (n >= 1000) { p999_latency_ns = latencies_ns[static_cast(n * 0.999)]; } + // P99.99 needs n >= 10000 to be statistically meaningful. + // When n < 10000 we fall back to the maximum sample so + // downstream consumers see a non-zero number (it would + // otherwise be silently 0.0, which is misleading). The + // print code already tags it with "(n<10000)" so the + // user knows it's an upper-bound proxy, not a true + // percentile. if (n >= 10000) { p9999_latency_ns = latencies_ns[static_cast(n * 0.9999)]; + } else { + p9999_latency_ns = latencies_ns.back(); } } @@ -966,87 +1164,103 @@ class StressBenchmark { (total_keys + FLAGS_num_threads - 1) / FLAGS_num_threads; for (size_t t = 0; t < FLAGS_num_threads; ++t) { - threads.emplace_back([&, t, keys_per_thread, total_keys]() { - bindToSocket(t % NR_SOCKETS); - char* my_buf = thread_buffers_[t].ptr; - - start_latch.arrive_and_wait(); - - size_t key_offset = t * keys_per_thread; - size_t key_idx = key_offset; - - if (FLAGS_batch_size <= 1) { - while (!stop_flag.load(std::memory_order_relaxed)) { - const std::string& key = all_keys[key_idx % total_keys]; - auto t0 = Clock::now(); - int64_t ret = - client_->get_into(key, my_buf, FLAGS_value_size); - auto t1 = Clock::now(); - int64_t latency_ns = ElapsedNanos(t0, t1); + // Per-thread client: each reader thread uses its own + // DummyClient in dummy mode (isolated SHM + RPC), or the + // shared RealClient in real mode. + std::shared_ptr thread_client; + if (is_dummy_) { + thread_client = dummy_clients_[t]; + } else { + thread_client = client_; + } + threads.emplace_back( + [&, t, keys_per_thread, total_keys, thread_client]() { + bindToSocket(t % NR_SOCKETS); + char* my_buf = thread_buffers_[t].ptr; - { - std::lock_guard lock( - latency_mutexes[t]); - thread_latencies[t].push_back(latency_ns); - } + start_latch.arrive_and_wait(); - if (ret < 0) { - global_failed.fetch_add(1, - std::memory_order_relaxed); - } else { - global_bytes.fetch_add(static_cast(ret), - std::memory_order_relaxed); - } - global_keys.fetch_add(1, std::memory_order_relaxed); - global_queries.fetch_add(1, std::memory_order_relaxed); - ++key_idx; - } - } else { - size_t per_key_buf = FLAGS_value_size; - while (!stop_flag.load(std::memory_order_relaxed)) { - std::vector keys; - std::vector bufs; - std::vector sizes; - keys.reserve(FLAGS_batch_size); - bufs.reserve(FLAGS_batch_size); - sizes.reserve(FLAGS_batch_size); + size_t key_offset = t * keys_per_thread; + size_t key_idx = key_offset; - for (size_t b = 0; b < FLAGS_batch_size; ++b) { + if (FLAGS_batch_size <= 1) { + while (!stop_flag.load(std::memory_order_relaxed)) { const std::string& key = all_keys[key_idx % total_keys]; - keys.push_back(key); - bufs.push_back(my_buf + b * per_key_buf); - sizes.push_back(FLAGS_value_size); - ++key_idx; - } - - auto t0 = Clock::now(); - auto results = - client_->batch_get_into(keys, bufs, sizes); - auto t1 = Clock::now(); - int64_t latency_ns = ElapsedNanos(t0, t1); - - { - std::lock_guard lock( - latency_mutexes[t]); - thread_latencies[t].push_back(latency_ns); - } + auto t0 = Clock::now(); + int64_t ret = thread_client->get_into( + key, my_buf, FLAGS_value_size); + auto t1 = Clock::now(); + int64_t latency_ns = ElapsedNanos(t0, t1); + + { + std::lock_guard lock( + latency_mutexes[t]); + thread_latencies[t].push_back(latency_ns); + } - for (size_t k = 0; k < results.size(); ++k) { - if (results[k] < 0) { + if (ret < 0) { global_failed.fetch_add( 1, std::memory_order_relaxed); } else { global_bytes.fetch_add( - static_cast(results[k]), + static_cast(ret), std::memory_order_relaxed); } global_keys.fetch_add(1, std::memory_order_relaxed); + global_queries.fetch_add(1, std::memory_order_relaxed); + ++key_idx; + } + } else { + size_t per_key_buf = FLAGS_value_size; + while (!stop_flag.load( + std::memory_order_relaxed)) { + std::vector keys; + std::vector bufs; + std::vector sizes; + keys.reserve(FLAGS_batch_size); + bufs.reserve(FLAGS_batch_size); + sizes.reserve(FLAGS_batch_size); + + for (size_t b = 0; b < FLAGS_batch_size; ++b) { + const std::string& key = + all_keys[key_idx % total_keys]; + keys.push_back(key); + bufs.push_back(my_buf + b * per_key_buf); + sizes.push_back(FLAGS_value_size); + ++key_idx; + } + + auto t0 = Clock::now(); + auto results = thread_client->batch_get_into( + keys, bufs, sizes); + auto t1 = Clock::now(); + int64_t latency_ns = ElapsedNanos(t0, t1); + + { + std::lock_guard lock( + latency_mutexes[t]); + for (size_t k = 0; k < results.size(); ++k) { + thread_latencies[t].push_back( + latency_ns / FLAGS_batch_size); + } + } + + for (size_t k = 0; k < results.size(); ++k) { + if (results[k] < 0) { + global_failed.fetch_add( + 1, std::memory_order_relaxed); + } else { + global_bytes.fetch_add( + static_cast(results[k]), + std::memory_order_relaxed); + } + global_keys.fetch_add(1, std::memory_order_relaxed); + } + global_queries.fetch_add(1, std::memory_order_relaxed); } - global_queries.fetch_add(1, std::memory_order_relaxed); } - } - }); + }); } auto bench_start = Clock::now(); @@ -1272,6 +1486,771 @@ class StressBenchmark { return 0; } + // ------------------------------------------------------------------------ + // client_rpc_bench: measure the round-trip latency of a single + // client-to-client RPC, with no master RPC, no SSD I/O, and no P2P + // data transfer. + // + // Multi-peer variant: in each round of the timed loop the bench + // visits every alive peer sequentially and reports per-peer + // statistics in addition to the AGGREGATE summary. Use + // --peer_rpc_addrs=ip1:port,ip2:port,... (preferred) or + // --peer_rpc_addr=ip:port (single-peer fallback) to point at the + // peer(s). Each peer must be started with + // --enable_ssd_offload=true so its offload_rpc_server_ is up. + // + // Implementation: spin up a standalone mooncake::ClientRequester in + // this process, and use it to call + // RealClient::batch_get_offload_object + // (the only client-side RPC handler that is already registered to the + // peer's offload_rpc_server_) with EMPTY keys and sizes. The peer: + // 1. receives the RPC request + // 2. runs BatchGet({},{}) on its FileStorage -- with empty inputs + // this is O(1) bookkeeping, no SSD read, no allocation loop + // 3. replies with BatchGetOffloadObjectResponse{ batch_id, + // pointers=, transfer_engine_addr=, + // gc_ttl_ms } + // We time the full call and treat the response's + // transfer_engine_addr as the "peer memory address" requested by the + // user. + // + // Why NOT use get_into / batch_get_into / execute_ranged_read / + // isExist: all of those go through the master (stage 1) or perform + // real data I/O (stages 2-5), so they do not isolate a single RPC + // hop. + // + // Prerequisite: every peer must be started with + // --enable_ssd_offload=true so its offload_rpc_server_ is up and + // the batch_get_offload_object handler is registered. + // ------------------------------------------------------------------------ + int RunClientRpcBench() { + LOG(INFO) << "=== CLIENT-TO-CLIENT RPC BENCHMARK (single " + "ClientRequester, batch_get_offload_object, multi-peer) ==="; + + if (FLAGS_num_threads == 0) { + LOG(ERROR) << "--num_threads must be > 0"; + return -1; + } + + auto peer_addrs = ParsePeerRpcAddrs(); + if (peer_addrs.empty()) { + LOG(ERROR) + << "Provide --peer_rpc_addrs=ip1:port,ip2:port,... or " + "--peer_rpc_addr=ip:port for client_rpc_bench. Each peer " + "must be started with --enable_ssd_offload=true."; + return -1; + } + LOG(INFO) << "Targeting " << peer_addrs.size() << " peer(s):"; + for (size_t i = 0; i < peer_addrs.size(); ++i) { + LOG(INFO) << " peer[" << i << "]: " << peer_addrs[i]; + } + + // One ClientRequester shared by all worker threads; the + // underlying coro_rpc client pool is thread-safe. + auto requester = std::make_shared(); + + // Build the request payload ONCE and reuse across all peers and + // all iterations so every measured call hits exactly the same + // path. When --ssd_key is set, the request actually reads from + // the peer's SSD storage; otherwise it is an empty request (no + // SSD I/O). + std::vector req_keys; + std::vector req_sizes; + if (!FLAGS_ssd_key.empty()) { + req_keys.push_back(FLAGS_ssd_key); + req_sizes.push_back(FLAGS_ssd_value_size); + LOG(INFO) << "Will request real SSD read: key=\"" + << FLAGS_ssd_key << "\" size=" + << FLAGS_ssd_value_size << "B. Make sure the peer " + << "has this key in its SSD (write it first via " + << "scenario=local_disk or similar on the peer)."; + } else { + LOG(INFO) << "Will request empty payload (no SSD I/O on peer). " + << "Set --ssd_key= to actually read SSD."; + } + + // Per-peer warmup. Peers whose warmup fully fails are kept in + // the list and marked warmup_ok=false so the timed loop skips + // them. We only abort if EVERY peer fails warmup. + std::vector peer_states(peer_addrs.size()); + for (size_t p = 0; p < peer_addrs.size(); ++p) { + peer_states[p].peer_addr = peer_addrs[p]; + size_t warmup_n = static_cast( + std::max(1, FLAGS_warmup_keys)); + size_t warmup_done = 0; + for (size_t i = 0; i < warmup_n; ++i) { + auto r = requester->batch_get_offload_object( + peer_addrs[p], req_keys, req_sizes); + if (r) { + ++warmup_done; + if (peer_states[p].peer_seen_addr.empty()) { + peer_states[p].peer_seen_addr = + r->transfer_engine_addr; + } + } + } + peer_states[p].warmup_ok = (warmup_done > 0); + LOG(INFO) << "Peer[" << p << "] " << peer_addrs[p] + << " warmup " << warmup_done << "/" << warmup_n + << ", transfer_engine_addr=\"" + << peer_states[p].peer_seen_addr << "\""; + if (!peer_states[p].warmup_ok) { + LOG(ERROR) << "Peer[" << p << "] " << peer_addrs[p] + << " warmup failed; this peer will be skipped."; + } + } + + size_t alive_peers = 0; + for (const auto& ps : peer_states) { + if (ps.warmup_ok) ++alive_peers; + } + if (alive_peers == 0) { + LOG(ERROR) << "All peer warmups failed. Aborting. Check that " + "the peers are reachable and were started with " + "--enable_ssd_offload=true."; + return -1; + } + LOG(INFO) << alive_peers << "/" << peer_states.size() + << " peer(s) ready for the timed loop."; + + if (FLAGS_duration == 0) { + return RunClientRpcBenchSinglePass(peer_states, req_keys, + req_sizes, requester); + } + return RunClientRpcBenchDuration(peer_states, req_keys, req_sizes, + requester); + } + + int RunClientRpcBenchSinglePass( + std::vector& peer_states, + const std::vector& req_keys, + const std::vector& req_sizes, + const std::shared_ptr& requester) { + const size_t num_peers = peer_states.size(); + const size_t per_thread = std::max(1, FLAGS_num_keys); + const size_t per_thread_rpcs = per_thread * num_peers; + + LOG(INFO) + << "Single-pass mode: " << FLAGS_num_threads << " threads x " + << per_thread << " rounds x " << num_peers + << " peers (skipping warmup-failed) = up to " + << (FLAGS_num_threads * per_thread_rpcs) + << " total RPCs. Measured: single client-to-client RPC RTT."; + + for (auto& ps : peer_states) { + if (ps.warmup_ok) { + ps.stats.InitThreads(FLAGS_num_threads, per_thread); + } + } + + std::latch start_latch(static_cast(FLAGS_num_threads)); + std::latch done_latch(static_cast(FLAGS_num_threads)); + std::vector threads; + threads.reserve(FLAGS_num_threads); + + for (size_t t = 0; t < FLAGS_num_threads; ++t) { + threads.emplace_back([&, t]() { + bindToSocket(t % NR_SOCKETS); + for (auto& ps : peer_states) { + if (ps.warmup_ok) { + ps.stats.GetThreadResult(t).latencies_ns.reserve( + per_thread); + } + } + start_latch.arrive_and_wait(); + + for (size_t k = 0; k < per_thread; ++k) { + // One round = one RPC per alive peer, sequentially. + for (size_t p = 0; p < num_peers; ++p) { + if (!peer_states[p].warmup_ok) continue; + const std::string& addr = peer_states[p].peer_addr; + + // --- single client-to-client RPC: send + // request, receive + // BatchGetOffloadObjectResponse --- + auto t0 = Clock::now(); + auto ret = requester->batch_get_offload_object( + addr, req_keys, req_sizes); + auto t1 = Clock::now(); + + int64_t lat_ns = ElapsedNanos(t0, t1); + ThreadResult& tr = + peer_states[p].stats.GetThreadResult(t); + tr.latencies_ns.push_back(lat_ns); + + if (!ret) { + ++tr.failed_ops; + LOG_EVERY_N(ERROR, 100) + << "batch_get_offload_object RPC to " + << addr << " failed: " + << mooncake::toString(ret.error()); + } else { + // Each successful response carries a + // BatchGetOffloadObjectResponse. "Bytes" in + // the stats sheet is the response size, + // which is a good proxy for "amount of data + // round-tripped" -- not real payload bytes. + tr.total_bytes += + static_cast( + ret->pointers.size() * sizeof(uint64_t)) + + ret->transfer_engine_addr.size() + + sizeof(ret->batch_id) + + sizeof(ret->gc_ttl_ms); + // Sanity: the response must have been served + // by the peer we asked for. Print the first + // successful peer's address per thread for + // end-to-end verification. + if (k == 0 && + peer_states[p].peer_seen_addr != + ret->transfer_engine_addr) { + LOG(WARNING) + << " [t=" << t << " p=" << p + << "] peer transfer_engine_addr " + << "mismatch: expected=\"" + << peer_states[p].peer_seen_addr + << "\" got=\"" + << ret->transfer_engine_addr << "\""; + } + } + ++tr.total_keys; + ++tr.total_queries; + } + } + + done_latch.arrive_and_wait(); + }); + } + + done_latch.wait(); + for (auto& th : threads) th.join(); + + // Per-peer prints, then AGGREGATE. + for (auto& ps : peer_states) { + if (!ps.warmup_ok) continue; + ps.stats.Finalize(); + std::string title = "CLIENT-TO-CLIENT RPC BENCHMARK [peer=" + + ps.peer_addr + "]"; + ps.stats.Print(title); + } + BenchmarkStats agg = MergePeerStats(peer_states); + agg.Print( + "CLIENT-TO-CLIENT RPC BENCHMARK [AGGREGATE across all alive " + "peers]"); + + return 0; + } + + int RunClientRpcBenchDuration( + std::vector& peer_states, + const std::vector& req_keys, + const std::vector& req_sizes, + const std::shared_ptr& requester) { + const size_t num_peers = peer_states.size(); + LOG(INFO) << "Duration mode: " << FLAGS_num_threads + << " threads continuously fire client-to-client RPCs to " + << num_peers << " peer(s) for " << FLAGS_duration + << "s, stats every " << FLAGS_statis_interval + << "s. keys=" << req_keys.size() << ", sizes=" + << req_sizes.size() << " (" + << (req_keys.empty() + ? "empty payload -> no SSD I/O on peer" + : "real SSD read on peer") + << "). Measured: single client-to-client RPC RTT."; + + for (auto& ps : peer_states) { + if (ps.warmup_ok) { + ps.InitDuration(FLAGS_num_threads); + } + } + + std::atomic stop_flag{false}; + std::latch start_latch(static_cast(FLAGS_num_threads)); + std::vector threads; + threads.reserve(FLAGS_num_threads); + + for (size_t t = 0; t < FLAGS_num_threads; ++t) { + threads.emplace_back([&, t]() { + bindToSocket(t % NR_SOCKETS); + start_latch.arrive_and_wait(); + + if (FLAGS_batch_size <= 1) { + // Single-RPC-per-iteration path. Each round + // visits every alive peer sequentially and times + // each RPC individually; the latency distribution + // is the steady-state RTT of one client-to-client + // RPC to that peer. + while ( + !stop_flag.load(std::memory_order_relaxed)) { + for (size_t p = 0; p < num_peers; ++p) { + if (!peer_states[p].warmup_ok) continue; + const std::string& addr = peer_states[p].peer_addr; + auto t0 = Clock::now(); + auto ret = requester->batch_get_offload_object( + addr, req_keys, req_sizes); + auto t1 = Clock::now(); + int64_t lat_ns = ElapsedNanos(t0, t1); + PeerBenchState& ps = peer_states[p]; + { + std::lock_guard lk( + ps.latency_mutexes[t]); + ps.thread_latencies[t].push_back(lat_ns); + } + if (!ret) { + ps.global_failed.fetch_add( + 1, std::memory_order_relaxed); + } else { + ps.global_bytes.fetch_add( + static_cast( + ret->pointers.size() * + sizeof(uint64_t)) + + ret->transfer_engine_addr.size() + + sizeof(ret->batch_id) + + sizeof(ret->gc_ttl_ms), + std::memory_order_relaxed); + } + ps.global_keys.fetch_add( + 1, std::memory_order_relaxed); + ps.global_queries.fetch_add( + 1, std::memory_order_relaxed); + } + } + } else { + // Batch path: each "batch" is FLAGS_batch_size + // rounds, and each round visits all alive peers. + // We time the whole batch and split equally across + // (batch_size * num_peers) RPCs for the per-RPC + // latency samples, then attribute FLAGS_batch_size + // samples (= per_rpc_ns each) to every alive peer + // to match the original single-peer semantics + // (one sample per RPC issued to that peer). + while ( + !stop_flag.load(std::memory_order_relaxed)) { + auto t0 = Clock::now(); + for (size_t b = 0; b < FLAGS_batch_size; ++b) { + for (size_t p = 0; p < num_peers; ++p) { + if (!peer_states[p].warmup_ok) continue; + const std::string& addr = + peer_states[p].peer_addr; + auto ret = requester->batch_get_offload_object( + addr, req_keys, req_sizes); + PeerBenchState& ps = peer_states[p]; + if (!ret) { + ps.global_failed.fetch_add( + 1, std::memory_order_relaxed); + } else { + ps.global_bytes.fetch_add( + static_cast( + ret->pointers.size() * + sizeof(uint64_t)) + + ret->transfer_engine_addr.size() + + sizeof(ret->batch_id) + + sizeof(ret->gc_ttl_ms), + std::memory_order_relaxed); + } + ps.global_keys.fetch_add( + 1, std::memory_order_relaxed); + } + } + auto t1 = Clock::now(); + int64_t batch_lat_ns = ElapsedNanos(t0, t1); + int64_t per_rpc_ns = batch_lat_ns / + (FLAGS_batch_size * num_peers); + for (auto& ps : peer_states) { + if (!ps.warmup_ok) continue; + std::lock_guard lk( + ps.latency_mutexes[t]); + for (size_t b = 0; b < FLAGS_batch_size; ++b) { + ps.thread_latencies[t].push_back(per_rpc_ns); + } + } + for (auto& ps : peer_states) { + if (!ps.warmup_ok) continue; + ps.global_queries.fetch_add( + 1, std::memory_order_relaxed); + } + } + } + }); + } + + auto bench_start = Clock::now(); + auto bench_end = bench_start + std::chrono::seconds(FLAGS_duration); + auto next_statis = + bench_start + std::chrono::seconds(FLAGS_statis_interval); + + // Per-peer interval state (one prev_* set per peer so the + // per-peer interval deltas are computed independently). + std::vector prev_keys(num_peers, 0); + std::vector prev_queries(num_peers, 0); + std::vector prev_bytes(num_peers, 0); + std::vector prev_failed(num_peers, 0); + std::vector prev_time(num_peers, bench_start); + std::vector interval_stats_list(num_peers); + + // Aggregate interval state. + size_t prev_keys_agg = 0; + size_t prev_queries_agg = 0; + size_t prev_bytes_agg = 0; + size_t prev_failed_agg = 0; + auto prev_time_agg = bench_start; + std::vector interval_agg_list; + + std::cout << "\n"; + std::cout << "========================================" + << "========================================\n"; + std::cout << " CLIENT-TO-CLIENT RPC DURATION BENCHMARK " + "[batch_get_offload_object, " << num_peers + << " peer(s)]\n"; + std::cout << "========================================" + << "========================================\n"; + std::cout << std::fixed << std::setprecision(2); + + while (Clock::now() < bench_end) { + auto now = Clock::now(); + if (now >= next_statis) { + // ---- Per-peer interval print ---- + for (size_t p = 0; p < num_peers; ++p) { + if (!peer_states[p].warmup_ok) continue; + PeerBenchState& ps = peer_states[p]; + size_t cur_keys = + ps.global_keys.load(std::memory_order_relaxed); + size_t cur_queries = + ps.global_queries.load(std::memory_order_relaxed); + size_t cur_bytes = + ps.global_bytes.load(std::memory_order_relaxed); + size_t cur_failed = + ps.global_failed.load(std::memory_order_relaxed); + + double interval_sec = + NanosToSec(ElapsedNanos(prev_time[p], now)); + size_t interval_keys = cur_keys - prev_keys[p]; + size_t interval_queries = + cur_queries - prev_queries[p]; + size_t interval_bytes = cur_bytes - prev_bytes[p]; + size_t interval_failed = + cur_failed - prev_failed[p]; + + double interval_mbps = + (interval_sec > 0) + ? (static_cast(interval_bytes) / MB) / + interval_sec + : 0; + double interval_qps = + (interval_sec > 0) + ? static_cast(interval_queries) / + interval_sec + : 0; + double interval_kps = + (interval_sec > 0) + ? static_cast(interval_keys) / + interval_sec + : 0; + + IntervalLatencyStats iv; + iv.throughput_mbps = interval_mbps; + iv.queries_per_sec = interval_qps; + iv.keys_per_sec = interval_kps; + for (size_t tt = 0; tt < FLAGS_num_threads; ++tt) { + std::lock_guard lk( + ps.latency_mutexes[tt]); + iv.latencies_ns.insert( + iv.latencies_ns.end(), + ps.thread_latencies[tt].begin(), + ps.thread_latencies[tt].end()); + ps.thread_latencies[tt].clear(); + } + iv.Finalize(); + interval_stats_list[p].Aggregate(iv); + + double total_sec = + NanosToSec(ElapsedNanos(bench_start, now)); + double total_mbps = + (total_sec > 0) + ? (static_cast(cur_bytes) / MB) / + total_sec + : 0; + double total_qps = + (total_sec > 0) + ? static_cast(cur_queries) / total_sec + : 0; + + std::cout << " [t=" << std::setw(6) << total_sec + << "s]" + << " peer[" << ps.peer_addr + << "] interval: " << interval_mbps + << " MB/s, " << interval_qps << " qps" + << " (failed=" << interval_failed << ")" + << " lat[us]: avg=" + << NanosToUs(iv.avg_latency_ns) + << ", P50=" << NanosToUs(iv.p50_latency_ns) + << ", P99=" << NanosToUs(iv.p99_latency_ns) + << " total: " << cur_queries << " queries, " + << cur_keys << " keys, " << total_mbps + << " MB/s, " << total_qps << " qps" + << " (failed=" << cur_failed << ")\n"; + + prev_keys[p] = cur_keys; + prev_queries[p] = cur_queries; + prev_bytes[p] = cur_bytes; + prev_failed[p] = cur_failed; + prev_time[p] = now; + } + + // ---- Aggregate interval print ---- + { + size_t cur_keys = 0, cur_queries = 0, cur_bytes = 0, + cur_failed = 0; + for (const auto& ps : peer_states) { + if (!ps.warmup_ok) continue; + cur_keys += + ps.global_keys.load(std::memory_order_relaxed); + cur_queries += + ps.global_queries.load(std::memory_order_relaxed); + cur_bytes += + ps.global_bytes.load(std::memory_order_relaxed); + cur_failed += + ps.global_failed.load(std::memory_order_relaxed); + } + double interval_sec = + NanosToSec(ElapsedNanos(prev_time_agg, now)); + size_t interval_keys = cur_keys - prev_keys_agg; + size_t interval_queries = + cur_queries - prev_queries_agg; + size_t interval_bytes = cur_bytes - prev_bytes_agg; + size_t interval_failed = + cur_failed - prev_failed_agg; + + double interval_mbps = + (interval_sec > 0) + ? (static_cast(interval_bytes) / MB) / + interval_sec + : 0; + double interval_qps = + (interval_sec > 0) + ? static_cast(interval_queries) / + interval_sec + : 0; + double interval_kps = + (interval_sec > 0) + ? static_cast(interval_keys) / + interval_sec + : 0; + + IntervalLatencyStats iv; + iv.throughput_mbps = interval_mbps; + iv.queries_per_sec = interval_qps; + iv.keys_per_sec = interval_kps; + for (auto& ps : peer_states) { + if (!ps.warmup_ok) continue; + for (size_t tt = 0; tt < FLAGS_num_threads; ++tt) { + std::lock_guard lk( + ps.latency_mutexes[tt]); + iv.latencies_ns.insert( + iv.latencies_ns.end(), + ps.thread_latencies[tt].begin(), + ps.thread_latencies[tt].end()); + ps.thread_latencies[tt].clear(); + } + } + iv.Finalize(); + interval_agg_list.push_back(iv); + + double total_sec = + NanosToSec(ElapsedNanos(bench_start, now)); + double total_mbps = + (total_sec > 0) + ? (static_cast(cur_bytes) / MB) / + total_sec + : 0; + double total_qps = + (total_sec > 0) + ? static_cast(cur_queries) / total_sec + : 0; + + std::cout << " [t=" << std::setw(6) << total_sec + << "s]" + << " AGGREGATE interval: " << interval_mbps + << " MB/s, " << interval_qps << " qps" + << " (failed=" << interval_failed << ")" + << " lat[us]: avg=" + << NanosToUs(iv.avg_latency_ns) + << ", P50=" << NanosToUs(iv.p50_latency_ns) + << ", P99=" << NanosToUs(iv.p99_latency_ns) + << " total: " << cur_queries << " queries, " + << cur_keys << " keys, " << total_mbps + << " MB/s, " << total_qps << " qps" + << " (failed=" << cur_failed << ")\n"; + + prev_keys_agg = cur_keys; + prev_queries_agg = cur_queries; + prev_bytes_agg = cur_bytes; + prev_failed_agg = cur_failed; + prev_time_agg = now; + } + + next_statis += std::chrono::seconds(FLAGS_statis_interval); + } + std::this_thread::sleep_for(std::chrono::milliseconds(100)); + } + + stop_flag.store(true, std::memory_order_relaxed); + for (auto& th : threads) th.join(); + + // ---- Per-peer final summary ---- + auto per_peer_final_time = Clock::now(); + for (size_t p = 0; p < num_peers; ++p) { + if (!peer_states[p].warmup_ok) continue; + PeerBenchState& ps = peer_states[p]; + double total_sec = + NanosToSec(ElapsedNanos(bench_start, per_peer_final_time)); + size_t final_keys = + ps.global_keys.load(std::memory_order_relaxed); + size_t final_queries = + ps.global_queries.load(std::memory_order_relaxed); + size_t final_bytes = + ps.global_bytes.load(std::memory_order_relaxed); + size_t final_failed = + ps.global_failed.load(std::memory_order_relaxed); + + double final_mbps = + (total_sec > 0) + ? (static_cast(final_bytes) / MB) / total_sec + : 0; + double final_qps = + (total_sec > 0) + ? static_cast(final_queries) / total_sec + : 0; + double final_kps = + (total_sec > 0) + ? static_cast(final_keys) / total_sec + : 0; + + const IntervalLatencyStats& overall = interval_stats_list[p]; + + std::cout << "\n FINAL SUMMARY [peer=" << ps.peer_addr << "]\n"; + std::cout << " Total time: " << total_sec << " s\n"; + std::cout << " Total queries: " << final_queries + << " (failed: " << final_failed << ")\n"; + std::cout << " Total keys: " << final_keys << "\n"; + std::cout << " Total data: " << FormatBytes(final_bytes) + << " (response bytes, not payload)\n"; + std::cout << " Throughput: " << final_mbps << " MB/s"; + if (final_mbps > 1024) + std::cout << " (" << final_mbps / 1024 << " GB/s)"; + std::cout << "\n"; + std::cout << " Keys/sec: " << final_kps << "\n"; + std::cout << " Queries/sec: " << final_qps << "\n"; + + if (overall.total_samples > 0) { + std::cout << "\n Latency (us) [n=" + << overall.total_samples << ", per-RPC]\n"; + std::cout << " Min: " << std::setw(12) + << NanosToUs(overall.min_latency_ns) << "\n"; + std::cout << " Avg: " << std::setw(12) + << NanosToUs(overall.avg_latency_ns) << "\n"; + std::cout << " P50: " << std::setw(12) + << NanosToUs(overall.p50_latency_ns) << "\n"; + std::cout << " P90: " << std::setw(12) + << NanosToUs(overall.p90_latency_ns) << "\n"; + std::cout << " P99: " << std::setw(12) + << NanosToUs(overall.p99_latency_ns) << "\n"; + std::cout << " P999: " << std::setw(12) + << NanosToUs(overall.p999_latency_ns); + if (overall.total_samples < 1000) std::cout << " (n<1000)"; + std::cout << "\n"; + std::cout << " P9999: " << std::setw(12) + << NanosToUs(overall.p9999_latency_ns); + if (overall.total_samples < 10000) + std::cout << " (n<10000)"; + std::cout << "\n"; + std::cout << " Max: " << std::setw(12) + << NanosToUs(overall.max_latency_ns) << "\n"; + } + } + + // ---- Aggregate final summary ---- + { + size_t final_keys = 0, final_queries = 0, final_bytes = 0, + final_failed = 0; + for (const auto& ps : peer_states) { + if (!ps.warmup_ok) continue; + final_keys += + ps.global_keys.load(std::memory_order_relaxed); + final_queries += + ps.global_queries.load(std::memory_order_relaxed); + final_bytes += + ps.global_bytes.load(std::memory_order_relaxed); + final_failed += + ps.global_failed.load(std::memory_order_relaxed); + } + double total_sec = + NanosToSec(ElapsedNanos(bench_start, per_peer_final_time)); + double final_mbps = + (total_sec > 0) + ? (static_cast(final_bytes) / MB) / total_sec + : 0; + double final_qps = + (total_sec > 0) + ? static_cast(final_queries) / total_sec + : 0; + double final_kps = + (total_sec > 0) + ? static_cast(final_keys) / total_sec + : 0; + + IntervalLatencyStats overall; + for (const auto& s : interval_agg_list) overall.Aggregate(s); + + std::cout << "\n FINAL SUMMARY [AGGREGATE across all alive " + "peers]\n"; + std::cout << " Total time: " << total_sec << " s\n"; + std::cout << " Total queries: " << final_queries + << " (failed: " << final_failed << ")\n"; + std::cout << " Total keys: " << final_keys << "\n"; + std::cout << " Total data: " << FormatBytes(final_bytes) + << " (response bytes, not payload)\n"; + std::cout << " Throughput: " << final_mbps << " MB/s"; + if (final_mbps > 1024) + std::cout << " (" << final_mbps / 1024 << " GB/s)"; + std::cout << "\n"; + std::cout << " Keys/sec: " << final_kps << "\n"; + std::cout << " Queries/sec: " << final_qps << "\n"; + + if (overall.total_samples > 0) { + std::cout << "\n Latency (us) [n=" + << overall.total_samples << ", per-RPC]\n"; + std::cout << " Min: " << std::setw(12) + << NanosToUs(overall.min_latency_ns) << "\n"; + std::cout << " Avg: " << std::setw(12) + << NanosToUs(overall.avg_latency_ns) << "\n"; + std::cout << " P50: " << std::setw(12) + << NanosToUs(overall.p50_latency_ns) << "\n"; + std::cout << " P90: " << std::setw(12) + << NanosToUs(overall.p90_latency_ns) << "\n"; + std::cout << " P99: " << std::setw(12) + << NanosToUs(overall.p99_latency_ns) << "\n"; + std::cout << " P999: " << std::setw(12) + << NanosToUs(overall.p999_latency_ns); + if (overall.total_samples < 1000) std::cout << " (n<1000)"; + std::cout << "\n"; + std::cout << " P9999: " << std::setw(12) + << NanosToUs(overall.p9999_latency_ns); + if (overall.total_samples < 10000) + std::cout << " (n<10000)"; + std::cout << "\n"; + std::cout << " Max: " << std::setw(12) + << NanosToUs(overall.max_latency_ns) << "\n"; + } + } + + std::cout << " Note: timed region is exactly the " + "batch_get_offload_object RPC round-trip; no master, " + "no P2P data, no SSD read.\n"; + std::cout << "========================================" + << "========================================\n\n"; + return 0; + } + int Run() { if (FLAGS_scenario == "local_memory") { return RunLocalMemory(); @@ -1283,6 +2262,8 @@ class StressBenchmark { return RunSegmentRead(); } else if (FLAGS_scenario == "list_segments") { return RunListSegments(); + } else if (FLAGS_scenario == "client_rpc_bench") { + return RunClientRpcBench(); } else if (FLAGS_scenario == "remote_memory" || FLAGS_scenario == "remote_disk") { if (FLAGS_role == "writer") { @@ -1297,6 +2278,72 @@ class StressBenchmark { } private: + // Per-peer benchmark state used by RunClientRpcBench* in multi-peer + // mode. Each peer owns its own per-thread results (single-pass) and + // its own per-thread latencies / global counters (duration), so the + // stats for every peer can be reported independently and then + // merged into an AGGREGATE summary. + struct PeerBenchState { + std::string peer_addr; + // transfer_engine_addr echoed back from the peer's first + // successful warmup RPC; used for end-to-end sanity checking. + std::string peer_seen_addr; + // Set to true if the per-peer warmup completed at least one + // successful RPC. Peers that fail warmup are kept in the list + // and skipped at timed-loop time. + bool warmup_ok = false; + + // Single-pass mode: per-thread results aggregated by + // BenchmarkStats::Finalize / Print. + BenchmarkStats stats; + + // Duration mode: per-thread latency samples and global + // counters. The worker loop visits every alive peer in + // sequence within each round, and writes each peer's samples / + // counters into its own PeerBenchState. + std::vector> thread_latencies; + std::vector latency_mutexes; + std::atomic global_keys{0}; + std::atomic global_queries{0}; + std::atomic global_bytes{0}; + std::atomic global_failed{0}; + + void InitDuration(size_t n_threads) { + thread_latencies.assign(n_threads, {}); + latency_mutexes.assign(n_threads, {}); + global_keys.store(0); + global_queries.store(0); + global_bytes.store(0); + global_failed.store(0); + } + }; + + // Merge per-peer per-thread results into a single BenchmarkStats + // for the AGGREGATE print. Only peers with warmup_ok == true are + // included. Called only by RunClientRpcBenchSinglePass. + static BenchmarkStats MergePeerStats( + const std::vector& peer_states) { + BenchmarkStats agg; + if (FLAGS_num_threads == 0) return agg; + agg.InitThreads(FLAGS_num_threads, /*expected_per_thread=*/0); + for (size_t t = 0; t < FLAGS_num_threads; ++t) { + ThreadResult& agg_tr = agg.GetThreadResult(t); + for (const auto& ps : peer_states) { + if (!ps.warmup_ok) continue; + const ThreadResult& pr_tr = ps.stats.GetThreadResult(t); + agg_tr.latencies_ns.insert(agg_tr.latencies_ns.end(), + pr_tr.latencies_ns.begin(), + pr_tr.latencies_ns.end()); + agg_tr.total_bytes += pr_tr.total_bytes; + agg_tr.total_keys += pr_tr.total_keys; + agg_tr.total_queries += pr_tr.total_queries; + agg_tr.failed_ops += pr_tr.failed_ops; + } + } + agg.Finalize(); + return agg; + } + static std::string MakeKey(size_t idx) { return "bench_key_" + std::to_string(idx); } @@ -1338,7 +2385,7 @@ class StressBenchmark { static_cast(FLAGS_num_keys)); for (size_t i = 0; i < warmup_end; ++i) { std::string key = MakeKey(i); - int64_t ret = client_->get_into(key, buffer_, FLAGS_value_size); + int64_t ret = main_buffer_client_->get_into(key, buffer_, FLAGS_value_size); if (ret < 0) { LOG(WARNING) << "Warmup get_into failed for key=" << key << " ret=" << ret; @@ -1351,7 +2398,8 @@ class StressBenchmark { void BatchReadWorker(size_t tid, size_t my_keys, size_t key_offset, BenchmarkStats& stats, std::latch& start_latch, std::latch& done_latch, - const std::function& key_func) { + const std::function& key_func, + std::shared_ptr thread_client) { bindToSocket(tid % NR_SOCKETS); ThreadResult& result = stats.GetThreadResult(tid); @@ -1372,7 +2420,8 @@ class StressBenchmark { std::string key = key_func(key_idx); auto t0 = Clock::now(); - int64_t ret = client_->get_into(key, my_buf, FLAGS_value_size); + int64_t ret = + thread_client->get_into(key, my_buf, FLAGS_value_size); auto t1 = Clock::now(); int64_t lat_ns = ElapsedNanos(t0, t1); @@ -1408,7 +2457,8 @@ class StressBenchmark { } auto t0 = Clock::now(); - auto results = client_->batch_get_into(key_list, bufs, sizes); + auto results = + thread_client->batch_get_into(key_list, bufs, sizes); auto t1 = Clock::now(); int64_t lat_ns = ElapsedNanos(t0, t1); @@ -1448,10 +2498,22 @@ class StressBenchmark { size_t my_keys = keys_per_thread + (t < remainder ? 1 : 0); size_t key_offset = t * keys_per_thread + std::min(t, remainder); - threads.emplace_back([&, t, my_keys, key_offset]() { - BatchReadWorker(t, my_keys, key_offset, stats, start_latch, - done_latch, key_func); - }); + // Per-thread client: a separate DummyClient for dummy mode + // (one per reader thread, isolated SHM and RPC), or the + // shared RealClient for real mode. + std::shared_ptr thread_client; + if (is_dummy_) { + thread_client = dummy_clients_[t]; + } else { + thread_client = client_; + } + + threads.emplace_back( + [&, t, my_keys, key_offset, thread_client]() { + BatchReadWorker(t, my_keys, key_offset, stats, + start_latch, done_latch, key_func, + thread_client); + }); } return threads; } @@ -1481,7 +2543,7 @@ class StressBenchmark { for (size_t i = 0; i < FLAGS_num_keys; ++i) { std::string key = MakeKey(i); - int64_t ret = client_->get_into(key, buffer_, FLAGS_value_size); + int64_t ret = main_buffer_client_->get_into(key, buffer_, FLAGS_value_size); if (ret < 0) { LOG(ERROR) << "Verify: get_into failed for key=" << key; ++errors; @@ -1498,9 +2560,92 @@ class StressBenchmark { return errors > 0 ? -1 : 0; } - std::shared_ptr client_; + // Real client (single, used for real mode only). + std::shared_ptr client_; + // For dummy mode: a dedicated DummyClient that owns the main buffer + // (buffer_). The single-threaded writer, warmup, and verify paths all + // operate on buffer_ via this client. + std::shared_ptr primary_dummy_client_; + // For dummy mode: one DummyClient per reader thread. Each thread's SHM + // segment and RPC channel are isolated to that thread's DummyClient, so + // multiple readers can run get_into / batch_get_into concurrently + // without contending on a single client (the per-thread setup matches + // the original Go test pattern that creates an independent store per + // goroutine / process). + std::vector> dummy_clients_; + // The client that owns the main buffer (buffer_). For real mode this is + // client_; for dummy mode this is primary_dummy_client_. Use this in + // single-threaded writer / warmup / verify code paths. + std::shared_ptr main_buffer_client_; char* buffer_; size_t buffer_size_; + // True if the underlying client is a DummyClient. DummyClient::register_buffer + // requires the memory to be inside a ShmHelper-managed segment (memfd+mmap), + // because the address+fd is later passed to the real client via IPC. A + // plain numa_alloc_local buffer would be rejected with "Buffer is not + // in any registered shared memory". Track this so Setup / destructor / + // AllocateThreadBuffers can pick the right allocator. + bool is_dummy_ = false; + + // Build and connect a single DummyClient. Returns nullptr on failure. + // Per-thread dummy clients are created with local_buffer_size=0 so they + // do not pre-allocate a SHM segment; AllocateThreadBuffers will then + // allocate the per-thread buffer via ShmHelper and register it. This + // gives every thread its own (SHM, RPC, IPC) triple and avoids sharing + // one client across threads. + std::shared_ptr CreateDummyClient(size_t local_buffer_size = 0) { + size_t mem_pool = FLAGS_dummy_mem_pool_size > 0 + ? FLAGS_dummy_mem_pool_size + : FLAGS_global_segment_size; + if (FLAGS_dummy_server_address.empty() || + FLAGS_dummy_ipc_socket_path.empty()) { + LOG(ERROR) + << "Dummy client requires non-empty --dummy_server_address " + "and --dummy_ipc_socket_path"; + return nullptr; + } + auto dummy = std::make_shared(); + int ret = dummy->setup_dummy(mem_pool, local_buffer_size, + FLAGS_dummy_server_address, + FLAGS_dummy_ipc_socket_path); + if (ret != 0) { + LOG(ERROR) << "DummyClient setup_dummy failed, ret=" << ret; + return nullptr; + } + return dummy; + } + + // Returns a NUMA-local buffer for RealClient, or a ShmHelper segment for + // DummyClient. Caller owns the buffer and must release it with + // FreeBuffer(). + char* AllocateBuffer(size_t size, int numa_node = -1) { + if (is_dummy_) { + try { + return static_cast( + mooncake::ShmHelper::getInstance()->allocate(size)); + } catch (const std::exception& e) { + LOG(ERROR) << "ShmHelper::allocate(" << size + << ") failed: " << e.what(); + return nullptr; + } + } + if (numa_node >= 0) { + return reinterpret_cast(numa_alloc_onnode(size, numa_node)); + } + return reinterpret_cast(numa_alloc_local(size)); + } + + // Counterpart of AllocateBuffer. + void FreeBuffer(char* ptr, size_t size) { + if (!ptr) return; + if (is_dummy_) { + if (mooncake::ShmHelper::getInstance()->free(ptr) != 0) { + LOG(WARNING) << "ShmHelper::free(" << ptr << ") failed"; + } + return; + } + numa_free(ptr, size); + } struct ThreadBuffer { char* ptr = nullptr; @@ -1511,32 +2656,146 @@ class StressBenchmark { int AllocateThreadBuffers(size_t num_threads) { thread_buffers_.resize(num_threads); + dummy_clients_.clear(); + dummy_clients_.reserve(num_threads); size_t per_buf_size = FLAGS_batch_size * FLAGS_value_size; for (size_t t = 0; t < num_threads; ++t) { int node = t % NR_SOCKETS; thread_buffers_[t].size = per_buf_size; thread_buffers_[t].numa_node = node; - thread_buffers_[t].ptr = - reinterpret_cast(numa_alloc_onnode(per_buf_size, node)); + + // Pick (or create) the client that will own this thread's + // buffer. For dummy mode, every thread gets its own + // DummyClient so its SHM and RPC channel are isolated; for + // real mode all threads share the single RealClient. + std::shared_ptr thread_client; + if (is_dummy_) { + auto dc = CreateDummyClient(/*local_buffer_size=*/0); + if (!dc) { + LOG(ERROR) << "Failed to create DummyClient for thread " + << t + << " (will roll back " << t + << " already-allocated thread(s))"; + RollbackThreadBuffers(t); + return -1; + } + dummy_clients_.push_back(dc); + thread_client = dc; + } else { + thread_client = client_; + } + + thread_buffers_[t].ptr = AllocateBuffer(per_buf_size, node); if (!thread_buffers_[t].ptr) { LOG(ERROR) << "Failed to allocate buffer for thread " << t - << " on NUMA node " << node; + << " on NUMA node " << node << " (" + << (is_dummy_ ? "ShmHelper" : "numa") + << "); will roll back " << t + << " already-allocated thread(s)"; + // The buffer failed to allocate, so this thread's dummy + // client owns no buffer; tear it down to release its + // RPC/IPC connection. + if (is_dummy_ && !dummy_clients_.empty()) { + try { + dummy_clients_.back()->tearDownAll(); + } catch (...) { + LOG(WARNING) + << "Failed to tearDownAll dummy client for " + << "thread " << t << ", ignoring"; + } + dummy_clients_.pop_back(); + } + RollbackThreadBuffers(t); return -1; } std::memset(thread_buffers_[t].ptr, 0, per_buf_size); - int ret = - client_->register_buffer(thread_buffers_[t].ptr, per_buf_size); + + int ret = thread_client->register_buffer( + thread_buffers_[t].ptr, per_buf_size); if (ret != 0) { LOG(ERROR) << "register_buffer failed for thread " << t - << " on NUMA node " << node; + << " on NUMA node " << node + << " (is_dummy=" << is_dummy_ + << "); will roll back " << t + << " already-allocated thread(s)"; + // Try to unregister with the same client (best effort). + try { + thread_client->unregister_buffer( + thread_buffers_[t].ptr); + } catch (...) { + LOG(WARNING) << "Best-effort unregister after " + "register_buffer failure for thread " + << t << " failed, ignoring"; + } + FreeBuffer(thread_buffers_[t].ptr, per_buf_size); + thread_buffers_[t].ptr = nullptr; + if (is_dummy_ && !dummy_clients_.empty()) { + try { + dummy_clients_.back()->tearDownAll(); + } catch (...) { + LOG(WARNING) + << "Failed to tearDownAll dummy client for " + << "thread " << t << ", ignoring"; + } + dummy_clients_.pop_back(); + } + RollbackThreadBuffers(t); return ret; } } LOG(INFO) << "Allocated " << num_threads << " thread buffers, each " - << per_buf_size / MB << " MB (NUMA-aware, " << NR_SOCKETS - << " sockets)"; + << per_buf_size / MB << " MB (" + << (is_dummy_ ? "per-thread ShmHelper, " + : "NUMA-aware, ") + << NR_SOCKETS << " sockets)"; return 0; } + + // Roll back all thread buffers in [0, count). Called when + // AllocateThreadBuffers fails partway through: we have to unregister + // each already-registered buffer with its owning dummy client (so the + // real client side releases the IPC fd) and free the SHM, then + // tearDownAll each dummy client (so its RPC + ping thread exit). + // After this call, thread_buffers_ and dummy_clients_ are empty. + void RollbackThreadBuffers(size_t count) { + for (size_t t = 0; t < count; ++t) { + auto& tb = thread_buffers_[t]; + if (!tb.ptr) continue; + std::shared_ptr thread_client; + if (is_dummy_ && t < dummy_clients_.size() && + dummy_clients_[t]) { + thread_client = dummy_clients_[t]; + } else { + thread_client = client_; + } + if (thread_client) { + try { + thread_client->unregister_buffer(tb.ptr); + } catch (...) { + LOG(WARNING) << "Rollback: failed to unregister " + "thread " + << t << " buffer, ignoring"; + } + } + FreeBuffer(tb.ptr, tb.size); + tb.ptr = nullptr; + } + // Free any DummyClient entries left (those whose buffers were + // never allocated or whose buffers we already cleaned up + // inline). tearDownAll stops the ping thread and closes the + // IPC / RPC channels. + for (auto& dc : dummy_clients_) { + if (!dc) continue; + try { + dc->tearDownAll(); + } catch (...) { + LOG(WARNING) << "Rollback: failed to tearDownAll dummy " + "client, ignoring"; + } + } + dummy_clients_.clear(); + thread_buffers_.clear(); + } }; int main(int argc, char* argv[]) { @@ -1553,6 +2812,11 @@ int main(int argc, char* argv[]) { LOG(INFO) << "Mooncake Stress Cluster Benchmark"; LOG(INFO) << " Scenario: " << FLAGS_scenario; LOG(INFO) << " Protocol: " << FLAGS_protocol; + LOG(INFO) << " Client type: " << FLAGS_client_type; + if (FLAGS_client_type == "dummy") { + LOG(INFO) << " Dummy server: " << FLAGS_dummy_server_address; + LOG(INFO) << " Dummy IPC path: " << FLAGS_dummy_ipc_socket_path; + } LOG(INFO) << " Value size: " << FLAGS_value_size / MB << " MB"; LOG(INFO) << " Num keys: " << FLAGS_num_keys; LOG(INFO) << " Batch size: " << FLAGS_batch_size; diff --git a/mooncake-store/go/examples/dummy_clients_test/main.go b/mooncake-store/go/examples/dummy_clients_test/main.go new file mode 100644 index 0000000000..069fbd53f3 --- /dev/null +++ b/mooncake-store/go/examples/dummy_clients_test/main.go @@ -0,0 +1,809 @@ +// Copyright 2024 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// multi_process_bench demonstrates multi-process benchmark with DummyClient. +package main + +/* +#cgo LDFLAGS: -L${SRCDIR}/../../../build/mooncake-store/src -Wl,-rpath,${SRCDIR}/../../../build/mooncake-store/src -lmooncake_store -lstdc++ -lnuma -lglog -lgflags -ljsoncpp -lcurl -luring +#cgo CFLAGS: -I/home/w00889253/Mooncake/mooncake-store/include + +#include +#include + +static int set_cpu_affinity(int pid, const unsigned long *mask, size_t size) { + return sched_setaffinity(pid, size, (cpu_set_t*)mask); +} + +static int get_cpu_affinity(int pid, unsigned long *mask, size_t size) { + return sched_getaffinity(pid, size, (cpu_set_t*)mask); +} +*/ +import "C" + +import ( + "flag" + "fmt" + "log" + "math" + "os" + "os/exec" + "runtime" + "strconv" + "strings" + "syscall" + "time" + "unsafe" + + "github.com/edsrzf/mmap-go" + + store "github.com/kvcache-ai/Mooncake/mooncake-store/go/mooncakestore" +) + +// CPU_SETSIZE is the maximum number of CPUs supported by Linux +const cpuSetSizeBytes = 128 + +// cpuBindResult 绑核结果信息 +type cpuBindResult struct { + Success bool + Pid int + TargetList string + ActualList string + BeforeCount int + AfterCount int + ErrorMessage string + ActualMatches bool +} + +// Global flags - 统一使用 uint64 类型避免类型不匹配问题 +var ( + flagWorkers = flag.Uint64("workers", 4, "Number of worker processes") + flagValueSize = flag.Uint64("value-size", 4*1024*1024, "Size of each value in bytes") + flagNumKeys = flag.Uint64("num-keys", 100, "Number of keys to write/read") + flagBatchSize = flag.Uint64("batch-size", 32, "Batch size for operations") + flagIsWorker = flag.Bool("worker", false, "Run as worker process") + flagWorkerID = flag.Uint64("worker-id", 0, "Worker process ID") + flagPhase = flag.String("phase", "write", "Phase: write/read/verify/all") + flagReplicaNum = flag.Uint64("replica-num", 1, "Number of replicas for each object") + flagWaitSeconds = flag.Uint64("wait-seconds", 5, "Seconds to wait between phases") + flagGlobalSize = flag.Uint64("global-size", 512*1024*1024, "Global segment size in bytes") + flagLocalBufferSize = flag.Uint64("local-buffer-size", 128*1024*1024, "Local buffer size in bytes") + flagPutBufferSize = flag.Uint64("put-buffer-size", 512*1024*1024, "Put buffer size in bytes") +) + +func main() { + flag.Parse() + + // 参数验证 + if err := validateConfig(); err != nil { + log.Fatalf("Configuration validation failed: %v", err) + } + log.Println("[INIT] Configuration validation passed") + + // 如果是子进程，执行子进程逻辑 + if *flagIsWorker { + log.Printf("[WORKER-%d] Starting worker process", *flagWorkerID) + runWorker(uint(*flagWorkerID)) + log.Printf("[WORKER-%d] Worker process completed", *flagWorkerID) + return + } + + // 主进程逻辑 + runMaster() +} + +// validateConfig 验证配置参数的有效性 +func validateConfig() error { + if *flagWorkers == 0 { + return fmt.Errorf("workers must be > 0, got %d", *flagWorkers) + } + if *flagValueSize == 0 { + return fmt.Errorf("value-size must be > 0, got %d", *flagValueSize) + } + if *flagNumKeys == 0 { + return fmt.Errorf("num-keys must be > 0, got %d", *flagNumKeys) + } + if *flagBatchSize == 0 { + return fmt.Errorf("batch-size must be > 0, got %d", *flagBatchSize) + } + if *flagGlobalSize == 0 { + return fmt.Errorf("global-size must be > 0, got %d", *flagGlobalSize) + } + if *flagLocalBufferSize == 0 { + return fmt.Errorf("local-buffer-size must be > 0, got %d", *flagLocalBufferSize) + } + if *flagPutBufferSize == 0 { + return fmt.Errorf("put-buffer-size must be > 0, got %d", *flagPutBufferSize) + } + if *flagPhase != "write" && *flagPhase != "read" && *flagPhase != "all" { + return fmt.Errorf("invalid phase: %s (must be write/read/all)", *flagPhase) + } + if *flagIsWorker && *flagWorkerID >= *flagWorkers { + return fmt.Errorf("invalid worker-id: %d (must be in [0, %d))", *flagWorkerID, *flagWorkers) + } + return nil +} + +// calculateKeyRange 计算每个 worker 负责的键范围 +func calculateKeyRange(workerID uint, totalWorkers, totalKeys uint64) (keyOffset uint64, keysPerWorker uint64) { + keysPerWorker = totalKeys / totalWorkers + keyOffset = uint64(workerID) * keysPerWorker + + // 处理余数，前 N 个 worker 多分配一个 key + remainder := totalKeys % totalWorkers + if uint64(workerID) < remainder { + keysPerWorker++ + keyOffset = uint64(workerID)*(totalKeys/totalWorkers) + uint64(workerID) + } else { + keyOffset = uint64(workerID)*(totalKeys/totalWorkers) + remainder + } + + return keyOffset, keysPerWorker +} + +func runMaster() { + log.Println("========================================") + log.Println(" Mooncake Multi-Process Benchmark ") + log.Println("========================================") + log.Println("[MASTER] Configuration Summary:") + log.Printf("[MASTER] Workers: %d", *flagWorkers) + log.Printf("[MASTER] Value Size: %d MB", *flagValueSize/(1024*1024)) + log.Printf("[MASTER] Num Keys: %d", *flagNumKeys) + log.Printf("[MASTER] Batch Size: %d", *flagBatchSize) + log.Printf("[MASTER] Replicas: %d", *flagReplicaNum) + log.Printf("[MASTER] Phase: %s", *flagPhase) + log.Printf("[MASTER] Wait Seconds: %d", *flagWaitSeconds) + log.Printf("[MASTER] Global Size: %d MB", *flagGlobalSize/(1024*1024)) + log.Printf("[MASTER] Local Buffer: %d MB", *flagLocalBufferSize/(1024*1024)) + log.Printf("[MASTER] Put Buffer: %d MB", *flagPutBufferSize/(1024*1024)) + log.Println("========================================") + + if *flagPhase == "all" || *flagPhase == "write" { + log.Println("\n--- Phase 1: WRITE PHASE ---") + runPhase("write") + log.Println("--- Write phase completed ---") + + if *flagWaitSeconds > 0 { + log.Printf("Waiting %d seconds...", *flagWaitSeconds) + time.Sleep(time.Duration(*flagWaitSeconds) * time.Second) + } + } + + if *flagPhase == "all" || *flagPhase == "read" { + log.Println("\n--- Phase 2: READ PHASE ---") + log.Println("(with data verification)") + runPhase("read") + log.Println("--- Read phase completed ---") + } + + log.Println("\n========================================") + log.Println(" Benchmark Completed ") + log.Println("========================================") +} + +func runPhase(phase string) { + var processes []*os.Process + startTime := time.Now() + + log.Printf("[MASTER] ===============================") + log.Printf("[MASTER] Starting %s phase", phase) + log.Printf("[MASTER] Workers: %d", *flagWorkers) + log.Printf("[MASTER] ===============================") + + // 启动所有 worker 进程 + for i := uint64(0); i < *flagWorkers; i++ { + log.Printf("[MASTER] Starting worker %d...", i) + cmd := exec.Command( + os.Args[0], + "-worker", + fmt.Sprintf("-worker-id=%d", i), + fmt.Sprintf("-workers=%d", *flagWorkers), + fmt.Sprintf("-value-size=%d", *flagValueSize), + fmt.Sprintf("-num-keys=%d", *flagNumKeys), + fmt.Sprintf("-batch-size=%d", *flagBatchSize), + fmt.Sprintf("-phase=%s", phase), + fmt.Sprintf("-replica-num=%d", *flagReplicaNum), + fmt.Sprintf("-wait-seconds=%d", *flagWaitSeconds), + fmt.Sprintf("-global-size=%d", *flagGlobalSize), + fmt.Sprintf("-local-buffer-size=%d", *flagLocalBufferSize), + fmt.Sprintf("-put-buffer-size=%d", *flagPutBufferSize), + ) + + cmd.Env = os.Environ() + cmd.Stdout = os.Stdout + cmd.Stderr = os.Stderr + + if err := cmd.Start(); err != nil { + log.Fatalf("[MASTER] Failed to start worker %d: %v", i, err) + } + + processes = append(processes, cmd.Process) + log.Printf("[MASTER] Worker %d started successfully (PID: %d)", i, cmd.Process.Pid) + } + + log.Printf("[MASTER] All %d workers started, waiting for completion...", *flagWorkers) + + // 等待所有进程完成 + var failedCount, successCount uint64 + for i, p := range processes { + log.Printf("[MASTER] Waiting for worker %d (PID: %d)...", i, p.Pid) + if _, err := p.Wait(); err != nil { + if exitErr, ok := err.(*exec.ExitError); ok { + if status, ok := exitErr.Sys().(syscall.WaitStatus); ok { + if status.Signaled() { + log.Printf("[MASTER] Worker %d (PID: %d) killed by signal %d", + i, p.Pid, status.Signal()) + } else { + log.Printf("[MASTER] Worker %d (PID: %d) exited with status %d: %v", + i, p.Pid, status.ExitStatus(), err) + } + } else { + log.Printf("[MASTER] Worker %d (PID: %d) failed: %v", i, p.Pid, err) + } + } else { + log.Printf("[MASTER] Worker %d (PID: %d) failed: %v", i, p.Pid, err) + } + failedCount++ + } else { + log.Printf("[MASTER] Worker %d (PID: %d) completed successfully", i, p.Pid) + successCount++ + } + } + + duration := time.Since(startTime) + log.Printf("[MASTER] ===============================") + log.Printf("[MASTER] %s phase summary:", phase) + log.Printf("[MASTER] Duration: %v", duration) + log.Printf("[MASTER] Success: %d/%d workers", successCount, *flagWorkers) + log.Printf("[MASTER] Failed: %d/%d workers", failedCount, *flagWorkers) + if failedCount > 0 { + log.Printf("[MASTER] WARNING: Some workers failed during %s phase", phase) + } + log.Printf("[MASTER] ===============================") +} + +func runWorker(workerID uint) { + log.Printf("[WORKER-%d] Starting worker initialization", workerID) + + // 1. 绑核 + log.Printf("[WORKER-%d] Step 1/5: Attempting to bind to CPU %d", workerID, workerID) + bindCPUForWorker(int(workerID)) + + // 2. 创建并配置 store + log.Printf("[WORKER-%d] Step 2/5: Creating DummyClient...", workerID) + s := createAndSetupStore(int(workerID)) + defer func() { + log.Printf("[WORKER-%d] Cleanup: Closing store...", workerID) + s.Close() + log.Printf("[WORKER-%d] Cleanup: Store closed successfully", workerID) + }() + + // 3. 分配并注册 buffer（返回 mmap.MMap 类型和指针） + bufBytes, _, cleanup := allocateAndRegisterBuffer(s, int(workerID)) + defer cleanup() // 关键：确保资源释放 + + // 4. 根据阶段执行操作 + log.Printf("[WORKER-%d] Step 4/5: Executing %s phase...", workerID, *flagPhase) + executePhase(s, bufBytes, int(workerID)) + + // 5. 完成 + log.Printf("[WORKER-%d] Step 5/5: Worker completed successfully", workerID) +} + +// createAndSetupStore 创建并配置 store +func createAndSetupStore(workerID int) *store.Store { + s, err := store.NewWithType(store.MOONCAKE_CLIENT_DUMMY) + if err != nil { + log.Fatalf("[WORKER-%d] Failed to create store: %v", workerID, err) + } + log.Printf("[WORKER-%d] DummyClient created successfully", workerID) + + metadataServer := envOrDefault("MC_METADATA_SERVER", "http://localhost:8080/metadata") + masterAddr := envOrDefault("MC_MASTER_ADDR", "localhost:50051") + localhostName := envOrDefault("MC_LOCALHOST", "localhost") + protocol := envOrDefault("MC_PROTOCOL", "tcp") + deviceName := envOrDefault("MC_DEVICENAME", "") + serverAddress := envOrDefault("MC_SERVERADDRESS", "") + ipcSocketPath := envOrDefault("MC_IPCSOCKETPATH", "") + + log.Printf("[WORKER-%d] Store configuration:", workerID) + log.Printf("[WORKER-%d] Metadata Server: %s", workerID, metadataServer) + log.Printf("[WORKER-%d] Master Address: %s", workerID, masterAddr) + log.Printf("[WORKER-%d] Protocol: %s", workerID, protocol) + log.Printf("[WORKER-%d] Global Size: %d MB", workerID, *flagGlobalSize/(1024*1024)) + log.Printf("[WORKER-%d] Local Buffer: %d MB", workerID, *flagLocalBufferSize/(1024*1024)) + log.Printf("[WORKER-%d] Put Buffer: %d MB", workerID, *flagPutBufferSize/(1024*1024)) + + err = s.Setup( + localhostName, + metadataServer, + *flagGlobalSize, + *flagLocalBufferSize, + protocol, + deviceName, + masterAddr, + *flagPutBufferSize, + serverAddress, + ipcSocketPath, + ) + if err != nil { + log.Fatalf("[WORKER-%d] Setup failed: %v", workerID, err) + } + log.Printf("[WORKER-%d] Store setup completed successfully", workerID) + + return s +} + +// allocateAndRegisterBuffer 分配并注册共享内存 buffer +func allocateAndRegisterBuffer(s *store.Store, workerID int) (buf []byte, ptr uintptr, cleanup func()) { + // 1. 优先使用 RealClient 已存在的预注册缓冲区（如果存在） + count, err := s.RegisteredBufferCount() + if err == nil && count > 0 { + bufInfo, err := s.RegisteredBufferAt(0) + if err == nil && bufInfo.Size > 0 { + ptr = bufInfo.Ptr + // 通过 unsafe 构造 []byte 切片，便于后续使用 + buf = unsafe.Slice((*byte)(unsafe.Pointer(ptr)), bufInfo.Size) + log.Printf("[WORKER-%d] Using pre-registered buffer: ptr=0x%x, size=%d bytes", workerID, ptr, bufInfo.Size) + // 清理函数：仅注销，不 unmapping（因为内存由 RealClient 管理） + cleanup = func() { + log.Printf("[WORKER-%d] Cleanup: Unregistering pre-registered buffer", workerID) + if err := s.UnregisterBuffer(ptr); err != nil { + log.Printf("[WORKER-%d] Warning: UnregisterBuffer failed: %v", workerID, err) + } + } + return + } + } + + // 2. 没有预注册缓冲区，则自行 mmap 匿名共享内存（RealClient 支持注册任意内存） + bufSize := *flagBatchSize * *flagValueSize + log.Printf("[WORKER-%d] No pre-registered buffer, allocating via mmap. Size: %d MB", workerID, bufSize/(1024*1024)) + + pageSize := int64(syscall.Getpagesize()) + alignedSize := (int64(bufSize) + pageSize - 1) / pageSize * pageSize + mmapBytes, err := syscall.Mmap(-1, 0, int(alignedSize), + syscall.PROT_READ|syscall.PROT_WRITE, + syscall.MAP_SHARED|syscall.MAP_ANONYMOUS) + if err != nil { + log.Fatalf("[WORKER-%d] Mmap failed: %v", workerID, err) + } + + ptr = uintptr(unsafe.Pointer(&mmapBytes[0])) + if err := s.RegisterBuffer(ptr, uint64(bufSize)); err != nil { + syscall.Munmap(mmapBytes) + log.Fatalf("[WORKER-%d] RegisterBuffer failed: %v", workerID, err) + } + log.Printf("[WORKER-%d] Buffer allocated via mmap: ptr=0x%x, size=%d bytes", workerID, ptr, bufSize) + + // 清理函数：注销并解除 mmap 映射 + cleanup = func() { + log.Printf("[WORKER-%d] Cleanup: Unregistering and unmapping buffer", workerID) + if err := s.UnregisterBuffer(ptr); err != nil { + log.Printf("[WORKER-%d] Warning: UnregisterBuffer failed: %v", workerID, err) + } + if err := syscall.Munmap(mmapBytes); err != nil { + log.Printf("[WORKER-%d] Warning: Munmap failed: %v", workerID, err) + } + } + return +} + +func allocateAndRegisterBufferViaMmap(s *store.Store, workerID int) (mmap.MMap, uintptr) { + bufSize := *flagBatchSize * *flagValueSize * 16 + log.Printf("[WORKER-%d] Buffer size required: %d MB", workerID, bufSize/(1024*1024)) + + pageSize := int64(syscall.Getpagesize()) + alignedSize := (int64(bufSize) + pageSize - 1) / pageSize * pageSize + log.Printf("[WORKER-%d] Aligned buffer size: %d bytes (page size: %d)", workerID, alignedSize, pageSize) + + filePath := fmt.Sprintf("/dev/shm/mooncake_dummy_client%d.dat", workerID) // 每个worker独立文件 + f, err := os.OpenFile(filePath, os.O_RDWR|os.O_CREATE, 0644) + if err != nil { + log.Fatalf("[WORKER-%d] Open file failed: %v", workerID, err) + } + defer f.Close() + + if err := f.Truncate(alignedSize); err != nil { + log.Fatalf("[WORKER-%d] Truncate file failed: %v", workerID, err) + } + + mm, err := mmap.MapRegion(f, int(alignedSize), mmap.RDWR, 0, 0) + if err != nil { + log.Fatalf("[WORKER-%d] mmap.MapRegion failed: %v", workerID, err) + } + + ptr := uintptr(unsafe.Pointer(&mm[0])) + if err := s.RegisterBuffer(ptr, uint64(bufSize)); err != nil { + mm.Unmap() + log.Fatalf("[WORKER-%d] RegisterBuffer failed: %v", workerID, err) + } + return mm, ptr +} + +// executePhase 根据阶段执行操作 +func executePhase(s *store.Store, buf []byte, workerID int) { + switch *flagPhase { + case "write": + log.Printf("[WORKER-%d] Starting write phase...", workerID) + writePhase(s, buf, workerID) + log.Printf("[WORKER-%d] Write phase completed", workerID) + case "read": + log.Printf("[WORKER-%d] Starting read phase...", workerID) + readPhase(s, buf, workerID) + log.Printf("[WORKER-%d] Read phase completed", workerID) + default: + log.Fatalf("[WORKER-%d] Unknown phase: %s", workerID, *flagPhase) + } +} + +func writePhase(s *store.Store, buf []byte, workerID int) { + keyOffset, keysPerWorker := calculateKeyRange(uint(workerID), *flagWorkers, *flagNumKeys) + valueSize := int(*flagValueSize) + + log.Printf("[WORKER-%d] Write phase: keys %d-%d (total %d keys)", + workerID, keyOffset, keyOffset+keysPerWorker-1, keysPerWorker) + + start := time.Now() + var written, failed uint64 + + for i := uint64(0); i < keysPerWorker; i++ { + keyIdx := keyOffset + i + key := MakeKey(keyIdx) + + // 填充数据 + FillBuffer(buf[:valueSize], keyIdx) + + // 使用 PutFrom 写入（零拷贝） + if err := s.PutFrom(key, uintptr(unsafe.Pointer(&buf[0])), *flagValueSize, nil); err != nil { + log.Printf("[WORKER-%d] PutFrom failed for key=%s: %v", workerID, key, err) + failed++ + } else { + written++ + } + + // 定期打印进度 + if (i+1)%10 == 0 || i == keysPerWorker-1 { + log.Printf("[WORKER-%d] Written %d/%d keys (key=%s)", + workerID, i+1, keysPerWorker, key) + } + } + + duration := time.Since(start) + throughput := float64(written*(*flagValueSize)) / duration.Seconds() / (1024 * 1024) + + log.Printf("[WORKER-%d] Write phase completed:", workerID) + log.Printf("[WORKER-%d] Written: %d", workerID, written) + log.Printf("[WORKER-%d] Failed: %d", workerID, failed) + log.Printf("[WORKER-%d] Duration: %v", workerID, duration) + log.Printf("[WORKER-%d] Throughput: %.2f MB/s", workerID, throughput) +} + +func readPhase(s *store.Store, buf []byte, workerID int) { + keyOffset, keysPerWorker := calculateKeyRange(uint(workerID), *flagWorkers, *flagNumKeys) + valueSize := int(*flagValueSize) + + log.Printf("[WORKER-%d] Read phase: keys %d-%d (total %d keys)", + workerID, keyOffset, keyOffset+keysPerWorker-1, keysPerWorker) + + start := time.Now() + var totalBytes, failedOps, verifyErrors, successOps, totalOps uint64 + + if *flagBatchSize > 1 { + // 批量读取 + batchSizeVal := *flagBatchSize + valueSize64 := *flagValueSize + + for i := uint64(0); i < keysPerWorker; i += batchSizeVal { + batchEnd := uint64(math.Min(float64(i+batchSizeVal), float64(keysPerWorker))) + batchSize := batchEnd - i + + batchStart := time.Now() + + for j := uint64(0); j < batchSize; j++ { + keyIdx := keyOffset + i + j + key := MakeKey(keyIdx) + offset := int(j * valueSize64) + + n, err := s.GetInto(key, uintptr(unsafe.Pointer(&buf[offset])), *flagValueSize) + totalOps++ + + if err != nil || n < 0 { + log.Printf("[WORKER-%d] GetInto failed for key=%s: %v", workerID, key, err) + failedOps++ + } else { + totalBytes += uint64(n) + + // 数据验证 + if !CheckBuffer(buf[offset:offset+valueSize], keyIdx) { + log.Printf("[WORKER-%d] Data mismatch for key=%s", workerID, key) + verifyErrors++ + } else { + successOps++ + } + } + } + + batchDuration := time.Since(batchStart) + if batchDuration.Seconds() > 0 { + batchThroughput := float64(batchSize*valueSize64) / batchDuration.Seconds() / (1024 * 1024) + log.Printf("[WORKER-%d] Batch %d-%d completed in %v (%.2f MB/s)", + workerID, i, batchEnd-1, batchDuration, batchThroughput) + } + } + } else { + // 单条读取 + for i := uint64(0); i < keysPerWorker; i++ { + keyIdx := keyOffset + i + key := MakeKey(keyIdx) + + n, err := s.GetInto(key, uintptr(unsafe.Pointer(&buf[0])), *flagValueSize) + totalOps++ + + if err != nil || n < 0 { + log.Printf("[WORKER-%d] GetInto failed for key=%s: %v", workerID, key, err) + failedOps++ + } else { + totalBytes += uint64(n) + + // 数据验证 + if !CheckBuffer(buf[:n], keyIdx) { + log.Printf("[WORKER-%d] Data mismatch for key=%s", workerID, key) + verifyErrors++ + } else { + successOps++ + } + } + + if (i+1)%10 == 0 { + log.Printf("[WORKER-%d] Read %d/%d keys", workerID, i+1, keysPerWorker) + } + } + } + + duration := time.Since(start) + throughput := float64(totalBytes) / duration.Seconds() / (1024 * 1024) + + log.Printf("[WORKER-%d] Read phase completed:", workerID) + log.Printf("[WORKER-%d] Total Bytes: %d", workerID, totalBytes) + log.Printf("[WORKER-%d] Total Ops: %d", workerID, totalOps) + log.Printf("[WORKER-%d] Success Ops: %d", workerID, successOps) + log.Printf("[WORKER-%d] Failed Ops: %d", workerID, failedOps) + log.Printf("[WORKER-%d] Verify Errors: %d", workerID, verifyErrors) + log.Printf("[WORKER-%d] Duration: %v", workerID, duration) + log.Printf("[WORKER-%d] Throughput: %.2f MB/s", workerID, throughput) +} + +// MakeKey 生成 key（模拟 C++ 的 MakeKey） +func MakeKey(idx uint64) string { + return fmt.Sprintf("bench_key_%d", idx) +} + +// FillBuffer 使用特定模式填充 buffer（模拟 C++ 的 FillBuffer） +func FillBuffer(buf []byte, seed uint64) { + pattern := seed * 0x9E3779B97F4A7C15 + numWords := len(buf) / 8 + + ptr := (*[1 << 30]uint64)(unsafe.Pointer(&buf[0])) + for i := 0; i < numWords; i++ { + ptr[i] = pattern + uint64(i) + } + + // 处理剩余字节 + remaining := len(buf) % 8 + if remaining > 0 { + lastWord := pattern + uint64(numWords) + for i := 0; i < remaining; i++ { + buf[numWords*8+i] = byte(lastWord >> (i * 8)) + } + } +} + +// CheckBuffer 验证数据完整性 +func CheckBuffer(buf []byte, seed uint64) bool { + if len(buf) == 0 { + return false + } + + pattern := seed * 0x9E3779B97F4A7C15 + numWords := len(buf) / 8 + + ptr := (*[1 << 30]uint64)(unsafe.Pointer(&buf[0])) + for i := 0; i < numWords; i++ { + if ptr[i] != pattern+uint64(i) { + log.Printf("CheckBuffer: Mismatch at word %d: expected %d, got %d", + i, pattern+uint64(i), ptr[i]) + return false + } + } + + // 验证剩余字节 + remaining := len(buf) % 8 + if remaining > 0 { + lastWord := pattern + uint64(numWords) + for i := 0; i < remaining; i++ { + expected := byte(lastWord >> (i * 8)) + if buf[numWords*8+i] != expected { + log.Printf("CheckBuffer: Mismatch at byte %d: expected %d, got %d", + numWords*8+i, expected, buf[numWords*8+i]) + return false + } + } + } + return true +} + +// bindCPUForWorker 为 worker 进程绑定 CPU +func bindCPUForWorker(workerID int) { + cpuList := fmt.Sprintf("%d", workerID) + result := bindProcessAffinity(cpuList) + + if result.Success { + log.Printf("[WORKER-%d] CPU binding succeeded: target=%s, actual=%s", + workerID, result.TargetList, result.ActualList) + } else { + log.Printf("[WORKER-%d] CPU binding failed: %s", workerID, result.ErrorMessage) + } +} + +// parseRangeToSet 解析 CPU 范围字符串到 set +func parseRangeToSet(s string) map[int]bool { + set := make(map[int]bool) + if s == "" { + return set + } + for _, part := range strings.Split(s, ",") { + part = strings.TrimSpace(part) + if part == "" { + continue + } + if strings.Contains(part, "-") { + r := strings.SplitN(part, "-", 2) + start, err1 := strconv.Atoi(strings.TrimSpace(r[0])) + end, err2 := strconv.Atoi(strings.TrimSpace(r[1])) + if err1 != nil || err2 != nil { + continue + } + for i := start; i <= end; i++ { + set[i] = true + } + } else { + cpu, err := strconv.Atoi(part) + if err == nil { + set[cpu] = true + } + } + } + return set +} + +// setToString 将 CPU 集合格式化为字符串 +func setToString(set map[int]bool) string { + if len(set) == 0 { + return "" + } + cpus := make([]int, 0, len(set)) + for cpu := range set { + cpus = append(cpus, cpu) + } + for i := 0; i < len(cpus); i++ { + for j := i + 1; j < len(cpus); j++ { + if cpus[j] < cpus[i] { + cpus[i], cpus[j] = cpus[j], cpus[i] + } + } + } + parts := make([]string, len(cpus)) + for i, cpu := range cpus { + parts[i] = strconv.Itoa(cpu) + } + return strings.Join(parts, ",") +} + +// getCurrentAffinityList 获取当前进程的 CPU 亲和性列表 +func getCurrentAffinityList() (string, int, error) { + mask := make([]C.ulong, cpuSetSizeBytes/8) + + result := C.get_cpu_affinity(0, &mask[0], C.size_t(cpuSetSizeBytes)) + if result != 0 { + return "", 0, fmt.Errorf("sched_getaffinity failed: result=%d", result) + } + + set := make(map[int]bool) + for i := 0; i < cpuSetSizeBytes*8; i++ { + wordIdx := i / (8 * 8) + bitIdx := i % (8 * 8) + if mask[wordIdx]&(1<= cpuSetSizeBytes*8 { + log.Printf("[CPU-BIND] Warning: CPU %d exceeds mask size %d, ignored", cpu, cpuSetSizeBytes*8) + continue + } + wordIdx := cpu / (8 * 8) + bitIdx := cpu % (8 * 8) + mask[wordIdx] |= 1 << C.ulong(bitIdx) + } + + cResult := C.set_cpu_affinity(0, &mask[0], C.size_t(cpuSetSizeBytes)) + if cResult != 0 { + result.ErrorMessage = fmt.Sprintf("sched_setaffinity failed: result=%d", cResult) + return result + } + + afterList, afterCount, err := getCurrentAffinityList() + if err != nil { + result.ErrorMessage = fmt.Sprintf("getaffinity after failed: %v", err) + return result + } + + result.Success = true + result.ActualList = afterList + result.AfterCount = afterCount + + actualSet := parseRangeToSet(afterList) + match := true + for cpu := range targetSet { + if !actualSet[cpu] { + match = false + break + } + } + if len(actualSet) != len(targetSet) { + match = false + } + result.ActualMatches = match + + return result +} + +// envOrDefault 获取环境变量，不存在则返回默认值 +func envOrDefault(key, defaultValue string) string { + if value := os.Getenv(key); value != "" { + return value + } + return defaultValue +} diff --git a/mooncake-store/go/mooncakestore/errors.go b/mooncake-store/go/mooncakestore/errors.go index 6c673e281f..c7cb0dcc02 100644 --- a/mooncake-store/go/mooncakestore/errors.go +++ b/mooncake-store/go/mooncakestore/errors.go @@ -31,4 +31,5 @@ var ( ErrBatchOp = errors.New("mooncakestore: batch operation failed") ErrHostname = errors.New("mooncakestore: get hostname failed") ErrInvalidArgument = errors.New("mooncakestore: invalid argument") + ErrBufferNotFound = errors.New("mooncakestore: buffer not found") ) diff --git a/mooncake-store/go/mooncakestore/store.go b/mooncake-store/go/mooncakestore/store.go index 1465bfa9c0..1a09d7e371 100644 --- a/mooncake-store/go/mooncakestore/store.go +++ b/mooncake-store/go/mooncakestore/store.go @@ -26,6 +26,11 @@ import ( "unsafe" ) +const ( + MOONCAKE_CLIENT_REAL = C.MOONCAKE_CLIENT_REAL + MOONCAKE_CLIENT_DUMMY = C.MOONCAKE_CLIENT_DUMMY +) + // Store wraps a Mooncake Store client handle. type Store struct { handle C.mooncake_store_t @@ -34,7 +39,13 @@ type Store struct { // New creates a new Store instance. Call Setup before performing operations, // and Close when done. func New() (*Store, error) { - h := C.mooncake_store_create() + return NewWithType(C.MOONCAKE_CLIENT_REAL) +} + +// NewWithType creates a new Store instance with the specified client type. +// Use MOONCAKE_CLIENT_REAL or MOONCAKE_CLIENT_DUMMY. +func NewWithType(clientType C.mooncake_client_type_t) (*Store, error) { + h := C.mooncake_store_create(clientType) if h == nil { return nil, ErrStoreNil } @@ -43,7 +54,7 @@ func New() (*Store, error) { // Setup initialises the store client and connects to the cluster. // -// Parameters: +// For real client (MOONCAKE_CLIENT_REAL), the relevant parameters are: // - localHostname: hostname/IP of this node // - metadataServer: metadata server URL (e.g. "http://host:8080/metadata") // - globalSegmentSize: size of the global memory segment in bytes @@ -51,9 +62,16 @@ func New() (*Store, error) { // - protocol: transport protocol ("tcp" or "rdma") // - deviceName: RDMA device name (empty for TCP or auto-discovery) // - masterServerAddr: master service address (e.g. "host:50051") +// +// For dummy client (MOONCAKE_CLIENT_DUMMY), the relevant parameters are: +// - localBufferSize: size of the local transfer buffer in bytes +// - memPoolSize: size of the memory pool in bytes +// - serverAddress: server address for dummy client +// - ipcSocketPath: IPC socket path for dummy client func (s *Store) Setup(localHostname, metadataServer string, globalSegmentSize, localBufferSize uint64, - protocol, deviceName, masterServerAddr string) error { + protocol, deviceName, masterServerAddr string, + memPoolSize uint64, serverAddress, ipcSocketPath string) error { if s.handle == nil { return ErrStoreNil } @@ -63,16 +81,21 @@ func (s *Store) Setup(localHostname, metadataServer string, cProtocol := C.CString(protocol) cDeviceName := C.CString(deviceName) cMasterAddr := C.CString(masterServerAddr) + cServerAddress := C.CString(serverAddress) + cIpcSocketPath := C.CString(ipcSocketPath) defer C.free(unsafe.Pointer(cLocalHostname)) defer C.free(unsafe.Pointer(cMetadataServer)) defer C.free(unsafe.Pointer(cProtocol)) defer C.free(unsafe.Pointer(cDeviceName)) defer C.free(unsafe.Pointer(cMasterAddr)) + defer C.free(unsafe.Pointer(cServerAddress)) + defer C.free(unsafe.Pointer(cIpcSocketPath)) ret := C.mooncake_store_setup(s.handle, cLocalHostname, cMetadataServer, C.uint64_t(globalSegmentSize), C.uint64_t(localBufferSize), - cProtocol, cDeviceName, cMasterAddr) + cProtocol, cDeviceName, cMasterAddr, + C.uint64_t(memPoolSize), cServerAddress, cIpcSocketPath) if ret != 0 { return ErrSetupFailed } @@ -538,3 +561,71 @@ func (s *Store) UnregisterBuffer(ptr uintptr) error { } return nil } + +// --------------------------------------------------------------------------- +// DummyClient-specific: Query registered buffers +// --------------------------------------------------------------------------- + +type RegisteredBufferInfo struct { + Ptr uintptr + Size uint64 + IsLocal bool +} + +func (s *Store) RegisteredBufferCount() (int, error) { + if s.handle == nil { + return 0, ErrStoreNil + } + count := C.mooncake_store_get_registered_buffer_count(s.handle) + return int(count), nil +} + +func (s *Store) RegisteredBufferAt(index int) (*RegisteredBufferInfo, error) { + if s.handle == nil { + return nil, ErrStoreNil + } + if index < 0 { + return nil, ErrInvalidArgument + } + var size C.size_t + ptr := C.mooncake_store_get_registered_buffer_at(s.handle, + C.size_t(index), &size) + if ptr == nil { + return nil, ErrBufferNotFound + } + return &RegisteredBufferInfo{ + Ptr: uintptr(ptr), + Size: uint64(size), + }, nil +} + +// UnregisterAllBuffers 注销所有已注册的 buffer +// 用于程序结束前清理资源，避免内存泄漏 +// 注意：此函数会遍历所有已注册 buffer 并逐个调用 UnregisterBuffer +func (s *Store) UnregisterAllBuffers() (int, error) { + if s.handle == nil { + return 0, ErrStoreNil + } + count, err := s.RegisteredBufferCount() + if err != nil { + return 0, err + } + unregistered := 0 + for i := 0; i < count; i++ { + bufInfo, err := s.RegisteredBufferAt(i) + if err != nil { + continue + } + if err := s.UnregisterBuffer(bufInfo.Ptr); err == nil { + unregistered++ + } + } + return unregistered, nil +} + +func (s *Store) IsHotCachePtr(ptr uintptr) bool { + if s.handle == nil || ptr == 0 { + return false + } + return C.mooncake_store_is_hot_cache_ptr(s.handle, unsafe.Pointer(ptr)) == 1 +} diff --git a/mooncake-store/include/dummy_client.h b/mooncake-store/include/dummy_client.h index d538708828..17cf07c0d2 100644 --- a/mooncake-store/include/dummy_client.h +++ b/mooncake-store/include/dummy_client.h @@ -130,6 +130,14 @@ class DummyClient : public PyClient { [[nodiscard]] std::string get_hostname() const; + struct RegisteredBufferInfo { + void *ptr; + size_t size; + bool is_local; + }; + + std::vector get_registered_buffers() const; + // Check if a pointer falls within the hot cache shm region bool is_hot_cache_ptr(const void *ptr) const { if (!hot_cache_base_) return false; diff --git a/mooncake-store/include/real_client.h b/mooncake-store/include/real_client.h index ed7be07634..52a62b6cc2 100644 --- a/mooncake-store/include/real_client.h +++ b/mooncake-store/include/real_client.h @@ -429,6 +429,10 @@ class RealClient : public PyClient { const std::vector &sizes, const ReplicateConfig &config, int32_t device_id, const UUID &client_id); + tl::expected put_from_dummy_helper( + const std::string &key, uint64_t dummy_buffer, size_t size, + const ReplicateConfig &config, int32_t device_id, const UUID &client_id); + std::vector> batch_put_from_multi_buffers_dummy_helper( const std::vector &keys, diff --git a/mooncake-store/include/store_c.h b/mooncake-store/include/store_c.h index b5646cd10f..c2e99a67fe 100644 --- a/mooncake-store/include/store_c.h +++ b/mooncake-store/include/store_c.h @@ -24,6 +24,12 @@ extern "C" { typedef void *mooncake_store_t; +enum mooncake_client_type { + MOONCAKE_CLIENT_REAL = 0, + MOONCAKE_CLIENT_DUMMY = 1, +}; +typedef enum mooncake_client_type mooncake_client_type_t; + struct mooncake_replicate_config { size_t replica_num; int with_soft_pin; @@ -45,7 +51,7 @@ typedef struct mooncake_replicate_config mooncake_replicate_config_t; // Lifecycle // --------------------------------------------------------------------------- -mooncake_store_t mooncake_store_create(); +mooncake_store_t mooncake_store_create(mooncake_client_type_t client_type); void mooncake_store_destroy(mooncake_store_t store); @@ -54,7 +60,10 @@ int mooncake_store_setup(mooncake_store_t store, const char *local_hostname, uint64_t global_segment_size, uint64_t local_buffer_size, const char *protocol, const char *device_name, - const char *master_server_addr); + const char *master_server_addr, + uint64_t mem_pool_size, + const char *server_address, + const char *ipc_socket_path); int mooncake_store_init_all(mooncake_store_t store, const char *protocol, const char *device_name, @@ -125,6 +134,17 @@ int mooncake_store_register_buffer(mooncake_store_t store, void *buffer, int mooncake_store_unregister_buffer(mooncake_store_t store, void *buffer); +// --------------------------------------------------------------------------- +// DummyClient-specific: Query registered buffers +// --------------------------------------------------------------------------- + +size_t mooncake_store_get_registered_buffer_count(mooncake_store_t store); + +void *mooncake_store_get_registered_buffer_at(mooncake_store_t store, + size_t index, size_t *size_out); + +int mooncake_store_is_hot_cache_ptr(mooncake_store_t store, const void *ptr); + #ifdef __cplusplus } #endif diff --git a/mooncake-store/src/dummy_client.cpp b/mooncake-store/src/dummy_client.cpp index c1a84a5c2a..ec767947c8 100644 --- a/mooncake-store/src/dummy_client.cpp +++ b/mooncake-store/src/dummy_client.cpp @@ -1106,6 +1106,20 @@ std::string DummyClient::get_hostname() const { return ""; } +std::vector +DummyClient::get_registered_buffers() const { + std::vector result; + if (!shm_helper_) return result; + + const auto& shms = shm_helper_->get_shms(); + for (const auto& shm : shms) { + if (shm->registered && shm->base_addr && shm->size > 0) { + result.push_back({shm->base_addr, shm->size, shm->is_local}); + } + } + return result; +} + std::vector DummyClient::batch_put_from( const std::vector& keys, const std::vector& buffer_ptrs, const std::vector& sizes, const ReplicateConfig& config) { @@ -1133,8 +1147,10 @@ std::vector DummyClient::batch_put_from( int DummyClient::put_from(const std::string& key, void* buffer, size_t size, const ReplicateConfig& config) { - // TODO: implement this function - return -1; + uint64_t buf_addr = reinterpret_cast(buffer); + auto result = invoke_rpc<&RealClient::put_from_dummy_helper, void>( + key, buf_addr, size, config, device_id_, client_id_); + return to_py_ret(result); } std::vector DummyClient::batch_get_into( diff --git a/mooncake-store/src/real_client.cpp b/mooncake-store/src/real_client.cpp index 4e67ecc78f..4d4305a9a3 100644 --- a/mooncake-store/src/real_client.cpp +++ b/mooncake-store/src/real_client.cpp @@ -4073,6 +4073,33 @@ RealClient::batch_put_from_dummy_helper( return batch_put_from_internal(keys, buffers_result.value(), sizes, config); } +tl::expected RealClient::put_from_dummy_helper( + const std::string &key, uint64_t dummy_buffer, size_t size, + const ReplicateConfig &config, int32_t device_id, const UUID &client_id) { +#ifdef USE_ASCEND_DIRECT + if (!ContextManager::getInstance().setCurrentContextByPhysicalId( + device_id)) { + LOG(ERROR) << "Failed to set context for physical device " << device_id; + return tl::unexpected(ErrorCode::INVALID_PARAMS); + } +#endif + + std::shared_lock lock(dummy_client_mutex_); + auto it = shm_contexts_.find(client_id); + if (it == shm_contexts_.end()) { + LOG(ERROR) << "client_id=" << client_id << ", error=shm_not_mapped"; + return tl::unexpected(ErrorCode::INVALID_PARAMS); + } + auto &context = it->second; + + auto buffers_result = map_dummy_addrs_to_real_ptrs( + context, {dummy_buffer}, {size}, client_id); + if (!buffers_result) { + return tl::unexpected(buffers_result.error()); + } + return put_from_internal(key, buffers_result.value()[0], size, config); +} + std::vector> RealClient::batch_put_from_internal( const std::vector &keys, const std::vector &buffers, const std::vector &sizes, const ReplicateConfig &config) { @@ -5903,11 +5930,29 @@ void RealClient::dummy_client_monitor_func() { // Update the client status to NEED_REMOUNT if (!expired_clients.empty()) { for (auto &client_id : expired_clients) { + { + std::shared_lock lock( + dummy_client_mutex_); + if (shm_contexts_.find(client_id) == + shm_contexts_.end()) { + LOG(INFO) + << "client_id=" << client_id + << ", action=shm_already_unmapped_by_other_path"; + continue; + } + } + // Unmap mapped_shms associated with this client + tl::expected result; if (globalConfig().ascend_agent_mode) { - ascend_unmap_shm_internal(client_id); + result = ascend_unmap_shm_internal(client_id); } else { - unmap_shm_internal(client_id); + result = unmap_shm_internal(client_id); + } + if (!result) { + // Client already unmapped (e.g., by other thread or earlier cleanup) + LOG(INFO) << "client_id=" << client_id + << ", action=client_already_unmapped"; } } } diff --git a/mooncake-store/src/real_client_main.cpp b/mooncake-store/src/real_client_main.cpp index 8f0a4ba80c..41519c1681 100644 --- a/mooncake-store/src/real_client_main.cpp +++ b/mooncake-store/src/real_client_main.cpp @@ -43,6 +43,7 @@ void RegisterClientRpcService(coro_rpc::coro_rpc_server &server, server.register_handler<&RealClient::getSize_internal>(&real_client); server.register_handler<&RealClient::batch_put_from_dummy_helper>( &real_client); + server.register_handler<&RealClient::put_from_dummy_helper>(&real_client); server.register_handler< &RealClient::batch_put_from_multi_buffers_dummy_helper>(&real_client); server.register_handler<&RealClient::upsert_dummy_helper>(&real_client); diff --git a/mooncake-store/src/store_c.cpp b/mooncake-store/src/store_c.cpp index a3a19086ef..45e3c7310f 100644 --- a/mooncake-store/src/store_c.cpp +++ b/mooncake-store/src/store_c.cpp @@ -21,6 +21,7 @@ #include #include +#include "dummy_client.h" #include "real_client.h" #include "replica.h" #include "types.h" @@ -31,7 +32,8 @@ namespace { // RealClient::create() returns shared_ptr and registers a weak_ptr in // ResourceTracker, so we must not let the shared_ptr die prematurely. struct StoreHandle { - std::shared_ptr client; + std::shared_ptr client; + mooncake_client_type_t client_type; }; inline const char *c_str_or(const char *s, const char *fallback) { @@ -62,7 +64,7 @@ StoreHandle *as_handle(mooncake_store_t store) { return static_cast(store); } -mooncake::RealClient *as_client(mooncake_store_t store) { +mooncake::PyClient *as_client(mooncake_store_t store) { return as_handle(store)->client.get(); } @@ -74,11 +76,17 @@ extern "C" { // Lifecycle // --------------------------------------------------------------------------- -mooncake_store_t mooncake_store_create() { +mooncake_store_t mooncake_store_create(mooncake_client_type_t client_type) { try { - auto client = mooncake::RealClient::create(); + std::shared_ptr client; + if (client_type == MOONCAKE_CLIENT_DUMMY) { + client = std::make_shared(); + } else { + client = mooncake::RealClient::create(); + } if (!client) return nullptr; - auto *handle = new (std::nothrow) StoreHandle{std::move(client)}; + auto *handle = + new (std::nothrow) StoreHandle{std::move(client), client_type}; if (!handle) return nullptr; return static_cast(handle); } catch (...) { @@ -103,14 +111,27 @@ int mooncake_store_setup(mooncake_store_t store, const char *local_hostname, uint64_t global_segment_size, uint64_t local_buffer_size, const char *protocol, const char *device_name, - const char *master_server_addr) { + const char *master_server_addr, + uint64_t mem_pool_size, + const char *server_address, + const char *ipc_socket_path) { if (!store) return -1; try { - return as_client(store)->setup_real( - c_str_or(local_hostname, ""), c_str_or(metadata_server, ""), - global_segment_size, local_buffer_size, c_str_or(protocol, "tcp"), - c_str_or(device_name, ""), - c_str_or(master_server_addr, "127.0.0.1:50051")); + auto *handle = as_handle(store); + if (handle->client_type == MOONCAKE_CLIENT_DUMMY) { + return handle->client->setup_dummy( + mem_pool_size, local_buffer_size, + c_str_or(server_address, ""), + c_str_or(ipc_socket_path, "")); + } else { + return handle->client->setup_real( + c_str_or(local_hostname, ""), + c_str_or(metadata_server, ""), + global_segment_size, local_buffer_size, + c_str_or(protocol, "tcp"), c_str_or(device_name, ""), + c_str_or(master_server_addr, "127.0.0.1:50051"), + nullptr, "", false, ""); + } } catch (...) { return -1; } @@ -358,4 +379,54 @@ int mooncake_store_unregister_buffer(mooncake_store_t store, void *buffer) { } } +// --------------------------------------------------------------------------- +// DummyClient-specific: Query registered buffers +// --------------------------------------------------------------------------- + +size_t mooncake_store_get_registered_buffer_count(mooncake_store_t store) { + if (!store) return 0; + try { + auto *handle = as_handle(store); + if (handle->client_type != MOONCAKE_CLIENT_DUMMY) { + return 0; + } + auto *dummy = static_cast(handle->client.get()); + return dummy->get_registered_buffers().size(); + } catch (...) { + return 0; + } +} + +void *mooncake_store_get_registered_buffer_at(mooncake_store_t store, + size_t index, size_t *size_out) { + if (!store) return nullptr; + try { + auto *handle = as_handle(store); + if (handle->client_type != MOONCAKE_CLIENT_DUMMY) { + return nullptr; + } + auto *dummy = static_cast(handle->client.get()); + auto buffers = dummy->get_registered_buffers(); + if (index >= buffers.size()) return nullptr; + if (size_out) *size_out = buffers[index].size; + return buffers[index].ptr; + } catch (...) { + return nullptr; + } +} + +int mooncake_store_is_hot_cache_ptr(mooncake_store_t store, const void *ptr) { + if (!store || !ptr) return 0; + try { + auto *handle = as_handle(store); + if (handle->client_type != MOONCAKE_CLIENT_DUMMY) { + return 0; + } + auto *dummy = static_cast(handle->client.get()); + return dummy->is_hot_cache_ptr(ptr) ? 1 : 0; + } catch (...) { + return 0; + } +} + } // extern "C" diff --git a/mooncake-store/tests/dummy_client_get_buffer_test.cpp b/mooncake-store/tests/dummy_client_get_buffer_test.cpp index b1b1dc5328..4ddd51532d 100644 --- a/mooncake-store/tests/dummy_client_get_buffer_test.cpp +++ b/mooncake-store/tests/dummy_client_get_buffer_test.cpp @@ -43,6 +43,7 @@ static void RegisterRpcHandlers(coro_rpc::coro_rpc_server &server, server.register_handler<&RealClient::getSize_internal>(&rc); server.register_handler<&RealClient::batch_get_into_dummy_helper>(&rc); server.register_handler<&RealClient::batch_put_from_dummy_helper>(&rc); + server.register_handler<&RealClient::put_from_dummy_helper>(&rc); server.register_handler<&RealClient::acquire_hot_cache>(&rc); server.register_handler<&RealClient::release_hot_cache>(&rc); server.register_handler<&RealClient::batch_acquire_hot_cache>(&rc);