feat(orchestrator): drain sandboxes during shutdown

Introduce a shared draingate.Gate (counter plus notification channel)
and use it in the sandbox factory and gRPC server to reject new
sandbox starts while draining, wait for in-flight starts, and drain
or force-stop live sandboxes before closers run. Forced shutdown
preserves buffered close errors on context cancellation and avoids
duplicate final-pass errors. Adds utils.WaitGroupWait to wait on a
WaitGroup with context cancellation.

The graceful drain phase is bounded by SHUTDOWN_DRAIN_TIMEOUT; when
it expires the drain escalates to a forced sandbox shutdown. By
default the drain waits forever, until sandboxes exit on their own or
a force-stop API call empties the node.

Author: wj-e2b

packages/orchestrator/pkg/cfg/model.go

@@ -105,6 +105,10 @@ type Config struct {
 	NBDPoolSize                 int               `env:"NBD_POOL_SIZE"                 envDefault:"64"`
 	Services                    []string          `env:"ORCHESTRATOR_SERVICES"         envDefault:"orchestrator"`
 	PersistentVolumeMounts      map[string]string `env:"PERSISTENT_VOLUME_MOUNTS"`
+	// ShutdownDrainTimeout bounds the graceful drain phase during shutdown.
+	// Zero (the default) drains forever: shutdown waits until sandboxes exit
+	// on their own or a force-stop API call empties the node.
+	ShutdownDrainTimeout time.Duration `env:"SHUTDOWN_DRAIN_TIMEOUT"`
 }
 
 // AdditionalClickhouseEndpoints returns the non-blank entries from

packages/orchestrator/pkg/draingate/gate.go

@@ -0,0 +1,133 @@
+package draingate
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"sync"
+)
+
+var ErrDraining = errors.New("drain gate is draining")
+
+type Gate struct {
+	initOnce  sync.Once
+	drainOnce sync.Once
+
+	mu       sync.Mutex
+	done     chan struct{}
+	draining bool
+	count    int
+	changed  chan struct{}
+}
+
+func New() *Gate {
+	g := &Gate{}
+	g.init()
+
+	return g
+}
+
+func (g *Gate) init() {
+	g.initOnce.Do(func() {
+		g.done = make(chan struct{})
+		g.changed = make(chan struct{})
+	})
+}
+
+func (g *Gate) StartDraining() bool {
+	if g == nil {
+		return false
+	}
+
+	g.init()
+	transitioned := false
+	g.drainOnce.Do(func() {
+		g.mu.Lock()
+		defer g.mu.Unlock()
+
+		g.draining = true
+		transitioned = true
+		close(g.done)
+	})
+
+	return transitioned
+}
+
+func (g *Gate) Draining() bool {
+	if g == nil {
+		return false
+	}
+
+	g.init()
+	select {
+	case <-g.done:
+		return true
+	default:
+		return false
+	}
+}
+
+func (g *Gate) Done() <-chan struct{} {
+	if g == nil {
+		return nil
+	}
+
+	g.init()
+
+	return g.done
+}
+
+func (g *Gate) Enter() (func(), error) {
+	if g == nil {
+		return func() {}, nil
+	}
+
+	g.init()
+	g.mu.Lock()
+	defer g.mu.Unlock()
+
+	if g.draining {
+		return nil, ErrDraining
+	}
+
+	g.count++
+	g.notifyChangeLocked()
+
+	return sync.OnceFunc(func() {
+		g.mu.Lock()
+		defer g.mu.Unlock()
+
+		g.count--
+		g.notifyChangeLocked()
+	}), nil
+}
+
+func (g *Gate) Wait(ctx context.Context) error {
+	if g == nil {
+		return nil
+	}
+
+	g.init()
+	for {
+		g.mu.Lock()
+		if g.count == 0 {
+			g.mu.Unlock()
+
+			return nil
+		}
+
+		changed := g.changed
+		g.mu.Unlock()
+
+		select {
+		case <-ctx.Done():
+			return fmt.Errorf("%w", ctx.Err())
+		case <-changed:
+		}
+	}
+}
+
+func (g *Gate) notifyChangeLocked() {
+	close(g.changed)
+	g.changed = make(chan struct{})
+}

packages/orchestrator/pkg/draingate/gate_test.go

@@ -0,0 +1,207 @@
+package draingate
+
+import (
+	"context"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestEnterReleaseAndWait(t *testing.T) {
+	t.Parallel()
+
+	g := New()
+	release, err := g.Enter()
+	require.NoError(t, err)
+
+	waitDone := make(chan error, 1)
+	go func() {
+		waitDone <- g.Wait(t.Context())
+	}()
+
+	requireNotDone(t, waitDone)
+	release()
+	release()
+	require.NoError(t, requireDone(t, waitDone))
+}
+
+func TestRejectAfterDrain(t *testing.T) {
+	t.Parallel()
+
+	g := New()
+	require.True(t, g.StartDraining())
+	require.False(t, g.StartDraining())
+	require.True(t, g.Draining())
+
+	select {
+	case <-g.Done():
+	default:
+		t.Fatal("Done channel was not closed")
+	}
+
+	release, err := g.Enter()
+	require.ErrorIs(t, err, ErrDraining)
+	require.Nil(t, release)
+}
+
+func TestWaitBlocksUntilAllReleases(t *testing.T) {
+	t.Parallel()
+
+	g := New()
+	release1, err := g.Enter()
+	require.NoError(t, err)
+	release2, err := g.Enter()
+	require.NoError(t, err)
+
+	g.StartDraining()
+	waitDone := make(chan error, 1)
+	go func() {
+		waitDone <- g.Wait(t.Context())
+	}()
+
+	requireNotDone(t, waitDone)
+	release1()
+	requireNotDone(t, waitDone)
+	release2()
+	require.NoError(t, requireDone(t, waitDone))
+}
+
+func TestWaitReturnsContextError(t *testing.T) {
+	t.Parallel()
+
+	g := New()
+	release, err := g.Enter()
+	require.NoError(t, err)
+	defer release()
+
+	ctx, cancel := context.WithCancel(t.Context())
+	cancel()
+
+	require.ErrorIs(t, g.Wait(ctx), context.Canceled)
+}
+
+func TestEnterDuringWaitWhileNotDrainingIsNotBlocked(t *testing.T) {
+	t.Parallel()
+
+	g := New()
+	release, err := g.Enter()
+	require.NoError(t, err)
+
+	waitDone := make(chan error, 1)
+	go func() {
+		waitDone <- g.Wait(t.Context())
+	}()
+	requireNotDone(t, waitDone)
+
+	entered := make(chan error, 1)
+	go func() {
+		secondRelease, err := g.Enter()
+		if err == nil {
+			secondRelease()
+		}
+		entered <- err
+	}()
+
+	require.NoError(t, requireDone(t, entered))
+	release()
+	require.NoError(t, requireDone(t, waitDone))
+}
+
+func TestConcurrentStress(t *testing.T) {
+	t.Parallel()
+
+	g := New()
+	start := make(chan struct{})
+	entered := make(chan func(), 100)
+	rejected := make(chan error, 100)
+	var wg sync.WaitGroup
+
+	for range 100 {
+		wg.Go(func() {
+			<-start
+			release, err := g.Enter()
+			if err != nil {
+				rejected <- err
+
+				return
+			}
+
+			entered <- release
+		})
+	}
+
+	close(start)
+	require.Eventually(t, func() bool {
+		return len(entered)+len(rejected) == 100
+	}, time.Second, time.Millisecond)
+
+	g.StartDraining()
+	wg.Wait()
+	close(entered)
+	close(rejected)
+
+	for err := range rejected {
+		require.ErrorIs(t, err, ErrDraining)
+	}
+
+	waitDone := make(chan error, 1)
+	go func() {
+		waitDone <- g.Wait(t.Context())
+	}()
+
+	for release := range entered {
+		release()
+	}
+	require.NoError(t, requireDone(t, waitDone))
+
+	_, err := g.Enter()
+	require.ErrorIs(t, err, ErrDraining)
+}
+
+func TestZeroValueAndNilGateAreSafe(t *testing.T) {
+	t.Parallel()
+
+	var g Gate
+	release, err := g.Enter()
+	require.NoError(t, err)
+	release()
+	require.NoError(t, g.Wait(t.Context()))
+	require.True(t, g.StartDraining())
+	require.True(t, g.Draining())
+
+	var nilGate *Gate
+	release, err = nilGate.Enter()
+	require.NoError(t, err)
+	release()
+	require.NoError(t, nilGate.Wait(t.Context()))
+	require.False(t, nilGate.StartDraining())
+	require.False(t, nilGate.Draining())
+	require.Nil(t, nilGate.Done())
+}
+
+func requireDone[T any](t *testing.T, ch <-chan T) T {
+	t.Helper()
+
+	select {
+	case got := <-ch:
+		return got
+	case <-time.After(time.Second):
+		t.Fatal("channel did not receive")
+
+		var zero T
+
+		return zero
+	}
+}
+
+func requireNotDone[T any](t *testing.T, ch <-chan T) {
+	t.Helper()
+
+	select {
+	case got := <-ch:
+		t.Fatalf("channel received unexpectedly: %v", got)
+	case <-time.After(25 * time.Millisecond):
+	}
+}