diff --git a/packages/orchestrator/pkg/sandbox/nbd/dispatch.go b/packages/orchestrator/pkg/sandbox/nbd/dispatch.go index 0644348134..b192eccfee 100644 --- a/packages/orchestrator/pkg/sandbox/nbd/dispatch.go +++ b/packages/orchestrator/pkg/sandbox/nbd/dispatch.go @@ -111,15 +111,21 @@ type Dispatch struct { shuttingDown bool shuttingDownLock sync.Mutex fatal chan error + // asyncWriteZeroes controls whether WRITE_ZEROES/TRIM commands are handled + // in a goroutine (true) instead of inline on the read loop (false). Handling + // them inline lets a blocked reply writer holding writeLock stall the whole + // read loop; see cmdWriteZeroes. + asyncWriteZeroes bool } -func NewDispatch(fp io.ReadWriter, prov Provider) *Dispatch { +func NewDispatch(fp io.ReadWriter, prov Provider, asyncWriteZeroes bool) *Dispatch { d := &Dispatch{ - responseHeader: make([]byte, 16), - fp: fp, - prov: prov, - provName: fmt.Sprintf("%T", prov), - fatal: make(chan error, 1), + responseHeader: make([]byte, 16), + fp: fp, + prov: prov, + provName: fmt.Sprintf("%T", prov), + fatal: make(chan error, 1), + asyncWriteZeroes: asyncWriteZeroes, } binary.BigEndian.PutUint32(d.responseHeader, NBDResponseMagic) @@ -432,19 +438,81 @@ func (d *Dispatch) cmdWrite(ctx context.Context, cmdHandle uint64, cmdFrom uint6 return nil } -// cmdWriteZeroes runs synchronously since WriteZeroesAt is cheap (mmap + bitmap). +// cmdWriteZeroes handles NBD WRITE_ZEROES and TRIM. The backend WriteZeroesAt +// call is cheap (mmap + bitmap), but the following writeResponse takes the +// shared writeLock and writes to the socket. Running this inline on the read +// loop (asyncWriteZeroes == false) means that if a concurrent reply writer is +// blocked inside writeResponse while holding writeLock (e.g. its socket write +// is blocked on a full send buffer), this command blocks the read loop on +// writeLock, the loop stops draining the socket, and the kernel eventually +// times out the connection (EIO). When asyncWriteZeroes is true the work runs +// in a goroutine like cmdRead/cmdWrite, so the read loop is never blocked. func (d *Dispatch) cmdWriteZeroes(ctx context.Context, cmdHandle uint64, cmdFrom uint64, cmdLength int64) error { - var respErr uint32 - if _, err := d.prov.WriteZeroesAt(int64(cmdFrom), cmdLength); err != nil { - respErr = 1 - logger.L().Error(ctx, "nbd backend write-zeroes failed", - zap.Error(err), - zap.String("nbd_provider", d.provName), - zap.Uint64("nbd_handle", cmdHandle), - zap.Uint64("nbd_offset", cmdFrom), - zap.Int64("nbd_length", cmdLength), - ) + performWriteZeroes := func() error { + // Run the backend call in a goroutine and select on ctx, mirroring + // cmdRead/cmdWrite, so a WriteZeroesAt that blocks cannot hang this + // goroutine (and the pendingResponses drain) during shutdown. The + // channel is buffered so the goroutine never leaks on the ctx.Done path. + errchan := make(chan error, 1) + go func() { + _, err := d.prov.WriteZeroesAt(int64(cmdFrom), cmdLength) + errchan <- err + }() + + var zeroErr error + select { + case <-ctx.Done(): + zeroErr = ctx.Err() + case zeroErr = <-errchan: + } + + var respErr uint32 + if zeroErr != nil { + respErr = 1 + logger.L().Error(ctx, "nbd backend write-zeroes failed", + zap.Error(zeroErr), + zap.String("nbd_provider", d.provName), + zap.Uint64("nbd_handle", cmdHandle), + zap.Uint64("nbd_offset", cmdFrom), + zap.Int64("nbd_length", cmdLength), + ) + } + + return d.writeResponse(respErr, cmdHandle, nil) } - return d.writeResponse(respErr, cmdHandle, nil) + if !d.asyncWriteZeroes { + return performWriteZeroes() + } + + d.shuttingDownLock.Lock() + if d.shuttingDown { + d.shuttingDownLock.Unlock() + + return ErrShuttingDown + } + + d.pendingResponses.Add(1) + d.shuttingDownLock.Unlock() + + go func() { + if err := performWriteZeroes(); err != nil { + select { + case d.fatal <- err: + default: + logger.L().Error(ctx, "nbd error cmd write-zeroes", + zap.Error(err), + zap.String("nbd_op", "write-zeroes"), + zap.String("nbd_provider", d.provName), + zap.Uint64("nbd_handle", cmdHandle), + zap.Uint64("nbd_offset", cmdFrom), + zap.Int64("nbd_length", cmdLength), + ) + } + } + + d.pendingResponses.Done() + }() + + return nil } diff --git a/packages/orchestrator/pkg/sandbox/nbd/dispatch_writezeroes_test.go b/packages/orchestrator/pkg/sandbox/nbd/dispatch_writezeroes_test.go new file mode 100644 index 0000000000..11b0eee7ac --- /dev/null +++ b/packages/orchestrator/pkg/sandbox/nbd/dispatch_writezeroes_test.go @@ -0,0 +1,195 @@ +//go:build linux + +package nbd + +import ( + "context" + "encoding/binary" + "io" + "sync" + "testing" + "time" +) + +// ctrlConn is a fake NBD socket driving the real Dispatch.Handle loop. +// +// - Read serves queued request bytes (one queued slice per Read), blocking +// until a request is queued (or reqCh is closed -> io.EOF). +// - Write fires firstWrite once, then blocks until gate is closed, simulating +// a full socket send buffer (the kernel not draining replies). +type ctrlConn struct { + reqCh chan []byte + rbuf []byte + gate chan struct{} + firstWrite chan struct{} + fwOnce sync.Once +} + +func (c *ctrlConn) Read(p []byte) (int, error) { + for len(c.rbuf) == 0 { + b, ok := <-c.reqCh + if !ok { + return 0, io.EOF + } + c.rbuf = b + } + n := copy(p, c.rbuf) + c.rbuf = c.rbuf[n:] + + return n, nil +} + +func (c *ctrlConn) Write(p []byte) (int, error) { + c.fwOnce.Do(func() { close(c.firstWrite) }) + <-c.gate + + return len(p), nil +} + +// stallProv is a minimal Provider. It records the offsets passed to ReadAt so a +// test can tell whether the read loop served a particular request, and signals +// the first WriteZeroesAt call. +type stallProv struct { + mu sync.Mutex + seen map[int64]bool + wz chan struct{} + wzOnce sync.Once +} + +func (m *stallProv) ReadAt(_ context.Context, p []byte, off int64) (int, error) { + m.mu.Lock() + m.seen[off] = true + m.mu.Unlock() + + return len(p), nil +} + +func (m *stallProv) sawRead(off int64) bool { + m.mu.Lock() + defer m.mu.Unlock() + + return m.seen[off] +} + +func (m *stallProv) Size(_ context.Context) (int64, error) { return 1 << 40, nil } +func (m *stallProv) WriteAt(p []byte, _ int64) (int, error) { return len(p), nil } + +func (m *stallProv) WriteZeroesAt(_, length int64) (int, error) { + m.wzOnce.Do(func() { close(m.wz) }) + + return int(length), nil +} + +func nbdRequest(typ uint16, handle, from uint64, length uint32) []byte { + b := make([]byte, 28) + binary.BigEndian.PutUint32(b[0:], NBDRequestMagic) + binary.BigEndian.PutUint16(b[4:], 0) // flags + binary.BigEndian.PutUint16(b[6:], typ) + binary.BigEndian.PutUint64(b[8:], handle) + binary.BigEndian.PutUint64(b[16:], from) + binary.BigEndian.PutUint32(b[24:], length) + + return b +} + +func waitForRead(p *stallProv, off int64, attempts int) bool { + for range attempts { + if p.sawRead(off) { + return true + } + time.Sleep(10 * time.Millisecond) + } + + return p.sawRead(off) +} + +// TestDispatchWriteZeroesReadLoopStall pins the head-of-line stall the +// nbd-async-write-zeroes flag fixes. +// +// A READ reply writer is made to block inside writeResponse while holding +// writeLock (simulating a full socket send buffer). A WRITE_ZEROES is then +// dispatched: +// - inline (asyncWriteZeroes=false): cmdWriteZeroes runs on the read loop and +// blocks acquiring writeLock, so the loop stops serving new requests. +// - async (asyncWriteZeroes=true): cmdWriteZeroes runs in a goroutine, so the +// read loop keeps serving. +// +// In both modes, once the blocked reply drains the loop must make progress +// again (matches the transient, self-clearing incidents this fix targets). +func TestDispatchWriteZeroesReadLoopStall(t *testing.T) { + t.Parallel() + + for _, tc := range []struct { + name string + asyncWriteZeroes bool + wantStall bool + }{ + {name: "inline_stalls_read_loop", asyncWriteZeroes: false, wantStall: true}, + {name: "async_keeps_read_loop_alive", asyncWriteZeroes: true, wantStall: false}, + } { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + conn := &ctrlConn{ + reqCh: make(chan []byte, 8), + gate: make(chan struct{}), + firstWrite: make(chan struct{}), + } + prov := &stallProv{seen: map[int64]bool{}, wz: make(chan struct{})} + d := NewDispatch(conn, prov, tc.asyncWriteZeroes) + + done := make(chan struct{}) + go func() { + _ = d.Handle(t.Context()) + close(done) + }() + // Wind the loop down deterministically at the end: unblock any + // blocked reply Write, then close the request stream so Handle's + // Read returns io.EOF and the goroutine exits. + t.Cleanup(func() { + select { + case <-conn.gate: + default: + close(conn.gate) + } + close(conn.reqCh) + <-done + }) + + const offProbe = int64(0xABCDEF000) + + // 1) READ -> its async reply writer blocks in Write holding writeLock. + conn.reqCh <- nbdRequest(NBDCmdRead, 1, 0, 512) + select { + case <-conn.firstWrite: + case <-time.After(2 * time.Second): + t.Fatal("read reply writer never reached the blocked socket Write") + } + + // 2) WRITE_ZEROES -> inline blocks the loop on writeLock; async does not. + conn.reqCh <- nbdRequest(NBDCmdWriteZeroes, 2, 4096, 4096) + select { + case <-prov.wz: + case <-time.After(2 * time.Second): + t.Fatal("WriteZeroesAt was never called") + } + + // 3) Probe with a fresh READ; a live loop serves it (records offProbe). + conn.reqCh <- nbdRequest(NBDCmdRead, 3, uint64(offProbe), 512) + + served := waitForRead(prov, offProbe, 50) // ~500ms + switch { + case tc.wantStall && served: + t.Fatal("expected the read loop to STALL while a reply write is blocked, but it served a new request") + case !tc.wantStall && !served: + t.Fatal("expected the read loop to keep serving with async WRITE_ZEROES, but it stalled") + } + + // Unblock replies; in both modes the loop must then serve the probe. + close(conn.gate) + if !waitForRead(prov, offProbe, 300) { // ~3s + t.Fatal("read loop did not serve the probe after replies drained") + } + }) + } +} diff --git a/packages/orchestrator/pkg/sandbox/nbd/path_direct.go b/packages/orchestrator/pkg/sandbox/nbd/path_direct.go index 212ed4fc6e..5d53540fe6 100644 --- a/packages/orchestrator/pkg/sandbox/nbd/path_direct.go +++ b/packages/orchestrator/pkg/sandbox/nbd/path_direct.go @@ -126,6 +126,7 @@ func (d *DirectPathMount) Open(ctx context.Context) (retDeviceIndex uint32, err d.dispatchers = make([]*Dispatch, 0) connections := d.featureFlags.IntFlag(ctx, featureflags.NBDConnectionsPerDevice) + asyncWriteZeroes := d.featureFlags.BoolFlag(ctx, featureflags.NBDAsyncWriteZeroesFlag) for i := range connections { // Create the socket pairs @@ -151,7 +152,7 @@ func (d *DirectPathMount) Open(ctx context.Context) (retDeviceIndex uint32, err } server.Close() - dispatch := NewDispatch(serverc, d.Backend) + dispatch := NewDispatch(serverc, d.Backend, asyncWriteZeroes) // Capture deviceIndex for the goroutine closure — it's reassigned on // each retry iteration of the outer for-loop (not a range loop, so // Go 1.22+ loop variable fix doesn't apply). diff --git a/packages/shared/pkg/featureflags/flags.go b/packages/shared/pkg/featureflags/flags.go index 0a20586aaa..403d62d433 100644 --- a/packages/shared/pkg/featureflags/flags.go +++ b/packages/shared/pkg/featureflags/flags.go @@ -266,6 +266,14 @@ var ( // NBDConnectionsPerDevice the number of NBD socket connections per device NBDConnectionsPerDevice = NewIntFlag("nbd-connections-per-device", 1) + // NBDAsyncWriteZeroesFlag, when enabled, handles NBD WRITE_ZEROES/TRIM + // commands in a goroutine instead of inline on the dispatch read loop. + // Inline handling can stall the read loop via head-of-line blocking on the + // shared write lock (when a reply writer is blocked on a full socket send + // buffer), which makes the kernel time out the NBD connection and surfaces + // as guest I/O errors. Disabled by default. + NBDAsyncWriteZeroesFlag = NewBoolFlag("nbd-async-write-zeroes", false) + // MemoryPrefetchMaxFetchWorkers is the maximum number of parallel fetch workers per sandbox for memory prefetching. // Fetching is I/O bound so we can have more parallelism. MemoryPrefetchMaxFetchWorkers = NewIntFlag("memory-prefetch-max-fetch-workers", 16)