perf: sync protocol and core hardening updates

This commit is contained in:
A 2026-07-11 19:48:26 +08:00
parent 152fed3b87
commit 4390ebf5a9
283 changed files with 29231 additions and 2295 deletions

View file

@ -3,6 +3,8 @@ package mtprotoedge
import (
"context"
"errors"
"sync"
"sync/atomic"
"testing"
"time"
@ -27,6 +29,35 @@ 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")
}
@ -78,6 +109,43 @@ func TestSessionManagerRegistry(t *testing.T) {
}
}
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)
@ -88,6 +156,7 @@ func TestSessionManagerBestEffortFanoutPreencodesOnce(t *testing.T) {
outbound: make(chan outboundOp, 1),
outboundControl: make(chan outboundOp, 1),
outboundStop: make(chan struct{}),
metrics: NopMetrics{},
}
c.userID.Store(userID)
c.userIDResolved.Store(true)
@ -115,6 +184,117 @@ func TestSessionManagerBestEffortFanoutPreencodesOnce(t *testing.T) {
}
}
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.terminal.Load() {
t.Fatal("healthy session was terminalized")
}
for i, c := range slow {
if !c.terminal.Load() {
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}
@ -130,6 +310,9 @@ func TestSessionManagerScopesSameSessionIDByAuthKey(t *testing.T) {
}
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)
@ -148,7 +331,7 @@ func TestSessionManagerScopesSameSessionIDByAuthKey(t *testing.T) {
sm.BindUserForAuthKey(raw1, 42, 300)
sm.BindUserForAuthKey(raw2, 42, 300)
sent, err := sm.PushToUserExceptAuthKeySession(context.Background(), 300, perm1, 42, proto.MessageFromServer, &tg.UpdatesTooLong{})
sent, err := sm.PushToUserExceptAuthKeySession(context.Background(), 300, raw1, 42, proto.MessageFromServer, &tg.UpdatesTooLong{})
if err != nil {
t.Fatalf("push except scoped session: %v", err)
}
@ -213,6 +396,151 @@ func TestSessionManagerCloseSessionsForBusinessAuthKeyClosesBoundTempAndRaw(t *t
}
}
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.terminal.Load() {
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}
@ -238,6 +566,63 @@ func TestSessionManagerBusinessAuthKeyIndexTracksRebind(t *testing.T) {
}
}
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.terminal.Load() || !slowTwo.terminal.Load() || slowOneTransport.closes != 1 || slowTwoTransport.closes != 1 {
t.Fatalf("slow connections not terminal/closed: one=%v/%d two=%v/%d",
slowOne.terminal.Load(), slowOneTransport.closes, slowTwo.terminal.Load(), slowTwoTransport.closes)
}
if healthy.terminal.Load() || healthyTransport.closes != 0 {
t.Fatalf("healthy connection was dropped: terminal=%v closes=%d", healthy.terminal.Load(), 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}
@ -368,8 +753,16 @@ func TestPushToSessionForAuthKeyImmediateBypassesReadinessQueue(t *testing.T) {
select {
case op := <-c.outbound:
if op.msg != msg {
t.Fatalf("enqueued msg = %T, want original update", op.msg)
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")
@ -383,6 +776,126 @@ func TestPushToSessionForAuthKeyImmediateBypassesReadinessQueue(t *testing.T) {
}
}
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.terminal.Load() {
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) {