544 lines
14 KiB
Go
544 lines
14 KiB
Go
package orbit
|
||
|
||
import (
|
||
"context"
|
||
"errors"
|
||
"math/rand/v2"
|
||
"sync"
|
||
"sync/atomic"
|
||
"time"
|
||
)
|
||
|
||
type runtimeAttempt struct {
|
||
assignment WorkerAssignment
|
||
startID, renewID, finishID string
|
||
guard *executionGuard
|
||
invoked, returned, lost bool
|
||
result WorkerFinish
|
||
nextTry, nextRenew time.Time
|
||
retryDelay time.Duration
|
||
}
|
||
type runtimeCommand struct {
|
||
kind string
|
||
id string
|
||
job *runtimeAttempt
|
||
capacity WorkerCapacity
|
||
conn *WorkerConnection
|
||
}
|
||
type runtimeReply struct {
|
||
command *runtimeCommand
|
||
result WorkerCommandResult
|
||
welcome WorkerWelcome
|
||
err error
|
||
}
|
||
type runtimeConnected struct {
|
||
conn *WorkerConnection
|
||
leases map[string]WorkerLease
|
||
err error
|
||
}
|
||
|
||
type workerRuntime struct {
|
||
worker *Worker
|
||
ctx context.Context
|
||
jobs map[string]*runtimeAttempt
|
||
order []string
|
||
cursor int
|
||
conn *WorkerConnection
|
||
command *runtimeCommand
|
||
connecting bool
|
||
connectedOnce bool
|
||
draining, drainConfirmed bool
|
||
capacity uint32
|
||
sequence uint64
|
||
heartbeat, dialAt time.Time
|
||
backoff time.Duration
|
||
results chan handlerCompletion
|
||
replies chan runtimeReply
|
||
connections chan runtimeConnected
|
||
wg sync.WaitGroup
|
||
}
|
||
|
||
// Run 持续消费直到 ctx 取消或不可恢复的认证/协议错误。取消后先排空并在
|
||
// ShutdownTimeout 内收集在途结果；正常收尾返回 nil，超时返回 *WorkerRunError。
|
||
// 已到期的处理器上下文不会因重连复活；每个已知尝试在本进程最多调用一次处理器。
|
||
func (w *Worker) Run(parent context.Context) error {
|
||
if !w.used.CompareAndSwap(false, true) {
|
||
return ErrWorkerAlreadyRun
|
||
}
|
||
defer close(w.events)
|
||
if parent.Err() != nil {
|
||
return nil
|
||
}
|
||
ctx, cancel := context.WithCancel(context.WithoutCancel(parent))
|
||
var shutdownExpired atomic.Bool
|
||
watchDone := make(chan struct{})
|
||
go func() {
|
||
defer close(watchDone)
|
||
select {
|
||
case <-ctx.Done():
|
||
return
|
||
case <-parent.Done():
|
||
}
|
||
timer := time.NewTimer(w.options.ShutdownTimeout)
|
||
defer timer.Stop()
|
||
select {
|
||
case <-ctx.Done():
|
||
case <-timer.C:
|
||
shutdownExpired.Store(true)
|
||
cancel()
|
||
}
|
||
}()
|
||
r := &workerRuntime{worker: w, ctx: ctx, jobs: make(map[string]*runtimeAttempt), backoff: w.options.ReconnectMinDelay, results: make(chan handlerCompletion, w.options.Capacity), replies: make(chan runtimeReply, 2), connections: make(chan runtimeConnected, 1)}
|
||
defer func() {
|
||
cancel()
|
||
if r.conn != nil {
|
||
r.conn.Close()
|
||
}
|
||
for _, j := range r.jobs {
|
||
if j.guard != nil {
|
||
j.guard.stop(context.Canceled)
|
||
}
|
||
}
|
||
r.wg.Wait()
|
||
// 连接可能在退出与交付结果之间成功，仍需关闭；不留下新会话连接。
|
||
select {
|
||
case c := <-r.connections:
|
||
if c.conn != nil {
|
||
c.conn.Close()
|
||
}
|
||
default:
|
||
}
|
||
<-watchDone
|
||
}()
|
||
tick := time.NewTicker(20 * time.Millisecond)
|
||
defer tick.Stop()
|
||
stop := parent.Done()
|
||
for {
|
||
if ctx.Err() != nil {
|
||
cause := ctx.Err()
|
||
if shutdownExpired.Load() {
|
||
cause = ErrWorkerShutdownTimeout
|
||
}
|
||
return r.failure(cause)
|
||
}
|
||
if parent.Err() != nil && !r.draining {
|
||
r.draining = true
|
||
stop = nil
|
||
w.emit(WorkerEvent{Type: "draining"})
|
||
}
|
||
r.observeExpired()
|
||
if r.conn != nil && r.conn.Err() != nil {
|
||
if err := r.disconnect(r.conn.Err()); err != nil {
|
||
return r.failure(err)
|
||
}
|
||
}
|
||
// drain ACK 之前的派发已被底层读取器排进队列，收齐后才能宣布排空。
|
||
if r.conn != nil {
|
||
for len(r.conn.assignments) > 0 {
|
||
if err := r.accept(<-r.conn.assignments); err != nil {
|
||
return r.failure(err)
|
||
}
|
||
}
|
||
}
|
||
if r.draining && len(r.jobs) == 0 && (r.drainConfirmed || !r.connectedOnce && !r.connecting) {
|
||
return nil
|
||
}
|
||
if r.conn == nil && !r.connecting && !time.Now().Before(r.dialAt) {
|
||
r.connect()
|
||
}
|
||
if r.conn != nil && r.command == nil {
|
||
if err := r.dispatch(); err != nil {
|
||
return r.failure(err)
|
||
}
|
||
}
|
||
var assignments <-chan WorkerAssignment
|
||
var closed <-chan struct{}
|
||
if r.conn != nil {
|
||
assignments = r.conn.assignments
|
||
closed = r.conn.Done()
|
||
}
|
||
select {
|
||
case <-ctx.Done():
|
||
case <-stop:
|
||
stop = nil
|
||
case <-tick.C:
|
||
case <-closed:
|
||
case a := <-assignments:
|
||
if r.conn.Err() == nil {
|
||
if err := r.accept(a); err != nil {
|
||
return r.failure(err)
|
||
}
|
||
}
|
||
case done := <-r.results:
|
||
if j := r.jobs[done.attemptID]; j != nil {
|
||
j.returned = true
|
||
j.result = done.result
|
||
j.nextTry = time.Time{}
|
||
}
|
||
case reply := <-r.replies:
|
||
if err := r.receive(reply); err != nil {
|
||
return r.failure(err)
|
||
}
|
||
case connected := <-r.connections:
|
||
if err := r.establish(connected); err != nil {
|
||
return r.failure(err)
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
func (r *workerRuntime) failure(err error) error {
|
||
return &WorkerRunError{Cause: err, Unconfirmed: len(r.jobs), Running: r.worker.running.Load()}
|
||
}
|
||
func (r *workerRuntime) event(j *runtimeAttempt, kind, code string) {
|
||
r.worker.emit(WorkerEvent{Type: kind, TaskID: j.assignment.Task.ID, AttemptID: j.assignment.AttemptID, Code: code})
|
||
}
|
||
func (r *workerRuntime) accept(a WorkerAssignment) error {
|
||
if j := r.jobs[a.AttemptID]; j != nil {
|
||
if j.assignment.Attempt() != a.Attempt() {
|
||
return ErrWorkerProtocol
|
||
}
|
||
return nil
|
||
}
|
||
// 运行中失去授权但不退出的处理器仍占物理槽；允许接收一轮已在网络中的预占，
|
||
// 不突破处理器容量，也不让待核对身份无限累积。
|
||
if len(r.jobs) >= 2*int(r.worker.options.Capacity) {
|
||
return ErrWorkerBackpressure
|
||
}
|
||
id, err := NewWorkerMessageID()
|
||
if err != nil {
|
||
return err
|
||
}
|
||
r.jobs[a.AttemptID] = &runtimeAttempt{assignment: a, startID: id}
|
||
r.order = append(r.order, a.AttemptID)
|
||
return nil
|
||
}
|
||
func (r *workerRuntime) remove(j *runtimeAttempt) {
|
||
if j.guard != nil {
|
||
j.guard.stop(context.Canceled)
|
||
}
|
||
delete(r.jobs, j.assignment.AttemptID)
|
||
for i, id := range r.order {
|
||
if id == j.assignment.AttemptID {
|
||
r.order = append(r.order[:i], r.order[i+1:]...)
|
||
break
|
||
}
|
||
}
|
||
if len(r.order) > 0 {
|
||
r.cursor %= len(r.order)
|
||
} else {
|
||
r.cursor = 0
|
||
}
|
||
}
|
||
func (r *workerRuntime) observeExpired() {
|
||
for _, j := range r.jobs {
|
||
if j.guard != nil && j.guard.ctx.Err() != nil && !j.lost {
|
||
j.lost = true
|
||
r.event(j, "lease_lost", "LEASE_LOST")
|
||
}
|
||
}
|
||
}
|
||
func (r *workerRuntime) desiredCapacity() uint32 {
|
||
if r.draining {
|
||
return 0
|
||
}
|
||
ghosts := uint32(0)
|
||
for _, j := range r.jobs {
|
||
if j.returned || j.lost {
|
||
ghosts++
|
||
}
|
||
}
|
||
if ghosts >= r.worker.options.Capacity {
|
||
return 0
|
||
}
|
||
return r.worker.options.Capacity - ghosts
|
||
}
|
||
|
||
func (r *workerRuntime) connect() {
|
||
r.connecting = true
|
||
if r.connectedOnce {
|
||
r.worker.reconnects.Add(1)
|
||
r.worker.emit(WorkerEvent{Type: "reconnecting"})
|
||
}
|
||
refs := make([]WorkerAttempt, 0, len(r.jobs))
|
||
for _, j := range r.jobs {
|
||
refs = append(refs, j.assignment.Attempt())
|
||
}
|
||
r.wg.Add(1)
|
||
go func() {
|
||
defer r.wg.Done()
|
||
ctx, stop := context.WithTimeout(r.ctx, 30*time.Second)
|
||
defer stop()
|
||
c, err := r.worker.client.DialWorker(ctx, r.worker.registration)
|
||
result := runtimeConnected{conn: c, err: err, leases: make(map[string]WorkerLease, len(refs))}
|
||
if err == nil {
|
||
for _, ref := range refs {
|
||
id, e := NewWorkerMessageID()
|
||
if e != nil {
|
||
result.err = e
|
||
break
|
||
}
|
||
lease, e := c.Resume(ctx, id, ref)
|
||
if e != nil {
|
||
result.err = e
|
||
break
|
||
}
|
||
result.leases[ref.AttemptID] = lease
|
||
}
|
||
}
|
||
if result.err != nil && c != nil {
|
||
c.Close()
|
||
result.conn = nil
|
||
}
|
||
select {
|
||
case r.connections <- result:
|
||
case <-r.ctx.Done():
|
||
if result.conn != nil {
|
||
result.conn.Close()
|
||
}
|
||
}
|
||
}()
|
||
}
|
||
func (r *workerRuntime) establish(c runtimeConnected) error {
|
||
r.connecting = false
|
||
if c.err != nil {
|
||
if !workerRetryable(c.err) {
|
||
return c.err
|
||
}
|
||
r.deferReconnect()
|
||
return nil
|
||
}
|
||
r.conn = c.conn
|
||
r.connectedOnce = true
|
||
r.sequence = 0
|
||
r.heartbeat = time.Time{}
|
||
r.capacity = ^uint32(0)
|
||
r.drainConfirmed = false
|
||
r.draining = r.draining || c.conn.Welcome().Draining
|
||
for id, lease := range c.leases {
|
||
j := r.jobs[id]
|
||
if j == nil {
|
||
continue
|
||
}
|
||
if j.guard != nil {
|
||
if !j.guard.update(lease) {
|
||
j.lost = true
|
||
r.event(j, "lease_lost", "LEASE_LOST")
|
||
}
|
||
j.nextRenew = time.Time{}
|
||
}
|
||
if !j.invoked && !lease.Active {
|
||
r.event(j, "lease_lost", "NOT_STARTED")
|
||
r.remove(j)
|
||
}
|
||
}
|
||
r.worker.emit(WorkerEvent{Type: "connected"})
|
||
return nil
|
||
}
|
||
func (r *workerRuntime) deferReconnect() {
|
||
min := r.worker.options.ReconnectMinDelay
|
||
r.dialAt = time.Now().Add(min + time.Duration(rand.Int64N(int64(r.backoff-min)+1)))
|
||
r.backoff = minDuration(r.backoff*2, r.worker.options.ReconnectMaxDelay)
|
||
}
|
||
func (r *workerRuntime) disconnect(err error) error {
|
||
if r.conn != nil {
|
||
r.conn.Close()
|
||
r.conn = nil
|
||
}
|
||
r.command = nil
|
||
if !workerRetryable(err) {
|
||
return err
|
||
}
|
||
r.deferReconnect()
|
||
return nil
|
||
}
|
||
func workerRetryable(err error) bool {
|
||
if errors.Is(err, ErrWorkerDisconnected) || errors.Is(err, ErrWorkerClosed) || errors.Is(err, ErrWorkerBackpressure) || errors.Is(err, context.DeadlineExceeded) {
|
||
return true
|
||
}
|
||
var control *WorkerError
|
||
if errors.As(err, &control) {
|
||
return control.Code == "STALE_SESSION" || control.Code == "UNAVAILABLE"
|
||
}
|
||
var httpErr *HTTPError
|
||
return errors.As(err, &httpErr) && (httpErr.StatusCode == 429 || httpErr.StatusCode >= 500)
|
||
}
|
||
func minDuration(a, b time.Duration) time.Duration {
|
||
if a < b {
|
||
return a
|
||
}
|
||
return b
|
||
}
|
||
|
||
func (r *workerRuntime) dispatch() error {
|
||
now := time.Now()
|
||
capacity := r.desiredCapacity()
|
||
if !now.Before(r.heartbeat) || capacity != r.capacity || r.draining && !r.drainConfirmed {
|
||
id, err := NewWorkerMessageID()
|
||
if err != nil {
|
||
return err
|
||
}
|
||
r.sequence++
|
||
r.launch(&runtimeCommand{kind: "capacity", id: id, capacity: WorkerCapacity{Sequence: r.sequence, Capacity: capacity, Drain: r.draining}})
|
||
return nil
|
||
}
|
||
// 结果优先释放占用；其后照顾临近续租，再批准新处理器。每类内部轮转。
|
||
for _, kind := range []string{"finish", "renew", "start"} {
|
||
for n := 0; n < len(r.order); n++ {
|
||
i := (r.cursor + n) % len(r.order)
|
||
j := r.jobs[r.order[i]]
|
||
if now.Before(j.nextTry) {
|
||
continue
|
||
}
|
||
var id *string
|
||
switch kind {
|
||
case "finish":
|
||
if !j.returned {
|
||
continue
|
||
}
|
||
id = &j.finishID
|
||
case "renew":
|
||
if !j.invoked || j.lost || j.returned || now.Before(j.nextRenew) {
|
||
continue
|
||
}
|
||
id = &j.renewID
|
||
case "start":
|
||
if j.invoked || r.worker.running.Load() >= int64(r.worker.options.Capacity) {
|
||
continue
|
||
}
|
||
id = &j.startID
|
||
}
|
||
if *id == "" {
|
||
var err error
|
||
*id, err = NewWorkerMessageID()
|
||
if err != nil {
|
||
return err
|
||
}
|
||
}
|
||
r.cursor = (i + 1) % len(r.order)
|
||
r.launch(&runtimeCommand{kind: kind, id: *id, job: j})
|
||
return nil
|
||
}
|
||
}
|
||
return nil
|
||
}
|
||
func (r *workerRuntime) launch(command *runtimeCommand) {
|
||
command.conn = r.conn
|
||
r.command = command
|
||
// 操作所需值在 goroutine 启动前复制；运行器事件循环独占可变尝试状态。
|
||
var ref WorkerAttempt
|
||
var finish WorkerFinish
|
||
if command.job != nil {
|
||
ref = command.job.assignment.Attempt()
|
||
finish = command.job.result
|
||
}
|
||
r.wg.Add(1)
|
||
go func() {
|
||
defer r.wg.Done()
|
||
reply := runtimeReply{command: command}
|
||
switch command.kind {
|
||
case "capacity":
|
||
reply.welcome, reply.err = command.conn.Capacity(r.ctx, command.id, command.capacity)
|
||
case "start":
|
||
reply.result, reply.err = command.conn.Start(r.ctx, command.id, ref)
|
||
case "renew":
|
||
reply.result, reply.err = command.conn.Renew(r.ctx, command.id, ref)
|
||
case "finish":
|
||
reply.result, reply.err = command.conn.Finish(r.ctx, command.id, finish)
|
||
}
|
||
select {
|
||
case r.replies <- reply:
|
||
case <-r.ctx.Done():
|
||
}
|
||
}()
|
||
}
|
||
|
||
func (r *workerRuntime) receive(reply runtimeReply) error {
|
||
command := reply.command
|
||
if r.command != command || command.conn != r.conn {
|
||
return nil
|
||
}
|
||
r.command = nil
|
||
j := command.job
|
||
if reply.err != nil {
|
||
var control *WorkerError
|
||
if j != nil && errors.As(reply.err, &control) {
|
||
switch control.Code {
|
||
case "LEASE_LOST":
|
||
if command.kind == "finish" {
|
||
return reply.err
|
||
}
|
||
if j.guard != nil {
|
||
j.guard.stop(ErrExecutionLeaseLost)
|
||
j.lost = true
|
||
r.event(j, "lease_lost", control.Code)
|
||
} else {
|
||
r.remove(j)
|
||
}
|
||
return nil
|
||
case "EXECUTION_PAUSED", "UNAVAILABLE":
|
||
if j.retryDelay == 0 {
|
||
j.retryDelay = 200 * time.Millisecond
|
||
} else {
|
||
j.retryDelay = minDuration(j.retryDelay*2, 5*time.Second)
|
||
}
|
||
j.nextTry = time.Now().Add(j.retryDelay)
|
||
return nil
|
||
}
|
||
}
|
||
if workerRetryable(reply.err) {
|
||
return r.disconnect(reply.err)
|
||
}
|
||
return reply.err
|
||
}
|
||
r.backoff = r.worker.options.ReconnectMinDelay
|
||
if command.kind == "capacity" {
|
||
r.capacity = command.capacity.Capacity
|
||
r.drainConfirmed = command.capacity.Drain && reply.welcome.Draining
|
||
r.heartbeat = time.Now().Add(minDuration(5*time.Second, reply.welcome.ExpiresAt.Sub(reply.welcome.ServerTime)/3))
|
||
return nil
|
||
}
|
||
if command.kind == "finish" {
|
||
if reply.result.Receipt.Disposition == "APPLIED" {
|
||
r.worker.reports.Add(1)
|
||
r.event(j, "result_applied", string(reply.result.Receipt.TaskState))
|
||
} else {
|
||
r.worker.evidence.Add(1)
|
||
r.event(j, "result_evidence", string(reply.result.Receipt.TaskState))
|
||
}
|
||
r.remove(j)
|
||
return nil
|
||
}
|
||
if command.kind == "start" {
|
||
guard, ok := newExecutionGuard(r.ctx, reply.result.Lease)
|
||
if !ok {
|
||
r.event(j, "lease_lost", "START_NOT_AUTHORIZED")
|
||
r.remove(j)
|
||
return nil
|
||
}
|
||
j.guard = guard
|
||
j.invoked = true
|
||
r.worker.running.Add(1)
|
||
task := j.assignment.Task
|
||
task.Payload = append([]byte(nil), task.Payload...)
|
||
task.Result = append([]byte(nil), task.Result...)
|
||
execution := WorkerExecution{Task: task, AttemptID: j.assignment.AttemptID, Epoch: j.assignment.Epoch}
|
||
handler := r.worker.handlers[WorkerCapability{Handler: task.Handler, HandlerVersion: task.HandlerVersion, PayloadVersion: task.PayloadVersion}]
|
||
r.event(j, "handler_started", "")
|
||
go r.worker.invoke(guard.ctx, handler, execution, j.assignment.Attempt(), r.results, r.ctx.Done())
|
||
} else {
|
||
r.worker.renewals.Add(1)
|
||
j.renewID = ""
|
||
if !j.guard.update(reply.result.Lease) {
|
||
j.lost = true
|
||
r.event(j, "lease_lost", "LEASE_LOST")
|
||
}
|
||
}
|
||
j.nextTry = time.Time{}
|
||
j.retryDelay = 0
|
||
if !j.lost {
|
||
j.nextRenew = time.Now().Add(max(20*time.Millisecond, time.Until(j.guard.until())/3))
|
||
}
|
||
return nil
|
||
}
|