Repository: e2b-dev/infra
Pull request #517: Improve NBD handling
Author: login: ValentaTomas
Labels: name: bug
description: Something isn't working
name: improvement
description: Improvement for current functionality
Created: 2025-03-31T05:17:08Z | Merged: 2025-04-28T14:49:59Z

We are experiencing a kernel NULL dereference that might be caused by incorrect NBD handling.


Diff:
diff --git a/packages/orchestrator/internal/sandbox/nbd/dispatch.go b/packages/orchestrator/internal/sandbox/nbd/dispatch.go
index 409430450b..2bc13a4562 100644
--- a/packages/orchestrator/internal/sandbox/nbd/dispatch.go
+++ b/packages/orchestrator/internal/sandbox/nbd/dispatch.go
@@ -3,6 +3,7 @@ package nbd
 import (
 	"context"
 	"encoding/binary"
+	"errors"
 	"fmt"
 	"io"
 	"sync"
@@ -10,6 +11,8 @@ import (
 	"go.uber.org/zap"
 )
 
+var ErrShuttingDown = errors.New("shutting down. Cannot serve any new requests")
+
 type Provider interface {
 	io.ReaderAt
 	io.WriterAt
@@ -55,7 +58,9 @@ type Dispatch struct {
 	writeLock        sync.Mutex
 	prov             Provider
 	pendingResponses sync.WaitGroup
-	pendingMu        sync.Mutex
+	shuttingDown     bool
+	shuttingDownLock sync.Mutex
+	fatal            chan error
 }
 
 func NewDispatch(ctx context.Context, fp io.ReadWriteCloser, prov Provider) *Dispatch {
@@ -64,15 +69,17 @@ func NewDispatch(ctx context.Context, fp io.ReadWriteCloser, prov Provider) *Dis
 		fp:             fp,
 		prov:           prov,
 		ctx:            ctx,
+		fatal:          make(chan error, 1),
 	}
 
 	binary.BigEndian.PutUint32(d.responseHeader, NBDResponseMagic)
 	return d
 }
 
-func (d *Dispatch) Wait() {
-	d.pendingMu.Lock()
-	defer d.pendingMu.Unlock()
+func (d *Dispatch) Drain() {
+	d.shuttingDownLock.Lock()
+	d.shuttingDown = true
+	defer d.shuttingDownLock.Unlock()
 
 	// Wait for any pending responses
 	d.pendingResponses.Wait()
@@ -124,9 +131,10 @@ func (d *Dispatch) Handle() error {
 		rp := 0
 	process:
 		for {
-
-			// If the context has been cancelled, quit
+			// Check if there is a fatal error from an async read/write to return
 			select {
+			case err := <-d.fatal:
+				return err
 			case <-d.ctx.Done():
 				return d.ctx.Err()
 			default:
@@ -173,14 +181,13 @@ func (d *Dispatch) Handle() error {
 					}
 				case NBDCmdTrim:
 					rp += 28
-					err = d.cmdTrim(request.Handle, request.From, request.Length)
+					err := d.cmdTrim(request.Handle, request.From, request.Length)
 					if err != nil {
 						return err
 					}
 				default:
 					return fmt.Errorf("nbd not implemented %d", request.Type)
 				}
-
 			} else {
 				break // Try again when we have more data...
 			}
@@ -194,9 +201,14 @@ func (d *Dispatch) Handle() error {
 }
 
 func (d *Dispatch) cmdRead(cmdHandle uint64, cmdFrom uint64, cmdLength uint32) error {
-	d.pendingMu.Lock()
-	d.pendingResponses.Add(1)
-	d.pendingMu.Unlock()
+	d.shuttingDownLock.Lock()
+	if !d.shuttingDown {
+		d.pendingResponses.Add(1)
+	} else {
+		d.shuttingDownLock.Unlock()
+		return ErrShuttingDown
+	}
+	d.shuttingDownLock.Unlock()
 
 	performRead := func(handle uint64, from uint64, length uint32) error {
 		errchan := make(chan error)
@@ -215,20 +227,21 @@ func (d *Dispatch) cmdRead(cmdHandle uint64, cmdFrom uint64, cmdLength uint32) e
 		case e = <-errchan:
 		}
 
-		errorValue := uint32(0)
 		if e != nil {
-			errorValue = 1
-			data = make([]byte, 0) // If there was an error, don't send data
+			return d.writeResponse(1, handle, []byte{})
 		}
-		return d.writeResponse(errorValue, handle, data)
+		return d.writeResponse(0, handle, data)
 	}
 
 	go func() {
 		err := performRead(cmdHandle, cmdFrom, cmdLength)
 		if err != nil {
-			zap.L().Error("nbd error cmd read", zap.Error(err))
+			select {
+			case d.fatal <- err:
+			default:
+				zap.L().Error("nbd error cmd read", zap.Error(err))
+			}
 		}
-
 		d.pendingResponses.Done()
 	}()
 
@@ -236,14 +249,19 @@ func (d *Dispatch) cmdRead(cmdHandle uint64, cmdFrom uint64, cmdLength uint32) e
 }
 
 func (d *Dispatch) cmdWrite(cmdHandle uint64, cmdFrom uint64, cmdData []byte) error {
-	d.pendingMu.Lock()
-	d.pendingResponses.Add(1)
-	d.pendingMu.Unlock()
+	d.shuttingDownLock.Lock()
+	if !d.shuttingDown {
+		d.pendingResponses.Add(1)
+	} else {
+		d.shuttingDownLock.Unlock()
+		return ErrShuttingDown
+	}
+	d.shuttingDownLock.Unlock()
 
-	go func() {
+	performWrite := func(handle uint64, from uint64, data []byte) error {
 		errchan := make(chan error)
 		go func() {
-			_, e := d.prov.WriteAt(cmdData, int64(cmdFrom))
+			_, e := d.prov.WriteAt(data, int64(from))
 			errchan <- e
 		}()
 
@@ -259,14 +277,20 @@ func (d *Dispatch) cmdWrite(cmdHandle uint64, cmdFrom uint64, cmdData []byte) er
 		if e != nil {
 			errorValue = 1
 		}
-		err := d.writeResponse(errorValue, cmdHandle, []byte{})
+		return d.writeResponse(errorValue, handle, []byte{})
+	}
+
+	go func() {
+		err := performWrite(cmdHandle, cmdFrom, cmdData)
 		if err != nil {
-			zap.L().Error("nbd error cmd write", zap.Error(err))
+			select {
+			case d.fatal <- err:
+			default:
+				zap.L().Error("nbd error cmd write", zap.Error(err))
+			}
 		}
-
 		d.pendingResponses.Done()
 	}()
-
 	return nil
 }
 
diff --git a/packages/orchestrator/internal/sandbox/nbd/path_direct_linux.go b/packages/orchestrator/internal/sandbox/nbd/path_direct_linux.go
index cb6a80a9e6..851b92df46 100644
--- a/packages/orchestrator/internal/sandbox/nbd/path_direct_linux.go
+++ b/packages/orchestrator/internal/sandbox/nbd/path_direct_linux.go
@@ -5,9 +5,14 @@ package nbd
 
 import (
 	"context"
+	"errors"
+	"fmt"
+	"io"
+	"math"
 	"net"
 	"os"
 	"strings"
+	"sync"
 	"syscall"
 	"time"
 
@@ -19,112 +24,159 @@ import (
 	"github.com/e2b-dev/infra/packages/shared/pkg/telemetry"
 )
 
+const (
+	connections    = 4
+	connectTimeout = 30 * time.Second
+
+	// disconnectTimeout should not be necessary if the disconnect is reliable
+	disconnectTimeout = 30 * time.Second
+)
+
 type DirectPathMount struct {
-	tracer      trace.Tracer
+	tracer   trace.Tracer
+	ctx      context.Context
+	cancelfn context.CancelFunc
+
+	devicePool *DevicePool
+
 	Backend     block.Device
-	ctx         context.Context
-	dispatcher  *Dispatch
-	conn        net.Conn
 	deviceIndex uint32
 	blockSize   uint64
-	cancelfn    context.CancelFunc
-	devicePool  *DevicePool
+
+	dispatchers []*Dispatch
+	socksClient []*os.File
+	socksServer []io.Closer
+
+	handlersWg sync.WaitGroup
 }
 
 func NewDirectPathMount(tracer trace.Tracer, b block.Device, devicePool *DevicePool) *DirectPathMount {
 	ctx, cancelfn := context.WithCancel(context.Background())
 
 	return &DirectPathMount{
-		tracer:     tracer,
-		Backend:    b,
-		ctx:        ctx,
-		cancelfn:   cancelfn,
-		blockSize:  4096,
-		devicePool: devicePool,
+		tracer:      tracer,
+		Backend:     b,
+		ctx:         ctx,
+		cancelfn:    cancelfn,
+		blockSize:   4096,
+		devicePool:  devicePool,
+		socksClient: make([]*os.File, 0),
+		socksServer: make([]io.Closer, 0),
+		deviceIndex: math.MaxUint32,
 	}
 }
 
-func (d *DirectPathMount) Open(ctx context.Context) (uint32, error) {
+func (d *DirectPathMount) Open(ctx context.Context) (retDeviceIndex uint32, err error) {
+	defer func() {
+		// Set the device index to the one returned, correctly capture error values
+		d.deviceIndex = retDeviceIndex
+		zap.L().Debug("opening direct path mount", zap.Uint32("device_index", d.deviceIndex), zap.Error(err))
+	}()
+
 	size, err := d.Backend.Size()
 	if err != nil {
-		return 0, err
+		return math.MaxUint32, err
 	}
 
-	for {
-		d.deviceIndex, err = d.devicePool.GetDevice(ctx)
-		if err != nil {
-			return 0, err
-		}
+	deviceIndex := uint32(math.MaxUint32)
 
-		// Create the socket pairs
-		sockPair, err := syscall.Socketpair(syscall.AF_UNIX, syscall.SOCK_STREAM, 0)
+	for {
+		deviceIndex, err = d.devicePool.GetDevice(ctx)
 		if err != nil {
-			return 0, err
+			return math.MaxUint32, err
 		}
 
-		client := os.NewFile(uintptr(sockPair[0]), "client")
-		server := os.NewFile(uintptr(sockPair[1]), "server")
-		d.conn, err = net.FileConn(server)
+		d.socksClient = make([]*os.File, 0)
+		d.socksServer = make([]io.Closer, 0)
+		d.dispatchers = make([]*Dispatch, 0)
 
-		if err != nil {
-			return 0, err
-		}
-		server.Close()
+		for i := 0; i < connections; i++ {
+			// Create the socket pairs
+			sockPair, err := syscall.Socketpair(syscall.AF_UNIX, syscall.SOCK_STREAM, 0)
+			if err != nil {
+				return math.MaxUint32, err
+			}
 
-		d.dispatcher = NewDispatch(d.ctx, d.conn, d.Backend)
-		// Start reading commands on the socket and dispatching them to our provider
-		go func() {
-			handleErr := d.dispatcher.Handle()
-			if handleErr != nil {
-				zap.L().Error("error handling NBD commands", zap.Error(handleErr))
+			client := os.NewFile(uintptr(sockPair[0]), "client")
+			server := os.NewFile(uintptr(sockPair[1]), "server")
+			serverc, err := net.FileConn(server)
+			if err != nil {
+				return math.MaxUint32, err
 			}
-		}()
+			server.Close()
+
+			dispatch := NewDispatch(d.ctx, serverc, d.Backend)
+			// Start reading commands on the socket and dispatching them to our provider
+			d.handlersWg.Add(1)
+			go func() {
+				defer d.handlersWg.Done()
+
+				handleErr := dispatch.Handle()
+				// The error is expected to happen if the nbd (socket connection) is closed
+				zap.L().Info("closing handler for NBD commands",
+					zap.Error(handleErr),
+					zap.Uint32("device_index", deviceIndex),
+					zap.Int("socket_index", i),
+				)
+			}()
+
+			d.socksServer = append(d.socksServer, serverc)
+			d.socksClient = append(d.socksClient, client)
+			d.dispatchers = append(d.dispatchers, dispatch)
+		}
 
 		var opts []nbdnl.ConnectOption
 		opts = append(opts, nbdnl.WithBlockSize(d.blockSize))
-		opts = append(opts, nbdnl.WithTimeout(5*time.Second))
-		opts = append(opts, nbdnl.WithDeadconnTimeout(5*time.Second))
+		opts = append(opts, nbdnl.WithTimeout(connectTimeout))
+		opts = append(opts, nbdnl.WithDeadconnTimeout(connectTimeout))
 
 		serverFlags := nbdnl.FlagHasFlags | nbdnl.FlagCanMulticonn
 
-		idx, err := nbdnl.Connect(d.deviceIndex, []*os.File{client}, uint64(size), 0, serverFlags, opts...)
+		idx, err := nbdnl.Connect(deviceIndex, d.socksClient, uint64(size), 0, serverFlags, opts...)
 		if err == nil {
-			d.deviceIndex = idx
+			// The idx should be the same as deviceIndex, because we are connecting to it,
+			// but we will use the one returned by nbdnl
+			deviceIndex = idx
+
 			break
 		}
 
+		zap.L().Error("error opening NBD, retrying", zap.Error(err), zap.Uint32("device_index", deviceIndex))
+
 		// Sometimes (rare), there seems to be a BADF error here. Lets just retry for now...
 		// Close things down and try again...
-		_ = client.Close()
-
-		connErr := d.conn.Close()
-		if connErr != nil {
-			zap.L().Error("error closing conn", zap.Error(connErr))
+		for _, sock := range d.socksClient {
+			sock.Close()
 		}
-
-		releaseErr := d.devicePool.ReleaseDevice(d.deviceIndex)
-		if releaseErr != nil {
-			zap.L().Error("error releasing device", zap.Error(releaseErr))
+		for _, sock := range d.socksServer {
+			sock.Close()
+		}
+		// Release the device back to the pool
+		err = d.devicePool.ReleaseDevice(deviceIndex)
+		if err != nil {
+			zap.L().Error("error opening NBD, error releasing device", zap.Error(err), zap.Uint32("device_index", deviceIndex))
 		}
-
-		d.deviceIndex = 0
 
 		if strings.Contains(err.Error(), "invalid argument") {
-			return 0, err
+			return math.MaxUint32, err
 		}
 
-		time.Sleep(25 * time.Millisecond)
+		select {
+		case <-d.ctx.Done():
+			return math.MaxUint32, errors.Join(err, d.ctx.Err())
+		case <-time.After(25 * time.Millisecond):
+		}
 	}
 
 	// Wait until it's connected...
 	for {
 		select {
 		case <-d.ctx.Done():
-			return 0, d.ctx.Err()
+			return math.MaxUint32, d.ctx.Err()
 		default:
 		}
 
-		s, err := nbdnl.Status(d.deviceIndex)
+		s, err := nbdnl.Status(deviceIndex)
 		if err == nil && s.Connected {
 			break
 		}
@@ -132,40 +184,80 @@ func (d *DirectPathMount) Open(ctx context.Context) (uint32, error) {
 		time.Sleep(100 * time.Nanosecond)
 	}
 
-	return d.deviceIndex, nil
+	return deviceIndex, nil
 }
 
 func (d *DirectPathMount) Close(ctx context.Context) error {
 	childCtx, childSpan := d.tracer.Start(ctx, "direct-path-mount-close")
 	defer childSpan.End()
 
+	var errs []error
+
+	idx := d.deviceIndex
+
 	// First cancel the context, which will stop waiting on pending readAt/writeAt...
 	telemetry.ReportEvent(childCtx, "canceling context")
 	d.cancelfn()
 
+	// Close all server socket pairs...
+	telemetry.ReportEvent(childCtx, "closing socket pairs server")
+	for _, v := range d.socksServer {
+		err := v.Close()
+		if err != nil {
+			errs = append(errs, fmt.Errorf("error closing server pair: %w", err))
+		}
+	}
+
+	// Now wait until the handlers return
+	telemetry.ReportEvent(childCtx, "await handlers return")
+	d.handlersWg.Wait()
+
 	// Now wait for any pending responses to be sent
 	telemetry.ReportEvent(childCtx, "waiting for pending responses")
-	if d.dispatcher != nil {
-		d.dispatcher.Wait()
+	for _, d := range d.dispatchers {
+		d.Drain()
 	}
 
-	// Now ask to disconnect
-	telemetry.ReportEvent(childCtx, "disconnecting NBD")
-	err := nbdnl.Disconnect(d.deviceIndex)
-	if err != nil {
-		return err
+	// Disconnect NBD
+	if idx != math.MaxUint32 {
+		err := disconnectNBDWithTimeout(childCtx, idx, disconnectTimeout)
+		if err != nil {
+			errs = append(errs, fmt.Errorf("error disconnecting NBD: %w", err))
+		}
 	}
 
-	// Close all the socket pairs...
-	telemetry.ReportEvent(childCtx, "closing socket pairs")
-	err = d.conn.Close()
+	// Close all client socket pairs...
+	telemetry.ReportEvent(childCtx, "closing socket pairs client")
+	for _, v := range d.socksClient {
+		err := v.Close()
+		if err != nil {
+			errs = append(errs, fmt.Errorf("error closing socket pair client: %w", err))
+		}
+	}
+
+	// Release the device back to the pool, retry if it is in use
+	if idx != math.MaxUint32 {
+		telemetry.ReportEvent(childCtx, "releasing device to the pool")
+		err := d.devicePool.ReleaseDeviceWithRetry(idx)
+		if err != nil {
+			errs = append(errs, fmt.Errorf("error releasing overlay device: %w", err))
+		}
+	}
+
+	return errors.Join(errs...)
+}
+
+func disconnectNBDWithTimeout(ctx context.Context, deviceIndex uint32, timeout time.Duration) error {
+	// Now ask to disconnect
+	telemetry.ReportEvent(ctx, "disconnecting NBD")
+	err := nbdnl.Disconnect(deviceIndex)
 	if err != nil {
 		return err
 	}
 
 	// Wait until it's completely disconnected...
-	telemetry.ReportEvent(childCtx, "waiting for complete disconnection")
-	ctxTimeout, cancel := context.WithTimeout(childCtx, 4*time.Second)
+	telemetry.ReportEvent(ctx, "waiting for complete disconnection")
+	ctxTimeout, cancel := context.WithTimeout(ctx, timeout)
 	defer cancel()
 	for {
 		select {
@@ -174,11 +266,10 @@ func (d *DirectPathMount) Close(ctx context.Context) error {
 		default:
 		}
 
-		s, err := nbdnl.Status(d.deviceIndex)
+		s, err := nbdnl.Status(deviceIndex)
 		if err == nil && !s.Connected {
 			break
 		}
-
 		time.Sleep(100 * time.Nanosecond)
 	}
 
diff --git a/packages/orchestrator/internal/sandbox/nbd/pool.go b/packages/orchestrator/internal/sandbox/nbd/pool.go
index 31c9f1e4e2..aa65b639e4 100644
--- a/packages/orchestrator/internal/sandbox/nbd/pool.go
+++ b/packages/orchestrator/internal/sandbox/nbd/pool.go
@@ -5,10 +5,10 @@ import (
 	"errors"
 	"fmt"
 	"os"
-	"regexp"
 	"strconv"
 	"strings"
 	"sync"
+	"time"
 
 	"github.com/bits-and-blooms/bitset"
 	"go.opentelemetry.io/otel/metric"
@@ -257,22 +257,32 @@ func (d *DevicePool) ReleaseDevice(idx DeviceSlot) error {
 	return nil
 }
 
-func GetDevicePath(slot DeviceSlot) DevicePath {
-	return fmt.Sprintf("/dev/nbd%d", slot)
-}
+// ReleaseDeviceWithRetry calls ReleaseDevice and retries if the device is in use.
+func (d *DevicePool) ReleaseDeviceWithRetry(idx DeviceSlot) error {
+	attempt := 0
+	for {
+		attempt++
+		err := d.ReleaseDevice(idx)
+		if errors.Is(err, ErrDeviceInUse{}) {
+			if attempt%100 == 0 {
+				zap.L().Error("error releasing device", zap.Int("attempt", attempt), zap.Error(err))
+			}
 
-var reSlot = regexp.MustCompile(`^/dev/nbd(\d+)$`)
+			time.Sleep(500 * time.Millisecond)
 
-func GetDeviceSlot(path DevicePath) (DeviceSlot, error) {
-	matches := reSlot.FindStringSubmatch(path)
-	if len(matches) != 2 {
-		return 0, fmt.Errorf("invalid nbd path: %s", path)
-	}
+			continue
+		}
 
-	slot, err := strconv.ParseUint(matches[1], 10, 0)
-	if err != nil {
-		return 0, fmt.Errorf("failed to parse slot from path: %w", err)
+		if err != nil {
+			return fmt.Errorf("error releasing device: %w", err)
+		}
+
+		break
 	}
 
-	return DeviceSlot(slot), nil
+	return nil
+}
+
+func GetDevicePath(slot DeviceSlot) DevicePath {
+	return fmt.Sprintf("/dev/nbd%d", slot)
 }
diff --git a/packages/orchestrator/internal/sandbox/rootfs/cow.go b/packages/orchestrator/internal/sandbox/rootfs/cow.go
index ff41fec681..f0b079d403 100644
--- a/packages/orchestrator/internal/sandbox/rootfs/cow.go
+++ b/packages/orchestrator/internal/sandbox/rootfs/cow.go
@@ -5,7 +5,6 @@ import (
 	"errors"
 	"fmt"
 	"io"
-	"time"
 
 	"github.com/bits-and-blooms/bitset"
 	"go.opentelemetry.io/otel/trace"
@@ -79,7 +78,12 @@ func (o *CowDevice) Export(parentCtx context.Context, out io.Writer, stopSandbox
 	}
 
 	// the error is already logged in go routine in SandboxCreate handler
-	go stopSandbox(childCtx)
+	go func() {
+		err := stopSandbox(childCtx)
+		if err != nil {
+			zap.L().Error("error stopping sandbox on cow export", zap.Error(err))
+		}
+	}()
 
 	select {
 	case <-o.finishedOperations:
@@ -122,44 +126,9 @@ func (o *CowDevice) Close(ctx context.Context) error {
 		errs = append(errs, fmt.Errorf("error closing overlay cache: %w", err))
 	}
 
-	devicePath, err := o.ready.Wait()
-	if err != nil {
-		errs = append(errs, fmt.Errorf("error getting overlay path: %w", err))
-
-		return errors.Join(errs...)
-	}
-
-	slot, err := nbd.GetDeviceSlot(devicePath)
-	if err != nil {
-		errs = append(errs, fmt.Errorf("error getting overlay slot: %w", err))
-
-		return errors.Join(errs...)
-	}
-
-	attempts := 0
-	for {
-		attempts++
-		err := o.devicePool.ReleaseDevice(slot)
-		if errors.Is(err, nbd.ErrDeviceInUse{}) {
-			if attempts%100 == 0 {
-				zap.L().Info("error releasing overlay device", zap.Int("attempts", attempts), zap.Error(err))
-			}
-
-			time.Sleep(500 * time.Millisecond)
-
-			continue
-		}
-
-		if err != nil {
-			return fmt.Errorf("error releasing overlay device: %w", err)
-		}
-
-		break
-	}
-
 	zap.L().Info("overlay device released")
 
-	return nil
+	return errors.Join(errs...)
 }
 
 func (o *CowDevice) Path() (string, error) {
diff --git a/packages/orchestrator/internal/sandbox/sandbox_linux.go b/packages/orchestrator/internal/sandbox/sandbox_linux.go
index 0cff94284f..808bf7e4f2 100644
--- a/packages/orchestrator/internal/sandbox/sandbox_linux.go
+++ b/packages/orchestrator/internal/sandbox/sandbox_linux.go
@@ -173,7 +173,9 @@ func NewSandbox(
 		childCtx, span := tracer.Start(ctx, "rootfs-overlay-close")
 		defer span.End()
 
-		rootfsOverlay.Close(childCtx)
+		if rootfsOverlayErr := rootfsOverlay.Close(childCtx); rootfsOverlayErr != nil {
+			return fmt.Errorf("failed to close overlay file: %w", rootfsOverlayErr)
+		}
 
 		return nil
 	})


Reviews:
user: login: ValentaTomas
body: Important changes:

### `dispatch.go`
- Reintroduce fatal channel for closing `Handle` on error
- Refactor pending lock/wg (refuse new cmds after close starts) 

### `path_direct_linux.go`
- Increase timeout for operations (4s -> 30s)
- Allow multiple socket connections
- Refactor device release on failure
- When retrying `Open` retry just connection establishment, without getting new device (not 100% sure this won't be a regression)
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-31T05:27:50Z
user: login: ValentaTomas
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-04-04T19:00:16Z
user: login: dobrac
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-04T19:21:28Z
user: login: ValentaTomas
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-04-04T19:37:27Z
user: login: jakubno
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-04-08T09:28:10Z
user: login: dobrac
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-08T09:33:31Z
user: login: dobrac
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-08T09:36:58Z
user: login: dobrac
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-08T09:38:25Z
user: login: dobrac
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-08T09:40:16Z
user: login: jakubno
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-04-08T09:54:47Z
user: login: dobrac
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-08T11:35:50Z
user: login: ValentaTomas
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-04-08T20:24:38Z
user: login: ValentaTomas
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-04-08T20:26:54Z
user: login: ValentaTomas
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-04-08T20:29:42Z
user: login: dobrac
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-09T07:17:53Z
user: login: dobrac
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-09T07:19:16Z
user: login: dobrac
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-09T13:17:52Z
user: login: dobrac
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-09T14:31:50Z
user: login: jakubno
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-04-09T14:58:19Z
user: login: jakubno
state: APPROVED
author_association: MEMBER
submitted_at: 2025-04-25T15:38:05Z
user: login: dobrac
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-28T10:20:54Z
user: login: jakubno
state: APPROVED
author_association: MEMBER
submitted_at: 2025-04-28T12:20:32Z

Review comments:
diff_hunk: @@ -156,16 +197,9 @@ func (d *DirectPathMount) Close(ctx context.Context) error {
 		return err
 	}
 
-	// Close all the socket pairs...
-	telemetry.ReportEvent(childCtx, "closing socket pairs")
-	err = d.conn.Close()
-	if err != nil {
-		return err
-	}
-
 	// Wait until it's completely disconnected...
 	telemetry.ReportEvent(childCtx, "waiting for complete disconnection")
-	ctxTimeout, cancel := context.WithTimeout(childCtx, 4*time.Second)
+	ctxTimeout, cancel := context.WithTimeout(childCtx, timeout)
path: packages/orchestrator/internal/sandbox/nbd/path_direct_linux.go
user: login: ValentaTomas
body: Let's not reuse the timeout that is also used for the NBD connection timeout.

Also I think ideally we would not use timeout here and the file would be 100% closed. If we can implement it reliably that way I think we should go for it.
created_at: 2025-04-04T19:00:16Z
author_association: MEMBER
diff_hunk: @@ -156,16 +197,9 @@ func (d *DirectPathMount) Close(ctx context.Context) error {
 		return err
 	}
 
-	// Close all the socket pairs...
-	telemetry.ReportEvent(childCtx, "closing socket pairs")
-	err = d.conn.Close()
-	if err != nil {
-		return err
-	}
-
 	// Wait until it's completely disconnected...
 	telemetry.ReportEvent(childCtx, "waiting for complete disconnection")
-	ctxTimeout, cancel := context.WithTimeout(childCtx, 4*time.Second)
+	ctxTimeout, cancel := context.WithTimeout(childCtx, timeout)
path: packages/orchestrator/internal/sandbox/nbd/path_direct_linux.go
user: login: dobrac
body: Changed to separate timeout. I agree it shouldn't be necessary, but I would leave it here for now, so we won't break it. Also with the increased time, we should see it in stats that the requests are timing out more clearly, as well as giving the disconnect more time.
created_at: 2025-04-04T19:21:28Z
author_association: CONTRIBUTOR
diff_hunk: @@ -156,16 +197,9 @@ func (d *DirectPathMount) Close(ctx context.Context) error {
 		return err
 	}
 
-	// Close all the socket pairs...
-	telemetry.ReportEvent(childCtx, "closing socket pairs")
-	err = d.conn.Close()
-	if err != nil {
-		return err
-	}
-
 	// Wait until it's completely disconnected...
 	telemetry.ReportEvent(childCtx, "waiting for complete disconnection")
-	ctxTimeout, cancel := context.WithTimeout(childCtx, 4*time.Second)
+	ctxTimeout, cancel := context.WithTimeout(childCtx, timeout)
path: packages/orchestrator/internal/sandbox/nbd/path_direct_linux.go
user: login: ValentaTomas
body: Agreed
created_at: 2025-04-04T19:37:27Z
author_association: MEMBER
diff_hunk: @@ -64,15 +69,17 @@ func NewDispatch(ctx context.Context, fp io.ReadWriteCloser, prov Provider) *Dis
 		fp:             fp,
 		prov:           prov,
 		ctx:            ctx,
+		fatal:          make(chan error, 1),
 	}
 
 	binary.BigEndian.PutUint32(d.responseHeader, NBDResponseMagic)
 	return d
 }
 
 func (d *Dispatch) Wait() {
path: packages/orchestrator/internal/sandbox/nbd/dispatch.go
user: login: jakubno
body: I would rename the method to Close/Shutdown
created_at: 2025-04-08T08:56:31Z
author_association: MEMBER
diff_hunk: @@ -41,74 +58,85 @@ func NewDirectPathMount(tracer trace.Tracer, b block.Device, devicePool *DeviceP
 		cancelfn:   cancelfn,
 		blockSize:  4096,
 		devicePool: devicePool,
+		socks:      make([]io.Closer, 0),
 	}
 }
 
-func (d *DirectPathMount) Open(ctx context.Context) (uint32, error) {
+func (d *DirectPathMount) Open(ctx context.Context) (deviceIndex uint32, err error) {
 	size, err := d.Backend.Size()
 	if err != nil {
 		return 0, err
 	}
 
-	for {
-		d.deviceIndex, err = d.devicePool.GetDevice(ctx)
-		if err != nil {
-			return 0, err
-		}
-
-		// Create the socket pairs
-		sockPair, err := syscall.Socketpair(syscall.AF_UNIX, syscall.SOCK_STREAM, 0)
-		if err != nil {
-			return 0, err
-		}
+	// TODO: Do we need retry getting the device or it is enough to retry the dispatchers?
path: packages/orchestrator/internal/sandbox/nbd/path_direct_linux.go
user: login: jakubno
body: What should be done?
created_at: 2025-04-08T09:03:44Z
author_association: MEMBER
diff_hunk: @@ -41,74 +58,85 @@ func NewDirectPathMount(tracer trace.Tracer, b block.Device, devicePool *DeviceP
 		cancelfn:   cancelfn,
 		blockSize:  4096,
 		devicePool: devicePool,
+		socks:      make([]io.Closer, 0),
 	}
 }
 
-func (d *DirectPathMount) Open(ctx context.Context) (uint32, error) {
+func (d *DirectPathMount) Open(ctx context.Context) (deviceIndex uint32, err error) {
 	size, err := d.Backend.Size()
 	if err != nil {
 		return 0, err
 	}
 
-	for {
-		d.deviceIndex, err = d.devicePool.GetDevice(ctx)
-		if err != nil {
-			return 0, err
-		}
-
-		// Create the socket pairs
-		sockPair, err := syscall.Socketpair(syscall.AF_UNIX, syscall.SOCK_STREAM, 0)
-		if err != nil {
-			return 0, err
-		}
+	// TODO: Do we need retry getting the device or it is enough to retry the dispatchers?
+	d.deviceIndex, err = d.devicePool.GetDevice(ctx)
+	if err != nil {
+		return 0, err
+	}
 
-		client := os.NewFile(uintptr(sockPair[0]), "client")
-		server := os.NewFile(uintptr(sockPair[1]), "server")
-		d.conn, err = net.FileConn(server)
+	for {
path: packages/orchestrator/internal/sandbox/nbd/path_direct_linux.go
user: login: jakubno
body: if you retry, you are not closing `d.dispatchers` and not cleaning and closing `d.socks`
created_at: 2025-04-08T09:17:55Z
author_association: MEMBER
diff_hunk: @@ -41,74 +58,85 @@ func NewDirectPathMount(tracer trace.Tracer, b block.Device, devicePool *DeviceP
 		cancelfn:   cancelfn,
 		blockSize:  4096,
 		devicePool: devicePool,
+		socks:      make([]io.Closer, 0),
 	}
 }
 
-func (d *DirectPathMount) Open(ctx context.Context) (uint32, error) {
+func (d *DirectPathMount) Open(ctx context.Context) (deviceIndex uint32, err error) {
 	size, err := d.Backend.Size()
 	if err != nil {
 		return 0, err
 	}
 
-	for {
-		d.deviceIndex, err = d.devicePool.GetDevice(ctx)
-		if err != nil {
-			return 0, err
-		}
-
-		// Create the socket pairs
-		sockPair, err := syscall.Socketpair(syscall.AF_UNIX, syscall.SOCK_STREAM, 0)
-		if err != nil {
-			return 0, err
-		}
+	// TODO: Do we need retry getting the device or it is enough to retry the dispatchers?
+	d.deviceIndex, err = d.devicePool.GetDevice(ctx)
+	if err != nil {
+		return 0, err
+	}
 
-		client := os.NewFile(uintptr(sockPair[0]), "client")
-		server := os.NewFile(uintptr(sockPair[1]), "server")
-		d.conn, err = net.FileConn(server)
+	for {
+		socks := make([]*os.File, 0)
+		d.dispatchers = make([]*Dispatch, 0)
+
+		for i := 0; i < connections; i++ {
+			// Create the socket pairs
+			sockPair, err := syscall.Socketpair(syscall.AF_UNIX, syscall.SOCK_STREAM, 0)
+			if err != nil {
+				return 0, err
+			}
 
-		if err != nil {
-			return 0, err
-		}
-		server.Close()
-
-		d.dispatcher = NewDispatch(d.ctx, d.conn, d.Backend)
-		// Start reading commands on the socket and dispatching them to our provider
-		go func() {
-			handleErr := d.dispatcher.Handle()
-			if handleErr != nil {
-				zap.L().Error("error handling NBD commands", zap.Error(handleErr))
+			client := os.NewFile(uintptr(sockPair[0]), "client")
+			server := os.NewFile(uintptr(sockPair[1]), "server")
+			serverc, err := net.FileConn(server)
+			if err != nil {
+				return 0, err
 			}
-		}()
+			server.Close()
+
+			dispatch := NewDispatch(d.ctx, serverc, d.Backend)
+			// Start reading commands on the socket and dispatching them to our provider
+			d.handlersWg.Add(1)
+			go func() {
+				defer d.handlersWg.Done()
+
+				handleErr := dispatch.Handle()
+				// The error is expected to happen if the nbd (socket connection) is closed
+				zap.L().Info("closing handler for NBD commands",
+					zap.Error(handleErr),
+					zap.Uint32("device_index", d.deviceIndex),
+					zap.Int("socket_index", i),
+				)
+			}()
+
+			d.socks = append(d.socks, serverc)
+			socks = append(socks, client)
+			d.dispatchers = append(d.dispatchers, dispatch)
+		}
 
 		var opts []nbdnl.ConnectOption
 		opts = append(opts, nbdnl.WithBlockSize(d.blockSize))
-		opts = append(opts, nbdnl.WithTimeout(5*time.Second))
-		opts = append(opts, nbdnl.WithDeadconnTimeout(5*time.Second))
+		opts = append(opts, nbdnl.WithTimeout(connectTimeout))
+		opts = append(opts, nbdnl.WithDeadconnTimeout(connectTimeout))
 
 		serverFlags := nbdnl.FlagHasFlags | nbdnl.FlagCanMulticonn
 
-		idx, err := nbdnl.Connect(d.deviceIndex, []*os.File{client}, uint64(size), 0, serverFlags, opts...)
+		idx, err := nbdnl.Connect(d.deviceIndex, socks, uint64(size), 0, serverFlags, opts...)
 		if err == nil {
+			// The idx should be the same as d.deviceIndex, because we are connecting to it,
+			// but we will use the one returned by nbdnl
 			d.deviceIndex = idx
+
 			break
 		}
 
+		zap.L().Error("error opening NBD, retrying", zap.Error(err), zap.Uint32("deviceIndex", d.deviceIndex))
+
 		// Sometimes (rare), there seems to be a BADF error here. Lets just retry for now...
 		// Close things down and try again...
-		_ = client.Close()
-
-		connErr := d.conn.Close()
-		if connErr != nil {
-			zap.L().Error("error closing conn", zap.Error(connErr))
+		for _, sock := range socks {
+			sock.Close()
 		}
 
-		releaseErr := d.devicePool.ReleaseDevice(d.deviceIndex)
path: packages/orchestrator/internal/sandbox/nbd/path_direct_linux.go
user: login: jakubno
body: there's no need to release because we didn't connect, do i understand it correctly right?
created_at: 2025-04-08T09:19:21Z
author_association: MEMBER
diff_hunk: @@ -64,15 +69,17 @@ func NewDispatch(ctx context.Context, fp io.ReadWriteCloser, prov Provider) *Dis
 		fp:             fp,
 		prov:           prov,
 		ctx:            ctx,
+		fatal:          make(chan error, 1),
 	}
 
 	binary.BigEndian.PutUint32(d.responseHeader, NBDResponseMagic)
 	return d
 }
 
 func (d *Dispatch) Wait() {
path: packages/orchestrator/internal/sandbox/nbd/dispatch.go
user: login: dobrac
body: renamed to Drain
created_at: 2025-04-08T09:33:31Z
author_association: CONTRIBUTOR
diff_hunk: @@ -41,74 +58,85 @@ func NewDirectPathMount(tracer trace.Tracer, b block.Device, devicePool *DeviceP
 		cancelfn:   cancelfn,
 		blockSize:  4096,
 		devicePool: devicePool,
+		socks:      make([]io.Closer, 0),
 	}
 }
 
-func (d *DirectPathMount) Open(ctx context.Context) (uint32, error) {
+func (d *DirectPathMount) Open(ctx context.Context) (deviceIndex uint32, err error) {
 	size, err := d.Backend.Size()
 	if err != nil {
 		return 0, err
 	}
 
-	for {
-		d.deviceIndex, err = d.devicePool.GetDevice(ctx)
-		if err != nil {
-			return 0, err
-		}
-
-		// Create the socket pairs
-		sockPair, err := syscall.Socketpair(syscall.AF_UNIX, syscall.SOCK_STREAM, 0)
-		if err != nil {
-			return 0, err
-		}
+	// TODO: Do we need retry getting the device or it is enough to retry the dispatchers?
+	d.deviceIndex, err = d.devicePool.GetDevice(ctx)
+	if err != nil {
+		return 0, err
+	}
 
-		client := os.NewFile(uintptr(sockPair[0]), "client")
-		server := os.NewFile(uintptr(sockPair[1]), "server")
-		d.conn, err = net.FileConn(server)
+	for {
path: packages/orchestrator/internal/sandbox/nbd/path_direct_linux.go
user: login: dobrac
body: socks should be closed, there is no drain await on dispatchers, but it's probably not necessary as it doesn't release any resources 
```
for _, sock := range socks {
	sock.Close()
}
```
		
		
created_at: 2025-04-08T09:36:58Z
author_association: CONTRIBUTOR
diff_hunk: @@ -41,74 +58,85 @@ func NewDirectPathMount(tracer trace.Tracer, b block.Device, devicePool *DeviceP
 		cancelfn:   cancelfn,
 		blockSize:  4096,
 		devicePool: devicePool,
+		socks:      make([]io.Closer, 0),
 	}
 }
 
-func (d *DirectPathMount) Open(ctx context.Context) (uint32, error) {
+func (d *DirectPathMount) Open(ctx context.Context) (deviceIndex uint32, err error) {
 	size, err := d.Backend.Size()
 	if err != nil {
 		return 0, err
 	}
 
-	for {
-		d.deviceIndex, err = d.devicePool.GetDevice(ctx)
-		if err != nil {
-			return 0, err
-		}
-
-		// Create the socket pairs
-		sockPair, err := syscall.Socketpair(syscall.AF_UNIX, syscall.SOCK_STREAM, 0)
-		if err != nil {
-			return 0, err
-		}
+	// TODO: Do we need retry getting the device or it is enough to retry the dispatchers?
+	d.deviceIndex, err = d.devicePool.GetDevice(ctx)
+	if err != nil {
+		return 0, err
+	}
 
-		client := os.NewFile(uintptr(sockPair[0]), "client")
-		server := os.NewFile(uintptr(sockPair[1]), "server")
-		d.conn, err = net.FileConn(server)
+	for {
+		socks := make([]*os.File, 0)
+		d.dispatchers = make([]*Dispatch, 0)
+
+		for i := 0; i < connections; i++ {
+			// Create the socket pairs
+			sockPair, err := syscall.Socketpair(syscall.AF_UNIX, syscall.SOCK_STREAM, 0)
+			if err != nil {
+				return 0, err
+			}
 
-		if err != nil {
-			return 0, err
-		}
-		server.Close()
-
-		d.dispatcher = NewDispatch(d.ctx, d.conn, d.Backend)
-		// Start reading commands on the socket and dispatching them to our provider
-		go func() {
-			handleErr := d.dispatcher.Handle()
-			if handleErr != nil {
-				zap.L().Error("error handling NBD commands", zap.Error(handleErr))
+			client := os.NewFile(uintptr(sockPair[0]), "client")
+			server := os.NewFile(uintptr(sockPair[1]), "server")
+			serverc, err := net.FileConn(server)
+			if err != nil {
+				return 0, err
 			}
-		}()
+			server.Close()
+
+			dispatch := NewDispatch(d.ctx, serverc, d.Backend)
+			// Start reading commands on the socket and dispatching them to our provider
+			d.handlersWg.Add(1)
+			go func() {
+				defer d.handlersWg.Done()
+
+				handleErr := dispatch.Handle()
+				// The error is expected to happen if the nbd (socket connection) is closed
+				zap.L().Info("closing handler for NBD commands",
+					zap.Error(handleErr),
+					zap.Uint32("device_index", d.deviceIndex),
+					zap.Int("socket_index", i),
+				)
+			}()
+
+			d.socks = append(d.socks, serverc)
+			socks = append(socks, client)
+			d.dispatchers = append(d.dispatchers, dispatch)
+		}
 
 		var opts []nbdnl.ConnectOption
 		opts = append(opts, nbdnl.WithBlockSize(d.blockSize))
-		opts = append(opts, nbdnl.WithTimeout(5*time.Second))
-		opts = append(opts, nbdnl.WithDeadconnTimeout(5*time.Second))
+		opts = append(opts, nbdnl.WithTimeout(connectTimeout))
+		opts = append(opts, nbdnl.WithDeadconnTimeout(connectTimeout))
 
 		serverFlags := nbdnl.FlagHasFlags | nbdnl.FlagCanMulticonn
 
-		idx, err := nbdnl.Connect(d.deviceIndex, []*os.File{client}, uint64(size), 0, serverFlags, opts...)
+		idx, err := nbdnl.Connect(d.deviceIndex, socks, uint64(size), 0, serverFlags, opts...)
 		if err == nil {
+			// The idx should be the same as d.deviceIndex, because we are connecting to it,
+			// but we will use the one returned by nbdnl
 			d.deviceIndex = idx
+
 			break
 		}
 
+		zap.L().Error("error opening NBD, retrying", zap.Error(err), zap.Uint32("deviceIndex", d.deviceIndex))
+
 		// Sometimes (rare), there seems to be a BADF error here. Lets just retry for now...
 		// Close things down and try again...
-		_ = client.Close()
-
-		connErr := d.conn.Close()
-		if connErr != nil {
-			zap.L().Error("error closing conn", zap.Error(connErr))
+		for _, sock := range socks {
+			sock.Close()
 		}
 
-		releaseErr := d.devicePool.ReleaseDevice(d.deviceIndex)
path: packages/orchestrator/internal/sandbox/nbd/path_direct_linux.go
user: login: dobrac
body: That goes to the TODO `// TODO: Do we need retry getting the device or it is enough to retry the dispatchers?`, so there is now an implementation, where the device is acquired at the beginning (Open) and released only on Close
created_at: 2025-04-08T09:38:25Z
author_association: CONTRIBUTOR
diff_hunk: @@ -41,74 +58,85 @@ func NewDirectPathMount(tracer trace.Tracer, b block.Device, devicePool *DeviceP
 		cancelfn:   cancelfn,
 		blockSize:  4096,
 		devicePool: devicePool,
+		socks:      make([]io.Closer, 0),
 	}
 }
 
-func (d *DirectPathMount) Open(ctx context.Context) (uint32, error) {
+func (d *DirectPathMount) Open(ctx context.Context) (deviceIndex uint32, err error) {
 	size, err := d.Backend.Size()
 	if err != nil {
 		return 0, err
 	}
 
-	for {
-		d.deviceIndex, err = d.devicePool.GetDevice(ctx)
-		if err != nil {
-			return 0, err
-		}
-
-		// Create the socket pairs
-		sockPair, err := syscall.Socketpair(syscall.AF_UNIX, syscall.SOCK_STREAM, 0)
-		if err != nil {
-			return 0, err
-		}
+	// TODO: Do we need retry getting the device or it is enough to retry the dispatchers?
path: packages/orchestrator/internal/sandbox/nbd/path_direct_linux.go
user: login: dobrac
body: @ValentaTomas asked here
> I'm not sure. @jakubno Do you remember if the device is unusable after fail or if just the connection was a problem?

I have no idea if it could cause any issue atm
created_at: 2025-04-08T09:40:16Z
author_association: CONTRIBUTOR
diff_hunk: @@ -41,74 +58,85 @@ func NewDirectPathMount(tracer trace.Tracer, b block.Device, devicePool *DeviceP
 		cancelfn:   cancelfn,
 		blockSize:  4096,
 		devicePool: devicePool,
+		socks:      make([]io.Closer, 0),
 	}
 }
 
-func (d *DirectPathMount) Open(ctx context.Context) (uint32, error) {
+func (d *DirectPathMount) Open(ctx context.Context) (deviceIndex uint32, err error) {
 	size, err := d.Backend.Size()
 	if err != nil {
 		return 0, err
 	}
 
-	for {
-		d.deviceIndex, err = d.devicePool.GetDevice(ctx)
-		if err != nil {
-			return 0, err
-		}
-
-		// Create the socket pairs
-		sockPair, err := syscall.Socketpair(syscall.AF_UNIX, syscall.SOCK_STREAM, 0)
-		if err != nil {
-			return 0, err
-		}
+	// TODO: Do we need retry getting the device or it is enough to retry the dispatchers?
+	d.deviceIndex, err = d.devicePool.GetDevice(ctx)
+	if err != nil {
+		return 0, err
+	}
 
-		client := os.NewFile(uintptr(sockPair[0]), "client")
-		server := os.NewFile(uintptr(sockPair[1]), "server")
-		d.conn, err = net.FileConn(server)
+	for {
path: packages/orchestrator/internal/sandbox/nbd/path_direct_linux.go
user: login: jakubno
body: fair for dispatchers
for socket [this](https://github.com/e2b-dev/infra/pull/517/files/0820688b5c8b4b0d17f10246e9e44e47ace61c8d#diff-f8e9ba3bee1cac243166476cd3fad1b2d83c48a0ebe6699d5193c98839a8fb1dR90) isn't closed 

```go
serverc, err := net.FileConn(server)
```
created_at: 2025-04-08T09:54:47Z
author_association: MEMBER
diff_hunk: @@ -41,74 +58,85 @@ func NewDirectPathMount(tracer trace.Tracer, b block.Device, devicePool *DeviceP
 		cancelfn:   cancelfn,
 		blockSize:  4096,
 		devicePool: devicePool,
+		socks:      make([]io.Closer, 0),
 	}
 }
 
-func (d *DirectPathMount) Open(ctx context.Context) (uint32, error) {
+func (d *DirectPathMount) Open(ctx context.Context) (deviceIndex uint32, err error) {
 	size, err := d.Backend.Size()
 	if err != nil {
 		return 0, err
 	}
 
-	for {
-		d.deviceIndex, err = d.devicePool.GetDevice(ctx)
-		if err != nil {
-			return 0, err
-		}
-
-		// Create the socket pairs
-		sockPair, err := syscall.Socketpair(syscall.AF_UNIX, syscall.SOCK_STREAM, 0)
-		if err != nil {
-			return 0, err
-		}
+	// TODO: Do we need retry getting the device or it is enough to retry the dispatchers?
+	d.deviceIndex, err = d.devicePool.GetDevice(ctx)
+	if err != nil {
+		return 0, err
+	}
 
-		client := os.NewFile(uintptr(sockPair[0]), "client")
-		server := os.NewFile(uintptr(sockPair[1]), "server")
-		d.conn, err = net.FileConn(server)
+	for {
path: packages/orchestrator/internal/sandbox/nbd/path_direct_linux.go
user: login: dobrac
body: added closing server and client sockets and retry and on nbd close
created_at: 2025-04-08T11:35:50Z
author_association: CONTRIBUTOR
diff_hunk: @@ -143,10 +180,23 @@ func (d *DirectPathMount) Close(ctx context.Context) error {
 	telemetry.ReportEvent(childCtx, "canceling context")
 	d.cancelfn()
 
+	// Close all server socket pairs...
+	telemetry.ReportEvent(childCtx, "closing socket pairs server")
+	for _, v := range d.socksServer {
+		err := v.Close()
+		if err != nil {
+			return err
path: packages/orchestrator/internal/sandbox/nbd/path_direct_linux.go
user: login: ValentaTomas
body: Should we return error here or report error and continue so we can finish the device release?
The error return was the previous behavior, so it should be ok to do this now, I just want to mention this as it might be problem in some cases.
created_at: 2025-04-08T20:24:37Z
author_association: MEMBER
diff_hunk: @@ -178,9 +221,40 @@ func (d *DirectPathMount) Close(ctx context.Context) error {
 		if err == nil && !s.Connected {
 			break
 		}
-
 		time.Sleep(100 * time.Nanosecond)
 	}
 
+	// Close all client socket pairs...
+	telemetry.ReportEvent(childCtx, "closing socket pairs client")
+	for _, v := range d.socksClient {
+		err := v.Close()
path: packages/orchestrator/internal/sandbox/nbd/path_direct_linux.go
user: login: ValentaTomas
body: I don't think we were previously closing the client side—could there be a chance of this returning error if the client side was closed automatically because the server side was closed?
created_at: 2025-04-08T20:26:53Z
author_association: MEMBER
diff_hunk: @@ -41,74 +58,85 @@ func NewDirectPathMount(tracer trace.Tracer, b block.Device, devicePool *DeviceP
 		cancelfn:   cancelfn,
 		blockSize:  4096,
 		devicePool: devicePool,
+		socks:      make([]io.Closer, 0),
 	}
 }
 
-func (d *DirectPathMount) Open(ctx context.Context) (uint32, error) {
+func (d *DirectPathMount) Open(ctx context.Context) (deviceIndex uint32, err error) {
 	size, err := d.Backend.Size()
 	if err != nil {
 		return 0, err
 	}
 
-	for {
-		d.deviceIndex, err = d.devicePool.GetDevice(ctx)
-		if err != nil {
-			return 0, err
-		}
-
-		// Create the socket pairs
-		sockPair, err := syscall.Socketpair(syscall.AF_UNIX, syscall.SOCK_STREAM, 0)
-		if err != nil {
-			return 0, err
-		}
+	// TODO: Do we need retry getting the device or it is enough to retry the dispatchers?
path: packages/orchestrator/internal/sandbox/nbd/path_direct_linux.go
user: login: ValentaTomas
body: @jakubno In the previous implementation when the NBD connection was not working we retried by trying to use a new NBD and new connection. 
Was using a new NBD necessary or could just a new connection work? I don't remember how many experiments we did with this.
created_at: 2025-04-08T20:29:42Z
author_association: MEMBER
diff_hunk: @@ -178,9 +221,40 @@ func (d *DirectPathMount) Close(ctx context.Context) error {
 		if err == nil && !s.Connected {
 			break
 		}
-
 		time.Sleep(100 * time.Nanosecond)
 	}
 
+	// Close all client socket pairs...
+	telemetry.ReportEvent(childCtx, "closing socket pairs client")
+	for _, v := range d.socksClient {
+		err := v.Close()
path: packages/orchestrator/internal/sandbox/nbd/path_direct_linux.go
user: login: dobrac
body: There wasn't, I've added it as the client sockets were closed only on retry and I haven't found them to be automatically closed. Based on testing, it doesn't throw any errors. But, we can just log the errors and don't return the function
created_at: 2025-04-09T07:17:52Z
author_association: CONTRIBUTOR
diff_hunk: @@ -143,10 +180,23 @@ func (d *DirectPathMount) Close(ctx context.Context) error {
 	telemetry.ReportEvent(childCtx, "canceling context")
 	d.cancelfn()
 
+	// Close all server socket pairs...
+	telemetry.ReportEvent(childCtx, "closing socket pairs server")
+	for _, v := range d.socksServer {
+		err := v.Close()
+		if err != nil {
+			return err
path: packages/orchestrator/internal/sandbox/nbd/path_direct_linux.go
user: login: dobrac
body: I'm not sure that if one of the Close errors out, if this will correctly finish and won't deadlock the Stop:
```
for _, d := range d.dispatchers {
	d.Drain()
}
```
created_at: 2025-04-09T07:19:16Z
author_association: CONTRIBUTOR
diff_hunk: @@ -178,9 +221,40 @@ func (d *DirectPathMount) Close(ctx context.Context) error {
 		if err == nil && !s.Connected {
 			break
 		}
-
 		time.Sleep(100 * time.Nanosecond)
 	}
 
+	// Close all client socket pairs...
+	telemetry.ReportEvent(childCtx, "closing socket pairs client")
+	for _, v := range d.socksClient {
+		err := v.Close()
path: packages/orchestrator/internal/sandbox/nbd/path_direct_linux.go
user: login: dobrac
body: Added "soft" error log here and removed the return statement
created_at: 2025-04-09T13:17:52Z
author_association: CONTRIBUTOR
diff_hunk: @@ -143,10 +180,23 @@ func (d *DirectPathMount) Close(ctx context.Context) error {
 	telemetry.ReportEvent(childCtx, "canceling context")
 	d.cancelfn()
 
+	// Close all server socket pairs...
+	telemetry.ReportEvent(childCtx, "closing socket pairs server")
+	for _, v := range d.socksServer {
+		err := v.Close()
+		if err != nil {
+			return err
path: packages/orchestrator/internal/sandbox/nbd/path_direct_linux.go
user: login: dobrac
body: changed to soft errors with additional checks
created_at: 2025-04-09T14:31:49Z
author_association: CONTRIBUTOR
diff_hunk: @@ -41,74 +58,85 @@ func NewDirectPathMount(tracer trace.Tracer, b block.Device, devicePool *DeviceP
 		cancelfn:   cancelfn,
 		blockSize:  4096,
 		devicePool: devicePool,
+		socks:      make([]io.Closer, 0),
 	}
 }
 
-func (d *DirectPathMount) Open(ctx context.Context) (uint32, error) {
+func (d *DirectPathMount) Open(ctx context.Context) (deviceIndex uint32, err error) {
 	size, err := d.Backend.Size()
 	if err != nil {
 		return 0, err
 	}
 
-	for {
-		d.deviceIndex, err = d.devicePool.GetDevice(ctx)
-		if err != nil {
-			return 0, err
-		}
-
-		// Create the socket pairs
-		sockPair, err := syscall.Socketpair(syscall.AF_UNIX, syscall.SOCK_STREAM, 0)
-		if err != nil {
-			return 0, err
-		}
+	// TODO: Do we need retry getting the device or it is enough to retry the dispatchers?
path: packages/orchestrator/internal/sandbox/nbd/path_direct_linux.go
user: login: jakubno
body: I don't remember anymore, I would return the logic that was there. I would add some logging we can monitor it and if there's no problem then we can move it to only socket pairs or delete it completely  
created_at: 2025-04-09T14:58:19Z
author_association: MEMBER
diff_hunk: @@ -19,153 +24,237 @@ import (
 	"github.com/e2b-dev/infra/packages/shared/pkg/telemetry"
 )
 
+const (
+	connections    = 4
+	connectTimeout = 30 * time.Second
+
+	// disconnectTimeout should not be necessary if the disconnect is reliable
+	disconnectTimeout = 30 * time.Second
+)
+
 type DirectPathMount struct {
-	tracer      trace.Tracer
+	tracer   trace.Tracer
+	ctx      context.Context
+	cancelfn context.CancelFunc
+
+	devicePool *DevicePool
+
 	Backend     block.Device
-	ctx         context.Context
-	dispatcher  *Dispatch
-	conn        net.Conn
 	deviceIndex uint32
 	blockSize   uint64
-	cancelfn    context.CancelFunc
-	devicePool  *DevicePool
+
+	dispatchers []*Dispatch
+	socksClient []*os.File
+	socksServer []io.Closer
+
+	handlersWg sync.WaitGroup
 }
 
 func NewDirectPathMount(tracer trace.Tracer, b block.Device, devicePool *DevicePool) *DirectPathMount {
 	ctx, cancelfn := context.WithCancel(context.Background())
 
 	return &DirectPathMount{
-		tracer:     tracer,
-		Backend:    b,
-		ctx:        ctx,
-		cancelfn:   cancelfn,
-		blockSize:  4096,
-		devicePool: devicePool,
+		tracer:      tracer,
+		Backend:     b,
+		ctx:         ctx,
+		cancelfn:    cancelfn,
+		blockSize:   4096,
+		devicePool:  devicePool,
+		socksClient: make([]*os.File, 0),
+		socksServer: make([]io.Closer, 0),
+		deviceIndex: math.MaxUint32,
 	}
 }
 
-func (d *DirectPathMount) Open(ctx context.Context) (uint32, error) {
+func (d *DirectPathMount) Open(ctx context.Context) (retDeviceIndex uint32, err error) {
+	defer func() {
+		// Set the device index to the one returned, correctly capture error values
+		d.deviceIndex = retDeviceIndex
+		zap.L().Debug("opening direct path mount", zap.Uint32("device_index", d.deviceIndex), zap.Error(err))
+	}()
+
 	size, err := d.Backend.Size()
 	if err != nil {
-		return 0, err
+		return math.MaxUint32, err
 	}
 
+	deviceIndex := uint32(math.MaxUint32)
+
 	for {
-		d.deviceIndex, err = d.devicePool.GetDevice(ctx)
+		deviceIndex, err = d.devicePool.GetDevice(ctx)
 		if err != nil {
-			return 0, err
+			return math.MaxUint32, err
 		}
 
-		// Create the socket pairs
-		sockPair, err := syscall.Socketpair(syscall.AF_UNIX, syscall.SOCK_STREAM, 0)
-		if err != nil {
-			return 0, err
-		}
+		d.socksClient = make([]*os.File, 0)
+		d.socksServer = make([]io.Closer, 0)
+		d.dispatchers = make([]*Dispatch, 0)
 
-		client := os.NewFile(uintptr(sockPair[0]), "client")
-		server := os.NewFile(uintptr(sockPair[1]), "server")
-		d.conn, err = net.FileConn(server)
+		for i := 0; i < connections; i++ {
+			// Create the socket pairs
+			sockPair, err := syscall.Socketpair(syscall.AF_UNIX, syscall.SOCK_STREAM, 0)
+			if err != nil {
+				return math.MaxUint32, err
+			}
 
-		if err != nil {
-			return 0, err
-		}
-		server.Close()
-
-		d.dispatcher = NewDispatch(d.ctx, d.conn, d.Backend)
-		// Start reading commands on the socket and dispatching them to our provider
-		go func() {
-			handleErr := d.dispatcher.Handle()
-			if handleErr != nil {
-				zap.L().Error("error handling NBD commands", zap.Error(handleErr))
+			client := os.NewFile(uintptr(sockPair[0]), "client")
+			server := os.NewFile(uintptr(sockPair[1]), "server")
+			serverc, err := net.FileConn(server)
+			if err != nil {
+				return math.MaxUint32, err
 			}
-		}()
+			server.Close()
+
+			dispatch := NewDispatch(d.ctx, serverc, d.Backend)
+			// Start reading commands on the socket and dispatching them to our provider
+			d.handlersWg.Add(1)
+			go func() {
+				defer d.handlersWg.Done()
+
+				handleErr := dispatch.Handle()
+				// The error is expected to happen if the nbd (socket connection) is closed
+				zap.L().Info("closing handler for NBD commands",
+					zap.Error(handleErr),
+					zap.Uint32("device_index", deviceIndex),
+					zap.Int("socket_index", i),
+				)
+			}()
+
+			d.socksServer = append(d.socksServer, serverc)
+			d.socksClient = append(d.socksClient, client)
+			d.dispatchers = append(d.dispatchers, dispatch)
+		}
 
 		var opts []nbdnl.ConnectOption
 		opts = append(opts, nbdnl.WithBlockSize(d.blockSize))
-		opts = append(opts, nbdnl.WithTimeout(5*time.Second))
-		opts = append(opts, nbdnl.WithDeadconnTimeout(5*time.Second))
+		opts = append(opts, nbdnl.WithTimeout(connectTimeout))
+		opts = append(opts, nbdnl.WithDeadconnTimeout(connectTimeout))
 
 		serverFlags := nbdnl.FlagHasFlags | nbdnl.FlagCanMulticonn
 
-		idx, err := nbdnl.Connect(d.deviceIndex, []*os.File{client}, uint64(size), 0, serverFlags, opts...)
+		idx, err := nbdnl.Connect(deviceIndex, d.socksClient, uint64(size), 0, serverFlags, opts...)
 		if err == nil {
-			d.deviceIndex = idx
+			// The idx should be the same as deviceIndex, because we are connecting to it,
+			// but we will use the one returned by nbdnl
+			deviceIndex = idx
+
 			break
 		}
 
+		zap.L().Error("error opening NBD, retrying", zap.Error(err), zap.Uint32("device_index", deviceIndex))
+
 		// Sometimes (rare), there seems to be a BADF error here. Lets just retry for now...
 		// Close things down and try again...
-		_ = client.Close()
-
-		connErr := d.conn.Close()
-		if connErr != nil {
-			zap.L().Error("error closing conn", zap.Error(connErr))
+		for _, sock := range d.socksClient {
+			sock.Close()
 		}
-
-		releaseErr := d.devicePool.ReleaseDevice(d.deviceIndex)
-		if releaseErr != nil {
-			zap.L().Error("error releasing device", zap.Error(releaseErr))
+		for _, sock := range d.socksServer {
+			sock.Close()
 		}
-
-		d.deviceIndex = 0
+		// Release the device back to the pool
+		d.devicePool.ReleaseDevice(idx)
path: packages/orchestrator/internal/sandbox/nbd/path_direct_linux.go
user: login: jakubno
body: shoudln't we at least log the error 
created_at: 2025-04-25T15:34:45Z
author_association: MEMBER
diff_hunk: @@ -19,153 +24,237 @@ import (
 	"github.com/e2b-dev/infra/packages/shared/pkg/telemetry"
 )
 
+const (
+	connections    = 4
+	connectTimeout = 30 * time.Second
+
+	// disconnectTimeout should not be necessary if the disconnect is reliable
+	disconnectTimeout = 30 * time.Second
+)
+
 type DirectPathMount struct {
-	tracer      trace.Tracer
+	tracer   trace.Tracer
+	ctx      context.Context
+	cancelfn context.CancelFunc
+
+	devicePool *DevicePool
+
 	Backend     block.Device
-	ctx         context.Context
-	dispatcher  *Dispatch
-	conn        net.Conn
 	deviceIndex uint32
 	blockSize   uint64
-	cancelfn    context.CancelFunc
-	devicePool  *DevicePool
+
+	dispatchers []*Dispatch
+	socksClient []*os.File
+	socksServer []io.Closer
+
+	handlersWg sync.WaitGroup
 }
 
 func NewDirectPathMount(tracer trace.Tracer, b block.Device, devicePool *DevicePool) *DirectPathMount {
 	ctx, cancelfn := context.WithCancel(context.Background())
 
 	return &DirectPathMount{
-		tracer:     tracer,
-		Backend:    b,
-		ctx:        ctx,
-		cancelfn:   cancelfn,
-		blockSize:  4096,
-		devicePool: devicePool,
+		tracer:      tracer,
+		Backend:     b,
+		ctx:         ctx,
+		cancelfn:    cancelfn,
+		blockSize:   4096,
+		devicePool:  devicePool,
+		socksClient: make([]*os.File, 0),
+		socksServer: make([]io.Closer, 0),
+		deviceIndex: math.MaxUint32,
 	}
 }
 
-func (d *DirectPathMount) Open(ctx context.Context) (uint32, error) {
+func (d *DirectPathMount) Open(ctx context.Context) (retDeviceIndex uint32, err error) {
path: packages/orchestrator/internal/sandbox/nbd/path_direct_linux.go
user: login: jakubno
body: why not just like this?
```suggestion
func (d *DirectPathMount) Open(ctx context.Context) (deviceIndex uint32, err error) {
```
created_at: 2025-04-25T15:36:22Z
author_association: MEMBER
diff_hunk: @@ -19,153 +24,237 @@ import (
 	"github.com/e2b-dev/infra/packages/shared/pkg/telemetry"
 )
 
+const (
+	connections    = 4
+	connectTimeout = 30 * time.Second
+
+	// disconnectTimeout should not be necessary if the disconnect is reliable
+	disconnectTimeout = 30 * time.Second
+)
+
 type DirectPathMount struct {
-	tracer      trace.Tracer
+	tracer   trace.Tracer
+	ctx      context.Context
+	cancelfn context.CancelFunc
+
+	devicePool *DevicePool
+
 	Backend     block.Device
-	ctx         context.Context
-	dispatcher  *Dispatch
-	conn        net.Conn
 	deviceIndex uint32
 	blockSize   uint64
-	cancelfn    context.CancelFunc
-	devicePool  *DevicePool
+
+	dispatchers []*Dispatch
+	socksClient []*os.File
+	socksServer []io.Closer
+
+	handlersWg sync.WaitGroup
 }
 
 func NewDirectPathMount(tracer trace.Tracer, b block.Device, devicePool *DevicePool) *DirectPathMount {
 	ctx, cancelfn := context.WithCancel(context.Background())
 
 	return &DirectPathMount{
-		tracer:     tracer,
-		Backend:    b,
-		ctx:        ctx,
-		cancelfn:   cancelfn,
-		blockSize:  4096,
-		devicePool: devicePool,
+		tracer:      tracer,
+		Backend:     b,
+		ctx:         ctx,
+		cancelfn:    cancelfn,
+		blockSize:   4096,
+		devicePool:  devicePool,
+		socksClient: make([]*os.File, 0),
+		socksServer: make([]io.Closer, 0),
+		deviceIndex: math.MaxUint32,
 	}
 }
 
-func (d *DirectPathMount) Open(ctx context.Context) (uint32, error) {
+func (d *DirectPathMount) Open(ctx context.Context) (retDeviceIndex uint32, err error) {
path: packages/orchestrator/internal/sandbox/nbd/path_direct_linux.go
user: login: dobrac
body: it would be in conflict with local variable name deviceIndex
created_at: 2025-04-28T10:20:53Z
author_association: CONTRIBUTOR

PR comments:
user: login: dobrac
created_at: 2025-04-04T14:15:48Z
body: Ready for a review. I've tested it using simulating traffic with pause/resume 3x for each sandbox for about a half an hour. I've not encountered any errors or undesired behavior.

> When retrying Open retry just connection establishment, without getting new device (not 100% sure this won't be a regression)

I'm not sure about this either, by default there is nbdnl.IndexAny so they are effectively changing the device index for each retry. If you are not sure either, I would remove it from this PR.  
author_association: CONTRIBUTOR
user: login: ValentaTomas
created_at: 2025-04-04T18:58:18Z
body: <!-- linear:replyTo:2728270159 -->

I'm not sure. @jakubno Do you remember if the device is unusable after fail or if just the connection was a problem?
author_association: MEMBER
user: login: ValentaTomas
created_at: 2025-04-04T21:41:01Z
body: Before we merge, can we do a simple benchmark to see how the number of connections affects the FS speed?
If it is significant, we will announce this.
author_association: MEMBER
user: login: dobrac
created_at: 2025-04-06T13:02:53Z
body: > Before we merge, can we do a simple benchmark to see how the number of connections affects the FS speed? If it is significant, we will announce this.

Simple test using the command `curl -sL https://yabs.sh | bash -s -- -ig` doesn't show any significant difference in disk speed when changing the number of socket connections.
author_association: CONTRIBUTOR
user: login: ValentaTomas
created_at: 2025-04-06T19:25:48Z
body: <!-- linear:replyTo:2728270159 -->

Because of the layers (FC->kernel->NBD) I think this might affect only specific tests like *reading* enough data to profit from parallelized connections, where the data must not be cached on the system level already (so data previously not accessed by this FC). The improvements to writes might be mostly hidden as that will be flushing in the background.
author_association: MEMBER
user: login: jakubno
created_at: 2025-04-07T07:15:04Z
body: What about downloading the initial template download?
author_association: MEMBER
user: login: ValentaTomas
created_at: 2025-04-07T18:10:47Z
body: <!-- linear:replyTo:2728270159 -->

That is related, because it can be a bottleneck, but this could probably allow additional parallelism of the download too and the read should be affected even if the template is downloaded.
author_association: MEMBER
user: login: ValentaTomas
created_at: 2025-04-08T20:32:26Z
body: Let's resolve the discussions and merge—we won't be doing more detailed benchmarking now.
author_association: MEMBER
user: login: dobrac
created_at: 2025-04-09T14:52:45Z
body: Done few additional changes:
- Added soft errors on close instead of hard return
- return errors from the cow close
- get new device for each try on NBD open for now (there is added logging, so we can monitor if it even happens)
- return math.MaxUint32 for errors in nbd to not confuse the deviceIndex 0 with error state (as we might still call Close on errored NBD connection)
author_association: CONTRIBUTOR
user: login: ValentaTomas
created_at: 2025-04-22T14:29:19Z
body: Let's update and merge this then (@jakubno review? :D)
author_association: MEMBER