feat: add standalone Worker protocol SDK and ACK-loss regression
This commit is contained in:
@@ -1,4 +1,4 @@
|
||||
// Package orbit is the standalone Go SDK for Orbit. This foundation release
|
||||
// implements health calls; task, Worker and workflow APIs are planned next.
|
||||
// The SDK does not import server models, migrations or storage dependencies.
|
||||
// Package orbit 是 Orbit 的独立 Go SDK，提供健康、任务、队列与 Worker 底层协议。
|
||||
// WorkerConnection 不自动重连、续租或运行处理器；自动运行器与流程 SDK 仍在开发。
|
||||
// SDK 不导入服务端模型、迁移或数据库依赖。
|
||||
package orbit
|
||||
|
||||
214
sdk/go/examples/worker-protocol/main.go
Normal file
214
sdk/go/examples/worker-protocol/main.go
Normal file
@@ -0,0 +1,214 @@
|
||||
// worker-protocol 演示一次文本规范化任务的真实执行与原消息重报。
|
||||
// 这是底层协议示例；持续运行、自动续租与优雅退出由后续 Worker 运行器提供。
|
||||
package main
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"os"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
orbit "git.yuntuo.hk/k2safe/orbit/sdk/go"
|
||||
)
|
||||
|
||||
func main() {
|
||||
if err := run(); err != nil {
|
||||
fmt.Fprintln(os.Stderr, err)
|
||||
os.Exit(1)
|
||||
}
|
||||
}
|
||||
|
||||
func run() error {
|
||||
token, queue := os.Getenv("ORBIT_TOKEN"), os.Getenv("ORBIT_QUEUE_ID")
|
||||
if path := os.Getenv("ORBIT_CREDENTIALS_FILE"); path != "" {
|
||||
raw, err := os.ReadFile(path)
|
||||
if err != nil {
|
||||
return fmt.Errorf("read private credentials: %w", err)
|
||||
}
|
||||
var credentials struct {
|
||||
Token string `json:"token"`
|
||||
QueueID string `json:"queue_id"`
|
||||
}
|
||||
if json.Unmarshal(raw, &credentials) != nil {
|
||||
return errors.New("invalid credentials file")
|
||||
}
|
||||
if token == "" {
|
||||
token = credentials.Token
|
||||
}
|
||||
if queue == "" {
|
||||
queue = credentials.QueueID
|
||||
}
|
||||
}
|
||||
client, err := orbit.NewClient(os.Getenv("ORBIT_URL"), orbit.Options{Token: token})
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
// 每次运行使用独立能力名，避免接走前一次示例遗留的任务。
|
||||
operation, err := orbit.NewWorkerMessageID()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
handler := "text.normalize." + operation
|
||||
registration, err := orbit.NewWorkerRegistration(queue, 1, []orbit.WorkerCapability{{Handler: handler, HandlerVersion: "v1", PayloadVersion: "v1"}})
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 25*time.Second)
|
||||
defer cancel()
|
||||
accepted, err := client.SubmitTask(ctx, "normalize-"+operation, orbit.SubmitTaskRequest{QueueID: queue, Handler: handler, HandlerVersion: "v1", PayloadVersion: "v1", Payload: json.RawMessage(`{"text":" Hello Orbit "}`), RecoverySafety: orbit.RecoveryIdempotent, TaskTimeoutMS: 60000, Retry: orbit.RetryPolicy{MaxAttempts: 1, AttemptTimeoutMS: 20000}})
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
session := &demoSession{client: client, registration: registration}
|
||||
defer session.close()
|
||||
if err := session.connect(ctx); err != nil {
|
||||
return err
|
||||
}
|
||||
capacityID, err := orbit.NewWorkerMessageID()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if _, err := session.conn.Capacity(ctx, capacityID, orbit.WorkerCapacity{Sequence: 1, Capacity: 1}); err != nil {
|
||||
return err
|
||||
}
|
||||
assignment, err := session.conn.NextAssignment(ctx)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if assignment.Task.ID != accepted.Task.ID {
|
||||
return errors.New("unexpected task in isolated example")
|
||||
}
|
||||
ref := assignment.Attempt()
|
||||
session.ref = &ref
|
||||
startID, err := orbit.NewWorkerMessageID()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
started, err := session.confirm(ctx, func(c *orbit.WorkerConnection) (orbit.WorkerCommandResult, error) { return c.Start(ctx, startID, ref) })
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
deadline, active := started.Lease.ExecutionDeadline()
|
||||
if !active {
|
||||
return errors.New("start receipt does not grant current execution authority")
|
||||
}
|
||||
// 业务代码只拿任务输入；不接收注册秘密或尝试秘密。处理器只调用这一次。
|
||||
work, stop := context.WithDeadline(ctx, deadline)
|
||||
result, handlerErr := normalize(work, assignment.Task.Payload)
|
||||
stop()
|
||||
finish := orbit.WorkerFinish{Ref: ref, Outcome: orbit.WorkerSucceeded, Result: result}
|
||||
if handlerErr != nil {
|
||||
// 普通异常不能证明外部效果未发生；默认报告 UNKNOWN，保留核查路径。
|
||||
finish.Outcome = orbit.WorkerUnknown
|
||||
finish.ErrorCode = "HANDLER_UNCONFIRMED"
|
||||
finish.Result = json.RawMessage(`{"reason":"handler did not confirm completion"}`)
|
||||
}
|
||||
finishID, err := orbit.NewWorkerMessageID()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
confirmed, err := session.confirm(ctx, func(c *orbit.WorkerConnection) (orbit.WorkerCommandResult, error) {
|
||||
return c.Finish(ctx, finishID, finish)
|
||||
})
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if confirmed.Receipt.Disposition != "APPLIED" {
|
||||
return errors.New("result saved as evidence only; inspect authoritative task state")
|
||||
}
|
||||
current, err := client.GetTask(ctx, assignment.Task.ID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if current.State != orbit.TaskSucceeded || current.AttemptCount != 1 || string(current.Result) != string(result) {
|
||||
return errors.New("task did not complete as expected")
|
||||
}
|
||||
drainID, err := orbit.NewWorkerMessageID()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
// 新连接尚未开启容量时，Sequence=2 仍是合法递增值。
|
||||
if _, err := session.conn.Capacity(ctx, drainID, orbit.WorkerCapacity{Sequence: 2, Capacity: 0, Drain: true}); err != nil {
|
||||
return err
|
||||
}
|
||||
fmt.Printf("task=%s final_state=%s attempts=%d result=%s\n", current.ID, current.State, current.AttemptCount, current.Result)
|
||||
return nil
|
||||
}
|
||||
|
||||
func normalize(ctx context.Context, payload json.RawMessage) (json.RawMessage, error) {
|
||||
if err := ctx.Err(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
var in struct {
|
||||
Text string `json:"text"`
|
||||
}
|
||||
if err := json.Unmarshal(payload, &in); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return json.Marshal(struct {
|
||||
Text string `json:"text"`
|
||||
}{Text: strings.ToUpper(strings.TrimSpace(in.Text))})
|
||||
}
|
||||
|
||||
type demoSession struct {
|
||||
client *orbit.Client
|
||||
registration orbit.WorkerRegistration
|
||||
conn *orbit.WorkerConnection
|
||||
ref *orbit.WorkerAttempt
|
||||
}
|
||||
|
||||
func (s *demoSession) close() {
|
||||
if s.conn != nil {
|
||||
s.conn.Close()
|
||||
}
|
||||
}
|
||||
func (s *demoSession) connect(ctx context.Context) error {
|
||||
s.close()
|
||||
c, err := s.client.DialWorker(ctx, s.registration)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
s.conn = c
|
||||
if s.ref != nil {
|
||||
id, err := orbit.NewWorkerMessageID()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
// 即使租约已失效，也可以重报原结果取得原回执或保存迟到证据。
|
||||
if _, err := c.Resume(ctx, id, *s.ref); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// confirm 只重报控制操作，不再次调用业务函数；同一操作闭包固定消息 ID 与内容。
|
||||
func (s *demoSession) confirm(ctx context.Context, call func(*orbit.WorkerConnection) (orbit.WorkerCommandResult, error)) (orbit.WorkerCommandResult, error) {
|
||||
var out orbit.WorkerCommandResult
|
||||
var err error
|
||||
for attempt := 0; attempt < 3; attempt++ {
|
||||
if attempt > 0 {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return out, ctx.Err()
|
||||
case <-time.After(time.Duration(attempt) * 100 * time.Millisecond):
|
||||
}
|
||||
if err = s.connect(ctx); err != nil {
|
||||
return out, err
|
||||
}
|
||||
}
|
||||
out, err = call(s.conn)
|
||||
if err == nil {
|
||||
return out, nil
|
||||
}
|
||||
var control *orbit.WorkerError
|
||||
retry := errors.Is(err, orbit.ErrWorkerDisconnected) || errors.Is(err, context.DeadlineExceeded) && ctx.Err() == nil || errors.As(err, &control) && control.Code == "STALE_SESSION"
|
||||
if !retry {
|
||||
return out, err
|
||||
}
|
||||
}
|
||||
return out, fmt.Errorf("control result still unconfirmed after bounded replay: %w", err)
|
||||
}
|
||||
@@ -1,3 +1,5 @@
|
||||
module git.yuntuo.hk/k2safe/orbit/sdk/go
|
||||
|
||||
go 1.27.1
|
||||
|
||||
require github.com/coder/websocket v1.8.15
|
||||
|
||||
2
sdk/go/go.sum
Normal file
2
sdk/go/go.sum
Normal file
@@ -0,0 +1,2 @@
|
||||
github.com/coder/websocket v1.8.15 h1:6B2JPeOGlpff2Uz6vOEH1Vzpi0iUz20A+lPVhPHtNUA=
|
||||
github.com/coder/websocket v1.8.15/go.mod h1:NX3SzP+inril6yawo5CQXx8+fk145lPDC6pumgx0mVg=
|
||||
138
sdk/go/worker_codec.go
Normal file
138
sdk/go/worker_codec.go
Normal file
@@ -0,0 +1,138 @@
|
||||
package orbit
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"io"
|
||||
"strconv"
|
||||
"unicode/utf8"
|
||||
|
||||
"github.com/coder/websocket"
|
||||
)
|
||||
|
||||
func parseWorkerFrame(kind websocket.MessageType, raw []byte) (workerFrame, error) {
|
||||
var f workerFrame
|
||||
if kind != websocket.MessageText || len(raw) > 512<<10 || validateWorkerJSON(raw) != nil {
|
||||
return f, ErrWorkerProtocol
|
||||
}
|
||||
// Envelope 字段严格区分大小写，防止 encoding/json 的大小写别名覆盖控制字段。
|
||||
var fields map[string]json.RawMessage
|
||||
if json.Unmarshal(raw, &fields) != nil || fields == nil {
|
||||
return f, ErrWorkerProtocol
|
||||
}
|
||||
for k := range fields {
|
||||
switch k {
|
||||
case "protocol_version", "type", "message_id", "reply_to", "session_id", "body", "error":
|
||||
default:
|
||||
return f, ErrWorkerProtocol
|
||||
}
|
||||
}
|
||||
if json.Unmarshal(raw, &f) != nil || f.Version != "1" || !validID(f.MessageID) || f.ReplyTo != "" && !validID(f.ReplyTo) || f.SessionID != "" && !validID(f.SessionID) {
|
||||
return f, ErrWorkerProtocol
|
||||
}
|
||||
if f.Type == "error" {
|
||||
if f.Error == nil || !workerName(f.Error.Code, 64) || len(f.Error.Message) > 1024 || len(f.Body) > 0 {
|
||||
return f, ErrWorkerProtocol
|
||||
}
|
||||
} else if f.Error != nil || len(f.Body) == 0 || bytes.Equal(f.Body, []byte("null")) {
|
||||
return f, ErrWorkerProtocol
|
||||
}
|
||||
return f, nil
|
||||
}
|
||||
|
||||
// validateWorkerJSON 拒绝重复键、过深嵌套及非法 Unicode，保留数字原文。
|
||||
// 数值的业务范围仍由服务端规范化器最终验证；SDK 不经由 float64 重编码。
|
||||
func validateWorkerJSON(raw []byte) error {
|
||||
if len(raw) == 0 || !utf8.Valid(raw) || !json.Valid(raw) || !workerUnicode(raw) {
|
||||
return ErrWorkerProtocol
|
||||
}
|
||||
d := json.NewDecoder(bytes.NewReader(raw))
|
||||
d.UseNumber()
|
||||
var value func(int) error
|
||||
value = func(depth int) error {
|
||||
if depth > 32 {
|
||||
return ErrWorkerProtocol
|
||||
}
|
||||
t, err := d.Token()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
switch t {
|
||||
case json.Delim('{'):
|
||||
keys := map[string]bool{}
|
||||
for d.More() {
|
||||
k, err := d.Token()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
s, ok := k.(string)
|
||||
if !ok || keys[s] {
|
||||
return ErrWorkerProtocol
|
||||
}
|
||||
keys[s] = true
|
||||
if err := value(depth + 1); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
_, err = d.Token()
|
||||
return err
|
||||
case json.Delim('['):
|
||||
for d.More() {
|
||||
if err := value(depth + 1); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
_, err = d.Token()
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}
|
||||
if err := value(0); err != nil {
|
||||
return err
|
||||
}
|
||||
if _, err := d.Token(); err != io.EOF {
|
||||
return ErrWorkerProtocol
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// JSON 的未配对 UTF-16 转义不能被静默替换为 U+FFFD。
|
||||
func workerUnicode(raw []byte) bool {
|
||||
inside := false
|
||||
for i := 0; i < len(raw); i++ {
|
||||
if raw[i] == '"' {
|
||||
inside = !inside
|
||||
continue
|
||||
}
|
||||
if !inside || raw[i] != '\\' {
|
||||
continue
|
||||
}
|
||||
i++
|
||||
if i >= len(raw) {
|
||||
return false
|
||||
}
|
||||
if raw[i] != 'u' {
|
||||
continue
|
||||
}
|
||||
if i+4 >= len(raw) {
|
||||
return false
|
||||
}
|
||||
n, err := strconv.ParseUint(string(raw[i+1:i+5]), 16, 16)
|
||||
if err != nil {
|
||||
return false
|
||||
}
|
||||
i += 4
|
||||
if n < 0xd800 || n > 0xdfff {
|
||||
continue
|
||||
}
|
||||
if n > 0xdbff || i+6 >= len(raw) || raw[i+1] != '\\' || raw[i+2] != 'u' {
|
||||
return false
|
||||
}
|
||||
low, err := strconv.ParseUint(string(raw[i+3:i+7]), 16, 16)
|
||||
if err != nil || low < 0xdc00 || low > 0xdfff {
|
||||
return false
|
||||
}
|
||||
i += 6
|
||||
}
|
||||
return !inside
|
||||
}
|
||||
418
sdk/go/worker_connection.go
Normal file
418
sdk/go/worker_connection.go
Normal file
@@ -0,0 +1,418 @@
|
||||
package orbit
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/coder/websocket"
|
||||
)
|
||||
|
||||
const workerSubprotocol = "pulse.worker.v1"
|
||||
const workerMaxPending = 32
|
||||
const workerSafetyMargin = 250 * time.Millisecond
|
||||
|
||||
var (
|
||||
ErrWorkerClosed = errors.New("Worker connection closed")
|
||||
ErrWorkerDisconnected = errors.New("Worker connection interrupted; reconcile before replay")
|
||||
ErrWorkerProtocol = errors.New("invalid Worker protocol response")
|
||||
ErrWorkerBusy = errors.New("Worker control request limit reached")
|
||||
ErrWorkerBackpressure = errors.New("Worker assignment buffer exhausted")
|
||||
)
|
||||
|
||||
// WorkerError 是服务端公开的控制错误；不会把错误解释为业务没有生效。
|
||||
type WorkerError struct {
|
||||
Code string `json:"code"` // Code 稳定分类，例如 LEASE_LOST、STALE_SESSION。
|
||||
Message string `json:"message"` // Message 公开说明，不应包含内部异常正文。
|
||||
}
|
||||
|
||||
func (e *WorkerError) Error() string { return "orbit Worker " + e.Code }
|
||||
|
||||
type workerFrame struct {
|
||||
Version string `json:"protocol_version"`
|
||||
Type string `json:"type"`
|
||||
MessageID string `json:"message_id"`
|
||||
ReplyTo string `json:"reply_to,omitempty"`
|
||||
SessionID string `json:"session_id,omitempty"`
|
||||
Body json.RawMessage `json:"body,omitempty"`
|
||||
Error *WorkerError `json:"error,omitempty"`
|
||||
}
|
||||
|
||||
// WorkerConnection 是一次连接的底层协议客户端，可并发调用控制方法。
|
||||
// 本类型不自动重连、不自动续租、不运行处理器。关闭连接不等于取消业务效果。
|
||||
// 最多 32 个在途控制请求、8 个未读派发；超时关闭连接以便按原消息身份核对。
|
||||
type WorkerConnection struct {
|
||||
conn *websocket.Conn
|
||||
ctx context.Context
|
||||
cancel context.CancelFunc
|
||||
welcome WorkerWelcome
|
||||
registration WorkerRegistration
|
||||
assignments chan WorkerAssignment
|
||||
done chan struct{}
|
||||
mu sync.Mutex
|
||||
err error
|
||||
pending map[string]chan workerFrame
|
||||
}
|
||||
|
||||
// DialWorker 使用 Authorization 和固定子协议连接；ctx 只控制握手。
|
||||
// 返回后调用方必须 Close。重连使用同一 registration，再逐个 Resume 旧尝试，
|
||||
// 完成核对后才发送第一次 Capacity 开启新派发。
|
||||
func (c *Client) DialWorker(ctx context.Context, registration WorkerRegistration) (*WorkerConnection, error) {
|
||||
if c.token == "" {
|
||||
return nil, errors.New("Orbit token is required for Worker APIs")
|
||||
}
|
||||
r, err := normalizeRegistration(registration)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
handshake, stop := context.WithTimeout(ctx, 8*time.Second)
|
||||
defer stop()
|
||||
header := http.Header{}
|
||||
header.Set("Authorization", "Bearer "+c.token)
|
||||
conn, response, err := websocket.Dial(handshake, c.baseURL+"/v1/worker/connect", &websocket.DialOptions{HTTPClient: c.http, HTTPHeader: header, Subprotocols: []string{workerSubprotocol}, CompressionMode: websocket.CompressionDisabled})
|
||||
if err != nil {
|
||||
if response != nil && response.StatusCode != http.StatusSwitchingProtocols {
|
||||
return nil, &HTTPError{StatusCode: response.StatusCode}
|
||||
}
|
||||
return nil, workerIOError(err)
|
||||
}
|
||||
ok := false
|
||||
defer func() {
|
||||
if !ok {
|
||||
conn.CloseNow()
|
||||
}
|
||||
}()
|
||||
if conn.Subprotocol() != workerSubprotocol {
|
||||
return nil, ErrWorkerProtocol
|
||||
}
|
||||
conn.SetReadLimit(512 << 10)
|
||||
id, err := NewWorkerMessageID()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
body, _ := json.Marshal(r)
|
||||
frame, _ := json.Marshal(workerFrame{Version: "1", Type: "worker.hello", MessageID: id, Body: body})
|
||||
if err = conn.Write(handshake, websocket.MessageText, frame); err != nil {
|
||||
return nil, workerIOError(err)
|
||||
}
|
||||
kind, raw, err := conn.Read(handshake)
|
||||
if err != nil {
|
||||
return nil, workerIOError(err)
|
||||
}
|
||||
f, err := parseWorkerFrame(kind, raw)
|
||||
if err != nil || f.ReplyTo != id {
|
||||
return nil, ErrWorkerProtocol
|
||||
}
|
||||
if f.Type == "error" {
|
||||
return nil, f.Error
|
||||
}
|
||||
var welcome WorkerWelcome
|
||||
if f.Type != "worker.welcome" || json.Unmarshal(f.Body, &welcome) != nil || !validWelcome(welcome) || welcome.SessionID != f.SessionID {
|
||||
return nil, ErrWorkerProtocol
|
||||
}
|
||||
lifetime, cancel := context.WithCancel(context.Background())
|
||||
w := &WorkerConnection{conn: conn, ctx: lifetime, cancel: cancel, welcome: welcome, registration: r, assignments: make(chan WorkerAssignment, 8), done: make(chan struct{}), pending: make(map[string]chan workerFrame)}
|
||||
ok = true
|
||||
go w.readLoop()
|
||||
return w, nil
|
||||
}
|
||||
|
||||
// Welcome 返回本次连接最初的确认快照；最新心跳确认由 Capacity 返回。
|
||||
func (w *WorkerConnection) Welcome() WorkerWelcome { return w.welcome }
|
||||
|
||||
// Done 在读取循环退出并关闭底层连接后关闭；可用于触发重新连接。
|
||||
func (w *WorkerConnection) Done() <-chan struct{} { return w.done }
|
||||
|
||||
// Err 返回连接关闭的原因；未关闭时返回 nil，不含帧正文或秘密。
|
||||
func (w *WorkerConnection) Err() error { w.mu.Lock(); defer w.mu.Unlock(); return w.err }
|
||||
|
||||
// Close 立即关闭传输并等待读取循环退出；排空需先调用 Capacity(Drain:true)。
|
||||
func (w *WorkerConnection) Close() error { w.fail(ErrWorkerClosed); <-w.done; return nil }
|
||||
|
||||
func (w *WorkerConnection) fail(err error) {
|
||||
w.mu.Lock()
|
||||
if w.err != nil {
|
||||
w.mu.Unlock()
|
||||
return
|
||||
}
|
||||
w.err = err
|
||||
w.cancel()
|
||||
w.mu.Unlock()
|
||||
_ = w.conn.CloseNow()
|
||||
}
|
||||
|
||||
func (w *WorkerConnection) readLoop() {
|
||||
defer close(w.done)
|
||||
for {
|
||||
read, stop := context.WithTimeout(w.ctx, 60*time.Second)
|
||||
kind, raw, err := w.conn.Read(read)
|
||||
stop()
|
||||
if err != nil {
|
||||
w.fail(workerIOError(err))
|
||||
return
|
||||
}
|
||||
f, err := parseWorkerFrame(kind, raw)
|
||||
if err != nil || f.SessionID != w.welcome.SessionID {
|
||||
w.fail(ErrWorkerProtocol)
|
||||
return
|
||||
}
|
||||
if f.Type == "task.assign" {
|
||||
var a WorkerAssignment
|
||||
if f.ReplyTo != "" || json.Unmarshal(f.Body, &a) != nil || !w.validAssignment(a) {
|
||||
w.fail(ErrWorkerProtocol)
|
||||
return
|
||||
}
|
||||
select {
|
||||
case w.assignments <- a:
|
||||
default:
|
||||
w.fail(ErrWorkerBackpressure)
|
||||
return
|
||||
}
|
||||
continue
|
||||
}
|
||||
if f.Type != "ack" && f.Type != "error" {
|
||||
w.fail(ErrWorkerProtocol)
|
||||
return
|
||||
}
|
||||
if f.ReplyTo == "" && f.Type == "error" {
|
||||
w.fail(f.Error)
|
||||
return
|
||||
}
|
||||
w.mu.Lock()
|
||||
waiter := w.pending[f.ReplyTo]
|
||||
w.mu.Unlock()
|
||||
if waiter == nil {
|
||||
w.fail(ErrWorkerProtocol)
|
||||
return
|
||||
}
|
||||
select {
|
||||
case waiter <- f:
|
||||
default:
|
||||
w.fail(ErrWorkerProtocol)
|
||||
return
|
||||
}
|
||||
if f.Type == "error" && (f.Error.Code == "STALE_SESSION" || f.Error.Code == "UNAUTHENTICATED" || f.Error.Code == "PERMISSION_DENIED") {
|
||||
w.fail(f.Error)
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// NextAssignment 等待预占通知；即使返回成功，也必须通过 Start 确认当前执行权。
|
||||
// ctx 到期只中断等待，不关闭连接；读取过慢触发有界缓冲保护时连接会关闭。
|
||||
func (w *WorkerConnection) NextAssignment(ctx context.Context) (WorkerAssignment, error) {
|
||||
if err := ctx.Err(); err != nil {
|
||||
return WorkerAssignment{}, err
|
||||
}
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return WorkerAssignment{}, ctx.Err()
|
||||
case <-w.ctx.Done():
|
||||
return WorkerAssignment{}, w.Err()
|
||||
case a := <-w.assignments:
|
||||
if err := w.Err(); err != nil {
|
||||
return WorkerAssignment{}, err
|
||||
}
|
||||
return a, nil
|
||||
}
|
||||
}
|
||||
|
||||
func (w *WorkerConnection) call(ctx context.Context, kind, id string, in, out any) (time.Time, error) {
|
||||
started := time.Now()
|
||||
if !validID(id) {
|
||||
return started, errors.New("invalid Worker message id")
|
||||
}
|
||||
if err := ctx.Err(); err != nil {
|
||||
return started, err
|
||||
}
|
||||
waiter := make(chan workerFrame, 1)
|
||||
w.mu.Lock()
|
||||
if w.err != nil {
|
||||
err := w.err
|
||||
w.mu.Unlock()
|
||||
return started, err
|
||||
}
|
||||
if len(w.pending) >= workerMaxPending {
|
||||
w.mu.Unlock()
|
||||
return started, ErrWorkerBusy
|
||||
}
|
||||
if _, ok := w.pending[id]; ok {
|
||||
w.mu.Unlock()
|
||||
return started, errors.New("Worker message id is already in flight")
|
||||
}
|
||||
w.pending[id] = waiter
|
||||
w.mu.Unlock()
|
||||
defer func() { w.mu.Lock(); delete(w.pending, id); w.mu.Unlock() }()
|
||||
// 先占用有界槽再分配编码缓冲，超额调用不先复制大结果。
|
||||
body, err := json.Marshal(in)
|
||||
if err != nil {
|
||||
return started, errors.New("invalid Worker request JSON")
|
||||
}
|
||||
raw, err := json.Marshal(workerFrame{Version: "1", Type: kind, MessageID: id, SessionID: w.welcome.SessionID, Body: body})
|
||||
if err != nil || len(raw) > 512<<10 {
|
||||
return started, errors.New("Worker request exceeds frame limit")
|
||||
}
|
||||
bounded, stop := context.WithDeadline(ctx, started.Add(10*time.Second))
|
||||
defer stop()
|
||||
// coder/websocket 串行化并发 Write；排队时间包含在操作总时限和租约余量中。
|
||||
if err = w.conn.Write(bounded, websocket.MessageText, raw); err != nil {
|
||||
err = workerIOError(err)
|
||||
w.fail(err)
|
||||
return started, err
|
||||
}
|
||||
var f workerFrame
|
||||
select {
|
||||
case f = <-waiter:
|
||||
case <-w.ctx.Done():
|
||||
// 致命控制错误先交付给相应调用，再关闭连接。
|
||||
select {
|
||||
case f = <-waiter:
|
||||
default:
|
||||
return started, w.Err()
|
||||
}
|
||||
case <-bounded.Done():
|
||||
w.fail(bounded.Err())
|
||||
return started, bounded.Err()
|
||||
}
|
||||
if f.Error != nil {
|
||||
return started, f.Error
|
||||
}
|
||||
if json.Unmarshal(f.Body, out) != nil {
|
||||
w.fail(ErrWorkerProtocol)
|
||||
return started, ErrWorkerProtocol
|
||||
}
|
||||
return started, nil
|
||||
}
|
||||
|
||||
// Capacity 确认绝对容量，同时续期会话。消息身份和 Sequence 在重报时保持不变。
|
||||
func (w *WorkerConnection) Capacity(ctx context.Context, id string, in WorkerCapacity) (WorkerWelcome, error) {
|
||||
var out WorkerWelcome
|
||||
if in.Sequence == 0 || in.Capacity > 256 {
|
||||
return out, errors.New("invalid Worker capacity or sequence")
|
||||
}
|
||||
_, err := w.call(ctx, "worker.capacity", id, in, &out)
|
||||
if err == nil && (!validWelcome(out) || out.SessionID != w.welcome.SessionID || out.Epoch != w.welcome.Epoch) {
|
||||
err = ErrWorkerProtocol
|
||||
w.fail(err)
|
||||
}
|
||||
return out, err
|
||||
}
|
||||
|
||||
// Resume 核对原尝试并迁移到当前会话，不创建尝试、不续期，也不授权再次运行处理器。
|
||||
func (w *WorkerConnection) Resume(ctx context.Context, id string, ref WorkerAttempt) (WorkerLease, error) {
|
||||
var out WorkerLease
|
||||
if !validWorkerAttempt(ref) {
|
||||
return out, errors.New("invalid Worker attempt proof")
|
||||
}
|
||||
started, err := w.call(ctx, "worker.resume", id, ref, &out)
|
||||
if err == nil {
|
||||
err = out.confirm(started)
|
||||
if err != nil {
|
||||
w.fail(err)
|
||||
}
|
||||
}
|
||||
return out, err
|
||||
}
|
||||
|
||||
// Start 持久化开始授权；只在返回 Lease.ExecutionDeadline 为真时运行一次处理器。
|
||||
func (w *WorkerConnection) Start(ctx context.Context, id string, ref WorkerAttempt) (WorkerCommandResult, error) {
|
||||
return w.command(ctx, "start", id, ref, ref)
|
||||
}
|
||||
|
||||
// Renew 续期已运行尝试；新续租用新消息身份，响应不明时原身份只核对原操作。
|
||||
func (w *WorkerConnection) Renew(ctx context.Context, id string, ref WorkerAttempt) (WorkerCommandResult, error) {
|
||||
return w.command(ctx, "renew", id, ref, ref)
|
||||
}
|
||||
|
||||
// Finish 保存业务结果/证据；迟到结果可返回 EVIDENCE_ONLY，不能视为任务成功终结。
|
||||
func (w *WorkerConnection) Finish(ctx context.Context, id string, in WorkerFinish) (WorkerCommandResult, error) {
|
||||
var out WorkerCommandResult
|
||||
if (in.Outcome != WorkerSucceeded && in.Outcome != WorkerFailed && in.Outcome != WorkerUnknown) || len(in.Result) > 256<<10 || bytes.Equal(bytes.TrimSpace(in.Result), []byte("null")) || validateWorkerJSON(in.Result) != nil {
|
||||
return out, errors.New("invalid Worker result outcome or JSON")
|
||||
}
|
||||
if in.Outcome == WorkerSucceeded && in.ErrorCode != "" || in.Outcome != WorkerSucceeded && !workerName(in.ErrorCode, 64) {
|
||||
return out, errors.New("Worker error code must match outcome")
|
||||
}
|
||||
return w.command(ctx, "finish", id, in.Ref, in)
|
||||
}
|
||||
func (w *WorkerConnection) command(ctx context.Context, op, id string, ref WorkerAttempt, in any) (WorkerCommandResult, error) {
|
||||
var out WorkerCommandResult
|
||||
if !validWorkerAttempt(ref) {
|
||||
return out, errors.New("invalid Worker attempt proof")
|
||||
}
|
||||
started, err := w.call(ctx, "task."+op, id, in, &out)
|
||||
if err == nil && (out.Receipt.Operation != op || out.Receipt.TaskRevision == 0 || !workerTaskState(out.Receipt.TaskState) || !workerAttemptState(out.Receipt.AttemptState) || out.Receipt.Disposition != "APPLIED" && out.Receipt.Disposition != "EVIDENCE_ONLY") {
|
||||
err = ErrWorkerProtocol
|
||||
}
|
||||
if err == nil {
|
||||
err = out.Lease.confirm(started)
|
||||
}
|
||||
if errors.Is(err, ErrWorkerProtocol) {
|
||||
w.fail(err)
|
||||
}
|
||||
return out, err
|
||||
}
|
||||
|
||||
func (w *WorkerConnection) validAssignment(a WorkerAssignment) bool {
|
||||
if !validWorkerAttempt(a.Attempt()) || a.Task.QueueID != w.registration.QueueID || !validID(a.Task.NamespaceID) || a.Task.State != TaskLeased || a.Task.Revision == 0 || a.LeaseUntil.IsZero() || a.ServerTime.IsZero() || len(a.Task.Payload) == 0 || bytes.Equal(bytes.TrimSpace(a.Task.Payload), []byte("null")) {
|
||||
return false
|
||||
}
|
||||
for _, c := range w.registration.Capabilities {
|
||||
if c.Handler == a.Task.Handler && c.HandlerVersion == a.Task.HandlerVersion && c.PayloadVersion == a.Task.PayloadVersion {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
func validWelcome(w WorkerWelcome) bool {
|
||||
return validID(w.SessionID) && w.Epoch > 0 && !w.ServerTime.IsZero() && w.ExpiresAt.After(w.ServerTime) && w.ExpiresAt.Sub(w.ServerTime) <= 5*time.Minute && w.LeaseDurationMS >= 1000 && w.LeaseDurationMS <= 120000
|
||||
}
|
||||
func (l *WorkerLease) confirm(started time.Time) error {
|
||||
l.deadline = time.Time{}
|
||||
if l.ServerTime.IsZero() || !workerAttemptState(l.State) {
|
||||
return ErrWorkerProtocol
|
||||
}
|
||||
if !l.Active {
|
||||
return nil
|
||||
}
|
||||
window := l.LeaseUntil.Sub(l.ServerTime)
|
||||
if (l.State != "LEASED" && l.State != "RUNNING") || window <= 0 || window > 2*time.Minute {
|
||||
return ErrWorkerProtocol
|
||||
}
|
||||
// 整个往返（包含 SDK 排队）都扣除；只使用服务端时间差，不比较两端墙钟。
|
||||
l.deadline = started.Add(window - workerSafetyMargin)
|
||||
return nil
|
||||
}
|
||||
func workerAttemptState(s string) bool {
|
||||
switch s {
|
||||
case "LEASED", "RUNNING", "SUCCEEDED", "FAILED", "UNKNOWN", "EXPIRED", "CANCELLED":
|
||||
return true
|
||||
}
|
||||
return false
|
||||
}
|
||||
func workerTaskState(s TaskState) bool {
|
||||
switch s {
|
||||
case TaskWaiting, TaskLeased, TaskRunning, TaskSucceeded, TaskFailed, TaskCancelled, TaskNeedsReview:
|
||||
return true
|
||||
}
|
||||
return false
|
||||
}
|
||||
func workerIOError(err error) error {
|
||||
if errors.Is(err, context.Canceled) {
|
||||
return context.Canceled
|
||||
}
|
||||
if errors.Is(err, context.DeadlineExceeded) {
|
||||
return context.DeadlineExceeded
|
||||
}
|
||||
return ErrWorkerDisconnected
|
||||
}
|
||||
|
||||
func (w *WorkerConnection) String() string {
|
||||
return fmt.Sprintf("WorkerConnection{SessionID:%s secrets:redacted}", w.welcome.SessionID)
|
||||
}
|
||||
func (w *WorkerConnection) GoString() string { return w.String() }
|
||||
424
sdk/go/worker_connection_test.go
Normal file
424
sdk/go/worker_connection_test.go
Normal file
@@ -0,0 +1,424 @@
|
||||
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)
|
||||
}
|
||||
}
|
||||
385
sdk/go/worker_integration_test.go
Normal file
385
sdk/go/worker_integration_test.go
Normal file
@@ -0,0 +1,385 @@
|
||||
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
|
||||
dropped chan string
|
||||
server *httptest.Server
|
||||
}
|
||||
|
||||
func newWorkerFaultGateway(t *testing.T, upstream string) *workerFaultGateway {
|
||||
t.Helper()
|
||||
g := &workerFaultGateway{dropped: make(chan string, 4)}
|
||||
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) {
|
||||
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()
|
||||
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) 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):
|
||||
}
|
||||
}
|
||||
}
|
||||
205
sdk/go/worker_types.go
Normal file
205
sdk/go/worker_types.go
Normal file
@@ -0,0 +1,205 @@
|
||||
package orbit
|
||||
|
||||
import (
|
||||
"crypto/rand"
|
||||
"encoding/base64"
|
||||
"encoding/hex"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"sort"
|
||||
"time"
|
||||
)
|
||||
|
||||
// WorkerCapability 精确声明一种处理器与参数版本，不进行静默版本降级。
|
||||
type WorkerCapability struct {
|
||||
Handler string `json:"handler"` // Handler 处理器名称，1–128 位 ASCII 标识。
|
||||
HandlerVersion string `json:"handler_version"` // HandlerVersion 处理器版本，1–64 位。
|
||||
PayloadVersion string `json:"payload_version"` // PayloadVersion 输入契约版本，1–32 位。
|
||||
}
|
||||
|
||||
// WorkerRegistration 是同一进程重连时必须复用的身份；进程重启使用新身份。
|
||||
// 包含秘密，禁止以 JSON 等形式写入日志；SDK 的 String/GoString 会隐藏秘密。
|
||||
type WorkerRegistration struct {
|
||||
InstanceID string `json:"instance_id"` // InstanceID 本进程随机标识，32 位小写十六进制。
|
||||
InstanceSecret string `json:"instance_secret"` // InstanceSecret 32 字节随机秘密的 base64url 编码，重连保留。
|
||||
QueueID string `json:"queue_id"` // QueueID 绑定队列，重连不可更换。
|
||||
Capabilities []WorkerCapability `json:"capabilities"` // Capabilities 1–32 种精确能力，重连集合固定。
|
||||
Capacity uint32 `json:"capacity"` // Capacity 总执行槽数量，1–256；不是空闲槽的增量。
|
||||
}
|
||||
|
||||
// NewWorkerRegistration 生成进程身份；同进程只创建一次，并在重连时复用返回值。
|
||||
func NewWorkerRegistration(queueID string, capacity uint32, capabilities []WorkerCapability) (WorkerRegistration, error) {
|
||||
id, err := NewWorkerMessageID()
|
||||
if err != nil {
|
||||
return WorkerRegistration{}, err
|
||||
}
|
||||
var secret [32]byte
|
||||
if _, err := rand.Read(secret[:]); err != nil {
|
||||
return WorkerRegistration{}, err
|
||||
}
|
||||
return normalizeRegistration(WorkerRegistration{InstanceID: id, InstanceSecret: base64.RawURLEncoding.EncodeToString(secret[:]), QueueID: queueID, Capacity: capacity, Capabilities: capabilities})
|
||||
}
|
||||
|
||||
// NewWorkerMessageID 生成一次控制操作的身份；响应不明时重报必须复用原值。
|
||||
func NewWorkerMessageID() (string, error) {
|
||||
var id [16]byte
|
||||
if _, err := rand.Read(id[:]); err != nil {
|
||||
return "", err
|
||||
}
|
||||
return hex.EncodeToString(id[:]), nil
|
||||
}
|
||||
|
||||
// WorkerWelcome 是连接/心跳确认，不提供任务执行授权。
|
||||
type WorkerWelcome struct {
|
||||
SessionID string `json:"session_id"` // SessionID 当前连接身份，重连变化。
|
||||
Epoch uint64 `json:"epoch,string"` // Epoch 会话代次，JSON 使用字符串防止精度丢失。
|
||||
ServerTime time.Time `json:"server_time"` // ServerTime 服务端本次确认时的数据库时钟。
|
||||
ExpiresAt time.Time `json:"expires_at"` // ExpiresAt 会话期限，不能替代任务租约。
|
||||
LeaseDurationMS uint64 `json:"lease_duration_ms"` // LeaseDurationMS 单次租约的最大窗口，毫秒。
|
||||
Draining bool `json:"draining"` // Draining 是否已持久化排空；同进程重连保留。
|
||||
}
|
||||
|
||||
// WorkerCapacity 使用绝对容量和单调序号；降低容量不撤销在途执行。
|
||||
type WorkerCapacity struct {
|
||||
Sequence uint64 `json:"sequence,string"` // Sequence 当前连接内从 1 开始递增，同序号必须同内容。
|
||||
Capacity uint32 `json:"capacity"` // Capacity 总槽数，0–256；零停止新增领取。
|
||||
Drain bool `json:"drain"` // Drain 为真后不可在同一进程身份下清除。
|
||||
}
|
||||
|
||||
// WorkerAssignment 是预占通知；必须 Start 并核验当前租约后才能运行处理器。
|
||||
// 含业务输入和执行秘密，禁止序列化到普通日志。
|
||||
type WorkerAssignment struct {
|
||||
Task Task `json:"task"` // Task 任务受理快照及业务输入。
|
||||
AttemptID string `json:"attempt_id"` // AttemptID 本次尝试标识。
|
||||
Epoch uint64 `json:"epoch,string"` // Epoch 任务执行代次，隔离过时尝试。
|
||||
LeaseSecret string `json:"lease_secret"` // LeaseSecret 本次尝试的秘密，业务处理器不应接收它。
|
||||
LeaseUntil time.Time `json:"lease_until"` // LeaseUntil 预占期限；不能独立授权执行。
|
||||
ServerTime time.Time `json:"server_time"` // ServerTime 领取时数据库时间。
|
||||
}
|
||||
|
||||
// Attempt 返回控制操作所需的证明；不要把证明转交业务处理器。
|
||||
func (a WorkerAssignment) Attempt() WorkerAttempt {
|
||||
return WorkerAttempt{TaskID: a.Task.ID, AttemptID: a.AttemptID, Epoch: a.Epoch, LeaseSecret: a.LeaseSecret}
|
||||
}
|
||||
|
||||
// WorkerAttempt 证明对某次尝试的控制权，不包含连接身份；重连保留原值。
|
||||
type WorkerAttempt struct {
|
||||
TaskID string `json:"task_id"` // TaskID 目标任务标识。
|
||||
AttemptID string `json:"attempt_id"` // AttemptID 原尝试标识，重连不新建。
|
||||
Epoch uint64 `json:"epoch,string"` // Epoch 原执行代次。
|
||||
LeaseSecret string `json:"lease_secret"` // LeaseSecret 原执行秘密，禁止普通日志输出。
|
||||
}
|
||||
|
||||
// WorkerFinish 将明确结果与未知效果分开；普通网络错误不等于业务明确失败。
|
||||
type WorkerFinish struct {
|
||||
Ref WorkerAttempt `json:"ref"` // Ref 对目标尝试的控制证明。
|
||||
Outcome WorkerOutcome `json:"outcome"` // Outcome 明确成功、明确失败或效果未知。
|
||||
Result json.RawMessage `json:"result"` // Result 非 null JSON，最多 256 KiB；保留精确数字文本。
|
||||
ErrorCode string `json:"error_code,omitempty"` // ErrorCode 失败/未知时必填稳定分类，成功时为空；禁止原始异常正文。
|
||||
}
|
||||
|
||||
// WorkerOutcome 描述业务效果；UNKNOWN 必须进入核查而不是盲目重跑。
|
||||
type WorkerOutcome string
|
||||
|
||||
const (
|
||||
WorkerSucceeded WorkerOutcome = "SUCCEEDED"
|
||||
WorkerFailed WorkerOutcome = "FAILED"
|
||||
WorkerUnknown WorkerOutcome = "UNKNOWN"
|
||||
)
|
||||
|
||||
// WorkerReceipt 是不可变历史回执，不能作为当前执行权。
|
||||
type WorkerReceipt struct {
|
||||
Operation string `json:"operation"` // Operation 原控制操作 start、renew 或 finish。
|
||||
TaskState TaskState `json:"task_state"` // TaskState 原操作提交时的任务状态。
|
||||
AttemptState string `json:"attempt_state"` // AttemptState 原操作提交时的尝试状态。
|
||||
TaskRevision uint64 `json:"task_revision,string"` // TaskRevision 原操作提交版本。
|
||||
LeaseUntil time.Time `json:"lease_until"` // LeaseUntil 历史期限；即便尚未到期也不能独立授权。
|
||||
Disposition string `json:"disposition"` // Disposition 为 APPLIED 或 EVIDENCE_ONLY。
|
||||
}
|
||||
|
||||
// WorkerLease 是本次响应重新核验的授权；必须调用 ExecutionDeadline 判断可否执行。
|
||||
type WorkerLease struct {
|
||||
Active bool `json:"active"` // Active 当前会话是否仍拥有有效尝试；LEASED 仅代表预占。
|
||||
State string `json:"state"` // State 当前尝试状态，只有 RUNNING 可执行。
|
||||
LeaseUntil time.Time `json:"lease_until"` // LeaseUntil 当前服务端期限，不直接和客户端墙钟比较。
|
||||
ServerTime time.Time `json:"server_time"` // ServerTime 本次核验的数据库时钟。
|
||||
deadline time.Time // deadline 扣除全请求耗时和安全余量后的本地单调期限，不参与 JSON。
|
||||
}
|
||||
|
||||
// ExecutionDeadline 返回保守的本地执行截止时间。false 时不能运行或继续处理器。
|
||||
// 仅 Start/Renew/Resume 的 SDK 响应具有本地期限；自行 JSON 解码的值不会授权执行。
|
||||
func (l WorkerLease) ExecutionDeadline() (time.Time, bool) {
|
||||
return l.deadline, l.Active && l.State == "RUNNING" && !l.deadline.IsZero() && time.Now().Before(l.deadline)
|
||||
}
|
||||
|
||||
// WorkerCommandResult 把历史回执与当前权限分开；Deduplicated 不表示可再次执行。
|
||||
type WorkerCommandResult struct {
|
||||
Receipt WorkerReceipt `json:"receipt"` // Receipt 对应稳定消息身份的原始回执。
|
||||
Lease WorkerLease `json:"lease"` // Lease 本次重新验证的当前授权。
|
||||
Deduplicated bool `json:"deduplicated"` // Deduplicated 是否返回已保存的回执。
|
||||
}
|
||||
|
||||
func (r WorkerRegistration) String() string {
|
||||
return fmt.Sprintf("WorkerRegistration{InstanceID:%s QueueID:%s Capacity:%d secrets:redacted}", r.InstanceID, r.QueueID, r.Capacity)
|
||||
}
|
||||
func (r WorkerRegistration) GoString() string { return r.String() }
|
||||
func (a WorkerAttempt) String() string {
|
||||
return fmt.Sprintf("WorkerAttempt{TaskID:%s AttemptID:%s Epoch:%d secret:redacted}", a.TaskID, a.AttemptID, a.Epoch)
|
||||
}
|
||||
func (a WorkerAttempt) GoString() string { return a.String() }
|
||||
func (a WorkerAssignment) String() string { return a.Attempt().String() }
|
||||
func (a WorkerAssignment) GoString() string { return a.String() }
|
||||
func (f WorkerFinish) String() string {
|
||||
return fmt.Sprintf("WorkerFinish{%s Outcome:%s data:redacted}", f.Ref, f.Outcome)
|
||||
}
|
||||
func (f WorkerFinish) GoString() string { return f.String() }
|
||||
|
||||
func workerName(s string, max int) bool {
|
||||
if len(s) < 1 || len(s) > max {
|
||||
return false
|
||||
}
|
||||
for i, c := range s {
|
||||
if !(c >= 'a' && c <= 'z' || c >= 'A' && c <= 'Z' || c >= '0' && c <= '9' || i > 0 && (c == '_' || c == '.' || c == ':' || c == '-')) {
|
||||
return false
|
||||
}
|
||||
}
|
||||
return true
|
||||
}
|
||||
func validWorkerSecret(s string) bool {
|
||||
if len(s) != 43 {
|
||||
return false
|
||||
}
|
||||
raw, err := base64.RawURLEncoding.Strict().DecodeString(s)
|
||||
return err == nil && len(raw) == 32
|
||||
}
|
||||
func normalizeRegistration(r WorkerRegistration) (WorkerRegistration, error) {
|
||||
if !validID(r.InstanceID) || !validID(r.QueueID) || !validWorkerSecret(r.InstanceSecret) || r.Capacity < 1 || r.Capacity > 256 || len(r.Capabilities) < 1 || len(r.Capabilities) > 32 {
|
||||
return WorkerRegistration{}, errors.New("invalid Worker identity, queue, secret, capacity or capabilities")
|
||||
}
|
||||
r.Capabilities = append([]WorkerCapability(nil), r.Capabilities...)
|
||||
for _, c := range r.Capabilities {
|
||||
if !workerName(c.Handler, 128) || !workerName(c.HandlerVersion, 64) || !workerName(c.PayloadVersion, 32) {
|
||||
return WorkerRegistration{}, errors.New("invalid Worker capability")
|
||||
}
|
||||
}
|
||||
sort.Slice(r.Capabilities, func(i, j int) bool {
|
||||
a, b := r.Capabilities[i], r.Capabilities[j]
|
||||
if a.Handler != b.Handler {
|
||||
return a.Handler < b.Handler
|
||||
}
|
||||
if a.HandlerVersion != b.HandlerVersion {
|
||||
return a.HandlerVersion < b.HandlerVersion
|
||||
}
|
||||
return a.PayloadVersion < b.PayloadVersion
|
||||
})
|
||||
for i := 1; i < len(r.Capabilities); i++ {
|
||||
if r.Capabilities[i] == r.Capabilities[i-1] {
|
||||
return WorkerRegistration{}, errors.New("duplicate Worker capability")
|
||||
}
|
||||
}
|
||||
return r, nil
|
||||
}
|
||||
func validWorkerAttempt(a WorkerAttempt) bool {
|
||||
return validID(a.TaskID) && validID(a.AttemptID) && a.Epoch > 0 && validWorkerSecret(a.LeaseSecret)
|
||||
}
|
||||
Reference in New Issue
Block a user