425 lines
15 KiB
Go
425 lines
15 KiB
Go
package orbit
|
||
|
||
import (
|
||
"context"
|
||
"encoding/json"
|
||
"errors"
|
||
"fmt"
|
||
"net/http"
|
||
"net/http/httptest"
|
||
"strings"
|
||
"sync"
|
||
"sync/atomic"
|
||
"testing"
|
||
"time"
|
||
|
||
"github.com/coder/websocket"
|
||
)
|
||
|
||
const workerTestQueue = "11111111111111111111111111111111"
|
||
const workerTestSession = "22222222222222222222222222222222"
|
||
const workerTestTask = "33333333333333333333333333333333"
|
||
const workerTestAttempt = "44444444444444444444444444444444"
|
||
|
||
func workerTestID(t *testing.T) string {
|
||
t.Helper()
|
||
s, err := NewWorkerMessageID()
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
return s
|
||
}
|
||
func workerTestRegistration(t *testing.T) WorkerRegistration {
|
||
t.Helper()
|
||
r, err := NewWorkerRegistration(workerTestQueue, 2, []WorkerCapability{{Handler: "text.normalize", HandlerVersion: "v1", PayloadVersion: "v1"}})
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
return r
|
||
}
|
||
func workerTestWelcome() WorkerWelcome {
|
||
return WorkerWelcome{SessionID: workerTestSession, Epoch: 1, ServerTime: time.Now().UTC(), ExpiresAt: time.Now().UTC().Add(45 * time.Second), LeaseDurationMS: 30000}
|
||
}
|
||
func workerTestAssignment(r WorkerRegistration) WorkerAssignment {
|
||
return WorkerAssignment{Task: Task{ID: workerTestTask, NamespaceID: workerTestQueue, QueueID: workerTestQueue, State: TaskLeased, Revision: 2, Handler: r.Capabilities[0].Handler, HandlerVersion: "v1", PayloadVersion: "v1", Payload: json.RawMessage(`{"amount":9007199254740993}`)}, AttemptID: workerTestAttempt, Epoch: 1, LeaseSecret: r.InstanceSecret, ServerTime: time.Now().UTC(), LeaseUntil: time.Now().UTC().Add(30 * time.Second)}
|
||
}
|
||
func workerTestRead(t *testing.T, ctx context.Context, c *websocket.Conn) workerFrame {
|
||
t.Helper()
|
||
kind, raw, err := c.Read(ctx)
|
||
if err != nil {
|
||
t.Errorf("read test frame: %v", err)
|
||
return workerFrame{}
|
||
}
|
||
var f workerFrame
|
||
if kind != websocket.MessageText || json.Unmarshal(raw, &f) != nil {
|
||
t.Error("invalid client frame")
|
||
}
|
||
return f
|
||
}
|
||
func workerTestWrite(t *testing.T, ctx context.Context, c *websocket.Conn, f workerFrame) {
|
||
t.Helper()
|
||
f.Version = "1"
|
||
f.SessionID = workerTestSession
|
||
f.MessageID = workerTestID(t)
|
||
raw, err := json.Marshal(f)
|
||
if err != nil {
|
||
t.Error(err)
|
||
return
|
||
}
|
||
if err := c.Write(ctx, websocket.MessageText, raw); err != nil {
|
||
t.Errorf("write test frame: %v", err)
|
||
}
|
||
}
|
||
func workerTestBody(t *testing.T, v any) json.RawMessage {
|
||
t.Helper()
|
||
raw, err := json.Marshal(v)
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
return raw
|
||
}
|
||
|
||
func workerMock(t *testing.T, r WorkerRegistration, serve func(context.Context, *websocket.Conn)) (*Client, *WorkerConnection) {
|
||
t.Helper()
|
||
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
||
t.Cleanup(cancel)
|
||
server := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
|
||
if req.URL.Path != "/v1/worker/connect" || req.Header.Get("Authorization") != "Bearer unit-test-token" || req.URL.RawQuery != "" {
|
||
t.Error("unexpected handshake")
|
||
}
|
||
c, err := websocket.Accept(w, req, &websocket.AcceptOptions{Subprotocols: []string{workerSubprotocol}})
|
||
if err != nil {
|
||
t.Error(err)
|
||
return
|
||
}
|
||
defer c.CloseNow()
|
||
f := workerTestRead(t, ctx, c)
|
||
var got WorkerRegistration
|
||
if f.Type != "worker.hello" || f.SessionID != "" || !validID(f.MessageID) || json.Unmarshal(f.Body, &got) != nil || got.InstanceSecret != r.InstanceSecret || got.InstanceID != r.InstanceID {
|
||
t.Error("invalid hello")
|
||
}
|
||
workerTestWrite(t, ctx, c, workerFrame{Type: "worker.welcome", ReplyTo: f.MessageID, Body: workerTestBody(t, workerTestWelcome())})
|
||
serve(ctx, c)
|
||
}))
|
||
t.Cleanup(func() { cancel(); server.Close() })
|
||
client, err := NewClient(server.URL, Options{Token: "unit-test-token", HTTPClient: server.Client()})
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
w, err := client.DialWorker(ctx, r)
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
t.Cleanup(func() { w.Close() })
|
||
return client, w
|
||
}
|
||
|
||
func TestWorkerProtocolMultiplexAndMonotonicLease(t *testing.T) {
|
||
r := workerTestRegistration(t)
|
||
a := workerTestAssignment(r)
|
||
_, w := workerMock(t, r, func(ctx context.Context, c *websocket.Conn) {
|
||
capacity := workerTestRead(t, ctx, c)
|
||
workerTestWrite(t, ctx, c, workerFrame{Type: "task.assign", Body: workerTestBody(t, a)})
|
||
workerTestWrite(t, ctx, c, workerFrame{Type: "ack", ReplyTo: capacity.MessageID, Body: workerTestBody(t, workerTestWelcome())})
|
||
first, second := workerTestRead(t, ctx, c), workerTestRead(t, ctx, c)
|
||
for _, f := range []workerFrame{second, first} {
|
||
// 模拟数据库墙钟相差 24 小时；SDK 只能根据时间差确定单调期限。
|
||
now := time.Now().UTC().Add(24 * time.Hour)
|
||
result := WorkerCommandResult{Receipt: WorkerReceipt{Operation: strings.TrimPrefix(f.Type, "task."), TaskState: TaskRunning, AttemptState: "RUNNING", TaskRevision: 9007199254740993, Disposition: "APPLIED", LeaseUntil: now.Add(30 * time.Second)}, Lease: WorkerLease{Active: true, State: "RUNNING", ServerTime: now, LeaseUntil: now.Add(30 * time.Second)}}
|
||
workerTestWrite(t, ctx, c, workerFrame{Type: "ack", ReplyTo: f.MessageID, Body: workerTestBody(t, result)})
|
||
}
|
||
<-ctx.Done()
|
||
})
|
||
ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
|
||
defer cancel()
|
||
if _, err := w.Capacity(ctx, workerTestID(t), WorkerCapacity{Sequence: 1, Capacity: 2}); err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
got, err := w.NextAssignment(ctx)
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
if string(got.Task.Payload) != string(a.Task.Payload) {
|
||
t.Fatal("payload precision changed")
|
||
}
|
||
var wg sync.WaitGroup
|
||
for range 2 {
|
||
wg.Add(1)
|
||
go func() {
|
||
defer wg.Done()
|
||
out, err := w.Start(ctx, workerTestID(t), a.Attempt())
|
||
if err != nil {
|
||
t.Error(err)
|
||
return
|
||
}
|
||
deadline, active := out.Lease.ExecutionDeadline()
|
||
remaining := time.Until(deadline)
|
||
if !active || remaining < 29*time.Second || remaining > 30*time.Second || out.Receipt.TaskRevision != 9007199254740993 {
|
||
t.Error("wrong ACK correlation, precise revision or monotonic deadline")
|
||
}
|
||
}()
|
||
}
|
||
wg.Wait()
|
||
}
|
||
|
||
func TestWorkerLateACKDoesNotAuthorize(t *testing.T) {
|
||
r := workerTestRegistration(t)
|
||
a := workerTestAssignment(r)
|
||
_, w := workerMock(t, r, func(ctx context.Context, c *websocket.Conn) {
|
||
f := workerTestRead(t, ctx, c)
|
||
now := time.Now().UTC()
|
||
result := WorkerCommandResult{Receipt: WorkerReceipt{Operation: "start", TaskState: TaskRunning, AttemptState: "RUNNING", TaskRevision: 3, Disposition: "APPLIED", LeaseUntil: now.Add(time.Hour)}, Lease: WorkerLease{State: "UNKNOWN", ServerTime: now, LeaseUntil: now.Add(-time.Second)}, Deduplicated: true}
|
||
workerTestWrite(t, ctx, c, workerFrame{Type: "ack", ReplyTo: f.MessageID, Body: workerTestBody(t, result)})
|
||
<-ctx.Done()
|
||
})
|
||
out, err := w.Start(context.Background(), workerTestID(t), a.Attempt())
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
if !out.Deduplicated {
|
||
t.Fatal("missing original receipt")
|
||
}
|
||
if _, active := out.Lease.ExecutionDeadline(); active {
|
||
t.Fatal("historical ACK authorized execution")
|
||
}
|
||
manual := WorkerLease{Active: true, State: "RUNNING", LeaseUntil: time.Now().Add(time.Minute), ServerTime: time.Now()}
|
||
if _, active := manual.ExecutionDeadline(); active {
|
||
t.Fatal("unconfirmed lease authorized execution")
|
||
}
|
||
}
|
||
|
||
func TestWorkerBoundedRequestAndCancellation(t *testing.T) {
|
||
r := workerTestRegistration(t)
|
||
received := make(chan workerFrame, 64)
|
||
_, w := workerMock(t, r, func(ctx context.Context, c *websocket.Conn) {
|
||
for {
|
||
_, raw, err := c.Read(ctx)
|
||
if err != nil {
|
||
return
|
||
}
|
||
var f workerFrame
|
||
if json.Unmarshal(raw, &f) != nil {
|
||
return
|
||
}
|
||
received <- f
|
||
}
|
||
})
|
||
ctx, cancel := context.WithCancel(context.Background())
|
||
errs := make(chan error, workerMaxPending)
|
||
id := workerTestID(t)
|
||
go func() { _, err := w.Capacity(ctx, id, WorkerCapacity{Sequence: 1, Capacity: 1}); errs <- err }()
|
||
<-received
|
||
if _, err := w.Capacity(ctx, id, WorkerCapacity{Sequence: 1, Capacity: 1}); err == nil {
|
||
t.Fatal("duplicate in-flight message accepted")
|
||
}
|
||
for i := 1; i < workerMaxPending; i++ {
|
||
go func(seq uint64) {
|
||
_, err := w.Capacity(ctx, workerTestID(t), WorkerCapacity{Sequence: seq, Capacity: 1})
|
||
errs <- err
|
||
}(uint64(i + 1))
|
||
}
|
||
for i := 1; i < workerMaxPending; i++ {
|
||
<-received
|
||
}
|
||
if _, err := w.Capacity(ctx, workerTestID(t), WorkerCapacity{Sequence: 100, Capacity: 1}); !errors.Is(err, ErrWorkerBusy) {
|
||
t.Fatalf("want bounded queue, got %v", err)
|
||
}
|
||
cancel()
|
||
for range workerMaxPending {
|
||
select {
|
||
case err := <-errs:
|
||
if err == nil {
|
||
t.Error("pending operation unexpectedly confirmed")
|
||
}
|
||
case <-time.After(time.Second):
|
||
t.Fatal("pending operation did not exit")
|
||
}
|
||
}
|
||
select {
|
||
case <-w.Done():
|
||
case <-time.After(time.Second):
|
||
t.Fatal("connection leaked on cancellation")
|
||
}
|
||
}
|
||
|
||
func TestWorkerNextAssignmentCancelAndBackpressure(t *testing.T) {
|
||
r := workerTestRegistration(t)
|
||
send := make(chan struct{})
|
||
_, w := workerMock(t, r, func(ctx context.Context, c *websocket.Conn) {
|
||
<-send
|
||
for i := 0; i < 9; i++ {
|
||
workerTestWrite(t, ctx, c, workerFrame{Type: "task.assign", Body: workerTestBody(t, workerTestAssignment(r))})
|
||
}
|
||
<-ctx.Done()
|
||
})
|
||
ctx, cancel := context.WithCancel(context.Background())
|
||
cancel()
|
||
if _, err := w.NextAssignment(ctx); !errors.Is(err, context.Canceled) {
|
||
t.Fatal(err)
|
||
}
|
||
if w.Err() != nil {
|
||
t.Fatal("cancelled wait closed connection")
|
||
}
|
||
close(send)
|
||
select {
|
||
case <-w.Done():
|
||
case <-time.After(time.Second):
|
||
t.Fatal("unbounded assignment reader")
|
||
}
|
||
if !errors.Is(w.Err(), ErrWorkerBackpressure) {
|
||
t.Fatal(w.Err())
|
||
}
|
||
if _, err := w.NextAssignment(context.Background()); !errors.Is(err, ErrWorkerBackpressure) {
|
||
t.Fatal("returned buffered work after close")
|
||
}
|
||
}
|
||
|
||
func TestWorkerRedirectDoesNotForwardCredential(t *testing.T) {
|
||
var reached atomic.Int32
|
||
destination := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { reached.Add(1) }))
|
||
defer destination.Close()
|
||
origin := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||
http.Redirect(w, r, destination.URL, http.StatusTemporaryRedirect)
|
||
}))
|
||
defer origin.Close()
|
||
c, _ := NewClient(origin.URL, Options{Token: "private-token"})
|
||
_, err := c.DialWorker(context.Background(), workerTestRegistration(t))
|
||
var httpErr *HTTPError
|
||
if !errors.As(err, &httpErr) || httpErr.StatusCode != 307 || reached.Load() != 0 {
|
||
t.Fatalf("redirect followed or wrong error: %v", err)
|
||
}
|
||
}
|
||
|
||
func TestWorkerCodecAndValidation(t *testing.T) {
|
||
for _, raw := range []string{`{"a":1,"a":2}`, `{"x":{"a":1,"a":2}}`, `{"x":"\ud800"}`, `{"x":"\udc00"}`, `{"x":"\ud800\u0041"}`, `{"x":"` + string([]byte{0xff}) + `"}`, strings.Repeat("[", 34) + "0" + strings.Repeat("]", 34), `{} {}`} {
|
||
if validateWorkerJSON([]byte(raw)) == nil {
|
||
t.Errorf("invalid JSON accepted: %q", raw)
|
||
}
|
||
}
|
||
for _, raw := range []string{`{"x":"\ud83d\ude00"}`, `{"x":"literal \\uD800 and \"quote"}`, `{"a":9007199254740993,"b":0.1234567890123456789012345,"c":[true,null]}`} {
|
||
if validateWorkerJSON([]byte(raw)) != nil {
|
||
t.Errorf("valid JSON rejected: %s", raw)
|
||
}
|
||
}
|
||
r := workerTestRegistration(t)
|
||
for _, v := range []any{r, workerTestAssignment(r), workerTestAssignment(r).Attempt(), WorkerFinish{Ref: workerTestAssignment(r).Attempt(), Result: json.RawMessage(`{"secret":"business-private"}`)}} {
|
||
for _, format := range []string{"%v", "%+v", "%#v"} {
|
||
out := fmt.Sprintf(format, v)
|
||
if strings.Contains(out, r.InstanceSecret) || strings.Contains(out, "business-private") {
|
||
t.Fatal("secret in formatted Worker value")
|
||
}
|
||
}
|
||
}
|
||
before := []WorkerCapability{{Handler: "b", HandlerVersion: "v1", PayloadVersion: "v1"}, {Handler: "a", HandlerVersion: "v1", PayloadVersion: "v1"}}
|
||
normalized, err := NewWorkerRegistration(workerTestQueue, 1, before)
|
||
if err != nil || normalized.Capabilities[0].Handler != "a" || before[0].Handler != "b" {
|
||
t.Fatal("normalization mutated caller capability slice")
|
||
}
|
||
if _, err := NewWorkerRegistration(workerTestQueue, 1, []WorkerCapability{before[0], before[0]}); err == nil {
|
||
t.Fatal("duplicate capability")
|
||
}
|
||
for _, window := range []time.Duration{0, -time.Second, 3 * time.Minute} {
|
||
lease := WorkerLease{Active: true, State: "RUNNING", ServerTime: time.Now()}
|
||
lease.LeaseUntil = lease.ServerTime.Add(window)
|
||
if lease.confirm(time.Now()) == nil {
|
||
t.Fatal("invalid active lease")
|
||
}
|
||
}
|
||
lease := WorkerLease{Active: true, State: "RUNNING", ServerTime: time.Now()}
|
||
lease.LeaseUntil = lease.ServerTime.Add(time.Second)
|
||
if err := lease.confirm(time.Now().Add(-time.Second)); err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
if _, ok := lease.ExecutionDeadline(); ok {
|
||
t.Fatal("request duration not deducted")
|
||
}
|
||
}
|
||
|
||
func TestWorkerFatalAndInvalidFrames(t *testing.T) {
|
||
for _, test := range []struct {
|
||
name string
|
||
frame workerFrame
|
||
want error
|
||
}{
|
||
{"revoked", workerFrame{Type: "error", Error: &WorkerError{Code: "UNAUTHENTICATED", Message: "credential expired"}}, &WorkerError{Code: "UNAUTHENTICATED"}},
|
||
{"unknown_reply", workerFrame{Type: "ack", ReplyTo: workerTestTask, Body: json.RawMessage(`{}`)}, ErrWorkerProtocol},
|
||
{"invalid_body", workerFrame{Type: "task.assign", Body: json.RawMessage(`{}`)}, ErrWorkerProtocol},
|
||
} {
|
||
t.Run(test.name, func(t *testing.T) {
|
||
r := workerTestRegistration(t)
|
||
_, w := workerMock(t, r, func(ctx context.Context, c *websocket.Conn) { workerTestWrite(t, ctx, c, test.frame); <-ctx.Done() })
|
||
select {
|
||
case <-w.Done():
|
||
case <-time.After(time.Second):
|
||
t.Fatal("invalid frame did not close")
|
||
}
|
||
var serverErr *WorkerError
|
||
if want, ok := test.want.(*WorkerError); ok {
|
||
if !errors.As(w.Err(), &serverErr) || serverErr.Code != want.Code {
|
||
t.Fatal(w.Err())
|
||
}
|
||
} else if !errors.Is(w.Err(), test.want) {
|
||
t.Fatal(w.Err())
|
||
}
|
||
})
|
||
}
|
||
}
|
||
|
||
func FuzzWorkerJSON(f *testing.F) {
|
||
for _, s := range []string{`{"x":"\ud83d\ude00"}`, `{"n":9007199254740993}`, `[true,null,{}]`, `{"a":1,"a":2}`} {
|
||
f.Add([]byte(s))
|
||
}
|
||
f.Fuzz(func(t *testing.T, raw []byte) {
|
||
if len(raw) > 512<<10 {
|
||
t.Skip()
|
||
}
|
||
if validateWorkerJSON(raw) == nil && !json.Valid(raw) {
|
||
t.Fatal("accepted invalid JSON")
|
||
}
|
||
})
|
||
}
|
||
|
||
func TestWorkerConnectionRejectsWrongSessionAndBinaryFrames(t *testing.T) {
|
||
for _, kind := range []websocket.MessageType{websocket.MessageText, websocket.MessageBinary} {
|
||
t.Run(fmt.Sprint(kind), func(t *testing.T) {
|
||
r := workerTestRegistration(t)
|
||
_, w := workerMock(t, r, func(ctx context.Context, c *websocket.Conn) {
|
||
f := workerFrame{Version: "1", Type: "task.assign", MessageID: workerTestID(t), SessionID: workerTestTask, Body: workerTestBody(t, workerTestAssignment(r))}
|
||
if kind == websocket.MessageBinary {
|
||
f.SessionID = workerTestSession
|
||
}
|
||
if err := c.Write(ctx, kind, workerTestBody(t, f)); err != nil {
|
||
t.Error(err)
|
||
}
|
||
<-ctx.Done()
|
||
})
|
||
select {
|
||
case <-w.Done():
|
||
case <-time.After(time.Second):
|
||
t.Fatal("invalid session or binary frame did not close")
|
||
}
|
||
if !errors.Is(w.Err(), ErrWorkerProtocol) {
|
||
t.Fatal(w.Err())
|
||
}
|
||
})
|
||
}
|
||
}
|
||
|
||
func TestWorkerHandshakeRequiresSubprotocol(t *testing.T) {
|
||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||
c, err := websocket.Accept(w, r, nil)
|
||
if err != nil {
|
||
return
|
||
}
|
||
defer c.CloseNow()
|
||
ctx, cancel := context.WithTimeout(r.Context(), time.Second)
|
||
defer cancel()
|
||
_, _, _ = c.Read(ctx)
|
||
}))
|
||
defer server.Close()
|
||
c, _ := NewClient(server.URL, Options{Token: "test-token"})
|
||
if _, err := c.DialWorker(context.Background(), workerTestRegistration(t)); !errors.Is(err, ErrWorkerProtocol) {
|
||
t.Fatalf("missing subprotocol accepted: %v", err)
|
||
}
|
||
}
|