feat: add dual-host IM pressure test tool
This commit is contained in:
richjohnrambo
@@ -0,0 +1,294 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"crypto/sha256"
|
||||
"encoding/hex"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"os"
|
||||
"sort"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
const manifestVersion = 1
|
||||
|
||||
type Config struct {
|
||||
Tenant string `json:"tenant"`
|
||||
HTTPBase string `json:"http_base"`
|
||||
APIKeyEnv string `json:"api_key_env"`
|
||||
PasswordEnv string `json:"password_env"`
|
||||
UsersPerMachine int `json:"users_per_machine"`
|
||||
Endpoints map[string][]string `json:"endpoints"`
|
||||
InsecureTLS bool `json:"insecure_tls"`
|
||||
PrepareWorkers int `json:"prepare_workers"`
|
||||
ConnectRamp string `json:"connect_ramp"`
|
||||
CommandTimeout string `json:"command_timeout"`
|
||||
P2PMessages int `json:"p2p_messages_per_user"`
|
||||
P2PInterval string `json:"p2p_interval"`
|
||||
DrainP2P string `json:"drain_p2p"`
|
||||
DrainGroup string `json:"drain_group"`
|
||||
}
|
||||
|
||||
func loadConfig(path string) (*Config, error) {
|
||||
raw, err := os.ReadFile(path)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
var cfg Config
|
||||
if err = json.Unmarshal(raw, &cfg); err != nil {
|
||||
return nil, fmt.Errorf("decode config: %w", err)
|
||||
}
|
||||
if cfg.Tenant == "" {
|
||||
cfg.Tenant = "006"
|
||||
}
|
||||
if cfg.UsersPerMachine <= 0 {
|
||||
cfg.UsersPerMachine = 1000
|
||||
}
|
||||
if cfg.PrepareWorkers <= 0 {
|
||||
cfg.PrepareWorkers = 20
|
||||
}
|
||||
if cfg.ConnectRamp == "" {
|
||||
cfg.ConnectRamp = "20s"
|
||||
}
|
||||
if cfg.CommandTimeout == "" {
|
||||
cfg.CommandTimeout = "20s"
|
||||
}
|
||||
if cfg.P2PMessages <= 0 {
|
||||
cfg.P2PMessages = 10
|
||||
}
|
||||
if cfg.P2PInterval == "" {
|
||||
cfg.P2PInterval = "1s"
|
||||
}
|
||||
if cfg.DrainP2P == "" {
|
||||
cfg.DrainP2P = "30s"
|
||||
}
|
||||
if cfg.DrainGroup == "" {
|
||||
cfg.DrainGroup = "60s"
|
||||
}
|
||||
if strings.TrimSpace(cfg.HTTPBase) == "" {
|
||||
return nil, errors.New("http_base is required")
|
||||
}
|
||||
for _, machine := range []string{"a", "b"} {
|
||||
if len(cfg.Endpoints[machine]) == 0 {
|
||||
return nil, fmt.Errorf("at least one endpoint is required for machine %s", machine)
|
||||
}
|
||||
}
|
||||
return &cfg, nil
|
||||
}
|
||||
|
||||
func (c *Config) apiKey() (string, error) {
|
||||
name := c.APIKeyEnv
|
||||
if name == "" {
|
||||
name = "IM_API_KEY"
|
||||
}
|
||||
value := strings.TrimSpace(os.Getenv(name))
|
||||
if value == "" {
|
||||
return "", fmt.Errorf("environment variable %s is required", name)
|
||||
}
|
||||
return value, nil
|
||||
}
|
||||
|
||||
func (c *Config) password() (string, error) {
|
||||
name := c.PasswordEnv
|
||||
if name == "" {
|
||||
name = "IM_LOAD_PASSWORD"
|
||||
}
|
||||
value := os.Getenv(name)
|
||||
if value == "" {
|
||||
return "", fmt.Errorf("environment variable %s is required", name)
|
||||
}
|
||||
return value, nil
|
||||
}
|
||||
|
||||
func configDuration(value string) (time.Duration, error) {
|
||||
d, err := time.ParseDuration(value)
|
||||
if err != nil || d <= 0 {
|
||||
return 0, fmt.Errorf("invalid positive duration %q", value)
|
||||
}
|
||||
return d, nil
|
||||
}
|
||||
|
||||
type Manifest struct {
|
||||
Version int `json:"version"`
|
||||
Tenant string `json:"tenant"`
|
||||
GroupTopic string `json:"group_topic,omitempty"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
Accounts []AccountRecord `json:"accounts"`
|
||||
}
|
||||
|
||||
type AccountRecord struct {
|
||||
Machine string `json:"machine"`
|
||||
Index int `json:"index"`
|
||||
Username string `json:"username"`
|
||||
UserID string `json:"user_id"`
|
||||
Token string `json:"token"`
|
||||
TokenExpires time.Time `json:"token_expires,omitempty"`
|
||||
P2PTopic string `json:"p2p_topic,omitempty"`
|
||||
Endpoint string `json:"endpoint,omitempty"`
|
||||
}
|
||||
|
||||
func newManifest(cfg *Config) *Manifest {
|
||||
m := &Manifest{Version: manifestVersion, Tenant: cfg.Tenant, CreatedAt: time.Now().UTC()}
|
||||
for _, machine := range []string{"a", "b"} {
|
||||
for i := 1; i <= cfg.UsersPerMachine; i++ {
|
||||
m.Accounts = append(m.Accounts, AccountRecord{
|
||||
Machine: machine,
|
||||
Index: i,
|
||||
Username: fmt.Sprintf("lt%s%04d", machine, i),
|
||||
Endpoint: cfg.Endpoints[machine][(i-1)%len(cfg.Endpoints[machine])],
|
||||
})
|
||||
}
|
||||
}
|
||||
return m
|
||||
}
|
||||
|
||||
func applyConfigEndpoints(cfg *Config, manifest *Manifest) error {
|
||||
if manifest.Tenant != cfg.Tenant {
|
||||
return fmt.Errorf("manifest tenant=%s config tenant=%s", manifest.Tenant, cfg.Tenant)
|
||||
}
|
||||
if len(manifest.Accounts) != cfg.UsersPerMachine*2 {
|
||||
return fmt.Errorf("manifest accounts=%d expected=%d", len(manifest.Accounts), cfg.UsersPerMachine*2)
|
||||
}
|
||||
for i := range manifest.Accounts {
|
||||
account := &manifest.Accounts[i]
|
||||
endpoints := cfg.Endpoints[account.Machine]
|
||||
if account.Index <= 0 || len(endpoints) == 0 {
|
||||
return fmt.Errorf("invalid endpoint assignment for %s", accountKey(account.Machine, account.Index))
|
||||
}
|
||||
account.Endpoint = endpoints[(account.Index-1)%len(endpoints)]
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func loadManifest(path string) (*Manifest, error) {
|
||||
raw, err := os.ReadFile(path)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
var manifest Manifest
|
||||
if err = json.Unmarshal(raw, &manifest); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if manifest.Version != manifestVersion {
|
||||
return nil, fmt.Errorf("unsupported manifest version %d", manifest.Version)
|
||||
}
|
||||
return &manifest, manifest.validate()
|
||||
}
|
||||
|
||||
func (m *Manifest) validate() error {
|
||||
seen := make(map[string]bool, len(m.Accounts))
|
||||
for _, account := range m.Accounts {
|
||||
key := accountKey(account.Machine, account.Index)
|
||||
if account.Machine != "a" && account.Machine != "b" {
|
||||
return fmt.Errorf("invalid machine %q", account.Machine)
|
||||
}
|
||||
if account.Index <= 0 || account.Username == "" || seen[key] {
|
||||
return fmt.Errorf("invalid or duplicate account %s", key)
|
||||
}
|
||||
seen[key] = true
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func saveManifest(path string, manifest *Manifest) error {
|
||||
if err := manifest.validate(); err != nil {
|
||||
return err
|
||||
}
|
||||
raw, err := json.MarshalIndent(manifest, "", " ")
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
tmp := path + ".tmp"
|
||||
if err = os.WriteFile(tmp, append(raw, '\n'), 0o600); err != nil {
|
||||
return err
|
||||
}
|
||||
return os.Rename(tmp, path)
|
||||
}
|
||||
|
||||
func accountKey(machine string, index int) string {
|
||||
return strings.ToLower(machine) + fmt.Sprintf("-%04d", index)
|
||||
}
|
||||
|
||||
func accountsForMachine(manifest *Manifest, machine string) []AccountRecord {
|
||||
var out []AccountRecord
|
||||
for _, account := range manifest.Accounts {
|
||||
if account.Machine == machine {
|
||||
out = append(out, account)
|
||||
}
|
||||
}
|
||||
sort.Slice(out, func(i, j int) bool { return out[i].Index < out[j].Index })
|
||||
return out
|
||||
}
|
||||
|
||||
func manifestAccountMap(manifest *Manifest) map[string]*AccountRecord {
|
||||
out := make(map[string]*AccountRecord, len(manifest.Accounts))
|
||||
for i := range manifest.Accounts {
|
||||
account := &manifest.Accounts[i]
|
||||
out[accountKey(account.Machine, account.Index)] = account
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func topicHash(topic string) string {
|
||||
sum := sha256.Sum256([]byte(topic))
|
||||
return hex.EncodeToString(sum[:6])
|
||||
}
|
||||
|
||||
type AcceptedMessage struct {
|
||||
RunID string `json:"run_id"`
|
||||
Scenario string `json:"scenario"`
|
||||
Kind string `json:"kind"`
|
||||
ClientMid string `json:"client_mid"`
|
||||
Topic string `json:"topic"`
|
||||
Seq int64 `json:"seq"`
|
||||
SenderMachine string `json:"sender_machine"`
|
||||
SenderIndex int `json:"sender_index"`
|
||||
PeerMachine string `json:"peer_machine,omitempty"`
|
||||
PeerIndex int `json:"peer_index,omitempty"`
|
||||
SentAt time.Time `json:"sent_at"`
|
||||
AckAt time.Time `json:"ack_at"`
|
||||
}
|
||||
|
||||
type AgentSummary struct {
|
||||
RunID string `json:"run_id"`
|
||||
Scenario string `json:"scenario"`
|
||||
Machine string `json:"machine"`
|
||||
StartedAt time.Time `json:"started_at"`
|
||||
LoadEndedAt time.Time `json:"load_ended_at"`
|
||||
FinishedAt time.Time `json:"finished_at"`
|
||||
ReadyConnections int `json:"ready_connections"`
|
||||
PublishAttempted int64 `json:"publish_attempted"`
|
||||
PublishAccepted int64 `json:"publish_accepted"`
|
||||
PublishRejected int64 `json:"publish_rejected"`
|
||||
ExpectedPeer int64 `json:"expected_peer"`
|
||||
ReceivedPeer int64 `json:"received_peer"`
|
||||
ExpectedEcho int64 `json:"expected_echo"`
|
||||
ReceivedEcho int64 `json:"received_echo"`
|
||||
ExpectedGroup int64 `json:"expected_group"`
|
||||
ReceivedGroup int64 `json:"received_group"`
|
||||
RealtimeLost int64 `json:"realtime_lost"`
|
||||
RealtimeLostStored int64 `json:"realtime_lost_but_persisted"`
|
||||
Duplicates int64 `json:"duplicates"`
|
||||
OutOfOrder int64 `json:"out_of_order"`
|
||||
Disconnects int64 `json:"disconnects"`
|
||||
UnexpectedMessages int64 `json:"unexpected_messages"`
|
||||
HistoryChecked int64 `json:"history_checked"`
|
||||
HistoryFound int64 `json:"history_found"`
|
||||
PersistenceLost int64 `json:"persistence_lost"`
|
||||
AckLatency HistogramSnapshot `json:"ack_latency"`
|
||||
DeliveryLatency HistogramSnapshot `json:"delivery_latency"`
|
||||
Errors []string `json:"errors,omitempty"`
|
||||
Missing []MissingDelivery `json:"missing,omitempty"`
|
||||
PersistenceMissing []string `json:"persistence_missing,omitempty"`
|
||||
AcceptedLocal []AcceptedMessage `json:"accepted_local,omitempty"`
|
||||
}
|
||||
|
||||
type MissingDelivery struct {
|
||||
ClientMid string `json:"client_mid"`
|
||||
Kind string `json:"kind"`
|
||||
TopicHash string `json:"topic_hash"`
|
||||
Seq int64 `json:"seq"`
|
||||
ReceiverIndex int `json:"receiver_index"`
|
||||
}
|
||||
Reference in New Issue
Block a user