411 lines
12 KiB
Go
411 lines
12 KiB
Go
package orbit
|
||
|
||
import (
|
||
"context"
|
||
"encoding/json"
|
||
"errors"
|
||
"net/http"
|
||
"net/http/httptest"
|
||
"os"
|
||
"strings"
|
||
"sync"
|
||
"sync/atomic"
|
||
"testing"
|
||
"time"
|
||
|
||
"github.com/coder/websocket"
|
||
)
|
||
|
||
// 由 scripts/test-mysql.sh 启动真实服务与隔离 MySQL；SDK 不导入服务端内部包。
|
||
func integrationWorkerClient(t *testing.T) (*Client, string, string) {
|
||
t.Helper()
|
||
path := os.Getenv("ORBIT_TEST_WORKER_CREDENTIALS_FILE")
|
||
base := os.Getenv("ORBIT_TEST_WORKER_URL")
|
||
if path == "" || base == "" {
|
||
t.Skip("requires isolated Orbit server and Worker credentials")
|
||
}
|
||
if !strings.HasPrefix(base, "http://127.0.0.1:") {
|
||
t.Fatal("Worker integration requires an isolated loopback server")
|
||
}
|
||
raw, err := os.ReadFile(path)
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
var credentials struct {
|
||
Token string `json:"token"`
|
||
QueueID string `json:"queue_id"`
|
||
}
|
||
if json.Unmarshal(raw, &credentials) != nil || credentials.Token == "" || !validID(credentials.QueueID) {
|
||
t.Fatal("invalid integration credentials")
|
||
}
|
||
c, err := NewClient(base, Options{Token: credentials.Token})
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
return c, credentials.QueueID, credentials.Token
|
||
}
|
||
|
||
// workerFaultGateway 在实际服务端确认后丢弃指定 ACK；SDK 到网关使用真实 TLS。
|
||
// 只丢响应，不伪造任务状态，测试凭据不进入 URL 或输出。
|
||
type workerFaultGateway struct {
|
||
drop atomic.Value
|
||
offline atomic.Bool
|
||
mu sync.Mutex
|
||
active map[*websocket.Conn]struct{}
|
||
dropped chan string
|
||
server *httptest.Server
|
||
}
|
||
|
||
func newWorkerFaultGateway(t *testing.T, upstream string) *workerFaultGateway {
|
||
t.Helper()
|
||
g := &workerFaultGateway{dropped: make(chan string, 4), active: make(map[*websocket.Conn]struct{})}
|
||
g.drop.Store("")
|
||
ctx, cancel := context.WithCancel(context.Background())
|
||
var handlers sync.WaitGroup
|
||
g.server = httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||
if g.offline.Load() {
|
||
http.Error(w, "gateway offline", 503)
|
||
return
|
||
}
|
||
handlers.Add(1)
|
||
defer handlers.Done()
|
||
bridge, stop := context.WithCancel(ctx)
|
||
defer stop()
|
||
header := http.Header{}
|
||
header.Set("Authorization", r.Header.Get("Authorization"))
|
||
up, _, err := websocket.Dial(bridge, upstream+"/v1/worker/connect", &websocket.DialOptions{HTTPHeader: header, Subprotocols: []string{workerSubprotocol}})
|
||
if err != nil {
|
||
http.Error(w, "upstream unavailable", 503)
|
||
return
|
||
}
|
||
defer up.CloseNow()
|
||
up.SetReadLimit(512 << 10)
|
||
down, err := websocket.Accept(w, r, &websocket.AcceptOptions{Subprotocols: []string{workerSubprotocol}})
|
||
if err != nil {
|
||
return
|
||
}
|
||
defer down.CloseNow()
|
||
g.mu.Lock()
|
||
g.active[down] = struct{}{}
|
||
g.mu.Unlock()
|
||
defer func() { g.mu.Lock(); delete(g.active, down); g.mu.Unlock() }()
|
||
if g.offline.Load() {
|
||
return
|
||
}
|
||
down.SetReadLimit(512 << 10)
|
||
var mu sync.Mutex
|
||
var forwarding sync.WaitGroup
|
||
forwarding.Add(2)
|
||
operations := map[string]string{}
|
||
finished := make(chan struct{}, 1)
|
||
go func() {
|
||
defer forwarding.Done()
|
||
defer func() { finished <- struct{}{} }()
|
||
for {
|
||
kind, raw, err := down.Read(bridge)
|
||
if err != nil {
|
||
return
|
||
}
|
||
var f workerFrame
|
||
if json.Unmarshal(raw, &f) != nil {
|
||
return
|
||
}
|
||
mu.Lock()
|
||
operations[f.MessageID] = f.Type
|
||
mu.Unlock()
|
||
if up.Write(bridge, kind, raw) != nil {
|
||
return
|
||
}
|
||
}
|
||
}()
|
||
go func() {
|
||
defer forwarding.Done()
|
||
defer stop()
|
||
for {
|
||
kind, raw, err := up.Read(bridge)
|
||
if err != nil {
|
||
return
|
||
}
|
||
var f workerFrame
|
||
if json.Unmarshal(raw, &f) != nil {
|
||
return
|
||
}
|
||
mu.Lock()
|
||
operation := operations[f.ReplyTo]
|
||
delete(operations, f.ReplyTo)
|
||
mu.Unlock()
|
||
if f.Type == "ack" && operation != "" && g.drop.CompareAndSwap(operation, "") {
|
||
g.dropped <- operation
|
||
return
|
||
}
|
||
if down.Write(bridge, kind, raw) != nil {
|
||
return
|
||
}
|
||
}
|
||
}()
|
||
select {
|
||
case <-bridge.Done():
|
||
case <-finished:
|
||
}
|
||
stop()
|
||
down.CloseNow()
|
||
up.CloseNow()
|
||
forwarding.Wait()
|
||
}))
|
||
t.Cleanup(func() { cancel(); g.server.Close(); handlers.Wait() })
|
||
return g
|
||
}
|
||
func (g *workerFaultGateway) disconnectActive() {
|
||
g.mu.Lock()
|
||
active := make([]*websocket.Conn, 0, len(g.active))
|
||
for c := range g.active {
|
||
active = append(active, c)
|
||
}
|
||
g.mu.Unlock()
|
||
for _, c := range active {
|
||
c.CloseNow()
|
||
}
|
||
}
|
||
func (g *workerFaultGateway) expectDrop(t *testing.T, operation string) {
|
||
t.Helper()
|
||
select {
|
||
case got := <-g.dropped:
|
||
if got != operation {
|
||
t.Fatal("wrong dropped operation")
|
||
}
|
||
case <-time.After(2 * time.Second):
|
||
t.Fatal("ACK was not dropped after upstream confirmation")
|
||
}
|
||
}
|
||
|
||
func TestIntegrationWorkerConnectionACKLoss(t *testing.T) {
|
||
api, queue, token := integrationWorkerClient(t)
|
||
ctx, cancel := context.WithTimeout(context.Background(), 25*time.Second)
|
||
defer cancel()
|
||
r, err := NewWorkerRegistration(queue, 1, []WorkerCapability{{Handler: "sdk.ack-loss", HandlerVersion: "v1", PayloadVersion: "v1"}})
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
g := newWorkerFaultGateway(t, api.baseURL)
|
||
client, err := NewClient(g.server.URL, Options{Token: token, HTTPClient: g.server.Client()})
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
dial := func() *WorkerConnection {
|
||
t.Helper()
|
||
w, err := client.DialWorker(ctx, r)
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
t.Cleanup(func() { w.Close() })
|
||
return w
|
||
}
|
||
accepted, err := api.SubmitTask(ctx, "sdk-ack-"+workerTestID(t), SubmitTaskRequest{QueueID: queue, Handler: r.Capabilities[0].Handler, HandlerVersion: "v1", PayloadVersion: "v1", Payload: json.RawMessage(`{"value":9007199254740993}`), TaskTimeoutMS: 60000, Retry: RetryPolicy{MaxAttempts: 1, AttemptTimeoutMS: 30000}})
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
w := dial()
|
||
if _, err := w.Capacity(ctx, workerTestID(t), WorkerCapacity{Sequence: 1, Capacity: 1}); err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
a, err := w.NextAssignment(ctx)
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
if a.Task.ID != accepted.Task.ID {
|
||
t.Fatal("wrong assigned task")
|
||
}
|
||
startID := workerTestID(t)
|
||
g.drop.Store("task.start")
|
||
if _, err := w.Start(ctx, startID, a.Attempt()); err == nil {
|
||
t.Fatal("lost start ACK reported success")
|
||
}
|
||
g.expectDrop(t, "task.start")
|
||
w.Close()
|
||
w = dial()
|
||
lease, err := w.Resume(ctx, workerTestID(t), a.Attempt())
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
if _, active := lease.ExecutionDeadline(); !active {
|
||
t.Fatal("committed start did not resume")
|
||
}
|
||
started, err := w.Start(ctx, startID, a.Attempt())
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
if !started.Deduplicated {
|
||
t.Fatal("start receipt was recreated")
|
||
}
|
||
if _, active := started.Lease.ExecutionDeadline(); !active {
|
||
t.Fatal("resume missing current authority")
|
||
}
|
||
// 处理器仅在首次确认当前授权后运行。之后连接恢复只核对，不再次运行它。
|
||
var handlerCalls int
|
||
handlerCalls++
|
||
if _, err := w.Capacity(ctx, workerTestID(t), WorkerCapacity{Sequence: 1, Capacity: 1}); err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
renewID := workerTestID(t)
|
||
first, err := w.Renew(ctx, renewID, a.Attempt())
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
second, err := w.Renew(ctx, renewID, a.Attempt())
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
if first.Deduplicated || !second.Deduplicated || first.Receipt.TaskRevision != second.Receipt.TaskRevision || !first.Receipt.LeaseUntil.Equal(second.Receipt.LeaseUntil) {
|
||
t.Fatal("duplicate renew changed original receipt")
|
||
}
|
||
finishID := workerTestID(t)
|
||
result := json.RawMessage(`{"normalized":"hello","sequence":9007199254740993,"ratio":0.1234567890123456789012345}`)
|
||
finish := WorkerFinish{Ref: a.Attempt(), Outcome: WorkerSucceeded, Result: result}
|
||
g.drop.Store("task.finish")
|
||
if _, err := w.Finish(ctx, finishID, finish); err == nil {
|
||
t.Fatal("lost finish ACK reported success")
|
||
}
|
||
g.expectDrop(t, "task.finish")
|
||
w.Close()
|
||
w = dial()
|
||
lease, err = w.Resume(ctx, workerTestID(t), a.Attempt())
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
if lease.Active || lease.State != "SUCCEEDED" {
|
||
t.Fatal("completed attempt resumed as executable")
|
||
}
|
||
confirmed, err := w.Finish(ctx, finishID, finish)
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
if !confirmed.Deduplicated || confirmed.Receipt.Disposition != "APPLIED" || confirmed.Receipt.TaskState != TaskSucceeded {
|
||
t.Fatal("missing durable finish receipt")
|
||
}
|
||
stale, err := w.Start(ctx, startID, a.Attempt())
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
if !stale.Deduplicated {
|
||
t.Fatal("missing old start receipt")
|
||
}
|
||
if _, active := stale.Lease.ExecutionDeadline(); active {
|
||
t.Fatal("old start ACK resurrected execution")
|
||
}
|
||
conflict := finish
|
||
conflict.Result = json.RawMessage(`{"normalized":"conflicting evidence"}`)
|
||
evidence, err := w.Finish(ctx, workerTestID(t), conflict)
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
if evidence.Receipt.Disposition != "EVIDENCE_ONLY" {
|
||
t.Fatal("conflicting result replaced authority")
|
||
}
|
||
current, err := api.GetTask(ctx, a.Task.ID)
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
if current.State != TaskSucceeded || current.AttemptCount != 1 || string(current.Result) != string(result) || handlerCalls != 1 {
|
||
t.Fatal("result, precision or execution count changed")
|
||
}
|
||
drain, err := w.Capacity(ctx, workerTestID(t), WorkerCapacity{Sequence: 1, Capacity: 0, Drain: true})
|
||
if err != nil || !drain.Draining {
|
||
t.Fatalf("drain: %v", err)
|
||
}
|
||
w.Close()
|
||
w = dial()
|
||
if !w.Welcome().Draining {
|
||
t.Fatal("drain did not persist across reconnect")
|
||
}
|
||
_, err = w.Capacity(ctx, workerTestID(t), WorkerCapacity{Sequence: 1, Capacity: 1})
|
||
var control *WorkerError
|
||
if !errors.As(err, &control) || control.Code != "STATE_CONFLICT" {
|
||
t.Fatalf("clearing persistent drain was not rejected: %v", err)
|
||
}
|
||
drain, err = w.Capacity(ctx, workerTestID(t), WorkerCapacity{Sequence: 1, Capacity: 0, Drain: true})
|
||
if err != nil || !drain.Draining {
|
||
t.Fatalf("draining heartbeat failed: %v", err)
|
||
}
|
||
}
|
||
|
||
func TestIntegrationWorkerConnectionExpiredLease(t *testing.T) {
|
||
api, queue, _ := integrationWorkerClient(t)
|
||
ctx, cancel := context.WithTimeout(context.Background(), 12*time.Second)
|
||
defer cancel()
|
||
r, err := NewWorkerRegistration(queue, 1, []WorkerCapability{{Handler: "sdk.expiry", HandlerVersion: "v1", PayloadVersion: "v1"}})
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
_, err = api.SubmitTask(ctx, "sdk-expiry-"+workerTestID(t), SubmitTaskRequest{QueueID: queue, Handler: "sdk.expiry", HandlerVersion: "v1", PayloadVersion: "v1", Payload: json.RawMessage(`{}`), TaskTimeoutMS: 30000, Retry: RetryPolicy{MaxAttempts: 1, AttemptTimeoutMS: 1100}})
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
w, err := api.DialWorker(ctx, r)
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
defer w.Close()
|
||
if _, err := w.Capacity(ctx, workerTestID(t), WorkerCapacity{Sequence: 1, Capacity: 1}); err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
a, err := w.NextAssignment(ctx)
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
startID := workerTestID(t)
|
||
out, err := w.Start(ctx, startID, a.Attempt())
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
deadline, active := out.Lease.ExecutionDeadline()
|
||
if !active {
|
||
t.Fatal("short valid lease was rejected")
|
||
}
|
||
// 用单调期限驱动协作取消；额外等待安全余量确保服务端授权也已过期。
|
||
work, stop := context.WithDeadline(ctx, deadline)
|
||
defer stop()
|
||
<-work.Done()
|
||
if !errors.Is(work.Err(), context.DeadlineExceeded) {
|
||
t.Fatal(work.Err())
|
||
}
|
||
select {
|
||
case <-ctx.Done():
|
||
t.Fatal(ctx.Err())
|
||
case <-time.After(350 * time.Millisecond):
|
||
}
|
||
stale, err := w.Start(ctx, startID, a.Attempt())
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
if !stale.Deduplicated {
|
||
t.Fatal("missing historical receipt")
|
||
}
|
||
if _, active := stale.Lease.ExecutionDeadline(); active {
|
||
t.Fatal("expired historical ACK authorized work")
|
||
}
|
||
late, err := w.Finish(ctx, workerTestID(t), WorkerFinish{Ref: a.Attempt(), Outcome: WorkerUnknown, Result: json.RawMessage(`{"reason":"handler_cancelled"}`), ErrorCode: "LEASE_EXPIRED"})
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
if late.Receipt.Disposition != "EVIDENCE_ONLY" {
|
||
t.Fatal("late report changed authoritative result")
|
||
}
|
||
for {
|
||
current, err := api.GetTask(ctx, a.Task.ID)
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
if current.State == TaskNeedsReview {
|
||
if current.AttemptCount != 1 || len(current.Result) != 0 {
|
||
t.Fatal("expired execution was replayed or evidence promoted")
|
||
}
|
||
break
|
||
}
|
||
select {
|
||
case <-ctx.Done():
|
||
t.Fatal("server did not recover expired execution")
|
||
case <-time.After(100 * time.Millisecond):
|
||
}
|
||
}
|
||
}
|