go-omni/lifecycle/phases.go

package lifecycle

import (
	"context"
	"errors"
	"fmt"
	"strings"
	"time"
	"unicode"

	"git.yuntuo.hk/k2safe/go-omni/internal/configdiag"
)

// 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 {
	for _, f := range []struct {
		name  string
		value time.Duration
	}{{"stopTimeout", c.StopTimeout}, {"drainTimeout", c.DrainTimeout}, {"cancelTimeout", c.CancelTimeout}} {
		if f.value < 0 {
			return configdiag.New(f.name, "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
}