feat: add generic phased background shutdown to managed services

task/phases_test.go

@@ -0,0 +1,60 @@
+package task_test
+
+import (
+	"context"
+	"errors"
+	"git.yuntuo.hk/k2safe/go-omni/lifecycle"
+	"git.yuntuo.hk/k2safe/go-omni/task"
+	"sync/atomic"
+	"testing"
+	"time"
+)
+
+func TestRunnerPhasedShutdown(t *testing.T) {
+	runtime := lifecycle.NewRuntime(context.Background())
+	runner, err := task.New(task.Config{Workers: 1, QueueSize: 1})
+	if err != nil {
+		t.Fatal(err)
+	}
+	var resourceClosed atomic.Bool
+	runtime.OnClose(func() error { resourceClosed.Store(true); return nil })
+	started, release, draining := make(chan struct{}), make(chan struct{}), make(chan struct{})
+	runtime.OnStop("jobs", runner.StopAccepting)
+	runtime.OnDrain("jobs", func(ctx context.Context) error { close(draining); return runner.Wait(ctx) })
+	future, err := runner.Submit(runtime.WorkContext(), "write", func(ctx context.Context) error {
+		close(started)
+		<-release
+		if resourceClosed.Load() {
+			return errors.New("resource already closed")
+		}
+		return ctx.Err()
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+	<-started
+	runtime.RequestShutdown(nil)
+	done := make(chan error, 1)
+	go func() { done <- runtime.Close() }()
+	select {
+	case <-draining:
+	case <-time.After(time.Second):
+		t.Fatal("did not drain")
+	}
+	if _, err := runner.Submit(context.Background(), "late", func(context.Context) error { return nil }); !errors.Is(err, task.ErrClosed) {
+		t.Fatal(err)
+	}
+	if resourceClosed.Load() {
+		t.Fatal("closed resource before task")
+	}
+	close(release)
+	if err := <-done; err != nil {
+		t.Fatal(err)
+	}
+	if err := future.Wait(context.Background()); err != nil {
+		t.Fatal(err)
+	}
+	if !resourceClosed.Load() {
+		t.Fatal("resource leaked")
+	}
+}