218 lines
6.9 KiB
Go
218 lines
6.9 KiB
Go
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)
|
||
}
|
||
}
|