[Store] Free NoF heartbeat probe buffers on segment unmount

[pjdurden]

Issue #4038 — Master retains NoF heartbeat probe resources after segment unmount

1. Root cause

The NoF heartbeat probe allocates a per-endpoint SPDK DMA buffer and caches it
forever.

SpdkWrapper::ProbeNofSegment() reads one block from the remote namespace to
decide whether a NoF segment is alive. The bounce buffer for that read comes
from SpdkWrapper::GetOrCreateProbeBuffer(), which memoizes it in

std::map<std::string, std::unique_ptr<ProbeBuffer>> probe_buffers_;  // keyed by te_endpoint

The only code that ever frees those buffers is SpdkWrapper::Cleanup(), i.e.
process shutdown.

On the master side, unmounting a NoF segment does tear down the heartbeat
bookkeeping — MasterService::UnmountNoFSegment(),
TryUnmountNoFSegmentByHeartbeat() and the stale-entry sweep in
NofHeartbeatThreadFunc() all nof_heartbeat_states_.erase(segment_id) — but
nothing told SpdkWrapper that the endpoint was gone. So every transport
endpoint the master ever probed kept a hugepage-backed DMA buffer for the
lifetime of the process, even after its segment was unmounted and could never
be probed again.

With mount/unmount churn (clients restarting, NoF nodes rotating, segments
auto-unmounted by the heartbeat itself after a device dies) probe_buffers_
grows monotonically, one entry per distinct endpoint, and the DMA pool is
never returned.

2. The fix and why

Release the probe resources for a transport endpoint at the moment the master
stops tracking the segment behind it.

SpdkWrapper gains ReleaseProbeResources(tr_str), which frees the cached
DMA buffer and drops the map entry. It is idempotent, so it is safe on the
repeated/racing unmount paths.

Freeing a DMA buffer is only safe if no NVMe request still targets it.
ProbeNofSegment() gives up after nof_heartbeat_probe_timeout_ms and returns
completion_timeout while the read is still outstanding, so the buffer can
still be written by the controller at that point. A ProbeBuffer::in_flight
flag (set when the buffer is handed to a probe, cleared on submit failure and
on completion) makes ReleaseProbeResources() skip — and warn about — exactly
those buffers instead of freeing memory a pending DMA targets. They are
reclaimed by Cleanup() as before, and reused/cleared if the endpoint is
remounted and probes succeed again.

MasterService gains a single choke point,
EraseNoFHeartbeatStateLocked(segment_id), which replaces the three raw
nof_heartbeat_states_.erase(...) calls, plus
ReleaseNoFProbeResourcesLocked(endpoints) used by the heartbeat thread's
stale-entry sweep (that sweep is what catches segments removed out-of-band,
e.g. by client offboarding). Releasing is skipped while any other tracked
segment still names the same endpoint, so a shared endpoint keeps its buffer
until its last segment is gone.

The actual release goes through an injectable NoFProbeReleaseFn, mirroring
the existing NoFProbeFn hook, so the behaviour is unit-testable without an
SPDK device (SetNoFProbeReleaseFnForTesting). Like nof_probe_fn_, it is
only wired to SpdkWrapper under USE_NOF.

3. Files changed

File	Change
mooncake-store/include/spdk/spdk_wrapper.h	Declare ReleaseProbeResources / MarkProbeBufferIdle; add ProbeBuffer::in_flight
mooncake-store/src/spdk/spdk_wrapper.cpp	Implement both; set/clear in_flight around the probe read
mooncake-store/include/master_service.h	NoFProbeReleaseFn + nof_probe_release_fn_, SetNoFProbeReleaseFnForTesting, the two new private helpers
mooncake-store/src/master_service.cpp	Default release fn, testing setter, helper implementations, route the four heartbeat-state erase sites through them
mooncake-store/tests/nof_heartbeat_test.cpp	Two regression tests

4. Risk / uncertainty

• Residual retention, deliberately. If the last probe of an endpoint timed
  out (completion_timeout), its buffer is kept rather than freed, because the
  controller may still DMA into it and SPDK keeps hugepages mapped, so a late
  write would silently corrupt whatever is allocated there next. Retention is
  then bounded by the number of distinct endpoints whose final probe hung,
  instead of every endpoint ever probed. Closing that last gap properly needs
  the probe's qpair to be reset/drained, which is a larger change to the
  OpenNofSegment handle cache and is out of scope here.
• Lock ordering. EraseNoFHeartbeatStateLocked runs under
  nof_heartbeat_mutex_ and reaches probe_buffers_mutex_; the release
  callback is copied out from under nof_probe_fn_mutex_ and invoked outside
  it. No path takes nof_heartbeat_mutex_ while holding either of the other
  two, so there is no inversion. ReleaseProbeResources only holds
  probe_buffers_mutex_ briefly and never during the probe's polling loop, so
  an RPC-thread unmount cannot stall behind a slow probe.
• Not verified against real hardware. The probe path needs USE_NOF=ON and
  an NVMe-oF target; neither SPDK nor the project's other build dependencies
  are available in this environment, so the SPDK-side change is reviewed and
  reasoned about, not executed. See below.
• The in_flight flag is plain bool guarded by probe_buffers_mutex_, not
  atomic; every read and write of it is inside that lock.

5. How I verified it

• What I could run. The two new MasterService helpers were extracted into
  a standalone C++20 translation unit with stub types and compiled with
  g++ -std=c++20 -Wall -Wextra, then exercised with assertions covering:
  sole-user endpoint is released; endpoint still used by another tracked
  segment is not; it is released once its last user goes away; erasing an
  unknown segment id is a no-op; and the batch path dedups, skips empty
  endpoints and skips still-tracked ones. All assertions pass.
• What I could not run. No full build or test-suite run. The repo's
  dependencies (glog, ylt, boost, SPDK) are not installed here, and
  nof_heartbeat_test is only built when USE_NOF=ON
  (mooncake-store/tests/CMakeLists.txt:136). scripts/code_format.sh refuses
  to run because it requires clang-format 20 and only 16 is present, so
  formatting was matched to the surrounding code by hand and checked against
  the 80-column limit (no new line exceeds it).
• Regression tests added to mooncake-store/tests/nof_heartbeat_test.cpp,
  both using the injected release hook:
  • UnmountReleasesNoFProbeResources — mount, wait until the heartbeat thread
    tracks the segment, call UnmountNoFSegment, assert the endpoint's probe
    resources were released and the heartbeat state is gone. Fails before the
    fix (nothing is ever released).
  • HeartbeatUnmountReleasesNoFProbeResources — a permanently failing probe
    drives the auto-unmount path; assert the release also happens there.
    Both snapshot the recorded endpoints under their own mutex before asserting,
    so an assertion never holds that mutex while calling back into the service.
• Manual review of every remaining nof_heartbeat_states_ mutation to
  confirm all four now route through the new helpers, and of the in_flight
  state machine for the submit-failure, completion and timeout paths.

[Copilot]

🟡 Changes recommended

Outstanding critical and moderate probe-lifecycle race issues must be addressed before approval.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Pull request overview

This PR releases NoF heartbeat probe DMA buffers when segments are unmounted, with lifecycle tracking and regression tests.

Changes:

• Adds probe-resource release and in-flight tracking.
• Centralizes heartbeat cleanup with shared-endpoint awareness.
• Adds explicit and heartbeat-triggered unmount tests.

File summaries

File	Summary	Findings
mooncake-store/tests/nof_heartbeat_test.cpp	Adds resource-release regression tests.	Nit (1 vote): Missing shared-endpoint coverage.
mooncake-store/src/spdk/spdk_wrapper.cpp	Implements buffer release and request tracking.	Critical (3 votes): Timed-out buffers may be reused while DMA remains outstanding. Moderate (1 vote each): Release lacks deferred cleanup after in-flight completion; probe/unmount handoff can recreate an unreleased buffer.
mooncake-store/src/master_service.cpp	Integrates cleanup into heartbeat-state removal.	Moderate (1 vote): Unmount can race with an already-selected heartbeat probe, leaving a buffer unreleased.
mooncake-store/include/spdk/spdk_wrapper.h	Declares probe lifecycle APIs and state.	None.
mooncake-store/include/master_service.h	Declares release callbacks and cleanup helpers.	None.

Review details

Suppressed comments (5)

mooncake-store/src/master_service.cpp:12343

• This release can race with the heartbeat thread, which selects probe_target and then probes outside nof_heartbeat_mutex_. If the thread has selected a target but has not yet called GetOrCreateProbeBuffer, an unmount can erase the old buffer here; the stale probe then creates a new buffer after its heartbeat state is gone, leaving no later state erase to release it. Coordinate endpoint release with in-progress/stale probes or perform a post-probe release.

    std::string te_endpoint = std::move(it->second.te_endpoint);
    nof_heartbeat_states_.erase(it);
    ReleaseNoFProbeResourcesLocked({te_endpoint});

mooncake-store/src/spdk/spdk_wrapper.cpp:444

• This early return races with the heartbeat thread's probe handoff: NofHeartbeatThreadFunc releases nof_heartbeat_mutex_ before ProbeNoFSegment reaches GetOrCreateProbeBuffer. If an unmount erases the state in that window, this sees no entry, but the already-selected probe then allocates and submits a new buffer; with no heartbeat state left, no later path releases it. Coordinate state removal with probe start or record a deferred release for probes already selected.

    if (it == probe_buffers_.end()) {
        return;

mooncake-store/src/spdk/spdk_wrapper.cpp:452

• When this branch skips an in-flight buffer, no follow-up release is recorded. The probe may complete afterward and MarkProbeBufferIdle only flips the flag; because the heartbeat state has already been erased, no later call reaches ReleaseProbeResources, so even a successfully completed request retains its DMA buffer until process shutdown. Add deferred-release state and perform the release from the completion/termination path once the endpoint is no longer tracked.

    if (it->second && it->second->in_flight) {
        // The last probe never completed, so the controller may still DMA
        // into this buffer. Keep it until Cleanup() tears the device down.
        LOG(WARNING) << "endpoint=" << tr_str
                     << ", action=skip_release_nof_probe_buffer"
                     << ", reason=probe_in_flight";
        return;

mooncake-store/src/spdk/spdk_wrapper.cpp:515

• After the loop exits with completion_timeout, this is the last place that can call MarkProbeBufferIdle for that probe. If the request completes later, ProbeReadComplete only sets done and recycles the context, so the buffer can remain marked in flight; an unmount that follows that late completion (before another probe resets the flag) will still skip it even though DMA has finished. Associate the endpoint or buffer with the request context and clear it from the completion callback while preserving the timeout case.

    if (probe_ctx->done.load(std::memory_order_acquire)) {
        // The request completed, so nothing references the buffer any more.
        MarkProbeBufferIdle(tr_str);
    }

mooncake-store/tests/nof_heartbeat_test.cpp:223

• The added tests cover only a single segment per endpoint. They do not exercise the still_tracked branch in ReleaseNoFProbeResourcesLocked, so a regression that releases shared endpoint resources on the first unmount would still pass. Add a test with two segments sharing one endpoint and assert release happens only after the second segment is removed.

TEST_F(NoFHeartbeatTest, UnmountReleasesNoFProbeResources) {

• Files reviewed: 5/5 changed files
• Comments generated: 1
• Review effort level: Lite

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[he-yufeng]

Read the whole diff carefully; this is the right shape. A few things I specifically checked:

Lock ordering holds: every caller of EraseNoFHeartbeatStateLocked / ReleaseNoFProbeResourcesLocked takes nof_heartbeat_mutex_ first and nof_probe_fn_mutex_ second, and I couldn't find a reverse path anywhere (the probe fns only ever take nof_probe_fn_mutex_ alone), so there's no deadlock cycle.

The in-flight discipline is the load-bearing part and it's right: a buffer whose probe timed out stays marked in_flight forever, so a late DMA landing after timeout can never write into freed memory. The cost is bounded (one buffer per endpoint, only when a probe never completed) and the log line makes it observable. Reuse through GetOrCreateProbeBuffer still works in that state, which matches the pre-existing memoization behavior.

The shared-endpoint refcounting via the any_of scan over remaining heartbeat states is correct for the "several segments, one transport endpoint" case, and the stale-entry sweep in NofHeartbeatThreadFunc (client offboarded without Unmount) was the gap I would have asked about; it's covered.

The two new tests exercise both release paths (explicit unmount and heartbeat-failure unmount) with mocked probes, so they run without NoF hardware.

One question, not blocking: on the timeout path the buffer stays in_flight, so ReleaseProbeResources will skip it at unmount and it lives until Cleanup(). That's the safe choice, just confirming it's intended to be permanent for the life of the device rather than retried on a later release.

COMMENT only since I can't exercise the real SPDK path locally; from the code and the mock-based tests this looks ready.