pro/im_server
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: support verified bulk load-test setup

Browse Source
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
Download Patch File Download Diff File Expand all files Collapse all files
loadtest/im-pressure/prepare.go
+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)