go-omni/lifecycle/runtime.go

package lifecycle

import (
	"context"
	"errors"
	"fmt"
	"sync"
)

// Runtime owns resources registered during dependency initialization. A server
// must finish draining before calling Close. Embed *Runtime in the project's svc.
// The logger and any other borrowed resources must not be registered here.
type Runtime struct {
	Lifecycle     *Controller
	mu            sync.Mutex
	cleanups      []func() error
	closed        bool
	once          sync.Once
	closeErr      error
	stops, drains []shutdownHook
	shutdown      ShutdownConfig
	work          context.Context
	cancelWork    context.CancelFunc
}

func NewRuntime(ctx context.Context) *Runtime {
	r, _ := NewRuntimeWithConfig(ctx, DefaultShutdownConfig())
	return r
}

// NewRuntimeWithConfig separates the stop notification from the lifetime of
// accepted background work. Zero timeout fields use documented defaults.
func NewRuntimeWithConfig(ctx context.Context, cfg ShutdownConfig) (*Runtime, error) {
	if err := cfg.Validate(); err != nil {
		return nil, err
	}
	if ctx == nil {
		ctx = context.Background()
	}
	work, cancel := context.WithCancel(context.WithoutCancel(ctx))
	return &Runtime{Lifecycle: New(ctx), shutdown: cfg.defaults(), work: work, cancelWork: cancel}, nil
}
func (r *Runtime) Context() context.Context  { return r.Lifecycle.Context() }
func (r *Runtime) RequestShutdown(err error) { r.Lifecycle.RequestShutdown(err) }

// WorkContext stays alive during HTTP draining and graceful background draining.
// It is canceled on a shutdown failure/deadline, or after successful draining.
// Context() remains the early stop notification; do not use it for accepted work.
func (r *Runtime) WorkContext() context.Context { return r.work }

// OnClose transfers cleanup ownership. Register immediately after creation,
// before exposing resources to other goroutines. Registration after Close is a
// programming error. Callbacks may request shutdown, but must not call Close.
func (r *Runtime) OnClose(cleanup func() error) {
	if cleanup == nil {
		return
	}
	r.mu.Lock()
	defer r.mu.Unlock()
	if r.closed {
		panic("lifecycle: register cleanup after Close")
	}
	r.cleanups = append(r.cleanups, cleanup)
}

// Close runs stop and drain barriers before releasing resources in reverse
// order, including after failed initialization. If quiescence is unproven it
// skips resource callbacks and returns ErrResourcesRetained. Once release is
// safe, all OnClose callbacks are attempted even if one fails or panics.
func (r *Runtime) Close() error {
	return r.close(false)
}

func (r *Runtime) close(retain bool) error {
	if r == nil {
		return nil
	}
	r.once.Do(func() {
		r.RequestShutdown(nil)
		r.mu.Lock()
		r.closed = true
		cleanups, stops, drains := r.cleanups, r.stops, r.drains
		r.cleanups = nil
		r.stops, r.drains = nil, nil
		r.mu.Unlock()
		stopErr := r.runPhase("stop", stops, r.shutdown.StopTimeout)
		var stopped *phaseError
		errors.As(stopErr, &stopped)
		for _, stop := range stops {
			found := false
			for _, drain := range drains {
				if drain.name == stop.name {
					found = true
					break
				}
			}
			if !found {
				stopErr = &phaseError{err: errors.Join(stopErr, &HookError{"drain", stop.name, errors.New("matching OnDrain required")}), unsafe: true}
			}
		}
		if stopErr != nil {
			r.cancelWork()
		}
		// A timed-out Stop can still be executing. Never overlap its own Drain;
		// other components can still establish their completion independently.
		eligible := make([]shutdownHook, 0, len(drains))
		for _, drain := range drains {
			if stopped == nil || !stopped.unjoined[drain.name] {
				eligible = append(eligible, drain)
			}
		}
		drainErr := r.runPhase("drain", eligible, r.shutdown.DrainTimeout)
		r.cancelWork()
		r.closeErr = errors.Join(stopErr, drainErr)
		if retain || unsafePhase(stopErr) || unsafePhase(drainErr) {
			r.closeErr = errors.Join(r.closeErr, ErrResourcesRetained)
			return
		}
		for i := len(cleanups) - 1; i >= 0; i-- {
			r.closeErr = errors.Join(r.closeErr, runCleanup(cleanups[i]))
		}
	})
	return r.closeErr
}

func runCleanup(fn func() error) (err error) {
	defer func() {
		if p := recover(); p != nil {
			err = fmt.Errorf("lifecycle: cleanup panic: %v", p)
		}
	}()
	return fn()
}

// Finish closes resources and preserves the operation, cleanup and shutdown
// errors. An incoming ErrResourcesRetained (for example, unfinished HTTP work)
// still runs stop/drain hooks, but permanently skips resource cleanup, including
// subsequent Close calls. Read the final cause after stopping the monitors.
func (r *Runtime) Finish(err error) error {
	if r == nil {
		return err
	}
	err = errors.Join(err, r.close(errors.Is(err, ErrResourcesRetained)))
	if cause := r.Lifecycle.Err(); cause != nil && !errors.Is(err, cause) {
		err = errors.Join(err, cause)
	}
	return err
}