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)) + } + } +}