feat: add bounded admission load tool
This commit is contained in:
300
sdk/go/cmd/orbit-load/main.go
Normal file
300
sdk/go/cmd/orbit-load/main.go
Normal file
@@ -0,0 +1,300 @@
|
||||
// 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
|
||||
}
|
||||
79
sdk/go/cmd/orbit-load/main_test.go
Normal file
79
sdk/go/cmd/orbit-load/main_test.go
Normal file
@@ -0,0 +1,79 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
orbit "git.yuntuo.hk/k2safe/orbit/sdk/go"
|
||||
)
|
||||
|
||||
func TestPercentileMSUsesNearestRankAndEmptyIsZero(t *testing.T) {
|
||||
values := []time.Duration{10 * time.Millisecond, 20 * time.Millisecond, 30 * time.Millisecond, 40 * time.Millisecond}
|
||||
if got := percentileMS(values, 0.50); got != 20 {
|
||||
t.Fatalf("p50=%v, want 20", got)
|
||||
}
|
||||
if got := percentileMS(values, 0.95); got != 40 {
|
||||
t.Fatalf("p95=%v, want 40", got)
|
||||
}
|
||||
if got := percentileMS(nil, 0.99); got != 0 {
|
||||
t.Fatalf("empty percentile=%v, want 0", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestParseConfigClampsWorkersAndRejectsInvalidQueue(t *testing.T) {
|
||||
queue := "0123456789abcdef0123456789abcdef"
|
||||
cfg, err := parseConfig([]string{"-url", "http://127.0.0.1:8080", "-token", "token", "-queue", queue, "-requests", "3", "-concurrency", "8"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if cfg.Concurrency != 3 {
|
||||
t.Fatalf("concurrency=%d, want clamp to requests", cfg.Concurrency)
|
||||
}
|
||||
if _, err := parseConfig([]string{"-url", "http://127.0.0.1:8080", "-token", "token", "-queue", "bad"}); err == nil {
|
||||
t.Fatal("invalid queue accepted")
|
||||
}
|
||||
}
|
||||
|
||||
func TestErrorCodeDoesNotExposeTransportText(t *testing.T) {
|
||||
if got := errorCode(&orbit.APIError{StatusCode: 409, Code: "IDEMPOTENCY_CONFLICT"}); got != "IDEMPOTENCY_CONFLICT" {
|
||||
t.Fatalf("api code=%q", got)
|
||||
}
|
||||
if got := errorCode(context.DeadlineExceeded); got != "TIMEOUT" {
|
||||
t.Fatalf("deadline code=%q", got)
|
||||
}
|
||||
if got := errorCode(errors.New("secret transport detail")); got != "TRANSPORT" {
|
||||
t.Fatalf("transport code=%q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRunSamplesUsesOneRequestPerSequenceWithoutRetry(t *testing.T) {
|
||||
var requests atomic.Int32
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.Method != http.MethodPost || r.URL.Path != "/v1/tasks" {
|
||||
t.Fatalf("unexpected request %s %s", r.Method, r.URL.Path)
|
||||
}
|
||||
requests.Add(1)
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_, _ = w.Write([]byte(`{"request_id":"0123456789abcdef0123456789abcdef","data":{"task":{},"deduplicated":false}}`))
|
||||
}))
|
||||
defer server.Close()
|
||||
client, err := orbit.NewClient(server.URL, orbit.Options{Token: "test-token"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
cfg := sampleConfig{QueueID: "0123456789abcdef0123456789abcdef", Requests: 7, Concurrency: 3, Timeout: time.Second, Handler: "load.noop", HandlerVersion: "v1", PayloadVersion: "v1"}
|
||||
results := runSamples(client, cfg, "0123456789abcdef0123456789abcdef")
|
||||
if len(results) != cfg.Requests || int(requests.Load()) != cfg.Requests {
|
||||
t.Fatalf("results=%d requests=%d, want %d", len(results), requests.Load(), cfg.Requests)
|
||||
}
|
||||
for _, result := range results {
|
||||
if result.code != "" {
|
||||
t.Fatalf("unexpected request failure: %s", result.code)
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user