orbit/sdk/go/worker_runtime_integration_test.go

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