feat: add automatic Worker runtime with lease recovery and bounded shutdown

2026-09-17 14:29:28 +08:00
parent 1753fcbccb
commit da123bba69
22 changed files with 1784 additions and 17 deletions
--- a/sdk/go/worker_runtime_integration_test.go
+++ b/sdk/go/worker_runtime_integration_test.go
@@ -0,0 +1,357 @@
+package orbit
+
+import (
+	"context"
+	"encoding/json"
+	"errors"
+	"sync/atomic"
+	"testing"
+	"time"
+)
+
+func runtimeTask(t *testing.T, ctx context.Context, c *Client, queue, handler string, timeout uint64) Task {
+	t.Helper()
+	out, err := c.SubmitTask(ctx, "runtime-"+workerTestID(t), SubmitTaskRequest{QueueID: queue, Handler: handler, HandlerVersion: "v1", PayloadVersion: "v1", Payload: json.RawMessage(`{"sequence":9007199254740993}`), TaskTimeoutMS: 120000, Retry: RetryPolicy{MaxAttempts: 1, AttemptTimeoutMS: timeout}})
+	if err != nil {
+		t.Fatal(err)
+	}
+	return out.Task
+}
+func runtimeHandler(name string, handle func(context.Context, WorkerExecution) (json.RawMessage, error)) WorkerHandler {
+	return WorkerHandler{Capability: WorkerCapability{Handler: name, HandlerVersion: "v1", PayloadVersion: "v1"}, Handle: handle}
+}
+func waitRuntimeState(t *testing.T, ctx context.Context, c *Client, id string, state TaskState) Task {
+	t.Helper()
+	for {
+		task, err := c.GetTask(ctx, id)
+		if err != nil {
+			t.Fatal(err)
+		}
+		if task.State == state {
+			return task
+		}
+		select {
+		case <-ctx.Done():
+			t.Fatalf("task did not reach %s", state)
+		case <-time.After(50 * time.Millisecond):
+		}
+	}
+}
+
+func TestIntegrationWorkerRuntimeACKLossRenewAndDrain(t *testing.T) {
+	api, queue, token := integrationWorkerClient(t)
+	ctx, cancel := context.WithTimeout(context.Background(), 55*time.Second)
+	defer cancel()
+	g := newWorkerFaultGateway(t, api.baseURL)
+	client, err := NewClient(g.server.URL, Options{Token: token, HTTPClient: g.server.Client()})
+	if err != nil {
+		t.Fatal(err)
+	}
+	name := "runtime.renew." + workerTestID(t)
+	entered := make(chan struct{})
+	var calls atomic.Int32
+	w, err := NewWorker(client, WorkerOptions{QueueID: queue, Capacity: 1, ShutdownTimeout: 45 * time.Second, Handlers: []WorkerHandler{runtimeHandler(name, func(work context.Context, e WorkerExecution) (json.RawMessage, error) {
+		calls.Add(1)
+		close(entered)
+		// 必须跨过服务端最初的 30 秒租约；只有自动续租成功才能获得 APPLIED。
+		select {
+		case <-work.Done():
+			return nil, work.Err()
+		case <-time.After(32 * time.Second):
+			return json.RawMessage(`{"sequence":9007199254740993,"done":true}`), nil
+		}
+	})}})
+	if err != nil {
+		t.Fatal(err)
+	}
+	task := runtimeTask(t, ctx, api, queue, name, 90000)
+	g.drop.Store("task.start")
+	runCtx, stop := context.WithCancel(ctx)
+	defer stop()
+	done := make(chan error, 1)
+	go func() { done <- w.Run(runCtx) }()
+	g.expectDrop(t, "task.start")
+	select {
+	case <-entered:
+	case err := <-done:
+		t.Fatalf("Worker stopped before handler: %v", err)
+	case <-ctx.Done():
+		t.Fatal("handler did not start")
+	}
+	g.drop.Store("task.finish")
+	stop() // 取消 Run 只进入排空，不能立即取消正在执行的函数。
+	select {
+	case err := <-done:
+		if err != nil {
+			t.Fatal(err)
+		}
+	case <-ctx.Done():
+		t.Fatal("Worker did not finish within grace")
+	}
+	g.expectDrop(t, "task.finish")
+	current := waitRuntimeState(t, ctx, api, task.ID, TaskSucceeded)
+	stats := w.Stats()
+	if calls.Load() != 1 || stats.Started != 1 || stats.Running != 0 || stats.ReportsConfirmed != 1 || stats.Reconnects < 2 || stats.RenewalsConfirmed < 2 || current.AttemptCount != 1 || string(current.Result) != `{"sequence":9007199254740993,"done":true}` {
+		t.Fatalf("automatic execution invariant failed: calls=%d stats=%+v attempts=%d", calls.Load(), stats, current.AttemptCount)
+	}
+	// 仍在等待的同类工作不能被已排空的进程领取。
+	waiting := runtimeTask(t, ctx, api, queue, name, 10000)
+	observed, err := api.GetTask(ctx, waiting.ID)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if observed.State != TaskWaiting {
+		t.Fatal("drained Worker accepted another task")
+	}
+	if _, err := api.CancelTask(ctx, observed.ID, "cleanup-"+workerTestID(t), CancelTaskRequest{ExpectedRevision: observed.Revision, Reason: "finish runtime drain regression"}); err != nil {
+		t.Fatal(err)
+	}
+}
+
+func TestIntegrationWorkerRuntimeUnknownAndUncooperativeShutdown(t *testing.T) {
+	api, queue, _ := integrationWorkerClient(t)
+	ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second)
+	defer cancel()
+	for _, test := range []struct {
+		name   string
+		handle func(context.Context, WorkerExecution) (json.RawMessage, error)
+		state  TaskState
+	}{
+		{"known", func(context.Context, WorkerExecution) (json.RawMessage, error) {
+			return nil, &WorkerFailure{Code: "INVALID_OBJECT"}
+		}, TaskFailed},
+		{"unknown", func(context.Context, WorkerExecution) (json.RawMessage, error) {
+			return nil, errors.New("external request timed out with private details")
+		}, TaskNeedsReview},
+		{"panic", func(context.Context, WorkerExecution) (json.RawMessage, error) { panic("private panic") }, TaskNeedsReview},
+	} {
+		t.Run(test.name, func(t *testing.T) {
+			name := "runtime." + test.name + "." + workerTestID(t)
+			w, err := NewWorker(api, WorkerOptions{QueueID: queue, Capacity: 1, Handlers: []WorkerHandler{runtimeHandler(name, test.handle)}})
+			if err != nil {
+				t.Fatal(err)
+			}
+			task := runtimeTask(t, ctx, api, queue, name, 5000)
+			run, stop := context.WithCancel(ctx)
+			done := make(chan error, 1)
+			go func() { done <- w.Run(run) }()
+			defer stop()
+			waitRuntimeState(t, ctx, api, task.ID, test.state)
+			stop()
+			select {
+			case err := <-done:
+				if err != nil {
+					t.Fatal(err)
+				}
+			case <-ctx.Done():
+				t.Fatal("classification Worker did not stop")
+			}
+		})
+	}
+	t.Run("uncooperative", func(t *testing.T) {
+		name := "runtime.uncooperative." + workerTestID(t)
+		entered, release, returned := make(chan struct{}), make(chan struct{}), make(chan struct{})
+		defer close(release)
+		w, err := NewWorker(api, WorkerOptions{QueueID: queue, Capacity: 1, ShutdownTimeout: 250 * time.Millisecond, Handlers: []WorkerHandler{runtimeHandler(name, func(context.Context, WorkerExecution) (json.RawMessage, error) {
+			close(entered)
+			<-release
+			close(returned)
+			return json.RawMessage(`{}`), nil
+		})}})
+		if err != nil {
+			t.Fatal(err)
+		}
+		runtimeTask(t, ctx, api, queue, name, 1500)
+		run, stop := context.WithCancel(ctx)
+		done := make(chan error, 1)
+		go func() { done <- w.Run(run) }()
+		defer stop()
+		select {
+		case <-entered:
+		case <-ctx.Done():
+			t.Fatal("handler did not enter")
+		}
+		began := time.Now()
+		stop()
+		select {
+		case err := <-done:
+			var stopped *WorkerRunError
+			if !errors.As(err, &stopped) || !errors.Is(err, ErrWorkerShutdownTimeout) || stopped.Running != 1 || stopped.Unconfirmed != 1 {
+				t.Fatalf("uncooperative handler hidden: %v", err)
+			}
+		case <-time.After(2 * time.Second):
+			t.Fatal("shutdown waited forever for handler")
+		}
+		if time.Since(began) > time.Second {
+			t.Fatal("shutdown exceeded bounded grace")
+		}
+		select {
+		case <-returned:
+			t.Fatal("SDK claimed to kill uncooperative handler")
+		default:
+		}
+	})
+}
+
+func TestIntegrationWorkerRuntimeLeaseLossCancelsAndDoesNotRevive(t *testing.T) {
+	api, queue, _ := integrationWorkerClient(t)
+	ctx, cancel := context.WithTimeout(context.Background(), 12*time.Second)
+	defer cancel()
+	name := "runtime.expired." + workerTestID(t)
+	cause := make(chan error, 1)
+	w, err := NewWorker(api, WorkerOptions{QueueID: queue, Capacity: 1, Handlers: []WorkerHandler{runtimeHandler(name, func(work context.Context, _ WorkerExecution) (json.RawMessage, error) {
+		<-work.Done()
+		cause <- context.Cause(work)
+		return nil, work.Err()
+	})}})
+	if err != nil {
+		t.Fatal(err)
+	}
+	task := runtimeTask(t, ctx, api, queue, name, 1400)
+	run, stop := context.WithCancel(ctx)
+	done := make(chan error, 1)
+	go func() { done <- w.Run(run) }()
+	defer stop()
+	select {
+	case err := <-cause:
+		if !errors.Is(err, ErrExecutionLeaseLost) {
+			t.Fatal(err)
+		}
+	case <-ctx.Done():
+		t.Fatal("handler did not observe lease cancellation")
+	}
+	waitRuntimeState(t, ctx, api, task.ID, TaskNeedsReview)
+	stop()
+	select {
+	case err := <-done:
+		if err != nil {
+			t.Fatal(err)
+		}
+	case <-ctx.Done():
+		t.Fatal("Worker did not stop")
+	}
+	if w.Stats().Started != 1 {
+		t.Fatal("expired handler was invoked again")
+	}
+}
+
+func TestIntegrationWorkerRuntimeOutageCancelsAndReplaysEvidence(t *testing.T) {
+	api, queue, token := integrationWorkerClient(t)
+	ctx, cancel := context.WithTimeout(context.Background(), 18*time.Second)
+	defer cancel()
+	g := newWorkerFaultGateway(t, api.baseURL)
+	client, err := NewClient(g.server.URL, Options{Token: token, HTTPClient: g.server.Client()})
+	if err != nil {
+		t.Fatal(err)
+	}
+	name := "runtime.outage." + workerTestID(t)
+	entered, cancelled := make(chan struct{}), make(chan error, 1)
+	w, err := NewWorker(client, WorkerOptions{QueueID: queue, Capacity: 1, ReconnectMaxDelay: time.Second, Handlers: []WorkerHandler{runtimeHandler(name, func(work context.Context, _ WorkerExecution) (json.RawMessage, error) {
+		close(entered)
+		<-work.Done()
+		cancelled <- context.Cause(work)
+		return nil, work.Err()
+	})}})
+	if err != nil {
+		t.Fatal(err)
+	}
+	task := runtimeTask(t, ctx, api, queue, name, 2000)
+	run, stop := context.WithCancel(ctx)
+	done := make(chan error, 1)
+	go func() { done <- w.Run(run) }()
+	defer stop()
+	select {
+	case <-entered:
+	case <-ctx.Done():
+		t.Fatal("handler did not start")
+	}
+	g.offline.Store(true)
+	g.disconnectActive()
+	select {
+	case cause := <-cancelled:
+		if !errors.Is(cause, ErrExecutionLeaseLost) {
+			t.Fatal(cause)
+		}
+	case <-ctx.Done():
+		t.Fatal("disconnected handler ignored local expiry")
+	}
+	waitRuntimeState(t, ctx, api, task.ID, TaskNeedsReview)
+	g.offline.Store(false)
+	for w.Stats().EvidenceConfirmed != 1 {
+		select {
+		case err := <-done:
+			t.Fatalf("Worker stopped before evidence replay: %v", err)
+		case <-ctx.Done():
+			t.Fatal("pending result was not replayed")
+		case <-time.After(50 * time.Millisecond):
+		}
+	}
+	stop()
+	select {
+	case err := <-done:
+		if err != nil {
+			t.Fatal(err)
+		}
+	case <-ctx.Done():
+		t.Fatal("Worker did not stop")
+	}
+	current, err := api.GetTask(ctx, task.ID)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if w.Stats().Started != 1 || current.AttemptCount != 1 || current.State != TaskNeedsReview || len(current.Result) != 0 {
+		t.Fatal("outage revived handler or promoted late evidence")
+	}
+}
+
+func TestIntegrationWorkerRuntimeConcurrency(t *testing.T) {
+	api, queue, _ := integrationWorkerClient(t)
+	ctx, cancel := context.WithTimeout(context.Background(), 18*time.Second)
+	defer cancel()
+	name := "runtime.concurrency." + workerTestID(t)
+	var running, peak atomic.Int32
+	w, err := NewWorker(api, WorkerOptions{QueueID: queue, Capacity: 2, Handlers: []WorkerHandler{runtimeHandler(name, func(work context.Context, _ WorkerExecution) (json.RawMessage, error) {
+		value := running.Add(1)
+		defer running.Add(-1)
+		for old := peak.Load(); value > old; old = peak.Load() {
+			if peak.CompareAndSwap(old, value) {
+				break
+			}
+		}
+		select {
+		case <-work.Done():
+			return nil, work.Err()
+		case <-time.After(300 * time.Millisecond):
+			return json.RawMessage(`{"done":true}`), nil
+		}
+	})}})
+	if err != nil {
+		t.Fatal(err)
+	}
+	var tasks []Task
+	for range 6 {
+		tasks = append(tasks, runtimeTask(t, ctx, api, queue, name, 10000))
+	}
+	run, stop := context.WithCancel(ctx)
+	done := make(chan error, 1)
+	go func() { done <- w.Run(run) }()
+	defer stop()
+	for _, task := range tasks {
+		current := waitRuntimeState(t, ctx, api, task.ID, TaskSucceeded)
+		if current.AttemptCount != 1 {
+			t.Fatal("unexpected new attempt")
+		}
+	}
+	stop()
+	select {
+	case err := <-done:
+		if err != nil {
+			t.Fatal(err)
+		}
+	case <-ctx.Done():
+		t.Fatal("Worker did not drain")
+	}
+	if peak.Load() != 2 || running.Load() != 0 || w.Stats().Started != 6 || w.Stats().ReportsConfirmed != 6 {
+		t.Fatalf("concurrency invariant failed: peak=%d stats=%+v", peak.Load(), w.Stats())
+	}
+}