orbit/sdk/go/worker_connection.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
	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), 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 := 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) || 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() }