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

task/task.go

@@ -227,10 +227,9 @@ func run(ctx context.Context, fn Func) (err error) {
 	return err
 }
 
-// Shutdown stops schedules and submission, then drains accepted jobs. Deadline
-// expiry cancels running/queued jobs and returns without waiting for code that
-// ignores cancellation. Such code cannot be forcibly stopped by Go.
-func (r *Runner) Shutdown(ctx context.Context) error {
+// StopAccepting stops schedules and submission without canceling accepted jobs.
+// Pair it with Wait to establish completion before releasing dependencies.
+func (r *Runner) StopAccepting(ctx context.Context) error {
 	if ctx == nil {
 		return errors.New("task: nil context")
 	}
@@ -242,6 +241,20 @@ func (r *Runner) Shutdown(ctx context.Context) error {
 		go func() { r.wg.Wait(); r.cancel(); close(r.done) }()
 	}
 	r.mu.Unlock()
+	return nil
+}
+
+// Wait joins all workers and schedules after StopAccepting. A canceled context
+// cancels work, but callers can wait again to establish actual completion.
+func (r *Runner) Wait(ctx context.Context) error {
+	if ctx == nil {
+		return errors.New("task: nil context")
+	}
+	select {
+	case <-r.done:
+		return nil
+	default:
+	}
 	select {
 	case <-r.done:
 		return nil
@@ -251,6 +264,14 @@ func (r *Runner) Shutdown(ctx context.Context) error {
 	}
 }
 
+// Shutdown stops admission and waits with the caller's deadline.
+func (r *Runner) Shutdown(ctx context.Context) error {
+	if err := r.StopAccepting(ctx); err != nil {
+		return err
+	}
+	return r.Wait(ctx)
+}
+
 func (r *Runner) Close() error {
 	ctx, cancel := context.WithTimeout(context.Background(), r.config.ShutdownTimeout)
 	defer cancel()