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

301 lines
8.8 KiB
Go

// Command orbit-load runs a bounded, non-retrying task admission sample.
// It is a measurement tool for an isolated environment, not a production
// scheduler and not a capacity promise.
package main
import (
"context"
"crypto/rand"
"encoding/hex"
"encoding/json"
"errors"
"flag"
"fmt"
"math"
"net/http"
"os"
"sort"
"strings"
"sync"
"time"
orbit "git.yuntuo.hk/k2safe/orbit/sdk/go"
)
type credentials struct {
Token string `json:"token"`
QueueID string `json:"queue_id"`
}
type sampleConfig struct {
URL string
Token string
QueueID string
Requests int
Concurrency int
Timeout time.Duration
Handler string
HandlerVersion string
PayloadVersion string
}
type sampleResult struct {
duration time.Duration
code string
}
type loadReport struct {
StartedAt time.Time `json:"started_at"`
FinishedAt time.Time `json:"finished_at"`
Requests int `json:"requests"`
Concurrency int `json:"concurrency"`
Success int `json:"success"`
Errors int `json:"errors"`
P50MS float64 `json:"p50_ms"`
P95MS float64 `json:"p95_ms"`
P99MS float64 `json:"p99_ms"`
MaxMS float64 `json:"max_ms"`
ErrorCodes map[string]int `json:"error_codes,omitempty"`
}
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{
HTTPClient: &http.Client{Timeout: cfg.Timeout},
Token: cfg.Token,
})
if err != nil {
return err
}
runID, err := newRunID()
if err != nil {
return err
}
started := time.Now().UTC()
results := runSamples(client, cfg, runID)
finished := time.Now().UTC()
report := makeReport(cfg, started, finished, results)
if err := json.NewEncoder(os.Stdout).Encode(report); err != nil {
return err
}
if report.Errors != 0 {
return errors.New("load sample had failed requests; inspect error_codes and rerun in an isolated environment")
}
return nil
}
func parseConfig(args []string) (sampleConfig, error) {
flags := flag.NewFlagSet("orbit-load", flag.ContinueOnError)
flags.SetOutput(os.Stderr)
cfg := sampleConfig{
URL: os.Getenv("ORBIT_URL"),
Token: os.Getenv("ORBIT_TOKEN"),
QueueID: os.Getenv("ORBIT_QUEUE_ID"),
Requests: 100,
Concurrency: 8,
Timeout: 10 * time.Second,
Handler: "load.noop",
HandlerVersion: "v1",
PayloadVersion: "v1",
}
credentialsPath := os.Getenv("ORBIT_CREDENTIALS_FILE")
flags.StringVar(&cfg.URL, "url", cfg.URL, "Orbit HTTP URL")
flags.StringVar(&cfg.Token, "token", cfg.Token, "Bearer token; prefer ORBIT_CREDENTIALS_FILE")
flags.StringVar(&cfg.QueueID, "queue", cfg.QueueID, "target queue ID")
flags.StringVar(&credentialsPath, "credentials", credentialsPath, "private bootstrap credentials JSON file")
flags.IntVar(&cfg.Requests, "requests", cfg.Requests, "number of independent admissions, 1..100000")
flags.IntVar(&cfg.Concurrency, "concurrency", cfg.Concurrency, "maximum concurrent HTTP requests, 1..1024")
flags.DurationVar(&cfg.Timeout, "timeout", cfg.Timeout, "per-request HTTP timeout")
flags.StringVar(&cfg.Handler, "handler", cfg.Handler, "handler name frozen in each task")
flags.StringVar(&cfg.HandlerVersion, "handler-version", cfg.HandlerVersion, "handler version frozen in each task")
flags.StringVar(&cfg.PayloadVersion, "payload-version", cfg.PayloadVersion, "payload contract version frozen in each task")
if err := flags.Parse(args); err != nil {
return sampleConfig{}, err
}
if flags.NArg() != 0 {
return sampleConfig{}, errors.New("orbit-load does not accept positional arguments")
}
if credentialsPath != "" {
file, err := os.ReadFile(credentialsPath)
if err != nil {
return sampleConfig{}, fmt.Errorf("read credentials: %w", err)
}
if len(file) > 16<<10 {
return sampleConfig{}, errors.New("credentials file exceeds 16 KiB")
}
var saved credentials
if err := json.Unmarshal(file, &saved); err != nil {
return sampleConfig{}, errors.New("invalid credentials file")
}
if cfg.Token == "" {
cfg.Token = saved.Token
}
if cfg.QueueID == "" {
cfg.QueueID = saved.QueueID
}
}
if cfg.URL == "" {
return sampleConfig{}, errors.New("-url or ORBIT_URL is required")
}
if cfg.Token == "" {
return sampleConfig{}, errors.New("-token, ORBIT_TOKEN or ORBIT_CREDENTIALS_FILE is required")
}
if !validID(cfg.QueueID) {
return sampleConfig{}, errors.New("-queue or ORBIT_QUEUE_ID must be a 32-character lowercase hexadecimal ID")
}
if cfg.Requests < 1 || cfg.Requests > 100000 {
return sampleConfig{}, errors.New("-requests must be between 1 and 100000")
}
if cfg.Concurrency < 1 || cfg.Concurrency > 1024 {
return sampleConfig{}, errors.New("-concurrency must be between 1 and 1024")
}
if cfg.Timeout <= 0 || cfg.Timeout > time.Hour {
return sampleConfig{}, errors.New("-timeout must be greater than zero and no more than one hour")
}
for name, value := range map[string]string{
"handler": cfg.Handler,
"handler-version": cfg.HandlerVersion,
"payload-version": cfg.PayloadVersion,
} {
if strings.TrimSpace(value) == "" || strings.ContainsAny(value, "\r\n") {
return sampleConfig{}, fmt.Errorf("-%s must be a non-empty single line", name)
}
}
if cfg.Concurrency > cfg.Requests {
cfg.Concurrency = cfg.Requests
}
return cfg, nil
}
func runSamples(client *orbit.Client, cfg sampleConfig, runID string) []sampleResult {
jobs := make(chan int)
results := make(chan sampleResult, cfg.Requests)
var wait sync.WaitGroup
for worker := 0; worker < cfg.Concurrency; worker++ {
wait.Add(1)
go func() {
defer wait.Done()
for sequence := range jobs {
start := time.Now()
requestCtx, cancel := context.WithTimeout(context.Background(), cfg.Timeout)
payload, _ := json.Marshal(struct {
RunID string `json:"run_id"`
Sequence int `json:"sequence"`
}{RunID: runID, Sequence: sequence})
_, err := client.SubmitTask(requestCtx, fmt.Sprintf("orbit-load-%s-%06d", runID, sequence), orbit.SubmitTaskRequest{
QueueID: cfg.QueueID,
Handler: cfg.Handler,
HandlerVersion: cfg.HandlerVersion,
PayloadVersion: cfg.PayloadVersion,
Payload: payload,
TaskTimeoutMS: 60 * 60 * 1000,
RecoverySafety: orbit.RecoveryReconcileOnly,
Retry: orbit.RetryPolicy{MaxAttempts: 1, AttemptTimeoutMS: 10 * 1000},
})
cancel()
code := ""
if err != nil {
code = errorCode(err)
}
results <- sampleResult{duration: time.Since(start), code: code}
}
}()
}
for sequence := 0; sequence < cfg.Requests; sequence++ {
jobs <- sequence
}
close(jobs)
wait.Wait()
close(results)
out := make([]sampleResult, 0, cfg.Requests)
for result := range results {
out = append(out, result)
}
return out
}
func makeReport(cfg sampleConfig, started, finished time.Time, results []sampleResult) loadReport {
durations := make([]time.Duration, 0, len(results))
errorCodes := make(map[string]int)
success := 0
for _, result := range results {
durations = append(durations, result.duration)
if result.code == "" {
success++
} else {
errorCodes[result.code]++
}
}
sort.Slice(durations, func(i, j int) bool { return durations[i] < durations[j] })
return loadReport{
StartedAt: started, FinishedAt: finished,
Requests: cfg.Requests, Concurrency: cfg.Concurrency,
Success: success, Errors: len(results) - success,
P50MS: percentileMS(durations, 0.50), P95MS: percentileMS(durations, 0.95),
P99MS: percentileMS(durations, 0.99), MaxMS: percentileMS(durations, 1),
ErrorCodes: errorCodes,
}
}
func percentileMS(values []time.Duration, fraction float64) float64 {
if len(values) == 0 {
return 0
}
index := int(math.Ceil(float64(len(values))*fraction)) - 1
if index < 0 {
index = 0
}
if index >= len(values) {
index = len(values) - 1
}
return float64(values[index]) / float64(time.Millisecond)
}
func errorCode(err error) string {
var apiErr *orbit.APIError
if errors.As(err, &apiErr) {
if apiErr.Code != "" {
return apiErr.Code
}
return fmt.Sprintf("HTTP_%d", apiErr.StatusCode)
}
var httpErr *orbit.HTTPError
if errors.As(err, &httpErr) {
return fmt.Sprintf("HTTP_%d", httpErr.StatusCode)
}
if errors.Is(err, context.DeadlineExceeded) {
return "TIMEOUT"
}
return "TRANSPORT"
}
func newRunID() (string, error) {
var bytes [16]byte
if _, err := rand.Read(bytes[:]); err != nil {
return "", fmt.Errorf("generate load run id: %w", err)
}
return hex.EncodeToString(bytes[:]), nil
}
func validID(value string) bool {
if len(value) != 32 {
return false
}
for _, ch := range value {
if !(ch >= '0' && ch <= '9' || ch >= 'a' && ch <= 'f') {
return false
}
}
return true
}