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) } }