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

fix: keep load agents connected before start

Browse Source
This commit is contained in:
richjohnrambo
2026-09-28 00:16:51 +08:00
parent 6b8ed4a2cc
commit 622fc50418
8 changed files with 352 additions and 58 deletions
Download Patch File Download Diff File Expand all files Collapse all files
loadtest/im-pressure/agent.go
+130 -23
View File
@@ -24,6 +24,7 @@ type agentRun struct {
runID string
scenario string
clients map[int]*WSClient
clientsMu sync.RWMutex
recorder *deliveryRecorder
ack *DurationHistogram
attempted atomic.Int64
@@ -132,27 +133,59 @@ func runAgent(ctx context.Context, cfg *Config, manifest *Manifest, machine, coo
return err
}
defer stopMetrics()
if err = run.connect(ctx, accounts, manifest.GroupTopic, apiKey, password); err != nil {
run.close()
return err
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()
start, err := coord.ready(ctx, AgentReady{Machine: machine, ReadyConnections: len(run.clients)})
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 alive := run.aliveConnections(); alive != len(run.clients) {
return fmt.Errorf("connections alive after coordinator wait=%d expected=%d", alive, len(run.clients))
}
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 {
@@ -177,7 +210,7 @@ func runAgent(ctx context.Context, cfg *Config, manifest *Manifest, machine, coo
base := AgentSummary{
RunID: start.RunID, Scenario: start.Scenario, Machine: machine, StartedAt: startedAt,
LoadEndedAt: loadEndedAt, FinishedAt: time.Now().UTC(), ExpectedConnections: cfg.UsersPerMachine,
ReadyConnections: len(run.clients),
ReadyConnections: run.clientCount(),
PublishAttempted: run.attempted.Load(), PublishAccepted: int64(len(local)), PublishRejected: run.rejected.Load(),
PublishErrors: run.publishErrorCounts(),
AckLatency: run.ack.Snapshot(), AcceptedLocal: local,
@@ -245,9 +278,9 @@ func (r *agentRun) reconcileHistory(ctx context.Context, ledger []AcceptedMessag
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.clients[messages[0].PeerIndex]
client := r.client(messages[0].PeerIndex)
if messages[0].Kind == "group" {
client = r.clients[1]
client = r.client(1)
}
if client == nil {
r.addError(fmt.Sprintf("history topic=%s has no verifier client", topicHash(topic)))
@@ -293,9 +326,9 @@ func (r *agentRun) historyQueriesForMessages(messages []AcceptedMessage) []histo
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.clients[topicMessages[0].PeerIndex]
client := r.client(topicMessages[0].PeerIndex)
if topicMessages[0].Kind == "group" {
client = r.clients[1]
client = r.client(1)
}
if client == nil {
continue
@@ -372,9 +405,14 @@ func (r *agentRun) connect(ctx context.Context, accounts []AccountRecord, groupT
interval := ramp / time.Duration(len(accounts))
var wg sync.WaitGroup
errCh := make(chan error, len(accounts))
var mu sync.Mutex
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()
@@ -383,9 +421,7 @@ func (r *agentRun) connect(ctx context.Context, accounts []AccountRecord, groupT
errCh <- connectErr
return
}
mu.Lock()
r.clients[account.Index] = client
mu.Unlock()
r.setClient(account.Index, client)
r.ready.Add(1)
}()
if interval > 0 {
@@ -402,12 +438,55 @@ func (r *agentRun) connect(ctx context.Context, accounts []AccountRecord, groupT
failures = append(failures, connectErr.Error())
}
}
if len(r.clients) != len(accounts) {
return fmt.Errorf("ready connections=%d expected=%d failures=%v", len(r.clients), len(accounts), failures)
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++ {
@@ -467,7 +546,7 @@ func (r *agentRun) execute(ctx context.Context, start time.Time, groupTopic stri
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.clients[index]
client := r.client(index)
if client == nil {
return fmt.Errorf("missing client %d", index)
}
@@ -499,7 +578,7 @@ func (r *agentRun) scheduleP2P(ctx context.Context, start time.Time, users, mess
}
func (r *agentRun) scheduleGroup(ctx context.Context, start time.Time, topic string, first, last, rounds int) error {
if first < 1 || last > len(r.clients) || first > last {
if first < 1 || last > r.clientCount() || first > last {
return fmt.Errorf("invalid group publisher range %d-%d", first, last)
}
count := last - first + 1
@@ -510,7 +589,7 @@ func (r *agentRun) scheduleGroup(ctx context.Context, start time.Time, topic str
}
var wg sync.WaitGroup
for index := first; index <= last; index++ {
client := r.clients[index]
client := r.client(index)
position := index - first
wg.Add(1)
go func(index, position int, client *WSClient) {
@@ -591,7 +670,7 @@ func (r *agentRun) acceptedCopy() []AcceptedMessage {
func (r *agentRun) aliveConnections() int {
alive := 0
for _, client := range r.clients {
for _, client := range r.clientsSnapshot() {
if client.alive() {
alive++
}
@@ -599,6 +678,34 @@ func (r *agentRun) aliveConnections() int {
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)
@@ -607,7 +714,7 @@ func (r *agentRun) drainDuration() (time.Duration, error) {
}
func (r *agentRun) close() {
for _, client := range r.clients {
for _, client := range r.clientsSnapshot() {
client.Close()
}
}