//go:build integration package pgstreambroker import ( "context" "fmt" "sync" "sync/atomic" "testing" "time" "github.com/centrifugal/centrifuge" "github.com/stretchr/testify/require" ) // newBenchPostgresStreamBroker is a slim broker constructor for benchmarks. // It uses BinaryData=true (so we don't pay JSON validation cost) and a // large outbox batch size for throughput tests. func newBenchPostgresStreamBroker(b *testing.B) (*PostgresStreamBroker, *centrifuge.Node, *testBrokerEventHandler) { node, err := centrifuge.New(centrifuge.Config{}) require.NoError(b, err) connString := getPostgresConnString(b) e, err := NewPostgresStreamBroker(node, PostgresStreamBrokerConfig{ DSN: connString, BinaryData: true, CleanupInterval: 1 * time.Minute, // long — don't interfere with bench StreamRetention: 24 * time.Hour, PartitionLookaheadDays: 1, PartitionRetentionDays: 7, Outbox: OutboxConfig{ BatchSize: 1000, // PollInterval uses default (50ms) — matches map broker bench // config so numbers are comparable. Lower poll intervals steal // PG capacity from the publish goroutines. }, }) require.NoError(b, err) ctx := context.Background() hardResetTestSchema(b, e) require.NoError(b, e.EnsureSchema(ctx)) cleanupTestTables(ctx, e) handler := &testBrokerEventHandler{} require.NoError(b, e.RegisterBrokerEventHandler(handler)) require.NoError(b, node.Run()) b.Cleanup(func() { _ = e.Close(context.Background()) _ = node.Shutdown(context.Background()) }) return e, node, handler } // BenchmarkPostgresStreamBroker_Publish measures the core publish path with // parallel goroutines across many channels (same structure as the map broker // Publish benchmark). This is the steady-state number to report. func BenchmarkPostgresStreamBroker_Publish(b *testing.B) { e, _, _ := newBenchPostgresStreamBroker(b) data := []byte("hello world") b.ResetTimer() b.ReportAllocs() var counter int64 b.RunParallel(func(pb *testing.PB) { for pb.Next() { i := atomic.AddInt64(&counter, 1) ch := fmt.Sprintf("bench_publish_%d", i%1000) _, err := e.Publish(ch, data, centrifuge.PublishOptions{ HistoryTTL: 10 * time.Minute, HistorySize: 1000, }) if err != nil { b.Fatalf("publish: %v", err) } } }) } // BenchmarkPostgresStreamBroker_OutboxLatency measures publish→delivery // end-to-end latency: time from Publish() return to HandlePublication() call // in the outbox worker. Reports the average latency via b.ReportMetric. func BenchmarkPostgresStreamBroker_OutboxLatency(b *testing.B) { e, _, handler := newBenchPostgresStreamBroker(b) channel := "bench_latency" data := []byte("hello world") // Pre-allocate a slice indexed by offset (1-based). The publish loop // stores the publish time at position res.Offset, then unlocks. The // handler reads at pub.Offset. Mutex serializes the store/load against // the publish call so the handler always sees the time before the // outbox worker delivers the row. publishTimes := make([]time.Time, b.N+2) var mu sync.Mutex delivered := make(chan time.Duration, b.N) handler.HandlePublicationFunc = func(ch string, pub *centrifuge.Publication, sp centrifuge.StreamPosition, delta bool, prevPub *centrifuge.Publication) error { if ch != channel { return nil } mu.Lock() var t time.Time if int(pub.Offset) < len(publishTimes) { t = publishTimes[pub.Offset] } mu.Unlock() if !t.IsZero() { delivered <- time.Since(t) } return nil } b.ResetTimer() b.ReportAllocs() for i := 0; i < b.N; i++ { // Hold the lock around publish + array store so the handler sees // the publish time before processing the outbox row. mu.Lock() now := time.Now() res, err := e.Publish(channel, data, centrifuge.PublishOptions{ HistoryTTL: 10 * time.Minute, HistorySize: 10000, }) if err != nil { mu.Unlock() b.Fatalf("publish: %v", err) } if int(res.Offset) < len(publishTimes) { publishTimes[res.Offset] = now } mu.Unlock() } // Wait for all deliveries. deadline := time.NewTimer(30 * time.Second) defer deadline.Stop() received := 0 totalLatency := time.Duration(0) for received < b.N { select { case d := <-delivered: totalLatency += d received++ case <-deadline.C: b.Fatalf("only received %d/%d deliveries within 30s", received, b.N) } } b.StopTimer() if b.N > 0 { avg := totalLatency / time.Duration(b.N) b.ReportMetric(float64(avg.Microseconds()), "µs/delivery") } } // BenchmarkPostgresStreamBroker_RetentionDrop measures partition drop time. // Pre-populates b.N partitions each with rowsPerPartition rows, then times // only the DROP TABLE calls (the create + populate happens with the timer // stopped). func BenchmarkPostgresStreamBroker_RetentionDrop(b *testing.B) { e, _, _ := newBenchPostgresStreamBroker(b) ctx := context.Background() const rowsPerPartition = 10000 yesterday := time.Now().UTC().AddDate(0, 0, -1) partNames := make([]string, b.N) // Stop the timer for the setup phase. b.StopTimer() for i := 0; i < b.N; i++ { partName := fmt.Sprintf("%s_bench_%d_%d", e.names.stream, time.Now().UnixNano(), i) // Use date ranges many years in the past to avoid collisions with // the lookahead-managed partitions. yearsAgo := yesterday.AddDate(-100-i, 0, 0) yearsAgoNext := yearsAgo.AddDate(0, 0, 1) _, err := e.pool.Exec(ctx, fmt.Sprintf( `CREATE TABLE IF NOT EXISTS %s PARTITION OF %s FOR VALUES FROM ('%s') TO ('%s')`, partName, e.names.stream, yearsAgo.Format("2006-01-02 00:00:00+00"), yearsAgoNext.Format("2006-01-02 00:00:00+00"), )) require.NoError(b, err) partNames[i] = partName // Pre-populate the partition. for r := 0; r < rowsPerPartition; r++ { _, err := e.pool.Exec(ctx, fmt.Sprintf( `INSERT INTO %s (channel, channel_offset, kind, data, created_at, shard_id) VALUES ($1, $2, 0, '\x'::bytea, $3, 0)`, e.names.stream, ), fmt.Sprintf("bench_drop_%d_%d", i, r), int64(r), yearsAgo.Add(time.Duration(r)*time.Microsecond)) require.NoError(b, err) } } b.ResetTimer() b.StartTimer() for i := 0; i < b.N; i++ { _, err := e.pool.Exec(ctx, fmt.Sprintf("DROP TABLE %s", partNames[i])) require.NoError(b, err) } }