feat(orchestrator): graceful template-build drain on shutdown

packages/orchestrator/pkg/template/server/drain_test.go

@@ -0,0 +1,46 @@
+//go:build linux
+
+package server
+
+import (
+	"context"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/e2b-dev/infra/packages/shared/pkg/logger"
+)
+
+func TestWaitReturnsErrorWhileBuildsInFlight(t *testing.T) {
+	t.Parallel()
+
+	s := &ServerStore{
+		logger: logger.NewNopLogger(),
+		wg:     &sync.WaitGroup{},
+	}
+	s.wg.Add(1)
+	defer s.wg.Done()
+
+	ctx, cancel := context.WithCancel(t.Context())
+	cancel()
+
+	require.ErrorIs(t, s.Wait(ctx), context.Canceled)
+}
+
+func TestWaitStopsAtContextDeadlineDuringGracePeriod(t *testing.T) {
+	t.Parallel()
+
+	s := &ServerStore{
+		logger: logger.NewNopLogger(),
+		wg:     &sync.WaitGroup{},
+	}
+
+	// No builds in flight, so the wait passes straight to the consumer grace
+	// period, which must be bounded by ctx rather than the full sleep.
+	ctx, cancel := context.WithTimeout(t.Context(), 50*time.Millisecond)
+	defer cancel()
+
+	require.ErrorIs(t, s.Wait(ctx), context.DeadlineExceeded)
+}

packages/orchestrator/pkg/template/server/main.go

@@ -27,12 +27,18 @@ import (
 	templatemanager "github.com/e2b-dev/infra/packages/shared/pkg/grpc/template-manager"
 	"github.com/e2b-dev/infra/packages/shared/pkg/logger"
 	"github.com/e2b-dev/infra/packages/shared/pkg/storage"
+	"github.com/e2b-dev/infra/packages/shared/pkg/utils"
 )
 
 type closeable interface {
 	Close() error
 }
 
+// consumerStatusCheckGracePeriod is how long the graceful drain waits, after
+// in-flight builds finish, for consumers to read the final build status before
+// the server shuts down.
+const consumerStatusCheckGracePeriod = 15 * time.Second
+
 type ServerStore struct {
 	templatemanager.UnimplementedTemplateServiceServer
 
@@ -149,21 +155,26 @@ func (s *ServerStore) Close(ctx context.Context) error {
 	}
 }
 
+// Wait gracefully drains in-flight template builds during shutdown. It waits
+// for running builds to finish, bounded by ctx, then gives consumers a grace
+// period to read the final build status. It returns ctx.Err() if ctx is
+// cancelled before the drain completes.
 func (s *ServerStore) Wait(ctx context.Context) error {
-	select {
-	case <-ctx.Done():
-		return errors.New("force exit, not waiting for builds to finish")
-	default:
-		s.logger.Info(ctx, "Waiting for all build jobs to finish", zap.Int64("active_builds", s.activeBuilds.Load()))
-		s.wg.Wait()
+	s.logger.Info(ctx, "Waiting for all build jobs to finish", zap.Int64("active_builds", s.activeBuilds.Load()))
+	if err := utils.WaitGroupWait(ctx, s.wg); err != nil {
+		return fmt.Errorf("waiting for template builds: %w", err)
+	}
 
-		if !env.IsLocal() {
-			s.logger.Info(ctx, "Waiting for consumers to check build status")
-			time.Sleep(15 * time.Second)
+	if !env.IsLocal() {
+		s.logger.Info(ctx, "Waiting for consumers to check build status")
+		select {
+		case <-time.After(consumerStatusCheckGracePeriod):
+		case <-ctx.Done():
+			return ctx.Err()
 		}
+	}
 
-		s.logger.Info(ctx, "Template build queue cleaned")
+	s.logger.Info(ctx, "Template build queue cleaned")
 
-		return nil
-	}
+	return nil
 }

packages/shared/pkg/utils/waitgroup.go

@@ -0,0 +1,36 @@
+package utils
+
+import (
+	"context"
+	"fmt"
+	"sync"
+)
+
+func WaitGroupWait(ctx context.Context, wg *sync.WaitGroup) error {
+	done := make(chan struct{})
+	go func() {
+		wg.Wait()
+		close(done)
+	}()
+
+	return waitDoneOrContext(ctx, done)
+}
+
+// waitDoneOrContext blocks until done is closed or ctx is cancelled.
+//
+// When both are ready, select chooses pseudo-randomly, so a re-check biases the
+// result toward done: a closed done channel means the wait group has actually
+// finished, which is a definitive success and must win over the context error.
+func waitDoneOrContext(ctx context.Context, done <-chan struct{}) error {
+	select {
+	case <-done:
+		return nil
+	case <-ctx.Done():
+		select {
+		case <-done:
+			return nil
+		default:
+			return fmt.Errorf("waiting for wait group: %w", ctx.Err())
+		}
+	}
+}

packages/shared/pkg/utils/waitgroup_test.go

@@ -0,0 +1,89 @@
+package utils
+
+import (
+	"context"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestWaitGroupWaitAlreadyDone(t *testing.T) {
+	t.Parallel()
+
+	var wg sync.WaitGroup
+
+	require.NoError(t, WaitGroupWait(t.Context(), &wg))
+}
+
+func TestWaitGroupWaitCompletesLater(t *testing.T) {
+	t.Parallel()
+
+	var wg sync.WaitGroup
+	wg.Add(1)
+
+	release := make(chan struct{})
+	done := make(chan error, 1)
+	go func() {
+		<-release
+		wg.Done()
+	}()
+	go func() {
+		done <- WaitGroupWait(t.Context(), &wg)
+	}()
+
+	select {
+	case err := <-done:
+		require.Failf(t, "WaitGroupWait returned before wait group completed", "err: %v", err)
+	case <-time.After(10 * time.Millisecond):
+	}
+
+	close(release)
+
+	select {
+	case err := <-done:
+		require.NoError(t, err)
+	case <-time.After(time.Second):
+		t.Fatal("WaitGroupWait did not return after wait group completed")
+	}
+}
+
+func TestWaitGroupWaitReturnsContextErrorWhileWaiting(t *testing.T) {
+	t.Parallel()
+
+	var wg sync.WaitGroup
+	wg.Add(1)
+	defer wg.Done()
+
+	ctx, cancel := context.WithCancel(t.Context())
+	cancel()
+
+	require.ErrorIs(t, WaitGroupWait(ctx, &wg), context.Canceled)
+}
+
+func TestWaitDoneOrContextPrefersDoneWhenBothReady(t *testing.T) {
+	t.Parallel()
+
+	ctx, cancel := context.WithCancel(t.Context())
+	cancel()
+
+	done := make(chan struct{})
+	close(done)
+
+	// Both done and ctx are ready; select would pick a branch pseudo-randomly.
+	// A closed done channel means the wait group finished, which must win over
+	// the context error, so loop enough times to surface a regression.
+	for range 1000 {
+		require.NoError(t, waitDoneOrContext(ctx, done))
+	}
+}
+
+func TestWaitDoneOrContextReturnsContextErrorWhenNotDone(t *testing.T) {
+	t.Parallel()
+
+	ctx, cancel := context.WithCancel(t.Context())
+	cancel()
+
+	require.ErrorIs(t, waitDoneOrContext(ctx, make(chan struct{})), context.Canceled)
+}