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 { 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 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. Read the final cause only after monitors have been stopped by Close. func (r *Runtime) Finish(err error) error { if r == nil { return err } err = errors.Join(err, r.Close()) if cause := r.Lifecycle.Err(); cause != nil && !errors.Is(err, cause) { err = errors.Join(err, cause) } return err }