feat: add standalone Worker protocol SDK and ACK-loss regression
This commit is contained in:
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() }
|
||||
Reference in New Issue
Block a user