diff --git a/loadtest/im-pressure/agent.go b/loadtest/im-pressure/agent.go index b3db924..47c96ef 100644 --- a/loadtest/im-pressure/agent.go +++ b/loadtest/im-pressure/agent.go @@ -254,19 +254,11 @@ func (r *agentRun) connect(ctx context.Context, accounts []AccountRecord, groupT wg.Add(1) go func() { defer wg.Done() - client, connectErr := connectWSClient(ctx, r.cfg, account, apiKey, password, r.recorder.onData, r.recorder.onDisconnect) + client, connectErr := r.connectAccountWithRetry(ctx, account, groupTopic, apiKey, password) 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() @@ -292,6 +284,34 @@ func (r *agentRun) connect(ctx context.Context, accounts []AccountRecord, groupT return nil } +func (r *agentRun) connectAccountWithRetry(ctx context.Context, account AccountRecord, groupTopic, apiKey, password string) (*WSClient, error) { + var lastErr error + for attempt := 1; attempt <= 6; attempt++ { + client, err := connectWSClient(ctx, r.cfg, account, apiKey, password, r.recorder.onData, r.recorder.onDisconnect) + if err == nil { + if err = client.Subscribe(ctx, account.P2PTopic); err == nil { + err = client.Subscribe(ctx, groupTopic) + } + } + if err == nil { + return client, nil + } + if client != nil { + client.Close() + } + lastErr = err + if ctx.Err() != nil { + return nil, ctx.Err() + } + delay := time.Duration(attempt*2) * time.Second + log.Printf("account=%s connect attempt=%d failed; retry_in=%s: %v", account.Username, attempt, delay, err) + if err = waitContext(ctx, delay); err != nil { + return nil, err + } + } + return nil, fmt.Errorf("connect %s failed after retries: %w", account.Username, lastErr) +} + func (r *agentRun) execute(ctx context.Context, start time.Time, groupTopic string) error { if groupTopic == "" { return fmt.Errorf("group topic is empty")