fix: keep load connections through verification

loadtest/im-pressure/agent.go

@@ -154,31 +154,31 @@ func runAgent(ctx context.Context, cfg *Config, manifest *Manifest, machine, coo
 		defer close(maintenanceDone)
 		run.maintainConnections(maintenanceCtx, accounts, manifest.GroupTopic, apiKey, password)
 	}()
+	var stopMaintenanceOnce sync.Once
+	stopConnectionMaintenance := func() {
+		stopMaintenanceOnce.Do(func() {
+			stopMaintenance()
+			<-maintenanceDone
+		})
+	}
+	defer stopConnectionMaintenance()
 	start, err := coord.ready(ctx, func() AgentReady {
 		return AgentReady{
 			Machine: machine, ReadyConnections: run.aliveConnections(), ExpectedConnections: len(accounts),
 		}
 	})
 	if err != nil {
-		stopMaintenance()
-		<-maintenanceDone
 		return err
 	}
 	if !runIDPattern.MatchString(start.RunID) || !validScenario(start.Scenario) {
-		stopMaintenance()
-		<-maintenanceDone
 		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 {
-		stopMaintenance()
-		<-maintenanceDone
 		return err
 	}
-	stopMaintenance()
-	<-maintenanceDone
 	if alive := run.aliveConnections(); alive != len(accounts) {
 		return fmt.Errorf("connections alive at scheduled start=%d expected=%d", alive, len(accounts))
 	}
@@ -241,6 +241,9 @@ func runAgent(ctx context.Context, cfg *Config, manifest *Manifest, machine, coo
 	if err != nil {
 		return err
 	}
+	// Both agents have completed history verification. Connection maintenance can
+	// now stop without making the other agent verify through stale sockets.
+	stopConnectionMaintenance()
 	summary.HistoryChecked = int64(history.checked)
 	summary.HistoryFound = int64(history.found)
 	summary.HistoryUnackedChecked = int64(history.unackedChecked)
@@ -601,29 +604,57 @@ func (r *agentRun) maintainConnections(ctx context.Context, accounts []AccountRe
 		if ctx.Err() != nil {
 			return
 		}
+		jobs := make(chan AccountRecord)
 		var wg sync.WaitGroup
-		for index, account := range byIndex {
+		workers := 16
+		if workers > len(accounts) {
+			workers = len(accounts)
+		}
+		for range workers {
+			wg.Add(1)
+			go func() {
+				defer wg.Done()
+				for account := range jobs {
+					index := account.Index
+					previous := r.client(index)
+					if previous != nil && previous.alive() {
+						continue
+					}
+					if previous != nil {
+						previous.Close()
+					}
+					client, err := r.connectAccountWithRetry(ctx, account, groupTopic, apiKey, password)
+					if err != nil {
+						if ctx.Err() == nil {
+							log.Printf("account=%s reconnect failed: %v", account.Username, err)
+						}
+						continue
+					}
+					r.setClient(index, client)
+					log.Printf("account=%s connection restored", account.Username)
+				}
+			}()
+		}
+		indexes := make([]int, 0, len(byIndex))
+		for index := range byIndex {
+			indexes = append(indexes, index)
+		}
+		sort.Ints(indexes)
+		for _, index := range indexes {
+			account := byIndex[index]
 			current := r.client(index)
 			if current != nil && current.alive() {
 				continue
 			}
-			wg.Add(1)
-			go func(index int, account AccountRecord, previous *WSClient) {
-				defer wg.Done()
-				if previous != nil {
-					previous.Close()
-				}
-				client, err := r.connectAccountWithRetry(ctx, account, groupTopic, apiKey, password)
-				if err != nil {
-					if ctx.Err() == nil {
-						log.Printf("account=%s pre-start reconnect failed: %v", account.Username, err)
-					}
-					return
-				}
-				r.setClient(index, client)
-				log.Printf("account=%s pre-start connection restored", account.Username)
-			}(index, account, current)
+			select {
+			case jobs <- account:
+			case <-ctx.Done():
+				close(jobs)
+				wg.Wait()
+				return
+			}
 		}
+		close(jobs)
 		wg.Wait()
 		select {
 		case <-ctx.Done():
@@ -714,12 +745,11 @@ func scaledUsers(total, numerator, denominator int) int {
 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.client(index)
-		if client == nil {
+		if r.client(index) == nil {
 			return fmt.Errorf("missing client %d", index)
 		}
 		wg.Add(1)
-		go func(index int, client *WSClient) {
+		go func(index int) {
 			defer wg.Done()
 			jitter := time.Duration((index*37)%200) * time.Millisecond
 			for messageIndex := 1; messageIndex <= messages; messageIndex++ {
@@ -728,6 +758,11 @@ func (r *agentRun) scheduleP2P(ctx context.Context, start time.Time, users, mess
 					return
 				}
 				mid := fmt.Sprintf("%s-p-%s-%04d-%02d", r.runID, r.machine, index, messageIndex)
+				client := r.client(index)
+				if client == nil || !client.alive() {
+					r.recordUnavailablePublish(mid)
+					continue
+				}
 				accepted, err := r.publish(ctx, client, client.account.P2PTopic, mid, messageIndex)
 				if err == nil {
 					accepted.PeerIndex = index
@@ -739,7 +774,7 @@ func (r *agentRun) scheduleP2P(ctx context.Context, start time.Time, users, mess
 					r.addAccepted(accepted)
 				}
 			}
-		}(index, client)
+		}(index)
 	}
 	wg.Wait()
 	return ctx.Err()
@@ -757,10 +792,9 @@ func (r *agentRun) scheduleGroup(ctx context.Context, start time.Time, topic str
 	}
 	var wg sync.WaitGroup
 	for index := first; index <= last; index++ {
-		client := r.client(index)
 		position := index - first
 		wg.Add(1)
-		go func(index, position int, client *WSClient) {
+		go func(index, position int) {
 			defer wg.Done()
 			for round := 1; round <= rounds; round++ {
 				target := start.Add(time.Duration(round-1)*roundDuration + time.Duration(position)*200*time.Millisecond + offset)
@@ -768,17 +802,32 @@ func (r *agentRun) scheduleGroup(ctx context.Context, start time.Time, topic str
 					return
 				}
 				mid := fmt.Sprintf("%s-g-%s-%04d-%02d", r.runID, r.machine, index, round)
+				client := r.client(index)
+				if client == nil || !client.alive() {
+					r.recordUnavailablePublish(mid)
+					continue
+				}
 				accepted, err := r.publish(ctx, client, topic, mid, round)
 				if err == nil {
 					r.addAccepted(accepted)
 				}
 			}
-		}(index, position, client)
+		}(index, position)
 	}
 	wg.Wait()
 	return ctx.Err()
 }
 
+func (r *agentRun) recordUnavailablePublish(mid string) {
+	r.attempted.Add(1)
+	r.publishSocket.Add(1)
+	r.mu.Lock()
+	if len(r.errors) < 1000 {
+		r.errors = append(r.errors, fmt.Sprintf("publish %s: websocket connection unavailable", mid))
+	}
+	r.mu.Unlock()
+}
+
 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)