im_server/loadtest/im-pressure/recorder.go

package main

import (
	"encoding/json"
	"fmt"
	"math/bits"
	"strings"
	"sync"
	"sync/atomic"
	"time"
)

type observation struct {
	topic string
	seq   int64
	bits  []uint64
}

type deliveryRecorder struct {
	machine     string
	localUsers  int
	latency     *DurationHistogram
	disconnects atomic.Int64
	received    atomic.Int64
	runID       atomic.Value
	historyMode atomic.Bool

	mu         sync.Mutex
	observed   map[string]*observation
	lastSeq    []map[string]int64
	duplicates int64
	outOfOrder int64
	malformed  int64
	history    map[string]bool
}

func newDeliveryRecorder(runID, machine string, localUsers int) *deliveryRecorder {
	lastSeq := make([]map[string]int64, localUsers)
	for i := range lastSeq {
		lastSeq[i] = make(map[string]int64)
	}
	recorder := &deliveryRecorder{
		machine: machine, localUsers: localUsers, latency: NewDurationHistogram(),
		observed: make(map[string]*observation), lastSeq: lastSeq, history: make(map[string]bool),
	}
	recorder.runID.Store(runID)
	return recorder
}

func (r *deliveryRecorder) setRunID(runID string) { r.runID.Store(runID) }

func (r *deliveryRecorder) onData(account AccountRecord, data dataPacket) {
	runID, _ := r.runID.Load().(string)
	mid, _ := data.Head["x-im-client-mid"].(string)
	if mid == "" || !strings.HasPrefix(mid, runID+"-") {
		return
	}
	var content struct {
		RunID  string `json:"run_id"`
		SentAt int64  `json:"sent_at_unix_nano"`
	}
	if json.Unmarshal(data.Content, &content) != nil || content.RunID != runID {
		r.mu.Lock()
		r.malformed++
		r.mu.Unlock()
		return
	}
	if r.historyMode.Load() {
		r.mu.Lock()
		r.history[mid] = true
		r.mu.Unlock()
		return
	}
	index := account.Index - 1
	if index < 0 || index >= r.localUsers {
		return
	}
	r.mu.Lock()
	defer r.mu.Unlock()
	obs := r.observed[mid]
	if obs == nil {
		obs = &observation{topic: data.Topic, seq: data.Seq, bits: make([]uint64, (r.localUsers+63)/64)}
		r.observed[mid] = obs
	}
	word, bit := index/64, uint(index%64)
	mask := uint64(1) << bit
	if obs.bits[word]&mask != 0 {
		r.duplicates++
		return
	}
	obs.bits[word] |= mask
	r.received.Add(1)
	if previous := r.lastSeq[index][data.Topic]; previous > 0 && data.Seq < previous {
		r.outOfOrder++
	}
	if data.Seq > r.lastSeq[index][data.Topic] {
		r.lastSeq[index][data.Topic] = data.Seq
	}
	if content.SentAt > 0 {
		r.latency.Observe(time.Since(time.Unix(0, content.SentAt)))
	}
}

func (r *deliveryRecorder) metricCounts() (received, duplicates, outOfOrder, disconnects int64) {
	r.mu.Lock()
	duplicates = r.duplicates
	outOfOrder = r.outOfOrder
	r.mu.Unlock()
	return r.received.Load(), duplicates, outOfOrder, r.disconnects.Load()
}

func (r *deliveryRecorder) beginHistory() { r.historyMode.Store(true) }

func (r *deliveryRecorder) historyFound(mid string) bool {
	r.mu.Lock()
	defer r.mu.Unlock()
	return r.history[mid]
}

func (r *deliveryRecorder) onDisconnect(_ AccountRecord, _ error) {
	r.disconnects.Add(1)
}

func (r *deliveryRecorder) summarize(base AgentSummary, ledger []AcceptedMessage) AgentSummary {
	base.DeliveryLatency = r.latency.Snapshot()
	base.Disconnects = r.disconnects.Load()
	runID, _ := r.runID.Load().(string)
	accepted := make(map[string]AcceptedMessage, len(ledger))
	const missingLimit = 100000

	r.mu.Lock()
	defer r.mu.Unlock()
	base.Duplicates = r.duplicates
	base.OutOfOrder = r.outOfOrder
	if r.malformed > 0 {
		base.Errors = append(base.Errors, fmt.Sprintf("malformed run messages=%d", r.malformed))
	}
	for _, message := range ledger {
		if message.RunID != runID {
			continue
		}
		accepted[message.ClientMid] = message
		obs := r.observed[message.ClientMid]
		switch message.Kind {
		case "group":
			base.ExpectedGroup += int64(r.localUsers)
			for index := 1; index <= r.localUsers; index++ {
				if observationHas(obs, index) {
					base.ReceivedGroup++
				} else if len(base.Missing) < missingLimit {
					base.Missing = append(base.Missing, missingFor(message, index))
				}
			}
		case "p2p":
			if message.PeerMachine == r.machine && message.PeerIndex > 0 {
				base.ExpectedPeer++
				if observationHas(obs, message.PeerIndex) {
					base.ReceivedPeer++
				} else if len(base.Missing) < missingLimit {
					base.Missing = append(base.Missing, missingFor(message, message.PeerIndex))
				}
			}
			if message.SenderMachine == r.machine && message.SenderIndex > 0 {
				base.ExpectedEcho++
				if observationHas(obs, message.SenderIndex) {
					base.ReceivedEcho++
				}
			}
		}
	}
	for mid, obs := range r.observed {
		if _, ok := accepted[mid]; !ok {
			base.UnexpectedMessages += bitsetCount(obs.bits)
		}
	}
	base.RealtimeLost = (base.ExpectedPeer - base.ReceivedPeer) + (base.ExpectedGroup - base.ReceivedGroup)
	if len(base.Missing) == missingLimit {
		base.Errors = append(base.Errors, "missing delivery details truncated at 100000")
	}
	return base
}

func (r *deliveryRecorder) countRealtimeLostStored(ledger []AcceptedMessage) int64 {
	runID, _ := r.runID.Load().(string)
	r.mu.Lock()
	defer r.mu.Unlock()
	var count int64
	for _, message := range ledger {
		if message.RunID != runID || !r.history[message.ClientMid] {
			continue
		}
		obs := r.observed[message.ClientMid]
		switch message.Kind {
		case "group":
			for index := 1; index <= r.localUsers; index++ {
				if !observationHas(obs, index) {
					count++
				}
			}
		case "p2p":
			if message.PeerMachine == r.machine && message.PeerIndex > 0 && !observationHas(obs, message.PeerIndex) {
				count++
			}
		}
	}
	return count
}

func observationHas(obs *observation, receiverIndex int) bool {
	if obs == nil || receiverIndex <= 0 {
		return false
	}
	index := receiverIndex - 1
	word, bit := index/64, uint(index%64)
	return word < len(obs.bits) && obs.bits[word]&(uint64(1)<<bit) != 0
}

func bitsetCount(words []uint64) int64 {
	var count int64
	for _, word := range words {
		count += int64(bits.OnesCount64(word))
	}
	return count
}

func missingFor(message AcceptedMessage, receiver int) MissingDelivery {
	return MissingDelivery{ClientMid: message.ClientMid, Kind: message.Kind, TopicHash: topicHash(message.Topic),
		Seq: message.Seq, ReceiverIndex: receiver}
}