551 lines
16 KiB
Go
551 lines
16 KiB
Go
package httpserver
|
||
|
||
import (
|
||
"bytes"
|
||
"context"
|
||
"encoding/json"
|
||
"errors"
|
||
"net/http"
|
||
"reflect"
|
||
"strconv"
|
||
"strings"
|
||
"sync"
|
||
"time"
|
||
|
||
"git.yuntuo.hk/k2safe/orbit/internal/dispatch"
|
||
"git.yuntuo.hk/k2safe/orbit/internal/fault"
|
||
"git.yuntuo.hk/k2safe/orbit/internal/jsonvalue"
|
||
"git.yuntuo.hk/k2safe/orbit/internal/jsonwire"
|
||
core "git.yuntuo.hk/k2safe/orbit/model"
|
||
"git.yuntuo.hk/k2safe/orbit/platform/auth"
|
||
"git.yuntuo.hk/k2safe/orbit/pulse/worker"
|
||
"github.com/coder/websocket"
|
||
)
|
||
|
||
const WorkerSubprotocol = "pulse.worker.v1"
|
||
|
||
// WorkerHub 管理节点内连接生命周期；任务归属和容量仍由 MySQL 决定。
|
||
type WorkerHub struct {
|
||
service *worker.Service
|
||
ctx context.Context
|
||
cancel context.CancelFunc
|
||
slots chan struct{}
|
||
mu sync.Mutex
|
||
closing bool
|
||
wg sync.WaitGroup
|
||
poll time.Duration
|
||
dispatch *dispatch.Pool
|
||
emptyPolls map[[32]byte]time.Time
|
||
}
|
||
|
||
func NewWorkerHub(parent context.Context, service *worker.Service, options ...HubOptions) *WorkerHub {
|
||
ctx, cancel := context.WithCancel(parent)
|
||
limits := resolveHubOptions(options)
|
||
h := &WorkerHub{service: service, ctx: ctx, cancel: cancel, slots: make(chan struct{}, limits.ConnectionLimit), poll: time.Second, emptyPolls: make(map[[32]byte]time.Time)}
|
||
h.dispatch = dispatch.New(ctx, limits.DispatchWorkers, limits.DispatchCapacity)
|
||
h.wg.Add(1)
|
||
go func() { defer h.wg.Done(); <-h.dispatch.Done() }()
|
||
return h
|
||
}
|
||
func RegisterWorkers(server *http.Server, hub *WorkerHub) {
|
||
server.Handler.(*http.ServeMux).Handle("GET /v1/worker/connect", hub)
|
||
}
|
||
|
||
// Shutdown 停止派发并关闭本节点连接；不清除持久化租约，由重連核对或到期恢复处理。
|
||
func (h *WorkerHub) Shutdown(ctx context.Context) error {
|
||
h.mu.Lock()
|
||
h.closing = true
|
||
h.cancel()
|
||
h.mu.Unlock()
|
||
done := make(chan struct{})
|
||
go func() { h.wg.Wait(); close(done) }()
|
||
select {
|
||
case <-done:
|
||
return nil
|
||
case <-ctx.Done():
|
||
return ctx.Err()
|
||
}
|
||
}
|
||
|
||
type workerFrame struct {
|
||
Version string `json:"protocol_version"`
|
||
Type string `json:"type"`
|
||
MessageID core.ID `json:"message_id"`
|
||
ReplyTo core.ID `json:"reply_to,omitempty"`
|
||
SessionID core.ID `json:"session_id,omitempty"`
|
||
Body json.RawMessage `json:"body,omitempty"`
|
||
Error *fault.Error `json:"error,omitempty"`
|
||
}
|
||
type outboundFrame struct {
|
||
data []byte
|
||
done chan error
|
||
}
|
||
|
||
func (h *WorkerHub) ServeHTTP(w http.ResponseWriter, r *http.Request) {
|
||
requestID, _ := core.NewID()
|
||
w.Header().Set("X-Request-ID", string(requestID))
|
||
w.Header().Set("Cache-Control", "no-store")
|
||
if r.URL.RawQuery != "" {
|
||
writeError(w, string(requestID), fault.Invalid("Worker credentials and options must not be placed in URL"))
|
||
return
|
||
}
|
||
values := r.Header.Values("Authorization")
|
||
if len(values) != 1 {
|
||
writeError(w, string(requestID), fault.Unauthorized)
|
||
return
|
||
}
|
||
scheme, value, ok := strings.Cut(values[0], " ")
|
||
if !ok || !strings.EqualFold(scheme, "Bearer") {
|
||
writeError(w, string(requestID), fault.Unauthorized)
|
||
return
|
||
}
|
||
token, err := auth.Parse(value)
|
||
if err != nil {
|
||
writeError(w, string(requestID), err)
|
||
return
|
||
}
|
||
supported := false
|
||
for _, header := range r.Header.Values("Sec-WebSocket-Protocol") {
|
||
for _, protocol := range strings.Split(header, ",") {
|
||
supported = supported || strings.TrimSpace(protocol) == WorkerSubprotocol
|
||
}
|
||
}
|
||
if !supported {
|
||
writeJSON(w, http.StatusUpgradeRequired, map[string]any{"request_id": requestID, "error": fault.New("PROTOCOL_REQUIRED", "pulse.worker.v1 subprotocol is required")})
|
||
return
|
||
}
|
||
h.mu.Lock()
|
||
if h.closing || h.ctx.Err() != nil {
|
||
h.mu.Unlock()
|
||
writeError(w, string(requestID), fault.Unavailable)
|
||
return
|
||
}
|
||
select {
|
||
case h.slots <- struct{}{}:
|
||
h.wg.Add(1)
|
||
default:
|
||
h.mu.Unlock()
|
||
writeJSON(w, 429, map[string]any{"request_id": requestID, "error": fault.New("CAPACITY_EXCEEDED", "node connection capacity exceeded")})
|
||
return
|
||
}
|
||
h.mu.Unlock()
|
||
defer func() { <-h.slots; h.wg.Done() }()
|
||
authorize, cancel := context.WithTimeout(r.Context(), 5*time.Second)
|
||
err = h.service.Authorize(authorize, token)
|
||
cancel()
|
||
if err != nil {
|
||
writeError(w, string(requestID), err)
|
||
return
|
||
}
|
||
conn, err := websocket.Accept(w, r, &websocket.AcceptOptions{Subprotocols: []string{WorkerSubprotocol}, CompressionMode: websocket.CompressionDisabled})
|
||
if err != nil {
|
||
return
|
||
}
|
||
defer conn.CloseNow()
|
||
conn.SetReadLimit(512 << 10)
|
||
ctx, cancel := context.WithCancel(h.ctx)
|
||
defer cancel()
|
||
helloCtx, helloCancel := context.WithTimeout(ctx, 5*time.Second)
|
||
kind, raw, err := conn.Read(helloCtx)
|
||
helloCancel()
|
||
if err != nil {
|
||
return
|
||
}
|
||
hello, decodeErr := decodeWorkerRequest(raw, true)
|
||
var input worker.ConnectInput
|
||
if kind != websocket.MessageText || decodeErr != nil || decodeWorkerJSON(hello.Body, &input) != nil {
|
||
_ = conn.Close(websocket.StatusPolicyViolation, "invalid worker hello")
|
||
return
|
||
}
|
||
call, cancelCall := context.WithTimeout(ctx, 8*time.Second)
|
||
binding, welcome, err := h.service.Connect(call, token, input)
|
||
cancelCall()
|
||
if err != nil {
|
||
_ = writeWorkerFrame(ctx, conn, workerFrame{Version: "1", Type: "error", MessageID: requestID, ReplyTo: hello.MessageID, Error: publicWorkerError(err)})
|
||
return
|
||
}
|
||
welcomeRaw, err := jsonwire.Marshal(welcome)
|
||
if err != nil {
|
||
return
|
||
}
|
||
if err := writeWorkerFrame(ctx, conn, workerFrame{Version: "1", Type: "worker.welcome", MessageID: requestID, ReplyTo: hello.MessageID, SessionID: binding.SessionID, Body: welcomeRaw}); err != nil {
|
||
return
|
||
}
|
||
// 有界读写队列与 SQL 控制循环分开，处理器绝不运行在这里。
|
||
incoming := make(chan workerFrame, 4)
|
||
outgoing := make(chan outboundFrame, 8)
|
||
var loops sync.WaitGroup
|
||
loops.Add(2)
|
||
go func() {
|
||
defer loops.Done()
|
||
defer cancel()
|
||
for {
|
||
readCtx, stop := context.WithTimeout(ctx, 60*time.Second)
|
||
kind, raw, err := conn.Read(readCtx)
|
||
stop()
|
||
if err != nil {
|
||
return
|
||
}
|
||
frame, decodeErr := decodeWorkerRequest(raw, false)
|
||
if kind != websocket.MessageText || decodeErr != nil || frame.SessionID != binding.SessionID {
|
||
_ = conn.Close(websocket.StatusPolicyViolation, "invalid worker frame")
|
||
return
|
||
}
|
||
select {
|
||
case incoming <- frame:
|
||
case <-ctx.Done():
|
||
return
|
||
}
|
||
}
|
||
}()
|
||
go func() {
|
||
defer loops.Done()
|
||
defer cancel()
|
||
for {
|
||
select {
|
||
case <-ctx.Done():
|
||
return
|
||
case frame := <-outgoing:
|
||
writeCtx, stop := context.WithTimeout(ctx, 5*time.Second)
|
||
err := conn.Write(writeCtx, websocket.MessageText, frame.data)
|
||
stop()
|
||
if frame.done != nil {
|
||
frame.done <- err
|
||
}
|
||
if err != nil {
|
||
return
|
||
}
|
||
}
|
||
}
|
||
}()
|
||
defer func() { cancel(); _ = conn.CloseNow(); loops.Wait() }()
|
||
send := func(frame workerFrame, wait bool) bool {
|
||
frame.Version = "1"
|
||
frame.SessionID = binding.SessionID
|
||
if frame.MessageID == "" {
|
||
var err error
|
||
frame.MessageID, err = core.NewID()
|
||
if err != nil {
|
||
return false
|
||
}
|
||
}
|
||
raw, err := jsonwire.Marshal(frame)
|
||
if err != nil {
|
||
return false
|
||
}
|
||
out := outboundFrame{data: raw}
|
||
if wait {
|
||
out.done = make(chan error, 1)
|
||
}
|
||
select {
|
||
case outgoing <- out:
|
||
default:
|
||
return false
|
||
}
|
||
if !wait {
|
||
return true
|
||
}
|
||
select {
|
||
case err := <-out.done:
|
||
return err == nil
|
||
case <-ctx.Done():
|
||
return false
|
||
}
|
||
}
|
||
capacity := input.Capacity
|
||
draining := welcome.Draining
|
||
eligible := false
|
||
leases := map[core.ID]time.Time{}
|
||
cancelSent := map[core.ID]bool{}
|
||
results := make(chan claimResult, 1)
|
||
pending, burstLeft := false, 0
|
||
pollKey := workerPollKey(token.NamespaceID, input)
|
||
canRequest := func() bool {
|
||
return eligible && !draining && capacity > 0 && len(leases) < int(capacity) && len(outgoing) < 4
|
||
}
|
||
requestClaim := func() {
|
||
if pending || burstLeft == 0 || !canRequest() {
|
||
return
|
||
}
|
||
pending = h.dispatch.Submit(dispatch.Job{Context: ctx, Application: string(binding.ApplicationID), Queue: string(token.NamespaceID) + "/" + string(input.QueueID), Run: func(jobCtx context.Context) {
|
||
results <- h.claimFair(jobCtx, token, binding, pollKey)
|
||
}})
|
||
if pending {
|
||
burstLeft--
|
||
}
|
||
}
|
||
ticker := time.NewTicker(h.poll)
|
||
defer ticker.Stop()
|
||
for {
|
||
select {
|
||
case <-ctx.Done():
|
||
return
|
||
case result := <-results:
|
||
pending = false
|
||
if result.err != nil {
|
||
if fault.Code(result.err) == "STALE_SESSION" || fault.Code(result.err) == "UNAUTHENTICATED" || fault.Code(result.err) == "PERMISSION_DENIED" {
|
||
send(workerFrame{Type: "error", Error: publicWorkerError(result.err)}, true)
|
||
return
|
||
}
|
||
burstLeft = 0
|
||
continue
|
||
}
|
||
if result.assignment == nil {
|
||
burstLeft = 0
|
||
continue
|
||
}
|
||
assignment := result.assignment
|
||
raw, err := jsonwire.Marshal(assignment)
|
||
if err != nil {
|
||
return
|
||
}
|
||
leases[assignment.AttemptID] = assignment.LeaseUntil
|
||
if !send(workerFrame{Type: "task.assign", Body: raw}, false) {
|
||
return
|
||
}
|
||
requestClaim()
|
||
case frame := <-incoming:
|
||
call, stop := context.WithTimeout(ctx, 8*time.Second)
|
||
var body any
|
||
var commandErr error
|
||
switch frame.Type {
|
||
case "worker.capacity":
|
||
var in worker.CapacityInput
|
||
if commandErr = decodeWorkerJSON(frame.Body, &in); commandErr == nil {
|
||
var out worker.Welcome
|
||
out, commandErr = h.service.Capacity(call, token, binding, in)
|
||
body = out
|
||
if commandErr == nil {
|
||
capacity = in.Capacity
|
||
draining = out.Draining
|
||
eligible = true
|
||
}
|
||
}
|
||
case "worker.resume":
|
||
var ref worker.AttemptRef
|
||
if commandErr = decodeWorkerJSON(frame.Body, &ref); commandErr == nil {
|
||
var out worker.LeaseStatus
|
||
out, commandErr = h.service.Resume(call, token, binding, ref)
|
||
body = out
|
||
if commandErr == nil {
|
||
if out.Active {
|
||
leases[ref.AttemptID] = out.LeaseUntil
|
||
} else {
|
||
delete(leases, ref.AttemptID)
|
||
}
|
||
}
|
||
}
|
||
case "task.start", "task.renew":
|
||
var ref worker.AttemptRef
|
||
if commandErr = decodeWorkerJSON(frame.Body, &ref); commandErr == nil {
|
||
var out worker.CommandResult
|
||
if frame.Type == "task.start" {
|
||
out, commandErr = h.service.Start(call, token, binding, frame.MessageID, ref)
|
||
} else {
|
||
out, commandErr = h.service.Renew(call, token, binding, frame.MessageID, ref)
|
||
}
|
||
body = out
|
||
if commandErr == nil {
|
||
if out.Lease.Active {
|
||
leases[ref.AttemptID] = out.Lease.LeaseUntil
|
||
} else {
|
||
delete(leases, ref.AttemptID)
|
||
}
|
||
}
|
||
}
|
||
case "task.finish":
|
||
var in worker.FinishInput
|
||
if commandErr = decodeWorkerJSON(frame.Body, &in); commandErr == nil {
|
||
var out worker.CommandResult
|
||
out, commandErr = h.service.Finish(call, token, binding, frame.MessageID, in)
|
||
body = out
|
||
if commandErr == nil && !out.Lease.Active {
|
||
delete(leases, in.Ref.AttemptID)
|
||
}
|
||
}
|
||
default:
|
||
commandErr = fault.Invalid("unsupported worker message type")
|
||
}
|
||
stop()
|
||
reply := workerFrame{Type: "ack", ReplyTo: frame.MessageID}
|
||
if commandErr != nil {
|
||
reply.Type = "error"
|
||
reply.Error = publicWorkerError(commandErr)
|
||
} else {
|
||
reply.Body, commandErr = jsonwire.Marshal(body)
|
||
if commandErr != nil {
|
||
return
|
||
}
|
||
}
|
||
fatal := commandErr != nil && (fault.Code(commandErr) == "STALE_SESSION" || fault.Code(commandErr) == "UNAUTHENTICATED" || fault.Code(commandErr) == "PERMISSION_DENIED")
|
||
if !send(reply, fatal) || fatal {
|
||
return
|
||
}
|
||
case <-ticker.C:
|
||
// 取消与领取容量、排空无关；通知丢失后仍可由 Resume/Renew 核对持久化意图。
|
||
if input.CancellationNotifications && len(leases) > 0 {
|
||
call, stop := context.WithTimeout(ctx, 8*time.Second)
|
||
items, err := h.service.Cancellations(call, token, binding)
|
||
stop()
|
||
if err != nil {
|
||
if fault.Code(err) == "STALE_SESSION" || fault.Code(err) == "UNAUTHENTICATED" || fault.Code(err) == "PERMISSION_DENIED" {
|
||
send(workerFrame{Type: "error", Error: publicWorkerError(err)}, true)
|
||
return
|
||
}
|
||
} else {
|
||
active := make(map[core.ID]bool, len(items))
|
||
for _, item := range items {
|
||
active[item.AttemptID] = true
|
||
if cancelSent[item.AttemptID] || len(outgoing) >= 4 {
|
||
continue
|
||
}
|
||
raw, err := jsonwire.Marshal(item)
|
||
if err != nil || !send(workerFrame{Type: "task.cancel", Body: raw}, false) {
|
||
return
|
||
}
|
||
cancelSent[item.AttemptID] = true
|
||
}
|
||
for id := range cancelSent {
|
||
if !active[id] {
|
||
delete(cancelSent, id)
|
||
}
|
||
}
|
||
}
|
||
}
|
||
// 本地记录仅减少空轮询；SQL 仍原子验证总容量、队列额度和有效会话。
|
||
now := time.Now()
|
||
for attempt, until := range leases {
|
||
if !until.After(now) {
|
||
delete(leases, attempt)
|
||
}
|
||
}
|
||
if !pending {
|
||
burstLeft = 4
|
||
requestClaim()
|
||
}
|
||
}
|
||
}
|
||
}
|
||
func publicWorkerError(err error) *fault.Error {
|
||
var known *fault.Error
|
||
if errors.As(err, &known) {
|
||
return known
|
||
}
|
||
return fault.Unavailable
|
||
}
|
||
func writeWorkerFrame(ctx context.Context, conn *websocket.Conn, frame workerFrame) error {
|
||
raw, err := jsonwire.Marshal(frame)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
bounded, cancel := context.WithTimeout(ctx, 5*time.Second)
|
||
defer cancel()
|
||
return conn.Write(bounded, websocket.MessageText, raw)
|
||
}
|
||
func decodeWorkerJSON(raw []byte, out any) error {
|
||
if len(raw) == 0 || len(raw) > 512<<10 {
|
||
return fault.Invalid("invalid worker message size")
|
||
}
|
||
if _, err := jsonvalue.Canonical(raw); err != nil {
|
||
return fault.Invalid("invalid worker JSON")
|
||
}
|
||
if err := exactFields(raw, reflect.TypeOf(out).Elem()); err != nil {
|
||
return fault.Invalid("unknown or invalid worker field")
|
||
}
|
||
if err := requiredWorkerFields(raw, reflect.TypeOf(out).Elem()); err != nil {
|
||
return err
|
||
}
|
||
decoder := json.NewDecoder(bytes.NewReader(raw))
|
||
decoder.DisallowUnknownFields()
|
||
if err := decoder.Decode(out); err != nil {
|
||
return fault.Invalid("invalid worker field value")
|
||
}
|
||
return nil
|
||
}
|
||
|
||
// decodeWorkerRequest 拒绝请求中伪装为响应的字段；空字符串/null 不能绕过方向约束。
|
||
func decodeWorkerRequest(raw []byte, hello bool) (workerFrame, error) {
|
||
var frame workerFrame
|
||
if err := decodeWorkerJSON(raw, &frame); err != nil {
|
||
return frame, err
|
||
}
|
||
if frame.Version != "1" || frame.MessageID.Validate() != nil || len(frame.Body) == 0 {
|
||
return frame, fault.Invalid("invalid worker envelope")
|
||
}
|
||
var fields map[string]json.RawMessage
|
||
_ = json.Unmarshal(raw, &fields)
|
||
if fields["error"] != nil || fields["reply_to"] != nil {
|
||
return frame, fault.Invalid("response fields are not allowed in worker requests")
|
||
}
|
||
if hello {
|
||
if frame.Type != "worker.hello" || fields["session_id"] != nil {
|
||
return frame, fault.Invalid("invalid worker hello envelope")
|
||
}
|
||
} else if frame.SessionID.Validate() != nil {
|
||
return frame, fault.Invalid("worker session is required")
|
||
}
|
||
return frame, nil
|
||
}
|
||
|
||
// requiredWorkerFields 校验必填字段及规范 uint64 字符串，不把省略布尔值等同于显式控制意图。
|
||
func requiredWorkerFields(raw []byte, t reflect.Type) error {
|
||
if t == reflect.TypeOf(json.RawMessage{}) || t == reflect.TypeOf(time.Time{}) {
|
||
return nil
|
||
}
|
||
if t.Kind() == reflect.Pointer {
|
||
return nil
|
||
}
|
||
if t.Kind() == reflect.Array || t.Kind() == reflect.Slice {
|
||
var values []json.RawMessage
|
||
if err := json.Unmarshal(raw, &values); err != nil {
|
||
return fault.Invalid("invalid worker array")
|
||
}
|
||
for _, value := range values {
|
||
if err := requiredWorkerFields(value, t.Elem()); err != nil {
|
||
return err
|
||
}
|
||
}
|
||
return nil
|
||
}
|
||
if t.Kind() != reflect.Struct {
|
||
return nil
|
||
}
|
||
var fields map[string]json.RawMessage
|
||
if json.Unmarshal(raw, &fields) != nil {
|
||
return fault.Invalid("invalid worker object")
|
||
}
|
||
for i := 0; i < t.NumField(); i++ {
|
||
field := t.Field(i)
|
||
tag := strings.Split(field.Tag.Get("json"), ",")
|
||
if tag[0] == "" || tag[0] == "-" {
|
||
continue
|
||
}
|
||
value, found := fields[tag[0]]
|
||
optional, numberString := false, false
|
||
for _, flag := range tag[1:] {
|
||
optional = optional || flag == "omitempty"
|
||
numberString = numberString || flag == "string"
|
||
}
|
||
if !found {
|
||
if optional {
|
||
continue
|
||
}
|
||
return fault.Invalid("missing required worker field")
|
||
}
|
||
if field.Type.Kind() == reflect.Uint64 && numberString {
|
||
var text string
|
||
if json.Unmarshal(value, &text) != nil {
|
||
return fault.Invalid("worker counter must be a decimal string")
|
||
}
|
||
n, err := strconv.ParseUint(text, 10, 64)
|
||
if err != nil || n == 0 || strconv.FormatUint(n, 10) != text {
|
||
return fault.Invalid("invalid worker counter")
|
||
}
|
||
}
|
||
if err := requiredWorkerFields(value, field.Type); err != nil {
|
||
return err
|
||
}
|
||
}
|
||
return nil
|
||
}
|