AdaAibaby · Oct 22, 2026
@@ -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

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