package lifecycle import ( "context" "errors" "fmt" "io" "log/slog" "os" "reflect" "strings" "sync" "sync/atomic" "syscall" "testing" "time" ) func newApp(t *testing.T, signals <-chan os.Signal, timeout time.Duration) *App { t.Helper() a, err := New(Config{Logger: slog.New(slog.NewTextHandler(io.Discard, nil)), Signals: signals, ShutdownTimeout: timeout}) if err != nil { t.Fatal(err) } return a } func add(t *testing.T, a *App, m Module) { t.Helper() if m.Setup == nil { m.Setup = func(context.Context) error { return nil } } if m.Teardown == nil { m.Teardown = func(context.Context) error { return nil } } if err := a.Register(m); err != nil { t.Fatal(err) } } func run(a *App, ctx context.Context) <-chan error { done := make(chan error, 1) go func() { done <- a.Run(ctx) }() return done } func receive(t *testing.T, done <-chan error) error { t.Helper() select { case err := <-done: return err case <-time.After(3 * time.Second): t.Fatal("Run blocked") return nil } } func wait(t *testing.T, ch <-chan struct{}) { t.Helper() select { case <-ch: case <-time.After(3 * time.Second): t.Fatal("callback blocked") } } func TestValidation(t *testing.T) { if _, err := New(Config{}); err == nil { t.Fatal("nil logger accepted") } logger := slog.New(slog.NewTextHandler(io.Discard, nil)) if _, err := New(Config{Logger: logger, ShutdownTimeout: -1}); err == nil { t.Fatal("negative timeout accepted") } a := newApp(t, nil, 0) if a.config.ShutdownTimeout != 30*time.Second { t.Fatal("wrong default") } if a.Register(Module{Name: "bad"}) == nil { t.Fatal("invalid module accepted") } add(t, a, Module{Name: "one"}) if a.Register(a.modules[0]) == nil { t.Fatal("duplicate accepted") } if a.Go(context.Background(), "early", func(context.Context) error { return nil }) == nil { t.Fatal("early worker accepted") } if a.Run(nil) == nil { t.Fatal("nil context accepted") } if err := a.Run(context.Background()); err != nil { t.Fatal(err) } if a.Run(context.Background()) == nil || a.Register(Module{}) == nil { t.Fatal("reuse accepted") } for _, tc := range []struct { ctx context.Context name string fn func(context.Context) error }{{nil, "x", func(context.Context) error { return nil }}, {context.Background(), "", func(context.Context) error { return nil }}, {context.Background(), "x", nil}} { if a.Go(tc.ctx, tc.name, tc.fn) == nil { t.Fatal("invalid worker accepted") } } } func TestOrderRollbackAndJoinedErrors(t *testing.T) { for _, fail := range []bool{false, true} { t.Run(map[bool]string{false: "normal", true: "rollback"}[fail], func(t *testing.T) { a := newApp(t, nil, 0) var order []string setupErr := errors.New("setup failed") first := errors.New("cleanup first") second := errors.New("cleanup second") for _, name := range []string{"one", "two", "three"} { add(t, a, Module{Name: name, Setup: func(context.Context) error { order = append(order, "setup "+name) if fail && name == "three" { return setupErr } return nil }, Teardown: func(context.Context) error { order = append(order, "stop "+name) if name == "one" { return first } if name == "two" { return second } return nil }}) } err := a.Run(context.Background()) want := []string{"setup one", "setup two", "setup three"} if !fail { want = append(want, "stop three") } want = append(want, "stop two", "stop one") if !reflect.DeepEqual(order, want) { t.Fatalf("order %v", order) } if !errors.Is(err, first) || !errors.Is(err, second) || errors.Is(err, setupErr) != fail { t.Fatalf("errors %v", err) } }) } } func TestAllWorkCompletion(t *testing.T) { for _, kind := range []string{"service", "worker"} { t.Run(kind, func(t *testing.T) { a := newApp(t, nil, 0) ready := make(chan struct{}) release := make(chan struct{}) var once sync.Once defer once.Do(func() { close(release) }) fn := func(context.Context) error { close(ready); <-release; return nil } m := Module{Name: "long"} if kind == "service" { m.Run = fn } else { m.Setup = func(ctx context.Context) error { return a.Go(ctx, "long", fn) } } add(t, a, m) add(t, a, Module{Name: "short", Run: func(context.Context) error { return nil }}) done := run(a, context.Background()) wait(t, ready) select { case err := <-done: t.Fatalf("early completion: %v", err) case <-time.After(10 * time.Millisecond): } once.Do(func() { close(release) }) if err := receive(t, done); err != nil { t.Fatal(err) } }) } } func TestWorkerFailureAndTeardownUnblocks(t *testing.T) { a := newApp(t, nil, time.Second) released := make(chan struct{}) ready := make(chan struct{}) failure := errors.New("worker failed") add(t, a, Module{Name: "resource", Setup: func(ctx context.Context) error { return a.Go(ctx, "waiting", func(ctx context.Context) error { close(ready); <-ctx.Done(); <-released; return ctx.Err() }) }, Teardown: func(ctx context.Context) error { if ctx.Err() != nil { t.Errorf("teardown canceled early: %v", ctx.Err()) } close(released) return nil }, Run: func(context.Context) error { <-ready; return failure }}) err := a.Run(context.Background()) if !errors.Is(err, failure) || errors.Is(err, context.Canceled) { t.Fatalf("errors %v", err) } b := newApp(t, nil, 0) add(t, b, Module{Name: "worker", Setup: func(ctx context.Context) error { return b.Go(ctx, "failing", func(context.Context) error { return failure }) }, Run: func(ctx context.Context) error { <-ctx.Done(); return ctx.Err() }}) if err = b.Run(context.Background()); !errors.Is(err, failure) { t.Fatalf("worker error %v", err) } } func TestSignalDuringSetup(t *testing.T) { aSignals := make(chan os.Signal, 2) a := newApp(t, aSignals, time.Second) entered := make(chan struct{}) rolled := false partialCleaned := false add(t, a, Module{Name: "ready", Teardown: func(context.Context) error { rolled = true; return nil }}) add(t, a, Module{Name: "partial", Setup: func(ctx context.Context) error { close(entered) <-ctx.Done() partialCleaned = true return errors.Join(ctx.Err(), nil) }, Teardown: func(context.Context) error { t.Error("failed setup torn down"); return nil }}) add(t, a, Module{Name: "never", Setup: func(context.Context) error { t.Error("late setup started"); return nil }}) done := run(a, context.Background()) wait(t, entered) aSignals <- syscall.SIGINT if err := receive(t, done); err != nil { t.Fatal(err) } if !rolled || !partialCleaned { t.Fatal("cleanup missing") } b := newApp(t, nil, 0) add(t, b, Module{Name: "initiator", Setup: func(context.Context) error { return context.Canceled }}) if err := b.Run(context.Background()); !errors.Is(err, context.Canceled) { t.Fatalf("initiator lost: %v", err) } } func TestMixedErrorsAndParentValues(t *testing.T) { type key struct{} parent, cancel := context.WithCancel(context.WithValue(context.Background(), key{}, "value")) defer cancel() a := newApp(t, nil, time.Second) ready := make(chan struct{}) independent := errors.New("independent") teardownErr := errors.New("teardown") add(t, a, Module{Name: "mixed", Run: func(ctx context.Context) error { close(ready) <-ctx.Done() return errors.Join(ctx.Err(), independent) }, Teardown: func(ctx context.Context) error { if ctx.Value(key{}) != "value" || ctx.Err() != nil { t.Error("shutdown context lost values or inherited cancellation") } return errors.Join(teardownErr, context.Canceled) }}) done := run(a, parent) wait(t, ready) cancel() err := receive(t, done) for _, want := range []error{context.Canceled, independent, teardownErr} { if !errors.Is(err, want) { t.Fatalf("missing %v: %v", want, err) } } if !strings.Contains(err.Error(), "mixed") { t.Fatal("missing diagnostic name") } } func TestBoundedHooksAndWorkers(t *testing.T) { for _, kind := range []string{"setup", "teardown", "worker"} { t.Run(kind, func(t *testing.T) { signals := make(chan os.Signal, 2) a := newApp(t, signals, 30*time.Millisecond) entered := make(chan struct{}) release := make(chan struct{}) exited := make(chan struct{}) defer func() { close(release); wait(t, exited) }() blocked := func(context.Context) error { close(entered); <-release; close(exited); return nil } var torn atomic.Int32 m := Module{Name: "blocked", Teardown: func(context.Context) error { torn.Add(1); return nil }} switch kind { case "setup": m.Setup = blocked case "teardown": m.Teardown = blocked case "worker": m.Setup = func(ctx context.Context) error { return a.Go(ctx, "blocked", blocked) } } add(t, a, m) done := run(a, context.Background()) wait(t, entered) start := time.Now() signals <- syscall.SIGTERM err := receive(t, done) if !errors.Is(err, context.DeadlineExceeded) { t.Fatalf("deadline missing: %v", err) } if time.Since(start) > time.Second { t.Fatal("shutdown exceeded bound") } if kind == "setup" && torn.Load() != 0 { t.Fatal("abandoned setup torn down") } }) } } func TestSecondSignalAndClosedSource(t *testing.T) { for _, trigger := range []string{"signal", "failure"} { t.Run(trigger, func(t *testing.T) { signals := make(chan os.Signal, 2) a := newApp(t, signals, time.Second) entered := make(chan struct{}) release := make(chan struct{}) exited := make(chan struct{}) defer func() { close(release); wait(t, exited) }() failure := errors.New("failed") ready := make(chan struct{}) add(t, a, Module{Name: "blocked", Run: func(ctx context.Context) error { close(ready) if trigger == "failure" { return failure } <-ctx.Done() return ctx.Err() }, Teardown: func(context.Context) error { close(entered); <-release; close(exited); return nil }}) done := run(a, context.Background()) wait(t, ready) if trigger == "signal" { signals <- syscall.SIGINT } wait(t, entered) if trigger == "failure" { signals <- syscall.SIGINT } signals <- syscall.SIGTERM err := receive(t, done) if !errors.Is(err, context.Canceled) || errors.Is(err, failure) != (trigger == "failure") { t.Fatalf("abandonment %v", err) } }) } signals := make(chan os.Signal) close(signals) a := newApp(t, signals, 0) if err := a.Run(context.Background()); err != nil { t.Fatal(err) } } func TestAdmissionRace(t *testing.T) { for range 50 { a := newApp(t, nil, time.Second) parent, cancel := context.WithCancel(context.Background()) ready := make(chan context.Context, 1) add(t, a, Module{Name: "service", Run: func(ctx context.Context) error { ready <- ctx; <-ctx.Done(); return ctx.Err() }}) done := run(a, parent) ctx := <-ready var accepted, finished atomic.Int32 var callers sync.WaitGroup for range 20 { callers.Go(func() { if err := a.Go(ctx, "racer", func(context.Context) error { finished.Add(1); return nil }); err == nil { accepted.Add(1) } }) } cancel() callers.Wait() err := receive(t, done) if !errors.Is(err, context.Canceled) { t.Fatalf("parent error %v", err) } if accepted.Load() != finished.Load() { t.Fatal("accepted worker escaped completion") } if a.Go(ctx, "late", func(context.Context) error { t.Error("late launch"); return nil }) == nil { t.Fatal("late admission") } } } func TestIndependentDescendantCancellation(t *testing.T) { for _, mixed := range []bool{false, true} { t.Run(map[bool]string{false: "cancellation-only", true: "mixed"}[mixed], func(t *testing.T) { signals := make(chan os.Signal, 2) a := newApp(t, signals, time.Second) canceled := make(chan struct{}) allowReturn := make(chan struct{}) other := errors.New("other") add(t, a, Module{Name: "owner", Setup: func(ctx context.Context) error { child, cancel := context.WithCancel(ctx) err := a.Go(child, "independent", func(ctx context.Context) error { cancel() close(canceled) <-allowReturn err := ctx.Err() if mixed { err = errors.Join(err, other) } return fmt.Errorf("worker result: %w", err) }) if err != nil { cancel() } return err }, Teardown: func(context.Context) error { close(allowReturn); return nil }}) done := run(a, context.Background()) wait(t, canceled) signals <- syscall.SIGTERM if err := receive(t, done); !errors.Is(err, context.Canceled) || errors.Is(err, other) != mixed { t.Fatalf("independent cancellation lost: %v", err) } }) } } func TestParentCancellationDuringShutdown(t *testing.T) { for _, boundary := range []string{"completion", "deadline", "second-signal"} { t.Run(boundary, func(t *testing.T) { iterations := 1 if boundary == "completion" { iterations = 1000 } for range iterations { parent, cancel := context.WithCancel(context.Background()) signals := make(chan os.Signal, 2) timeout := time.Second if boundary == "deadline" { timeout = 20 * time.Millisecond } a := newApp(t, signals, timeout) exited := make(chan struct{}) add(t, a, Module{Name: "last", Teardown: func(ctx context.Context) error { defer close(exited) cancel() switch boundary { case "deadline": <-ctx.Done() case "second-signal": signals <- syscall.SIGINT signals <- syscall.SIGTERM <-ctx.Done() } return nil }}) err := a.Run(parent) cancel() wait(t, exited) if !errors.Is(err, context.Canceled) { t.Fatalf("parent cancellation lost: %v", err) } if errors.Is(err, context.DeadlineExceeded) != (boundary == "deadline") { t.Fatalf("unexpected shutdown result: %v", err) } parentErrors := 0 for _, child := range err.(interface{ Unwrap() []error }).Unwrap() { if child == context.Canceled { parentErrors++ } } if parentErrors != 1 { t.Fatalf("parent error recorded %d times: %v", parentErrors, err) } } }) } } func TestSuccessfulInterruptedSetupRollsBack(t *testing.T) { signals := make(chan os.Signal, 2) a := newApp(t, signals, time.Second) entered := make(chan struct{}) var order []string for _, name := range []string{"first", "interrupted"} { add(t, a, Module{Name: name, Setup: func(ctx context.Context) error { if name == "interrupted" { close(entered) <-ctx.Done() } return nil }, Teardown: func(context.Context) error { order = append(order, name); return nil }}) } done := run(a, context.Background()) wait(t, entered) signals <- syscall.SIGINT if err := receive(t, done); err != nil { t.Fatal(err) } if !reflect.DeepEqual(order, []string{"interrupted", "first"}) { t.Fatalf("rollback %v", order) } } func TestSharedDeadline(t *testing.T) { a := newApp(t, nil, 80*time.Millisecond) var firstDeadline time.Time entered := make(chan struct{}) release := make(chan struct{}) exited := make(chan struct{}) defer func() { close(release); wait(t, exited) }() var skipped atomic.Bool add(t, a, Module{Name: "never", Teardown: func(context.Context) error { skipped.Store(true); return nil }}) add(t, a, Module{Name: "blocked", Teardown: func(ctx context.Context) error { deadline, ok := ctx.Deadline() if !ok || !deadline.Equal(firstDeadline) { t.Error("deadline reset between hooks") } close(entered) <-release close(exited) return nil }}) add(t, a, Module{Name: "first", Teardown: func(ctx context.Context) error { firstDeadline, _ = ctx.Deadline() time.Sleep(20 * time.Millisecond) return nil }}) done := run(a, context.Background()) wait(t, entered) if err := receive(t, done); !errors.Is(err, context.DeadlineExceeded) { t.Fatalf("deadline %v", err) } if skipped.Load() { t.Fatal("hook started after abandonment") } } func TestNormalCompletionAdmissionRace(t *testing.T) { for range 50 { a := newApp(t, nil, time.Second) ready := make(chan context.Context, 1) release := make(chan struct{}) add(t, a, Module{Name: "service", Run: func(ctx context.Context) error { ready <- ctx; <-release; return nil }}) done := run(a, context.Background()) ctx := <-ready var accepted, finished atomic.Int32 var callers sync.WaitGroup for range 20 { callers.Go(func() { if a.Go(ctx, "racer", func(context.Context) error { finished.Add(1); return nil }) == nil { accepted.Add(1) } }) } close(release) callers.Wait() if err := receive(t, done); err != nil { t.Fatal(err) } if accepted.Load() != finished.Load() { t.Fatal("normal completion lost worker") } } } func TestStartupWaitsForSetupDespiteEmptyWorkers(t *testing.T) { a := newApp(t, nil, time.Second) entered := make(chan struct{}) release := make(chan struct{}) defer func() { select { case <-release: default: close(release) } }() add(t, a, Module{Name: "setup", Setup: func(context.Context) error { close(entered); <-release; return nil }}) done := run(a, context.Background()) wait(t, entered) if a.Run(context.Background()) == nil { t.Fatal("concurrent Run accepted") } if a.Register(Module{}) == nil { t.Fatal("registration remained open") } select { case err := <-done: t.Fatalf("ended during setup: %v", err) default: } close(release) if err := receive(t, done); err != nil { t.Fatal(err) } } func TestMixedShutdownErrorPreservesCause(t *testing.T) { signals := make(chan os.Signal, 2) a := newApp(t, signals, time.Second) ready := make(chan struct{}) independent := &os.PathError{Op: "read", Path: "queue", Err: io.ErrUnexpectedEOF} add(t, a, Module{Name: "mixed", Run: func(ctx context.Context) error { close(ready) <-ctx.Done() return fmt.Errorf("outer: %w", errors.Join(fmt.Errorf("cancellation: %w", ctx.Err()), independent)) }}) done := run(a, context.Background()) wait(t, ready) signals <- syscall.SIGTERM err := receive(t, done) var pathErr *os.PathError if !errors.As(err, &pathErr) || pathErr != independent || !errors.Is(err, io.ErrUnexpectedEOF) { t.Fatalf("independent cause lost: %v", err) } if !errors.Is(err, context.Canceled) { t.Fatalf("mixed wrapper was reconstructed: %v", err) } } // matchingError exercises wrapper-owned matching independently of its child. type matchingError struct { child error match error } func (e *matchingError) Error() string { return "custom: " + e.child.Error() } func (e *matchingError) Unwrap() error { return e.child } func (e *matchingError) Is(target error) bool { return target == e.match } func TestMixedShutdownErrorPreservesWrapper(t *testing.T) { for _, kind := range []string{"path", "custom"} { t.Run(kind, func(t *testing.T) { signals := make(chan os.Signal, 2) a := newApp(t, signals, time.Second) ready := make(chan struct{}) match := errors.New("custom match") var original error add(t, a, Module{Name: "mixed", Run: func(ctx context.Context) error { close(ready) <-ctx.Done() joined := errors.Join(ctx.Err(), io.ErrUnexpectedEOF) if kind == "path" { original = &os.PathError{Op: "read", Path: "queue", Err: joined} } else { original = &matchingError{child: joined, match: match} } return original }}) done := run(a, context.Background()) wait(t, ready) signals <- syscall.SIGTERM err := receive(t, done) if !errors.Is(err, original) || !errors.Is(err, io.ErrUnexpectedEOF) { t.Fatalf("wrapper identity or inner failure lost: %v", err) } if kind == "path" { var pathErr *os.PathError if !errors.As(err, &pathErr) || pathErr != original { t.Fatalf("typed wrapper lost: %v", err) } } else { var custom *matchingError if !errors.As(err, &custom) || custom != original || !errors.Is(err, match) { t.Fatalf("custom matching lost: %v", err) } } }) } } func TestCancellationOnly(t *testing.T) { for _, err := range []error{ context.Canceled, fmt.Errorf("wrapped: %w", context.Canceled), &os.PathError{Op: "read", Path: "queue", Err: errors.Join(context.Canceled, fmt.Errorf("wrapped: %w", context.Canceled))}, } { if !cancellationOnly(err, context.Canceled) { t.Errorf("expected cancellation not recognized: %v", err) } } for _, err := range []error{nil, io.ErrUnexpectedEOF, context.DeadlineExceeded, errors.Join(context.Canceled, io.ErrUnexpectedEOF)} { if cancellationOnly(err, context.Canceled) { t.Errorf("independent error classified as cancellation: %v", err) } } }