feat: add generic phased background shutdown to managed services

lifecycle/phases.go

@@ -0,0 +1,184 @@
+package lifecycle
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"strings"
+	"time"
+	"unicode"
+)
+
+// ErrResourcesRetained means background quiescence could not be established.
+// OnClose callbacks were deliberately skipped. Close is not retried; an owning
+// process should exit unsuccessfully instead of continuing with a stopped Runtime.
+var ErrResourcesRetained = errors.New("lifecycle: background shutdown incomplete; shared resources retained")
+
+type ShutdownConfig struct {
+	StopTimeout   time.Duration `json:"stopTimeout"`
+	DrainTimeout  time.Duration `json:"drainTimeout"`
+	CancelTimeout time.Duration `json:"cancelTimeout"`
+}
+
+func DefaultShutdownConfig() ShutdownConfig {
+	return ShutdownConfig{StopTimeout: 5 * time.Second, DrainTimeout: 10 * time.Second, CancelTimeout: 5 * time.Second}
+}
+func (c ShutdownConfig) Validate() error {
+	if c.StopTimeout < 0 || c.DrainTimeout < 0 || c.CancelTimeout < 0 {
+		return errors.New("lifecycle: shutdown timeouts must be nonnegative")
+	}
+	return nil
+}
+func (c ShutdownConfig) defaults() ShutdownConfig {
+	d := DefaultShutdownConfig()
+	if c.StopTimeout == 0 {
+		c.StopTimeout = d.StopTimeout
+	}
+	if c.DrainTimeout == 0 {
+		c.DrainTimeout = d.DrainTimeout
+	}
+	if c.CancelTimeout == 0 {
+		c.CancelTimeout = d.CancelTimeout
+	}
+	return c
+}
+
+// HookError identifies the failed phase/component and preserves the original error.
+type HookError struct {
+	Phase, Name string
+	Err         error
+}
+
+func (e *HookError) Error() string { return fmt.Sprintf("lifecycle %s %q: %v", e.Phase, e.Name, e.Err) }
+func (e *HookError) Unwrap() error { return e.Err }
+
+// phaseError keeps diagnostics separate from proof that every hook succeeded.
+// A deadline is still reported when all callbacks finish during cancellation grace.
+type phaseError struct {
+	err      error
+	unsafe   bool
+	unjoined map[string]bool
+}
+
+func (e *phaseError) Error() string { return e.err.Error() }
+func (e *phaseError) Unwrap() error { return e.err }
+func unsafePhase(err error) bool    { var p *phaseError; return errors.As(err, &p) && p.unsafe }
+
+type shutdownHook struct {
+	name string
+	fn   func(context.Context) error
+}
+
+// OnStop registers a prompt, idempotent stop-admission callback. It must not
+// close resources used by accepted work. All stop hooks precede all drain hooks.
+// Hooks within a phase run concurrently: combine dependent components in one hook.
+// Every OnStop requires a matching OnDrain with the same name.
+func (r *Runtime) OnStop(name string, fn func(context.Context) error) { r.addHook(&r.stops, name, fn) }
+
+// OnDrain registers proof of quiescence. Return nil only after the component can
+// no longer use shared resources, including final acknowledgements/cleanup.
+// Deadline cancellation asks the component to stop; returning ctx.Err alone does
+// not prove completion. A callback may finish during the cancellation grace.
+func (r *Runtime) OnDrain(name string, fn func(context.Context) error) {
+	r.addHook(&r.drains, name, fn)
+}
+
+func (r *Runtime) addHook(list *[]shutdownHook, name string, fn func(context.Context) error) {
+	if fn == nil || name == "" || len(name) > 128 || strings.TrimSpace(name) != name || strings.IndexFunc(name, unicode.IsControl) >= 0 {
+		panic("lifecycle: valid hook name and function required")
+	}
+	r.mu.Lock()
+	defer r.mu.Unlock()
+	if r.closed {
+		panic("lifecycle: register hook after Close")
+	}
+	for _, h := range *list {
+		if h.name == name {
+			panic("lifecycle: duplicate phase hook name")
+		}
+	}
+	*list = append(*list, shutdownHook{name, fn})
+}
+
+func callHook(ctx context.Context, h shutdownHook) (err error) {
+	defer func() {
+		if p := recover(); p != nil {
+			if cause, ok := p.(error); ok {
+				err = fmt.Errorf("hook panic: %w", cause)
+			} else {
+				err = fmt.Errorf("hook panic: %v", p)
+			}
+		}
+	}()
+	return h.fn(ctx)
+}
+
+func (r *Runtime) runPhase(phase string, hooks []shutdownHook, timeout time.Duration) error {
+	if len(hooks) == 0 {
+		return nil
+	}
+	ctx, cancel := context.WithTimeout(context.WithoutCancel(r.work), timeout)
+	defer cancel()
+	type result struct {
+		index int
+		err   error
+	}
+	results := make(chan result, len(hooks))
+	pending := make([]bool, len(hooks))
+	for i, h := range hooks {
+		pending[i] = true
+		go func(i int, h shutdownHook) { results <- result{i, callHook(ctx, h)} }(i, h)
+	}
+	remaining := len(hooks)
+	var errs error
+	unsafe := false
+	accept := func(v result) {
+		pending[v.index] = false
+		remaining--
+		if v.err != nil {
+			unsafe = true
+			errs = errors.Join(errs, &HookError{phase, hooks[v.index].name, v.err})
+			r.cancelWork()
+		}
+	}
+	for remaining > 0 {
+		select {
+		case v := <-results:
+			accept(v)
+		case <-ctx.Done():
+			r.cancelWork()
+			errs = errors.Join(errs, &HookError{phase, "deadline", ctx.Err()})
+			grace := time.NewTimer(r.shutdown.CancelTimeout)
+			defer grace.Stop()
+			for remaining > 0 {
+				select {
+				case v := <-results:
+					accept(v)
+				case <-grace.C:
+					// Collect callbacks already returned at the deadline before marking pending ones.
+					for {
+						select {
+						case v := <-results:
+							accept(v)
+						default:
+							goto collected
+						}
+					}
+				collected:
+					unjoined := make(map[string]bool)
+					for i, active := range pending {
+						if active {
+							unjoined[hooks[i].name] = true
+							errs = errors.Join(errs, &HookError{phase, hooks[i].name, context.DeadlineExceeded})
+						}
+					}
+					return &phaseError{err: errs, unsafe: unsafe || remaining > 0, unjoined: unjoined}
+				}
+			}
+		}
+	}
+	if errs != nil {
+		return &phaseError{err: errs, unsafe: unsafe}
+	}
+	return nil
+}