go/lifecycle/lifecycle.go

// Package lifecycle supervises ordered module setup, concurrent services and
// explicitly admitted workers within one shutdown budget. It owns cancellation
// and reverse teardown, not signal subscriptions, process exit or raw goroutines.
// Callbacks and the injected logger must cooperate: abandonment bounds Run's
// return, but cannot kill callbacks or reclaim their resources.
package lifecycle

import (
	"context"
	"errors"
	"fmt"
	"log/slog"
	"os"
	"sync"
	"time"
)

// Module describes one resource owner. Setup acquires resources; Teardown releases
// them after successful setup. A failing Setup owns its partial cleanup. Optional
// Run starts only after all setups succeed. Callbacks must honor cancellation.
type Module struct {
	// Name identifies the module in errors.
	Name string
	// Setup acquires resources before any service starts.
	Setup func(context.Context) error
	// Teardown releases resources after successful Setup.
	Teardown func(context.Context) error
	// Run optionally serves until completion or cancellation.
	Run func(context.Context) error
}

// Config supplies application policy. Logger is required and must not block.
// Zero ShutdownTimeout selects 30 seconds; negative values are invalid.
// The caller owns Signals, its subscription and cleanup, filtering to SIGINT and
// SIGTERM and buffering at least two signals. Nil or closed disables reception.
type Config struct {
	// Logger receives lifecycle diagnostics and must not block.
	Logger *slog.Logger
	// Signals is a caller-owned, filtered signal source.
	Signals <-chan os.Signal
	// ShutdownTimeout bounds all shutdown work together.
	ShutdownTimeout time.Duration
}

// App is a single-use coordinator constructed by New. Do not copy it. Register
// modules before Run; Go may be called concurrently during setup and running.
// App stores no context. After abandonment, callers own process exit and must not
// reuse resources whose callbacks have not returned.
type App struct {
	config        Config
	mu            sync.Mutex
	modules       []Module
	used          bool
	accepting     bool
	stopping      bool
	active        int
	cancel        context.CancelCauseFunc
	shutdownCause error
	stoppedAt     time.Time
	errs          []error
	wake          chan struct{}
}

// New validates and captures config without subscribing to signals or starting work.
func New(config Config) (*App, error) {
	if config.Logger == nil {
		return nil, errors.New("lifecycle: logger is required")
	}
	if config.ShutdownTimeout < 0 {
		return nil, errors.New("lifecycle: negative shutdown timeout")
	}
	if config.ShutdownTimeout == 0 {
		config.ShutdownTimeout = 30 * time.Second
	}
	return &App{config: config, wake: make(chan struct{}, 1)}, nil
}

// Register captures a uniquely named module with required setup and teardown.
// Registration closes permanently when Run begins, including a failed run.
func (a *App) Register(module Module) error {
	a.mu.Lock()
	defer a.mu.Unlock()
	if a.used {
		return errors.New("lifecycle: registration is closed")
	}
	if module.Name == "" || module.Setup == nil || module.Teardown == nil {
		return errors.New("lifecycle: module needs name, setup and teardown")
	}
	for _, m := range a.modules {
		if m.Name == module.Name {
			return fmt.Errorf("lifecycle: duplicate module %q", module.Name)
		}
	}
	a.modules = append(a.modules, module)
	return nil
}

// Go immediately launches and tracks fn, or returns an admission error without
// launching it. ctx must be a callback's lifecycle context or a descendant;
// ancestry is a caller precondition and cannot generally be checked. Detached
// contexts forfeit graceful cancellation. Names are diagnostic and may repeat.
// Callers must handle rejection; raw goroutines are outside App's guarantee.
func (a *App) Go(ctx context.Context, name string, fn func(context.Context) error) error {
	if ctx == nil || name == "" || fn == nil {
		return errors.New("lifecycle: worker needs context, name and function")
	}
	a.mu.Lock()
	defer a.mu.Unlock()
	if !a.accepting {
		return errors.New("lifecycle: worker admission is closed")
	}
	a.launchLocked(ctx, "worker", name, fn)
	return nil
}

func (a *App) notifyLocked() {
	select {
	case a.wake <- struct{}{}:
	default:
	}
}

func (a *App) stopLocked() {
	if !a.stopping {
		// invariant: closure and admission share the lock, so every accepted worker is counted.
		a.accepting = false
		a.stopping = true
		a.stoppedAt = time.Now()
		a.cancel(a.shutdownCause)
	}
	a.notifyLocked()
}

// cancellationOnly recognizes expected cancellation without rebuilding error trees.
func cancellationOnly(err, cancellation error) bool {
	if err == nil {
		return false
	}
	if err == cancellation {
		return true
	}
	if many, ok := err.(interface{ Unwrap() []error }); ok {
		children := many.Unwrap()
		for _, child := range children {
			if !cancellationOnly(child, cancellation) {
				return false
			}
		}
		return len(children) > 0
	}
	if one, ok := err.(interface{ Unwrap() error }); ok {
		return cancellationOnly(one.Unwrap(), cancellation)
	}
	return false
}

func (a *App) recordLocked(ctx context.Context, phase, name string, err error) {
	if a.stopping && context.Cause(ctx) == a.shutdownCause && cancellationOnly(err, ctx.Err()) {
		err = nil
	}
	if err != nil {
		a.errs = append(a.errs, fmt.Errorf("%s %q: %w", phase, name, err))
		a.stopLocked()
	}
}

func (a *App) launchLocked(ctx context.Context, phase, name string, fn func(context.Context) error) {
	a.active++
	go func() {
		err := fn(ctx)
		a.mu.Lock()
		defer a.mu.Unlock()
		a.recordLocked(ctx, phase, name, err)
		a.active--
		a.notifyLocked()
	}()
}

// Run sets up modules in registration order, then runs services concurrently.
// Every service and admitted worker must finish for normal completion. The first
// failure, parent cancellation or signal cancels their shared context and closes
// admission before reverse teardown. A fresh context retains parent values and
// gives teardown and callback completion one deadline from the shutdown trigger.
// In-flight setup is awaited within that budget before rollback; abandoned setup
// is never torn down later. Deadline or second signal abandons outstanding work.
// Observed failures are joined; expected callback cancellation during established
// shutdown is omitted only for cancellation-only error trees. Mixed errors remain
// intact to preserve wrapper types, identity and matching, so may still match
// cancellation. Initiating failures and every teardown error remain.
// Cancellation is omitted only when context.Cause proves App caused it; parent
// propagation and independently canceled descendants are retained conservatively.
// Run never recovers panics, closes Signals or exits the process.
func (a *App) Run(parent context.Context) error {
	if parent == nil {
		return errors.New("lifecycle: nil run context")
	}
	a.mu.Lock()
	if a.used {
		a.mu.Unlock()
		return errors.New("lifecycle: app already run")
	}
	ctx, cancel := context.WithCancelCause(parent)
	a.cancel = cancel
	a.shutdownCause = errors.New("lifecycle shutdown")
	a.used = true
	a.accepting = true
	modules := a.modules
	a.mu.Unlock()
	defer cancel(nil)
	signals := a.config.Signals
	signalCount := 0
	var shutdown context.Context
	var shutdownCancel context.CancelFunc
	var setupDone <-chan error
	completed := 0
	var teardownDone <-chan error
	teardownIndex := -1
	startup := true
	parentDone := parent.Done()
	receiveSignal := func(sig os.Signal) error {
		signalCount++
		if signalCount == 2 {
			if shutdownCancel != nil {
				shutdownCancel()
			}
			return fmt.Errorf("second signal %v: %w", sig, context.Canceled)
		}
		a.mu.Lock()
		a.stopLocked()
		a.mu.Unlock()
		return nil
	}

	for {
		// invariant: queued abandonment is observed before starting another hook.
		select {
		case sig, ok := <-signals:
			if !ok {
				signals = nil
			} else if err := receiveSignal(sig); err != nil {
				return a.result(err)
			}
			continue
		default:
		}

		a.mu.Lock()
		if !a.stopping && parent.Err() != nil {
			a.errs = append(a.errs, parent.Err())
			a.stopLocked()
			parentDone = nil
		}
		if !a.stopping && startup && setupDone == nil {
			if completed < len(modules) {
				m := modules[completed]
				// invariant: buffered hook results cannot strand a sender after abandonment.
				result := make(chan error, 1)
				setupDone = result
				go func() {
					err := m.Setup(ctx)
					a.mu.Lock()
					a.recordLocked(ctx, "setup", m.Name, err)
					a.mu.Unlock()
					result <- err
				}()
			} else {
				startup = false
				for _, m := range modules {
					if m.Run != nil {
						a.launchLocked(ctx, "run", m.Name, m.Run)
					}
				}
			}
		}
		if !a.stopping && !startup && a.active == 0 {
			a.stopLocked()
		}
		if a.stopping && shutdown == nil {
			created, release := context.WithDeadline(context.WithoutCancel(parent), a.stoppedAt.Add(a.config.ShutdownTimeout))
			shutdown, shutdownCancel = created, release
			defer release()
			teardownIndex = completed - 1
			a.config.Logger.Debug("lifecycle shutdown started")
		}
		stopping := a.stopping
		active := a.active
		a.mu.Unlock()

		if stopping {
			if err := shutdown.Err(); err != nil {
				return a.result(err)
			}
			if setupDone == nil && teardownDone == nil {
				if teardownIndex >= 0 {
					m := modules[teardownIndex]
					result := make(chan error, 1)
					teardownDone = result
					go func() { result <- m.Teardown(shutdown) }()
				} else if active == 0 {
					return a.result(nil)
				}
			}
		}
		var deadline <-chan struct{}
		if shutdown != nil {
			deadline = shutdown.Done()
		}
		select {
		case err := <-setupDone:
			setupDone = nil
			if err == nil {
				completed++
				if stopping {
					teardownIndex = completed - 1
				}
			}
		case err := <-teardownDone:
			teardownDone = nil
			if err != nil {
				a.mu.Lock()
				a.errs = append(a.errs, fmt.Errorf("teardown %q: %w", modules[teardownIndex].Name, err))
				a.mu.Unlock()
			}
			teardownIndex--
		case <-a.wake:
		case <-parentDone:
			a.mu.Lock()
			a.errs = append(a.errs, parent.Err())
			a.stopLocked()
			a.mu.Unlock()
			parentDone = nil
		case sig, ok := <-signals:
			if !ok {
				signals = nil
				continue
			}
			if err := receiveSignal(sig); err != nil {
				return a.result(err)
			}
		case <-deadline:
			return a.result(context.DeadlineExceeded)
		}
	}
}

func (a *App) result(extra error) error {
	a.mu.Lock()
	defer a.mu.Unlock()
	return errors.Join(append(a.errs, extra)...)
}