im_server/loadtest/im-pressure/agent.go

package main

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"log"
	"os"
	"path/filepath"
	"regexp"
	"sort"
	"strings"
	"sync"
	"sync/atomic"
	"time"
)

var runIDPattern = regexp.MustCompile(`^[a-zA-Z0-9_-]{1,24}$`)

type agentRun struct {
	cfg            *Config
	machine        string
	runID          string
	scenario       string
	clients        map[int]*WSClient
	clientsMu      sync.RWMutex
	recorder       *deliveryRecorder
	ack            *DurationHistogram
	attempted      atomic.Int64
	ready          atomic.Int64
	publish403     atomic.Int64
	publish502     atomic.Int64
	publish503     atomic.Int64
	publishTimeout atomic.Int64
	publishSocket  atomic.Int64
	publishOther   atomic.Int64

	mu       sync.Mutex
	accepted []AcceptedMessage
	errors   []string
}

func validScenario(scenario string) bool {
	return scenario == "p2p" || scenario == "group" || scenario == "mixed"
}

func preflightP2P(ctx context.Context, cfg *Config, manifest *Manifest, local []AccountRecord, apiKey string) error {
	accounts := manifestAccountMap(manifest)
	api := newAPIClient(cfg, apiKey)
	jobs := make(chan AccountRecord)
	errCh := make(chan error, len(local))
	var wg sync.WaitGroup
	workers := cfg.PreflightWorkers
	if workers > len(local) {
		workers = len(local)
	}
	for i := 0; i < workers; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			for source := range jobs {
				peerMachine := "a"
				if source.Machine == "a" {
					peerMachine = "b"
				}
				peer := accounts[accountKey(peerMachine, source.Index)]
				if peer == nil {
					errCh <- fmt.Errorf("p2p preflight %s-%04d: peer is missing", source.Machine, source.Index)
					continue
				}
				conversation, err := api.ensureP2P(ctx, source, *peer)
				if err == nil && (conversation.Relationship != "friend" || conversation.Topic != source.P2PTopic) {
					err = fmt.Errorf("relationship=%s topic_match=%t", conversation.Relationship, conversation.Topic == source.P2PTopic)
				}
				if err != nil {
					errCh <- fmt.Errorf("p2p preflight %s-%04d: %w", source.Machine, source.Index, err)
				}
			}
		}()
	}
	for _, account := range local {
		jobs <- account
	}
	close(jobs)
	wg.Wait()
	close(errCh)
	failures := make([]string, 0, len(errCh))
	for err := range errCh {
		if len(failures) < 20 {
			failures = append(failures, err.Error())
		}
	}
	if len(failures) > 0 {
		return fmt.Errorf("p2p permission preflight failed: %s", strings.Join(failures, "; "))
	}
	return nil
}

func runAgent(ctx context.Context, cfg *Config, manifest *Manifest, machine, coordinatorURL, output, metricsListen string) error {
	if machine != "a" && machine != "b" {
		return fmt.Errorf("machine must be a or b")
	}
	apiKey, err := cfg.apiKey()
	if err != nil {
		return err
	}
	password := ""
	if name := cfg.PasswordEnv; name != "" {
		password = os.Getenv(name)
	} else {
		password = os.Getenv("IM_LOAD_PASSWORD")
	}
	accounts := accountsForMachine(manifest, machine)
	if len(accounts) != cfg.UsersPerMachine {
		return fmt.Errorf("machine %s accounts=%d expected=%d", machine, len(accounts), cfg.UsersPerMachine)
	}
	for i := range accounts {
		accounts[i].Endpoint = cfg.Endpoints[machine][(accounts[i].Index-1)%len(cfg.Endpoints[machine])]
	}
	if !cfg.SkipP2PPreflight {
		if err = preflightP2P(ctx, cfg, manifest, accounts, apiKey); err != nil {
			return err
		}
	}
	coord := newCoordinatorClient(coordinatorURL)
	placeholderRunID := "pending"
	run := &agentRun{cfg: cfg, machine: machine, clients: make(map[int]*WSClient), ack: NewDurationHistogram()}
	run.recorder = newDeliveryRecorder(placeholderRunID, machine, len(accounts))
	stopMetrics, err := startAgentMetrics(metricsListen, run)
	if err != nil {
		return err
	}
	defer stopMetrics()
	for {
		err = run.connect(ctx, accounts, manifest.GroupTopic, apiKey, password)
		if err == nil {
			break
		}
		if ctx.Err() != nil {
			run.close()
			return ctx.Err()
		}
		log.Printf("agent=%s initial connection set incomplete; retrying: %v", machine, err)
		if err = waitContext(ctx, 2*time.Second); err != nil {
			run.close()
			return err
		}
	}
	defer run.close()
	maintenanceCtx, stopMaintenance := context.WithCancel(ctx)
	maintenanceDone := make(chan struct{})
	go func() {
		defer close(maintenanceDone)
		run.maintainConnections(maintenanceCtx, accounts, manifest.GroupTopic, apiKey, password)
	}()
	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))
	}
	// Connection churn while waiting for the other machine is outside the load
	// measurement window. From this point onward every disconnect is a failure.
	run.recorder.resetDisconnects()
	startedAt := time.Now().UTC()
	log.Printf("agent=%s scenario=%s started", machine, start.Scenario)
	if err = run.execute(ctx, start.StartAt, manifest.GroupTopic); err != nil {
		return err
	}
	loadEndedAt := time.Now().UTC()
	local := run.acceptedCopy()
	if err = coord.submitLedger(ctx, machine, local); err != nil {
		return err
	}
	ledger, err := coord.ledger(ctx)
	if err != nil {
		return err
	}
	drain, err := run.drainDuration()
	if err != nil {
		return err
	}
	if err = waitContext(ctx, drain); err != nil {
		return err
	}
	observedLocal := run.recorder.unackedObservedMessages(ledger, start.Scenario)
	if err = coord.submitObserved(ctx, machine, observedLocal); err != nil {
		return err
	}
	observedLedger, err := coord.observed(ctx)
	if err != nil {
		return err
	}
	verificationLedger := mergeMessageLedgers(ledger, observedLedger)
	base := AgentSummary{
		RunID: start.RunID, Scenario: start.Scenario, Machine: machine, StartedAt: startedAt,
		LoadEndedAt: loadEndedAt, FinishedAt: time.Now().UTC(), ExpectedConnections: cfg.UsersPerMachine,
		ReadyConnections: run.clientCount(),
		PublishAttempted: run.attempted.Load(), PublishAccepted: int64(len(local)),
		PublishRejected: run.explicitRejected(), PublishOutcomeUnknown: run.outcomeUnknown(),
		PublishErrors: run.publishErrorCounts(),
		AckLatency:    run.ack.Snapshot(), AcceptedLocal: local,
	}
	run.mu.Lock()
	base.Errors = append(base.Errors, run.errors...)
	run.mu.Unlock()
	summary := run.recorder.summarize(base, ledger)
	historySettle, _ := configDuration(cfg.HistorySettle)
	log.Printf("agent=%s waiting %s before history verification", machine, historySettle)
	if err = waitContext(ctx, historySettle); err != nil {
		return err
	}
	history := run.reconcileHistory(ctx, verificationLedger, ledger)
	if err = coord.submitHistory(ctx, machine, history.foundIDs); err != nil {
		return err
	}
	storedMessageIDs, err := coord.history(ctx)
	if err != nil {
		return err
	}
	summary.HistoryChecked = int64(history.checked)
	summary.HistoryFound = int64(history.found)
	summary.HistoryUnackedChecked = int64(history.unackedChecked)
	summary.HistoryUnackedFound = int64(history.unackedFound)
	persistenceMissing := history.missing
	summary.PersistenceLost = int64(len(persistenceMissing))
	summary.RealtimeLostStored = run.recorder.countRealtimeLostStored(ledger, storedMessageIDs)
	if len(persistenceMissing) > 10000 {
		summary.PersistenceMissing = append([]string(nil), persistenceMissing[:10000]...)
		summary.Errors = append(summary.Errors, "persistence missing details truncated at 10000")
	} else {
		summary.PersistenceMissing = persistenceMissing
	}
	summary.FinishedAt = time.Now().UTC()
	if err = writeAgentSummary(output, summary); err != nil {
		return err
	}
	if err = coord.submitSummary(ctx, summary); err != nil {
		return err
	}
	log.Printf("agent=%s accepted=%d expected_peer=%d received_peer=%d expected_group=%d received_group=%d lost=%d duplicates=%d",
		machine, summary.PublishAccepted, summary.ExpectedPeer, summary.ReceivedPeer,
		summary.ExpectedGroup, summary.ReceivedGroup, summary.RealtimeLost, summary.Duplicates)
	return nil
}

type historyQuery struct {
	client   *WSClient
	topic    string
	since    int64
	before   int64
	limit    int
	expected []string
}

type historyResult struct {
	checked        int
	found          int
	unackedChecked int
	unackedFound   int
	missing        []string
	foundIDs       []string
}

func (r *agentRun) reconcileHistory(ctx context.Context, ledger, ackLedger []AcceptedMessage) historyResult {
	expected := make(map[string]AcceptedMessage)
	acked := make(map[string]bool, len(ackLedger))
	for _, message := range ackLedger {
		acked[message.ClientMid] = true
	}
	byTopic := make(map[string][]AcceptedMessage)
	for _, message := range ledger {
		if message.RunID != r.runID {
			continue
		}
		if message.Kind == "p2p" && message.PeerMachine == r.machine {
			expected[message.ClientMid] = message
			byTopic[message.Topic] = append(byTopic[message.Topic], message)
		} else if message.Kind == "group" && r.machine == "a" {
			expected[message.ClientMid] = message
			byTopic[message.Topic] = append(byTopic[message.Topic], message)
		}
	}
	if len(expected) == 0 {
		return historyResult{}
	}
	r.recorder.beginHistory()
	queries := make([]historyQuery, 0, len(byTopic))
	for topic, messages := range byTopic {
		sort.Slice(messages, func(i, j int) bool { return messages[i].Seq < messages[j].Seq })
		client := r.client(messages[0].PeerIndex)
		if messages[0].Kind == "group" {
			client = r.client(1)
		}
		if client == nil {
			r.addError(fmt.Sprintf("history topic=%s has no verifier client", topicHash(topic)))
			continue
		}
		pageSize := r.historyPageSize()
		for start := 0; start < len(messages); start += pageSize {
			end := start + pageSize
			if end > len(messages) {
				end = len(messages)
			}
			queries = append(queries, newHistoryQuery(client, topic, messages[start:end], pageSize))
		}
	}
	historyTimeout, _ := configDuration(r.cfg.HistoryTimeout)
	historyCtx, cancel := context.WithTimeout(ctx, historyTimeout)
	defer cancel()
	r.issueHistoryQueries(historyCtx, queries)
	_ = waitContext(historyCtx, 2*time.Second)
	for {
		missing := missingHistory(r.recorder, expected)
		if len(missing) == 0 || historyCtx.Err() != nil {
			break
		}
		retryQueries := r.historyQueriesForMessages(missing)
		r.issueHistoryQueries(historyCtx, retryQueries)
		_ = waitContext(historyCtx, 5*time.Second)
	}
	missing := missingHistory(r.recorder, expected)
	missingIDs := make([]string, 0, len(missing))
	for _, message := range missing {
		missingIDs = append(missingIDs, message.ClientMid)
	}
	sort.Strings(missingIDs)
	result := historyResult{checked: len(expected), found: len(expected) - len(missing), missing: missingIDs,
		foundIDs: make([]string, 0, len(expected)-len(missing))}
	for mid := range expected {
		if r.recorder.historyFound(mid) {
			result.foundIDs = append(result.foundIDs, mid)
		}
		if acked[mid] {
			continue
		}
		result.unackedChecked++
		if r.recorder.historyFound(mid) {
			result.unackedFound++
		}
	}
	sort.Strings(result.foundIDs)
	return result
}

func (r *agentRun) historyQueriesForMessages(messages []AcceptedMessage) []historyQuery {
	byTopic := make(map[string][]AcceptedMessage)
	for _, message := range messages {
		byTopic[message.Topic] = append(byTopic[message.Topic], message)
	}
	queries := make([]historyQuery, 0, len(byTopic))
	for topic, topicMessages := range byTopic {
		sort.Slice(topicMessages, func(i, j int) bool { return topicMessages[i].Seq < topicMessages[j].Seq })
		client := r.client(topicMessages[0].PeerIndex)
		if topicMessages[0].Kind == "group" {
			client = r.client(1)
		}
		if client == nil {
			continue
		}
		pageSize := r.historyPageSize()
		for start := 0; start < len(topicMessages); start += pageSize {
			end := start + pageSize
			if end > len(topicMessages) {
				end = len(topicMessages)
			}
			queries = append(queries, newHistoryQuery(client, topic, topicMessages[start:end], pageSize))
		}
	}
	return queries
}

func (r *agentRun) issueHistoryQueries(ctx context.Context, queries []historyQuery) {
	byTopic := make(map[string][]historyQuery)
	for _, query := range queries {
		byTopic[query.topic] = append(byTopic[query.topic], query)
	}
	workers := r.cfg.HistoryWorkers
	if workers > len(byTopic) {
		workers = len(byTopic)
	}
	if workers == 0 {
		return
	}
	jobs := make(chan []historyQuery)
	interval := 500 * time.Millisecond
	if r.cfg != nil && r.cfg.HistoryInterval != "" {
		interval = mustDuration(r.cfg.HistoryInterval)
	}
	pageWait := 3 * time.Second
	if r.cfg != nil && r.cfg.HistoryPageWait != "" {
		pageWait = mustDuration(r.cfg.HistoryPageWait)
	}
	var wg sync.WaitGroup
	for i := 0; i < workers; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			for topicQueries := range jobs {
				for index, query := range topicQueries {
					if index > 0 {
						if err := waitContext(ctx, interval); err != nil {
							return
						}
					}
					if err := query.client.GetData(ctx, query.topic, query.since, query.before, query.limit); err != nil && ctx.Err() == nil {
						r.addError(fmt.Sprintf("history get topic=%s range=%d-%d: %v", topicHash(query.topic), query.since, query.before, err))
						continue
					}
					r.waitHistoryPage(ctx, query.expected, pageWait)
				}
			}
		}()
	}
	topics := make([]string, 0, len(byTopic))
	for topic := range byTopic {
		topics = append(topics, topic)
	}
	sort.Strings(topics)
	for _, topic := range topics {
		select {
		case jobs <- byTopic[topic]:
		case <-ctx.Done():
			close(jobs)
			wg.Wait()
			return
		}
	}
	close(jobs)
	wg.Wait()
}

func newHistoryQuery(client *WSClient, topic string, messages []AcceptedMessage, limit int) historyQuery {
	expected := make([]string, 0, len(messages))
	for _, message := range messages {
		expected = append(expected, message.ClientMid)
	}
	since := messages[0].Seq
	before := messages[len(messages)-1].Seq + 1
	// P2P topics interleave both senders. The requested range must leave room
	// for non-target messages between the expected sequence numbers.
	if span := before - since; span > int64(limit) {
		limit = int(span)
	}
	if limit > 200 {
		limit = 200
	}
	return historyQuery{client: client, topic: topic, since: since,
		before: before, limit: limit, expected: expected}
}

func (r *agentRun) waitHistoryPage(ctx context.Context, expected []string, timeout time.Duration) {
	deadline := time.NewTimer(timeout)
	defer deadline.Stop()
	ticker := time.NewTicker(25 * time.Millisecond)
	defer ticker.Stop()
	for {
		if r.recorder.historyFoundAll(expected) {
			return
		}
		select {
		case <-ctx.Done():
			return
		case <-deadline.C:
			return
		case <-ticker.C:
		}
	}
}

func (r *agentRun) historyPageSize() int {
	if r.cfg != nil && r.cfg.HistoryPageSize > 0 {
		return r.cfg.HistoryPageSize
	}
	return 10
}

func missingHistory(recorder *deliveryRecorder, expected map[string]AcceptedMessage) []AcceptedMessage {
	missing := make([]AcceptedMessage, 0)
	for mid, message := range expected {
		if !recorder.historyFound(mid) {
			missing = append(missing, message)
		}
	}
	return missing
}

func mergeMessageLedgers(accepted, observed []AcceptedMessage) []AcceptedMessage {
	byID := make(map[string]AcceptedMessage, len(accepted)+len(observed))
	for _, message := range observed {
		if message.ClientMid == "" {
			continue
		}
		current, exists := byID[message.ClientMid]
		currentIsPeer := current.PeerMachine != "" && current.PeerMachine != current.SenderMachine
		candidateIsPeer := message.PeerMachine != "" && message.PeerMachine != message.SenderMachine
		if !exists || (!currentIsPeer && candidateIsPeer) {
			byID[message.ClientMid] = message
		}
	}
	// ACK-derived records carry authoritative sender and timing fields.
	for _, message := range accepted {
		if message.ClientMid != "" {
			byID[message.ClientMid] = message
		}
	}
	merged := make([]AcceptedMessage, 0, len(byID))
	for _, message := range byID {
		merged = append(merged, message)
	}
	sort.Slice(merged, func(i, j int) bool { return merged[i].ClientMid < merged[j].ClientMid })
	return merged
}

func (r *agentRun) addError(message string) {
	r.mu.Lock()
	defer r.mu.Unlock()
	if len(r.errors) < 1000 {
		r.errors = append(r.errors, message)
	}
}

func (r *agentRun) connect(ctx context.Context, accounts []AccountRecord, groupTopic, apiKey, password string) error {
	ramp, err := configDuration(r.cfg.ConnectRamp)
	if err != nil {
		return err
	}
	interval := ramp / time.Duration(len(accounts))
	var wg sync.WaitGroup
	errCh := make(chan error, len(accounts))
	for _, account := range accounts {
		account := account
		if client := r.client(account.Index); client != nil {
			if client.alive() {
				continue
			}
			client.Close()
		}
		wg.Add(1)
		go func() {
			defer wg.Done()
			client, connectErr := r.connectAccountWithRetry(ctx, account, groupTopic, apiKey, password)
			if connectErr != nil {
				errCh <- connectErr
				return
			}
			r.setClient(account.Index, client)
			r.ready.Add(1)
		}()
		if interval > 0 {
			if err = waitContext(ctx, interval); err != nil {
				break
			}
		}
	}
	wg.Wait()
	close(errCh)
	var failures []string
	for connectErr := range errCh {
		if len(failures) < 20 {
			failures = append(failures, connectErr.Error())
		}
	}
	if r.aliveConnections() != len(accounts) {
		return fmt.Errorf("ready connections=%d expected=%d failures=%v", r.aliveConnections(), len(accounts), failures)
	}
	return nil
}

func (r *agentRun) maintainConnections(ctx context.Context, accounts []AccountRecord, groupTopic, apiKey, password string) {
	byIndex := make(map[int]AccountRecord, len(accounts))
	for _, account := range accounts {
		byIndex[account.Index] = account
	}
	ticker := time.NewTicker(time.Second)
	defer ticker.Stop()
	for {
		if ctx.Err() != nil {
			return
		}
		var wg sync.WaitGroup
		for index, account := range byIndex {
			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)
		}
		wg.Wait()
		select {
		case <-ctx.Done():
			return
		case <-ticker.C:
		}
	}
}

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")
	}
	switch r.scenario {
	case "p2p":
		return r.scheduleP2P(ctx, start, r.cfg.UsersPerMachine, r.cfg.P2PMessages, mustDuration(r.cfg.P2PInterval))
	case "group":
		return r.scheduleGroup(ctx, start, groupTopic, 1, scaledUsers(r.cfg.UsersPerMachine, 1, 2), 2)
	case "mixed":
		p2pUsers := scaledUsers(r.cfg.UsersPerMachine, 3, 10)
		groupSenders := scaledUsers(r.cfg.UsersPerMachine, 1, 4)
		if p2pUsers+groupSenders > r.cfg.UsersPerMachine {
			groupSenders = r.cfg.UsersPerMachine - p2pUsers
		}
		if groupSenders <= 0 {
			return fmt.Errorf("users_per_machine=%d is too small for mixed scenario", r.cfg.UsersPerMachine)
		}
		var wg sync.WaitGroup
		errCh := make(chan error, 2)
		wg.Add(2)
		go func() { defer wg.Done(); errCh <- r.scheduleP2P(ctx, start, p2pUsers, 5, 20*time.Second) }()
		go func() {
			defer wg.Done()
			errCh <- r.scheduleGroup(ctx, start, groupTopic, p2pUsers+1, p2pUsers+groupSenders, 2)
		}()
		wg.Wait()
		close(errCh)
		for err := range errCh {
			if err != nil {
				return err
			}
		}
		return nil
	default:
		return fmt.Errorf("unknown scenario %s", r.scenario)
	}
}

func scaledUsers(total, numerator, denominator int) int {
	count := total * numerator / denominator
	if count < 1 {
		return 1
	}
	if count > total {
		return total
	}
	return count
}

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 {
			return fmt.Errorf("missing client %d", index)
		}
		wg.Add(1)
		go func(index int, client *WSClient) {
			defer wg.Done()
			jitter := time.Duration((index*37)%200) * time.Millisecond
			for messageIndex := 1; messageIndex <= messages; messageIndex++ {
				target := start.Add(time.Duration(messageIndex-1)*interval + jitter)
				if waitUntil(ctx, target) != nil {
					return
				}
				mid := fmt.Sprintf("%s-p-%s-%04d-%02d", r.runID, r.machine, index, messageIndex)
				accepted, err := r.publish(ctx, client, client.account.P2PTopic, mid, messageIndex)
				if err == nil {
					accepted.PeerIndex = index
					if r.machine == "a" {
						accepted.PeerMachine = "b"
					} else {
						accepted.PeerMachine = "a"
					}
					r.addAccepted(accepted)
				}
			}
		}(index, client)
	}
	wg.Wait()
	return ctx.Err()
}

func (r *agentRun) scheduleGroup(ctx context.Context, start time.Time, topic string, first, last, rounds int) error {
	if first < 1 || last > r.clientCount() || first > last {
		return fmt.Errorf("invalid group publisher range %d-%d", first, last)
	}
	count := last - first + 1
	roundDuration := time.Duration(count) * 200 * time.Millisecond
	offset := time.Duration(0)
	if r.machine == "b" {
		offset = 100 * time.Millisecond
	}
	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) {
			defer wg.Done()
			for round := 1; round <= rounds; round++ {
				target := start.Add(time.Duration(round-1)*roundDuration + time.Duration(position)*200*time.Millisecond + offset)
				if waitUntil(ctx, target) != nil {
					return
				}
				mid := fmt.Sprintf("%s-g-%s-%04d-%02d", r.runID, r.machine, index, round)
				accepted, err := r.publish(ctx, client, topic, mid, round)
				if err == nil {
					r.addAccepted(accepted)
				}
			}
		}(index, position, client)
	}
	wg.Wait()
	return ctx.Err()
}

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)
	if err != nil {
		r.classifyPublishError(err)
		r.mu.Lock()
		if len(r.errors) < 1000 {
			r.errors = append(r.errors, fmt.Sprintf("publish %s: %v", mid, err))
		}
		r.mu.Unlock()
		return accepted, err
	}
	r.ack.Observe(accepted.AckAt.Sub(accepted.SentAt))
	return accepted, nil
}

func (r *agentRun) classifyPublishError(err error) {
	var ctrlErr *ctrlError
	switch {
	case errors.As(err, &ctrlErr) && ctrlErr.Code == 403:
		r.publish403.Add(1)
	case errors.As(err, &ctrlErr) && ctrlErr.Code == 502:
		r.publish502.Add(1)
	case errors.As(err, &ctrlErr) && ctrlErr.Code == 503:
		r.publish503.Add(1)
	case strings.Contains(err.Error(), "ctrl timeout") || errors.Is(err, context.DeadlineExceeded):
		r.publishTimeout.Add(1)
	case strings.Contains(err.Error(), "websocket") || strings.Contains(err.Error(), "broken pipe") ||
		strings.Contains(err.Error(), "connection reset") || strings.Contains(err.Error(), "closed network connection"):
		r.publishSocket.Add(1)
	default:
		r.publishOther.Add(1)
	}
}

func (r *agentRun) publishErrorCounts() map[string]int64 {
	return map[string]int64{
		"403": r.publish403.Load(), "502": r.publish502.Load(), "503": r.publish503.Load(),
		"ctrl_timeout": r.publishTimeout.Load(), "socket": r.publishSocket.Load(), "other": r.publishOther.Load(),
	}
}

func (r *agentRun) explicitRejected() int64 {
	return r.publish403.Load() + r.publish502.Load() + r.publish503.Load() + r.publishOther.Load()
}

func (r *agentRun) outcomeUnknown() int64 {
	return r.publishTimeout.Load() + r.publishSocket.Load()
}

func (r *agentRun) addAccepted(message AcceptedMessage) {
	r.mu.Lock()
	r.accepted = append(r.accepted, message)
	r.mu.Unlock()
}

func (r *agentRun) acceptedCopy() []AcceptedMessage {
	r.mu.Lock()
	defer r.mu.Unlock()
	out := append([]AcceptedMessage(nil), r.accepted...)
	sort.Slice(out, func(i, j int) bool { return out[i].ClientMid < out[j].ClientMid })
	return out
}

func (r *agentRun) aliveConnections() int {
	alive := 0
	for _, client := range r.clientsSnapshot() {
		if client.alive() {
			alive++
		}
	}
	return alive
}

func (r *agentRun) client(index int) *WSClient {
	r.clientsMu.RLock()
	defer r.clientsMu.RUnlock()
	return r.clients[index]
}

func (r *agentRun) setClient(index int, client *WSClient) {
	r.clientsMu.Lock()
	r.clients[index] = client
	r.clientsMu.Unlock()
}

func (r *agentRun) clientCount() int {
	r.clientsMu.RLock()
	defer r.clientsMu.RUnlock()
	return len(r.clients)
}

func (r *agentRun) clientsSnapshot() []*WSClient {
	r.clientsMu.RLock()
	defer r.clientsMu.RUnlock()
	clients := make([]*WSClient, 0, len(r.clients))
	for _, client := range r.clients {
		clients = append(clients, client)
	}
	return clients
}

func (r *agentRun) drainDuration() (time.Duration, error) {
	if r.scenario == "p2p" {
		return configDuration(r.cfg.DrainP2P)
	}
	return configDuration(r.cfg.DrainGroup)
}

func (r *agentRun) close() {
	clients := r.clientsSnapshot()
	sort.Slice(clients, func(i, j int) bool { return clients[i].account.Index < clients[j].account.Index })
	interval, initialDelay := closeSchedule(r.cfg, r.machine, len(clients))
	if initialDelay > 0 {
		time.Sleep(initialDelay)
	}
	for index, client := range clients {
		client.Close()
		if interval > 0 && index < len(clients)-1 {
			time.Sleep(interval)
		}
	}
}

func closeSchedule(cfg *Config, machine string, clients int) (interval, initialDelay time.Duration) {
	if cfg == nil || cfg.CloseRamp == "" || clients <= 1 {
		return 0, 0
	}
	// The two logical agents share one global ramp. Interleaving their slots avoids
	// doubling the disconnect rate when both agents begin cleanup at the same time.
	interval = 2 * mustDuration(cfg.CloseRamp) / time.Duration(clients-1)
	if machine == "b" {
		initialDelay = interval / 2
	}
	return interval, initialDelay
}

func waitUntil(ctx context.Context, target time.Time) error {
	delay := time.Until(target)
	if delay <= 0 {
		return nil
	}
	return waitContext(ctx, delay)
}

func mustDuration(value string) time.Duration {
	d, err := time.ParseDuration(value)
	if err != nil || d <= 0 {
		return time.Second
	}
	return d
}

func writeAgentSummary(path string, summary AgentSummary) error {
	raw, err := json.MarshalIndent(summary, "", "  ")
	if err != nil {
		return err
	}
	if dir := filepath.Dir(path); dir != "." {
		if err = os.MkdirAll(dir, 0o750); err != nil {
			return err
		}
	}
	return os.WriteFile(path, append(raw, '\n'), 0o600)
}