468 lines
16 KiB
Go
468 lines
16 KiB
Go
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
|
||
cancellations chan WorkerCancellation
|
||
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, _ := marshalWire(r)
|
||
frame, _ := marshalWire(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), cancellations: make(chan WorkerCancellation, 32), 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 == "task.cancel" {
|
||
var c WorkerCancellation
|
||
if !w.registration.CancellationNotifications || f.ReplyTo != "" || json.Unmarshal(f.Body, &c) != nil || !validID(c.TaskID) || !validID(c.AttemptID) || c.Epoch == 0 || c.RequestedAt.IsZero() {
|
||
w.fail(ErrWorkerProtocol)
|
||
return
|
||
}
|
||
select {
|
||
case w.cancellations <- c:
|
||
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
|
||
}
|
||
}
|
||
|
||
// NextCancellation 等待特定尝试的停止提示；需注册 CancellationNotifications 并持续读取。
|
||
// 提示允许迟到或重复；只取消匹配 TaskID/AttemptID/Epoch 的处理器，结果仍通过 Finish 确认。
|
||
func (w *WorkerConnection) NextCancellation(ctx context.Context) (WorkerCancellation, error) {
|
||
if !w.registration.CancellationNotifications {
|
||
return WorkerCancellation{}, errors.New("cancellation notifications not enabled")
|
||
}
|
||
if err := ctx.Err(); err != nil {
|
||
return WorkerCancellation{}, err
|
||
}
|
||
select {
|
||
case <-ctx.Done():
|
||
return WorkerCancellation{}, ctx.Err()
|
||
case <-w.ctx.Done():
|
||
return WorkerCancellation{}, w.Err()
|
||
case c := <-w.cancellations:
|
||
if err := w.Err(); err != nil {
|
||
return WorkerCancellation{}, err
|
||
}
|
||
return c, 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 := marshalWire(in)
|
||
if err != nil {
|
||
return started, errors.New("invalid Worker request JSON")
|
||
}
|
||
raw, err := marshalWire(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.Retryable && in.Outcome != WorkerFailed {
|
||
return out, errors.New("retryable is only valid for known failure")
|
||
}
|
||
if (in.Outcome != WorkerSucceeded && in.Outcome != WorkerFailed && in.Outcome != WorkerUnknown && in.Outcome != WorkerCancelled) || 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 a.Task.RetryOf != nil && (!validID(*a.Task.RetryOf) || *a.Task.RetryOf == a.Task.ID) {
|
||
return false
|
||
}
|
||
if a.Task.Priority < 0 || a.Task.Priority > 9 {
|
||
return false
|
||
}
|
||
if (a.Task.ConcurrencyKey == "") != (a.ConcurrencyToken == 0) || (a.Task.ConcurrencyKey != "" && (!w.registration.ConcurrencyKeys || !validConcurrencyKey(a.Task.ConcurrencyKey))) {
|
||
return false
|
||
}
|
||
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() }
|