From e146c4e0c955ac4de9effdc7c3488656a7ea46cd Mon Sep 17 00:00:00 2001 From: Nikita Kalyazin Date: Fri, 19 Jun 2026 17:29:21 +0100 Subject: [PATCH 1/3] refactor(orch): thread async-write-zeroes option through NBD dispatch Add an asyncWriteZeroes field to Dispatch and a corresponding NewDispatch parameter, and split cmdWriteZeroes into a shared performWriteZeroes closure that can run either inline or in a goroutine. All callers pass false, so behavior is unchanged: WRITE_ZEROES/TRIM still run inline on the read loop. This prepares for gating the async path behind a feature flag. Signed-off-by: Nikita Kalyazin Co-Authored-By: Claude Opus 4.8 --- .../orchestrator/pkg/sandbox/nbd/dispatch.go | 104 +++++++++++++++--- .../pkg/sandbox/nbd/path_direct.go | 2 +- 2 files changed, 87 insertions(+), 19 deletions(-) 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/path_direct.go b/packages/orchestrator/pkg/sandbox/nbd/path_direct.go index 212ed4fc6e..3ca84737d7 100644 --- a/packages/orchestrator/pkg/sandbox/nbd/path_direct.go +++ b/packages/orchestrator/pkg/sandbox/nbd/path_direct.go @@ -151,7 +151,7 @@ func (d *DirectPathMount) Open(ctx context.Context) (retDeviceIndex uint32, err } server.Close() - dispatch := NewDispatch(serverc, d.Backend) + dispatch := NewDispatch(serverc, d.Backend, false) // 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). From 71ca49b590a4193f0381bbd17fb4e75163fda195 Mon Sep 17 00:00:00 2001 From: Nikita Kalyazin Date: Fri, 19 Jun 2026 17:30:46 +0100 Subject: [PATCH 2/3] feat(orch): handle NBD WRITE_ZEROES off the read loop behind a flag Add the nbd-async-write-zeroes feature flag (default false), evaluate it per device mount, and pass it to NewDispatch. When enabled, WRITE_ZEROES and TRIM are handled in a goroutine instead of inline on the dispatch read loop. Inline handling shares writeLock with the asynchronous read/write reply writers. If a reply write blocks on a full socket send buffer while holding writeLock, a WRITE_ZEROES processed on the read loop blocks acquiring the lock; the loop then stops draining the socket and the kernel times out the NBD connection, returning EIO to the guest and crashing Firecracker on the next flush. Moving the work off the loop keeps the socket drained. Disabled by default for a gradual rollout. Signed-off-by: Nikita Kalyazin Co-Authored-By: Claude Opus 4.8 --- packages/orchestrator/pkg/sandbox/nbd/path_direct.go | 3 ++- packages/shared/pkg/featureflags/flags.go | 8 ++++++++ 2 files changed, 10 insertions(+), 1 deletion(-) diff --git a/packages/orchestrator/pkg/sandbox/nbd/path_direct.go b/packages/orchestrator/pkg/sandbox/nbd/path_direct.go index 3ca84737d7..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, false) + 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) From fce95075f0a3e924f74d9375817ad60ca8210354 Mon Sep 17 00:00:00 2001 From: Nikita Kalyazin Date: Fri, 19 Jun 2026 17:34:54 +0100 Subject: [PATCH 3/3] test(orch): cover NBD WRITE_ZEROES read-loop head-of-line stall Drive the real Dispatch loop over an in-memory socket whose reply Write blocks, simulating a full kernel send buffer. With a READ reply blocked while holding writeLock, dispatch a WRITE_ZEROES and probe with a fresh READ: - inline (flag off): the loop stalls and stops serving the probe - async (flag on): the loop keeps serving the probe Both modes recover once the blocked reply drains, matching the transient, self-clearing incidents this fix targets. Signed-off-by: Nikita Kalyazin Co-Authored-By: Claude Opus 4.8 --- .../sandbox/nbd/dispatch_writezeroes_test.go | 195 ++++++++++++++++++ 1 file changed, 195 insertions(+) create mode 100644 packages/orchestrator/pkg/sandbox/nbd/dispatch_writezeroes_test.go 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") + } + }) + } +}