go-omni/lifecycle/phases_test.go

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