feat(store): dedicated heartbeat RPC server to prevent head-of-line blocking by staugust · Pull Request #3541 · kvcache-ai/Mooncake

staugust · 2026-08-20T02:57:01Z

Description

Summary

Add an optional dedicated heartbeat RPC server to the Mooncake Store master, so heartbeat requests are no longer head-of-line-blocked by heavy metadata RPCs on the main server.

When --heartbeat_rpc_port is set (> 0), the master spins up a separate coro_rpc_server (plain TCP, no IB) that only serves Heartbeat. The client routes heartbeat traffic to <master_host>:heartbeat_rpc_port while all other RPCs continue to use the main server.

Smooth migration with dual-mode

A --heartbeat_keep_on_main flag allows Heartbeat to be registered on both the dedicated and main servers simultaneously. This lets legacy clients (unaware of the new port) and migrated clients coexist during rollout. Set false once all clients are migrated to drop Heartbeat from the main server.

Key changes

Master: serve Heartbeat on a dedicated coro_rpc_server when configured
Client: route Heartbeat to the dedicated port via a separate connection pool
Config plumbing: expose heartbeat_rpc_port through CLI flags, config dict, Python bindings, and environment variable (MOONCAKE_HEARTBEAT_RPC_PORT)
Dual-mode: heartbeat_keep_on_main keeps Heartbeat on the main server too during migration
Tests: verify routing (dedicated port, main port, dual-mode, and legacy fallback)

Module

Type of Change

How Has This Been Tested?

Test commands:

ctest -R 'master_service_test|non_ha_reconnect_test|ha_integration_test|client_integration_test'

Test results:

Unit tests pass
Integration tests pass (if applicable)
Manual testing done (describe below)

Manual tests

Firstly, start master with dedicate heartbeat rpc server with command:

  mooncake_master \
    --rpc_port=50051 \
    --rpc_thread_num=1 \
    --client_ttl=3 \
    --heartbeat_rpc_port=50099 \
    --heartbeat_keep_on_main \
    --heartbeat_rpc_thread_num=1 \
    --enable_http_metadata_server \
    --http_metadata_server_port=8080 &> master.log

Then, start real_client with dedicate heartbeat rpc port:

  GLOG_v=1 /usr/local/bin/mooncake_client \
    --host=127.0.0.1 \
    --master_server_address=127.0.0.1:50051 \
    --metadata_server=http://127.0.0.1:8080/metadata \
    --protocol=tcp \
    --http_port=19002 \
    --heartbeat_rpc_port=50099 \
    --global_segment_size=16MB \
    --port=50052 &> heartbeat.log

Finally, running master_bench to flood mooncake_master:

for i in $(seq 24); do
  echo client $i started
/home/local/workspace/augusto.yjh/kvpool/Mooncake/build/mooncake-store/benchmarks/master_bench \
    --master_server=127.0.0.1:50051 \
    --num_clients 16 \
    --num_threads=4 \
    --operation=BatchPut \
    --duration=120 &> bench.log &
done
wait

From the log of mooncake_master, we can see that there're many client disconnection log lines.

And mooncake_client that uses dedicate heartbeat rpc port is not disconnected.

Logs of mooncake_client show that, with dedicate heartbeat rpc port, its heartbeat rpc is not blocked by rps from master_bench instances.

Checklist

I have performed a self-review of my own code
I have formatted my code using ./scripts/code_format.sh
I have run pre-commit on the files changed in this PR and all hooks pass
I have updated the documentation (if applicable)
I have added tests to prove my changes are effective
For changes >500 LOC: I have filed an RFC issue

AI Assistance Disclosure

No AI tools were used
AI tools were used (specify below)
claude code with deepseek v4 pro

Extend InProcMaster/InProcP2PMaster test helpers and InProcMasterConfig to optionally start the dedicated heartbeat coro_rpc_server (mirroring master.cpp/ ha_helper), and register the main service with include_heartbeat=false when enabled. Add heartbeat_dedicated_port_test with four deterministic cases: - Heartbeat routed to the dedicated port succeeds. - Heartbeat is NOT registered on the main port (include_heartbeat=false). - Non-heartbeat RPCs (ServiceReady) still served on the main port. - With the dedicated server disabled (port=0), Heartbeat falls back to the main port (legacy behavior). Co-Authored-By: Claude <noreply@anthropic.com>

… builder Thread the dedicated heartbeat port through the standard real-client deployment path so heartbeats can be isolated from heavy metadata RPCs without custom code. - ClientConfigBuilder::build_{centralized,p2p}_real_client gain a trailing heartbeat_rpc_port param (default 0, backward compatible) that sets the RealClientConfigBase field. - real_client_main gains --heartbeat_rpc_port flag, passed to both builder call sites. When 0 (default) behavior is unchanged. When >0 it must match the master's --heartbeat_rpc_port; the real client then sends heartbeats to <master host>:heartbeat_rpc_port. The pybind/integrated RealClient::create path and the JSON config-map overload still default to 0 (follow-up). Co-Authored-By: Claude <noreply@anthropic.com>

…dict Make the in-process store client (the path sglang HiCache uses in integrated mode) configurable for the dedicated heartbeat port, closing the engine-side gap. - client_config_builder.h: add DictCommon::kHeartbeatRpcPort key ("heartbeat_rpc_port") + default 0; both config-map overloads (build_centralized_real_client / build_p2p_real_client) parse it and forward to the positional builder. - store_py.cpp: positional setup() / P2P setup gain a heartbeat_rpc_port kwarg (default 0), applied to the config after build; the dict-based setup()/setup_p2p_real_client() paths pick it up automatically via the map key. Docstrings updated. So an engine can configure it either way: store.setup({... "heartbeat_rpc_port": "50099" ...}) # dict/JSON path store.setup(..., heartbeat_rpc_port=50099) # positional path Default 0 = legacy behavior (heartbeat on the main RPC port). Must match the master's --heartbeat_rpc_port. Note: sglang-repo-side config plumbing must forward the heartbeat_rpc_port key to Mooncake's setup dict (out of this repo's scope). Co-Authored-By: Claude <noreply@anthropic.com>

…onfig Close the engine-service path: MooncakeStoreService (which uses the pybind store in integrated mode) can now configure the dedicated heartbeat port end-to-end from JSON config or env, and forwards it to store.setup(). - mooncake_config.py: add heartbeat_rpc_port field (default 0); read from JSON ("heartbeat_rpc_port") and from MOONCAKE_HEARTBEAT_RPC_PORT env. - mooncake_store_service.py: pass heartbeat_rpc_port= to store.setup(). - test_mooncake_config.py: cover default, from-file, and from-env. Default 0 = legacy behavior. Must match the master's --heartbeat_rpc_port. Co-Authored-By: Claude <noreply@anthropic.com>

Enable a zero-downtime upgrade to the dedicated-heartbeat-port mode. Problem: once the master enables heartbeat_rpc_port=P it removes Heartbeat from the main server, so any not-yet-migrated client (heartbeat_rpc_port=0) that still sends heartbeat to the main port fails. Coordinating a master+clients flip across a fleet is brittle. Add a transitional master flag heartbeat_keep_on_main (default false). When true and heartbeat_rpc_port > 0, Heartbeat is ALSO kept registered on the main server while the dedicated heartbeat server runs, so legacy clients (heartbeat_rpc_port=0) and migrated clients (heartbeat_rpc_port=P) both work. Set it back to false once all clients are migrated to drop heartbeat from the main server and finish isolation. - master_config.h: add heartbeat_keep_on_main to MasterConfig, supervisor config (optional, copied in FromConfig), and InProcMasterConfig/builder. - master.cpp / ha_helper.cpp: add gflag + config read; main server registers Heartbeat when (!dedicated || keep_on_main). Dedicated server still always registers Heartbeat when dedicated is on. - test helpers (InProcMaster / InProcP2PMaster): thread keep_on_main into the include_heartbeat decision. - heartbeat_dedicated_port_test: add HeartbeatDedicatedPortDualModeTest asserting heartbeat succeeds on BOTH the dedicated and the main port in dual mode. Rollout: Phase0 deploy new code (flag=0) -> Phase1 master heartbeat_rpc_port=P + heartbeat_keep_on_main=true -> Phase2 migrate clients to P in batches -> Phase3 set heartbeat_keep_on_main=false. Each phase is independently rollback-able by reverting the flag. Co-Authored-By: Claude <noreply@anthropic.com>

CentralizedMasterClient has no ServiceReady() method (ServiceReady is a server-side handler). Replace the non-heartbeat-RPC check with ExistKey, which is a client-exposed data-path RPC. ExistKey does a metadata lookup (returns false for a missing key, no client registration required), so has_value() proves the RPC reached the main-port handler. Co-Authored-By: Claude <noreply@anthropic.com>

SegmentClient constructor called MountSegment before registering the client with the master. Since commit 48d3a47, MountSegment checks ClientManager for the client_id and returns CLIENT_NOT_FOUND if the client hasn't been explicitly registered. Add RegisterClient() call after Connect() and before MountSegment() to fix this regression. Co-Authored-By: Claude <noreply@anthropic.com>

The heartbeat_rpc_port / heartbeat_rpc_thread_num / heartbeat_keep_on_main flags were registered and read in InitMasterConf (config-file path) but missing from LoadConfigFromCmdline (CLI path). Consequence: without --config_path, InitMasterConf never runs and LoadConfigFromCmdline didn't read these flags, so e.g. `mooncake_master --heartbeat_rpc_port=50099` (no config file) was silently ignored -- master_config.heartbeat_rpc_port stayed 0 and no dedicated heartbeat server started. With a config file, an explicit CLI value also couldn't override a file-present value. Add the three GetCommandLineFlagInfo blocks mirroring every other flag, and print the three values in the startup config log so the effective setting is visible. Co-Authored-By: Claude <noreply@anthropic.com>

Reformat overlong function call lines to respect column limits. No logic changes - pure reformatting of static_cast, RegisterRpcService, invoke_rpc_via, and builder chains across client_config_builder.h, ha_helper.cpp, master.cpp, master_client.cpp, rpc_service.cpp, and tests. Co-Authored-By: Claude <noreply@anthropic.com>

Drop the heartbeat_keep_on_main rollout flag entirely. With a dedicated heartbeat_rpc_port configured, the main RPC server no longer serves Heartbeat (main_includes_heartbeat = !dedicated_heartbeat); the port=0 default retains the legacy fallback of serving Heartbeat on the main server so existing default-config deployments and the reconnect/HA tests keep working. Removes the dual-mode test; keeps the dedicated-port and legacy-fallback routing tests. Co-Authored-By: Claude <noreply@anthropic.com>

Apply clang-format line wrapping and fix missing trailing newline in heartbeat_dedicated_port_test.cpp. Co-Authored-By: Claude <noreply@anthropic.com>

The d746bc3 refactor that removed heartbeat_keep_on_main rewrote the set_heartbeat_keep_on_main method body into set_heartbeat_rpc_thread_num instead of deleting it. A set_heartbeat_rpc_thread_num already existed immediately above, leaving two identical member-function definitions in InProcMasterConfigBuilder — a C++ redefinition error that breaks any translation unit including master_config.h. Remove the stray duplicate; the original (95124e2) is kept. Co-Authored-By: Claude <noreply@anthropic.com>

When a client is configured with a dedicated heartbeat_rpc_port but the master was started without one (legacy fallback), Connect previously succeeded by probing only the main server, then heartbeats silently TCP-refused forever on the heartbeat timer until the client got reaped (client_live_ttl expiry, segment reclaim). Detect this at Connect time: - Register ServiceReady on the dedicated heartbeat server (RegisterHeartbeatRpcService) so clients have a lightweight probe handler. ServiceReady is trivial (returns version) and does not undermine the HoL-isolation the dedicated server provides. - In MasterClient::Connect, after the main ServiceReady + version check, when heartbeat_rpc_port_ > 0 probe the heartbeat endpoint (heartbeat_accessor_) with ServiceReady. On failure return the new HEARTBEAT_RPC_UNREACHABLE error code and clear the cached address (fail-fast / strict mode). heartbeat_rpc_port_ == 0 skips the probe, so legacy deployments are unaffected. Add HeartbeatPortMismatchFailsConnect test covering the misconfig; the existing HeartbeatRoutedToDedicatedPort test now also implicitly proves ServiceReady is registered on the dedicated port (Connect probe would fail otherwise). Co-Authored-By: Claude <noreply@anthropic.com>

The main-pool ServiceReady probe in MasterClient::Connect retries once when is_same_addr to force a stale connection pool to re-establish (master_client.cpp:140-145). The dedicated-heartbeat probe added in b205bdb had no such retry, so on a reconnect to the same master_addr the reused heartbeat_accessor_ pool could hold a stale connection and the probe would fail with HEARTBEAT_RPC_UNREACHABLE even though the heartbeat server was fine — a false-positive fail-fast. Mirror the main path: retry the heartbeat probe once, but only when is_same_addr (the stale-pool scenario). When !is_same_addr a fresh pool was just obtained, so a failure is a real client/master mismatch and stays fail-fast. Co-Authored-By: Claude <noreply@anthropic.com>

Fill three coverage gaps around the dedicated heartbeat port feature: 1. Reconnect after master restart (is_same_addr=true stale-pool retry): new HeartbeatDedicatedPortReconnectTest starts a dedicated-port master, Connects, stops+restarts the master on the same ports, then Connects again to the same address. Both the main and heartbeat pools are stale after the restart; Connect must rebuild and succeed (not spuriously return HEARTBEAT_RPC_UNREACHABLE). Also asserts Heartbeat still lands afterward. This guards the is_same_addr retry added for the heartbeat probe in 95ab6dc, which previously had zero coverage. 2. P2P mode: parameterize the routing tests (1-3), the legacy-fallback test (4), and the mismatch fail-fast test (5) over CENTRALIZED/P2P via a small HeartbeatTestMaster wrapper + MakeClient factory. Both modes share RegisterHeartbeatRpcService and MasterClient, but this guards against P2P's RegisterP2PRpcService registration path breaking the dedicated-port routing and exercises P2PMasterClient end-to-end. 3. Error code serialization: assert toString(HEARTBEAT_RPC_UNREACHABLE) == "HEARTBEAT_RPC_UNREACHABLE" and the toInt/fromInt round-trip. Co-Authored-By: Claude <noreply@anthropic.com>

…tbeatServiceReady The previous fail-fast (b205bdb) only caught one direction: client configured with a dedicated heartbeat port while the master ran in legacy mode. The reverse — client legacy (heartbeat_rpc_port=0) while the master enabled a dedicated port — still silently starved: Connect succeeded by probing only the main server, then heartbeats failed forever on a main server that had dropped the Heartbeat handler, until the client was reaped. Add a HeartbeatServiceReady RPC on the main server that reports the master's heartbeat_rpc_port. Connect queries it and compares against the client's heartbeat_rpc_port_, catching both mismatch directions at startup: - client dedicated, master legacy (forward) -> HEARTBEAT_ROUTING_MISMATCH - client legacy, master dedicated (reverse) -> HEARTBEAT_ROUTING_MISMATCH - both dedicated, ports differ -> dedicated reachability probe fails -> HEARTBEAT_RPC_UNREACHABLE - both legacy / both dedicated & reachable -> OK HEARTBEAT_ROUTING_MISMATCH (new, -902) covers any dedicated-vs-legacy mismatch; HEARTBEAT_RPC_UNREACHABLE (-901) is retained for a configured-but-unreachable dedicated server. Server: - WrappedMasterServiceConfig carries heartbeat_rpc_port (propagated from MasterConfig and MasterServiceSupervisorConfig, incl. HA path). - WrappedMasterService stores it and implements HeartbeatServiceReady, registered on the main server (always) so legacy clients can reach it. HA registration reuses the same path. Client: - MasterClient::Connect calls HeartbeatServiceReady (main pool) after the version handshake; on RPC failure returns the error (requires master-first rollout, consistent with the dedicated-port feature's coordinated deployment). The dedicated-server ServiceReady reachability probe (with is_same_addr stale-pool retry) now runs only when both sides are dedicated. Tests (heartbeat_dedicated_port_test, parameterized CENTRALIZED/P2P): - Reverse mismatch: legacy client vs dedicated master fails Connect. - Dedicated-port value mismatch: both dedicated, wrong port -> HEARTBEAT_RPC_UNREACHABLE. - Forward mismatch expectation updated to HEARTBEAT_ROUTING_MISMATCH. - Error-code round-trip covers both new codes. Co-Authored-By: Claude <noreply@anthropic.com>

codecov-commenter · 2026-08-20T19:31:09Z

⚠️ Please install the to ensure uploads and comments are reliably processed by Codecov.

Codecov Report

❌ Patch coverage is 74.39024% with 84 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
mooncake-store/src/master.cpp	0.00%	30 Missing ⚠️
mooncake-store/src/ha_helper.cpp	0.00%	24 Missing ⚠️
mooncake-store/benchmarks/master_bench.cpp	0.00%	7 Missing ⚠️
mooncake-store/src/master_client.cpp	86.36%	6 Missing ⚠️
...cake-store/tests/heartbeat_dedicated_port_test.cpp	95.23%	6 Missing ⚠️
mooncake-integration/store/store_py.cpp	0.00%	5 Missing ⚠️
mooncake-store/include/master_config.h	55.55%	4 Missing ⚠️
mooncake-store/tests/test_p2p_server_helpers.h	96.00%	1 Missing ⚠️
mooncake-store/tests/test_server_helpers.h	96.00%	1 Missing ⚠️

📢 Thoughts on this report? Let us know!

…Connect A non-coroutine RPC handler runs on the io thread that read it. When every main io thread is inside a long handler -- a synchronous ClearInvalidHandles after an unmount, a large batch operation -- a Ping waiting on any of those connections is not even read, so a live client stops being observed and the master moves it to SUSPECTED and then OFFLINE. No lock is involved; the heartbeat simply shares its threads with the data plane. Add --heartbeat_rpc_port (default 0, non-HA mode). When set, the master starts a second coro_rpc_server with one io thread that serves only Ping and ServiceReady, and advertises its port through a new GetHeartbeatRpcPort RPC. MasterClient asks for that port on Connect and sends Ping there once the server answers ServiceReady. Rollout needs no coordination: the main server keeps serving Ping, so a client that predates this change, a master that does not advertise a port, or an advertised server that is unreachable all leave Ping on the main server. This is the design of kvcache-ai#3541 on the P2P branch, applied to main's Ping, with the port discovered from the master instead of configured on every client.

staugust and others added 11 commits August 20, 2026 09:04

add heartbeat port for client&master

1e781c4

move invoke_rpc_via to protected fields

b356cfe

staugust requested review from ShangmingCai, XucSh, YiXR, stmatengss and ykwd as code owners August 20, 2026 02:57

github-actions Bot added run-ci Store Installation Integration Tests labels Aug 20, 2026

staugust changed the title ~~P2P mooncake store~~ feat(store): dedicated heartbeat RPC server to prevent head-of-line blocking Aug 20, 2026

staugust and others added 7 commits August 20, 2026 15:48

style(store): format heartbeat dedicated port test

79ded4d

Apply clang-format line wrapping and fix missing trailing newline in heartbeat_dedicated_port_test.cpp. Co-Authored-By: Claude <noreply@anthropic.com>

wanyue-wy merged commit f67d03f into kvcache-ai:P2P-Mooncake-Store Aug 20, 2026
23 of 25 checks passed

Navigation Menu

Repository navigation

feat(store): dedicated heartbeat RPC server to prevent head-of-line blocking - #3541

feat(store): dedicated heartbeat RPC server to prevent head-of-line blocking#3541
wanyue-wy merged 18 commits into
kvcache-ai:P2P-Mooncake-Storefrom
staugust:P2P-Mooncake-Store

staugust commented Aug 20, 2026 •

edited

Loading

Uh oh!

codecov-commenter commented Aug 20, 2026

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Conversation

staugust commented Aug 20, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Module

Type of Change

How Has This Been Tested?

Manual tests

Checklist

AI Assistance Disclosure

Uh oh!

codecov-commenter commented Aug 20, 2026

Codecov Report

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

staugust commented Aug 20, 2026 •

edited

Loading