feat: add generic phased background shutdown to managed services

2026-09-25 00:35:59 +08:00
parent adcb6f60b6
commit 16deead70e
37 changed files with 1330 additions and 68 deletions
--- a/lifecycle/phases_test.go
+++ b/lifecycle/phases_test.go
@@ -0,0 +1,234 @@
+package lifecycle
+
+import (
+	"context"
+	"errors"
+	"reflect"
+	"sync"
+	"sync/atomic"
+	"testing"
+	"time"
+)
+
+func testRuntime(t *testing.T) *Runtime {
+	t.Helper()
+	r, err := NewRuntimeWithConfig(context.Background(), ShutdownConfig{StopTimeout: time.Second, DrainTimeout: time.Second, CancelTimeout: time.Second})
+	if err != nil {
+		t.Fatal(err)
+	}
+	return r
+}
+
+func TestPhasesBarrierAndWorkLifetime(t *testing.T) {
+	parent, cancel := context.WithCancel(context.WithValue(context.Background(), struct{}{}, "value"))
+	defer cancel()
+	r := NewRuntime(parent)
+	var stops, drains atomic.Int32
+	var mu sync.Mutex
+	var closed []int
+	started := make(chan struct{}, 2)
+	release := make(chan struct{})
+	// Deliberately interleave registration. Phases take precedence over order.
+	r.OnClose(func() error { mu.Lock(); defer mu.Unlock(); closed = append(closed, 1); return nil })
+	for _, name := range []string{"producer", "consumer"} {
+		r.OnStop(name, func(ctx context.Context) error {
+			if r.WorkContext().Err() != nil || ctx.Err() != nil {
+				return errors.New("work canceled before drain")
+			}
+			stops.Add(1)
+			started <- struct{}{}
+			<-release
+			return nil
+		})
+		r.OnDrain(name, func(ctx context.Context) error {
+			if stops.Load() != 2 || r.WorkContext().Err() != nil || ctx.Err() != nil {
+				return errors.New("stage barrier failed")
+			}
+			drains.Add(1)
+			return nil
+		})
+	}
+	r.OnClose(func() error {
+		if drains.Load() != 2 {
+			return errors.New("dependency closed early")
+		}
+		mu.Lock()
+		defer mu.Unlock()
+		closed = append(closed, 2)
+		return nil
+	})
+	cancel()
+	if r.WorkContext().Err() != nil || r.WorkContext().Value(struct{}{}) != "value" {
+		t.Fatal("work context inherited cancellation or lost values")
+	}
+	done := make(chan error, 1)
+	go func() { done <- r.Finish(nil) }()
+	for i := 0; i < 2; i++ {
+		select {
+		case <-started:
+		case <-time.After(time.Second):
+			t.Fatal("stop callbacks did not run concurrently")
+		}
+	}
+	if drains.Load() != 0 {
+		t.Fatal("drain ran before stop barrier")
+	}
+	close(release)
+	if err := <-done; err != nil {
+		t.Fatal(err)
+	}
+	if !reflect.DeepEqual(closed, []int{2, 1}) || r.WorkContext().Err() == nil {
+		t.Fatal(closed)
+	}
+}
+
+func TestPhasesFailuresRetainDependenciesAndCauses(t *testing.T) {
+	for _, phase := range []string{"stop", "drain"} {
+		for _, panics := range []bool{false, true} {
+			t.Run(phase+map[bool]string{false: "Error", true: "Panic"}[panics], func(t *testing.T) {
+				r := testRuntime(t)
+				cause := errors.New("component cause")
+				fatal := errors.New("fatal monitor")
+				cleanup := atomic.Int32{}
+				other := atomic.Int32{}
+				r.OnClose(func() error { cleanup.Add(1); return nil })
+				hook := func(context.Context) error {
+					r.RequestShutdown(fatal)
+					if panics {
+						panic(cause)
+					}
+					return cause
+				}
+				if phase == "stop" {
+					r.OnStop("broken", hook)
+				} else {
+					r.OnDrain("broken", hook)
+				}
+				r.OnDrain("other", func(context.Context) error { other.Add(1); return nil })
+				err := r.Finish(nil)
+				var detail *HookError
+				if !errors.Is(err, ErrResourcesRetained) || !errors.Is(err, cause) || !errors.Is(err, fatal) || !errors.As(err, &detail) || detail.Phase != phase {
+					t.Fatal(err)
+				}
+				if cleanup.Load() != 0 || other.Load() != 1 || r.WorkContext().Err() == nil {
+					t.Fatal("unsafe cleanup or omitted drain")
+				}
+				if !errors.Is(r.Close(), ErrResourcesRetained) || other.Load() != 1 {
+					t.Fatal("close retried hooks")
+				}
+			})
+		}
+	}
+}
+
+func TestPhasesDeadlineCancelsWorkAndJoinsBeforeClose(t *testing.T) {
+	r, err := NewRuntimeWithConfig(context.Background(), ShutdownConfig{DrainTimeout: 20 * time.Millisecond, CancelTimeout: time.Second})
+	if err != nil {
+		t.Fatal(err)
+	}
+	var joined atomic.Bool
+	var closed atomic.Bool
+	r.OnDrain("worker", func(ctx context.Context) error {
+		<-ctx.Done()
+		<-r.WorkContext().Done()
+		// Cleanup can still use shared resources after forced cancellation.
+		if closed.Load() {
+			return errors.New("dependency closed during cancellation")
+		}
+		joined.Store(true)
+		return nil
+	})
+	r.OnClose(func() error {
+		if !joined.Load() {
+			t.Error("worker not joined")
+		}
+		closed.Store(true)
+		return nil
+	})
+	err = r.Close()
+	if !errors.Is(err, context.DeadlineExceeded) || errors.Is(err, ErrResourcesRetained) || !closed.Load() {
+		t.Fatal(err)
+	}
+}
+
+func TestPhasesUncooperativeHookIsBoundedAndNeverClosesDependencies(t *testing.T) {
+	for _, phase := range []string{"stop", "drain"} {
+		t.Run(phase, func(t *testing.T) {
+			r, err := NewRuntimeWithConfig(context.Background(), ShutdownConfig{StopTimeout: 20 * time.Millisecond, DrainTimeout: 20 * time.Millisecond, CancelTimeout: 20 * time.Millisecond})
+			if err != nil {
+				t.Fatal(err)
+			}
+			release := make(chan struct{})
+			finished := make(chan struct{})
+			var closed atomic.Int32
+			hook := func(context.Context) error { defer close(finished); <-release; return nil }
+			if phase == "stop" {
+				r.OnStop("stuck", hook)
+				r.OnDrain("stuck", func(context.Context) error { t.Error("Drain overlapped unfinished Stop"); return nil })
+			} else {
+				r.OnDrain("stuck", hook)
+			}
+			r.OnClose(func() error { closed.Add(1); return nil })
+			began := time.Now()
+			err = r.Finish(errors.New("listen failed"))
+			close(release)
+			<-finished
+			if time.Since(began) > time.Second || !errors.Is(err, ErrResourcesRetained) || !errors.Is(err, context.DeadlineExceeded) || closed.Load() != 0 {
+				t.Fatal(err)
+			}
+			_ = r.Close()
+			if closed.Load() != 0 {
+				t.Fatal("late cleanup raced work")
+			}
+		})
+	}
+}
+
+func TestPhasesConcurrentFinishAndValidation(t *testing.T) {
+	if _, err := NewRuntimeWithConfig(nil, ShutdownConfig{DrainTimeout: -1}); err == nil {
+		t.Fatal("negative timeout accepted")
+	}
+	r := testRuntime(t)
+	var called atomic.Int32
+	r.OnStop("worker", func(context.Context) error { called.Add(1); return nil })
+	r.OnDrain("worker", func(context.Context) error { called.Add(1); return nil })
+	r.OnClose(func() error { called.Add(1); return nil })
+	var wg sync.WaitGroup
+	for i := 0; i < 20; i++ {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			if err := r.Finish(nil); err != nil {
+				t.Error(err)
+			}
+		}()
+	}
+	wg.Wait()
+	if called.Load() != 3 {
+		t.Fatal(called.Load())
+	}
+	for _, fn := range []func(){func() { r.OnStop("late", func(context.Context) error { return nil }) }, func() { NewRuntime(nil).OnDrain("bad\nname", func(context.Context) error { return nil }) }, func() {
+		x := NewRuntime(nil)
+		x.OnStop("duplicate", func(context.Context) error { return nil })
+		x.OnStop("duplicate", func(context.Context) error { return nil })
+	}} {
+		func() {
+			defer func() {
+				if recover() == nil {
+					t.Error("invalid registration accepted")
+				}
+			}()
+			fn()
+		}()
+	}
+}
+
+func TestMissingDrainRetainsResources(t *testing.T) {
+	r := testRuntime(t)
+	r.OnStop("worker", func(context.Context) error { return nil })
+	closed := false
+	r.OnClose(func() error { closed = true; return nil })
+	if err := r.Close(); !errors.Is(err, ErrResourcesRetained) || closed {
+		t.Fatal("missing completion proof accepted", err)
+	}
+}