orbit/sdk/go/worker_connection_test.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)
	}
}