301 lines
8.8 KiB
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
|
|
}
|