feat: support verified bulk load-test setup

This commit is contained in:
richjohnrambo
2026-09-27 01:34:39 +08:00
parent 021169555a
commit b3b5dd55a6
8 changed files with 272 additions and 47 deletions
+2
View File
@@ -1,5 +1,7 @@
/im-pressure
/config.json
/config.local*.json
/.env.test
/manifest*.json
/results/
/*-summary.json
+2
View File
@@ -70,6 +70,8 @@ export IM_LOAD_PASSWORD='...'
修改 `config.json`:`http_base` 指向 IM HTTP API;`endpoints.a/b` 分别填写两台机器应直连的 WS Pod/LB 地址,多个地址时用户会轮询分配。正式测试保持 `insecure_tls=false`。
测试环境要求短信验证注册时,将 `registration_phone_prefix` 配置为 7 位数字。工具为机器 A 追加后缀 `0001` 起的序号,为机器 B 追加后缀 `5001` 起的序号,生成互不冲突的 11 位测试手机号;验证码从 `verification_code_env` 指定的环境变量读取。正式环境不得使用固定测试验证码。
`endpoints.a` 是机器 A 的 1000 个用户使用的 WS Pod/LB 地址,`endpoints.b` 是机器 B 的 1000 个用户使用的地址。要验证跨 Pod 投递时,应让 A、B 连接到不同 Pod 或不同的定向入口;如果都填同一个普通 LB,连接落在哪个 Pod 将由 LB 决定。
机器 B 同样需要设置 `IM_API_KEY` 和 `IM_LOAD_PASSWORD`。密码用于 Token 失效时自动回退 Basic 登录,不能只在机器 A 设置。
+2
View File
@@ -3,6 +3,8 @@
"http_base": "https://im-api.example.com",
"api_key_env": "IM_API_KEY",
"password_env": "IM_LOAD_PASSWORD",
"verification_code_env": "IM_VERIFICATION_CODE",
"registration_phone_prefix": "",
"users_per_machine": 1000,
"endpoints": {
"a": ["wss://im-ws-pod-a.example.com/v0/channels"],
+41 -3
View File
@@ -18,6 +18,7 @@ import (
type apiClient struct {
baseURL string
apiKey string
tenant string
client *http.Client
seq atomic.Uint64
}
@@ -47,8 +48,8 @@ func newAPIClient(cfg *Config, apiKey string) *apiClient {
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},
baseURL: strings.TrimRight(cfg.HTTPBase, "/"), apiKey: apiKey, tenant: cfg.Tenant,
client: &http.Client{Timeout: 2 * time.Minute, Transport: transport},
}
}
@@ -66,6 +67,7 @@ func (a *apiClient) do(ctx context.Context, method, path, token, idempotencyKey
return err
}
req.Header.Set("X-MeetingIM-APIKey", a.apiKey)
req.Header.Set("X-Tenant-Code", a.tenant)
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 != "" {
@@ -107,6 +109,40 @@ type privacyResponse struct {
Version int `json:"version"`
}
type registrationCredential struct {
ChallengeID string
VerificationToken string
}
func (a *apiClient) createRegistrationCredential(ctx context.Context, account AccountRecord, code string) (*registrationCredential, error) {
if account.Phone == "" {
return nil, nil
}
var challenge struct {
ChallengeID string `json:"challenge_id"`
}
key := fmt.Sprintf("loadtest-verification-%s-%d", account.Username, time.Now().UnixNano())
body := map[string]any{"scene": "register", "registration_type": 0, "tenant_code": a.tenant, "phone": account.Phone}
if err := a.do(ctx, http.MethodPost, "/v1/im/auth/verification-codes", "", key, body, &challenge); err != nil {
return nil, fmt.Errorf("create verification for %s: %w", account.Username, err)
}
if challenge.ChallengeID == "" {
return nil, fmt.Errorf("create verification for %s: missing challenge_id", account.Username)
}
var verified struct {
VerificationToken string `json:"verification_token"`
}
key = fmt.Sprintf("loadtest-verify-%s-%d", account.Username, time.Now().UnixNano())
path := "/v1/im/auth/verification-codes/" + url.PathEscape(challenge.ChallengeID) + "/verify"
if err := a.do(ctx, http.MethodPost, path, "", key, map[string]any{"code": code}, &verified); err != nil {
return nil, fmt.Errorf("verify code for %s: %w", account.Username, err)
}
if verified.VerificationToken == "" {
return nil, fmt.Errorf("verify code for %s: missing verification_token", account.Username)
}
return &registrationCredential{ChallengeID: challenge.ChallengeID, VerificationToken: verified.VerificationToken}, nil
}
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 {
@@ -211,8 +247,10 @@ type groupInvitationResponse struct {
func (a *apiClient) inviteGroup(ctx context.Context, owner AccountRecord, topic string, batch int, userIDs []string) (*groupInvitationResponse, error) {
var result groupInvitationResponse
membersHash := topicHash(strings.Join(userIDs, ","))
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)
fmt.Sprintf("loadtest-group-%s-invite-%03d-%s", topicHash(topic), batch, membersHash),
map[string]any{"user_ids": userIDs, "message": "loadtest"}, &result)
return &result, err
}
+69
View File
@@ -29,6 +29,8 @@ type Config struct {
P2PInterval string `json:"p2p_interval"`
DrainP2P string `json:"drain_p2p"`
DrainGroup string `json:"drain_group"`
PhonePrefix string `json:"registration_phone_prefix"`
VerifyCodeEnv string `json:"verification_code_env"`
}
func loadConfig(path string) (*Config, error) {
@@ -75,6 +77,14 @@ func loadConfig(path string) (*Config, error) {
return nil, fmt.Errorf("at least one endpoint is required for machine %s", machine)
}
}
if cfg.PhonePrefix != "" {
if len(cfg.PhonePrefix) != 7 || strings.Trim(cfg.PhonePrefix, "0123456789") != "" {
return nil, errors.New("registration_phone_prefix must contain exactly 7 digits")
}
if cfg.UsersPerMachine > 4999 {
return nil, errors.New("registration phone suffix supports at most 4999 users per machine")
}
}
return &cfg, nil
}
@@ -102,6 +112,18 @@ func (c *Config) password() (string, error) {
return value, nil
}
func (c *Config) verificationCode() (string, error) {
name := c.VerifyCodeEnv
if name == "" {
name = "IM_VERIFICATION_CODE"
}
value := strings.TrimSpace(os.Getenv(name))
if c.PhonePrefix != "" && 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 {
@@ -122,6 +144,7 @@ type AccountRecord struct {
Machine string `json:"machine"`
Index int `json:"index"`
Username string `json:"username"`
Phone string `json:"phone,omitempty"`
UserID string `json:"user_id"`
Token string `json:"token"`
TokenExpires time.Time `json:"token_expires,omitempty"`
@@ -133,10 +156,12 @@ 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++ {
phone, _ := registrationPhone(cfg, machine, i)
m.Accounts = append(m.Accounts, AccountRecord{
Machine: machine,
Index: i,
Username: fmt.Sprintf("lt%s%04d", machine, i),
Phone: phone,
Endpoint: cfg.Endpoints[machine][(i-1)%len(cfg.Endpoints[machine])],
})
}
@@ -158,10 +183,54 @@ func applyConfigEndpoints(cfg *Config, manifest *Manifest) error {
return fmt.Errorf("invalid endpoint assignment for %s", accountKey(account.Machine, account.Index))
}
account.Endpoint = endpoints[(account.Index-1)%len(endpoints)]
phone, err := registrationPhone(cfg, account.Machine, account.Index)
if err != nil {
return err
}
account.Phone = phone
}
return nil
}
func reconcileManifest(cfg *Config, existing *Manifest) (*Manifest, error) {
if existing.Tenant != cfg.Tenant {
return nil, fmt.Errorf("manifest tenant=%s config tenant=%s", existing.Tenant, cfg.Tenant)
}
desired := newManifest(cfg)
if len(existing.Accounts) > len(desired.Accounts) {
return nil, fmt.Errorf("refusing to shrink manifest accounts from %d to %d", len(existing.Accounts), len(desired.Accounts))
}
desired.GroupTopic = existing.GroupTopic
desired.CreatedAt = existing.CreatedAt
old := manifestAccountMap(existing)
for i := range desired.Accounts {
account := &desired.Accounts[i]
if previous := old[accountKey(account.Machine, account.Index)]; previous != nil {
account.UserID = previous.UserID
account.Token = previous.Token
account.TokenExpires = previous.TokenExpires
account.P2PTopic = previous.P2PTopic
}
}
return desired, nil
}
func registrationPhone(cfg *Config, machine string, index int) (string, error) {
if cfg.PhonePrefix == "" {
return "", nil
}
ordinal := index
if machine == "b" {
ordinal += 5000
} else if machine != "a" {
return "", fmt.Errorf("invalid machine %q", machine)
}
if index <= 0 || ordinal > 9999 {
return "", fmt.Errorf("invalid phone ordinal %d", ordinal)
}
return fmt.Sprintf("%s%04d", cfg.PhonePrefix, ordinal), nil
}
func loadManifest(path string) (*Manifest, error) {
raw, err := os.ReadFile(path)
if err != nil {
+36
View File
@@ -43,3 +43,39 @@ func TestApplyConfigEndpointsRefreshesManifest(t *testing.T) {
t.Fatalf("endpoints were not refreshed: %+v", manifest.Accounts)
}
}
func TestRegistrationPhoneUsesDistinctMachineRanges(t *testing.T) {
cfg := &Config{UsersPerMachine: 1000, PhonePrefix: "1990927"}
a, err := registrationPhone(cfg, "a", 1)
if err != nil || a != "19909270001" {
t.Fatalf("machine A phone=%q err=%v", a, err)
}
b, err := registrationPhone(cfg, "b", 1)
if err != nil || b != "19909275001" {
t.Fatalf("machine B phone=%q err=%v", b, err)
}
last, err := registrationPhone(cfg, "b", 1000)
if err != nil || last != "19909276000" {
t.Fatalf("last phone=%q err=%v", last, err)
}
}
func TestReconcileManifestExpandsSmokeRun(t *testing.T) {
existing := &Manifest{Version: manifestVersion, Tenant: "allrich", GroupTopic: "grp-one", Accounts: []AccountRecord{
{Machine: "a", Index: 1, Username: "lta0001", UserID: "usr-a", Token: "token-a", P2PTopic: "p2p-one"},
{Machine: "b", Index: 1, Username: "ltb0001", UserID: "usr-b", Token: "token-b", P2PTopic: "p2p-one"},
}}
cfg := &Config{Tenant: "allrich", UsersPerMachine: 2, PhonePrefix: "1990927", Endpoints: map[string][]string{
"a": {"wss://a"}, "b": {"wss://b"},
}}
expanded, err := reconcileManifest(cfg, existing)
if err != nil {
t.Fatal(err)
}
if len(expanded.Accounts) != 4 || expanded.GroupTopic != "grp-one" || expanded.Accounts[0].Token != "token-a" {
t.Fatalf("unexpected expanded manifest: %+v", expanded)
}
if expanded.Accounts[3].Phone != "19909275002" {
t.Fatalf("machine B phone=%q", expanded.Accounts[3].Phone)
}
}
+105 -40
View File
@@ -9,6 +9,7 @@ import (
"sort"
"sync"
"sync/atomic"
"time"
)
func runPrepare(ctx context.Context, cfg *Config, manifestPath string) error {
@@ -20,21 +21,28 @@ func runPrepare(ctx context.Context, cfg *Config, manifestPath string) error {
if err != nil {
return err
}
verificationCode, err := cfg.verificationCode()
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))
manifest, err = reconcileManifest(cfg, existing)
if err != nil {
return err
}
log.Printf("resuming manifest with %d existing accounts, target=%d", len(existing.Accounts), len(manifest.Accounts))
}
if err = applyConfigEndpoints(cfg, manifest); err != nil {
return err
}
if err = prepareAccounts(ctx, cfg, manifest, apiKey, password); err != nil {
api := newAPIClient(cfg, apiKey)
if err = prepareAccounts(ctx, cfg, manifest, api, password, verificationCode); 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
}
@@ -50,12 +58,16 @@ func runPrepare(ctx context.Context, cfg *Config, manifestPath string) error {
return runVerify(ctx, cfg, manifest, api)
}
func prepareAccounts(ctx context.Context, cfg *Config, manifest *Manifest, apiKey, password string) error {
func prepareAccounts(ctx context.Context, cfg *Config, manifest *Manifest, api *apiClient, password, verificationCode 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 account.UserID != "" && account.Token != "" && account.TokenExpires.After(time.Now().Add(5*time.Minute)) {
existing.Add(1)
return nil
}
client, wasCreated, err := connectForPrepareWithRetry(ctx, cfg, account, api, password, verificationCode)
if err != nil {
return fmt.Errorf("account %s: %w", account.Username, err)
}
@@ -84,6 +96,27 @@ func prepareAccounts(ctx context.Context, cfg *Config, manifest *Manifest, apiKe
return err
}
func connectForPrepareWithRetry(ctx context.Context, cfg *Config, account AccountRecord, api *apiClient,
password, verificationCode string) (*WSClient, bool, error) {
var lastErr error
for attempt := 1; attempt <= 6; attempt++ {
client, created, err := connectForPrepare(ctx, cfg, account, api, password, verificationCode)
if err == nil {
return client, created, nil
}
lastErr = err
if ctx.Err() != nil {
return nil, false, ctx.Err()
}
delay := time.Duration(attempt*5) * time.Second
log.Printf("account=%s prepare attempt=%d failed; retry_in=%s: %v", account.Username, attempt, delay, err)
if err = waitContext(ctx, delay); err != nil {
return nil, false, err
}
}
return nil, false, fmt.Errorf("account %s prepare failed after retries: %w", account.Username, lastErr)
}
func prepareFriendships(ctx context.Context, cfg *Config, manifest *Manifest, api *apiClient) error {
accounts := manifestAccountMap(manifest)
pairs := make([]AccountRecord, 0, cfg.UsersPerMachine)
@@ -92,42 +125,47 @@ func prepareFriendships(ctx context.Context, cfg *Config, manifest *Manifest, ap
}
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)
return retryPrepareOperation(ctx, fmt.Sprintf("friend-pair-%04d", a.Index), func() error {
b := *accounts[accountKey("b", a.Index)]
if a.P2PTopic != "" && b.P2PTopic == a.P2PTopic {
return nil
}
} 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)
if a.UserID == "" || b.UserID == "" {
return fmt.Errorf("pair %d has missing user ID", a.Index)
}
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)
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
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)
@@ -137,7 +175,9 @@ func prepareFriendships(ctx context.Context, cfg *Config, manifest *Manifest, ap
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)
return retryPrepareOperation(ctx, "privacy-"+account.Username, func() error {
return api.disableGroupInviteVerification(ctx, account)
})
}); err != nil {
return fmt.Errorf("prepare group invite privacy: %w", err)
}
@@ -163,7 +203,12 @@ func prepareGroup(ctx context.Context, cfg *Config, manifest *Manifest, api *api
if end > len(memberIDs) {
end = len(memberIDs)
}
response, err := api.inviteGroup(ctx, owner, manifest.GroupTopic, batch, memberIDs[offset:end])
var response *groupInvitationResponse
err := retryPrepareOperation(ctx, fmt.Sprintf("group-invite-%03d", batch), func() error {
var inviteErr error
response, inviteErr = api.inviteGroup(ctx, owner, manifest.GroupTopic, batch, memberIDs[offset:end])
return inviteErr
})
if err != nil {
return fmt.Errorf("invite group batch %d: %w", batch, err)
}
@@ -184,6 +229,26 @@ func prepareGroup(ctx context.Context, cfg *Config, manifest *Manifest, api *api
return nil
}
func retryPrepareOperation(ctx context.Context, label string, operation func() error) error {
var lastErr error
for attempt := 1; attempt <= 6; attempt++ {
if err := operation(); err == nil {
return nil
} else {
lastErr = err
}
if ctx.Err() != nil {
return ctx.Err()
}
delay := time.Duration(attempt*5) * time.Second
log.Printf("operation=%s attempt=%d failed; retry_in=%s: %v", label, attempt, delay, lastErr)
if err := waitContext(ctx, delay); err != nil {
return err
}
}
return fmt.Errorf("operation %s failed after retries: %w", label, lastErr)
}
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)
+15 -4
View File
@@ -119,12 +119,12 @@ func connectWSClient(ctx context.Context, cfg *Config, account AccountRecord, ap
return client, nil
}
func connectForPrepare(ctx context.Context, cfg *Config, account AccountRecord, apiKey, password string) (*WSClient, bool, error) {
func connectForPrepare(ctx context.Context, cfg *Config, account AccountRecord, api *apiClient, password, verificationCode string) (*WSClient, bool, error) {
timeout, err := configDuration(cfg.CommandTimeout)
if err != nil {
return nil, false, err
}
endpoint, err := websocketEndpoint(account.Endpoint, apiKey)
endpoint, err := websocketEndpoint(account.Endpoint, api.apiKey)
if err != nil {
return nil, false, err
}
@@ -156,11 +156,22 @@ func connectForPrepare(ctx context.Context, cfg *Config, account AccountRecord,
client.Close()
return nil, false, loginErr
}
ctrl, err = client.command(ctx, "acc", map[string]any{"acc": map[string]any{
credential, credentialErr := api.createRegistrationCredential(ctx, account, verificationCode)
if credentialErr != nil {
client.Close()
return nil, false, credentialErr
}
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 credential != nil {
acc["phone"] = account.Phone
acc["cred"] = []map[string]any{{"meth": "tel", "val": account.Phone,
"resp": credential.VerificationToken, "params": map[string]any{"challenge_id": credential.ChallengeID}}}
}
ctrl, err = client.command(ctx, "acc", map[string]any{"acc": acc})
if err != nil {
client.Close()
return nil, false, err