snapshot of https://github.com/centrifugal/centrifugo.git (master, 2026-09-16; blobs > 131072B omitted)

2026-09-16 01:17:21 +05:00
commit b352a9a30d
472 changed files with 108375 additions and 0 deletions
--- a/internal/pgmapbroker/pgmapbroker_bench_test.go
+++ b/internal/pgmapbroker/pgmapbroker_bench_test.go
@@ -0,0 +1,570 @@
+//go:build integration
+
+package pgmapbroker
+
+import (
+	"context"
+	"fmt"
+	"sync"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"github.com/centrifugal/centrifuge"
+)
+
+// setupBench creates a broker with Persistent mode for standard benchmarks.
+func setupBench(b *testing.B) (*PostgresMapBroker, func()) {
+	b.Helper()
+	connString := getPostgresConnString(b)
+	node, _ := centrifuge.New(centrifuge.Config{
+		Map: centrifuge.MapConfig{
+			GetMapChannelOptions: func(channel string) centrifuge.MapChannelOptions {
+				return centrifuge.MapChannelOptions{
+					Mode: centrifuge.MapModePersistent,
+				}
+			},
+		},
+	})
+	broker, err := NewPostgresMapBroker(node, PostgresMapBrokerConfig{
+		DSN:                    connString,
+		BinaryData:             true,
+		PartitionRetentionDays: 7,
+	})
+	if err != nil {
+		b.Fatal(err)
+	}
+	ctx := context.Background()
+	if err := broker.EnsureSchema(ctx); err != nil {
+		b.Fatal(err)
+	}
+	_ = broker.RegisterEventHandler(nil)
+	cleanupTestTables(ctx, broker)
+	return broker, func() {
+		_ = broker.Close(context.Background())
+		_ = node.Shutdown(context.Background())
+	}
+}
+
+// setupBenchOutbox creates a broker with outbox workers running.
+func setupBenchOutbox(b *testing.B, handler centrifuge.BrokerEventHandler) (*PostgresMapBroker, func()) {
+	b.Helper()
+	connString := getPostgresConnString(b)
+	node, _ := centrifuge.New(centrifuge.Config{
+		Map: centrifuge.MapConfig{
+			GetMapChannelOptions: func(channel string) centrifuge.MapChannelOptions {
+				return centrifuge.MapChannelOptions{
+					Mode: centrifuge.MapModePersistent,
+				}
+			},
+		},
+	})
+	if err := node.Run(); err != nil {
+		b.Fatal(err)
+	}
+	broker, err := NewPostgresMapBroker(node, PostgresMapBrokerConfig{
+		DSN:                    connString,
+		BinaryData:             true,
+		PoolSize:               32,
+		NumShards:              8,
+		PartitionRetentionDays: 7,
+		Outbox: OutboxConfig{
+			PollInterval: 5 * time.Millisecond,
+			BatchSize:    1000,
+		},
+	})
+	if err != nil {
+		b.Fatal(err)
+	}
+	ctx := context.Background()
+	if err := broker.EnsureSchema(ctx); err != nil {
+		b.Fatal(err)
+	}
+	cleanupTestTables(ctx, broker)
+	_ = broker.RegisterEventHandler(handler)
+	return broker, func() {
+		_ = broker.Close(context.Background())
+		_ = node.Shutdown(context.Background())
+	}
+}
+
+// ============================================================================
+// Publish benchmarks
+// ============================================================================
+
+// BenchmarkPostgresMapBroker_Publish measures the core map publish path:
+// shard lock → meta UPSERT → state UPSERT → stream INSERT → NOTIFY.
+// All publish benchmarks go through the same cf_map_publish SQL function.
+func BenchmarkPostgresMapBroker_Publish(b *testing.B) {
+	broker, cleanup := setupBench(b)
+	defer cleanup()
+	ctx := context.Background()
+
+	b.ReportAllocs()
+	b.ResetTimer()
+	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)
+			key := fmt.Sprintf("key%d", i%1000)
+			_, err := broker.Publish(ctx, ch, key, centrifuge.MapPublishOptions{
+				Data: []byte("x"),
+			})
+			if err != nil {
+				b.Fatal(err)
+			}
+		}
+	})
+}
+
+// BenchmarkPostgresMapBroker_PublishIdempotent measures publish with
+// idempotency key check (extra lookup + conditional insert).
+func BenchmarkPostgresMapBroker_PublishIdempotent(b *testing.B) {
+	broker, cleanup := setupBench(b)
+	defer cleanup()
+	ctx := context.Background()
+
+	b.ReportAllocs()
+	b.ResetTimer()
+	var counter int64
+	b.RunParallel(func(pb *testing.PB) {
+		for pb.Next() {
+			i := atomic.AddInt64(&counter, 1)
+			ch := fmt.Sprintf("bench_idempotent_%d", i%1000)
+			key := fmt.Sprintf("key_%d", i)
+			_, err := broker.Publish(ctx, ch, "", centrifuge.MapPublishOptions{
+				Data:                []byte(key),
+				IdempotencyKey:      key,
+				IdempotentResultTTL: 60 * time.Second,
+			})
+			if err != nil {
+				b.Fatal(err)
+			}
+		}
+	})
+}
+
+// BenchmarkPostgresMapBroker_PublishCAS measures Compare-And-Swap: read
+// current state + conditional publish. Includes both the read and write
+// round-trips.
+func BenchmarkPostgresMapBroker_PublishCAS(b *testing.B) {
+	broker, cleanup := setupBench(b)
+	defer cleanup()
+	ctx := context.Background()
+	channel := "bench_cas"
+
+	// Seed the key.
+	_, err := broker.Publish(ctx, channel, "counter", centrifuge.MapPublishOptions{
+		Data: []byte("0"),
+	})
+	if err != nil {
+		b.Fatal(err)
+	}
+
+	b.ReportAllocs()
+	b.ResetTimer()
+	b.RunParallel(func(pb *testing.PB) {
+		for pb.Next() {
+			stateRes, err := broker.ReadState(ctx, channel, centrifuge.MapReadStateOptions{
+				Key: "counter",
+			})
+			if err != nil || len(stateRes.Publications) == 0 {
+				continue
+			}
+			pos := centrifuge.StreamPosition{
+				Offset: stateRes.Publications[0].Offset,
+				Epoch:  stateRes.Position.Epoch,
+			}
+			_, _ = broker.Publish(ctx, channel, "counter", centrifuge.MapPublishOptions{
+				Data:             []byte("updated"),
+				ExpectedPosition: &pos,
+			})
+		}
+	})
+}
+
+// ============================================================================
+// Read benchmarks
+// ============================================================================
+
+// BenchmarkPostgresMapBroker_ReadStream measures reading publications from
+// the stream table (the change log used for recovery).
+func BenchmarkPostgresMapBroker_ReadStream(b *testing.B) {
+	broker, cleanup := setupBench(b)
+	defer cleanup()
+	ctx := context.Background()
+	channel := "bench_read_stream"
+
+	// Seed 100 publications.
+	for i := 0; i < 100; i++ {
+		_, err := broker.Publish(ctx, channel, fmt.Sprintf("k%d", i), centrifuge.MapPublishOptions{
+			Data: []byte("data"),
+		})
+		if err != nil {
+			b.Fatal(err)
+		}
+	}
+
+	b.ReportAllocs()
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		_, err := broker.ReadStream(ctx, channel, centrifuge.MapReadStreamOptions{
+			Filter: centrifuge.StreamFilter{Limit: 50},
+		})
+		if err != nil {
+			b.Fatal(err)
+		}
+	}
+}
+
+// BenchmarkPostgresMapBroker_ReadState measures reading the full key-value
+// snapshot of a channel.
+func BenchmarkPostgresMapBroker_ReadState(b *testing.B) {
+	broker, cleanup := setupBench(b)
+	defer cleanup()
+	ctx := context.Background()
+	channel := "bench_read_state"
+
+	for i := 0; i < 100; i++ {
+		_, err := broker.Publish(ctx, channel, fmt.Sprintf("k%d", i), centrifuge.MapPublishOptions{
+			Data: []byte("data"),
+		})
+		if err != nil {
+			b.Fatal(err)
+		}
+	}
+
+	b.ReportAllocs()
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		_, err := broker.ReadState(ctx, channel, centrifuge.MapReadStateOptions{
+			Limit: 100,
+		})
+		if err != nil {
+			b.Fatal(err)
+		}
+	}
+}
+
+// BenchmarkPostgresMapBroker_ReadStatePaginated measures paginated state
+// reads (typical for large channels with cursor-based iteration).
+func BenchmarkPostgresMapBroker_ReadStatePaginated(b *testing.B) {
+	broker, cleanup := setupBench(b)
+	defer cleanup()
+	ctx := context.Background()
+	channel := "bench_read_state_paged"
+
+	for i := 0; i < 100; i++ {
+		_, err := broker.Publish(ctx, channel, fmt.Sprintf("k%d", i), centrifuge.MapPublishOptions{
+			Data: []byte("data"),
+		})
+		if err != nil {
+			b.Fatal(err)
+		}
+	}
+
+	b.ReportAllocs()
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		_, err := broker.ReadState(ctx, channel, centrifuge.MapReadStateOptions{
+			Limit: 10,
+		})
+		if err != nil {
+			b.Fatal(err)
+		}
+	}
+}
+
+// BenchmarkPostgresMapBroker_Stats measures the channel statistics query.
+func BenchmarkPostgresMapBroker_Stats(b *testing.B) {
+	broker, cleanup := setupBench(b)
+	defer cleanup()
+	ctx := context.Background()
+	channel := "bench_stats"
+
+	for i := 0; i < 50; i++ {
+		_, _ = broker.Publish(ctx, channel, fmt.Sprintf("k%d", i), centrifuge.MapPublishOptions{
+			Data: []byte("data"),
+		})
+	}
+
+	b.ReportAllocs()
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		_, err := broker.Stats(ctx, channel)
+		if err != nil {
+			b.Fatal(err)
+		}
+	}
+}
+
+// BenchmarkPostgresMapBroker_Remove measures key removal.
+func BenchmarkPostgresMapBroker_Remove(b *testing.B) {
+	broker, cleanup := setupBench(b)
+	defer cleanup()
+	ctx := context.Background()
+
+	b.ReportAllocs()
+	b.ResetTimer()
+	var counter int64
+	b.RunParallel(func(pb *testing.PB) {
+		for pb.Next() {
+			i := atomic.AddInt64(&counter, 1)
+			ch := fmt.Sprintf("bench_remove_%d", i%1000)
+			key := fmt.Sprintf("key%d", i%1000)
+			// Ensure key exists, then remove it.
+			_, _ = broker.Publish(ctx, ch, key, centrifuge.MapPublishOptions{
+				Data: []byte("x"),
+			})
+			_, err := broker.Remove(ctx, ch, key, centrifuge.MapRemoveOptions{})
+			if err != nil {
+				b.Fatal(err)
+			}
+		}
+	})
+}
+
+// ============================================================================
+// Outbox delivery benchmarks
+// ============================================================================
+
+// BenchmarkPostgresMapBroker_OutboxPublish measures publish latency with
+// the outbox worker running (parallel publishers across many channels).
+func BenchmarkPostgresMapBroker_OutboxPublish(b *testing.B) {
+	broker, cleanup := setupBenchOutbox(b, nil)
+	defer cleanup()
+	ctx := context.Background()
+
+	b.ReportAllocs()
+	b.ResetTimer()
+	var counter int64
+	b.RunParallel(func(pb *testing.PB) {
+		for pb.Next() {
+			i := atomic.AddInt64(&counter, 1)
+			ch := fmt.Sprintf("bench_outbox_%d", i%1000)
+			key := fmt.Sprintf("key%d", i%1000)
+			_, err := broker.Publish(ctx, ch, key, centrifuge.MapPublishOptions{
+				Data: []byte("x"),
+			})
+			if err != nil {
+				b.Fatal(err)
+			}
+		}
+	})
+}
+
+// BenchmarkPostgresMapBroker_OutboxThroughput measures end-to-end outbox
+// throughput: 8 publisher goroutines + outbox workers delivering to a handler.
+// Verifies per-channel offset ordering as a side effect.
+func BenchmarkPostgresMapBroker_OutboxThroughput(b *testing.B) {
+	ctx := context.Background()
+	prefix := fmt.Sprintf("bench_tp_%d_", time.Now().UnixNano())
+
+	var delivered int64
+
+	broker, cleanup := setupBenchOutbox(b, &testBrokerEventHandler{
+		HandlePublicationFunc: func(ch string, pub *centrifuge.Publication, sp centrifuge.StreamPosition, delta bool, prevPub *centrifuge.Publication) error {
+			// Only count publications from this bench run's channels.
+			if len(ch) < len(prefix) || ch[:len(prefix)] != prefix {
+				return nil
+			}
+			atomic.AddInt64(&delivered, 1)
+			return nil
+		},
+	})
+	defer cleanup()
+	time.Sleep(100 * time.Millisecond) // let workers start
+
+	b.ReportAllocs()
+	b.ResetTimer()
+
+	var published int64
+	var counter int64
+	var wg sync.WaitGroup
+	const numPublishers = 8
+	perPublisher := b.N / numPublishers
+	remainder := b.N % numPublishers
+	for p := 0; p < numPublishers; p++ {
+		n := perPublisher
+		if p < remainder {
+			n++
+		}
+		if n == 0 {
+			continue
+		}
+		wg.Add(1)
+		go func(count int) {
+			defer wg.Done()
+			for j := 0; j < count; j++ {
+				i := atomic.AddInt64(&counter, 1)
+				ch := fmt.Sprintf("%sch%d", prefix, i%100)
+				key := fmt.Sprintf("key%d", i)
+				_, err := broker.Publish(ctx, ch, key, centrifuge.MapPublishOptions{
+					Data: []byte("x"),
+				})
+				if err != nil {
+					b.Error(err)
+					return
+				}
+				atomic.AddInt64(&published, 1)
+			}
+		}(n)
+	}
+	wg.Wait()
+
+	want := atomic.LoadInt64(&published)
+	deadline := time.NewTimer(5 * time.Second)
+	defer deadline.Stop()
+	for {
+		if atomic.LoadInt64(&delivered) >= want {
+			break
+		}
+		select {
+		case <-deadline.C:
+			b.Fatalf("timeout: delivered %d/%d", atomic.LoadInt64(&delivered), want)
+		default:
+			time.Sleep(time.Millisecond)
+		}
+	}
+	b.StopTimer()
+}
+
+// BenchmarkPostgresMapBroker_OutboxLatency measures publish-to-delivery
+// latency via the outbox. Reports average latency as a custom metric.
+func BenchmarkPostgresMapBroker_OutboxLatency(b *testing.B) {
+	ctx := context.Background()
+	channel := fmt.Sprintf("bench_latency_%d", time.Now().UnixNano())
+
+	publishTimes := make([]time.Time, b.N+2)
+	delivered := make(chan time.Duration, b.N)
+	var mu sync.Mutex
+
+	broker, cleanup := setupBenchOutbox(b, &testBrokerEventHandler{
+		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(sp.Offset) < len(publishTimes) {
+				t = publishTimes[sp.Offset]
+			}
+			mu.Unlock()
+			if !t.IsZero() {
+				delivered <- time.Since(t)
+			}
+			return nil
+		},
+	})
+	defer cleanup()
+	time.Sleep(100 * time.Millisecond) // let workers start
+
+	b.ReportAllocs()
+	b.ResetTimer()
+
+	for i := 0; i < b.N; i++ {
+		mu.Lock()
+		now := time.Now()
+		res, err := broker.Publish(ctx, channel, fmt.Sprintf("key%d", i), centrifuge.MapPublishOptions{
+			Data: []byte("x"),
+		})
+		if err != nil {
+			mu.Unlock()
+			b.Fatal(err)
+		}
+		if int(res.Position.Offset) < len(publishTimes) {
+			publishTimes[res.Position.Offset] = now
+		}
+		mu.Unlock()
+	}
+
+	deadline := time.NewTimer(30 * time.Second)
+	defer deadline.Stop()
+	var totalLatency time.Duration
+	received := 0
+	for received < b.N {
+		select {
+		case d := <-delivered:
+			totalLatency += d
+			received++
+		case <-deadline.C:
+			b.Fatalf("timeout: received %d/%d", received, b.N)
+		}
+	}
+	b.StopTimer()
+
+	if b.N > 0 {
+		avg := totalLatency / time.Duration(b.N)
+		b.ReportMetric(float64(avg.Microseconds()), "µs/delivery")
+	}
+}
+
+// ============================================================================
+// Cleanup benchmark
+// ============================================================================
+
+// BenchmarkPostgresMapBroker_Cleanup measures TTL-based key expiration
+// throughput via the expire_keys SQL function. Sub-benchmarks cover
+// different key counts.
+func BenchmarkPostgresMapBroker_Cleanup(b *testing.B) {
+	for _, numKeys := range []int{100, 1000} {
+		b.Run(fmt.Sprintf("keys_%d", numKeys), func(b *testing.B) {
+			connString := getPostgresConnString(b)
+			node, _ := centrifuge.New(centrifuge.Config{
+				Map: centrifuge.MapConfig{
+					GetMapChannelOptions: func(channel string) centrifuge.MapChannelOptions {
+						return centrifuge.MapChannelOptions{
+							Mode:   centrifuge.MapModeRecoverable,
+							KeyTTL: time.Millisecond,
+						}
+					},
+				},
+			})
+			broker, err := NewPostgresMapBroker(node, PostgresMapBrokerConfig{
+				DSN:                    connString,
+				BinaryData:             true,
+				PartitionRetentionDays: 7,
+			})
+			if err != nil {
+				b.Fatal(err)
+			}
+			ctx := context.Background()
+			if err := broker.EnsureSchema(ctx); err != nil {
+				b.Fatal(err)
+			}
+			_ = broker.RegisterEventHandler(nil)
+			defer func() {
+				_ = broker.Close(context.Background())
+				_ = node.Shutdown(context.Background())
+			}()
+
+			b.ReportAllocs()
+			b.ResetTimer()
+			for i := 0; i < b.N; i++ {
+				b.StopTimer()
+				channel := fmt.Sprintf("bench_cleanup_%d_%d", time.Now().UnixNano(), i)
+				for k := 0; k < numKeys; k++ {
+					_, err := broker.Publish(ctx, channel, fmt.Sprintf("key_%d", k), centrifuge.MapPublishOptions{
+						Data: []byte("data"),
+					})
+					if err != nil {
+						b.Fatal(err)
+					}
+				}
+				time.Sleep(5 * time.Millisecond)
+				b.StartTimer()
+
+				numShards := broker.conf.NumShards
+				var metaTTL *string
+				_, err := broker.pool.Exec(ctx, fmt.Sprintf(`
+					SELECT %s($1, $2, $3::interval, $4)
+				`, broker.names.expireKeys), 1000, numShards, metaTTL, channel)
+				if err != nil {
+					b.Fatal(err)
+				}
+			}
+			b.ReportMetric(float64(numKeys), "keys/op")
+		})
+	}
+}