fix: align group management APIs
This commit is contained in:
richjohnrambo
@@ -0,0 +1,377 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"context"
|
||||
"crypto/tls"
|
||||
"encoding/base64"
|
||||
"encoding/csv"
|
||||
"encoding/json"
|
||||
"flag"
|
||||
"fmt"
|
||||
"log"
|
||||
"math"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"os"
|
||||
"sort"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
"github.com/gorilla/websocket"
|
||||
)
|
||||
|
||||
type account struct {
|
||||
Username string
|
||||
Password string
|
||||
Token string
|
||||
}
|
||||
|
||||
type client struct {
|
||||
id int
|
||||
account account
|
||||
wsURL string
|
||||
conn *websocket.Conn
|
||||
sendMu sync.Mutex
|
||||
stats *probeStats
|
||||
}
|
||||
|
||||
type probeStats struct {
|
||||
expected atomic.Int64
|
||||
received atomic.Int64
|
||||
acks atomic.Int64
|
||||
errors atomic.Int64
|
||||
|
||||
mu sync.Mutex
|
||||
latencies []time.Duration
|
||||
seen map[string]struct{}
|
||||
}
|
||||
|
||||
func main() {
|
||||
var wsList string
|
||||
var apiKey string
|
||||
var accountsPath string
|
||||
var topic string
|
||||
var subscribers int
|
||||
var publishers int
|
||||
var rate int
|
||||
var duration time.Duration
|
||||
var ramp time.Duration
|
||||
var insecure bool
|
||||
|
||||
flag.StringVar(&wsList, "ws", "ws://localhost:6060/v0/channels", "Comma-separated websocket endpoints, one or more pods/LBs.")
|
||||
flag.StringVar(&apiKey, "apikey", "", "Tinode API key appended to websocket URL.")
|
||||
flag.StringVar(&accountsPath, "accounts", "loadtest/users.csv", "CSV: username,password[,token].")
|
||||
flag.StringVar(&topic, "topic", "", "Group topic to subscribe and publish to, e.g. grpABC.")
|
||||
flag.IntVar(&subscribers, "subscribers", 100, "Number of websocket subscribers.")
|
||||
flag.IntVar(&publishers, "publishers", 1, "Number of publishers selected from subscribers.")
|
||||
flag.IntVar(&rate, "rate", 10, "Total publish rate, messages per second.")
|
||||
flag.DurationVar(&duration, "duration", time.Minute, "Publish duration.")
|
||||
flag.DurationVar(&ramp, "ramp", 10*time.Second, "Subscriber connection ramp duration.")
|
||||
flag.BoolVar(&insecure, "insecure", false, "Skip TLS verification for wss endpoints.")
|
||||
flag.Parse()
|
||||
|
||||
if topic == "" || subscribers <= 0 || publishers <= 0 || rate <= 0 || duration <= 0 {
|
||||
log.Fatal("topic, subscribers, publishers, rate, and duration are required and must be positive")
|
||||
}
|
||||
endpoints := parseEndpoints(wsList, apiKey)
|
||||
accounts, err := readAccounts(accountsPath)
|
||||
if err != nil {
|
||||
log.Fatal(err)
|
||||
}
|
||||
if len(accounts) < subscribers {
|
||||
log.Fatalf("need at least %d accounts, got %d", subscribers, len(accounts))
|
||||
}
|
||||
if publishers > subscribers {
|
||||
publishers = subscribers
|
||||
}
|
||||
|
||||
stats := &probeStats{seen: make(map[string]struct{})}
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
defer cancel()
|
||||
|
||||
clients := make([]*client, subscribers)
|
||||
var wg sync.WaitGroup
|
||||
for i := 0; i < subscribers; i++ {
|
||||
c := &client{id: i, account: accounts[i], wsURL: endpoints[i%len(endpoints)], stats: stats}
|
||||
clients[i] = c
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
if err := c.connectAndSubscribe(ctx, topic, insecure); err != nil {
|
||||
stats.errors.Add(1)
|
||||
log.Printf("client %d connect failed: %v", c.id, err)
|
||||
}
|
||||
}()
|
||||
if ramp > 0 {
|
||||
time.Sleep(ramp / time.Duration(subscribers))
|
||||
}
|
||||
}
|
||||
wg.Wait()
|
||||
active := 0
|
||||
for _, c := range clients {
|
||||
if c.conn != nil {
|
||||
active++
|
||||
}
|
||||
}
|
||||
if active == 0 {
|
||||
log.Fatal("no active websocket clients")
|
||||
}
|
||||
log.Printf("subscribed clients=%d publishers=%d endpoints=%d topic=%s", active, publishers, len(endpoints), topic)
|
||||
|
||||
publishDone := make(chan struct{})
|
||||
go publishLoop(ctx, clients[:publishers], topic, rate, duration, stats, publishDone)
|
||||
<-publishDone
|
||||
log.Printf("publish finished; waiting 10s for fanout drain")
|
||||
time.Sleep(10 * time.Second)
|
||||
cancel()
|
||||
for _, c := range clients {
|
||||
if c.conn != nil {
|
||||
_ = c.conn.Close()
|
||||
}
|
||||
}
|
||||
printSummary(stats, active)
|
||||
}
|
||||
|
||||
func (c *client) connectAndSubscribe(ctx context.Context, topic string, insecure bool) error {
|
||||
dialer := websocket.DefaultDialer
|
||||
if insecure {
|
||||
dialer.TLSClientConfig = &tls.Config{InsecureSkipVerify: true}
|
||||
}
|
||||
conn, _, err := dialer.DialContext(ctx, c.wsURL, http.Header{})
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
c.conn = conn
|
||||
if err = c.sendAndWaitCtrl("hi", map[string]any{"hi": map[string]any{"id": c.msgID("hi"), "ver": "0.22.8", "ua": "realtime-fanout-probe/1.0"}}); err != nil {
|
||||
return err
|
||||
}
|
||||
if c.account.Token != "" {
|
||||
if err = c.sendAndWaitCtrl("login", map[string]any{"login": map[string]any{"id": c.msgID("login"), "scheme": "token", "secret": c.account.Token}}); err != nil {
|
||||
return err
|
||||
}
|
||||
} else {
|
||||
secret := base64.StdEncoding.EncodeToString([]byte(c.account.Username + ":" + c.account.Password))
|
||||
if err = c.sendAndWaitCtrl("login", map[string]any{"login": map[string]any{"id": c.msgID("login"), "scheme": "basic", "secret": secret}}); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
if err = c.sendAndWaitCtrl("sub", map[string]any{"sub": map[string]any{"id": c.msgID("sub"), "topic": topic, "get": map[string]any{"what": "desc"}}}); err != nil {
|
||||
return err
|
||||
}
|
||||
go c.readLoop(ctx)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (c *client) sendAndWaitCtrl(name string, msg map[string]any) error {
|
||||
c.sendMu.Lock()
|
||||
if err := c.conn.WriteJSON(msg); err != nil {
|
||||
c.sendMu.Unlock()
|
||||
return err
|
||||
}
|
||||
c.sendMu.Unlock()
|
||||
deadline := time.Now().Add(15 * time.Second)
|
||||
for time.Now().Before(deadline) {
|
||||
var raw map[string]json.RawMessage
|
||||
if err := c.conn.ReadJSON(&raw); err != nil {
|
||||
return err
|
||||
}
|
||||
if ctrlRaw, ok := raw["ctrl"]; ok {
|
||||
var ctrl struct {
|
||||
Code int `json:"code"`
|
||||
Text string `json:"text"`
|
||||
}
|
||||
_ = json.Unmarshal(ctrlRaw, &ctrl)
|
||||
if ctrl.Code >= 200 && ctrl.Code < 300 {
|
||||
return nil
|
||||
}
|
||||
return fmt.Errorf("%s ctrl code=%d text=%s", name, ctrl.Code, ctrl.Text)
|
||||
}
|
||||
}
|
||||
return fmt.Errorf("%s ctrl timeout", name)
|
||||
}
|
||||
|
||||
func (c *client) readLoop(ctx context.Context) {
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
default:
|
||||
}
|
||||
var raw map[string]json.RawMessage
|
||||
if err := c.conn.ReadJSON(&raw); err != nil {
|
||||
return
|
||||
}
|
||||
if _, ok := raw["ctrl"]; ok {
|
||||
c.stats.acks.Add(1)
|
||||
continue
|
||||
}
|
||||
dataRaw, ok := raw["data"]
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
var data struct {
|
||||
Content json.RawMessage `json:"content"`
|
||||
}
|
||||
if err := json.Unmarshal(dataRaw, &data); err != nil || len(data.Content) == 0 {
|
||||
continue
|
||||
}
|
||||
var content struct {
|
||||
ProbeID string `json:"probe_id"`
|
||||
SentAt int64 `json:"sent_at_unix_nano"`
|
||||
}
|
||||
if err := json.Unmarshal(data.Content, &content); err != nil || content.ProbeID == "" || content.SentAt <= 0 {
|
||||
continue
|
||||
}
|
||||
c.stats.record(content.ProbeID, time.Since(time.Unix(0, content.SentAt)))
|
||||
}
|
||||
}
|
||||
|
||||
func publishLoop(ctx context.Context, publishers []*client, topic string, rate int, duration time.Duration, stats *probeStats, done chan<- struct{}) {
|
||||
defer close(done)
|
||||
interval := time.Second / time.Duration(rate)
|
||||
if interval <= 0 {
|
||||
interval = time.Nanosecond
|
||||
}
|
||||
ticker := time.NewTicker(interval)
|
||||
defer ticker.Stop()
|
||||
deadline := time.Now().Add(duration)
|
||||
seq := int64(0)
|
||||
for time.Now().Before(deadline) {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case <-ticker.C:
|
||||
}
|
||||
n := atomic.AddInt64(&seq, 1)
|
||||
publisher := publishers[int(n)%len(publishers)]
|
||||
probeID := fmt.Sprintf("probe-%d-%d", time.Now().UnixNano(), n)
|
||||
msg := map[string]any{"pub": map[string]any{
|
||||
"id": publisher.msgID("pub-" + strconv.FormatInt(n, 10)),
|
||||
"topic": topic,
|
||||
"head": map[string]any{"mime": "application/json"},
|
||||
"content": map[string]any{
|
||||
"probe_id": probeID,
|
||||
"sent_at_unix_nano": time.Now().UnixNano(),
|
||||
"seq": n,
|
||||
},
|
||||
}}
|
||||
publisher.sendMu.Lock()
|
||||
err := publisher.conn.WriteJSON(msg)
|
||||
publisher.sendMu.Unlock()
|
||||
if err != nil {
|
||||
stats.errors.Add(1)
|
||||
log.Printf("publish failed client=%d seq=%d: %v", publisher.id, n, err)
|
||||
continue
|
||||
}
|
||||
stats.expected.Add(1)
|
||||
}
|
||||
}
|
||||
|
||||
func (s *probeStats) record(probeID string, latency time.Duration) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
s.received.Add(1)
|
||||
s.seen[probeID] = struct{}{}
|
||||
s.latencies = append(s.latencies, latency)
|
||||
}
|
||||
|
||||
func printSummary(stats *probeStats, activeSubscribers int) {
|
||||
stats.mu.Lock()
|
||||
latencies := append([]time.Duration(nil), stats.latencies...)
|
||||
unique := len(stats.seen)
|
||||
stats.mu.Unlock()
|
||||
sort.Slice(latencies, func(i, j int) bool { return latencies[i] < latencies[j] })
|
||||
published := stats.expected.Load()
|
||||
expectedDeliveries := published * int64(activeSubscribers)
|
||||
received := stats.received.Load()
|
||||
lost := expectedDeliveries - received
|
||||
lossRate := 0.0
|
||||
if expectedDeliveries > 0 {
|
||||
lossRate = float64(lost) / float64(expectedDeliveries) * 100
|
||||
}
|
||||
fmt.Println("==== realtime fanout probe summary ====")
|
||||
fmt.Printf("published=%d active_subscribers=%d expected_deliveries=%d received=%d lost=%d loss_rate=%.4f%% unique_messages_seen=%d errors=%d ctrl_acks=%d\n",
|
||||
published, activeSubscribers, expectedDeliveries, received, lost, lossRate, unique, stats.errors.Load(), stats.acks.Load())
|
||||
fmt.Printf("latency p50=%s p95=%s p99=%s max=%s samples=%d\n",
|
||||
percentile(latencies, 50), percentile(latencies, 95), percentile(latencies, 99), maxDuration(latencies), len(latencies))
|
||||
}
|
||||
|
||||
func percentile(values []time.Duration, p int) time.Duration {
|
||||
if len(values) == 0 {
|
||||
return 0
|
||||
}
|
||||
idx := int(math.Ceil(float64(p)/100*float64(len(values)))) - 1
|
||||
if idx < 0 {
|
||||
idx = 0
|
||||
}
|
||||
if idx >= len(values) {
|
||||
idx = len(values) - 1
|
||||
}
|
||||
return values[idx]
|
||||
}
|
||||
|
||||
func maxDuration(values []time.Duration) time.Duration {
|
||||
if len(values) == 0 {
|
||||
return 0
|
||||
}
|
||||
return values[len(values)-1]
|
||||
}
|
||||
|
||||
func (c *client) msgID(suffix string) string {
|
||||
return fmt.Sprintf("rtp-%d-%s", c.id, suffix)
|
||||
}
|
||||
|
||||
func parseEndpoints(raw string, apiKey string) []string {
|
||||
parts := strings.Split(raw, ",")
|
||||
out := make([]string, 0, len(parts))
|
||||
for _, part := range parts {
|
||||
endpoint := strings.TrimSpace(part)
|
||||
if endpoint == "" {
|
||||
continue
|
||||
}
|
||||
u, err := url.Parse(endpoint)
|
||||
if err != nil {
|
||||
log.Fatalf("bad ws endpoint %q: %v", endpoint, err)
|
||||
}
|
||||
q := u.Query()
|
||||
if apiKey != "" && q.Get("apikey") == "" {
|
||||
q.Set("apikey", apiKey)
|
||||
}
|
||||
u.RawQuery = q.Encode()
|
||||
out = append(out, u.String())
|
||||
}
|
||||
if len(out) == 0 {
|
||||
log.Fatal("no websocket endpoints")
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func readAccounts(path string) ([]account, error) {
|
||||
file, err := os.Open(path)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer file.Close()
|
||||
reader := csv.NewReader(bufio.NewReader(file))
|
||||
reader.FieldsPerRecord = -1
|
||||
rows, err := reader.ReadAll()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
accounts := make([]account, 0, len(rows))
|
||||
for _, row := range rows {
|
||||
if len(row) < 2 || strings.HasPrefix(strings.TrimSpace(row[0]), "#") {
|
||||
continue
|
||||
}
|
||||
acc := account{Username: strings.TrimSpace(row[0]), Password: strings.TrimSpace(row[1])}
|
||||
if len(row) > 2 {
|
||||
acc.Token = strings.TrimSpace(row[2])
|
||||
}
|
||||
accounts = append(accounts, acc)
|
||||
}
|
||||
return accounts, nil
|
||||
}
|
||||
Reference in New Issue
Block a user