orbit/examples/batch/main.go

// batch 演示批量部分受理、稳定身份核对和释放配额后仅补交失败项。
package main

import (
	"context"
	"crypto/rand"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"os"
	"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 := os.Getenv("ORBIT_TOKEN")
	if token == "" && os.Getenv("ORBIT_CREDENTIALS_FILE") != "" {
		raw, err := os.ReadFile(os.Getenv("ORBIT_CREDENTIALS_FILE"))
		if err != nil {
			return fmt.Errorf("read credentials: %w", err)
		}
		var c struct {
			Token string `json:"token"`
		}
		if json.Unmarshal(raw, &c) != nil {
			return fmt.Errorf("invalid credentials file")
		}
		token = c.Token
	}
	api, err := orbit.NewClient(os.Getenv("ORBIT_URL"), orbit.Options{Token: token})
	if err != nil {
		return err
	}
	ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second)
	defer cancel()
	var random [16]byte
	if _, err := rand.Read(random[:]); err != nil {
		return err
	}
	key := "batch-example-" + hex.EncodeToString(random[:])
	q, err := api.CreateQueue(ctx, key+"-queue", orbit.CreateQueueRequest{Name: key, MaxConcurrency: 2, RatePerSecond: 10, MaxUnfinishedTasks: 2})
	if err != nil {
		return err
	}
	if _, err := api.ConfigureQueue(ctx, q.Queue.ID, key+"-pause", orbit.ConfigureQueueRequest{ExpectedRevision: q.Queue.Revision, Reason: "isolated batch acceptance example", Paused: true, MaxConcurrency: 2, RatePerSecond: 10, MaxUnfinishedTasks: 2}); err != nil {
		return err
	}
	request := orbit.SubmitTasksRequest{}
	// 演示运行生成独立身份；生产必须先将每项键与完整参数持久化到业务库，再发送。
	for i := 0; i < 3; i++ {
		payload := json.RawMessage(fmt.Sprintf(`{"source_ref":"objects/input-%d.txt","sequence":9007199254740993}`, i))
		request.Items = append(request.Items, orbit.SubmitTasksItem{IdempotencyKey: fmt.Sprintf("%s-%d", key, i), Task: orbit.SubmitTaskRequest{QueueID: q.Queue.ID, Handler: "file.convert", HandlerVersion: "v1", PayloadVersion: "v1", Payload: payload, TaskTimeoutMS: 60000}})
	}
	first, err := api.SubmitTasks(ctx, request)
	if err != nil {
		return err
	}
	if first.Items[0].Status != orbit.BatchAccepted || first.Items[1].Status != orbit.BatchAccepted || first.Items[2].Status != orbit.BatchRejected || first.Items[2].Error.Code != "QUEUE_FULL" {
		return fmt.Errorf("unexpected partial acceptance")
	}
	fmt.Println("accepted=2 rejected=1 code=QUEUE_FULL")
	replay, err := api.SubmitTasks(ctx, request)
	if err != nil {
		return err
	}
	for i := 0; i < 2; i++ {
		if replay.Items[i].Status != orbit.BatchAccepted || !replay.Items[i].Receipt.Deduplicated || replay.Items[i].Receipt.TaskID != first.Items[i].Receipt.TaskID {
			return fmt.Errorf("batch replay changed identity")
		}
	}
	if replay.Items[2].Status != orbit.BatchRejected || replay.Items[2].Error.Code != "QUEUE_FULL" {
		return fmt.Errorf("backpressure changed unexpectedly")
	}
	fmt.Println("replayed=2 same_task_ids=true")
	cancelTask := func(id string) error {
		current, err := api.GetTask(ctx, id)
		if err != nil {
			return err
		}
		out, err := api.CancelTask(ctx, id, key+"-cancel-"+id, orbit.CancelTaskRequest{ExpectedRevision: current.Revision, Reason: "finish batch example"})
		if err != nil {
			return err
		}
		if out.Task.State != orbit.TaskCancelled {
			return fmt.Errorf("expected confirmed cancellation")
		}
		return nil
	}
	if err := cancelTask(first.Items[0].Receipt.TaskID); err != nil {
		return err
	}
	// 子批次索引重新从0开始；键和参数仍来自最初被拒项，不重新生成。
	retry, err := api.SubmitTasks(ctx, orbit.SubmitTasksRequest{Items: request.Items[2:]})
	if err != nil {
		return err
	}
	if retry.Items[0].Status != orbit.BatchAccepted || retry.Items[0].Receipt.Deduplicated {
		return fmt.Errorf("failed item was not independently accepted")
	}
	fmt.Println("retried_items=1 same_identity=true accepted=true")
	for _, id := range []string{first.Items[1].Receipt.TaskID, retry.Items[0].Receipt.TaskID} {
		if err := cancelTask(id); err != nil {
			return err
		}
	}
	final, err := api.GetQueue(ctx, q.Queue.ID)
	if err != nil {
		return err
	}
	if final.UnfinishedTasks != 0 {
		return fmt.Errorf("unfinished tasks remain")
	}
	fmt.Printf("queue=%s paused=%t unfinished=%d\n", final.ID, final.Paused, final.UnfinishedTasks)
	return nil
}