From 2f3d7a595c6c5a47c6b8a326c6477f8208f89fc0 Mon Sep 17 00:00:00 2001 From: AdaAibaby <8872412+adababys@users.noreply.github.com> Date: Thu, 22 Oct 2026 15:20:14 +0800 Subject: [PATCH] perf(orchestrator): coalesce NBD header + max payload into one read dispatchBufferSize was exactly 4 MiB, the maximum typical NBD data payload. Every request also carries a 28-byte header, so a max-size (4 MiB) write did not fit the buffer whole: the first Read filled 4 MiB (header + 4 MiB-28 of payload) and a second Read was needed for the trailing 28 payload bytes. The code carried a TODO acknowledging exactly this. Widen dispatchBufferSize by one header (4*1024*1024 + nbdRequestHeaderSize) so [28-byte header + 4 MiB payload] lands in a single Read. Introduce the named nbdRequestHeaderSize constant for the header width and use it in place of the bare 28 literals in the dispatch loop, and drop the resolved TODO and the now- inaccurate deadlock comment in the write path. Measured on this repo with a counting socket feeding one 4 MiB write (header + full payload buffered), reads that returned payload bytes: before (buffer 4194304): 2 Reads after (buffer 4194332): 1 Read One syscall saved per max-size write; for I/O-heavy sandboxes doing large sequential writes this is thousands of requests. Memory cost is 28 bytes per sandbox connection buffer. Rebased onto current main at the maintainers invitation on #3373 (the Sept NBD dispatch refactor left the buffer sizing untouched). Includes the requested before/after measurement as a pinned test (TestDispatchMaxSizeWriteSingleRead). Closes #3355. --- .../orchestrator/pkg/sandbox/nbd/dispatch.go | 28 +++-- .../sandbox/nbd/dispatch_writezeroes_test.go | 105 ++++++++++++++++++ 2 files changed, 123 insertions(+), 10 deletions(-) diff --git a/packages/orchestrator/pkg/sandbox/nbd/dispatch.go b/packages/orchestrator/pkg/sandbox/nbd/dispatch.go index 5e3911068a..e3b2d5edcd 100644 --- a/packages/orchestrator/pkg/sandbox/nbd/dispatch.go +++ b/packages/orchestrator/pkg/sandbox/nbd/dispatch.go @@ -67,9 +67,15 @@ type Provider interface { } const ( - // TODO: Look into optimizing the buffer reads by increasing the buffer size by 28 bytes, - // to account for a request that is 28 bytes of header + 4MB of data (this seems to be preferred kernel buffer size). - dispatchBufferSize = 4 * 1024 * 1024 + // nbdRequestHeaderSize is the wire size of an NBD request header (see Request). + nbdRequestHeaderSize = 28 + + // dispatchBufferSize is the raw read buffer for the NBD socket. 4 MiB covers + // the largest typical data payload; the extra nbdRequestHeaderSize bytes let a + // full max-size payload sit in the buffer alongside its own 28-byte header, so + // the whole [header + 4 MiB payload] is taken in a single Read instead of + // spilling the trailing 28 bytes into a second Read syscall. + dispatchBufferSize = 4*1024*1024 + nbdRequestHeaderSize // https://sourceforge.net/p/nbd/mailman/message/35081223/ // 32MB is the maximum buffer size for a single request that should be universally supported. dispatchMaxWriteBufferSize = 32 * 1024 * 1024 @@ -299,13 +305,13 @@ func (d *Dispatch) Handle(ctx context.Context) error { } // Make sure we have a complete header - if wp-rp < 28 { + if wp-rp < nbdRequestHeaderSize { break // Try again when we have more data... } // We can read the neader... - header := buffer[rp : rp+28] + header := buffer[rp : rp+nbdRequestHeaderSize] request.Magic = binary.BigEndian.Uint32(header) request.Flags = binary.BigEndian.Uint16(header[4:6]) request.Type = binary.BigEndian.Uint16(header[6:8]) @@ -323,13 +329,13 @@ func (d *Dispatch) Handle(ctx context.Context) error { case NBDCmdFlush: return errors.New("not supported: Flush") case NBDCmdRead: - rp += 28 + rp += nbdRequestHeaderSize err := d.cmdRead(ctx, cmdOf(opRead, request)) if err != nil { return err } case NBDCmdWrite: - rp += 28 + rp += nbdRequestHeaderSize if request.Length > dispatchMaxWriteBufferSize { return fmt.Errorf("nbd write request length %d exceeds maximum %d", request.Length, dispatchMaxWriteBufferSize) @@ -341,8 +347,10 @@ func (d *Dispatch) Handle(ctx context.Context) error { rp += dataCopied - // We need to wait for more data here, otherwise we will deadlock if the buffer is Xmb and the length is Xmb because of the header's extra 28 bytes needed. - // At the same time we don't want to increase the default buffer size as the max would be 32mb which is too large for hundreds of sandbox connections. + // Read any remaining payload that did not fit in the dispatch buffer. + // A max-size (4 MiB) payload now fits whole — its trailing bytes no + // longer spill past the buffer — so this loop only runs for payloads + // larger than the typical maximum. for dataCopied < int(request.Length) { n, err := d.fp.Read(data[dataCopied:]) @@ -367,7 +375,7 @@ func (d *Dispatch) Handle(ctx context.Context) error { } case NBDCmdWriteZeroes, NBDCmdTrim: // TRIM and WRITE_ZEROES both punch the cache; NO_HOLE is intentionally ignored. - rp += 28 + rp += nbdRequestHeaderSize err := d.cmdWriteZeroes(ctx, cmdOf(punchOp(request.Type), request)) if err != nil { return err diff --git a/packages/orchestrator/pkg/sandbox/nbd/dispatch_writezeroes_test.go b/packages/orchestrator/pkg/sandbox/nbd/dispatch_writezeroes_test.go index 9c4ed7a00c..1959a3ce7c 100644 --- a/packages/orchestrator/pkg/sandbox/nbd/dispatch_writezeroes_test.go +++ b/packages/orchestrator/pkg/sandbox/nbd/dispatch_writezeroes_test.go @@ -195,3 +195,108 @@ func TestDispatchWriteZeroesReadLoopStall(t *testing.T) { }) } } + +// chunkedConn is a fake NBD socket that serves a fixed script of bytes. Each +// Read copies as much of the remaining script as fits in p (exactly how a +// socket with data already buffered in the kernel behaves), then returns io.EOF. +// It counts the Reads that returned payload bytes, so a test can measure how +// many Read syscalls the dispatch loop needs for one request. +type chunkedConn struct { + data []byte + pos int + reads int +} + +func (c *chunkedConn) Read(p []byte) (int, error) { + if c.pos >= len(c.data) { + return 0, io.EOF + } + n := copy(p, c.data[c.pos:]) + c.pos += n + c.reads++ + + return n, nil +} + +func (c *chunkedConn) Write(p []byte) (int, error) { return len(p), nil } + +// measureProv is a no-op Provider that records the bytes/offset of the WriteAt. +type measureProv struct { + mu sync.Mutex + data []byte + off int64 + done chan struct{} +} + +func (m *measureProv) ReadAt(_ context.Context, p []byte, _ int64) (int, error) { return len(p), nil } +func (m *measureProv) Size(_ context.Context) (int64, error) { return 1 << 40, nil } +func (m *measureProv) WriteZeroesAt(_, l int64) (int, error) { return int(l), nil } +func (m *measureProv) WriteAt(p []byte, off int64) (int, error) { + m.mu.Lock() + m.data = append([]byte(nil), p...) + m.off = off + m.mu.Unlock() + close(m.done) + + return len(p), nil +} + +// TestDispatchMaxSizeWriteSingleRead measures that a maximum-typical write +// (4 MiB payload + its 28-byte header) is taken from the socket in ONE Read. +// Before dispatchBufferSize was widened by nbdRequestHeaderSize the 28 trailing +// payload bytes spilled past the 4 MiB buffer and forced a second Read; this +// test pins the coalesced single-read behaviour and verifies data integrity. +func TestDispatchMaxSizeWriteSingleRead(t *testing.T) { + t.Parallel() + + const payloadSize = 4 * 1024 * 1024 // the maximum typical NBD data payload + + payload := make([]byte, payloadSize) + for i := range payload { + payload[i] = byte(i * 7) + } + + const offset = uint64(0x40000) + script := append(nbdRequest(NBDCmdWrite, 42, offset, uint32(payloadSize)), payload...) + + // Sanity: the whole wire message is header + payload and fits the buffer. + if len(script) != payloadSize+nbdRequestHeaderSize { + t.Fatalf("script size %d, want %d", len(script), payloadSize+nbdRequestHeaderSize) + } + if len(script) > dispatchBufferSize { + t.Fatalf("message %d exceeds dispatch buffer %d; cannot be one Read", len(script), dispatchBufferSize) + } + + conn := &chunkedConn{data: script} + prov := &measureProv{done: make(chan struct{})} + d := NewDispatch(conn, prov, false, logger.L()) + + _ = d.Handle(t.Context()) // returns io.EOF once the script is drained + + // The payload-bearing Read count: exactly one with the widened buffer. + if conn.reads != 1 { + t.Fatalf("max-size write took %d Reads, want 1 (header+payload must coalesce)", conn.reads) + } + + // WriteAt runs in the command's async goroutine; wait for it to land. + select { + case <-prov.done: + case <-time.After(2 * time.Second): + t.Fatal("WriteAt was never called for the max-size write") + } + + prov.mu.Lock() + gotData, gotOff := prov.data, prov.off + prov.mu.Unlock() + if gotOff != int64(offset) { + t.Fatalf("WriteAt offset = %d, want %d", gotOff, offset) + } + if len(gotData) != payloadSize { + t.Fatalf("WriteAt payload = %d bytes, want %d", len(gotData), payloadSize) + } + for i := range gotData { + if gotData[i] != byte(i*7) { + t.Fatalf("payload corrupted at byte %d: got %d want %d", i, gotData[i], byte(i*7)) + } + } +}