go/lifecycle/lifecycle_test.go

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) {
	signals := make(chan os.Signal, 2)
	a := newApp(t, signals, time.Second)
	canceled := make(chan struct{})
	allowReturn := make(chan struct{})
	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
			return fmt.Errorf("worker result: %w", errors.Join(ctx.Err(), errors.New("other")))
		})
		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) || !strings.Contains(err.Error(), "other") {
		t.Fatalf("independent cancellation lost: %v", 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)
		}
	}
}