fix: harden high-fanout delivery under load

This commit is contained in:
richjohnrambo
2026-09-30 23:04:54 +08:00
parent 2b631ab176
commit 19584d0939
14 changed files with 491 additions and 90 deletions
+25 -2
View File
@@ -26,6 +26,7 @@ type agentRun struct {
clients map[int]*WSClient
clientsMu sync.RWMutex
recorder *deliveryRecorder
callbackQueue *callbackQueueMetrics
ack *DurationHistogram
attempted atomic.Int64
ready atomic.Int64
@@ -125,7 +126,8 @@ func runAgent(ctx context.Context, cfg *Config, manifest *Manifest, machine, coo
}
coord := newCoordinatorClient(coordinatorURL)
placeholderRunID := "pending"
run := &agentRun{cfg: cfg, machine: machine, clients: make(map[int]*WSClient), ack: NewDurationHistogram()}
run := &agentRun{cfg: cfg, machine: machine, clients: make(map[int]*WSClient), ack: NewDurationHistogram(),
callbackQueue: &callbackQueueMetrics{}}
run.recorder = newDeliveryRecorder(placeholderRunID, machine, len(accounts))
stopMetrics, err := startAgentMetrics(metricsListen, run)
if err != nil {
@@ -206,6 +208,9 @@ func runAgent(ctx context.Context, cfg *Config, manifest *Manifest, machine, coo
if err = waitContext(ctx, drain); err != nil {
return err
}
if err = waitForCallbackQueue(ctx, run.callbackQueue, 30*time.Second); err != nil {
return err
}
observedLocal := run.recorder.unackedObservedMessages(ledger, start.Scenario)
if err = coord.submitObserved(ctx, machine, observedLocal); err != nil {
return err
@@ -223,6 +228,7 @@ func runAgent(ctx context.Context, cfg *Config, manifest *Manifest, machine, coo
PublishRejected: run.explicitRejected(), PublishOutcomeUnknown: run.outcomeUnknown(),
PublishErrors: run.publishErrorCounts(),
AckLatency: run.ack.Snapshot(), AcceptedLocal: local,
CallbackQueue: run.callbackQueue.snapshot(),
}
run.mu.Lock()
base.Errors = append(base.Errors, run.errors...)
@@ -270,6 +276,23 @@ func runAgent(ctx context.Context, cfg *Config, manifest *Manifest, machine, coo
return nil
}
func waitForCallbackQueue(ctx context.Context, queue *callbackQueueMetrics, timeout time.Duration) error {
deadline := time.NewTimer(timeout)
defer deadline.Stop()
ticker := time.NewTicker(10 * time.Millisecond)
defer ticker.Stop()
for !queue.isDrained() {
select {
case <-ctx.Done():
return ctx.Err()
case <-deadline.C:
return fmt.Errorf("data callback queue did not drain within %s: current=%d processing=%d",
timeout, queue.current.Load(), queue.processing.Load())
case <-ticker.C:
}
}
return nil
}
type historyQuery struct {
client *WSClient
topic string
@@ -667,7 +690,7 @@ func (r *agentRun) maintainConnections(ctx context.Context, accounts []AccountRe
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)
client, err := connectWSClient(ctx, r.cfg, account, apiKey, password, r.recorder.onData, r.recorder.onDisconnect, r.callbackQueue)
if err == nil {
if err = client.Subscribe(ctx, account.P2PTopic); err == nil {
err = client.Subscribe(ctx, groupTopic)
+21
View File
@@ -63,6 +63,27 @@ func startAgentMetrics(address string, run *agentRun) (func(), error) {
_, _, _, disconnects := run.recorder.metricCounts()
return float64(disconnects)
}),
prometheus.NewGaugeFunc(prometheus.GaugeOpts{Name: "im_load_agent_callback_queue_current", Help: "Data callbacks currently waiting to start.", ConstLabels: labels}, func() float64 {
return float64(run.callbackQueue.current.Load())
}),
prometheus.NewGaugeFunc(prometheus.GaugeOpts{Name: "im_load_agent_callback_queue_max", Help: "Peak data callbacks waiting to start.", ConstLabels: labels}, func() float64 {
return float64(run.callbackQueue.maximum.Load())
}),
prometheus.NewGaugeFunc(prometheus.GaugeOpts{Name: "im_load_agent_callback_processing", Help: "Data callbacks currently processing.", ConstLabels: labels}, func() float64 {
return float64(run.callbackQueue.processing.Load())
}),
prometheus.NewCounterFunc(prometheus.CounterOpts{Name: "im_load_agent_callback_queue_enqueued_total", Help: "Data callbacks enqueued.", ConstLabels: labels}, func() float64 {
return float64(run.callbackQueue.enqueued.Load())
}),
prometheus.NewCounterFunc(prometheus.CounterOpts{Name: "im_load_agent_callback_queue_completed_total", Help: "Data callbacks which completed processing.", ConstLabels: labels}, func() float64 {
return float64(run.callbackQueue.completed.Load())
}),
prometheus.NewCounterFunc(prometheus.CounterOpts{Name: "im_load_agent_callback_queue_wait_seconds_total", Help: "Cumulative time spent waiting to start data callbacks.", ConstLabels: labels}, func() float64 {
return float64(run.callbackQueue.waitNanos.Load()) / float64(time.Second)
}),
prometheus.NewGaugeFunc(prometheus.GaugeOpts{Name: "im_load_agent_callback_queue_wait_seconds_max", Help: "Maximum time spent waiting to start a data callback.", ConstLabels: labels}, func() float64 {
return float64(run.callbackQueue.maxWait.Load()) / float64(time.Second)
}),
)
mux := http.NewServeMux()
mux.Handle("/metrics", promhttp.HandlerFor(registry, promhttp.HandlerOpts{}))
+70
View File
@@ -0,0 +1,70 @@
package main
import (
"sync/atomic"
"time"
)
type callbackQueueMetrics struct {
current atomic.Int64
processing atomic.Int64
maximum atomic.Int64
enqueued atomic.Int64
completed atomic.Int64
waitNanos atomic.Int64
maxWait atomic.Int64
}
func (m *callbackQueueMetrics) enqueue() {
if m == nil {
return
}
current := m.current.Add(1)
m.enqueued.Add(1)
updateAtomicMax(&m.maximum, current)
}
func (m *callbackQueueMetrics) startCallback(queuedAt time.Time) {
if m == nil {
return
}
wait := time.Since(queuedAt).Nanoseconds()
if wait < 0 {
wait = 0
}
m.current.Add(-1)
m.processing.Add(1)
m.waitNanos.Add(wait)
updateAtomicMax(&m.maxWait, wait)
}
func (m *callbackQueueMetrics) finishCallback() {
if m == nil {
return
}
m.processing.Add(-1)
m.completed.Add(1)
}
func (m *callbackQueueMetrics) snapshot() CallbackQueueSnapshot {
if m == nil {
return CallbackQueueSnapshot{}
}
return CallbackQueueSnapshot{
Current: m.current.Load(), Processing: m.processing.Load(), Maximum: m.maximum.Load(), Enqueued: m.enqueued.Load(),
Completed: m.completed.Load(), WaitSecondsTotal: float64(m.waitNanos.Load()) / float64(time.Second),
WaitSecondsMax: float64(m.maxWait.Load()) / float64(time.Second),
}
}
func (m *callbackQueueMetrics) isDrained() bool {
return m == nil || (m.current.Load() == 0 && m.processing.Load() == 0)
}
func updateAtomicMax(target *atomic.Int64, value int64) {
for current := target.Load(); value > current; current = target.Load() {
if target.CompareAndSwap(current, value) {
return
}
}
}
@@ -0,0 +1,47 @@
package main
import (
"context"
"testing"
"time"
)
func TestCallbackQueueMetricsTrackWaitingAndProcessing(t *testing.T) {
metrics := &callbackQueueMetrics{}
metrics.enqueue()
metrics.enqueue()
metrics.startCallback(time.Now().Add(-10 * time.Millisecond))
snapshot := metrics.snapshot()
if snapshot.Current != 1 || snapshot.Processing != 1 || snapshot.Maximum != 2 || snapshot.Enqueued != 2 {
t.Fatalf("unexpected active queue snapshot: %+v", snapshot)
}
if snapshot.WaitSecondsTotal <= 0 || snapshot.WaitSecondsMax <= 0 || metrics.isDrained() {
t.Fatalf("wait or drain state not tracked: %+v", snapshot)
}
metrics.finishCallback()
metrics.startCallback(time.Now())
metrics.finishCallback()
if snapshot = metrics.snapshot(); snapshot.Current != 0 || snapshot.Processing != 0 || snapshot.Completed != 2 || !metrics.isDrained() {
t.Fatalf("unexpected drained queue snapshot: %+v", snapshot)
}
}
func TestWaitForCallbackQueueWaitsForProcessingCallback(t *testing.T) {
metrics := &callbackQueueMetrics{}
metrics.enqueue()
metrics.startCallback(time.Now())
go func() {
time.Sleep(20 * time.Millisecond)
metrics.finishCallback()
}()
started := time.Now()
if err := waitForCallbackQueue(context.Background(), metrics, time.Second); err != nil {
t.Fatal(err)
}
if time.Since(started) < 15*time.Millisecond {
t.Fatal("wait returned before the callback finished")
}
}
+53 -42
View File
@@ -385,48 +385,59 @@ type AcceptedMessage struct {
}
type AgentSummary struct {
RunID string `json:"run_id"`
Scenario string `json:"scenario"`
Machine string `json:"machine"`
StartedAt time.Time `json:"started_at"`
LoadEndedAt time.Time `json:"load_ended_at"`
FinishedAt time.Time `json:"finished_at"`
ExpectedConnections int `json:"expected_connections"`
ReadyConnections int `json:"ready_connections"`
PublishAttempted int64 `json:"publish_attempted"`
PublishAccepted int64 `json:"publish_accepted"`
PublishRejected int64 `json:"publish_rejected"`
PublishOutcomeUnknown int64 `json:"publish_outcome_unknown"`
PublishErrors map[string]int64 `json:"publish_errors,omitempty"`
ExpectedPeer int64 `json:"expected_peer"`
ReceivedPeer int64 `json:"received_peer"`
ExpectedEcho int64 `json:"expected_echo"`
ReceivedEcho int64 `json:"received_echo"`
ExpectedGroup int64 `json:"expected_group"`
ReceivedGroup int64 `json:"received_group"`
RealtimeLost int64 `json:"realtime_lost"`
RealtimeLostStored int64 `json:"realtime_lost_but_persisted"`
Duplicates int64 `json:"duplicates"`
OutOfOrder int64 `json:"out_of_order"`
OrderingViolations []OrderingViolation `json:"ordering_violations,omitempty"`
Disconnects int64 `json:"disconnects"`
DisconnectReasons map[string]int64 `json:"disconnect_reasons,omitempty"`
DisconnectSamples []DisconnectSample `json:"disconnect_samples,omitempty"`
UnexpectedMessages int64 `json:"unexpected_messages"`
HistoryChecked int64 `json:"history_checked"`
HistoryFound int64 `json:"history_found"`
HistoryUnackedChecked int64 `json:"history_unacked_observed_checked"`
HistoryUnackedFound int64 `json:"history_unacked_observed_found"`
PersistenceLost int64 `json:"persistence_lost"`
AckLatency HistogramSnapshot `json:"ack_latency"`
DeliveryLatency HistogramSnapshot `json:"delivery_latency"`
DeliveryLatencyLocal HistogramSnapshot `json:"delivery_latency_local"`
DeliveryLatencyCross HistogramSnapshot `json:"delivery_latency_cross_machine"`
CrossClockSkewSamples int64 `json:"cross_machine_negative_latency_samples"`
Errors []string `json:"errors,omitempty"`
Missing []MissingDelivery `json:"missing,omitempty"`
PersistenceMissing []string `json:"persistence_missing,omitempty"`
AcceptedLocal []AcceptedMessage `json:"accepted_local,omitempty"`
RunID string `json:"run_id"`
Scenario string `json:"scenario"`
Machine string `json:"machine"`
StartedAt time.Time `json:"started_at"`
LoadEndedAt time.Time `json:"load_ended_at"`
FinishedAt time.Time `json:"finished_at"`
ExpectedConnections int `json:"expected_connections"`
ReadyConnections int `json:"ready_connections"`
PublishAttempted int64 `json:"publish_attempted"`
PublishAccepted int64 `json:"publish_accepted"`
PublishRejected int64 `json:"publish_rejected"`
PublishOutcomeUnknown int64 `json:"publish_outcome_unknown"`
PublishErrors map[string]int64 `json:"publish_errors,omitempty"`
ExpectedPeer int64 `json:"expected_peer"`
ReceivedPeer int64 `json:"received_peer"`
ExpectedEcho int64 `json:"expected_echo"`
ReceivedEcho int64 `json:"received_echo"`
ExpectedGroup int64 `json:"expected_group"`
ReceivedGroup int64 `json:"received_group"`
RealtimeLost int64 `json:"realtime_lost"`
RealtimeLostStored int64 `json:"realtime_lost_but_persisted"`
Duplicates int64 `json:"duplicates"`
OutOfOrder int64 `json:"out_of_order"`
OrderingViolations []OrderingViolation `json:"ordering_violations,omitempty"`
Disconnects int64 `json:"disconnects"`
DisconnectReasons map[string]int64 `json:"disconnect_reasons,omitempty"`
DisconnectSamples []DisconnectSample `json:"disconnect_samples,omitempty"`
UnexpectedMessages int64 `json:"unexpected_messages"`
HistoryChecked int64 `json:"history_checked"`
HistoryFound int64 `json:"history_found"`
HistoryUnackedChecked int64 `json:"history_unacked_observed_checked"`
HistoryUnackedFound int64 `json:"history_unacked_observed_found"`
PersistenceLost int64 `json:"persistence_lost"`
AckLatency HistogramSnapshot `json:"ack_latency"`
DeliveryLatency HistogramSnapshot `json:"delivery_latency"`
DeliveryLatencyLocal HistogramSnapshot `json:"delivery_latency_local"`
DeliveryLatencyCross HistogramSnapshot `json:"delivery_latency_cross_machine"`
CrossClockSkewSamples int64 `json:"cross_machine_negative_latency_samples"`
CallbackQueue CallbackQueueSnapshot `json:"callback_queue"`
Errors []string `json:"errors,omitempty"`
Missing []MissingDelivery `json:"missing,omitempty"`
PersistenceMissing []string `json:"persistence_missing,omitempty"`
AcceptedLocal []AcceptedMessage `json:"accepted_local,omitempty"`
}
type CallbackQueueSnapshot struct {
Current int64 `json:"current"`
Processing int64 `json:"processing"`
Maximum int64 `json:"maximum"`
Enqueued int64 `json:"enqueued"`
Completed int64 `json:"completed"`
WaitSecondsTotal float64 `json:"wait_seconds_total"`
WaitSecondsMax float64 `json:"wait_seconds_max"`
}
type DisconnectSample struct {
+65 -19
View File
@@ -17,6 +17,16 @@ import (
const observationShardCount = 64
type observation struct {
topic string
seq int64
senderMachine string
senderIndex int
scenario string
sentAt int64
bits []atomic.Uint64
}
type observationView struct {
topic string
seq int64
senderMachine string
@@ -121,23 +131,26 @@ func (r *deliveryRecorder) onData(account AccountRecord, data dataPacket) {
return
}
shard := r.shard(mid)
shard.mu.Lock()
shard.mu.RLock()
obs := shard.observed[mid]
shard.mu.RUnlock()
if obs == nil {
obs = &observation{topic: data.Topic, seq: data.Seq, senderMachine: content.SenderMachine,
senderIndex: content.SenderIndex, scenario: content.Scenario, sentAt: content.SentAt,
bits: make([]uint64, (r.localUsers+63)/64)}
shard.observed[mid] = obs
shard.mu.Lock()
obs = shard.observed[mid]
if obs == nil {
obs = &observation{topic: data.Topic, seq: data.Seq, senderMachine: content.SenderMachine,
senderIndex: content.SenderIndex, scenario: content.Scenario, sentAt: content.SentAt,
bits: make([]atomic.Uint64, (r.localUsers+63)/64)}
shard.observed[mid] = obs
}
shard.mu.Unlock()
}
word, bit := index/64, uint(index%64)
mask := uint64(1) << bit
if obs.bits[word]&mask != 0 {
shard.mu.Unlock()
if !markObservationBit(obs.bits, word, mask) {
r.duplicates.Add(1)
return
}
obs.bits[word] |= mask
shard.mu.Unlock()
r.received.Add(1)
sequence := &r.sequences[index]
sequence.mu.Lock()
@@ -235,7 +248,7 @@ func (r *deliveryRecorder) unackedObservedMessages(ledger []AcceptedMessage, sce
}
message := AcceptedMessage{RunID: runID, Scenario: observedScenario, Kind: kind,
ClientMid: mid, Topic: obs.topic, Seq: obs.seq, SenderMachine: obs.senderMachine,
SenderIndex: obs.senderIndex, PeerMachine: r.machine, PeerIndex: firstObservedIndex(obs.bits)}
SenderIndex: obs.senderIndex, PeerMachine: r.machine, PeerIndex: firstObservedAtomicIndex(obs.bits)}
if obs.sentAt > 0 {
message.SentAt = time.Unix(0, obs.sentAt).UTC()
}
@@ -256,6 +269,27 @@ func firstObservedIndex(words []uint64) int {
return 0
}
func firstObservedAtomicIndex(words []atomic.Uint64) int {
for wordIndex := range words {
if word := words[wordIndex].Load(); word != 0 {
return wordIndex*64 + bits.TrailingZeros64(word) + 1
}
}
return 0
}
func markObservationBit(words []atomic.Uint64, word int, mask uint64) bool {
for {
current := words[word].Load()
if current&mask != 0 {
return false
}
if words[word].CompareAndSwap(current, current|mask) {
return true
}
}
}
func (r *deliveryRecorder) onDisconnect(account AccountRecord, err error) {
category := disconnectCategory(err)
r.disconnectMu.Lock()
@@ -370,7 +404,7 @@ func (r *deliveryRecorder) summarize(base AgentSummary, ledger []AcceptedMessage
shard.mu.RLock()
for mid, obs := range shard.observed {
if _, ok := accepted[mid]; !ok {
base.UnexpectedMessages += bitsetCount(obs.bits)
base.UnexpectedMessages += atomicBitsetCount(obs.bits)
}
}
shard.mu.RUnlock()
@@ -419,17 +453,21 @@ func (r *deliveryRecorder) shard(mid string) *observationShard {
return &r.shards[hash%uint32(len(r.shards))]
}
func (r *deliveryRecorder) observationSnapshot(mid string) *observation {
func (r *deliveryRecorder) observationSnapshot(mid string) *observationView {
shard := r.shard(mid)
shard.mu.RLock()
obs := shard.observed[mid]
if obs != nil {
copyOfObservation := *obs
copyOfObservation.bits = append([]uint64(nil), obs.bits...)
obs = &copyOfObservation
}
shard.mu.RUnlock()
return obs
if obs == nil {
return nil
}
view := &observationView{topic: obs.topic, seq: obs.seq, senderMachine: obs.senderMachine,
senderIndex: obs.senderIndex, scenario: obs.scenario, sentAt: obs.sentAt,
bits: make([]uint64, len(obs.bits))}
for i := range obs.bits {
view.bits[i] = obs.bits[i].Load()
}
return view
}
func cloneCountMap(source map[string]int64) map[string]int64 {
@@ -443,7 +481,7 @@ func cloneCountMap(source map[string]int64) map[string]int64 {
return clone
}
func observationHas(obs *observation, receiverIndex int) bool {
func observationHas(obs *observationView, receiverIndex int) bool {
if obs == nil || receiverIndex <= 0 {
return false
}
@@ -452,6 +490,14 @@ func observationHas(obs *observation, receiverIndex int) bool {
return word < len(obs.bits) && obs.bits[word]&(uint64(1)<<bit) != 0
}
func atomicBitsetCount(words []atomic.Uint64) int64 {
var count int64
for i := range words {
count += int64(bits.OnesCount64(words[i].Load()))
}
return count
}
func bitsetCount(words []uint64) int64 {
var count int64
for _, word := range words {
+17 -6
View File
@@ -35,6 +35,11 @@ type dataPacket struct {
Content json.RawMessage `json:"content"`
}
type queuedDataPacket struct {
packet dataPacket
queuedAt time.Time
}
type ctrlError struct {
Code int
Text string
@@ -52,18 +57,19 @@ type WSClient struct {
commandTimeout time.Duration
onData func(AccountRecord, dataPacket)
onDisconnect func(AccountRecord, error)
callbackQueue *callbackQueueMetrics
sendMu sync.Mutex
mu sync.Mutex
seq uint64
waits map[string]chan ctrlPacket
done chan struct{}
data chan dataPacket
data chan queuedDataPacket
closed atomic.Bool
}
func connectWSClient(ctx context.Context, cfg *Config, account AccountRecord, apiKey, password string,
onData func(AccountRecord, dataPacket), onDisconnect func(AccountRecord, error)) (*WSClient, error) {
onData func(AccountRecord, dataPacket), onDisconnect func(AccountRecord, error), callbackQueue *callbackQueueMetrics) (*WSClient, error) {
timeout, err := configDuration(cfg.CommandTimeout)
if err != nil {
return nil, err
@@ -85,7 +91,8 @@ func connectWSClient(ctx context.Context, cfg *Config, account AccountRecord, ap
}
client := &WSClient{
account: account, conn: conn, commandTimeout: timeout, onData: onData, onDisconnect: onDisconnect,
waits: make(map[string]chan ctrlPacket), done: make(chan struct{}),
callbackQueue: callbackQueue,
waits: make(map[string]chan ctrlPacket), done: make(chan struct{}),
}
client.startDataProcessor()
go client.readLoop()
@@ -382,7 +389,9 @@ func (c *WSClient) readLoop() {
if c.data == nil {
c.onData(c.account, data)
} else {
c.data <- data
queued := queuedDataPacket{packet: data, queuedAt: time.Now()}
c.callbackQueue.enqueue()
c.data <- queued
}
}
}
@@ -393,10 +402,12 @@ func (c *WSClient) startDataProcessor() {
if c.onData == nil || c.data != nil {
return
}
c.data = make(chan dataPacket, dataCallbackQueueSize)
c.data = make(chan queuedDataPacket, dataCallbackQueueSize)
go func() {
for data := range c.data {
c.onData(c.account, data)
c.callbackQueue.startCallback(data.queuedAt)
c.onData(c.account, data.packet)
c.callbackQueue.finishCallback()
}
}()
}