feat: serialize keyed tasks across queues with durable fencing tokens

sdk/go/keyed_concurrency_integration_test.go

@@ -0,0 +1,217 @@
+package orbit
+
+import (
+	"context"
+	"encoding/json"
+	"errors"
+	"sync"
+	"sync/atomic"
+	"testing"
+	"time"
+)
+
+func TestIntegrationWorkerRuntimeKeyedConcurrency(t *testing.T) {
+	api, _, token := integrationWorkerClient(t)
+	ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second)
+	defer cancel()
+	name := "keyed." + workerTestID(t)
+	gateway := newWorkerFaultGateway(t, api.baseURL)
+	client, err := NewClient(gateway.server.URL, Options{Token: token, HTTPClient: gateway.server.Client()})
+	if err != nil {
+		t.Fatal(err)
+	}
+	var active, calls atomic.Int32
+	var mu sync.Mutex
+	tokens := map[string]uint64{}
+	var last uint64
+	handler := runtimeHandler(name, func(ctx context.Context, e WorkerExecution) (json.RawMessage, error) {
+		calls.Add(1)
+		if active.Add(1) != 1 {
+			t.Error("same resource handlers overlapped")
+		}
+		defer active.Add(-1)
+		mu.Lock()
+		if e.Task.ConcurrencyKey != name || e.ConcurrencyToken <= last {
+			t.Error("key token was lost, duplicated or reset")
+		}
+		last = e.ConcurrencyToken
+		tokens[e.Task.ID] = e.ConcurrencyToken
+		mu.Unlock()
+		select {
+		case <-ctx.Done():
+			return nil, ctx.Err()
+		case <-time.After(100 * time.Millisecond):
+		}
+		return e.Task.Payload, nil
+	})
+	var tasks []Task
+	var workers []*Worker
+	for i := 0; i < 2; i++ {
+		q, err := api.CreateQueue(ctx, workerTestID(t), CreateQueueRequest{Name: name + workerTestID(t), MaxConcurrency: 4, RatePerSecond: 100})
+		if err != nil {
+			t.Fatal(err)
+		}
+		for j := 0; j < 2; j++ {
+			out, err := api.SubmitTask(ctx, workerTestID(t), SubmitTaskRequest{QueueID: q.Queue.ID, Handler: name, HandlerVersion: "v1", PayloadVersion: "v1", Payload: json.RawMessage(`{"n":9007199254740993}`), ConcurrencyKey: name})
+			if err != nil {
+				t.Fatal(err)
+			}
+			tasks = append(tasks, out.Task)
+		}
+		w, err := NewWorker(client, WorkerOptions{QueueID: q.Queue.ID, Capacity: 4, Handlers: []WorkerHandler{handler}})
+		if err != nil {
+			t.Fatal(err)
+		}
+		workers = append(workers, w)
+	}
+	// 在真实 Start 提交后丢 ACK；恢复必须沿用原 key/token，不能重跑处理器。
+	gateway.drop.Store("task.start")
+	run, stop := context.WithCancel(ctx)
+	done := make(chan error, 2)
+	for _, w := range workers {
+		go func() { done <- w.Run(run) }()
+	}
+	defer func() {
+		stop()
+		for range workers {
+			if err := <-done; err != nil {
+				t.Error(err)
+			}
+		}
+	}()
+	for _, task := range tasks {
+		finished := waitRuntimeState(t, ctx, api, task.ID, TaskSucceeded)
+		if finished.AttemptCount != 1 || finished.ConcurrencyKey != name {
+			t.Fatal("keyed task lost identity")
+		}
+		history, err := api.ListTaskAttempts(ctx, task.ID, HistoryPageRequest{})
+		if err != nil || len(history.Items) != 1 {
+			t.Fatal("missing keyed history", err)
+		}
+		mu.Lock()
+		want := tokens[task.ID]
+		mu.Unlock()
+		if want == 0 || history.Items[0].ConcurrencyToken != want {
+			t.Fatal("history token differs from handler")
+		}
+		detail, err := api.GetTaskAttempt(ctx, task.ID, history.Items[0].ID)
+		if err != nil || detail.ConcurrencyToken != want {
+			t.Fatal("detail token differs from handler", err)
+		}
+	}
+	gateway.expectDrop(t, "task.start")
+	if calls.Load() != 4 || active.Load() != 0 {
+		t.Fatal("incorrect handler count")
+	}
+}
+
+func TestIntegrationWorkerProcessKeyedConcurrency(t *testing.T) {
+	a := startOrbitProcess(t)
+	_, _, token := integrationWorkerClient(t)
+	api := processClient(t, a, token)
+	ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second)
+	defer cancel()
+	name := "key-crash." + workerTestID(t)
+	var registrations []WorkerRegistration
+	var tasks []Task
+	for i := 0; i < 2; i++ {
+		q, err := api.CreateQueue(ctx, workerTestID(t), CreateQueueRequest{Name: name + workerTestID(t), MaxConcurrency: 2, RatePerSecond: 100})
+		if err != nil {
+			t.Fatal(err)
+		}
+		r, err := NewWorkerRegistration(q.Queue.ID, 1, []WorkerCapability{{Handler: name, HandlerVersion: "v1", PayloadVersion: "v1"}})
+		if err != nil {
+			t.Fatal(err)
+		}
+		r.ConcurrencyKeys = true
+		registrations = append(registrations, r)
+		out, err := api.SubmitTask(ctx, workerTestID(t), SubmitTaskRequest{QueueID: q.Queue.ID, Handler: name, HandlerVersion: "v1", PayloadVersion: "v1", Payload: json.RawMessage(`{}`), ConcurrencyKey: name})
+		if err != nil {
+			t.Fatal(err)
+		}
+		tasks = append(tasks, out.Task)
+	}
+	dial := func(api *Client, r WorkerRegistration) *WorkerConnection {
+		t.Helper()
+		w, err := api.DialWorker(ctx, r)
+		if err != nil {
+			t.Fatal(err)
+		}
+		t.Cleanup(func() { w.Close() })
+		return w
+	}
+	capacity := func(w *WorkerConnection) {
+		t.Helper()
+		if _, err := w.Capacity(ctx, workerTestID(t), WorkerCapacity{Sequence: 1, Capacity: 1}); err != nil {
+			t.Fatal(err)
+		}
+	}
+	w1 := dial(api, registrations[0])
+	capacity(w1)
+	first, err := w1.NextAssignment(ctx)
+	if err != nil || first.ConcurrencyToken == 0 {
+		t.Fatal("first key not granted", err)
+	}
+	started, err := w1.Start(ctx, workerTestID(t), first.Attempt())
+	if err != nil {
+		t.Fatal(err)
+	}
+	if _, active := started.Lease.ExecutionDeadline(); !active {
+		t.Fatal("no current execution authorization")
+	}
+	a.kill(t)
+	b := startOrbitProcess(t)
+	api = processClient(t, b, token)
+	if a.nodeID == b.nodeID {
+		t.Fatal("node id reused")
+	}
+	w1 = dial(api, registrations[0])
+	lease, err := w1.Resume(ctx, workerTestID(t), first.Attempt())
+	if err != nil {
+		t.Fatal(err)
+	}
+	if _, active := lease.ExecutionDeadline(); !active {
+		t.Fatal("original attempt not resumed")
+	}
+	w2 := dial(api, registrations[1])
+	capacity(w2)
+	brief, stop := context.WithTimeout(ctx, 300*time.Millisecond)
+	_, err = w2.NextAssignment(brief)
+	stop()
+	if !errors.Is(err, context.DeadlineExceeded) {
+		t.Fatal("restart lost key exclusion", err)
+	}
+	waiting, err := api.GetTask(ctx, tasks[1].ID)
+	if err != nil || waiting.AttemptCount != 0 {
+		t.Fatal("blocked key consumed attempt", err)
+	}
+	history, err := api.GetTaskAttempt(ctx, tasks[0].ID, first.AttemptID)
+	if err != nil || history.ConcurrencyToken != first.ConcurrencyToken {
+		t.Fatal("restart changed key token", err)
+	}
+	finishID := workerTestID(t)
+	finish := WorkerFinish{Ref: first.Attempt(), Outcome: WorkerSucceeded, Result: json.RawMessage(`{}`)}
+	if _, err = w1.Finish(ctx, finishID, finish); err != nil {
+		t.Fatal(err)
+	}
+	second, err := w2.NextAssignment(ctx)
+	if err != nil || second.ConcurrencyToken <= first.ConcurrencyToken {
+		t.Fatal("restart reset sequence", err)
+	}
+	if _, err = w2.Start(ctx, workerTestID(t), second.Attempt()); err != nil {
+		t.Fatal(err)
+	}
+	// 第二次许可已经开始，旧完成回执的重放不能释放它。
+	if out, err := w1.Finish(ctx, finishID, finish); err != nil || !out.Deduplicated {
+		t.Fatal("duplicate finish changed ownership", err)
+	}
+	if _, err = w2.Renew(ctx, workerTestID(t), second.Attempt()); err != nil {
+		t.Fatal(err)
+	}
+	if _, err = w2.Finish(ctx, workerTestID(t), WorkerFinish{Ref: second.Attempt(), Outcome: WorkerSucceeded, Result: json.RawMessage(`{}`)}); err != nil {
+		t.Fatal(err)
+	}
+	for _, task := range tasks {
+		waitRuntimeState(t, ctx, api, task.ID, TaskSucceeded)
+	}
+}