diff --git a/packages/orchestrator/pkg/template/server/drain_test.go b/packages/orchestrator/pkg/template/server/drain_test.go new file mode 100644 index 0000000000..1fbbac1b66 --- /dev/null +++ b/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) +} diff --git a/packages/orchestrator/pkg/template/server/main.go b/packages/orchestrator/pkg/template/server/main.go index 2013af14fa..2b1c9fa40e 100644 --- a/packages/orchestrator/pkg/template/server/main.go +++ b/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 } diff --git a/packages/shared/pkg/utils/waitgroup.go b/packages/shared/pkg/utils/waitgroup.go new file mode 100644 index 0000000000..1b99d857e8 --- /dev/null +++ b/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()) + } + } +} diff --git a/packages/shared/pkg/utils/waitgroup_test.go b/packages/shared/pkg/utils/waitgroup_test.go new file mode 100644 index 0000000000..99f2869c99 --- /dev/null +++ b/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) +}