fix: retry load-test websocket setup

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")