[pjdurden]

NoF client: unaligned buffer lengths and addresses (issue #4029)

1. Root cause

TransferSubmitter::submitSpdkNofOperation() in
mooncake-store/src/transfer_task.cpp gated every NoF request behind a hard
alignment check:

if (block_size == INVALID_BLOCK_SIZE ||
    handle.buffer_address_ % block_size != 0 || size % block_size != 0 ||
    reinterpret_cast<std::uintptr_t>(ptr) % block_size != 0) {
    LOG(ERROR) << ... << " is not aligned to block size " << block_size;
    return std::nullopt;
}

The request was then turned 1:1 into a single SpdkNofTask of
size / block_size blocks reading/writing straight out of the caller's
buffer.

NVMe commands are addressed in whole logical blocks, so the length of the
I/O genuinely has to be a block multiple and the disk offset genuinely has
to start on a block boundary. But neither of those is true of the caller's
memory: upstream users (LMCache, SGLang) hand in logical object sizes such as
513 B on a 512 B namespace, and buffers whose virtual address is not a
multiple of the block size. The client had no way to build a padded, block
aligned I/O out of such a request without reading or writing past the ends of
the caller's buffer, so it simply refused the request and the put/get failed
with TRANSFER_FAIL.

2. The fix and why

Split the problem into "what I/Os does this request expand into" (pure
arithmetic, testable) and "issue them" (SPDK specific).

New module nof_io_layout.{h,cpp}. BuildNoFIoLayout() maps
(block_size, disk_offset, buffer_address, size) onto a list of block aligned
NoFIoSegments, each either

• kDirect — DMA runs straight against the caller's buffer, or
• kStaged — DMA runs against a block sized bounce buffer whose first
  payload_size bytes mirror a range of the caller's buffer.

The tiers follow the strategy sketched in the issue, limited to what the
current single-buffer SpdkWrapper::SubmitRequest() can express:

caller's request	layout
block aligned address, block multiple length	one kDirect segment (unchanged, zero copies)
block aligned address, partial trailing block	kDirect prefix + one kStaged trailing block (at most block_size - 1 bytes copied)
unaligned address	one kStaged segment covering the padded length

A misaligned disk offset is still rejected: the object's first byte would
not be addressable by an NVMe command and the client cannot repair a
master-side allocation.

submitSpdkNofOperation() now builds the layout instead of rejecting, and
for the staged cases allocates the bounce buffers with
SpdkWrapper::Alloc(bytes, block_size) (spdk_zmalloc, i.e. DMA capable and
zero filled). On write the valid bytes are copied in and the padding stays
zero, so the caller's buffer is never read past its logical end. On read the
padded blocks land in the bounce buffer and only payload_size bytes are
copied back, so the caller's buffer is never written past its logical end. The
fully aligned path still submits a single task with no copies at all.

All bounce buffers are allocated before the first task is queued: once a task
reaches the worker pool its completion callback is guaranteed to run, so
bailing out halfway would leave the shared state waiting forever.

SpdkNofStagedOperationState (in transfer_task.h) joins the parts. It
owns the bounce buffers (RAII via SpdkNofStagingBuffer), counts the
outstanding parts down, ORs their failures together, performs the read
copy-backs, and only then publishes the single result the existing
TransferFuture waits on. SpdkNofOperationState::set_completed() became
virtual so the worker pool keeps calling it unchanged.

The logical object size, the returned length and any checksum computed over
the caller's buffer are untouched — only the on-wire I/O length is padded.

3. Files changed

file	change
mooncake-store/include/nof_io_layout.h	new — NoFIoSegment / NoFIoLayout / BuildNoFIoLayout()
mooncake-store/src/nof_io_layout.cpp	new — layout construction and its bounds/overflow checks
mooncake-store/src/transfer_task.cpp	submitSpdkNofOperation() builds and issues the layout instead of rejecting unaligned requests
mooncake-store/include/transfer_task.h	SpdkNofStagingBuffer, SpdkNofStagedOperationState; SpdkNofOperationState::set_completed() made virtual
mooncake-store/src/CMakeLists.txt	compile nof_io_layout.cpp (unconditionally, it has no SPDK dependency)
mooncake-store/tests/nof_io_layout_test.cpp	new — unit tests for the layout
mooncake-store/tests/CMakeLists.txt	register nof_io_layout_test

4. Risk and uncertainty

• Depends on master-side padding (PR [Bugfix] Align NoF allocation offsets to the namespace block size #4024, still open). Writing a padded
  trailing block overwrites the bytes between the logical end of the object and
  the end of its last LBA. That is only safe if the allocator reserved those
  bytes for this object, which is exactly what [Bugfix] Align NoF allocation offsets to the namespace block size #4024 does ("use the namespace
  block size as the minimum allocation unit for NoF"). Until [Bugfix] Align NoF allocation offsets to the namespace block size #4024 lands, a
  master that hands out byte-granular NoF offsets will also hand out offsets
  that are not block aligned, and BuildNoFIoLayout() rejects those — so the
  padded write path is effectively unreachable before [Bugfix] Align NoF allocation offsets to the namespace block size #4024. Reads are safe
  either way because the extra bytes are discarded. This ordering should be
  stated in the PR: merge [Bugfix] Align NoF allocation offsets to the namespace block size #4024 first.
• No SGL/iovec path. SpdkWrapper::SubmitRequest() only takes one
  contiguous pointer, so the "direct SGL with a separate padding descriptor"
  and "TCP readv/writev" tiers from the issue are not implemented here;
  they need new spdk_nvme_ns_cmd_readv/writev wrapper entry points and NVMe
  SGL-capability probing. This change keeps the common cases cheap (fully
  aligned: zero copies; aligned address with a partial tail: at most
  block_size - 1 bytes copied) and falls back to full staging only when the
  buffer address itself is unusable. The byte-granular-SGL case the issue calls
  out (don't copy just because the address is not 4-byte aligned) still copies
  here — it is the follow-up that the wrapper work unblocks.
• Bounce buffers are allocated per request via spdk_zmalloc. For the
  tail-only case that is one block; for the full-staging case it is the whole
  padded object, which for a large object with an unaligned address is a real
  allocation plus a full memcpy on the submit path. A pooled/chunked staging
  allocator would bound this and is worth a follow-up if unaligned addresses
  turn out to be common.
• Unverified against real hardware. USE_NOF is OFF by default and needs
  SPDK/DPDK plus an NVMe-oF target, neither of which is available here, so the
  SPDK-facing code path was not compiled or executed against a device. See
  below for what was verified instead.
• set_completed() becoming virtual adds a vtable dispatch per NoF completion;
  negligible next to the I/O itself.

5. How I verified it

• mooncake-store/tests/nof_io_layout_test.cpp — 9 tests over 512 B and
  4096 B namespaces: aligned requests stay direct with zero staging; a 513 B
  object becomes a direct 512 B block plus a 1 block staged tail carrying a
  single payload byte; sub-block objects and unaligned addresses are fully
  staged; the segments always tile the caller's buffer exactly (covered == size, no segment claims more payload than its blocks hold); misaligned disk
  offsets, zero block size, zero length, and block-count overflow are rejected.
  Built and run standalone (the repo's full build needs etcd/RDMA deps that
  are not installed here):

  g++ -std=c++20 -Wall -Wextra -fsanitize=address,undefined \
      -I<fused-gtest> -Imooncake-store/include \
      mooncake-store/tests/nof_io_layout_test.cpp \
      mooncake-store/src/nof_io_layout.cpp gtest-all.cc gtest_main.cc -lpthread
  

  → [ PASSED ] 9 tests., no warnings, clean under ASan/UBSan.

• Because the SPDK path cannot be built here, the new submitSpdkNofOperation()
  body and SpdkNofStagedOperationState were replayed verbatim in a throwaway
  harness (/tmp/nof_submit_sim.cpp) against a stub SpdkWrapper and an
  in-memory namespace: 135 write+read round trips over sizes
  {1, 100, 511, 512, 513, 1024, 1025, 4096, 5000} × buffer skews {0, 1, 3, 7,
  64} × disk offsets {0, 512, 4096}, under ASan/UBSan, asserting that
  (a) readback matches byte for byte, (b) guard bytes on both sides of the
  caller's buffer are untouched by reads and writes, (c) no write touches a
  byte outside [offset, offset + padded_size), and (d) staging allocations
  and frees balance (270/270, no leak). All passed.

• clang-format (repo .clang-format, Google/4-space/80-col) run over every
  touched C++ file. pre-commit itself is not installed in this environment.

[Copilot]

🟡 Changes recommended

Critical counter, LBA, and padded-extent safety issues remain, along with missing staging integration coverage.

Get a fresh assessment by requesting another Copilot review.

Pull request overview

This PR adds block-aligned NoF I/O handling for unaligned buffers and logical lengths using direct or staged transfers.

Changes:

• Adds validated direct/staged I/O layout construction.
• Adds DMA bounce-buffer ownership, padding, copy-back, and completion aggregation.
• Registers the implementation and layout unit tests.

File summaries

File	Summary and final review comments
mooncake-store/tests/nof_io_layout_test.cpp	Adds layout tests.
mooncake-store/tests/CMakeLists.txt	Registers the layout test.
mooncake-store/src/transfer_task.cpp	Integrates direct/staged submission. Critical (2 votes): padded I/O may exceed handle.size_ with stale descriptors. Moderate (1 vote): staging integration lacks automated coverage.
mooncake-store/src/nof_io_layout.cpp	Builds validated layouts. Critical (3 votes): block counts can overflow the signed worker counter. Critical (1 vote): base-LBA plus block-count overflow is unchecked.
mooncake-store/src/CMakeLists.txt	Compiles the layout implementation.
mooncake-store/include/transfer_task.h	Adds staging ownership and completion aggregation. Moderate (1 vote): add tests for copy-backs, padding, completion, and failure aggregation.
mooncake-store/include/nof_io_layout.h	Defines layout structures and API.

Review details

Suppressed comments (2)

mooncake-store/include/transfer_task.h:426

• The committed tests cover only BuildNoFIoLayout; transfer_task_test.cpp explicitly leaves TransferSubmitter integration uncovered. Consequently the new state and SPDK path have no CI coverage for read copy-backs, zero padding, all-parts completion, or failure aggregation, so regressions in the safety-critical staging behavior can pass the 9 layout tests; please check in a fake-worker/SPDK harness or equivalent unit tests.

    void set_completed(ErrorCode error_code) override {
        if (error_code != ErrorCode::OK) {
            failed_.store(true, std::memory_order_relaxed);
        }
        if (pending_parts_.fetch_sub(1, std::memory_order_acq_rel) > 1) {
            return;
        }

        const bool failed = failed_.load(std::memory_order_relaxed);
        if (!failed) {
            for (const auto& copy_back : copy_backs_) {
                std::memcpy(copy_back.dest, copy_back.source, copy_back.size);
            }
        }
        SpdkNofOperationState::set_completed(failed ? ErrorCode::TRANSFER_FAIL
                                                    : ErrorCode::OK);

mooncake-store/src/transfer_task.cpp:1504

• The checked-in tests cover only BuildNoFIoLayout; the existing transfer-task test explicitly does not cover TransferSubmitter integration. Consequently, the new multi-part state, read copy-back, and failure aggregation in this path have no automated regression test, while the SPDK path is disabled by default. Add a hardware-independent test seam or state test covering successful copy-back and a failed part.

    auto state = std::make_shared<SpdkNofStagedOperationState>(
        static_cast<int>(layout.segments.size()), std::move(staging),
        std::move(copy_backs));

• Files reviewed: 7/7 changed files
• Comments generated: 3
• 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 change end to end; the three invariants I care most about all hold, so the direction looks right.

Write-path padding is safe, and I checked why. The staged write fills the block with spdk_zmalloc (zeroed, confirmed at spdk_wrapper.cpp:173 — the "z" variant), then memcpys only payload_size bytes, so the padding is deterministically zero. That alone would not be enough if the trailing bytes belonged to a neighbor, but #4024 already reserves the completing LBA on the master side (a 513B object owns at least 1024B on disk), so writing the full padded block only touches the object's own reservation, and the logical size keeps reads capped at payload_size. Deterministic padding also keeps full-block checksums reproducible.

No half-queued bail. All staging buffers are allocated before the first task is queued, so an allocation failure can't strand the shared state waiting on segments that will never complete — the failure mode I would have flagged otherwise.

Layout builder is conservative in the right places. Misaligned disk_offset is rejected instead of repaired (alignment is an allocation property), and the unaligned-address case goes full-staging rather than pretending a tail copy can fix it — matching the issue's own note that a tail-only copy cannot fix an incompatible starting address.

Two things worth an explicit answer, mostly about scope rather than this diff:

1. Coverage. nof_io_layout_test pins the pure builder well, but nothing exercises the staged integration (submit → bounce → completion → copy-back), including the failure path where a segment completes with an error. If that needs a real SPDK env, saying so in the PR is fine; if any of it can run against the wrapper's mock hooks, that path is the one most likely to break silently.
2. Scope vs the issue. #4029 asks for capability-based selection — RDMA SGL with a zero-padding entry or tail-only SGL, and the TCP readv/writev iovec path — not staging alone. This PR implements staging only, and the whole-payload staging on a block-misaligned address means misaligned buffers always pay a full copy even where a byte-granular SGL could have avoided it. Is the intent to land this as stage one with the SGL/iovec paths following, or to keep this PR strictly to staging? Either is defensible; the PR description should just say which.

CI is green across the board.

[cage-goat]

For RDMA, could we prefer direct SGL when possible to avoid memcpy overhead, keeping staging as a fallback?
The preferred order, subject to memory registration and descriptor limits, would be:

1. No SGL alignment or granularity requirement: use the original payload with separate padding.
2. DWORD alignment and granularity required: transfer the DWORD-aligned prefix directly and stage only the trailing bytes as needed.
3. Use full staging when buffer alignment.

For example, a 513 B object on a 512 B namespace could use the original 513 B buffer plus 511 B of separate padding without payload copies. If DWORD alignment is required and the buffer address is aligned, only the final 1 B needs copying into a padded tail buffer.

[Enigmo-x]

@cage-goat
I agree with your opinion. In our current non-block alignment scenario, this modification is used. The submission interface of SPDK is changed from spdk_nvme_ns_cmd_read to spdk_nvme_ns_cmd_readv.

[pjdurden]

merged current main in (1959141). #3424 moved the NoF path behind NVMeoFInitiator, so the padded layout is ported onto it: segments go out as byte ranges through NofTask, bounce buffers come from the DMA allocator (hugepage_memory_alloc), and the staged state now derives from NofOperationState and keeps the first error detail. main's SGL relaxation is kept by passing the initiator's pointer alignment into the layout, so a DWORD aligned buffer with SGL still goes direct; added a layout test for that. no change to the overall approach here, the direct-SGL idea above is still open.