owpengram-server/internal/mtprotoedge/session_manager_test.go

991 lines
33 KiB
Go
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

package mtprotoedge
import (
"context"
"errors"
"sync"
"sync/atomic"
"testing"
"time"
"go.uber.org/zap/zaptest"
"github.com/gotd/td/bin"
"github.com/gotd/td/mt"
"github.com/gotd/td/proto"
"github.com/gotd/td/tg"
)
type countingOutboundEncoder struct {
count *int
}
func (e *countingOutboundEncoder) Encode(b *bin.Buffer) error {
*e.count++
return (&tg.UpdatesTooLong{}).Encode(b)
}
type closeCountingTransport struct {
closes int
}
type slowCloseTransport struct {
delay time.Duration
release <-chan struct{}
done chan struct{}
once sync.Once
closes atomic.Int32
}
func newSlowCloseTransport(delay time.Duration, release <-chan struct{}) *slowCloseTransport {
return &slowCloseTransport{delay: delay, release: release, done: make(chan struct{})}
}
func (*slowCloseTransport) Send(context.Context, *bin.Buffer) error {
return errors.New("test transport send")
}
func (*slowCloseTransport) Recv(context.Context, *bin.Buffer) error {
return errors.New("test transport recv")
}
func (t *slowCloseTransport) Close() error {
t.closes.Add(1)
if t.release != nil {
<-t.release
} else if t.delay > 0 {
time.Sleep(t.delay)
}
t.once.Do(func() { close(t.done) })
return nil
}
func (t *closeCountingTransport) Send(context.Context, *bin.Buffer) error {
return errors.New("test transport send")
}
func (t *closeCountingTransport) Recv(context.Context, *bin.Buffer) error {
return errors.New("test transport recv")
}
func (t *closeCountingTransport) Close() error {
t.closes++
return nil
}
// TestSessionManagerRegistry 验证注册表的注册/注销/查找语义(不涉及网络发送)。
func TestSessionManagerRegistry(t *testing.T) {
sm := NewSessionManager(zaptest.NewLogger(t))
c := &Conn{sessionID: 42, authKeyID: [8]byte{1, 2, 3}}
c.receivesUpdates.Store(true)
sm.Register(c)
if got := sm.Online(); got != 1 {
t.Fatalf("online = %d, want 1", got)
}
sm.BindAuthKey(42, [8]byte{1, 2, 3})
sm.BindUser(42, 100)
if userID, ok := sm.UserID(42); !ok || userID != 100 {
t.Fatalf("cached user = %d ok %v, want 100/true", userID, ok)
}
sm.BindAuthKey(42, [8]byte{9})
if userID, ok := sm.UserID(42); ok || userID != 0 {
t.Fatalf("cached user after auth key switch = %d ok %v, want 0/false", userID, ok)
}
if userID, resolved := sm.UserIDResolved(42); resolved || userID != 0 {
t.Fatalf("resolved user after auth key switch = %d resolved %v, want unresolved", userID, resolved)
}
sm.BindUser(42, 0)
if userID, resolved := sm.UserIDResolved(42); !resolved || userID != 0 {
t.Fatalf("negative user cache = %d resolved %v, want 0/true", userID, resolved)
}
sm.Unregister(c)
if got := sm.Online(); got != 0 {
t.Fatalf("online after unregister = %d, want 0", got)
}
err := sm.PushToSession(context.Background(), 42, proto.MessageFromServer, &tg.UpdatesTooLong{})
if !errors.Is(err, ErrSessionNotFound) {
t.Fatalf("push to missing session err = %v, want ErrSessionNotFound", err)
}
}
func TestSessionManagerReplacementClosesOldPhysicalTransport(t *testing.T) {
sm := NewSessionManager(zaptest.NewLogger(t))
raw := [8]byte{1, 2, 3}
oldTransport := &closeCountingTransport{}
old := &Conn{sessionID: 42, authKeyID: raw, transport: oldTransport}
replacement := &Conn{sessionID: 42, authKeyID: raw}
sm.Register(old)
sm.Register(replacement)
if oldTransport.closes != 1 {
t.Fatalf("old transport closes = %d, want 1", oldTransport.closes)
}
// 旧 serveConn 稍后退出时不得把 replacement 从索引删掉。
sm.Unregister(old)
if got, ok := sm.bySession[sessionKey{authKeyID: raw, sessionID: 42}]; !ok || got != replacement {
t.Fatal("old unregister removed the replacement connection")
}
}
func TestSessionManagerDestroyClosesPhysicalTransport(t *testing.T) {
sm := NewSessionManager(zaptest.NewLogger(t))
raw := [8]byte{4, 5, 6}
physical := &closeCountingTransport{}
c := &Conn{sessionID: 77, authKeyID: raw, transport: physical}
sm.Register(c)
if !sm.DestroySessionForAuthKey(raw, 77) {
t.Fatal("DestroySessionForAuthKey returned false")
}
if physical.closes != 1 {
t.Fatalf("destroyed transport closes = %d, want 1", physical.closes)
}
if sm.Online() != 0 {
t.Fatalf("online after destroy = %d, want 0", sm.Online())
}
}
func TestSessionManagerBestEffortFanoutPreencodesOnce(t *testing.T) {
sm := NewSessionManager(zaptest.NewLogger(t))
const userID = int64(100)
for i := 0; i < 2; i++ {
c := &Conn{
sessionID: int64(i + 1),
authKeyID: [8]byte{byte(i + 1)},
outbound: make(chan outboundOp, 1),
outboundControl: make(chan outboundOp, 1),
outboundStop: make(chan struct{}),
metrics: NopMetrics{},
}
c.userID.Store(userID)
c.userIDResolved.Store(true)
c.receivesUpdates.Store(true)
sm.Register(c)
}
encodes := 0
sent, err := sm.PushToUserExceptSessionBestEffort(
context.Background(),
userID,
0,
proto.MessageFromServer,
&countingOutboundEncoder{count: &encodes},
0,
)
if err != nil {
t.Fatalf("push: %v", err)
}
if sent != 2 {
t.Fatalf("sent = %d, want 2", sent)
}
if encodes != 1 {
t.Fatalf("encoded %d times, want 1", encodes)
}
}
func TestSessionManagerPendingFanoutSharesOneEncodedBodyAndBudget(t *testing.T) {
sm := NewSessionManager(zaptest.NewLogger(t))
const userID = int64(102)
keys := make([]sessionKey, 0, 2)
for i := 0; i < 2; i++ {
c := &Conn{sessionID: int64(i + 1), authKeyID: [8]byte{byte(i + 1)}}
c.userID.Store(userID)
c.userIDResolved.Store(true)
sm.Register(c)
keys = append(keys, connSessionKey(c))
}
encodes := 0
msg := &countingOutboundEncoder{count: &encodes}
sent, err := sm.PushToUserExceptSession(context.Background(), userID, 0, proto.MessageFromServer, msg)
if err != nil {
t.Fatalf("push: %v", err)
}
if sent != 2 || encodes != 1 {
t.Fatalf("pending fanout = sent:%d encodes:%d, want 2/1", sent, encodes)
}
sm.mu.Lock()
first := sm.pending[keys[0]][0]
second := sm.pending[keys[1]][0]
if first.encoded != second.encoded || first.reservation != second.reservation {
sm.mu.Unlock()
t.Fatal("pending sessions did not share encoded body/reservation")
}
wantBytes := int64(len(first.encoded.body))
sm.deletePendingLocked(keys[0])
if got := sm.pendingBudget.snapshot(); got != wantBytes {
sm.mu.Unlock()
t.Fatalf("budget after first session drop = %d, want shared body %d", got, wantBytes)
}
sm.deletePendingLocked(keys[1])
sm.mu.Unlock()
if got := sm.pendingBudget.snapshot(); got != 0 {
t.Fatalf("budget after last session drop = %d, want 0", got)
}
}
func TestSessionManagerBestEffortFanoutUsesOneBudgetAndDropsOnlySlowConsumers(t *testing.T) {
sm := NewSessionManager(zaptest.NewLogger(t))
const userID = int64(101)
// 三个满队列模拟三个慢设备;没有 outbound actor确保队列在测试期间不会自行排空。
slow := make([]*Conn, 0, 3)
for i := 0; i < 3; i++ {
tr := &closeCountingTransport{}
c := &Conn{
sessionID: int64(i + 1),
authKeyID: [8]byte{byte(i + 1)},
transport: tr,
metrics: NopMetrics{},
outbound: make(chan outboundOp, 1),
outboundControl: make(chan outboundOp, 1),
outboundStop: make(chan struct{}),
}
c.outbound <- outboundOp{}
c.userID.Store(userID)
c.userIDResolved.Store(true)
c.receivesUpdates.Store(true)
sm.Register(c)
slow = append(slow, c)
}
healthy := &Conn{
sessionID: 99,
authKeyID: [8]byte{99},
metrics: NopMetrics{},
outbound: make(chan outboundOp, 1),
outboundControl: make(chan outboundOp, 1),
outboundStop: make(chan struct{}),
}
healthy.userID.Store(userID)
healthy.userIDResolved.Store(true)
healthy.receivesUpdates.Store(true)
sm.Register(healthy)
const budget = 40 * time.Millisecond
start := time.Now()
sent, err := sm.PushToUserExceptSessionBestEffort(
context.Background(), userID, 0, proto.MessageFromServer, &tg.UpdatesTooLong{}, budget,
)
elapsed := time.Since(start)
if err != nil {
t.Fatalf("push: %v", err)
}
if sent != 1 {
t.Fatalf("sent = %d, want only healthy session", sent)
}
if elapsed >= 3*budget {
t.Fatalf("fan-out waited %v; want one shared %v budget, not one per slow session", elapsed, budget)
}
if got := len(healthy.outbound); got != 1 {
t.Fatalf("healthy queued ops = %d, want 1", got)
}
if healthy.isRetired() {
t.Fatal("healthy session was terminalized")
}
for i, c := range slow {
if !c.isRetired() {
t.Fatalf("slow session %d was not terminalized", i)
}
if tr := c.transport.(*closeCountingTransport); tr.closes != 1 {
t.Fatalf("slow session %d transport closes = %d, want 1", i, tr.closes)
}
}
}
func TestSessionManagerScopesSameSessionIDByAuthKey(t *testing.T) {
sm := NewSessionManager(zaptest.NewLogger(t))
raw1 := [8]byte{1}
raw2 := [8]byte{2}
perm1 := [8]byte{9}
c1 := &Conn{sessionID: 42, authKeyID: raw1}
c2 := &Conn{sessionID: 42, authKeyID: raw2}
sm.Register(c1)
sm.Register(c2)
if got := sm.Online(); got != 2 {
t.Fatalf("online = %d, want 2", got)
}
sm.BindAuthKeyForSession(raw1, 42, perm1)
// 两条 PFS/raw 连接可以解析到同一业务 perm key 且复用同一个 session_id
// 精确排除必须只匹配 raw1不能按 business key 把 raw2 一并排除。
sm.BindAuthKeyForSession(raw2, 42, perm1)
sm.BindUserForAuthKey(raw1, 42, 100)
sm.BindUserForAuthKey(raw2, 42, 200)
if userID, ok := sm.UserIDForAuthKey(raw1, 42); !ok || userID != 100 {
t.Fatalf("scoped user raw1 = %d ok %v, want 100/true", userID, ok)
}
if userID, ok := sm.UserIDForAuthKey(raw2, 42); !ok || userID != 200 {
t.Fatalf("scoped user raw2 = %d ok %v, want 200/true", userID, ok)
}
if _, ok := sm.UserID(42); ok {
t.Fatal("legacy UserID unexpectedly resolved ambiguous session_id")
}
if err := sm.PushToSession(context.Background(), 42, proto.MessageFromServer, &tg.UpdatesTooLong{}); !errors.Is(err, ErrSessionAmbiguous) {
t.Fatalf("ambiguous push err = %v, want ErrSessionAmbiguous", err)
}
sm.BindUserForAuthKey(raw1, 42, 300)
sm.BindUserForAuthKey(raw2, 42, 300)
sent, err := sm.PushToUserExceptAuthKeySession(context.Background(), 300, raw1, 42, proto.MessageFromServer, &tg.UpdatesTooLong{})
if err != nil {
t.Fatalf("push except scoped session: %v", err)
}
if sent != 1 {
t.Fatalf("pushed to %d sessions, want 1", sent)
}
if _, ok := sm.pending[sessionKey{authKeyID: raw1, sessionID: 42}]; ok {
t.Fatal("excluded session received pending push")
}
if got := len(sm.pending[sessionKey{authKeyID: raw2, sessionID: 42}]); got != 1 {
t.Fatalf("raw2 pending pushes = %d, want 1", got)
}
if !sm.DestroySessionForAuthKey(raw1, 42) {
t.Fatal("scoped destroy did not remove raw1 session")
}
if _, ok := sm.AuthKeyIDForSession(raw1, 42); ok {
t.Fatal("raw1 session survived scoped destroy")
}
if _, ok := sm.AuthKeyIDForSession(raw2, 42); !ok {
t.Fatal("same session_id under raw2 was removed by scoped destroy")
}
}
func TestSessionManagerCloseSessionsForBusinessAuthKeyClosesBoundTempAndRaw(t *testing.T) {
sm := NewSessionManager(zaptest.NewLogger(t))
rawTemp := [8]byte{1}
perm := [8]byte{9}
otherRaw := [8]byte{2}
otherPerm := [8]byte{8}
tempTransport := &closeCountingTransport{}
permTransport := &closeCountingTransport{}
otherTransport := &closeCountingTransport{}
cTemp := &Conn{sessionID: 11, authKeyID: rawTemp, transport: tempTransport}
cPerm := &Conn{sessionID: 12, authKeyID: perm, transport: permTransport}
cOther := &Conn{sessionID: 13, authKeyID: otherRaw}
cOther.transport = otherTransport
sm.Register(cTemp)
sm.Register(cPerm)
sm.Register(cOther)
sm.BindAuthKeyForSession(rawTemp, 11, perm)
sm.BindAuthKeyForSession(perm, 12, perm)
sm.BindAuthKeyForSession(otherRaw, 13, otherPerm)
sm.BindUserForAuthKey(rawTemp, 11, 100)
sm.BindUserForAuthKey(perm, 12, 100)
sm.BindUserForAuthKey(otherRaw, 13, 200)
if closed := sm.CloseSessionsForBusinessAuthKey(perm); closed != 2 {
t.Fatalf("closed sessions = %d, want 2", closed)
}
if tempTransport.closes != 1 || permTransport.closes != 1 {
t.Fatalf("transport closes temp=%d perm=%d, want 1/1", tempTransport.closes, permTransport.closes)
}
if otherTransport.closes != 0 {
t.Fatalf("other transport closes = %d, want 0", otherTransport.closes)
}
if got := sm.Online(); got != 1 {
t.Fatalf("online after close = %d, want 1", got)
}
if _, ok := sm.AuthKeyIDForSession(rawTemp, 11); ok {
t.Fatal("temp session still indexed after business auth key close")
}
if _, ok := sm.AuthKeyIDForSession(perm, 12); ok {
t.Fatal("raw perm session still indexed after business auth key close")
}
if userID, ok := sm.UserIDForAuthKey(otherRaw, 13); !ok || userID != 200 {
t.Fatalf("other session user = %d ok %v, want 200/true", userID, ok)
}
if closed := sm.CloseSessionsForBusinessAuthKey(perm); closed != 0 {
t.Fatalf("second close = %d, want 0", closed)
}
}
func TestSessionManagerCloseSessionsRunsSlowPhysicalClosesConcurrently(t *testing.T) {
sm := NewSessionManager(zaptest.NewLogger(t))
business := [8]byte{9, 9, 9}
const sessions = 8
const closeDelay = 75 * time.Millisecond
transports := make([]*slowCloseTransport, 0, sessions)
for i := 0; i < sessions; i++ {
raw := [8]byte{byte(i + 1)}
tr := newSlowCloseTransport(closeDelay, nil)
c := &Conn{sessionID: int64(i + 1), authKeyID: raw, transport: tr}
sm.Register(c)
sm.BindAuthKeyForSession(raw, c.sessionID, business)
transports = append(transports, tr)
}
started := time.Now()
if got := sm.CloseSessionsForBusinessAuthKey(business); got != sessions {
t.Fatalf("closed sessions = %d, want %d", got, sessions)
}
elapsed := time.Since(started)
// A serial implementation takes ~600ms. Leave ample Windows/CI scheduling margin while
// still proving that the per-Conn delay is not multiplied by the session count.
if elapsed >= 4*closeDelay {
t.Fatalf("batch close elapsed = %v, want concurrent closes near %v", elapsed, closeDelay)
}
for i, tr := range transports {
select {
case <-tr.done:
default:
t.Fatalf("transport %d close had not completed when batch returned", i)
}
if got := tr.closes.Load(); got != 1 {
t.Fatalf("transport %d closes = %d, want 1", i, got)
}
}
}
func TestSessionManagerCloseRawSessionsExceptRunsConcurrentlyAndPreservesExcluded(t *testing.T) {
sm := NewSessionManager(zaptest.NewLogger(t))
raw := [8]byte{6, 6, 6}
const sessions = 7
const excludedSession = int64(4)
const closeDelay = 60 * time.Millisecond
transports := make([]*slowCloseTransport, 0, sessions)
for i := 0; i < sessions; i++ {
tr := newSlowCloseTransport(closeDelay, nil)
c := &Conn{sessionID: int64(i + 1), authKeyID: raw, transport: tr}
sm.Register(c)
transports = append(transports, tr)
}
started := time.Now()
if got, want := sm.CloseSessionsForRawAuthKeyExcept(raw, excludedSession), sessions-1; got != want {
t.Fatalf("closed sessions = %d, want %d", got, want)
}
if elapsed := time.Since(started); elapsed >= 4*closeDelay {
t.Fatalf("raw-key batch close elapsed = %v, want concurrent closes near %v", elapsed, closeDelay)
}
for i, tr := range transports {
sessionID := int64(i + 1)
if sessionID == excludedSession {
if got := tr.closes.Load(); got != 0 {
t.Fatalf("excluded transport closes = %d, want 0", got)
}
continue
}
select {
case <-tr.done:
default:
t.Fatalf("transport for session %d had not closed", sessionID)
}
}
if _, ok := sm.bySession[sessionKey{authKeyID: raw, sessionID: excludedSession}]; !ok {
t.Fatal("excluded session was removed from the registry")
}
// Clean up the deliberately preserved connection without making the assertion path depend
// on test process teardown.
if !sm.DestroySessionForAuthKey(raw, excludedSession) {
t.Fatal("cleanup destroy of excluded session failed")
}
}
func TestForceCloseBatchTimeoutStillClosesProducerAndRPCGates(t *testing.T) {
release := make(chan struct{})
const sessions = 4
scheduler := newInboundRPCScheduler(1, 16, 1<<20)
defer scheduler.stop(time.Second)
conns := make([]*Conn, 0, sessions)
transports := make([]*slowCloseTransport, 0, sessions)
for i := 0; i < sessions; i++ {
tr := newSlowCloseTransport(0, release)
c := &Conn{
transport: tr,
metrics: NopMetrics{},
outbound: make(chan outboundOp, 1),
outboundControl: make(chan outboundOp, 1),
outboundStop: make(chan struct{}),
}
c.startInboundRPCScheduler(scheduler, 1, 1, time.Second)
if err := c.enqueueInboundRPC(context.Background(), inboundRPC{
method: "shutdown.budget",
size: 32,
}); err != nil {
t.Fatalf("enqueue queued RPC %d: %v", i, err)
}
conns = append(conns, c)
transports = append(transports, tr)
}
started := time.Now()
if completed := forceCloseConnBatch(conns, 40*time.Millisecond); completed {
t.Fatal("blocked transport close batch unexpectedly completed")
}
if elapsed := time.Since(started); elapsed > 250*time.Millisecond {
t.Fatalf("timed batch close blocked for %v", elapsed)
}
for i, c := range conns {
if !c.isRetired() {
t.Fatalf("connection %d producer gate remains open after batch timeout", i)
}
select {
case <-c.outboundStop:
default:
t.Fatalf("connection %d outbound stop was not published", i)
}
select {
case <-c.rpcRootCtx.Done():
default:
t.Fatalf("connection %d RPC root remains open after batch timeout", i)
}
}
if tasks, bytes := scheduler.budgetSnapshot(); tasks != 0 || bytes != 0 {
t.Fatalf("RPC budget after batch gate close = tasks:%d bytes:%d, want zero", tasks, bytes)
}
close(release)
for i, tr := range transports {
select {
case <-tr.done:
case <-time.After(time.Second):
t.Fatalf("transport %d did not finish after release", i)
}
}
}
func TestSessionManagerBusinessAuthKeyIndexTracksRebind(t *testing.T) {
sm := NewSessionManager(zaptest.NewLogger(t))
raw := [8]byte{1}
oldPerm := [8]byte{7}
newPerm := [8]byte{8}
c := &Conn{sessionID: 21, authKeyID: raw}
sm.Register(c)
sm.BindAuthKeyForSession(raw, 21, oldPerm)
sm.BindAuthKeyForSession(raw, 21, newPerm)
if closed := sm.CloseSessionsForBusinessAuthKey(oldPerm); closed != 0 {
t.Fatalf("close old business auth key = %d, want 0", closed)
}
if got := sm.Online(); got != 1 {
t.Fatalf("online after closing old key = %d, want 1", got)
}
if closed := sm.CloseSessionsForBusinessAuthKey(newPerm); closed != 1 {
t.Fatalf("close new business auth key = %d, want 1", closed)
}
if got := sm.Online(); got != 0 {
t.Fatalf("online after closing new key = %d, want 0", got)
}
}
func TestPushToUserAuthKeyUsesOneDeadlineAndDropsOnlySlowPFSConnections(t *testing.T) {
sm := NewSessionManager(zaptest.NewLogger(t))
business := [8]byte{9, 9}
const userID = int64(100)
newConn := func(raw [8]byte, sessionID int64, queueFull bool) (*Conn, *closeCountingTransport) {
transport := &closeCountingTransport{}
c := &Conn{
authKeyID: raw,
sessionID: sessionID,
metrics: NopMetrics{},
transport: transport,
outbound: make(chan outboundOp, 1),
outboundControl: make(chan outboundOp, 1),
outboundStop: make(chan struct{}),
}
c.receivesUpdates.Store(true)
if queueFull {
c.outbound <- outboundOp{}
}
sm.Register(c)
sm.BindAuthKeyForSession(raw, sessionID, business)
sm.BindUserForAuthKey(raw, sessionID, userID)
return c, transport
}
slowOne, slowOneTransport := newConn([8]byte{1}, 11, true)
slowTwo, slowTwoTransport := newConn([8]byte{2}, 12, true)
healthy, healthyTransport := newConn([8]byte{3}, 13, false)
ctx, cancel := context.WithTimeout(context.Background(), 25*time.Millisecond)
defer cancel()
started := time.Now()
sent, err := sm.PushToUserAuthKey(ctx, userID, business, proto.MessageFromServer, &tg.UpdatesTooLong{})
elapsed := time.Since(started)
if err != nil {
t.Fatalf("PushToUserAuthKey: %v", err)
}
if sent != 1 {
t.Fatalf("sent = %d, want only healthy connection", sent)
}
if elapsed > 100*time.Millisecond {
t.Fatalf("elapsed = %v, want one shared deadline rather than per-session waits", elapsed)
}
if !slowOne.isRetired() || !slowTwo.isRetired() || slowOneTransport.closes != 1 || slowTwoTransport.closes != 1 {
t.Fatalf("slow connections not terminal/closed: one=%v/%d two=%v/%d",
slowOne.isRetired(), slowOneTransport.closes, slowTwo.isRetired(), slowTwoTransport.closes)
}
if healthy.isRetired() || healthyTransport.closes != 0 {
t.Fatalf("healthy connection was dropped: lifecycle=%v closes=%d", healthy.lifecycleState(), healthyTransport.closes)
}
select {
case <-healthy.outbound:
default:
t.Fatal("healthy PFS connection did not receive best-effort enqueue")
}
}
func TestSessionManagerChannelInterestIndex(t *testing.T) {
sm := NewSessionManager(zaptest.NewLogger(t))
raw := [8]byte{1, 2, 3}
c := &Conn{sessionID: 42, authKeyID: raw}
sm.Register(c)
sm.BindUserForAuthKey(raw, 42, 100)
sm.TrackChannelInterest(raw, 42, 100, []int64{10, 10, 20})
if got := sm.OnlineChannelUserIDs(10, 10); len(got) != 1 || got[0] != 100 {
t.Fatalf("channel 10 online users = %v, want [100]", got)
}
sm.TrackChannelInterest(raw, 42, 100, []int64{20})
if got := sm.OnlineChannelUserIDs(10, 10); len(got) != 0 {
t.Fatalf("channel 10 after viewer switch = %v, want empty", got)
}
if got := sm.OnlineChannelUserIDs(20, 10); len(got) != 1 || got[0] != 100 {
t.Fatalf("channel 20 after viewer switch = %v, want [100]", got)
}
sm.TrackChannelInterest(raw, 42, 100, []int64{10})
if got := sm.OnlineChannelMemberUserIDs(10, 10); len(got) != 0 {
t.Fatalf("channel 10 online members before membership sync = %v, want empty", got)
}
sm.SetSessionChannelMemberships(raw, 42, 100, []int64{10, 30}, sm.ChannelMembershipGeneration(raw, 42))
if got := sm.OnlineChannelMemberUserIDs(10, 10); len(got) != 1 || got[0] != 100 {
t.Fatalf("channel 10 online members = %v, want [100]", got)
}
if got := sm.OnlineChannelUserIDs(30, 10); len(got) != 0 {
t.Fatalf("channel 30 viewers = %v, want empty", got)
}
if got := sm.OnlineUserIDsForCandidates([]int64{0, 200, 100, 100}, 10); len(got) != 1 || got[0] != 100 {
t.Fatalf("candidate online users = %v, want [100]", got)
}
sm.BindUserForAuthKey(raw, 42, 200)
if got := sm.OnlineChannelUserIDs(10, 10); len(got) != 0 {
t.Fatalf("channel interest after user switch = %v, want empty", got)
}
if got := sm.OnlineChannelMemberUserIDs(10, 10); len(got) != 0 {
t.Fatalf("channel membership after user switch = %v, want empty", got)
}
sm.TrackChannelInterest(raw, 42, 200, []int64{10})
if got := sm.OnlineChannelUserIDs(10, 10); len(got) != 1 || got[0] != 200 {
t.Fatalf("channel 10 after re-track = %v, want [200]", got)
}
sm.AddUserChannelMembership(200, 10)
if got := sm.OnlineChannelMemberUserIDs(10, 10); len(got) != 1 || got[0] != 200 {
t.Fatalf("channel 10 membership after add = %v, want [200]", got)
}
sm.RemoveUserChannelMembership(200, 10)
if got := sm.OnlineChannelMemberUserIDs(10, 10); len(got) != 0 {
t.Fatalf("channel membership after remove = %v, want empty", got)
}
sm.ClearChannelInterest(raw, 42, 200)
if got := sm.OnlineChannelUserIDs(10, 10); len(got) != 0 {
t.Fatalf("channel interest after explicit clear = %v, want empty", got)
}
sm.Unregister(c)
if got := sm.OnlineChannelUserIDs(10, 10); len(got) != 0 {
t.Fatalf("channel interest after unregister = %v, want empty", got)
}
if got := sm.OnlineChannelMemberUserIDs(10, 10); len(got) != 0 {
t.Fatalf("channel membership after unregister = %v, want empty", got)
}
}
func TestSessionManagerClearsChannelIndexesOnAuthAndReadinessChanges(t *testing.T) {
sm := NewSessionManager(zaptest.NewLogger(t))
raw := [8]byte{1, 2, 3}
business := [8]byte{8}
c := &Conn{sessionID: 42, authKeyID: raw}
sm.Register(c)
sm.BindAuthKeyForSession(raw, 42, business)
sm.BindUserForAuthKey(raw, 42, 100)
track := func() {
sm.TrackChannelInterest(raw, 42, 100, []int64{10})
sm.SetSessionChannelMemberships(raw, 42, 100, []int64{10}, sm.ChannelMembershipGeneration(raw, 42))
if got := sm.OnlineChannelUserIDs(10, 10); len(got) != 1 || got[0] != 100 {
t.Fatalf("channel viewers before cleanup = %v, want [100]", got)
}
if got := sm.OnlineChannelMemberUserIDs(10, 10); len(got) != 1 || got[0] != 100 {
t.Fatalf("channel members before cleanup = %v, want [100]", got)
}
}
assertCleared := func(label string) {
if got := sm.OnlineChannelUserIDs(10, 10); len(got) != 0 {
t.Fatalf("%s viewers = %v, want empty", label, got)
}
if got := sm.OnlineChannelMemberUserIDs(10, 10); len(got) != 0 {
t.Fatalf("%s members = %v, want empty", label, got)
}
}
track()
sm.SetReceivesUpdatesForAuthKey(raw, 42, false)
assertCleared("after receivesUpdates=false")
track()
sm.BindAuthKeyForSession(raw, 42, [8]byte{9})
assertCleared("after business auth key change")
sm.BindAuthKeyForSession(raw, 42, business)
sm.BindUserForAuthKey(raw, 42, 100)
track()
if n := sm.UnbindAuthKey(business); n != 1 {
t.Fatalf("UnbindAuthKey count = %d, want 1", n)
}
assertCleared("after unbind auth key")
}
func TestPushToSessionForAuthKeyImmediateBypassesReadinessQueue(t *testing.T) {
sm := NewSessionManager(zaptest.NewLogger(t))
raw := [8]byte{1, 2, 3}
c := &Conn{
sessionID: 42,
authKeyID: raw,
outbound: make(chan outboundOp, 1),
outboundControl: make(chan outboundOp, 1),
outboundStop: make(chan struct{}),
}
sm.Register(c)
msg := &tg.UpdateShort{Update: &tg.UpdateLoginToken{}, Date: 1700000000}
if err := sm.PushToSessionForAuthKeyImmediate(context.Background(), raw, 42, proto.MessageFromServer, msg); err != nil {
t.Fatalf("immediate push: %v", err)
}
select {
case op := <-c.outbound:
defer op.releaseReservation(c.outboundTrackedBudget)
if op.encoded == nil {
t.Fatal("immediate push did not retain its encoded body")
}
var got tg.UpdateShort
if err := got.Decode(&bin.Buffer{Buf: op.encoded.body}); err != nil {
t.Fatalf("decode enqueued update: %v", err)
}
if _, ok := got.Update.(*tg.UpdateLoginToken); !ok || got.Date != msg.Date {
t.Fatalf("enqueued update = %+v, want login token date %d", got, msg.Date)
}
case <-time.After(time.Second):
t.Fatal("immediate push was not enqueued")
}
sm.mu.RLock()
pending := len(sm.pending[sessionKey{authKeyID: raw, sessionID: 42}])
sm.mu.RUnlock()
if pending != 0 {
t.Fatalf("pending pushes = %d, want 0", pending)
}
}
func TestPendingPushBodiesUseGlobalByteBudgetAndReleaseOnDrop(t *testing.T) {
sm := NewSessionManager(zaptest.NewLogger(t))
msg := &tg.UpdateShort{Update: &tg.UpdateLoginToken{}, Date: 1700000000}
encoded, err := encodeOutboundMessage(msg)
if err != nil {
t.Fatalf("encode pending fixture: %v", err)
}
sm.pendingBudget = newOutboundTrackedBudget(int64(len(encoded.body)))
key := sessionKey{authKeyID: [8]byte{9}, sessionID: 77}
sm.mu.Lock()
first := sm.queueLocked(key, proto.MessageFromServer, msg)
second := sm.queueLocked(key, proto.MessageFromServer, msg)
sm.mu.Unlock()
if !first || second {
t.Fatalf("pending queue results = first %v second %v, want true/false at byte cap", first, second)
}
if got := sm.pendingBudget.snapshot(); got != int64(len(encoded.body)) {
t.Fatalf("pending body budget = %d, want %d", got, len(encoded.body))
}
sm.mu.Lock()
sm.deletePendingLocked(key)
sm.mu.Unlock()
if got := sm.pendingBudget.snapshot(); got != 0 {
t.Fatalf("pending body budget after drop = %d, want zero", got)
}
}
func TestPendingFlushGlobalBodyPressureDoesNotTerminateHealthyConnection(t *testing.T) {
sm := NewSessionManager(zaptest.NewLogger(t))
key := sessionKey{authKeyID: [8]byte{6}, sessionID: 66}
c := &Conn{
authKeyID: key.authKeyID,
sessionID: key.sessionID,
outbound: make(chan outboundOp, 1),
outboundControl: make(chan outboundOp, 1),
outboundStop: make(chan struct{}),
metrics: NopMetrics{},
outboundTrackedBudget: newOutboundTrackedBudget(1),
}
const userID = int64(606)
c.userID.Store(userID)
c.userIDResolved.Store(true)
sm.Register(c)
msg := &tg.UpdateShort{Update: &tg.UpdateLoginToken{}, Date: 1700000000}
sm.mu.Lock()
if !sm.queueLocked(key, proto.MessageFromServer, msg) {
sm.mu.Unlock()
t.Fatal("queue pending push")
}
sm.flushing[key] = true
sm.mu.Unlock()
// Enter at the final retry so the test exercises the durable-difference fallback without
// waiting for the production backoff timer.
sm.runFlush(c, key, userID, maxFlushAttempts-1)
if c.isRetired() {
t.Fatal("shared body pressure terminated a healthy pending-flush connection")
}
if !c.receivesUpdates.Load() {
t.Fatal("pending flush did not activate difference fallback after bounded retries")
}
if got := sm.pendingBudget.snapshot(); got != 0 {
t.Fatalf("pending budget after fallback = %d, want zero", got)
}
}
func TestPendingPushBudgetSurvivesTakeAndReturnsAcrossOverflowAndUnregister(t *testing.T) {
sm := NewSessionManager(zaptest.NewLogger(t))
msg := &tg.UpdateShort{Update: &tg.UpdateLoginToken{}, Date: 1700000000}
encoded, err := encodeOutboundMessage(msg)
if err != nil {
t.Fatalf("encode pending fixture: %v", err)
}
bytesPerPush := int64(len(encoded.body))
sm.pendingBudget = newOutboundTrackedBudget(bytesPerPush * (maxPendingPushesPerSession + 8))
key := sessionKey{authKeyID: [8]byte{7}, sessionID: 55}
c := &Conn{authKeyID: key.authKeyID, sessionID: key.sessionID}
sm.Register(c)
sm.mu.Lock()
for i := 0; i < maxPendingPushesPerSession+5; i++ {
if !sm.queueLocked(key, proto.MessageFromServer, msg) {
sm.mu.Unlock()
t.Fatalf("queue pending push %d unexpectedly failed", i)
}
}
if got, want := sm.pendingBudget.snapshot(), bytesPerPush*maxPendingPushesPerSession; got != want {
sm.mu.Unlock()
t.Fatalf("budget after overflow replacement = %d, want %d", got, want)
}
batch := sm.takePendingLocked(key, true)
sm.mu.Unlock()
if len(batch) != maxPendingPushesPerSession {
t.Fatalf("taken pending pushes = %d, want %d", len(batch), maxPendingPushesPerSession)
}
// take transfers ownership to runFlush; deleting the map entry must not release bodies while
// the batch still references them.
if got, want := sm.pendingBudget.snapshot(), bytesPerPush*maxPendingPushesPerSession; got != want {
t.Fatalf("budget after take = %d, want transferred ownership %d", got, want)
}
releaseQueuedPushes(batch)
if got := sm.pendingBudget.snapshot(); got != 0 {
t.Fatalf("budget after taken batch release = %d, want 0", got)
}
sm.mu.Lock()
if !sm.queueLocked(key, proto.MessageFromServer, msg) {
sm.mu.Unlock()
t.Fatal("queue before unregister failed")
}
sm.mu.Unlock()
sm.Unregister(c)
if got := sm.pendingBudget.snapshot(); got != 0 {
t.Fatalf("budget after unregister = %d, want 0", got)
}
}
// TestSessionManagerPush 验证主动推送端到端:两个 client 连接握手并建立 session 后,
// server 经 PushToSession / PushToUser 主动向其推送client 收到。
func TestSessionManagerPush(t *testing.T) {
const dc = 2
addr, pub, srv := startTestServer(t, Options{DC: dc})
conn1, auth1, cipher1 := dialHandshake(t, addr, dc, pub)
conn2, auth2, cipher2 := dialHandshake(t, addr, dc, pub)
// 各发一个 ping 建立 session触发注册并清掉 new_session_created/pong/ack
msgGen := proto.NewMessageIDGen(time.Now)
sendEncrypted(t, conn1, cipher1, auth1, msgGen.New(proto.MessageFromClient), &mt.PingRequest{PingID: 1})
collectReplies(t, conn1, cipher1, auth1.AuthKey, mt.PongTypeID)
sendEncrypted(t, conn2, cipher2, auth2, msgGen.New(proto.MessageFromClient), &mt.PingRequest{PingID: 2})
collectReplies(t, conn2, cipher2, auth2.AuthKey, mt.PongTypeID)
if got := srv.Conns().Online(); got != 2 {
t.Fatalf("online = %d, want 2", got)
}
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
// 1) PushToSessionsession2 尚未进入 updates 同步入口时先暂存ready 后下发。
if err := srv.Conns().PushToSession(ctx, auth2.SessionID, proto.MessageFromServer, &tg.UpdatesTooLong{}); err != nil {
t.Fatalf("push to session: %v", err)
}
srv.Conns().SetReceivesUpdates(auth2.SessionID, true)
r2 := collectReplies(t, conn2, cipher2, auth2.AuthKey, tg.UpdatesTooLongTypeID)
mustHave(t, r2, tg.UpdatesTooLongTypeID, "pushed updates on conn2")
// 2) BindUser + PushToUser按 user 维度推送给 conn1。
srv.Conns().BindUser(auth1.SessionID, 100)
srv.Conns().SetReceivesUpdates(auth1.SessionID, true)
sent, err := srv.Conns().PushToUser(ctx, 100, proto.MessageFromServer, &tg.UpdatesTooLong{})
if err != nil {
t.Fatalf("push to user: %v", err)
}
if sent != 1 {
t.Fatalf("pushed to %d conns, want 1", sent)
}
r1 := collectReplies(t, conn1, cipher1, auth1.AuthKey, tg.UpdatesTooLongTypeID)
mustHave(t, r1, tg.UpdatesTooLongTypeID, "pushed updates on conn1")
// 3) PushToUserExceptSession模拟 SyncUpdatesNotMe跳过当前 session。
srv.Conns().BindUser(auth1.SessionID, 200)
srv.Conns().BindUser(auth2.SessionID, 200)
sent, err = srv.Conns().PushToUserExceptSession(ctx, 200, auth2.SessionID, proto.MessageFromServer, &tg.UpdatesTooLong{})
if err != nil {
t.Fatalf("push to user except session: %v", err)
}
if sent != 1 {
t.Fatalf("pushed to %d conns, want 1 after excluding current session", sent)
}
r1 = collectReplies(t, conn1, cipher1, auth1.AuthKey, tg.UpdatesTooLongTypeID)
mustHave(t, r1, tg.UpdatesTooLongTypeID, "pushed not-me updates on conn1")
}
func BenchmarkSessionManagerOnlineCandidateFilter(b *testing.B) {
sm := NewSessionManager(zaptest.NewLogger(b))
const online = 200_000
rawPrefix := [8]byte{9}
for i := 1; i <= online; i++ {
raw := rawPrefix
raw[1] = byte(i)
raw[2] = byte(i >> 8)
raw[3] = byte(i >> 16)
raw[4] = byte(i >> 24)
c := &Conn{sessionID: int64(i), authKeyID: raw}
sm.Register(c)
sm.BindUserForAuthKey(raw, int64(i), int64(i))
}
candidates := make([]int64, 0, 500)
for i := 0; i < 500; i++ {
candidates = append(candidates, int64(i*97+1))
}
b.ResetTimer()
for i := 0; i < b.N; i++ {
got := sm.OnlineUserIDsForCandidates(candidates, 500)
if len(got) == 0 {
b.Fatal("no candidates matched")
}
}
}