feat: add automatic Worker runtime with lease recovery and bounded shutdown
This commit is contained in:
543
sdk/go/worker_runtime.go
Normal file
543
sdk/go/worker_runtime.go
Normal file
@@ -0,0 +1,543 @@
|
||||
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
|
||||
}
|
||||
Reference in New Issue
Block a user