main.go

// bench is a network-path load generator for the broker, in the spirit of
// kafka-producer-perf-test. It drives the real gRPC client API (pkg/client),
// so results include serialization, network, and server overhead — unlike the
// in-process benchmarks under tests/benchmark.
//
// Usage:
//
//	go run ./cmd/bench -addr localhost:8080 -mode produce -workers 8 -batch 100 -msg-size 1024 -duration 30s
//	go run ./cmd/bench -addr localhost:8080 -mode consume -workers 8 -batch 500
package main

import (
	"context"
	"encoding/json"
	"flag"
	"fmt"
	"math/rand"
	"os"
	"sort"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"my-broker/pkg/client"
)

type config struct {
	addr       string
	topic      string
	partitions int
	workers    int
	msgSize    int
	batch      int
	duration   time.Duration
	warmup     time.Duration
	mode       string
	jsonOut    bool
}

type report struct {
	Mode          string  `json:"mode"`
	Addr          string  `json:"addr"`
	Topic         string  `json:"topic"`
	Partitions    int     `json:"partitions"`
	Workers       int     `json:"workers"`
	MsgSizeBytes  int     `json:"msg_size_bytes"`
	BatchSize     int     `json:"batch_size"`
	DurationSecs  float64 `json:"duration_secs"`
	Messages      int64   `json:"messages"`
	RPCs          int64   `json:"rpcs"`
	Errors        int64   `json:"errors"`
	MsgsPerSec    float64 `json:"msgs_per_sec"`
	MBPerSec      float64 `json:"mb_per_sec"`
	LatencyP50Ms  float64 `json:"latency_p50_ms"`
	LatencyP95Ms  float64 `json:"latency_p95_ms"`
	LatencyP99Ms  float64 `json:"latency_p99_ms"`
	LatencyMaxMs  float64 `json:"latency_max_ms"`
}

func main() {
	var cfg config
	flag.StringVar(&cfg.addr, "addr", "localhost:8080", "broker gRPC address")
	flag.StringVar(&cfg.topic, "topic", "bench", "topic to produce/consume")
	flag.IntVar(&cfg.partitions, "partitions", 8, "partition count (topic is created if missing)")
	flag.IntVar(&cfg.workers, "workers", 8, "concurrent workers, each with its own connection")
	flag.IntVar(&cfg.msgSize, "msg-size", 1024, "message size in bytes (produce mode)")
	flag.IntVar(&cfg.batch, "batch", 100, "records per RPC (1 = single-record produce)")
	flag.DurationVar(&cfg.duration, "duration", 30*time.Second, "measured duration (produce) / deadline (consume)")
	flag.DurationVar(&cfg.warmup, "warmup", 3*time.Second, "warmup period excluded from produce results")
	flag.StringVar(&cfg.mode, "mode", "produce", "produce or consume")
	flag.BoolVar(&cfg.jsonOut, "json", false, "emit results as JSON")
	flag.Parse()

	if cfg.batch < 1 || cfg.workers < 1 || cfg.partitions < 1 || cfg.msgSize < 1 {
		fmt.Fprintln(os.Stderr, "batch, workers, partitions, and msg-size must be >= 1")
		os.Exit(1)
	}

	ctx := context.Background()

	admin, err := client.New(cfg.addr)
	if err != nil {
		fatal("connect %s: %v", cfg.addr, err)
	}
	defer admin.Close()

	if err := admin.CreateTopic(ctx, cfg.topic, cfg.partitions); err != nil &&
		!strings.Contains(err.Error(), "exists") {
		fatal("create topic: %v", err)
	}

	var rep report
	switch cfg.mode {
	case "produce":
		rep = runProduce(ctx, cfg)
	case "consume":
		rep = runConsume(ctx, cfg)
	default:
		fatal("unknown mode %q (want produce or consume)", cfg.mode)
	}

	if cfg.jsonOut {
		enc := json.NewEncoder(os.Stdout)
		enc.SetIndent("", "  ")
		enc.Encode(rep)
		return
	}
	printReport(rep)
}

// worker holds per-goroutine counters, merged after the run.
type worker struct {
	messages  int64
	rpcs      int64
	errors    int64
	latencies []time.Duration
}

func runProduce(ctx context.Context, cfg config) report {
	payload := make([]byte, cfg.msgSize)
	rand.New(rand.NewSource(42)).Read(payload)
	batch := make([][]byte, cfg.batch)
	for i := range batch {
		batch[i] = payload
	}

	workers := make([]*worker, cfg.workers)
	clients := make([]*client.Client, cfg.workers)
	for i := range clients {
		c, err := client.New(cfg.addr)
		if err != nil {
			fatal("worker %d connect: %v", i, err)
		}
		defer c.Close()
		clients[i] = c
		workers[i] = &worker{}
	}

	var recording atomic.Bool
	stop := make(chan struct{})
	var wg sync.WaitGroup

	for i := 0; i < cfg.workers; i++ {
		wg.Add(1)
		go func(w *worker, c *client.Client, partition int) {
			defer wg.Done()
			for {
				select {
				case <-stop:
					return
				default:
				}
				start := time.Now()
				var err error
				if cfg.batch == 1 {
					_, err = c.ProduceWithPartition(ctx, cfg.topic, partition, nil, payload)
				} else {
					_, _, err = c.ProduceBatchToPartition(ctx, cfg.topic, partition, batch)
				}
				if !recording.Load() {
					continue
				}
				if err != nil {
					atomic.AddInt64(&w.errors, 1)
					time.Sleep(10 * time.Millisecond)
					continue
				}
				w.latencies = append(w.latencies, time.Since(start))
				w.rpcs++
				w.messages += int64(cfg.batch)
			}
		}(workers[i], clients[i], i%cfg.partitions)
	}

	time.Sleep(cfg.warmup)
	recording.Store(true)
	measureStart := time.Now()
	time.Sleep(cfg.duration)
	close(stop)
	wg.Wait()
	elapsed := time.Since(measureStart)

	return summarize("produce", cfg, workers, elapsed)
}

func runConsume(ctx context.Context, cfg config) report {
	workers := make([]*worker, cfg.workers)
	clients := make([]*client.Client, cfg.workers)
	for i := range clients {
		c, err := client.New(cfg.addr)
		if err != nil {
			fatal("worker %d connect: %v", i, err)
		}
		defer c.Close()
		clients[i] = c
		workers[i] = &worker{}
	}

	// Partitions are assigned round-robin; each worker drains its partitions
	// from offset 0 until it catches up or the deadline passes.
	deadline := time.Now().Add(cfg.duration)
	var wg sync.WaitGroup
	measureStart := time.Now()

	for i := 0; i < cfg.workers; i++ {
		wg.Add(1)
		go func(w *worker, c *client.Client, workerID int) {
			defer wg.Done()
			for p := workerID; p < cfg.partitions; p += cfg.workers {
				offset := uint64(0)
				for time.Now().Before(deadline) {
					start := time.Now()
					recs, next, err := c.ConsumeBatchFromPartition(ctx, cfg.topic, p, offset, cfg.batch)
					if err != nil {
						atomic.AddInt64(&w.errors, 1)
						time.Sleep(10 * time.Millisecond)
						continue
					}
					if len(recs) == 0 {
						break // caught up
					}
					w.latencies = append(w.latencies, time.Since(start))
					w.rpcs++
					w.messages += int64(len(recs))
					offset = next
				}
			}
		}(workers[i], clients[i], i)
	}
	wg.Wait()
	elapsed := time.Since(measureStart)

	return summarize("consume", cfg, workers, elapsed)
}

func summarize(mode string, cfg config, workers []*worker, elapsed time.Duration) report {
	var messages, rpcs, errors int64
	var all []time.Duration
	for _, w := range workers {
		messages += w.messages
		rpcs += w.rpcs
		errors += w.errors
		all = append(all, w.latencies...)
	}
	sort.Slice(all, func(i, j int) bool { return all[i] < all[j] })

	secs := elapsed.Seconds()
	return report{
		Mode:         mode,
		Addr:         cfg.addr,
		Topic:        cfg.topic,
		Partitions:   cfg.partitions,
		Workers:      cfg.workers,
		MsgSizeBytes: cfg.msgSize,
		BatchSize:    cfg.batch,
		DurationSecs: secs,
		Messages:     messages,
		RPCs:         rpcs,
		Errors:       errors,
		MsgsPerSec:   float64(messages) / secs,
		MBPerSec:     float64(messages) * float64(cfg.msgSize) / secs / (1 << 20),
		LatencyP50Ms: ms(pct(all, 0.50)),
		LatencyP95Ms: ms(pct(all, 0.95)),
		LatencyP99Ms: ms(pct(all, 0.99)),
		LatencyMaxMs: ms(pct(all, 1.0)),
	}
}

func pct(sorted []time.Duration, p float64) time.Duration {
	if len(sorted) == 0 {
		return 0
	}
	idx := int(p*float64(len(sorted))) - 1
	if idx < 0 {
		idx = 0
	}
	if idx >= len(sorted) {
		idx = len(sorted) - 1
	}
	return sorted[idx]
}

func ms(d time.Duration) float64 {
	return float64(d.Microseconds()) / 1000
}

func printReport(r report) {
	fmt.Printf("=== BENCH RESULTS (%s) ===\n", r.Mode)
	fmt.Printf("config      addr=%s topic=%s partitions=%d workers=%d batch=%d msg=%dB\n",
		r.Addr, r.Topic, r.Partitions, r.Workers, r.BatchSize, r.MsgSizeBytes)
	fmt.Printf("duration    %.1fs\n", r.DurationSecs)
	fmt.Printf("messages    %s\n", comma(r.Messages))
	fmt.Printf("rpcs        %s   errors %d\n", comma(r.RPCs), r.Errors)
	fmt.Printf("throughput  %s msg/s\n", comma(int64(r.MsgsPerSec)))
	fmt.Printf("bandwidth   %.1f MB/s\n", r.MBPerSec)
	fmt.Printf("rpc latency p50 %.2fms  p95 %.2fms  p99 %.2fms  max %.2fms\n",
		r.LatencyP50Ms, r.LatencyP95Ms, r.LatencyP99Ms, r.LatencyMaxMs)
}

func comma(n int64) string {
	s := fmt.Sprintf("%d", n)
	if len(s) <= 3 {
		return s
	}
	var b strings.Builder
	lead := len(s) % 3
	if lead > 0 {
		b.WriteString(s[:lead])
	}
	for i := lead; i < len(s); i += 3 {
		if b.Len() > 0 {
			b.WriteByte(',')
		}
		b.WriteString(s[i : i+3])
	}
	return b.String()
}

func fatal(format string, args ...any) {
	fmt.Fprintf(os.Stderr, format+"\n", args...)
	os.Exit(1)
}