Initial open source release

internal/mtprotoedge/session_manager_bench_test.go

@@ -0,0 +1,215 @@
+package mtprotoedge
+
+import (
+	"context"
+	"fmt"
+	"os"
+	"runtime"
+	"sort"
+	"strconv"
+	"sync"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"go.uber.org/zap"
+
+	"github.com/gotd/td/proto"
+	"github.com/gotd/td/tg"
+)
+
+// 连接层 fan-out / churn 压测：聚焦 SessionManager 的锁争用，不走真实 socket / 加密。
+//
+// 构造的 Conn 故意不 startOutbound：pushToUser 持锁快照 byUser 后，锁外对每个 conn 调 c.Send，
+// 此时 c.outbound==nil 立即返回 ErrConnClosed（见 outbound.go），因此测量集中在「持锁段 + 分发开销」，
+// 即分片要消除的全局锁热点。每条连接仅结构体内存（无 1024 容量的 outbound channel、无 goroutine），
+// 故可注册到 20 万规模。
+//
+// 用法：
+//
+//	go test ./internal/mtprotoedge/ -run '^$' -bench BenchmarkSessionManager -benchmem -cpu 1,4,8
+//	go test ./internal/mtprotoedge/ -run '^$' -bench BenchmarkSessionManagerPushConcurrent -mutexprofile mu.out
+//	TELESRV_LOAD_CONNS=200000 go test ./internal/mtprotoedge/ -run TestSessionManagerFanoutThroughput -v -timeout 300s
+
+func benchConn(sessionID int64, authKeyID [8]byte, userID int64) *Conn {
+	c := &Conn{sessionID: sessionID, authKeyID: authKeyID}
+	if userID != 0 {
+		c.userID.Store(userID)
+		c.userIDResolved.Store(true)
+	}
+	c.receivesUpdates.Store(true) // 走 fanout 的「收集 conns→锁外 Send」分支，而非 pending 暂存
+	return c
+}
+
+func authKeyIDFromInt(v uint64) [8]byte {
+	var id [8]byte
+	for i := 0; i < 8; i++ {
+		id[i] = byte(v >> (8 * i))
+	}
+	return id
+}
+
+// seedSessions 注册 conns 个连接，每个 user 绑定 connsPerUser 个连接（模拟多设备）。
+// 返回注册的 userID 列表（去重、有序范围 [1, userCount]）。
+func seedSessions(sm *SessionManager, conns, connsPerUser int) (userCount int) {
+	if connsPerUser < 1 {
+		connsPerUser = 1
+	}
+	for i := 0; i < conns; i++ {
+		userID := int64(i/connsPerUser) + 1
+		sm.Register(benchConn(int64(i)+1, authKeyIDFromInt(uint64(i)+1), userID))
+	}
+	return (conns + connsPerUser - 1) / connsPerUser
+}
+
+// BenchmarkSessionManagerPushConcurrent 模拟 20 万在线下的真实热点：大量 goroutine 并发对
+// 不同 user pushToUser，全部抢同一把全局锁。-mutexprofile 会把 SessionManager.mu 顶上来。
+func BenchmarkSessionManagerPushConcurrent(b *testing.B) {
+	const conns = 200_000
+	const connsPerUser = 2
+	sm := NewSessionManager(zap.NewNop())
+	userCount := seedSessions(sm, conns, connsPerUser)
+	msg := &tg.UpdatesTooLong{}
+	ctx := context.Background()
+
+	b.ResetTimer()
+	b.RunParallel(func(pb *testing.PB) {
+		var n uint64
+		for pb.Next() {
+			n++
+			userID := int64(n%uint64(userCount)) + 1
+			_, _ = sm.PushToUser(ctx, userID, proto.MessageFromServer, msg)
+		}
+	})
+}
+
+// BenchmarkSessionManagerRegisterChurn 测连接建立/断开的锁成本：并发 Register+Unregister。
+// 20 万在线意味着持续的 connect/disconnect churn，每次都抢全局写锁。
+func BenchmarkSessionManagerRegisterChurn(b *testing.B) {
+	sm := NewSessionManager(zap.NewNop())
+	var seq atomic.Uint64
+
+	b.ResetTimer()
+	b.RunParallel(func(pb *testing.PB) {
+		for pb.Next() {
+			id := seq.Add(1)
+			c := benchConn(int64(id), authKeyIDFromInt(id), int64(id))
+			sm.Register(c)
+			sm.Unregister(c)
+		}
+	})
+}
+
+// BenchmarkSessionManagerPushFanoutWidth 测单次 push 的 fanout 广度成本：一个 user 绑定很多连接，
+// 单次 PushToUser 要持锁遍历全部。真实私聊 user 设备数少（2-4），此为上界参考。
+func BenchmarkSessionManagerPushFanoutWidth(b *testing.B) {
+	for _, width := range []int{1, 4, 16, 64} {
+		b.Run(fmt.Sprintf("width=%d", width), func(b *testing.B) {
+			sm := NewSessionManager(zap.NewNop())
+			const userID = 1
+			for i := 0; i < width; i++ {
+				sm.Register(benchConn(int64(i)+1, authKeyIDFromInt(uint64(i)+1), userID))
+			}
+			msg := &tg.UpdatesTooLong{}
+			ctx := context.Background()
+			b.ResetTimer()
+			for i := 0; i < b.N; i++ {
+				_, _ = sm.PushToUser(ctx, userID, proto.MessageFromServer, msg)
+			}
+		})
+	}
+}
+
+// TestSessionManagerFanoutThroughput 是数据驱动吞吐测：注册 N 连接后，P 个 goroutine 持续并发
+// push，测全局锁下的实际 push 吞吐与 p99。默认小规模冒烟；设 TELESRV_LOAD_CONNS 放大到 20 万。
+func TestSessionManagerFanoutThroughput(t *testing.T) {
+	conns := envIntDefault("TELESRV_LOAD_CONNS", 20_000)
+	connsPerUser := envIntDefault("TELESRV_LOAD_CONNS_PER_USER", 2)
+	workers := envIntDefault("TELESRV_LOAD_PUSH_WORKERS", 0) // 0 → GOMAXPROCS
+	duration := time.Duration(envIntDefault("TELESRV_LOAD_SECONDS", 3)) * time.Second
+	if workers <= 0 {
+		workers = runtime.GOMAXPROCS(0)
+	}
+
+	sm := NewSessionManager(zap.NewNop())
+	t0 := time.Now()
+	userCount := seedSessions(sm, conns, connsPerUser)
+	seedWall := time.Since(t0)
+	if got := sm.Online(); got != conns {
+		t.Fatalf("online = %d, want %d", got, conns)
+	}
+
+	msg := &tg.UpdatesTooLong{}
+	ctx := context.Background()
+	var ops atomic.Int64
+	perWorkerLat := make([][]time.Duration, workers)
+
+	deadline := time.Now().Add(duration)
+	var wg sync.WaitGroup
+	start := time.Now()
+	for w := 0; w < workers; w++ {
+		wg.Add(1)
+		go func(w int) {
+			defer wg.Done()
+			lat := make([]time.Duration, 0, 1<<16)
+			var n uint64
+			for time.Now().Before(deadline) {
+				// 批量 256 次再查一次时钟，降低 time.Now 占比。
+				for j := 0; j < 256; j++ {
+					n++
+					userID := int64(n%uint64(userCount)) + 1
+					s := time.Now()
+					_, _ = sm.PushToUser(ctx, userID, proto.MessageFromServer, msg)
+					lat = append(lat, time.Since(s))
+				}
+				ops.Add(256)
+			}
+			perWorkerLat[w] = lat
+		}(w)
+	}
+	wg.Wait()
+	wall := time.Since(start)
+
+	all := make([]time.Duration, 0, ops.Load())
+	for _, l := range perWorkerLat {
+		all = append(all, l...)
+	}
+	sortDurations(all)
+	total := ops.Load()
+	thr := float64(total) / wall.Seconds()
+
+	t.Logf("==== session_manager fan-out throughput ====")
+	t.Logf("config:  conns=%d connsPerUser=%d users=%d pushWorkers=%d dur=%s seed=%s",
+		conns, connsPerUser, userCount, workers, duration, seedWall.Round(time.Millisecond))
+	t.Logf("push:    %d ops in %s -> %.0f push/s", total, wall.Round(time.Millisecond), thr)
+	t.Logf("push.lat p50=%s p90=%s p99=%s max=%s",
+		pct(all, 50), pct(all, 90), pct(all, 99), pct(all, 100))
+	t.Logf("=============================================")
+}
+
+func pct(sorted []time.Duration, p int) time.Duration {
+	if len(sorted) == 0 {
+		return 0
+	}
+	idx := (p*len(sorted))/100 - 1
+	if idx < 0 {
+		idx = 0
+	}
+	if idx >= len(sorted) {
+		idx = len(sorted) - 1
+	}
+	return sorted[idx]
+}
+
+func sortDurations(d []time.Duration) {
+	sort.Slice(d, func(i, j int) bool { return d[i] < d[j] })
+}
+
+func envIntDefault(key string, def int) int {
+	if v := os.Getenv(key); v != "" {
+		if n, err := strconv.Atoi(v); err == nil {
+			return n
+		}
+	}
+	return def
+}