orbit/examples/worker/main.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))})
}