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

// orbit-exec-load 在隔离环境测量真实 HTTP 受理、WSS Worker 和持久化完成链路。
// 处理器只有有界等待及回显，不访问外部业务系统。
package main

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

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

type config struct {
	URL, Token, QueueID                      string
	Requests, Concurrency, Workers, Capacity int
	Timeout, RunTimeout, PollInterval        time.Duration
	Durations                                []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("execution 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: 100, Concurrency: 8, Workers: 2, Capacity: 4, Timeout: 5 * time.Second, RunTimeout: 2 * time.Minute, PollInterval: 250 * time.Millisecond}
	flags := flag.NewFlagSet("orbit-exec-load", flag.ContinueOnError)
	var credentials, durations 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, "tasks, 1..10000")
	flags.IntVar(&cfg.Concurrency, "concurrency", cfg.Concurrency, "maximum concurrent submission/verification requests, 1..128")
	flags.IntVar(&cfg.Workers, "workers", cfg.Workers, "independent WSS Workers in this process, 1..32")
	flags.IntVar(&cfg.Capacity, "capacity", cfg.Capacity, "concurrent handlers per Worker, 1..256; total at most 1024")
	flags.DurationVar(&cfg.Timeout, "timeout", cfg.Timeout, "per HTTP request timeout, 1ms..1m")
	flags.DurationVar(&cfg.RunTimeout, "run-timeout", cfg.RunTimeout, "whole sample deadline, 1s..1h, plus at most 5s Worker shutdown")
	flags.DurationVar(&cfg.PollInterval, "poll-interval", cfg.PollInterval, "per active verifier polling interval, 100ms..10s")
	flags.StringVar(&durations, "durations", "0ms", "1..16 round-robin simulated handler durations, e.g. 0ms,100ms,2s; whole milliseconds, 0..60s")
	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 > 10000 || cfg.Concurrency < 1 || cfg.Concurrency > 128 || cfg.Workers < 1 || cfg.Workers > 32 || cfg.Capacity < 1 || cfg.Capacity > 256 || cfg.Workers*cfg.Capacity > 1024 {
		return config{}, errors.New("request, concurrency or Worker count is out of bounds")
	}
	if cfg.Timeout < time.Millisecond || cfg.Timeout > time.Minute || cfg.RunTimeout < time.Second || cfg.RunTimeout > time.Hour || cfg.PollInterval < 100*time.Millisecond || cfg.PollInterval > 10*time.Second {
		return config{}, errors.New("timeout or poll interval is out of bounds")
	}
	parts := strings.Split(durations, ",")
	if len(parts) > 16 {
		return config{}, errors.New("at most 16 handler durations are allowed")
	}
	for _, part := range parts {
		d, err := time.ParseDuration(strings.TrimSpace(part))
		if err != nil || d < 0 || d > time.Minute || d%time.Millisecond != 0 || d+time.Second > cfg.RunTimeout {
			return config{}, errors.New("handler duration must be whole milliseconds in 0..60s and allow at least 1s within run-timeout")
		}
		cfg.Durations = append(cfg.Durations, d)
	}
	if cfg.Concurrency > cfg.Requests {
		cfg.Concurrency = cfg.Requests
	}
	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
}