134 lines
4.1 KiB
Go
134 lines
4.1 KiB
Go
// worker 演示自动 Worker：声明处理器，Run 管理连接/心跳/续租/恢复，取消后排空。
|
||
package main
|
||
|
||
import (
|
||
"context"
|
||
"encoding/json"
|
||
"errors"
|
||
"fmt"
|
||
"os"
|
||
"strings"
|
||
"time"
|
||
|
||
orbit "git.yuntuo.hk/k2safe/orbit/sdk/go"
|
||
)
|
||
|
||
func main() {
|
||
if err := run(); err != nil {
|
||
fmt.Fprintln(os.Stderr, err)
|
||
os.Exit(1)
|
||
}
|
||
}
|
||
func run() error {
|
||
token, queue := os.Getenv("ORBIT_TOKEN"), os.Getenv("ORBIT_QUEUE_ID")
|
||
if path := os.Getenv("ORBIT_CREDENTIALS_FILE"); path != "" {
|
||
raw, err := os.ReadFile(path)
|
||
if err != nil {
|
||
return fmt.Errorf("read private credentials: %w", err)
|
||
}
|
||
var credentials struct {
|
||
Token string `json:"token"`
|
||
QueueID string `json:"queue_id"`
|
||
}
|
||
if json.Unmarshal(raw, &credentials) != nil {
|
||
return errors.New("invalid credentials file")
|
||
}
|
||
if token == "" {
|
||
token = credentials.Token
|
||
}
|
||
if queue == "" {
|
||
queue = credentials.QueueID
|
||
}
|
||
}
|
||
client, err := orbit.NewClient(os.Getenv("ORBIT_URL"), orbit.Options{Token: token})
|
||
if err != nil {
|
||
return err
|
||
}
|
||
operation, err := orbit.NewWorkerMessageID()
|
||
if err != nil {
|
||
return err
|
||
}
|
||
handler := "text.normalize.runtime." + operation
|
||
w, err := orbit.NewWorker(client, orbit.WorkerOptions{QueueID: queue, Capacity: 2, ShutdownTimeout: 5 * time.Second, Handlers: []orbit.WorkerHandler{{Capability: orbit.WorkerCapability{Handler: handler, HandlerVersion: "v1", PayloadVersion: "v1"}, Handle: normalize}}})
|
||
if err != nil {
|
||
return err
|
||
}
|
||
ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second)
|
||
defer cancel()
|
||
// 示例先受理，再启动消费者；生产通常是独立的长驻 Worker 进程。
|
||
var tasks []orbit.Task
|
||
for i, text := range []string{" hello orbit ", " reusable worker ", " durable result "} {
|
||
payload, _ := json.Marshal(struct {
|
||
Text string `json:"text"`
|
||
}{text})
|
||
accepted, err := client.SubmitTask(ctx, fmt.Sprintf("runtime-%s-%d", operation, i), orbit.SubmitTaskRequest{QueueID: queue, Handler: handler, HandlerVersion: "v1", PayloadVersion: "v1", Payload: payload, TaskTimeoutMS: 60000, Retry: orbit.RetryPolicy{MaxAttempts: 1, AttemptTimeoutMS: 10000}, RecoverySafety: orbit.RecoveryIdempotent})
|
||
if err != nil {
|
||
return err
|
||
}
|
||
tasks = append(tasks, accepted.Task)
|
||
}
|
||
workerCtx, stop := context.WithCancel(ctx)
|
||
done := make(chan error, 1)
|
||
go func() { done <- w.Run(workerCtx) }()
|
||
// 任一调用失败时也等待有界排空，不在 main 返回时静默丢掉执行中的 Worker。
|
||
waited := false
|
||
defer func() {
|
||
stop()
|
||
if !waited {
|
||
<-done
|
||
}
|
||
}()
|
||
for _, task := range tasks {
|
||
for {
|
||
current, err := client.GetTask(ctx, task.ID)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
if current.State == orbit.TaskSucceeded {
|
||
fmt.Printf("task=%s state=%s result=%s\n", current.ID, current.State, current.Result)
|
||
break
|
||
}
|
||
if current.State == orbit.TaskFailed || current.State == orbit.TaskNeedsReview {
|
||
return fmt.Errorf("task %s requires attention: %s", current.ID, current.State)
|
||
}
|
||
select {
|
||
case err := <-done:
|
||
waited = true
|
||
return fmt.Errorf("Worker stopped before completion: %v", err)
|
||
case <-ctx.Done():
|
||
return ctx.Err()
|
||
case <-time.After(100 * time.Millisecond):
|
||
}
|
||
}
|
||
}
|
||
stop()
|
||
err = <-done
|
||
waited = true
|
||
if err != nil {
|
||
return err
|
||
}
|
||
stats := w.Stats()
|
||
if stats.Started != 3 || stats.ReportsConfirmed != 3 || stats.Running != 0 {
|
||
return errors.New("unexpected Worker completion counts")
|
||
}
|
||
fmt.Printf("worker_drained=true started=%d confirmed=%d running=%d\n", stats.Started, stats.ReportsConfirmed, stats.Running)
|
||
return nil
|
||
}
|
||
func normalize(ctx context.Context, execution orbit.WorkerExecution) (json.RawMessage, error) {
|
||
var in struct {
|
||
Text string `json:"text"`
|
||
}
|
||
if err := json.Unmarshal(execution.Task.Payload, &in); err != nil {
|
||
return nil, &orbit.WorkerFailure{Code: "INVALID_TEXT"}
|
||
}
|
||
// 模拟可取消的短工作；真实 I/O 也必须使用传入的 ctx。
|
||
select {
|
||
case <-ctx.Done():
|
||
return nil, ctx.Err()
|
||
case <-time.After(200 * time.Millisecond):
|
||
}
|
||
return json.Marshal(struct {
|
||
Text string `json:"text"`
|
||
}{Text: strings.ToUpper(strings.TrimSpace(in.Text))})
|
||
}
|