feat(lifecycle): supervise modules and managed workers

lifecycle/lifecycle_test.go

@@ -0,0 +1,547 @@
+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("expected cancellation retained: %v", err)
+	}
+}