feat: add dual-host IM pressure test tool
This commit is contained in:
richjohnrambo
@@ -0,0 +1,7 @@
|
||||
/im-pressure
|
||||
/config.json
|
||||
/manifest*.json
|
||||
/results/
|
||||
/*-summary.json
|
||||
/report.json
|
||||
/metrics.jsonl
|
||||
@@ -0,0 +1,123 @@
|
||||
# IM 双机压力测试工具
|
||||
|
||||
该工具按商户 `006` 的固定模型运行：两台压测机各承载 1000 个 WebSocket 用户，共 2000 个用户、1000 对好友和一个 2000 人群。
|
||||
|
||||
## 统计口径
|
||||
|
||||
- 发送成功以 `{ctrl code=2xx, params.seq>0}` 为准，形成两台机器共享的 ACK 台账。
|
||||
- 实时到达只统计同时匹配本轮 `run_id` 和 `x-im-client-mid` 的 `{data}`。
|
||||
- 群消息设置 `noecho=false`，每条服务端接受的群消息期望 2000 个连接各收到一次。
|
||||
- 实时 drain 结束后通过 `{get what=data}` 按 ACK 的 Topic/Seq 回查历史，区分跨 Pod 实时漏投与数据库持久化缺失。
|
||||
- `manifest.json` 含登录 token，文件权限为 `0600`，不得提交到 Git 或放入结果包。
|
||||
|
||||
## 构建与配置
|
||||
|
||||
两台机器使用同一版本二进制、配置和 manifest：
|
||||
|
||||
```bash
|
||||
go build -o loadtest/im-pressure/im-pressure ./loadtest/im-pressure
|
||||
cp loadtest/im-pressure/config.example.json loadtest/im-pressure/config.json
|
||||
export IM_API_KEY='...'
|
||||
export IM_LOAD_PASSWORD='...'
|
||||
```
|
||||
|
||||
修改 `config.json`：`http_base` 指向 IM HTTP API；`endpoints.a/b` 分别填写两台机器应直连的 WS Pod/LB 地址，多个地址时用户会轮询分配。正式测试保持 `insecure_tls=false`。
|
||||
|
||||
## 一次性准备
|
||||
|
||||
只在一台机器执行。命令可重复运行，会复用已有账号、好友和群：
|
||||
|
||||
```bash
|
||||
./loadtest/im-pressure/im-pressure prepare \
|
||||
-config loadtest/im-pressure/config.json \
|
||||
-manifest loadtest/im-pressure/manifest.json
|
||||
|
||||
./loadtest/im-pressure/im-pressure verify \
|
||||
-config loadtest/im-pressure/config.json \
|
||||
-manifest loadtest/im-pressure/manifest.json
|
||||
```
|
||||
|
||||
将 `config.json` 和 `manifest.json` 安全复制到机器 B，并执行 `chmod 600 loadtest/im-pressure/manifest.json`。
|
||||
|
||||
## 双机运行
|
||||
|
||||
先在机器 A 启动协调器。`run-id` 每轮必须唯一且不超过 24 个字符：
|
||||
|
||||
```bash
|
||||
mkdir -p loadtest/im-pressure/results/p2p-001
|
||||
./loadtest/im-pressure/im-pressure coordinator \
|
||||
-listen :8090 -run-id p2p-001 -scenario p2p \
|
||||
-output loadtest/im-pressure/results/p2p-001/coordinator.json
|
||||
```
|
||||
|
||||
然后分别启动两个 agent。`COORDINATOR_IP` 必须是机器 B 可访问的机器 A 内网地址。
|
||||
|
||||
机器 A：
|
||||
|
||||
```bash
|
||||
./loadtest/im-pressure/im-pressure agent \
|
||||
-config loadtest/im-pressure/config.json \
|
||||
-manifest loadtest/im-pressure/manifest.json \
|
||||
-machine a -coordinator http://COORDINATOR_IP:8090 \
|
||||
-metrics-listen :9105 \
|
||||
-output loadtest/im-pressure/results/p2p-001/agent-a-summary.json
|
||||
```
|
||||
|
||||
机器 B：
|
||||
|
||||
```bash
|
||||
./loadtest/im-pressure/im-pressure agent \
|
||||
-config loadtest/im-pressure/config.json \
|
||||
-manifest loadtest/im-pressure/manifest.json \
|
||||
-machine b -coordinator http://COORDINATOR_IP:8090 \
|
||||
-metrics-listen :9105 \
|
||||
-output loadtest/im-pressure/results/p2p-001/agent-b-summary.json
|
||||
```
|
||||
|
||||
两个 agent 都建立 1000 条连接后，协调器下发统一的 UTC `start_at`。场景可替换为：
|
||||
|
||||
- `p2p`：1000 对用户两侧同时发送，默认每人 10 条、间隔 1 秒。
|
||||
- `group`：每台机器 500 个发送者，每人 2 条；全局每 100ms 一条，共 2000 条。
|
||||
- `mixed`：300 对私聊用户每人 5 条，同时每台 250 个群发送者各发 2 条。
|
||||
|
||||
## 合并报告
|
||||
|
||||
把两个 agent summary 放到同一台机器：
|
||||
|
||||
```bash
|
||||
./loadtest/im-pressure/im-pressure merge \
|
||||
-a loadtest/im-pressure/results/p2p-001/agent-a-summary.json \
|
||||
-b loadtest/im-pressure/results/p2p-001/agent-b-summary.json \
|
||||
-output loadtest/im-pressure/results/p2p-001/report.json
|
||||
```
|
||||
|
||||
`report.json` 给出 ACK、实时投递、重复、乱序、断连和历史持久化核验结果。当前正确性门禁要求上述错误计数全部为 0。
|
||||
|
||||
每个 agent 在 `:9105/metrics` 暴露连接、发送、接收、重复、乱序、断连以及 Go/process 指标；端口冲突时可通过 `-metrics-listen` 修改，传空字符串可关闭。
|
||||
|
||||
## 指标采集
|
||||
|
||||
`collect` 可并行抓取 Prometheus、Redis exporter、Kafka exporter、mysqld exporter、kube-state-metrics/cAdvisor 和压测机 node-exporter，原始响应按时间写入 JSONL：
|
||||
|
||||
```bash
|
||||
./loadtest/im-pressure/im-pressure collect \
|
||||
-duration 10m -interval 5s \
|
||||
-url im=http://PROMETHEUS_HOST/api/v1/query?query=up \
|
||||
-url redis=http://REDIS_EXPORTER/metrics \
|
||||
-url kafka=http://KAFKA_EXPORTER/metrics \
|
||||
-url mysql=http://MYSQLD_EXPORTER/metrics \
|
||||
-url node-a=http://NODE_A_EXPORTER/metrics \
|
||||
-url node-b=http://NODE_B_EXPORTER/metrics \
|
||||
-output loadtest/im-pressure/results/p2p-001/metrics.jsonl
|
||||
```
|
||||
|
||||
URL 含 `&` 时需用单引号包住整个 `name=url` 参数。建议起压前 2 分钟开始采集，结束后继续 5 分钟，以观察 Kafka lag、数据库落库追平时间及资源恢复。
|
||||
|
||||
## 运行前检查
|
||||
|
||||
```bash
|
||||
ulimit -n 100000
|
||||
curl -fsS http://COORDINATOR_IP:8090/v1/status
|
||||
```
|
||||
|
||||
两台压测机应完成时钟同步，并确认压测出口、Ingress、Pod 与 Redis/MySQL/Kafka 连接数限制可以容纳本轮负载。协调器没有鉴权，只能监听在受控压测网络中。
|
||||
@@ -0,0 +1,461 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"log"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"regexp"
|
||||
"sort"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
)
|
||||
|
||||
var runIDPattern = regexp.MustCompile(`^[a-zA-Z0-9_-]{1,24}$`)
|
||||
|
||||
type agentRun struct {
|
||||
cfg *Config
|
||||
machine string
|
||||
runID string
|
||||
scenario string
|
||||
clients map[int]*WSClient
|
||||
recorder *deliveryRecorder
|
||||
ack *DurationHistogram
|
||||
attempted atomic.Int64
|
||||
rejected atomic.Int64
|
||||
ready atomic.Int64
|
||||
|
||||
mu sync.Mutex
|
||||
accepted []AcceptedMessage
|
||||
errors []string
|
||||
}
|
||||
|
||||
func validScenario(scenario string) bool {
|
||||
return scenario == "p2p" || scenario == "group" || scenario == "mixed"
|
||||
}
|
||||
|
||||
func runAgent(ctx context.Context, cfg *Config, manifest *Manifest, machine, coordinatorURL, output, metricsListen string) error {
|
||||
if machine != "a" && machine != "b" {
|
||||
return fmt.Errorf("machine must be a or b")
|
||||
}
|
||||
apiKey, err := cfg.apiKey()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
password := ""
|
||||
if name := cfg.PasswordEnv; name != "" {
|
||||
password = os.Getenv(name)
|
||||
} else {
|
||||
password = os.Getenv("IM_LOAD_PASSWORD")
|
||||
}
|
||||
accounts := accountsForMachine(manifest, machine)
|
||||
if len(accounts) != cfg.UsersPerMachine {
|
||||
return fmt.Errorf("machine %s accounts=%d expected=%d", machine, len(accounts), cfg.UsersPerMachine)
|
||||
}
|
||||
for i := range accounts {
|
||||
accounts[i].Endpoint = cfg.Endpoints[machine][(accounts[i].Index-1)%len(cfg.Endpoints[machine])]
|
||||
}
|
||||
coord := newCoordinatorClient(coordinatorURL)
|
||||
placeholderRunID := "pending"
|
||||
run := &agentRun{cfg: cfg, machine: machine, clients: make(map[int]*WSClient), ack: NewDurationHistogram()}
|
||||
run.recorder = newDeliveryRecorder(placeholderRunID, machine, len(accounts))
|
||||
stopMetrics, err := startAgentMetrics(metricsListen, run)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer stopMetrics()
|
||||
if err = run.connect(ctx, accounts, manifest.GroupTopic, apiKey, password); err != nil {
|
||||
run.close()
|
||||
return err
|
||||
}
|
||||
defer run.close()
|
||||
start, err := coord.ready(ctx, AgentReady{Machine: machine, ReadyConnections: len(run.clients)})
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if !runIDPattern.MatchString(start.RunID) || !validScenario(start.Scenario) {
|
||||
return fmt.Errorf("invalid coordinator start: %+v", start)
|
||||
}
|
||||
run.runID = start.RunID
|
||||
run.scenario = start.Scenario
|
||||
run.recorder.setRunID(start.RunID)
|
||||
if err = waitUntil(ctx, start.StartAt); err != nil {
|
||||
return err
|
||||
}
|
||||
startedAt := time.Now().UTC()
|
||||
log.Printf("agent=%s scenario=%s started", machine, start.Scenario)
|
||||
if err = run.execute(ctx, start.StartAt, manifest.GroupTopic); err != nil {
|
||||
return err
|
||||
}
|
||||
loadEndedAt := time.Now().UTC()
|
||||
local := run.acceptedCopy()
|
||||
if err = coord.submitLedger(ctx, machine, local); err != nil {
|
||||
return err
|
||||
}
|
||||
ledger, err := coord.ledger(ctx)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
drain, err := run.drainDuration()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if err = waitContext(ctx, drain); err != nil {
|
||||
return err
|
||||
}
|
||||
base := AgentSummary{
|
||||
RunID: start.RunID, Scenario: start.Scenario, Machine: machine, StartedAt: startedAt,
|
||||
LoadEndedAt: loadEndedAt, FinishedAt: time.Now().UTC(), ReadyConnections: len(run.clients),
|
||||
PublishAttempted: run.attempted.Load(), PublishAccepted: int64(len(local)), PublishRejected: run.rejected.Load(),
|
||||
AckLatency: run.ack.Snapshot(), AcceptedLocal: local,
|
||||
}
|
||||
run.mu.Lock()
|
||||
base.Errors = append(base.Errors, run.errors...)
|
||||
run.mu.Unlock()
|
||||
summary := run.recorder.summarize(base, ledger)
|
||||
checked, found, persistenceMissing := run.reconcileHistory(ctx, ledger)
|
||||
summary.HistoryChecked = int64(checked)
|
||||
summary.HistoryFound = int64(found)
|
||||
summary.PersistenceLost = int64(len(persistenceMissing))
|
||||
summary.RealtimeLostStored = run.recorder.countRealtimeLostStored(ledger)
|
||||
if len(persistenceMissing) > 10000 {
|
||||
summary.PersistenceMissing = append([]string(nil), persistenceMissing[:10000]...)
|
||||
summary.Errors = append(summary.Errors, "persistence missing details truncated at 10000")
|
||||
} else {
|
||||
summary.PersistenceMissing = persistenceMissing
|
||||
}
|
||||
summary.FinishedAt = time.Now().UTC()
|
||||
if err = writeAgentSummary(output, summary); err != nil {
|
||||
return err
|
||||
}
|
||||
if err = coord.submitSummary(ctx, summary); err != nil {
|
||||
return err
|
||||
}
|
||||
log.Printf("agent=%s accepted=%d expected_peer=%d received_peer=%d expected_group=%d received_group=%d lost=%d duplicates=%d",
|
||||
machine, summary.PublishAccepted, summary.ExpectedPeer, summary.ReceivedPeer,
|
||||
summary.ExpectedGroup, summary.ReceivedGroup, summary.RealtimeLost, summary.Duplicates)
|
||||
return nil
|
||||
}
|
||||
|
||||
type historyQuery struct {
|
||||
client *WSClient
|
||||
topic string
|
||||
since int64
|
||||
before int64
|
||||
limit int
|
||||
}
|
||||
|
||||
func (r *agentRun) reconcileHistory(ctx context.Context, ledger []AcceptedMessage) (int, int, []string) {
|
||||
expected := make(map[string]AcceptedMessage)
|
||||
byTopic := make(map[string][]AcceptedMessage)
|
||||
for _, message := range ledger {
|
||||
if message.RunID != r.runID {
|
||||
continue
|
||||
}
|
||||
if message.Kind == "p2p" && message.PeerMachine == r.machine {
|
||||
expected[message.ClientMid] = message
|
||||
byTopic[message.Topic] = append(byTopic[message.Topic], message)
|
||||
} else if message.Kind == "group" && r.machine == "a" {
|
||||
expected[message.ClientMid] = message
|
||||
byTopic[message.Topic] = append(byTopic[message.Topic], message)
|
||||
}
|
||||
}
|
||||
if len(expected) == 0 {
|
||||
return 0, 0, nil
|
||||
}
|
||||
r.recorder.beginHistory()
|
||||
queries := make([]historyQuery, 0, len(byTopic))
|
||||
for topic, messages := range byTopic {
|
||||
sort.Slice(messages, func(i, j int) bool { return messages[i].Seq < messages[j].Seq })
|
||||
client := r.clients[messages[0].PeerIndex]
|
||||
if messages[0].Kind == "group" {
|
||||
client = r.clients[1]
|
||||
}
|
||||
if client == nil {
|
||||
r.addError(fmt.Sprintf("history topic=%s has no verifier client", topicHash(topic)))
|
||||
continue
|
||||
}
|
||||
for start := 0; start < len(messages); start += 200 {
|
||||
end := start + 200
|
||||
if end > len(messages) {
|
||||
end = len(messages)
|
||||
}
|
||||
queries = append(queries, historyQuery{client: client, topic: topic,
|
||||
since: messages[start].Seq, before: messages[end-1].Seq + 1, limit: 256})
|
||||
}
|
||||
}
|
||||
for _, query := range queries {
|
||||
if err := query.client.GetData(ctx, query.topic, query.since, query.before, query.limit); err != nil {
|
||||
r.addError(fmt.Sprintf("history get topic=%s range=%d-%d: %v", topicHash(query.topic), query.since, query.before, err))
|
||||
}
|
||||
}
|
||||
_ = waitContext(ctx, 2*time.Second)
|
||||
deadline := time.Now().Add(120 * time.Second)
|
||||
for {
|
||||
missing := missingHistory(r.recorder, expected)
|
||||
if len(missing) == 0 || time.Now().After(deadline) {
|
||||
break
|
||||
}
|
||||
if waitContext(ctx, 5*time.Second) != nil {
|
||||
break
|
||||
}
|
||||
for _, message := range missing {
|
||||
client := r.clients[message.PeerIndex]
|
||||
if message.Kind == "group" {
|
||||
client = r.clients[1]
|
||||
}
|
||||
if client != nil {
|
||||
_ = client.GetData(ctx, message.Topic, message.Seq, message.Seq+1, 1)
|
||||
}
|
||||
}
|
||||
_ = waitContext(ctx, 500*time.Millisecond)
|
||||
}
|
||||
missing := missingHistory(r.recorder, expected)
|
||||
missingIDs := make([]string, 0, len(missing))
|
||||
for _, message := range missing {
|
||||
missingIDs = append(missingIDs, message.ClientMid)
|
||||
}
|
||||
sort.Strings(missingIDs)
|
||||
return len(expected), len(expected) - len(missing), missingIDs
|
||||
}
|
||||
|
||||
func missingHistory(recorder *deliveryRecorder, expected map[string]AcceptedMessage) []AcceptedMessage {
|
||||
missing := make([]AcceptedMessage, 0)
|
||||
for mid, message := range expected {
|
||||
if !recorder.historyFound(mid) {
|
||||
missing = append(missing, message)
|
||||
}
|
||||
}
|
||||
return missing
|
||||
}
|
||||
|
||||
func (r *agentRun) addError(message string) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
if len(r.errors) < 1000 {
|
||||
r.errors = append(r.errors, message)
|
||||
}
|
||||
}
|
||||
|
||||
func (r *agentRun) connect(ctx context.Context, accounts []AccountRecord, groupTopic, apiKey, password string) error {
|
||||
ramp, err := configDuration(r.cfg.ConnectRamp)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
interval := ramp / time.Duration(len(accounts))
|
||||
var wg sync.WaitGroup
|
||||
errCh := make(chan error, len(accounts))
|
||||
var mu sync.Mutex
|
||||
for _, account := range accounts {
|
||||
account := account
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
client, connectErr := connectWSClient(ctx, r.cfg, account, apiKey, password, r.recorder.onData, r.recorder.onDisconnect)
|
||||
if connectErr != nil {
|
||||
errCh <- connectErr
|
||||
return
|
||||
}
|
||||
if connectErr = client.Subscribe(ctx, account.P2PTopic); connectErr == nil {
|
||||
connectErr = client.Subscribe(ctx, groupTopic)
|
||||
}
|
||||
if connectErr != nil {
|
||||
client.Close()
|
||||
errCh <- fmt.Errorf("subscribe %s: %w", account.Username, connectErr)
|
||||
return
|
||||
}
|
||||
mu.Lock()
|
||||
r.clients[account.Index] = client
|
||||
mu.Unlock()
|
||||
r.ready.Add(1)
|
||||
}()
|
||||
if interval > 0 {
|
||||
if err = waitContext(ctx, interval); err != nil {
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
wg.Wait()
|
||||
close(errCh)
|
||||
var failures []string
|
||||
for connectErr := range errCh {
|
||||
if len(failures) < 20 {
|
||||
failures = append(failures, connectErr.Error())
|
||||
}
|
||||
}
|
||||
if len(r.clients) != len(accounts) {
|
||||
return fmt.Errorf("ready connections=%d expected=%d failures=%v", len(r.clients), len(accounts), failures)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *agentRun) execute(ctx context.Context, start time.Time, groupTopic string) error {
|
||||
if groupTopic == "" {
|
||||
return fmt.Errorf("group topic is empty")
|
||||
}
|
||||
switch r.scenario {
|
||||
case "p2p":
|
||||
return r.scheduleP2P(ctx, start, r.cfg.UsersPerMachine, r.cfg.P2PMessages, mustDuration(r.cfg.P2PInterval))
|
||||
case "group":
|
||||
return r.scheduleGroup(ctx, start, groupTopic, 1, 500, 2)
|
||||
case "mixed":
|
||||
var wg sync.WaitGroup
|
||||
errCh := make(chan error, 2)
|
||||
wg.Add(2)
|
||||
go func() { defer wg.Done(); errCh <- r.scheduleP2P(ctx, start, 300, 5, 20*time.Second) }()
|
||||
go func() { defer wg.Done(); errCh <- r.scheduleGroup(ctx, start, groupTopic, 301, 550, 2) }()
|
||||
wg.Wait()
|
||||
close(errCh)
|
||||
for err := range errCh {
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
default:
|
||||
return fmt.Errorf("unknown scenario %s", r.scenario)
|
||||
}
|
||||
}
|
||||
|
||||
func (r *agentRun) scheduleP2P(ctx context.Context, start time.Time, users, messages int, interval time.Duration) error {
|
||||
var wg sync.WaitGroup
|
||||
for index := 1; index <= users; index++ {
|
||||
client := r.clients[index]
|
||||
if client == nil {
|
||||
return fmt.Errorf("missing client %d", index)
|
||||
}
|
||||
wg.Add(1)
|
||||
go func(index int, client *WSClient) {
|
||||
defer wg.Done()
|
||||
jitter := time.Duration((index*37)%200) * time.Millisecond
|
||||
for messageIndex := 1; messageIndex <= messages; messageIndex++ {
|
||||
target := start.Add(time.Duration(messageIndex-1)*interval + jitter)
|
||||
if waitUntil(ctx, target) != nil {
|
||||
return
|
||||
}
|
||||
mid := fmt.Sprintf("%s-p-%s-%04d-%02d", r.runID, r.machine, index, messageIndex)
|
||||
accepted, err := r.publish(ctx, client, client.account.P2PTopic, mid, messageIndex)
|
||||
if err == nil {
|
||||
accepted.PeerIndex = index
|
||||
if r.machine == "a" {
|
||||
accepted.PeerMachine = "b"
|
||||
} else {
|
||||
accepted.PeerMachine = "a"
|
||||
}
|
||||
r.addAccepted(accepted)
|
||||
}
|
||||
}
|
||||
}(index, client)
|
||||
}
|
||||
wg.Wait()
|
||||
return ctx.Err()
|
||||
}
|
||||
|
||||
func (r *agentRun) scheduleGroup(ctx context.Context, start time.Time, topic string, first, last, rounds int) error {
|
||||
if first < 1 || last > len(r.clients) || first > last {
|
||||
return fmt.Errorf("invalid group publisher range %d-%d", first, last)
|
||||
}
|
||||
count := last - first + 1
|
||||
roundDuration := time.Duration(count) * 200 * time.Millisecond
|
||||
offset := time.Duration(0)
|
||||
if r.machine == "b" {
|
||||
offset = 100 * time.Millisecond
|
||||
}
|
||||
var wg sync.WaitGroup
|
||||
for index := first; index <= last; index++ {
|
||||
client := r.clients[index]
|
||||
position := index - first
|
||||
wg.Add(1)
|
||||
go func(index, position int, client *WSClient) {
|
||||
defer wg.Done()
|
||||
for round := 1; round <= rounds; round++ {
|
||||
target := start.Add(time.Duration(round-1)*roundDuration + time.Duration(position)*200*time.Millisecond + offset)
|
||||
if waitUntil(ctx, target) != nil {
|
||||
return
|
||||
}
|
||||
mid := fmt.Sprintf("%s-g-%s-%04d-%02d", r.runID, r.machine, index, round)
|
||||
accepted, err := r.publish(ctx, client, topic, mid, round)
|
||||
if err == nil {
|
||||
r.addAccepted(accepted)
|
||||
}
|
||||
}
|
||||
}(index, position, client)
|
||||
}
|
||||
wg.Wait()
|
||||
return ctx.Err()
|
||||
}
|
||||
|
||||
func (r *agentRun) publish(ctx context.Context, client *WSClient, topic, mid string, messageIndex int) (AcceptedMessage, error) {
|
||||
r.attempted.Add(1)
|
||||
accepted, err := client.Publish(ctx, topic, mid, r.runID, r.scenario, messageIndex)
|
||||
if err != nil {
|
||||
r.rejected.Add(1)
|
||||
r.mu.Lock()
|
||||
if len(r.errors) < 1000 {
|
||||
r.errors = append(r.errors, fmt.Sprintf("publish %s: %v", mid, err))
|
||||
}
|
||||
r.mu.Unlock()
|
||||
return accepted, err
|
||||
}
|
||||
r.ack.Observe(accepted.AckAt.Sub(accepted.SentAt))
|
||||
return accepted, nil
|
||||
}
|
||||
|
||||
func (r *agentRun) addAccepted(message AcceptedMessage) {
|
||||
r.mu.Lock()
|
||||
r.accepted = append(r.accepted, message)
|
||||
r.mu.Unlock()
|
||||
}
|
||||
|
||||
func (r *agentRun) acceptedCopy() []AcceptedMessage {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
out := append([]AcceptedMessage(nil), r.accepted...)
|
||||
sort.Slice(out, func(i, j int) bool { return out[i].ClientMid < out[j].ClientMid })
|
||||
return out
|
||||
}
|
||||
|
||||
func (r *agentRun) drainDuration() (time.Duration, error) {
|
||||
if r.scenario == "p2p" {
|
||||
return configDuration(r.cfg.DrainP2P)
|
||||
}
|
||||
return configDuration(r.cfg.DrainGroup)
|
||||
}
|
||||
|
||||
func (r *agentRun) close() {
|
||||
for _, client := range r.clients {
|
||||
client.Close()
|
||||
}
|
||||
}
|
||||
|
||||
func waitUntil(ctx context.Context, target time.Time) error {
|
||||
delay := time.Until(target)
|
||||
if delay <= 0 {
|
||||
return nil
|
||||
}
|
||||
return waitContext(ctx, delay)
|
||||
}
|
||||
|
||||
func mustDuration(value string) time.Duration {
|
||||
d, err := time.ParseDuration(value)
|
||||
if err != nil || d <= 0 {
|
||||
return time.Second
|
||||
}
|
||||
return d
|
||||
}
|
||||
|
||||
func writeAgentSummary(path string, summary AgentSummary) error {
|
||||
raw, err := json.MarshalIndent(summary, "", " ")
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if dir := filepath.Dir(path); dir != "." {
|
||||
if err = os.MkdirAll(dir, 0o750); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return os.WriteFile(path, append(raw, '\n'), 0o600)
|
||||
}
|
||||
@@ -0,0 +1,68 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"net"
|
||||
"net/http"
|
||||
"time"
|
||||
|
||||
"github.com/prometheus/client_golang/prometheus"
|
||||
"github.com/prometheus/client_golang/prometheus/collectors"
|
||||
"github.com/prometheus/client_golang/prometheus/promhttp"
|
||||
)
|
||||
|
||||
func startAgentMetrics(address string, run *agentRun) (func(), error) {
|
||||
if address == "" {
|
||||
return func() {}, nil
|
||||
}
|
||||
listener, err := net.Listen("tcp", address)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("listen for agent metrics on %s: %w", address, err)
|
||||
}
|
||||
labels := prometheus.Labels{"machine": run.machine}
|
||||
registry := prometheus.NewRegistry()
|
||||
registry.MustRegister(
|
||||
collectors.NewGoCollector(),
|
||||
collectors.NewProcessCollector(collectors.ProcessCollectorOpts{}),
|
||||
prometheus.NewGaugeFunc(prometheus.GaugeOpts{Name: "im_load_agent_ready_connections", Help: "Current ready WebSocket connections.", ConstLabels: labels}, func() float64 {
|
||||
return float64(run.ready.Load())
|
||||
}),
|
||||
prometheus.NewCounterFunc(prometheus.CounterOpts{Name: "im_load_agent_publish_attempted_total", Help: "Publish commands attempted.", ConstLabels: labels}, func() float64 {
|
||||
return float64(run.attempted.Load())
|
||||
}),
|
||||
prometheus.NewCounterFunc(prometheus.CounterOpts{Name: "im_load_agent_publish_accepted_total", Help: "Publish commands acknowledged by the server.", ConstLabels: labels}, func() float64 {
|
||||
run.mu.Lock()
|
||||
defer run.mu.Unlock()
|
||||
return float64(len(run.accepted))
|
||||
}),
|
||||
prometheus.NewCounterFunc(prometheus.CounterOpts{Name: "im_load_agent_publish_rejected_total", Help: "Publish commands rejected or timed out.", ConstLabels: labels}, func() float64 {
|
||||
return float64(run.rejected.Load())
|
||||
}),
|
||||
prometheus.NewCounterFunc(prometheus.CounterOpts{Name: "im_load_agent_received_total", Help: "Unique matching data deliveries received.", ConstLabels: labels}, func() float64 {
|
||||
received, _, _, _ := run.recorder.metricCounts()
|
||||
return float64(received)
|
||||
}),
|
||||
prometheus.NewCounterFunc(prometheus.CounterOpts{Name: "im_load_agent_duplicates_total", Help: "Duplicate matching deliveries received.", ConstLabels: labels}, func() float64 {
|
||||
_, duplicates, _, _ := run.recorder.metricCounts()
|
||||
return float64(duplicates)
|
||||
}),
|
||||
prometheus.NewCounterFunc(prometheus.CounterOpts{Name: "im_load_agent_out_of_order_total", Help: "Per-connection topic sequence regressions.", ConstLabels: labels}, func() float64 {
|
||||
_, _, outOfOrder, _ := run.recorder.metricCounts()
|
||||
return float64(outOfOrder)
|
||||
}),
|
||||
prometheus.NewCounterFunc(prometheus.CounterOpts{Name: "im_load_agent_disconnects_total", Help: "Unexpected WebSocket disconnects.", ConstLabels: labels}, func() float64 {
|
||||
_, _, _, disconnects := run.recorder.metricCounts()
|
||||
return float64(disconnects)
|
||||
}),
|
||||
)
|
||||
mux := http.NewServeMux()
|
||||
mux.Handle("/metrics", promhttp.HandlerFor(registry, promhttp.HandlerOpts{}))
|
||||
server := &http.Server{Handler: mux, ReadHeaderTimeout: 5 * time.Second}
|
||||
go func() { _ = server.Serve(listener) }()
|
||||
return func() {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
|
||||
defer cancel()
|
||||
_ = server.Shutdown(ctx)
|
||||
}, nil
|
||||
}
|
||||
@@ -0,0 +1,114 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
type namedURLs []string
|
||||
|
||||
func (n *namedURLs) String() string { return strings.Join(*n, ",") }
|
||||
func (n *namedURLs) Set(value string) error {
|
||||
if !strings.Contains(value, "=") {
|
||||
return fmt.Errorf("URL must use name=https://host/path")
|
||||
}
|
||||
*n = append(*n, value)
|
||||
return nil
|
||||
}
|
||||
|
||||
type metricSnapshot struct {
|
||||
Timestamp time.Time `json:"timestamp"`
|
||||
Name string `json:"name"`
|
||||
URL string `json:"url"`
|
||||
Status int `json:"status"`
|
||||
Duration int64 `json:"duration_ms"`
|
||||
Body string `json:"body,omitempty"`
|
||||
Error string `json:"error,omitempty"`
|
||||
}
|
||||
|
||||
func runCollector(ctx context.Context, sources namedURLs, output string, interval, duration time.Duration) error {
|
||||
if len(sources) == 0 {
|
||||
return fmt.Errorf("at least one --url name=url is required")
|
||||
}
|
||||
if interval <= 0 || duration <= 0 {
|
||||
return fmt.Errorf("interval and duration must be positive")
|
||||
}
|
||||
if dir := filepath.Dir(output); dir != "." {
|
||||
if err := os.MkdirAll(dir, 0o750); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
file, err := os.OpenFile(output, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0o600)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer file.Close()
|
||||
writer := bufio.NewWriterSize(file, 256<<10)
|
||||
defer writer.Flush()
|
||||
client := &http.Client{Timeout: 10 * time.Second}
|
||||
deadline := time.Now().Add(duration)
|
||||
for {
|
||||
snapshots := collectOnce(ctx, client, sources)
|
||||
for _, snapshot := range snapshots {
|
||||
raw, _ := json.Marshal(snapshot)
|
||||
if _, err = writer.Write(append(raw, '\n')); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
if err = writer.Flush(); err != nil {
|
||||
return err
|
||||
}
|
||||
if time.Now().Add(interval).After(deadline) {
|
||||
return nil
|
||||
}
|
||||
if err = waitContext(ctx, interval); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func collectOnce(ctx context.Context, client *http.Client, sources namedURLs) []metricSnapshot {
|
||||
out := make([]metricSnapshot, len(sources))
|
||||
var wg sync.WaitGroup
|
||||
for i, source := range sources {
|
||||
i, source := i, source
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
parts := strings.SplitN(source, "=", 2)
|
||||
snapshot := metricSnapshot{Timestamp: time.Now().UTC(), Name: parts[0], URL: parts[1]}
|
||||
started := time.Now()
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, parts[1], nil)
|
||||
if err == nil {
|
||||
var response *http.Response
|
||||
response, err = client.Do(req)
|
||||
if err == nil {
|
||||
snapshot.Status = response.StatusCode
|
||||
raw, readErr := io.ReadAll(io.LimitReader(response.Body, 16<<20))
|
||||
_ = response.Body.Close()
|
||||
if readErr != nil {
|
||||
err = readErr
|
||||
} else {
|
||||
snapshot.Body = string(raw)
|
||||
}
|
||||
}
|
||||
}
|
||||
snapshot.Duration = time.Since(started).Milliseconds()
|
||||
if err != nil {
|
||||
snapshot.Error = err.Error()
|
||||
}
|
||||
out[i] = snapshot
|
||||
}()
|
||||
}
|
||||
wg.Wait()
|
||||
return out
|
||||
}
|
||||
@@ -0,0 +1,19 @@
|
||||
{
|
||||
"tenant": "006",
|
||||
"http_base": "https://im-api.example.com",
|
||||
"api_key_env": "IM_API_KEY",
|
||||
"password_env": "IM_LOAD_PASSWORD",
|
||||
"users_per_machine": 1000,
|
||||
"endpoints": {
|
||||
"a": ["wss://im-ws-pod-a.example.com/v0/channels"],
|
||||
"b": ["wss://im-ws-pod-b.example.com/v0/channels"]
|
||||
},
|
||||
"insecure_tls": false,
|
||||
"prepare_workers": 20,
|
||||
"connect_ramp": "20s",
|
||||
"command_timeout": "20s",
|
||||
"p2p_messages_per_user": 10,
|
||||
"p2p_interval": "1s",
|
||||
"drain_p2p": "30s",
|
||||
"drain_group": "60s"
|
||||
}
|
||||
@@ -0,0 +1,323 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"log"
|
||||
"net/http"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"sort"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
type StartInfo struct {
|
||||
RunID string `json:"run_id"`
|
||||
Scenario string `json:"scenario"`
|
||||
StartAt time.Time `json:"start_at"`
|
||||
}
|
||||
|
||||
type AgentReady struct {
|
||||
Machine string `json:"machine"`
|
||||
ReadyConnections int `json:"ready_connections"`
|
||||
}
|
||||
|
||||
type coordinatorState struct {
|
||||
mu sync.Mutex
|
||||
start StartInfo
|
||||
startDelay time.Duration
|
||||
expectedAgents int
|
||||
ready map[string]AgentReady
|
||||
ledgers map[string][]AcceptedMessage
|
||||
summaries map[string]AgentSummary
|
||||
complete chan struct{}
|
||||
completeOnce sync.Once
|
||||
}
|
||||
|
||||
func runCoordinator(ctx context.Context, listen, runID, scenario, output string, expectedAgents int, startDelay time.Duration) error {
|
||||
if !runIDPattern.MatchString(runID) {
|
||||
return fmt.Errorf("run-id must match %s", runIDPattern.String())
|
||||
}
|
||||
if !validScenario(scenario) {
|
||||
return fmt.Errorf("unsupported scenario %q", scenario)
|
||||
}
|
||||
if expectedAgents <= 0 {
|
||||
expectedAgents = 2
|
||||
}
|
||||
state := &coordinatorState{
|
||||
start: StartInfo{RunID: runID, Scenario: scenario}, startDelay: startDelay, expectedAgents: expectedAgents,
|
||||
ready: make(map[string]AgentReady), ledgers: make(map[string][]AcceptedMessage),
|
||||
summaries: make(map[string]AgentSummary), complete: make(chan struct{}),
|
||||
}
|
||||
mux := http.NewServeMux()
|
||||
mux.HandleFunc("/v1/ready", state.handleReady)
|
||||
mux.HandleFunc("/v1/start", state.handleStart)
|
||||
mux.HandleFunc("/v1/ledger", state.handleLedger)
|
||||
mux.HandleFunc("/v1/summary", state.handleSummary)
|
||||
mux.HandleFunc("/v1/status", state.handleStatus)
|
||||
server := &http.Server{Addr: listen, Handler: mux, ReadHeaderTimeout: 5 * time.Second}
|
||||
errCh := make(chan error, 1)
|
||||
go func() {
|
||||
log.Printf("coordinator listening=%s run_id=%s scenario=%s expected_agents=%d", listen, runID, scenario, expectedAgents)
|
||||
if err := server.ListenAndServe(); err != nil && err != http.ErrServerClosed {
|
||||
errCh <- err
|
||||
}
|
||||
}()
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
case err := <-errCh:
|
||||
return err
|
||||
case <-state.complete:
|
||||
log.Printf("all agent summaries received")
|
||||
}
|
||||
shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
defer cancel()
|
||||
_ = server.Shutdown(shutdownCtx)
|
||||
if output != "" {
|
||||
if err := state.writeOutput(output); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *coordinatorState) handleReady(w http.ResponseWriter, r *http.Request) {
|
||||
if r.Method != http.MethodPost {
|
||||
http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
|
||||
return
|
||||
}
|
||||
var ready AgentReady
|
||||
if !decodeLimitedJSON(w, r, &ready) || (ready.Machine != "a" && ready.Machine != "b") || ready.ReadyConnections <= 0 {
|
||||
if ready.Machine != "a" && ready.Machine != "b" {
|
||||
http.Error(w, "invalid machine", http.StatusBadRequest)
|
||||
}
|
||||
return
|
||||
}
|
||||
s.mu.Lock()
|
||||
s.ready[ready.Machine] = ready
|
||||
if len(s.ready) == s.expectedAgents && s.start.StartAt.IsZero() {
|
||||
s.start.StartAt = time.Now().UTC().Add(s.startDelay)
|
||||
log.Printf("all agents ready; start_at=%s", s.start.StartAt.Format(time.RFC3339Nano))
|
||||
}
|
||||
start := s.start
|
||||
s.mu.Unlock()
|
||||
writeJSON(w, http.StatusOK, start)
|
||||
}
|
||||
|
||||
func (s *coordinatorState) handleStart(w http.ResponseWriter, r *http.Request) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
if s.start.StartAt.IsZero() {
|
||||
writeJSON(w, http.StatusAccepted, map[string]any{"ready_agents": len(s.ready), "expected_agents": s.expectedAgents})
|
||||
return
|
||||
}
|
||||
writeJSON(w, http.StatusOK, s.start)
|
||||
}
|
||||
|
||||
func (s *coordinatorState) handleLedger(w http.ResponseWriter, r *http.Request) {
|
||||
machine := strings.ToLower(strings.TrimSpace(r.URL.Query().Get("machine")))
|
||||
switch r.Method {
|
||||
case http.MethodPost:
|
||||
if machine != "a" && machine != "b" {
|
||||
http.Error(w, "invalid machine", http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
var messages []AcceptedMessage
|
||||
if !decodeLimitedJSON(w, r, &messages) {
|
||||
return
|
||||
}
|
||||
s.mu.Lock()
|
||||
s.ledgers[machine] = messages
|
||||
count := len(s.ledgers)
|
||||
s.mu.Unlock()
|
||||
writeJSON(w, http.StatusOK, map[string]int{"submitted_agents": count})
|
||||
case http.MethodGet:
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
if len(s.ledgers) < s.expectedAgents {
|
||||
writeJSON(w, http.StatusAccepted, map[string]int{"submitted_agents": len(s.ledgers), "expected_agents": s.expectedAgents})
|
||||
return
|
||||
}
|
||||
var merged []AcceptedMessage
|
||||
for _, messages := range s.ledgers {
|
||||
merged = append(merged, messages...)
|
||||
}
|
||||
sort.Slice(merged, func(i, j int) bool { return merged[i].ClientMid < merged[j].ClientMid })
|
||||
writeJSON(w, http.StatusOK, merged)
|
||||
default:
|
||||
http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
|
||||
}
|
||||
}
|
||||
|
||||
func (s *coordinatorState) handleSummary(w http.ResponseWriter, r *http.Request) {
|
||||
if r.Method != http.MethodPost {
|
||||
http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
|
||||
return
|
||||
}
|
||||
var summary AgentSummary
|
||||
if !decodeLimitedJSON(w, r, &summary) {
|
||||
return
|
||||
}
|
||||
if summary.Machine != "a" && summary.Machine != "b" {
|
||||
http.Error(w, "invalid machine", http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
s.mu.Lock()
|
||||
s.summaries[summary.Machine] = summary
|
||||
complete := len(s.summaries) == s.expectedAgents
|
||||
s.mu.Unlock()
|
||||
writeJSON(w, http.StatusOK, map[string]any{"complete": complete})
|
||||
if complete {
|
||||
s.completeOnce.Do(func() {
|
||||
go func() {
|
||||
time.Sleep(500 * time.Millisecond)
|
||||
close(s.complete)
|
||||
}()
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func (s *coordinatorState) handleStatus(w http.ResponseWriter, _ *http.Request) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
writeJSON(w, http.StatusOK, map[string]any{
|
||||
"start": s.start, "ready": s.ready, "ledger_agents": len(s.ledgers), "summary_agents": len(s.summaries),
|
||||
})
|
||||
}
|
||||
|
||||
func (s *coordinatorState) writeOutput(path string) error {
|
||||
s.mu.Lock()
|
||||
data := struct {
|
||||
Start StartInfo `json:"start"`
|
||||
Ready map[string]AgentReady `json:"ready"`
|
||||
Ledgers map[string][]AcceptedMessage `json:"ledgers"`
|
||||
Summaries map[string]AgentSummary `json:"summaries"`
|
||||
}{s.start, s.ready, s.ledgers, s.summaries}
|
||||
s.mu.Unlock()
|
||||
raw, err := json.MarshalIndent(data, "", " ")
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if err = os.MkdirAll(filepath.Dir(path), 0o750); err != nil && filepath.Dir(path) != "." {
|
||||
return err
|
||||
}
|
||||
return os.WriteFile(path, append(raw, '\n'), 0o600)
|
||||
}
|
||||
|
||||
func decodeLimitedJSON(w http.ResponseWriter, r *http.Request, output any) bool {
|
||||
dec := json.NewDecoder(io.LimitReader(r.Body, 32<<20))
|
||||
if err := dec.Decode(output); err != nil {
|
||||
http.Error(w, err.Error(), http.StatusBadRequest)
|
||||
return false
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
func writeJSON(w http.ResponseWriter, status int, value any) {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(status)
|
||||
_ = json.NewEncoder(w).Encode(value)
|
||||
}
|
||||
|
||||
type coordinatorClient struct {
|
||||
base string
|
||||
client *http.Client
|
||||
}
|
||||
|
||||
func newCoordinatorClient(base string) *coordinatorClient {
|
||||
return &coordinatorClient{base: strings.TrimRight(base, "/"), client: &http.Client{Timeout: 35 * time.Second}}
|
||||
}
|
||||
|
||||
func (c *coordinatorClient) ready(ctx context.Context, ready AgentReady) (StartInfo, error) {
|
||||
var start StartInfo
|
||||
if _, err := c.request(ctx, http.MethodPost, "/v1/ready", ready, &start); err != nil {
|
||||
return start, err
|
||||
}
|
||||
for start.StartAt.IsZero() {
|
||||
if err := waitContext(ctx, time.Second); err != nil {
|
||||
return start, err
|
||||
}
|
||||
status, err := c.request(ctx, http.MethodGet, "/v1/start", nil, &start)
|
||||
if err != nil {
|
||||
return start, err
|
||||
}
|
||||
if status != http.StatusOK {
|
||||
start = StartInfo{}
|
||||
}
|
||||
}
|
||||
return start, nil
|
||||
}
|
||||
|
||||
func (c *coordinatorClient) submitLedger(ctx context.Context, machine string, messages []AcceptedMessage) error {
|
||||
_, err := c.request(ctx, http.MethodPost, "/v1/ledger?machine="+machine, messages, nil)
|
||||
return err
|
||||
}
|
||||
|
||||
func (c *coordinatorClient) ledger(ctx context.Context) ([]AcceptedMessage, error) {
|
||||
for {
|
||||
var messages []AcceptedMessage
|
||||
status, err := c.request(ctx, http.MethodGet, "/v1/ledger", nil, &messages)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if status == http.StatusOK {
|
||||
return messages, nil
|
||||
}
|
||||
if err = waitContext(ctx, time.Second); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (c *coordinatorClient) submitSummary(ctx context.Context, summary AgentSummary) error {
|
||||
_, err := c.request(ctx, http.MethodPost, "/v1/summary", summary, nil)
|
||||
return err
|
||||
}
|
||||
|
||||
func (c *coordinatorClient) request(ctx context.Context, method, path string, input, output any) (int, error) {
|
||||
var body io.Reader
|
||||
if input != nil {
|
||||
raw, err := json.Marshal(input)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
body = strings.NewReader(string(raw))
|
||||
}
|
||||
req, err := http.NewRequestWithContext(ctx, method, c.base+path, body)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
if input != nil {
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
}
|
||||
resp, err := c.client.Do(req)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
if resp.StatusCode >= 400 {
|
||||
raw, _ := io.ReadAll(io.LimitReader(resp.Body, 4096))
|
||||
return resp.StatusCode, fmt.Errorf("coordinator HTTP %d: %s", resp.StatusCode, strings.TrimSpace(string(raw)))
|
||||
}
|
||||
if output != nil && resp.StatusCode == http.StatusOK {
|
||||
if err = json.NewDecoder(resp.Body).Decode(output); err != nil {
|
||||
return resp.StatusCode, err
|
||||
}
|
||||
}
|
||||
return resp.StatusCode, nil
|
||||
}
|
||||
|
||||
func waitContext(ctx context.Context, delay time.Duration) error {
|
||||
timer := time.NewTimer(delay)
|
||||
defer timer.Stop()
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return ctx.Err()
|
||||
case <-timer.C:
|
||||
return nil
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,60 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestCoordinatorClientWaitsForBothAgentsAndMergesLedger(t *testing.T) {
|
||||
state := &coordinatorState{
|
||||
start: StartInfo{RunID: "run1", Scenario: "p2p"}, startDelay: time.Second, expectedAgents: 2,
|
||||
ready: make(map[string]AgentReady), ledgers: make(map[string][]AcceptedMessage),
|
||||
summaries: make(map[string]AgentSummary), complete: make(chan struct{}),
|
||||
}
|
||||
server := httptest.NewServer(testCoordinatorMux(state))
|
||||
defer server.Close()
|
||||
client := newCoordinatorClient(server.URL)
|
||||
|
||||
aDone := make(chan StartInfo, 1)
|
||||
go func() {
|
||||
start, _ := client.ready(context.Background(), AgentReady{Machine: "a", ReadyConnections: 1000})
|
||||
aDone <- start
|
||||
}()
|
||||
time.Sleep(50 * time.Millisecond)
|
||||
startB, err := client.ready(context.Background(), AgentReady{Machine: "b", ReadyConnections: 1000})
|
||||
if err != nil || startB.StartAt.IsZero() {
|
||||
t.Fatalf("machine B ready: start=%+v err=%v", startB, err)
|
||||
}
|
||||
select {
|
||||
case startA := <-aDone:
|
||||
if !startA.StartAt.Equal(startB.StartAt) {
|
||||
t.Fatalf("different start times A=%s B=%s", startA.StartAt, startB.StartAt)
|
||||
}
|
||||
case <-time.After(2 * time.Second):
|
||||
t.Fatal("machine A did not receive start barrier")
|
||||
}
|
||||
|
||||
if err = client.submitLedger(context.Background(), "a", []AcceptedMessage{{ClientMid: "b"}}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err = client.submitLedger(context.Background(), "b", []AcceptedMessage{{ClientMid: "a"}}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
ledger, err := client.ledger(context.Background())
|
||||
if err != nil || len(ledger) != 2 || ledger[0].ClientMid != "a" || ledger[1].ClientMid != "b" {
|
||||
t.Fatalf("ledger=%+v err=%v", ledger, err)
|
||||
}
|
||||
}
|
||||
|
||||
func testCoordinatorMux(state *coordinatorState) *http.ServeMux {
|
||||
mux := http.NewServeMux()
|
||||
mux.HandleFunc("/v1/ready", state.handleReady)
|
||||
mux.HandleFunc("/v1/start", state.handleStart)
|
||||
mux.HandleFunc("/v1/ledger", state.handleLedger)
|
||||
mux.HandleFunc("/v1/summary", state.handleSummary)
|
||||
mux.HandleFunc("/v1/status", state.handleStatus)
|
||||
return mux
|
||||
}
|
||||
@@ -0,0 +1,64 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"sort"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
var latencyBounds = []time.Duration{
|
||||
time.Millisecond, 3 * time.Millisecond, 5 * time.Millisecond, 10 * time.Millisecond,
|
||||
25 * time.Millisecond, 50 * time.Millisecond, 100 * time.Millisecond, 200 * time.Millisecond,
|
||||
500 * time.Millisecond, time.Second, 2 * time.Second, 5 * time.Second, 10 * time.Second,
|
||||
}
|
||||
|
||||
type DurationHistogram struct {
|
||||
mu sync.Mutex
|
||||
count int64
|
||||
sum time.Duration
|
||||
max time.Duration
|
||||
bucket []int64
|
||||
}
|
||||
|
||||
type HistogramSnapshot struct {
|
||||
Count int64 `json:"count"`
|
||||
SumMS float64 `json:"sum_ms"`
|
||||
MaxMS float64 `json:"max_ms"`
|
||||
Buckets map[string]int64 `json:"buckets"`
|
||||
}
|
||||
|
||||
func NewDurationHistogram() *DurationHistogram {
|
||||
return &DurationHistogram{bucket: make([]int64, len(latencyBounds)+1)}
|
||||
}
|
||||
|
||||
func (h *DurationHistogram) Observe(value time.Duration) {
|
||||
if value < 0 {
|
||||
return
|
||||
}
|
||||
h.mu.Lock()
|
||||
defer h.mu.Unlock()
|
||||
h.count++
|
||||
h.sum += value
|
||||
if value > h.max {
|
||||
h.max = value
|
||||
}
|
||||
idx := sort.Search(len(latencyBounds), func(i int) bool { return value <= latencyBounds[i] })
|
||||
h.bucket[idx]++
|
||||
}
|
||||
|
||||
func (h *DurationHistogram) Snapshot() HistogramSnapshot {
|
||||
h.mu.Lock()
|
||||
defer h.mu.Unlock()
|
||||
buckets := make(map[string]int64, len(h.bucket))
|
||||
for i, count := range h.bucket {
|
||||
label := "+Inf"
|
||||
if i < len(latencyBounds) {
|
||||
label = latencyBounds[i].String()
|
||||
}
|
||||
buckets[label] = count
|
||||
}
|
||||
return HistogramSnapshot{
|
||||
Count: h.count, SumMS: float64(h.sum) / float64(time.Millisecond),
|
||||
MaxMS: float64(h.max) / float64(time.Millisecond), Buckets: buckets,
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,259 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"crypto/tls"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
)
|
||||
|
||||
type apiClient struct {
|
||||
baseURL string
|
||||
apiKey string
|
||||
client *http.Client
|
||||
seq atomic.Uint64
|
||||
}
|
||||
|
||||
type apiEnvelope struct {
|
||||
Code int `json:"code"`
|
||||
Data json.RawMessage `json:"data"`
|
||||
Reason string `json:"reason"`
|
||||
Message string `json:"message"`
|
||||
RequestID string `json:"request_id"`
|
||||
}
|
||||
|
||||
type apiError struct {
|
||||
Status int
|
||||
Code int
|
||||
Reason string
|
||||
Message string
|
||||
}
|
||||
|
||||
func (e *apiError) Error() string {
|
||||
return fmt.Sprintf("HTTP %d code=%d reason=%s message=%s", e.Status, e.Code, e.Reason, e.Message)
|
||||
}
|
||||
|
||||
func newAPIClient(cfg *Config, apiKey string) *apiClient {
|
||||
transport := http.DefaultTransport.(*http.Transport).Clone()
|
||||
if cfg.InsecureTLS {
|
||||
transport.TLSClientConfig = &tls.Config{InsecureSkipVerify: true} // #nosec G402: explicit load-test option.
|
||||
}
|
||||
return &apiClient{
|
||||
baseURL: strings.TrimRight(cfg.HTTPBase, "/"), apiKey: apiKey,
|
||||
client: &http.Client{Timeout: 30 * time.Second, Transport: transport},
|
||||
}
|
||||
}
|
||||
|
||||
func (a *apiClient) do(ctx context.Context, method, path, token, idempotencyKey string, body any, output any) error {
|
||||
var reader io.Reader
|
||||
if body != nil {
|
||||
raw, err := json.Marshal(body)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
reader = bytes.NewReader(raw)
|
||||
}
|
||||
req, err := http.NewRequestWithContext(ctx, method, a.baseURL+path, reader)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
req.Header.Set("X-MeetingIM-APIKey", a.apiKey)
|
||||
req.Header.Set("X-IM-Schema", "1")
|
||||
req.Header.Set("X-Request-ID", fmt.Sprintf("loadtest-%d-%d", time.Now().UnixNano(), a.seq.Add(1)))
|
||||
if token != "" {
|
||||
req.Header.Set("Authorization", "Bearer "+token)
|
||||
}
|
||||
if idempotencyKey != "" {
|
||||
req.Header.Set("Idempotency-Key", idempotencyKey)
|
||||
}
|
||||
if body != nil {
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
}
|
||||
resp, err := a.client.Do(req)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
raw, err := io.ReadAll(io.LimitReader(resp.Body, 8<<20))
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
var envelope apiEnvelope
|
||||
if err = json.Unmarshal(raw, &envelope); err != nil {
|
||||
return fmt.Errorf("decode HTTP %d response: %w", resp.StatusCode, err)
|
||||
}
|
||||
if resp.StatusCode < 200 || resp.StatusCode >= 300 || envelope.Code != 0 {
|
||||
return &apiError{Status: resp.StatusCode, Code: envelope.Code, Reason: envelope.Reason, Message: envelope.Message}
|
||||
}
|
||||
if output != nil && len(envelope.Data) > 0 {
|
||||
if err = json.Unmarshal(envelope.Data, output); err != nil {
|
||||
return fmt.Errorf("decode response data: %w", err)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
type privacyResponse struct {
|
||||
FriendRequestVerification bool `json:"friend_request_verification"`
|
||||
GroupInviteVerification bool `json:"group_invite_verification"`
|
||||
Version int `json:"version"`
|
||||
}
|
||||
|
||||
func (a *apiClient) disableGroupInviteVerification(ctx context.Context, account AccountRecord) error {
|
||||
var privacy privacyResponse
|
||||
if err := a.do(ctx, http.MethodGet, "/v1/im/me/privacy", account.Token, "", nil, &privacy); err != nil {
|
||||
return err
|
||||
}
|
||||
if !privacy.GroupInviteVerification {
|
||||
return nil
|
||||
}
|
||||
body := map[string]any{"group_invite_verification": false, "version": privacy.Version}
|
||||
return a.do(ctx, http.MethodPatch, "/v1/im/me/privacy", account.Token,
|
||||
"loadtest-privacy-"+account.Username, body, &privacy)
|
||||
}
|
||||
|
||||
type friendRequestResponse struct {
|
||||
RequestID string `json:"request_id"`
|
||||
Status string `json:"status"`
|
||||
Topic string `json:"topic"`
|
||||
}
|
||||
|
||||
type friendRequestListResponse struct {
|
||||
Items []struct {
|
||||
RequestID string `json:"request_id"`
|
||||
ApplicantID string `json:"applicant_id"`
|
||||
Status string `json:"status"`
|
||||
} `json:"items"`
|
||||
NextCursor string `json:"next_cursor"`
|
||||
}
|
||||
|
||||
func (a *apiClient) findPendingFriendRequest(ctx context.Context, target AccountRecord, applicantID string) (string, error) {
|
||||
cursor := ""
|
||||
for {
|
||||
path := "/v1/im/friend-requests?direction=inbound&status=pending&limit=100"
|
||||
if cursor != "" {
|
||||
path += "&cursor=" + url.QueryEscape(cursor)
|
||||
}
|
||||
var page friendRequestListResponse
|
||||
if err := a.do(ctx, http.MethodGet, path, target.Token, "", nil, &page); err != nil {
|
||||
return "", err
|
||||
}
|
||||
for _, item := range page.Items {
|
||||
if item.ApplicantID == applicantID && item.Status == "pending" {
|
||||
return item.RequestID, nil
|
||||
}
|
||||
}
|
||||
if page.NextCursor == "" || page.NextCursor == cursor {
|
||||
return "", nil
|
||||
}
|
||||
cursor = page.NextCursor
|
||||
}
|
||||
}
|
||||
|
||||
func (a *apiClient) createFriendRequest(ctx context.Context, source, target AccountRecord) (*friendRequestResponse, error) {
|
||||
var result friendRequestResponse
|
||||
body := map[string]any{"target_user_id": target.UserID, "message": "loadtest", "source": "search", "source_topic": ""}
|
||||
err := a.do(ctx, http.MethodPost, "/v1/im/friend-requests", source.Token,
|
||||
fmt.Sprintf("loadtest-friend-%04d", source.Index), body, &result)
|
||||
return &result, err
|
||||
}
|
||||
|
||||
func (a *apiClient) acceptFriendRequest(ctx context.Context, target AccountRecord, requestID string) (*friendRequestResponse, error) {
|
||||
var result friendRequestResponse
|
||||
err := a.do(ctx, http.MethodPost, "/v1/im/friend-requests/"+url.PathEscape(requestID)+"/accept", target.Token,
|
||||
"loadtest-friend-accept-"+requestID, map[string]any{}, &result)
|
||||
return &result, err
|
||||
}
|
||||
|
||||
type p2pResponse struct {
|
||||
Topic string `json:"topic"`
|
||||
Relationship string `json:"relationship"`
|
||||
}
|
||||
|
||||
func (a *apiClient) ensureP2P(ctx context.Context, source, target AccountRecord) (*p2pResponse, error) {
|
||||
var result p2pResponse
|
||||
err := a.do(ctx, http.MethodPost, "/v1/im/conversations/p2p", source.Token,
|
||||
fmt.Sprintf("loadtest-p2p-%s-%04d", source.Machine, source.Index),
|
||||
map[string]any{"peer_user_id": target.UserID}, &result)
|
||||
return &result, err
|
||||
}
|
||||
|
||||
type groupCreateResponse struct {
|
||||
Topic string `json:"topic"`
|
||||
MemberCount int `json:"member_count"`
|
||||
}
|
||||
|
||||
func (a *apiClient) createGroup(ctx context.Context, owner AccountRecord) (*groupCreateResponse, error) {
|
||||
var result groupCreateResponse
|
||||
body := map[string]any{"name": "IM 006 Load Test", "description": "im-pressure", "member_user_ids": []string{}}
|
||||
err := a.do(ctx, http.MethodPost, "/v1/im/groups", owner.Token, "loadtest-group-create", body, &result)
|
||||
return &result, err
|
||||
}
|
||||
|
||||
type groupInvitationItem struct {
|
||||
UserID string `json:"user_id"`
|
||||
Status string `json:"status"`
|
||||
Reason string `json:"reason"`
|
||||
}
|
||||
|
||||
type groupInvitationResponse struct {
|
||||
Topic string `json:"topic"`
|
||||
Items []groupInvitationItem `json:"items"`
|
||||
}
|
||||
|
||||
func (a *apiClient) inviteGroup(ctx context.Context, owner AccountRecord, topic string, batch int, userIDs []string) (*groupInvitationResponse, error) {
|
||||
var result groupInvitationResponse
|
||||
err := a.do(ctx, http.MethodPost, "/v1/im/groups/"+url.PathEscape(topic)+"/invitations", owner.Token,
|
||||
fmt.Sprintf("loadtest-group-%s-invite-%03d", topicHash(topic), batch), map[string]any{"user_ids": userIDs, "message": "loadtest"}, &result)
|
||||
return &result, err
|
||||
}
|
||||
|
||||
type groupMembersResponse struct {
|
||||
Items []struct {
|
||||
UserID string `json:"user_id"`
|
||||
} `json:"items"`
|
||||
NextCursor string `json:"next_cursor"`
|
||||
HasMore bool `json:"has_more"`
|
||||
}
|
||||
|
||||
func (a *apiClient) countGroupMembers(ctx context.Context, account AccountRecord, topic string) (int, error) {
|
||||
cursor := ""
|
||||
seen := make(map[string]bool)
|
||||
for {
|
||||
path := "/v1/im/groups/" + url.PathEscape(topic) + "/members?limit=100"
|
||||
if cursor != "" {
|
||||
path += "&cursor=" + url.QueryEscape(cursor)
|
||||
}
|
||||
var page groupMembersResponse
|
||||
if err := a.do(ctx, http.MethodGet, path, account.Token, "", nil, &page); err != nil {
|
||||
return 0, err
|
||||
}
|
||||
for _, item := range page.Items {
|
||||
seen[item.UserID] = true
|
||||
}
|
||||
if page.NextCursor == "" || (!page.HasMore && len(page.Items) < 100) {
|
||||
break
|
||||
}
|
||||
if page.NextCursor == cursor {
|
||||
return 0, fmt.Errorf("group member cursor did not advance: %s", cursor)
|
||||
}
|
||||
cursor = page.NextCursor
|
||||
}
|
||||
return len(seen), nil
|
||||
}
|
||||
|
||||
func queryInt(value string, fallback int) int {
|
||||
n, err := strconv.Atoi(value)
|
||||
if err != nil {
|
||||
return fallback
|
||||
}
|
||||
return n
|
||||
}
|
||||
@@ -0,0 +1,146 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"context"
|
||||
"flag"
|
||||
"fmt"
|
||||
"log"
|
||||
"os"
|
||||
"os/signal"
|
||||
"syscall"
|
||||
"time"
|
||||
)
|
||||
|
||||
func main() {
|
||||
log.SetFlags(log.LstdFlags | log.Lmicroseconds | log.LUTC)
|
||||
ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
|
||||
defer stop()
|
||||
if len(os.Args) < 2 {
|
||||
usage()
|
||||
os.Exit(2)
|
||||
}
|
||||
var err error
|
||||
switch os.Args[1] {
|
||||
case "prepare":
|
||||
err = prepareCommand(ctx, os.Args[2:])
|
||||
case "verify":
|
||||
err = verifyCommand(ctx, os.Args[2:])
|
||||
case "coordinator":
|
||||
err = coordinatorCommand(ctx, os.Args[2:])
|
||||
case "agent":
|
||||
err = agentCommand(ctx, os.Args[2:])
|
||||
case "merge":
|
||||
err = mergeCommand(os.Args[2:])
|
||||
case "collect":
|
||||
err = collectCommand(ctx, os.Args[2:])
|
||||
default:
|
||||
usage()
|
||||
os.Exit(2)
|
||||
}
|
||||
if err != nil {
|
||||
log.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
func prepareCommand(ctx context.Context, args []string) error {
|
||||
fs := flag.NewFlagSet("prepare", flag.ContinueOnError)
|
||||
configPath := fs.String("config", "loadtest/im-pressure/config.example.json", "JSON configuration path")
|
||||
manifestPath := fs.String("manifest", "loadtest/im-pressure/manifest.json", "manifest output path")
|
||||
if err := fs.Parse(args); err != nil {
|
||||
return err
|
||||
}
|
||||
cfg, err := loadConfig(*configPath)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return runPrepare(ctx, cfg, *manifestPath)
|
||||
}
|
||||
|
||||
func verifyCommand(ctx context.Context, args []string) error {
|
||||
fs := flag.NewFlagSet("verify", flag.ContinueOnError)
|
||||
configPath := fs.String("config", "loadtest/im-pressure/config.example.json", "JSON configuration path")
|
||||
manifestPath := fs.String("manifest", "loadtest/im-pressure/manifest.json", "manifest path")
|
||||
if err := fs.Parse(args); err != nil {
|
||||
return err
|
||||
}
|
||||
cfg, err := loadConfig(*configPath)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
manifest, err := loadManifest(*manifestPath)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
apiKey, err := cfg.apiKey()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return runVerify(ctx, cfg, manifest, newAPIClient(cfg, apiKey))
|
||||
}
|
||||
|
||||
func coordinatorCommand(ctx context.Context, args []string) error {
|
||||
fs := flag.NewFlagSet("coordinator", flag.ContinueOnError)
|
||||
listen := fs.String("listen", ":8090", "coordinator listen address")
|
||||
runID := fs.String("run-id", "", "unique run ID, max 24 characters")
|
||||
scenario := fs.String("scenario", "", "p2p, group, or mixed")
|
||||
output := fs.String("output", "loadtest/im-pressure/results/coordinator.json", "coordinator result path")
|
||||
expected := fs.Int("expected-agents", 2, "number of agents")
|
||||
delay := fs.Duration("start-delay", 30*time.Second, "delay after all agents are ready")
|
||||
if err := fs.Parse(args); err != nil {
|
||||
return err
|
||||
}
|
||||
return runCoordinator(ctx, *listen, *runID, *scenario, *output, *expected, *delay)
|
||||
}
|
||||
|
||||
func agentCommand(ctx context.Context, args []string) error {
|
||||
fs := flag.NewFlagSet("agent", flag.ContinueOnError)
|
||||
configPath := fs.String("config", "loadtest/im-pressure/config.example.json", "JSON configuration path")
|
||||
manifestPath := fs.String("manifest", "loadtest/im-pressure/manifest.json", "manifest path")
|
||||
machine := fs.String("machine", "", "a or b")
|
||||
coordinatorURL := fs.String("coordinator", "", "coordinator base URL")
|
||||
output := fs.String("output", "loadtest/im-pressure/results/agent-summary.json", "agent result path")
|
||||
metricsListen := fs.String("metrics-listen", ":9105", "Prometheus listen address; empty disables it")
|
||||
if err := fs.Parse(args); err != nil {
|
||||
return err
|
||||
}
|
||||
if *coordinatorURL == "" {
|
||||
return fmt.Errorf("--coordinator is required")
|
||||
}
|
||||
cfg, err := loadConfig(*configPath)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
manifest, err := loadManifest(*manifestPath)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return runAgent(ctx, cfg, manifest, *machine, *coordinatorURL, *output, *metricsListen)
|
||||
}
|
||||
|
||||
func mergeCommand(args []string) error {
|
||||
fs := flag.NewFlagSet("merge", flag.ContinueOnError)
|
||||
a := fs.String("a", "loadtest/im-pressure/results/agent-a-summary.json", "machine A summary")
|
||||
b := fs.String("b", "loadtest/im-pressure/results/agent-b-summary.json", "machine B summary")
|
||||
output := fs.String("output", "loadtest/im-pressure/results/report.json", "combined report")
|
||||
if err := fs.Parse(args); err != nil {
|
||||
return err
|
||||
}
|
||||
return runMerge(*a, *b, *output)
|
||||
}
|
||||
|
||||
func collectCommand(ctx context.Context, args []string) error {
|
||||
fs := flag.NewFlagSet("collect", flag.ContinueOnError)
|
||||
var urls namedURLs
|
||||
fs.Var(&urls, "url", "repeatable name=https://host/metrics source")
|
||||
output := fs.String("output", "loadtest/im-pressure/results/metrics.jsonl", "snapshot output")
|
||||
interval := fs.Duration("interval", 5*time.Second, "collection interval")
|
||||
duration := fs.Duration("duration", 10*time.Minute, "collection duration")
|
||||
if err := fs.Parse(args); err != nil {
|
||||
return err
|
||||
}
|
||||
return runCollector(ctx, urls, *output, *interval, *duration)
|
||||
}
|
||||
|
||||
func usage() {
|
||||
fmt.Fprintln(os.Stderr, "usage: im-pressure <prepare|verify|coordinator|agent|merge|collect> [flags]")
|
||||
}
|
||||
@@ -0,0 +1,114 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"sort"
|
||||
"time"
|
||||
)
|
||||
|
||||
type CombinedReport struct {
|
||||
RunID string `json:"run_id"`
|
||||
Scenario string `json:"scenario"`
|
||||
GeneratedAt time.Time `json:"generated_at"`
|
||||
Agents []AgentSummary `json:"agents"`
|
||||
ReadyConnections int `json:"ready_connections"`
|
||||
PublishAttempted int64 `json:"publish_attempted"`
|
||||
PublishAccepted int64 `json:"publish_accepted"`
|
||||
PublishRejected int64 `json:"publish_rejected"`
|
||||
ExpectedDeliveries int64 `json:"expected_deliveries"`
|
||||
ReceivedDeliveries int64 `json:"received_deliveries"`
|
||||
RealtimeLost int64 `json:"realtime_lost"`
|
||||
RealtimeLostStored int64 `json:"realtime_lost_but_persisted"`
|
||||
ArrivalRate float64 `json:"arrival_rate"`
|
||||
LossRate float64 `json:"loss_rate"`
|
||||
Duplicates int64 `json:"duplicates"`
|
||||
OutOfOrder int64 `json:"out_of_order"`
|
||||
Disconnects int64 `json:"disconnects"`
|
||||
HistoryChecked int64 `json:"history_checked"`
|
||||
HistoryFound int64 `json:"history_found"`
|
||||
PersistenceLost int64 `json:"persistence_lost"`
|
||||
Missing []MissingDelivery `json:"missing,omitempty"`
|
||||
Passed bool `json:"passed"`
|
||||
FailureReasons []string `json:"failure_reasons,omitempty"`
|
||||
}
|
||||
|
||||
func runMerge(aPath, bPath, output string) error {
|
||||
a, err := readAgentSummary(aPath)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
b, err := readAgentSummary(bPath)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if a.RunID != b.RunID || a.Scenario != b.Scenario {
|
||||
return fmt.Errorf("agent summaries do not belong to the same run")
|
||||
}
|
||||
report := CombinedReport{RunID: a.RunID, Scenario: a.Scenario, GeneratedAt: time.Now().UTC(), Agents: []AgentSummary{a, b}, Passed: true}
|
||||
for _, summary := range report.Agents {
|
||||
report.ReadyConnections += summary.ReadyConnections
|
||||
report.PublishAttempted += summary.PublishAttempted
|
||||
report.PublishAccepted += summary.PublishAccepted
|
||||
report.PublishRejected += summary.PublishRejected
|
||||
report.ExpectedDeliveries += summary.ExpectedPeer + summary.ExpectedGroup
|
||||
report.ReceivedDeliveries += summary.ReceivedPeer + summary.ReceivedGroup
|
||||
report.RealtimeLost += summary.RealtimeLost
|
||||
report.RealtimeLostStored += summary.RealtimeLostStored
|
||||
report.Duplicates += summary.Duplicates
|
||||
report.OutOfOrder += summary.OutOfOrder
|
||||
report.Disconnects += summary.Disconnects
|
||||
report.HistoryChecked += summary.HistoryChecked
|
||||
report.HistoryFound += summary.HistoryFound
|
||||
report.PersistenceLost += summary.PersistenceLost
|
||||
report.Missing = append(report.Missing, summary.Missing...)
|
||||
if len(summary.Errors) > 0 {
|
||||
report.FailureReasons = append(report.FailureReasons, fmt.Sprintf("agent %s errors=%d", summary.Machine, len(summary.Errors)))
|
||||
}
|
||||
}
|
||||
if report.ExpectedDeliveries > 0 {
|
||||
report.ArrivalRate = float64(report.ReceivedDeliveries) / float64(report.ExpectedDeliveries)
|
||||
report.LossRate = float64(report.RealtimeLost) / float64(report.ExpectedDeliveries)
|
||||
}
|
||||
if report.ReadyConnections != 2000 {
|
||||
report.FailureReasons = append(report.FailureReasons, fmt.Sprintf("ready connections=%d expected=2000", report.ReadyConnections))
|
||||
}
|
||||
if report.PublishRejected > 0 || report.RealtimeLost > 0 || report.PersistenceLost > 0 || report.Duplicates > 0 || report.OutOfOrder > 0 || report.Disconnects > 0 {
|
||||
report.FailureReasons = append(report.FailureReasons, "message correctness gate failed")
|
||||
}
|
||||
report.Passed = len(report.FailureReasons) == 0
|
||||
sort.Slice(report.Missing, func(i, j int) bool {
|
||||
if report.Missing[i].ClientMid == report.Missing[j].ClientMid {
|
||||
return report.Missing[i].ReceiverIndex < report.Missing[j].ReceiverIndex
|
||||
}
|
||||
return report.Missing[i].ClientMid < report.Missing[j].ClientMid
|
||||
})
|
||||
return writeJSONFile(output, report)
|
||||
}
|
||||
|
||||
func readAgentSummary(path string) (AgentSummary, error) {
|
||||
var summary AgentSummary
|
||||
raw, err := os.ReadFile(path)
|
||||
if err != nil {
|
||||
return summary, err
|
||||
}
|
||||
if err = json.Unmarshal(raw, &summary); err != nil {
|
||||
return summary, err
|
||||
}
|
||||
return summary, nil
|
||||
}
|
||||
|
||||
func writeJSONFile(path string, value any) error {
|
||||
raw, err := json.MarshalIndent(value, "", " ")
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if dir := filepath.Dir(path); dir != "." {
|
||||
if err = os.MkdirAll(dir, 0o750); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return os.WriteFile(path, append(raw, '\n'), 0o600)
|
||||
}
|
||||
@@ -0,0 +1,49 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestRunMergeComputesArrivalAndPersistenceGate(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
a := AgentSummary{RunID: "run1", Scenario: "p2p", Machine: "a", ReadyConnections: 1000,
|
||||
PublishAttempted: 10, PublishAccepted: 10, ExpectedPeer: 10, ReceivedPeer: 9,
|
||||
RealtimeLost: 1, HistoryChecked: 10, HistoryFound: 10}
|
||||
b := AgentSummary{RunID: "run1", Scenario: "p2p", Machine: "b", ReadyConnections: 1000,
|
||||
PublishAttempted: 10, PublishAccepted: 10, ExpectedPeer: 10, ReceivedPeer: 10,
|
||||
HistoryChecked: 10, HistoryFound: 9, PersistenceLost: 1}
|
||||
aPath, bPath, reportPath := filepath.Join(dir, "a.json"), filepath.Join(dir, "b.json"), filepath.Join(dir, "report.json")
|
||||
writeSummaryForTest(t, aPath, a)
|
||||
writeSummaryForTest(t, bPath, b)
|
||||
if err := runMerge(aPath, bPath, reportPath); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var report CombinedReport
|
||||
raw, err := os.ReadFile(reportPath)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err = json.Unmarshal(raw, &report); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if report.ExpectedDeliveries != 20 || report.ReceivedDeliveries != 19 || report.ArrivalRate != 0.95 || report.LossRate != 0.05 {
|
||||
t.Fatalf("unexpected delivery totals: %+v", report)
|
||||
}
|
||||
if report.Passed || report.PersistenceLost != 1 {
|
||||
t.Fatalf("correctness gate did not fail: %+v", report)
|
||||
}
|
||||
}
|
||||
|
||||
func writeSummaryForTest(t *testing.T, path string, summary AgentSummary) {
|
||||
t.Helper()
|
||||
raw, err := json.Marshal(summary)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err = os.WriteFile(path, raw, 0o600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
@@ -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"`
|
||||
}
|
||||
@@ -0,0 +1,45 @@
|
||||
package main
|
||||
|
||||
import "testing"
|
||||
|
||||
func TestNewManifestDistributesUsersAndEndpoints(t *testing.T) {
|
||||
cfg := &Config{Tenant: "006", UsersPerMachine: 3, Endpoints: map[string][]string{
|
||||
"a": {"ws://a1", "ws://a2"}, "b": {"ws://b1"},
|
||||
}}
|
||||
manifest := newManifest(cfg)
|
||||
if len(manifest.Accounts) != 6 {
|
||||
t.Fatalf("accounts=%d", len(manifest.Accounts))
|
||||
}
|
||||
if got := accountsForMachine(manifest, "a"); got[0].Username != "lta0001" || got[1].Endpoint != "ws://a2" || got[2].Endpoint != "ws://a1" {
|
||||
t.Fatalf("unexpected A accounts: %+v", got)
|
||||
}
|
||||
if err := manifest.validate(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestManifestRejectsDuplicateAccount(t *testing.T) {
|
||||
manifest := &Manifest{Version: manifestVersion, Accounts: []AccountRecord{
|
||||
{Machine: "a", Index: 1, Username: "lta0001"},
|
||||
{Machine: "a", Index: 1, Username: "duplicate"},
|
||||
}}
|
||||
if err := manifest.validate(); err == nil {
|
||||
t.Fatal("expected duplicate account validation error")
|
||||
}
|
||||
}
|
||||
|
||||
func TestApplyConfigEndpointsRefreshesManifest(t *testing.T) {
|
||||
cfg := &Config{Tenant: "006", UsersPerMachine: 1, Endpoints: map[string][]string{
|
||||
"a": {"ws://new-a"}, "b": {"ws://new-b"},
|
||||
}}
|
||||
manifest := &Manifest{Tenant: "006", Accounts: []AccountRecord{
|
||||
{Machine: "a", Index: 1, Username: "lta0001", Endpoint: "ws://old"},
|
||||
{Machine: "b", Index: 1, Username: "ltb0001", Endpoint: "ws://old"},
|
||||
}}
|
||||
if err := applyConfigEndpoints(cfg, manifest); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if manifest.Accounts[0].Endpoint != "ws://new-a" || manifest.Accounts[1].Endpoint != "ws://new-b" {
|
||||
t.Fatalf("endpoints were not refreshed: %+v", manifest.Accounts)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,253 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"log"
|
||||
"net/http"
|
||||
"sort"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
)
|
||||
|
||||
func runPrepare(ctx context.Context, cfg *Config, manifestPath string) error {
|
||||
apiKey, err := cfg.apiKey()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
password, err := cfg.password()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
manifest := newManifest(cfg)
|
||||
if existing, loadErr := loadManifest(manifestPath); loadErr == nil {
|
||||
manifest = existing
|
||||
log.Printf("resuming manifest with %d accounts", len(manifest.Accounts))
|
||||
}
|
||||
if err = applyConfigEndpoints(cfg, manifest); err != nil {
|
||||
return err
|
||||
}
|
||||
if err = prepareAccounts(ctx, cfg, manifest, apiKey, password); err != nil {
|
||||
return err
|
||||
}
|
||||
if err = saveManifest(manifestPath, manifest); err != nil {
|
||||
return err
|
||||
}
|
||||
api := newAPIClient(cfg, apiKey)
|
||||
if err = prepareFriendships(ctx, cfg, manifest, api); err != nil {
|
||||
return err
|
||||
}
|
||||
if err = saveManifest(manifestPath, manifest); err != nil {
|
||||
return err
|
||||
}
|
||||
if err = prepareGroup(ctx, cfg, manifest, api); err != nil {
|
||||
return err
|
||||
}
|
||||
if err = saveManifest(manifestPath, manifest); err != nil {
|
||||
return err
|
||||
}
|
||||
return runVerify(ctx, cfg, manifest, api)
|
||||
}
|
||||
|
||||
func prepareAccounts(ctx context.Context, cfg *Config, manifest *Manifest, apiKey, password string) error {
|
||||
var created atomic.Int64
|
||||
var existing atomic.Int64
|
||||
var mu sync.Mutex
|
||||
err := parallelAccounts(ctx, manifest.Accounts, cfg.PrepareWorkers, func(account AccountRecord) error {
|
||||
client, wasCreated, err := connectForPrepare(ctx, cfg, account, apiKey, password)
|
||||
if err != nil {
|
||||
return fmt.Errorf("account %s: %w", account.Username, err)
|
||||
}
|
||||
updated := client.Account()
|
||||
client.Close()
|
||||
mu.Lock()
|
||||
defer mu.Unlock()
|
||||
for i := range manifest.Accounts {
|
||||
if manifest.Accounts[i].Machine == account.Machine && manifest.Accounts[i].Index == account.Index {
|
||||
manifest.Accounts[i].UserID = updated.UserID
|
||||
manifest.Accounts[i].Token = updated.Token
|
||||
manifest.Accounts[i].TokenExpires = updated.TokenExpires
|
||||
break
|
||||
}
|
||||
}
|
||||
if wasCreated {
|
||||
created.Add(1)
|
||||
} else {
|
||||
existing.Add(1)
|
||||
}
|
||||
return nil
|
||||
})
|
||||
if err == nil {
|
||||
log.Printf("accounts ready=%d created=%d existing=%d", len(manifest.Accounts), created.Load(), existing.Load())
|
||||
}
|
||||
return err
|
||||
}
|
||||
|
||||
func prepareFriendships(ctx context.Context, cfg *Config, manifest *Manifest, api *apiClient) error {
|
||||
accounts := manifestAccountMap(manifest)
|
||||
pairs := make([]AccountRecord, 0, cfg.UsersPerMachine)
|
||||
for i := 1; i <= cfg.UsersPerMachine; i++ {
|
||||
pairs = append(pairs, *accounts[accountKey("a", i)])
|
||||
}
|
||||
var mu sync.Mutex
|
||||
err := parallelAccounts(ctx, pairs, cfg.PrepareWorkers, func(a AccountRecord) error {
|
||||
b := *accounts[accountKey("b", a.Index)]
|
||||
if a.UserID == "" || b.UserID == "" {
|
||||
return fmt.Errorf("pair %d has missing user ID", a.Index)
|
||||
}
|
||||
request, requestErr := api.createFriendRequest(ctx, a, b)
|
||||
if requestErr == nil && request.Status == "pending" {
|
||||
if _, requestErr = api.acceptFriendRequest(ctx, b, request.RequestID); requestErr != nil {
|
||||
return fmt.Errorf("accept pair %d: %w", a.Index, requestErr)
|
||||
}
|
||||
} else if requestErr != nil {
|
||||
var apiErr *apiError
|
||||
if !errors.As(requestErr, &apiErr) || apiErr.Status != http.StatusConflict {
|
||||
return fmt.Errorf("friend pair %d: %w", a.Index, requestErr)
|
||||
}
|
||||
pendingID, findErr := api.findPendingFriendRequest(ctx, b, a.UserID)
|
||||
if findErr != nil {
|
||||
return fmt.Errorf("find pending friend pair %d: %w", a.Index, findErr)
|
||||
}
|
||||
if pendingID != "" {
|
||||
if _, acceptErr := api.acceptFriendRequest(ctx, b, pendingID); acceptErr != nil {
|
||||
return fmt.Errorf("resume friend pair %d: %w", a.Index, acceptErr)
|
||||
}
|
||||
}
|
||||
}
|
||||
conversation, err := api.ensureP2P(ctx, a, b)
|
||||
if err != nil {
|
||||
return fmt.Errorf("ensure P2P pair %d: %w", a.Index, err)
|
||||
}
|
||||
if conversation.Relationship != "friend" || conversation.Topic == "" {
|
||||
return fmt.Errorf("pair %d relationship=%s topic=%q", a.Index, conversation.Relationship, conversation.Topic)
|
||||
}
|
||||
mu.Lock()
|
||||
accounts[accountKey("a", a.Index)].P2PTopic = conversation.Topic
|
||||
accounts[accountKey("b", a.Index)].P2PTopic = conversation.Topic
|
||||
mu.Unlock()
|
||||
return nil
|
||||
})
|
||||
if err == nil {
|
||||
log.Printf("friendships ready=%d", cfg.UsersPerMachine)
|
||||
}
|
||||
return err
|
||||
}
|
||||
|
||||
func prepareGroup(ctx context.Context, cfg *Config, manifest *Manifest, api *apiClient) error {
|
||||
if err := parallelAccounts(ctx, manifest.Accounts, cfg.PrepareWorkers, func(account AccountRecord) error {
|
||||
return api.disableGroupInviteVerification(ctx, account)
|
||||
}); err != nil {
|
||||
return fmt.Errorf("prepare group invite privacy: %w", err)
|
||||
}
|
||||
accounts := manifestAccountMap(manifest)
|
||||
owner := *accounts[accountKey("a", 1)]
|
||||
if manifest.GroupTopic == "" {
|
||||
group, err := api.createGroup(ctx, owner)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
manifest.GroupTopic = group.Topic
|
||||
log.Printf("created load-test group topic_hash=%s", topicHash(group.Topic))
|
||||
}
|
||||
memberIDs := make([]string, 0, len(manifest.Accounts)-1)
|
||||
for _, account := range manifest.Accounts {
|
||||
if account.UserID != owner.UserID {
|
||||
memberIDs = append(memberIDs, account.UserID)
|
||||
}
|
||||
}
|
||||
sort.Strings(memberIDs)
|
||||
for offset, batch := 0, 1; offset < len(memberIDs); offset, batch = offset+100, batch+1 {
|
||||
end := offset + 100
|
||||
if end > len(memberIDs) {
|
||||
end = len(memberIDs)
|
||||
}
|
||||
response, err := api.inviteGroup(ctx, owner, manifest.GroupTopic, batch, memberIDs[offset:end])
|
||||
if err != nil {
|
||||
return fmt.Errorf("invite group batch %d: %w", batch, err)
|
||||
}
|
||||
for _, item := range response.Items {
|
||||
if item.Status != "joined" && item.Status != "already_member" {
|
||||
return fmt.Errorf("invite user %s status=%s reason=%s", item.UserID, item.Status, item.Reason)
|
||||
}
|
||||
}
|
||||
}
|
||||
count, err := api.countGroupMembers(ctx, owner, manifest.GroupTopic)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if count != len(manifest.Accounts) {
|
||||
return fmt.Errorf("group has %d members, expected %d", count, len(manifest.Accounts))
|
||||
}
|
||||
log.Printf("group ready members=%d topic_hash=%s", count, topicHash(manifest.GroupTopic))
|
||||
return nil
|
||||
}
|
||||
|
||||
func runVerify(ctx context.Context, cfg *Config, manifest *Manifest, api *apiClient) error {
|
||||
if len(manifest.Accounts) != cfg.UsersPerMachine*2 {
|
||||
return fmt.Errorf("manifest accounts=%d expected=%d", len(manifest.Accounts), cfg.UsersPerMachine*2)
|
||||
}
|
||||
for _, account := range manifest.Accounts {
|
||||
if account.UserID == "" || account.Token == "" || account.P2PTopic == "" {
|
||||
return fmt.Errorf("account %s is incomplete", account.Username)
|
||||
}
|
||||
}
|
||||
if manifest.GroupTopic == "" {
|
||||
return errors.New("manifest group topic is empty")
|
||||
}
|
||||
owner := manifestAccountMap(manifest)[accountKey("a", 1)]
|
||||
count, err := api.countGroupMembers(ctx, *owner, manifest.GroupTopic)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if count != len(manifest.Accounts) {
|
||||
return fmt.Errorf("group member count=%d expected=%d", count, len(manifest.Accounts))
|
||||
}
|
||||
log.Printf("verify passed accounts=%d pairs=%d group_members=%d", len(manifest.Accounts), cfg.UsersPerMachine, count)
|
||||
return nil
|
||||
}
|
||||
|
||||
func parallelAccounts(ctx context.Context, accounts []AccountRecord, workers int, fn func(AccountRecord) error) error {
|
||||
if workers <= 0 {
|
||||
workers = 1
|
||||
}
|
||||
workCtx, cancel := context.WithCancel(ctx)
|
||||
defer cancel()
|
||||
jobs := make(chan AccountRecord)
|
||||
errCh := make(chan error, 1)
|
||||
var wg sync.WaitGroup
|
||||
for i := 0; i < workers; i++ {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
for account := range jobs {
|
||||
if err := fn(account); err != nil {
|
||||
select {
|
||||
case errCh <- err:
|
||||
default:
|
||||
}
|
||||
cancel()
|
||||
return
|
||||
}
|
||||
}
|
||||
}()
|
||||
}
|
||||
go func() {
|
||||
defer close(jobs)
|
||||
for _, account := range accounts {
|
||||
select {
|
||||
case jobs <- account:
|
||||
case <-workCtx.Done():
|
||||
return
|
||||
}
|
||||
}
|
||||
}()
|
||||
wg.Wait()
|
||||
select {
|
||||
case err := <-errCh:
|
||||
return err
|
||||
default:
|
||||
return ctx.Err()
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,229 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"math/bits"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
)
|
||||
|
||||
type observation struct {
|
||||
topic string
|
||||
seq int64
|
||||
bits []uint64
|
||||
}
|
||||
|
||||
type deliveryRecorder struct {
|
||||
machine string
|
||||
localUsers int
|
||||
latency *DurationHistogram
|
||||
disconnects atomic.Int64
|
||||
received atomic.Int64
|
||||
runID atomic.Value
|
||||
historyMode atomic.Bool
|
||||
|
||||
mu sync.Mutex
|
||||
observed map[string]*observation
|
||||
lastSeq []map[string]int64
|
||||
duplicates int64
|
||||
outOfOrder int64
|
||||
malformed int64
|
||||
history map[string]bool
|
||||
}
|
||||
|
||||
func newDeliveryRecorder(runID, machine string, localUsers int) *deliveryRecorder {
|
||||
lastSeq := make([]map[string]int64, localUsers)
|
||||
for i := range lastSeq {
|
||||
lastSeq[i] = make(map[string]int64)
|
||||
}
|
||||
recorder := &deliveryRecorder{
|
||||
machine: machine, localUsers: localUsers, latency: NewDurationHistogram(),
|
||||
observed: make(map[string]*observation), lastSeq: lastSeq, history: make(map[string]bool),
|
||||
}
|
||||
recorder.runID.Store(runID)
|
||||
return recorder
|
||||
}
|
||||
|
||||
func (r *deliveryRecorder) setRunID(runID string) { r.runID.Store(runID) }
|
||||
|
||||
func (r *deliveryRecorder) onData(account AccountRecord, data dataPacket) {
|
||||
runID, _ := r.runID.Load().(string)
|
||||
mid, _ := data.Head["x-im-client-mid"].(string)
|
||||
if mid == "" || !strings.HasPrefix(mid, runID+"-") {
|
||||
return
|
||||
}
|
||||
var content struct {
|
||||
RunID string `json:"run_id"`
|
||||
SentAt int64 `json:"sent_at_unix_nano"`
|
||||
}
|
||||
if json.Unmarshal(data.Content, &content) != nil || content.RunID != runID {
|
||||
r.mu.Lock()
|
||||
r.malformed++
|
||||
r.mu.Unlock()
|
||||
return
|
||||
}
|
||||
if r.historyMode.Load() {
|
||||
r.mu.Lock()
|
||||
r.history[mid] = true
|
||||
r.mu.Unlock()
|
||||
return
|
||||
}
|
||||
index := account.Index - 1
|
||||
if index < 0 || index >= r.localUsers {
|
||||
return
|
||||
}
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
obs := r.observed[mid]
|
||||
if obs == nil {
|
||||
obs = &observation{topic: data.Topic, seq: data.Seq, bits: make([]uint64, (r.localUsers+63)/64)}
|
||||
r.observed[mid] = obs
|
||||
}
|
||||
word, bit := index/64, uint(index%64)
|
||||
mask := uint64(1) << bit
|
||||
if obs.bits[word]&mask != 0 {
|
||||
r.duplicates++
|
||||
return
|
||||
}
|
||||
obs.bits[word] |= mask
|
||||
r.received.Add(1)
|
||||
if previous := r.lastSeq[index][data.Topic]; previous > 0 && data.Seq < previous {
|
||||
r.outOfOrder++
|
||||
}
|
||||
if data.Seq > r.lastSeq[index][data.Topic] {
|
||||
r.lastSeq[index][data.Topic] = data.Seq
|
||||
}
|
||||
if content.SentAt > 0 {
|
||||
r.latency.Observe(time.Since(time.Unix(0, content.SentAt)))
|
||||
}
|
||||
}
|
||||
|
||||
func (r *deliveryRecorder) metricCounts() (received, duplicates, outOfOrder, disconnects int64) {
|
||||
r.mu.Lock()
|
||||
duplicates = r.duplicates
|
||||
outOfOrder = r.outOfOrder
|
||||
r.mu.Unlock()
|
||||
return r.received.Load(), duplicates, outOfOrder, r.disconnects.Load()
|
||||
}
|
||||
|
||||
func (r *deliveryRecorder) beginHistory() { r.historyMode.Store(true) }
|
||||
|
||||
func (r *deliveryRecorder) historyFound(mid string) bool {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
return r.history[mid]
|
||||
}
|
||||
|
||||
func (r *deliveryRecorder) onDisconnect(_ AccountRecord, _ error) {
|
||||
r.disconnects.Add(1)
|
||||
}
|
||||
|
||||
func (r *deliveryRecorder) summarize(base AgentSummary, ledger []AcceptedMessage) AgentSummary {
|
||||
base.DeliveryLatency = r.latency.Snapshot()
|
||||
base.Disconnects = r.disconnects.Load()
|
||||
runID, _ := r.runID.Load().(string)
|
||||
accepted := make(map[string]AcceptedMessage, len(ledger))
|
||||
const missingLimit = 100000
|
||||
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
base.Duplicates = r.duplicates
|
||||
base.OutOfOrder = r.outOfOrder
|
||||
if r.malformed > 0 {
|
||||
base.Errors = append(base.Errors, fmt.Sprintf("malformed run messages=%d", r.malformed))
|
||||
}
|
||||
for _, message := range ledger {
|
||||
if message.RunID != runID {
|
||||
continue
|
||||
}
|
||||
accepted[message.ClientMid] = message
|
||||
obs := r.observed[message.ClientMid]
|
||||
switch message.Kind {
|
||||
case "group":
|
||||
base.ExpectedGroup += int64(r.localUsers)
|
||||
for index := 1; index <= r.localUsers; index++ {
|
||||
if observationHas(obs, index) {
|
||||
base.ReceivedGroup++
|
||||
} else if len(base.Missing) < missingLimit {
|
||||
base.Missing = append(base.Missing, missingFor(message, index))
|
||||
}
|
||||
}
|
||||
case "p2p":
|
||||
if message.PeerMachine == r.machine && message.PeerIndex > 0 {
|
||||
base.ExpectedPeer++
|
||||
if observationHas(obs, message.PeerIndex) {
|
||||
base.ReceivedPeer++
|
||||
} else if len(base.Missing) < missingLimit {
|
||||
base.Missing = append(base.Missing, missingFor(message, message.PeerIndex))
|
||||
}
|
||||
}
|
||||
if message.SenderMachine == r.machine && message.SenderIndex > 0 {
|
||||
base.ExpectedEcho++
|
||||
if observationHas(obs, message.SenderIndex) {
|
||||
base.ReceivedEcho++
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
for mid, obs := range r.observed {
|
||||
if _, ok := accepted[mid]; !ok {
|
||||
base.UnexpectedMessages += bitsetCount(obs.bits)
|
||||
}
|
||||
}
|
||||
base.RealtimeLost = (base.ExpectedPeer - base.ReceivedPeer) + (base.ExpectedGroup - base.ReceivedGroup)
|
||||
if len(base.Missing) == missingLimit {
|
||||
base.Errors = append(base.Errors, "missing delivery details truncated at 100000")
|
||||
}
|
||||
return base
|
||||
}
|
||||
|
||||
func (r *deliveryRecorder) countRealtimeLostStored(ledger []AcceptedMessage) int64 {
|
||||
runID, _ := r.runID.Load().(string)
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
var count int64
|
||||
for _, message := range ledger {
|
||||
if message.RunID != runID || !r.history[message.ClientMid] {
|
||||
continue
|
||||
}
|
||||
obs := r.observed[message.ClientMid]
|
||||
switch message.Kind {
|
||||
case "group":
|
||||
for index := 1; index <= r.localUsers; index++ {
|
||||
if !observationHas(obs, index) {
|
||||
count++
|
||||
}
|
||||
}
|
||||
case "p2p":
|
||||
if message.PeerMachine == r.machine && message.PeerIndex > 0 && !observationHas(obs, message.PeerIndex) {
|
||||
count++
|
||||
}
|
||||
}
|
||||
}
|
||||
return count
|
||||
}
|
||||
|
||||
func observationHas(obs *observation, receiverIndex int) bool {
|
||||
if obs == nil || receiverIndex <= 0 {
|
||||
return false
|
||||
}
|
||||
index := receiverIndex - 1
|
||||
word, bit := index/64, uint(index%64)
|
||||
return word < len(obs.bits) && obs.bits[word]&(uint64(1)<<bit) != 0
|
||||
}
|
||||
|
||||
func bitsetCount(words []uint64) int64 {
|
||||
var count int64
|
||||
for _, word := range words {
|
||||
count += int64(bits.OnesCount64(word))
|
||||
}
|
||||
return count
|
||||
}
|
||||
|
||||
func missingFor(message AcceptedMessage, receiver int) MissingDelivery {
|
||||
return MissingDelivery{ClientMid: message.ClientMid, Kind: message.Kind, TopicHash: topicHash(message.Topic),
|
||||
Seq: message.Seq, ReceiverIndex: receiver}
|
||||
}
|
||||
@@ -0,0 +1,55 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestDeliveryRecorderSummarizesPeerGroupAndDuplicates(t *testing.T) {
|
||||
recorder := newDeliveryRecorder("run1", "a", 2)
|
||||
peer := AcceptedMessage{RunID: "run1", Kind: "p2p", ClientMid: "run1-p-b-0001-01", Topic: "p2p-one", Seq: 10,
|
||||
SenderMachine: "b", SenderIndex: 1, PeerMachine: "a", PeerIndex: 1}
|
||||
group := AcceptedMessage{RunID: "run1", Kind: "group", ClientMid: "run1-g-b-0001-01", Topic: "grp-one", Seq: 20,
|
||||
SenderMachine: "b", SenderIndex: 1}
|
||||
|
||||
recorder.onData(AccountRecord{Index: 1}, testData(peer, time.Now()))
|
||||
recorder.onData(AccountRecord{Index: 1}, testData(peer, time.Now()))
|
||||
recorder.onData(AccountRecord{Index: 1}, testData(group, time.Now()))
|
||||
|
||||
summary := recorder.summarize(AgentSummary{}, []AcceptedMessage{peer, group})
|
||||
if summary.ExpectedPeer != 1 || summary.ReceivedPeer != 1 {
|
||||
t.Fatalf("peer expected/received = %d/%d", summary.ExpectedPeer, summary.ReceivedPeer)
|
||||
}
|
||||
if summary.ExpectedGroup != 2 || summary.ReceivedGroup != 1 || summary.RealtimeLost != 1 {
|
||||
t.Fatalf("group expected/received/lost = %d/%d/%d", summary.ExpectedGroup, summary.ReceivedGroup, summary.RealtimeLost)
|
||||
}
|
||||
if summary.Duplicates != 1 || len(summary.Missing) != 1 || summary.Missing[0].ReceiverIndex != 2 {
|
||||
t.Fatalf("duplicates=%d missing=%+v", summary.Duplicates, summary.Missing)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDeliveryRecorderSeparatesHistoryFromRealtime(t *testing.T) {
|
||||
recorder := newDeliveryRecorder("run1", "a", 1)
|
||||
message := AcceptedMessage{RunID: "run1", Kind: "p2p", ClientMid: "run1-p-b-0001-01", Topic: "p2p-one", Seq: 10,
|
||||
SenderMachine: "b", SenderIndex: 1, PeerMachine: "a", PeerIndex: 1}
|
||||
recorder.beginHistory()
|
||||
recorder.onData(AccountRecord{Index: 1}, testData(message, time.Now()))
|
||||
|
||||
summary := recorder.summarize(AgentSummary{}, []AcceptedMessage{message})
|
||||
if summary.ReceivedPeer != 0 || summary.RealtimeLost != 1 {
|
||||
t.Fatalf("history response polluted realtime result: %+v", summary)
|
||||
}
|
||||
if !recorder.historyFound(message.ClientMid) {
|
||||
t.Fatal("history message was not recorded")
|
||||
}
|
||||
if got := recorder.countRealtimeLostStored([]AcceptedMessage{message}); got != 1 {
|
||||
t.Fatalf("realtime lost but persisted=%d", got)
|
||||
}
|
||||
}
|
||||
|
||||
func testData(message AcceptedMessage, sentAt time.Time) dataPacket {
|
||||
content, _ := json.Marshal(map[string]any{"run_id": message.RunID, "sent_at_unix_nano": sentAt.UnixNano()})
|
||||
return dataPacket{Topic: message.Topic, Seq: message.Seq,
|
||||
Head: map[string]any{"x-im-client-mid": message.ClientMid}, Content: content}
|
||||
}
|
||||
@@ -0,0 +1,365 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/tls"
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
"github.com/gorilla/websocket"
|
||||
)
|
||||
|
||||
type ctrlPacket struct {
|
||||
ID string `json:"id"`
|
||||
Topic string `json:"topic"`
|
||||
Code int `json:"code"`
|
||||
Text string `json:"text"`
|
||||
Params map[string]json.RawMessage `json:"params"`
|
||||
}
|
||||
|
||||
type dataPacket struct {
|
||||
Topic string `json:"topic"`
|
||||
From string `json:"from"`
|
||||
Timestamp time.Time `json:"ts"`
|
||||
Seq int64 `json:"seq"`
|
||||
Head map[string]any `json:"head"`
|
||||
Content json.RawMessage `json:"content"`
|
||||
}
|
||||
|
||||
type ctrlError struct {
|
||||
Code int
|
||||
Text string
|
||||
}
|
||||
|
||||
func (e *ctrlError) Error() string {
|
||||
return fmt.Sprintf("ctrl code=%d text=%s", e.Code, e.Text)
|
||||
}
|
||||
|
||||
type WSClient struct {
|
||||
account AccountRecord
|
||||
conn *websocket.Conn
|
||||
commandTimeout time.Duration
|
||||
onData func(AccountRecord, dataPacket)
|
||||
onDisconnect func(AccountRecord, error)
|
||||
|
||||
sendMu sync.Mutex
|
||||
mu sync.Mutex
|
||||
seq uint64
|
||||
waits map[string]chan ctrlPacket
|
||||
done chan struct{}
|
||||
closed atomic.Bool
|
||||
}
|
||||
|
||||
func connectWSClient(ctx context.Context, cfg *Config, account AccountRecord, apiKey, password string,
|
||||
onData func(AccountRecord, dataPacket), onDisconnect func(AccountRecord, error)) (*WSClient, error) {
|
||||
timeout, err := configDuration(cfg.CommandTimeout)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
endpoint, err := websocketEndpoint(account.Endpoint, apiKey)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
dialer := *websocket.DefaultDialer
|
||||
if cfg.InsecureTLS {
|
||||
dialer.TLSClientConfig = &tls.Config{InsecureSkipVerify: true} // #nosec G402: explicit load-test option.
|
||||
}
|
||||
conn, response, err := dialer.DialContext(ctx, endpoint, http.Header{})
|
||||
if err != nil {
|
||||
if response != nil {
|
||||
return nil, fmt.Errorf("websocket dial status=%d: %w", response.StatusCode, err)
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
client := &WSClient{
|
||||
account: account, conn: conn, commandTimeout: timeout, onData: onData, onDisconnect: onDisconnect,
|
||||
waits: make(map[string]chan ctrlPacket), done: make(chan struct{}),
|
||||
}
|
||||
go client.readLoop()
|
||||
if _, err = client.command(ctx, "hi", map[string]any{"hi": map[string]any{
|
||||
"id": client.nextID("hi"), "ver": "0.25", "ua": "im-pressure/1.0",
|
||||
"platf": "loadtest", "tenant": cfg.Tenant,
|
||||
}}); err != nil {
|
||||
client.Close()
|
||||
return nil, fmt.Errorf("hi %s: %w", account.Username, err)
|
||||
}
|
||||
login := map[string]any{"id": client.nextID("login")}
|
||||
if account.Token != "" {
|
||||
login["scheme"] = "token"
|
||||
login["secret"] = account.Token
|
||||
} else {
|
||||
login["scheme"] = "basic"
|
||||
login["secret"] = base64.StdEncoding.EncodeToString([]byte(account.Username + ":" + password))
|
||||
login["login_type"] = 0
|
||||
login["tenant_code"] = cfg.Tenant
|
||||
}
|
||||
ctrl, err := client.command(ctx, "login", map[string]any{"login": login})
|
||||
if err != nil && account.Token != "" && password != "" {
|
||||
login["id"] = client.nextID("login-basic")
|
||||
login["scheme"] = "basic"
|
||||
login["secret"] = base64.StdEncoding.EncodeToString([]byte(account.Username + ":" + password))
|
||||
login["login_type"] = 0
|
||||
login["tenant_code"] = cfg.Tenant
|
||||
ctrl, err = client.command(ctx, "login", map[string]any{"login": login})
|
||||
}
|
||||
if err != nil {
|
||||
client.Close()
|
||||
return nil, fmt.Errorf("login %s: %w", account.Username, err)
|
||||
}
|
||||
client.applyLogin(ctrl)
|
||||
return client, nil
|
||||
}
|
||||
|
||||
func connectForPrepare(ctx context.Context, cfg *Config, account AccountRecord, apiKey, password string) (*WSClient, bool, error) {
|
||||
timeout, err := configDuration(cfg.CommandTimeout)
|
||||
if err != nil {
|
||||
return nil, false, err
|
||||
}
|
||||
endpoint, err := websocketEndpoint(account.Endpoint, apiKey)
|
||||
if err != nil {
|
||||
return nil, false, err
|
||||
}
|
||||
dialer := *websocket.DefaultDialer
|
||||
if cfg.InsecureTLS {
|
||||
dialer.TLSClientConfig = &tls.Config{InsecureSkipVerify: true} // #nosec G402: explicit load-test option.
|
||||
}
|
||||
conn, _, err := dialer.DialContext(ctx, endpoint, http.Header{})
|
||||
if err != nil {
|
||||
return nil, false, err
|
||||
}
|
||||
client := &WSClient{account: account, conn: conn, commandTimeout: timeout, waits: make(map[string]chan ctrlPacket), done: make(chan struct{})}
|
||||
go client.readLoop()
|
||||
if _, err = client.command(ctx, "hi", map[string]any{"hi": map[string]any{
|
||||
"id": client.nextID("hi"), "ver": "0.25", "ua": "im-pressure-prepare/1.0", "tenant": cfg.Tenant,
|
||||
}}); err != nil {
|
||||
client.Close()
|
||||
return nil, false, err
|
||||
}
|
||||
secret := base64.StdEncoding.EncodeToString([]byte(account.Username + ":" + password))
|
||||
ctrl, loginErr := client.command(ctx, "login", map[string]any{"login": map[string]any{
|
||||
"id": client.nextID("login"), "scheme": "basic", "secret": secret,
|
||||
"login_type": 0, "tenant_code": cfg.Tenant,
|
||||
}})
|
||||
created := false
|
||||
if loginErr != nil {
|
||||
var ce *ctrlError
|
||||
if !errors.As(loginErr, &ce) || ce.Code != http.StatusUnauthorized {
|
||||
client.Close()
|
||||
return nil, false, loginErr
|
||||
}
|
||||
ctrl, err = client.command(ctx, "acc", map[string]any{"acc": map[string]any{
|
||||
"id": client.nextID("acc"), "user": "new", "scheme": "basic", "username": account.Username,
|
||||
"registration_type": 0, "tenant_code": cfg.Tenant, "secret": secret, "login": true,
|
||||
"desc": map[string]any{"public": map[string]any{"fn": strings.ToUpper(account.Username)}},
|
||||
}})
|
||||
if err != nil {
|
||||
client.Close()
|
||||
return nil, false, err
|
||||
}
|
||||
created = true
|
||||
}
|
||||
client.applyLogin(ctrl)
|
||||
return client, created, nil
|
||||
}
|
||||
|
||||
func websocketEndpoint(raw, apiKey string) (string, error) {
|
||||
u, err := url.Parse(strings.TrimSpace(raw))
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if u.Scheme != "ws" && u.Scheme != "wss" {
|
||||
return "", fmt.Errorf("unsupported websocket scheme %q", u.Scheme)
|
||||
}
|
||||
q := u.Query()
|
||||
if q.Get("apikey") == "" {
|
||||
q.Set("apikey", apiKey)
|
||||
}
|
||||
u.RawQuery = q.Encode()
|
||||
return u.String(), nil
|
||||
}
|
||||
|
||||
func (c *WSClient) applyLogin(ctrl ctrlPacket) {
|
||||
if raw := ctrl.Params["user"]; len(raw) > 0 {
|
||||
_ = json.Unmarshal(raw, &c.account.UserID)
|
||||
}
|
||||
if raw := ctrl.Params["token"]; len(raw) > 0 {
|
||||
_ = json.Unmarshal(raw, &c.account.Token)
|
||||
}
|
||||
if raw := ctrl.Params["expires"]; len(raw) > 0 {
|
||||
var value string
|
||||
if json.Unmarshal(raw, &value) == nil {
|
||||
c.account.TokenExpires, _ = time.Parse(time.RFC3339Nano, value)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (c *WSClient) Account() AccountRecord { return c.account }
|
||||
|
||||
func (c *WSClient) Subscribe(ctx context.Context, topic string) error {
|
||||
if topic == "" {
|
||||
return errors.New("empty topic")
|
||||
}
|
||||
id := c.nextID("sub")
|
||||
_, err := c.command(ctx, "sub", map[string]any{"sub": map[string]any{
|
||||
"id": id, "topic": topic, "get": map[string]any{"what": "desc"},
|
||||
}})
|
||||
return err
|
||||
}
|
||||
|
||||
func (c *WSClient) GetData(ctx context.Context, topic string, since, before int64, limit int) error {
|
||||
if topic == "" || since <= 0 || before <= since || limit <= 0 {
|
||||
return errors.New("invalid get data range")
|
||||
}
|
||||
id := c.nextID("get-data")
|
||||
_, err := c.command(ctx, "get", map[string]any{"get": map[string]any{
|
||||
"id": id, "topic": topic, "what": "data",
|
||||
"data": map[string]any{"since": since, "before": before, "limit": limit},
|
||||
}})
|
||||
return err
|
||||
}
|
||||
|
||||
func (c *WSClient) Publish(ctx context.Context, topic, clientMid, runID, scenario string, messageIndex int) (AcceptedMessage, error) {
|
||||
id := c.nextID("pub")
|
||||
sentAt := time.Now().UTC()
|
||||
ctrl, err := c.command(ctx, "pub", map[string]any{"pub": map[string]any{
|
||||
"id": id, "topic": topic, "noecho": false,
|
||||
"head": map[string]any{"mime": "application/json", "x-im-schema": 1, "x-im-type": "text", "x-im-client-mid": clientMid},
|
||||
"content": map[string]any{"text": "loadtest", "run_id": runID, "scenario": scenario,
|
||||
"sender_index": c.account.Index, "message_index": messageIndex, "sent_at_unix_nano": sentAt.UnixNano()},
|
||||
}})
|
||||
if err != nil {
|
||||
return AcceptedMessage{}, err
|
||||
}
|
||||
seq, err := ctrlInt64(ctrl.Params["seq"])
|
||||
if err != nil || seq <= 0 {
|
||||
return AcceptedMessage{}, fmt.Errorf("publish ACK has invalid seq: %w", err)
|
||||
}
|
||||
kind := "p2p"
|
||||
if strings.HasPrefix(topic, "grp") {
|
||||
kind = "group"
|
||||
}
|
||||
return AcceptedMessage{RunID: runID, Scenario: scenario, Kind: kind, ClientMid: clientMid,
|
||||
Topic: topic, Seq: seq, SenderMachine: c.account.Machine, SenderIndex: c.account.Index,
|
||||
SentAt: sentAt, AckAt: time.Now().UTC()}, nil
|
||||
}
|
||||
|
||||
func ctrlInt64(raw json.RawMessage) (int64, error) {
|
||||
if len(raw) == 0 {
|
||||
return 0, errors.New("missing value")
|
||||
}
|
||||
var number json.Number
|
||||
if err := json.Unmarshal(raw, &number); err == nil {
|
||||
return strconv.ParseInt(number.String(), 10, 64)
|
||||
}
|
||||
var text string
|
||||
if err := json.Unmarshal(raw, &text); err != nil {
|
||||
return 0, err
|
||||
}
|
||||
return strconv.ParseInt(text, 10, 64)
|
||||
}
|
||||
|
||||
func (c *WSClient) command(ctx context.Context, _ string, message map[string]any) (ctrlPacket, error) {
|
||||
id := nestedMessageID(message)
|
||||
if id == "" {
|
||||
return ctrlPacket{}, errors.New("command id is required")
|
||||
}
|
||||
wait := make(chan ctrlPacket, 1)
|
||||
c.mu.Lock()
|
||||
c.waits[id] = wait
|
||||
c.mu.Unlock()
|
||||
c.sendMu.Lock()
|
||||
err := c.conn.WriteJSON(message)
|
||||
c.sendMu.Unlock()
|
||||
if err != nil {
|
||||
c.removeWait(id)
|
||||
return ctrlPacket{}, err
|
||||
}
|
||||
timer := time.NewTimer(c.commandTimeout)
|
||||
defer timer.Stop()
|
||||
select {
|
||||
case ctrl := <-wait:
|
||||
if ctrl.Code < 200 || ctrl.Code >= 300 {
|
||||
return ctrl, &ctrlError{Code: ctrl.Code, Text: ctrl.Text}
|
||||
}
|
||||
return ctrl, nil
|
||||
case <-timer.C:
|
||||
c.removeWait(id)
|
||||
return ctrlPacket{}, fmt.Errorf("ctrl timeout id=%s", id)
|
||||
case <-ctx.Done():
|
||||
c.removeWait(id)
|
||||
return ctrlPacket{}, ctx.Err()
|
||||
case <-c.done:
|
||||
c.removeWait(id)
|
||||
return ctrlPacket{}, errors.New("websocket closed")
|
||||
}
|
||||
}
|
||||
|
||||
func nestedMessageID(message map[string]any) string {
|
||||
for _, value := range message {
|
||||
if body, ok := value.(map[string]any); ok {
|
||||
if id, ok := body["id"].(string); ok {
|
||||
return id
|
||||
}
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func (c *WSClient) nextID(prefix string) string {
|
||||
n := atomic.AddUint64(&c.seq, 1)
|
||||
return fmt.Sprintf("%s-%s-%04d-%d", prefix, c.account.Machine, c.account.Index, n)
|
||||
}
|
||||
|
||||
func (c *WSClient) removeWait(id string) {
|
||||
c.mu.Lock()
|
||||
delete(c.waits, id)
|
||||
c.mu.Unlock()
|
||||
}
|
||||
|
||||
func (c *WSClient) readLoop() {
|
||||
defer close(c.done)
|
||||
for {
|
||||
var envelope map[string]json.RawMessage
|
||||
if err := c.conn.ReadJSON(&envelope); err != nil {
|
||||
if !c.closed.Load() && c.onDisconnect != nil {
|
||||
c.onDisconnect(c.account, err)
|
||||
}
|
||||
return
|
||||
}
|
||||
if raw := envelope["ctrl"]; len(raw) > 0 {
|
||||
var ctrl ctrlPacket
|
||||
if json.Unmarshal(raw, &ctrl) == nil && ctrl.ID != "" {
|
||||
c.mu.Lock()
|
||||
wait := c.waits[ctrl.ID]
|
||||
delete(c.waits, ctrl.ID)
|
||||
c.mu.Unlock()
|
||||
if wait != nil {
|
||||
wait <- ctrl
|
||||
}
|
||||
}
|
||||
continue
|
||||
}
|
||||
if raw := envelope["data"]; len(raw) > 0 && c.onData != nil {
|
||||
var data dataPacket
|
||||
if json.Unmarshal(raw, &data) == nil {
|
||||
c.onData(c.account, data)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (c *WSClient) Close() {
|
||||
if c == nil || !c.closed.CompareAndSwap(false, true) {
|
||||
return
|
||||
}
|
||||
_ = c.conn.Close()
|
||||
}
|
||||
Reference in New Issue
Block a user