orbit/sdk/go/cmd/orbit-schedule-load/main.go

// orbit-schedule-load 在隔离环境测量固定间隔计划的持久化触发、真实Worker执行和来源核对。
// 计划模板只包含固定JSON输入，处理器没有外部业务副作用。
package main

import (
	"context"
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"io"
	"net/http"
	"os"
	"os/signal"
	"syscall"
	"time"

	orbit "git.yuntuo.hk/k2safe/orbit/sdk/go"
)

type config struct {
	URL, Token, QueueID                    string
	Requests, Schedules, Workers, Capacity int
	Interval, Timeout, RunTimeout          time.Duration
	PollInterval, StartDelay               time.Duration
}

func main() {
	if err := run(); err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}
}

func run() error {
	cfg, err := parseConfig(os.Args[1:])
	if err != nil {
		return err
	}
	client, err := orbit.NewClient(cfg.URL, orbit.Options{Token: cfg.Token, HTTPClient: &http.Client{Timeout: cfg.Timeout}})
	if err != nil {
		return errors.New("invalid client configuration")
	}
	runID, err := orbit.NewWorkerMessageID()
	if err != nil {
		return errors.New("cannot generate run identity")
	}
	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
	defer stop()
	report := measure(ctx, cfg, client, runID)
	if err := json.NewEncoder(os.Stdout).Encode(report); err != nil {
		return err
	}
	if !report.Complete {
		return errors.New("schedule sample incomplete; inspect report and original run identities")
	}
	return nil
}

func parseConfig(args []string) (config, error) {
	cfg := config{URL: os.Getenv("ORBIT_URL"), Token: os.Getenv("ORBIT_TOKEN"), QueueID: os.Getenv("ORBIT_QUEUE_ID"), Requests: 6, Schedules: 1, StartDelay: 5 * time.Second, Workers: 2, Capacity: 2, Interval: time.Second, Timeout: 5 * time.Second, RunTimeout: 2 * time.Minute, PollInterval: 250 * time.Millisecond}
	flags := flag.NewFlagSet("orbit-schedule-load", flag.ContinueOnError)
	var credentials string
	flags.StringVar(&cfg.URL, "url", cfg.URL, "Orbit HTTP(S) URL")
	flags.StringVar(&credentials, "credentials", os.Getenv("ORBIT_CREDENTIALS_FILE"), "private credentials file for an isolated namespace")
	flags.StringVar(&cfg.QueueID, "queue", cfg.QueueID, "existing isolated queue; configuration is not changed")
	flags.IntVar(&cfg.Requests, "requests", cfg.Requests, "expected occurrences per schedule, 1..100")
	flags.IntVar(&cfg.Schedules, "schedules", cfg.Schedules, "schedules sharing the same trigger times, 1..32")
	flags.DurationVar(&cfg.StartDelay, "start-delay", cfg.StartDelay, "lead time for setup, 1s..1m; late setup fails the sample")
	flags.IntVar(&cfg.Workers, "workers", cfg.Workers, "independent WSS Workers in this process, 1..16")
	flags.IntVar(&cfg.Capacity, "capacity", cfg.Capacity, "concurrent handlers per Worker, 1..128; total at most 256")
	flags.DurationVar(&cfg.Interval, "interval", cfg.Interval, "fixed-rate interval in whole milliseconds, 1s..1m")
	flags.DurationVar(&cfg.Timeout, "timeout", cfg.Timeout, "per HTTP request timeout, 1ms..1m")
	flags.DurationVar(&cfg.RunTimeout, "run-timeout", cfg.RunTimeout, "whole sample deadline, 5s..1h, plus at most 5s Worker shutdown and 5s schedule cleanup")
	flags.DurationVar(&cfg.PollInterval, "poll-interval", cfg.PollInterval, "occurrence polling interval, 100ms..10s")
	if err := flags.Parse(args); err != nil {
		return config{}, err
	}
	if flags.NArg() != 0 {
		return config{}, errors.New("positional arguments are not supported")
	}
	if credentials != "" {
		f, err := os.Open(credentials)
		if err != nil {
			return config{}, errors.New("cannot open private credentials file")
		}
		raw, err := io.ReadAll(io.LimitReader(f, (16<<10)+1))
		_ = f.Close()
		var saved struct {
			Token   string `json:"token"`
			QueueID string `json:"queue_id"`
		}
		if err != nil || len(raw) > 16<<10 || json.Unmarshal(raw, &saved) != nil {
			return config{}, errors.New("invalid credentials file (maximum 16 KiB)")
		}
		if cfg.Token == "" {
			cfg.Token = saved.Token
		}
		if cfg.QueueID == "" {
			cfg.QueueID = saved.QueueID
		}
	}
	if cfg.URL == "" || cfg.Token == "" || !validID(cfg.QueueID) {
		return config{}, errors.New("URL, token and a 32-character lowercase hexadecimal queue ID are required")
	}
	if cfg.Requests < 1 || cfg.Requests > 100 || cfg.Schedules < 1 || cfg.Schedules > 32 || cfg.Workers < 1 || cfg.Workers > 16 || cfg.Capacity < 1 || cfg.Capacity > 128 || cfg.Workers*cfg.Capacity > 256 {
		return config{}, errors.New("request or Worker count is out of bounds")
	}
	if cfg.Interval < time.Second || cfg.Interval > time.Minute || cfg.Interval%time.Millisecond != 0 || cfg.StartDelay < time.Second || cfg.StartDelay > time.Minute || cfg.StartDelay%time.Millisecond != 0 || cfg.Timeout < time.Millisecond || cfg.Timeout > time.Minute || cfg.RunTimeout < 5*time.Second || cfg.RunTimeout > time.Hour || cfg.PollInterval < 100*time.Millisecond || cfg.PollInterval > 10*time.Second {
		return config{}, errors.New("interval, timeout or poll interval is out of bounds")
	}
	return cfg, nil
}

func validID(id string) bool {
	if len(id) != 32 {
		return false
	}
	for _, c := range id {
		if !(c >= '0' && c <= '9' || c >= 'a' && c <= 'f') {
			return false
		}
	}
	return true
}

func classify(err error) string {
	if errors.Is(err, context.DeadlineExceeded) {
		return "TIMEOUT"
	}
	if errors.Is(err, context.Canceled) {
		return "CANCELLED"
	}
	var api *orbit.APIError
	if errors.As(err, &api) {
		switch api.StatusCode {
		case 401, 403:
			return "AUTH"
		case 409:
			return "CONFLICT"
		case 429:
			return "BACKPRESSURE"
		case 400, 422:
			return "INVALID_REQUEST"
		}
		return "HTTP"
	}
	var h *orbit.HTTPError
	if errors.As(err, &h) {
		return "HTTP"
	}
	return "TRANSPORT"
}