diff --git a/main.go b/main.go index 09e226b..8c45cdc 100644 --- a/main.go +++ b/main.go @@ -47,7 +47,20 @@ func shutdownAll(ctx context.Context, cancelRoot context.CancelFunc, steps ...sh defer cancelRoot() var shutdownErr error for _, step := range steps { - if err := step.run(ctx); err != nil { + done := make(chan error, 1) + started := make(chan struct{}) + go func() { + close(started) + done <- step.run(ctx) + }() + <-started + var err error + select { + case err = <-done: + case <-ctx.Done(): + err = ctx.Err() + } + if err != nil { log.Printf("[main] %s shutdown: %v", step.name, err) shutdownErr = errors.Join(shutdownErr, fmt.Errorf("%s shutdown: %w", step.name, err)) } diff --git a/main_test.go b/main_test.go index 5902be7..ee483e1 100644 --- a/main_test.go +++ b/main_test.go @@ -70,6 +70,48 @@ func TestShutdownAllContinuesAfterTimeoutWithSharedDeadline(t *testing.T) { } } +func TestShutdownAllContinuesWhenStepIgnoresDeadline(t *testing.T) { + ctx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond) + defer cancel() + + blocked := make(chan struct{}) + defer close(blocked) + cleanupStarted := make(chan struct{}) + rootCtx, rootCancel := context.WithCancel(context.Background()) + result := make(chan error, 1) + go func() { + result <- shutdownAll(ctx, rootCancel, + shutdownStep{name: "blocked", run: func(context.Context) error { + <-blocked + return nil + }}, + shutdownStep{name: "cleanup", run: func(context.Context) error { + close(cleanupStarted) + return nil + }}, + ) + }() + + select { + case err := <-result: + if !errors.Is(err, context.DeadlineExceeded) { + t.Fatalf("shutdownAll() error = %v, want deadline exceeded", err) + } + case <-time.After(250 * time.Millisecond): + t.Fatal("shutdownAll did not enforce the shared deadline") + } + select { + case <-cleanupStarted: + default: + t.Fatal("cleanup after blocked step was not started") + } + select { + case <-rootCtx.Done(): + default: + t.Fatal("root context was not canceled") + } +} + func recordShutdown(order *[]string, contexts *[]context.Context, name string, err error) func(context.Context) error { return func(ctx context.Context) error { *order = append(*order, name)