astra/owpengram-server
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

fix: sync scoped connection and outbox exclusion updates

Browse source
This commit is contained in:
A 2026-07-12 12:01:53 +08:00
parent aa21bd04e1
commit cbccd6a8d9
58 changed files with 919 additions and 1435 deletions
Download patch file Download diff file Expand all files Collapse all files

54
internal/mtprotoedge/conn.go
View file

@ -82,12 +82,9 @@ type Conn struct {
outboundControlBudgetOnce sync.Once
outboundScratchPool *outboundScratchPool
outboundScratchOnce sync.Once
// terminal 表示该 logical Conn 已停止接受新的出站操作。写失败时由
// outbound actor 置位并只发停止信号，不能在 actor 内等待自身退出。
terminal atomic.Bool
// lifecycle is a monotonic activation state machine. In particular, retired
// never transitions back to claiming/active; this closes the stale-read-loop
// ABA where an evicted Conn observed "not active" and registered itself again.
// lifecycle is the sole monotonic activation/retirement state machine.
// retired never transitions back to claiming/active; one atomic state avoids
// contradictory activation and shutdown observations.
lifecycle atomic.Uint32
transportClose sync.Once
@ -138,14 +135,6 @@ type Conn struct {
membershipGen atomic.Int64
// createdAt 是连接建立时刻，供同 auth_key session 数触顶时驱逐真正最旧的连接。
createdAt time.Time
// keyDestroyed 标记本连接的 auth_key 已被 destroy_auth_key 删除。serveConn 对已建立
// 连接复用缓存密钥跳过每帧 AuthKeyStore 回查；置位后强制回落到 Get→AuthKeyNotFound，
// 维持「destroy_auth_key 发起连接下一帧自然失效」契约。只由 destroy_auth_key 处理器置位。
keyDestroyed atomic.Bool
// lastSessionSaveUnix 是上次把本连接 session 持久化到 SessionStore 的 unix 秒，用于把
// 每帧 Save 去抖到固定间隔——session 持久化是软状态（生产无热读路径，仅观测/未来用）。
// 只由单连接的读循环 goroutine 访问。
lastSessionSaveUnix atomic.Int64
// clientLayer 是本连接协商的 TL layer（invokeWithLayer/initConnection），由 handleRPC
// 在每次 Dispatch 后从 RPC 注册表刷新。出站(rpc_result/push)按此把 227 对象降级给老客户端；
// 0 表示尚未协商，按 canonical(227) 处理=不降级。
@ -159,8 +148,29 @@ func (c *Conn) lifecycleState() connLifecycle {
return connLifecycle(c.lifecycle.Load())
}
func (c *Conn) isRetired() bool {
return c == nil || c.lifecycleState() == connLifecycleRetired
}
// retire irreversibly fences the logical connection. The caller that wins the
// transition may additionally own one-shot physical cleanup.
func (c *Conn) retire() bool {
if c == nil {
return false
}
for {
state := c.lifecycle.Load()
if connLifecycle(state) == connLifecycleRetired {
return false
}
if c.lifecycle.CompareAndSwap(state, uint32(connLifecycleRetired)) {
return true
}
}
}
func (c *Conn) beginActivationClaim() bool {
if c == nil || c.terminal.Load() || !c.isPhysicalTransportCurrentOpen() {
if c == nil || !c.isPhysicalTransportCurrentOpen() {
return false
}
if !c.lifecycle.CompareAndSwap(uint32(connLifecycleProvisional), uint32(connLifecycleClaiming)) {
@ -169,31 +179,31 @@ func (c *Conn) beginActivationClaim() bool {
// Physical close can win after the pre-check but before the lifecycle CAS.
// Do not let a doomed claimant enter SessionManager and retire a healthy old
// owner for the same logical session.
if c.terminal.Load() || !c.isPhysicalTransportCurrentOpen() {
c.lifecycle.Store(uint32(connLifecycleRetired))
if c.lifecycleState() != connLifecycleClaiming || !c.isPhysicalTransportCurrentOpen() {
c.retire()
return false
}
return true
}
func (c *Conn) publishActivation() bool {
if c == nil || c.terminal.Load() || !c.isPhysicalTransportCurrentOpen() {
if c == nil || !c.isPhysicalTransportCurrentOpen() {
return false
}
if !c.lifecycle.CompareAndSwap(uint32(connLifecycleClaiming), uint32(connLifecycleActive)) {
return false
}
// A concurrent transport failure can retire the Conn between the first
// terminal check and the CAS. Never let that intermediate active value escape.
if c.terminal.Load() || !c.isPhysicalTransportCurrentOpen() {
c.lifecycle.Store(uint32(connLifecycleRetired))
// physical check and the CAS. Never let that intermediate active value escape.
if c.lifecycleState() != connLifecycleActive || !c.isPhysicalTransportCurrentOpen() {
c.retire()
return false
}
return true
}
func (c *Conn) isActive() bool {
return c != nil && !c.terminal.Load() && c.lifecycleState() == connLifecycleActive
return c != nil && c.lifecycleState() == connLifecycleActive && c.isPhysicalTransportCurrentOpen()
}
// transferTransportOwnership hands this Conn's physical socket to the next

180
internal/mtprotoedge/encrypted.go
View file

@ -10,8 +10,6 @@ import (
"fmt"
"io"
"math"
"sync"
"sync/atomic"
"time"
"go.uber.org/zap"
@ -231,7 +229,7 @@ func (s *Server) handleEncrypted(ctx context.Context, tc transport.Conn, cs *con
}
return current, errActivationAuthKeyRejected
}
if current.terminal.Load() || !current.isPhysicalTransportCurrentOpen() {
if current.isRetired() || !current.isPhysicalTransportCurrentOpen() {
return current, ErrConnClosed
}
}
@ -258,7 +256,6 @@ func (s *Server) handleEncrypted(ctx context.Context, tc transport.Conn, cs *con
cs.createdFloor = plan.logicalMin
}
plan.commitState(cs)
s.maybePersistSession(ctx, current, frame.sessionID, key.ID, serverSalt)
if err := s.executeInboundPlan(ctx, cs, current, plan); err != nil {
return current, err
@ -280,34 +277,6 @@ func (s *Server) handleEncrypted(ctx context.Context, tc transport.Conn, cs *con
return current, nil
}
// sessionSaveMinInterval 是单连接持久化 session 记录的最小间隔。把原本「每帧一次 Redis SET」
// 去抖到固定间隔——session 是软状态（生产无热读路径），只需周期刷新 last_seen/续 TTL。
const sessionSaveMinInterval = 30 * time.Second
// maybePersistSession 按 sessionSaveMinInterval 去抖持久化 session，失败只告警不断连。
// 原实现每帧同步 Save 且失败即断连：N 连接×帧率的 Redis 写放大 + Redis 抖动级联断连。
func (s *Server) maybePersistSession(ctx context.Context, c *Conn, sessionID int64, authKeyID [8]byte, salt int64) {
if c == nil {
return
}
now := s.clock.Now().Unix()
if last := c.lastSessionSaveUnix.Load(); last != 0 && now-last < int64(sessionSaveMinInterval/time.Second) {
return
}
c.lastSessionSaveUnix.Store(now)
if err := s.sessions.Save(ctx, store.SessionData{
ID: sessionID,
AuthKeyID: authKeyID,
Salt: salt,
LastSeen: now,
}); err != nil {
s.log.Warn("Persist session failed (non-fatal)",
zap.Int64("session_id", sessionID),
zap.Error(err),
)
}
}
func sendQuickAckIfRequested(ctx context.Context, tc transport.Conn, key crypto.AuthKey, plaintext []byte, writeTimeout time.Duration) error {
q, ok := tc.(quickAckTransport)
if !ok || !q.ConsumeQuickAckRequested() {
@ -343,23 +312,6 @@ func clientQuickAckToken(key crypto.AuthKey, plaintext []byte) uint32 {
return binary.LittleEndian.Uint32(sum[:4]) &^ quickAckResponseFlag
}
// dispatch 处理一条明文消息：解包 container/gzip，处理服务消息，其余转 RPC 路由。
// content-related 消息（ping、RPC）的 msg_id 会收集到 acks 以便统一确认。
func (s *Server) dispatch(ctx context.Context, cs *connState, c *Conn, msgID int64, seqNo int32, b *bin.Buffer, acks *[]int64) error {
plan, err := s.preflightInbound(cs, msgID, seqNo, b.Buf)
if err != nil {
var bad *dispatchBadMsgError
if errors.As(err, &bad) && c != nil {
return s.sendBadMsg(ctx, c, bad.msgID, bad.seqNo, bad.code)
}
return err
}
defer plan.close()
plan.commitState(cs)
*acks = append(*acks, plan.ackIDs...)
return s.executeInboundPlan(ctx, cs, c, plan)
}
// dispatchBadMsgError carries a protocol-level rejection discovered during the
// side-effect-free wrapper/container preflight. The caller emits the single
// bad_msg_notification only after the whole container has been inspected.
@ -625,8 +577,8 @@ func mergeStateInfo(primary, fallback []byte) []byte {
return info
}
// enqueueRPC 把一条 RPC 请求交给连接的 inbound 调度器。typeID 由 dispatch 传入
// （已 PeekID 过一次），method 只解析一次并随任务透传，避免同一请求三处重复 PeekID/typeName。
// enqueueRPC 重试一个旧 owner 未发布结果的请求。正常收包统一走 container batch；
// 这里也用长度为 1 的 batch，避免维护第二套预算/commit 状态机。
func (s *Server) enqueueRPC(ctx context.Context, c *Conn, msgID int64, typeID uint32, request *bin.Buffer) error {
method := s.typeName(typeID)
claim, err := s.rpcResults.Acquire(c.authKeyID, c.sessionID, msgID)
@ -666,13 +618,13 @@ func (s *Server) enqueueRPC(ctx context.Context, c *Conn, msgID int64, typeID ui
}()
// 两级条数/字节预算必须先于 Copy：对抗客户端不能用大量满尺寸请求在“判断队列满”
// 之前制造一轮无上限的临时 body 分配。reservation 在 commit/abort 间唯一持有预算。
reservation, err := c.reserveInboundRPC(ctx, method, request.Len())
reservation, err := c.reserveInboundRPCBatch(ctx, []inboundRPCSpec{{method: method, size: request.Len()}})
if err != nil {
return s.handleInboundRPCAdmissionError(ctx, c, msgID, method, err)
}
defer reservation.abort()
body := request.Copy()
err = reservation.commit(s.newInboundRPCTask(c, msgID, method, body, owner))
err = reservation.commit([]inboundRPC{s.newInboundRPCTask(c, msgID, method, body, owner)})
transferred = err == nil
return s.handleInboundRPCAdmissionError(ctx, c, msgID, method, err)
}
@ -681,14 +633,10 @@ func (s *Server) enqueueRPC(ctx context.Context, c *Conn, msgID int64, typeID ui
// single-message and atomic container-batch admission paths. body must already
// be an independently owned, budgeted copy.
func (s *Server) newInboundRPCTask(c *Conn, msgID int64, method string, body []byte, owner *rpcResultOwnerLease) inboundRPC {
responseGate := newRPCResponseGate()
timeoutResponse := func() {
if !responseGate.tryTimeout() {
return
}
defer responseGate.finish()
// 原 task context 已到期，使用有界的新 context 回显明确的可重试超时；
// 500 保持 TDesktop 默认重试语义，错误名区分于容量型 FLOOD_WAIT。
// 只有尚未进入 handler 的排队请求会走这里。运行中的请求只取消
// context，等 handler 收敛后再决定成功或 RPC_TIMEOUT，避免客户端用
// 新 msg_id 重试时与旧业务提交并发。
writeTimeout := c.writeTimeout
if writeTimeout <= 0 || writeTimeout > 5*time.Second {
writeTimeout = 5 * time.Second
@ -716,13 +664,6 @@ func (s *Server) newInboundRPCTask(c *Conn, msgID int64, method string, body []b
if owner == nil {
return
}
// A running deadline callback owns the same terminal flight but executes
// on context.AfterFunc's goroutine. Let its bounded delivery complete before
// deciding that the task produced no result; otherwise release could Abort
// the owner while RPC_TIMEOUT is already inside the physical writer.
if !responseGate.wait(requiredControlMaxWait + time.Second) {
c.fenceUndeliveredRPCResult()
}
if owner.Abort() {
// connState already remembers this request. If a committed task exits
// without publishing any terminal rpc_result, a same-Conn retransmit
@ -734,7 +675,7 @@ func (s *Server) newInboundRPCTask(c *Conn, msgID int64, method string, body []b
run: func(taskCtx context.Context) error {
// body 是预算成功后生成的独立副本，且每个任务只 run 一次，
// 无需再 append 拷贝；直接复用，省掉一份 inbound 在途内存。
if err := s.handleRPC(taskCtx, c, msgID, method, &bin.Buffer{Buf: body}, responseGate); err != nil {
if err := s.handleRPC(taskCtx, c, msgID, method, &bin.Buffer{Buf: body}); err != nil {
fields := []zap.Field{
zap.Int64("msg_id", msgID),
zap.String("auth_key_id", c.authKeyHex),
@ -770,15 +711,11 @@ func (s *Server) handleInboundRPCAdmissionError(ctx context.Context, c *Conn, ms
}
// handleRPC 把明文 RPC 请求交给 RPC 路由，并将结果或错误包成 rpc_result 回发。
func (s *Server) handleRPC(ctx context.Context, c *Conn, msgID int64, method string, b *bin.Buffer, responseGate *rpcResponseGate) error {
func (s *Server) handleRPC(ctx context.Context, c *Conn, msgID int64, method string, b *bin.Buffer) error {
if s.rpc == nil {
if err := ctx.Err(); err != nil {
return err
}
if responseGate != nil && !responseGate.tryNormal() {
return context.DeadlineExceeded
}
defer responseGate.finish()
s.log.Warn("No RPC handler configured", zap.String("method", method))
return s.sendResult(ctx, c, msgID, &mt.RPCError{
ErrorCode: 500,
@ -817,29 +754,29 @@ func (s *Server) handleRPC(ctx context.Context, c *Conn, msgID int64, method str
fields = dbtrace.AppendZapFields(fields, "", dbStats.Snapshot())
if ctxErr := ctx.Err(); ctxErr != nil {
// The old physical Conn may have been replaced after the business transaction
// committed. Never write with the expired context itself: a fenced generation
// publishes cache-only for its replacement, while a still-live generation uses
// a fresh bounded delivery context. A deadline callback that already won
// responseGate has published RPC_TIMEOUT and prevents this late result from
// overwriting it.
// Only a successful business result proves useful work completed. Errors
// observed after cancellation may themselves be cancellation-derived or
// transient and must remain retryable rather than poisoning the replay cache.
// A running request owns its terminal response until Dispatch returns. If
// useful work completed despite cancellation, preserve that success; otherwise
// a deadline becomes RPC_TIMEOUT only now, after the handler has converged.
// Plain connection cancellation remains retryable on the replacement.
var terminal bin.Encoder
runPostResponse := false
if err == nil && result != nil {
terminal = result
runPostResponse = true
} else if errors.Is(ctxErr, context.DeadlineExceeded) {
terminal = &mt.RPCError{ErrorCode: 500, ErrorMessage: "RPC_TIMEOUT"}
}
if terminal != nil && (responseGate == nil || responseGate.tryNormal()) {
defer responseGate.finish()
if c.terminal.Load() || !c.isPhysicalTransportCurrentOpen() {
if terminal != nil {
if c.isRetired() || !c.isPhysicalTransportCurrentOpen() {
// Replacement/shutdown already fenced this logical generation. Cache-only
// publication is safe and lets the replacement join the completed flight.
if encoded, encodeErr := s.encodeRPCResult(c, msgID, terminal); encodeErr != nil {
s.log.Warn("Encode canceled RPC result for replay failed", append(fields, zap.Error(encodeErr))...)
} else {
s.storeRPCResult(c, msgID, encoded)
postresponse.Run(context.WithoutCancel(ctx))
if runPostResponse {
postresponse.Run(context.WithoutCancel(ctx))
}
}
} else {
// An individual RPC deadline can expire while the physical connection is
@ -854,7 +791,7 @@ func (s *Server) handleRPC(ctx context.Context, c *Conn, msgID int64, method str
cancel()
if sendErr != nil {
s.log.Debug("Send canceled RPC result failed", append(fields, zap.Error(sendErr))...)
} else {
} else if runPostResponse {
postresponse.Run(context.WithoutCancel(ctx))
}
}
@ -866,13 +803,6 @@ func (s *Server) handleRPC(ctx context.Context, c *Conn, msgID int64, method str
s.log.Info("RPC canceled", cancelFields...)
return ctxErr
}
// A deadline callback may have already emitted RPC_TIMEOUT while Dispatch was returning.
// Claim the single normal-response slot before serializing any success/error rpc_result.
if responseGate != nil && !responseGate.tryNormal() {
s.log.Info("RPC result suppressed after timeout", fields...)
return context.DeadlineExceeded
}
defer responseGate.finish()
if err != nil {
var rpcErr *tgerr.Error
@ -898,57 +828,6 @@ func (s *Server) handleRPC(ctx context.Context, c *Conn, msgID int64, method str
return nil
}
// rpcResponseGate guarantees exactly one terminal rpc_result per request. A running deadline
// races legitimately with a handler completing at the boundary; whichever path claims state
// first owns the response, and the other path becomes a no-op.
type rpcResponseGate struct {
state atomic.Uint32
done chan struct{}
once sync.Once
}
func newRPCResponseGate() *rpcResponseGate {
return &rpcResponseGate{done: make(chan struct{})}
}
func (g *rpcResponseGate) tryNormal() bool {
return g == nil || g.state.CompareAndSwap(0, 1)
}
func (g *rpcResponseGate) tryTimeout() bool {
return g != nil && g.state.CompareAndSwap(0, 2)
}
func (g *rpcResponseGate) finish() {
if g == nil || g.done == nil {
return
}
g.once.Do(func() { close(g.done) })
}
// wait returns once there is no response owner, the winning response has
// completed deliver-or-fence publication, or the bounded safety deadline wins.
func (g *rpcResponseGate) wait(timeout time.Duration) bool {
if g == nil || g.state.Load() == 0 {
return true
}
if g.done == nil {
return false
}
if timeout <= 0 {
<-g.done
return true
}
timer := time.NewTimer(timeout)
defer timer.Stop()
select {
case <-g.done:
return true
case <-timer.C:
return false
}
}
// sendResult 把 RPC 结果包成 rpc_result 并加密回发。
func (s *Server) sendResult(ctx context.Context, c *Conn, reqMsgID int64, result bin.Encoder) error {
if result == nil {
@ -1135,13 +1014,6 @@ func (s *Server) sendDestroySession(ctx context.Context, c *Conn, sessionID int6
removed := false
if sessionID != c.sessionID {
removed = s.conns.DestroySessionForAuthKey(c.authKeyID, sessionID)
if err := s.sessions.Delete(ctx, sessionID); err != nil {
s.log.Debug("Delete session record failed",
zap.String("auth_key_id", c.authKeyHex),
zap.Int64("session_id", sessionID),
zap.Error(err),
)
}
}
if removed {
return c.Send(ctx, proto.MessageServerResponse, &mt.DestroySessionOk{SessionID: sessionID})
@ -1277,10 +1149,6 @@ func (cs *connState) validateSeq(msgID int64, seqNo int32, content bool) int {
return 0
}
func (cs *connState) track(msgID int64, seqNo int32, content bool, state byte) {
cs.trackInbound(msgID, seqNo, content, false, state)
}
func (cs *connState) trackInbound(msgID int64, seqNo int32, content, service bool, state byte) {
cs.seen[msgID] = clientMsgRecord{
state: state,

12
internal/mtprotoedge/encrypted_test.go
View file

@ -1,6 +1,7 @@
package mtprotoedge
import (
"context"
"testing"
"time"
@ -414,7 +415,7 @@ func TestPingDelayDisconnectOddSeqAccepted(t *testing.T) {
// 避免 TDesktop 清理旧 key 时落到业务 RPC fallback。
func TestDestroyAuthKey(t *testing.T) {
const dc = 2
addr, pub, _ := startTestServer(t, Options{DC: dc})
addr, pub, srv := startTestServer(t, Options{DC: dc})
conn, auth, cipher := dialHandshake(t, addr, dc, pub)
clientMsgID := proto.NewMessageIDGen(time.Now)
@ -423,6 +424,15 @@ func TestDestroyAuthKey(t *testing.T) {
replies := collectReplies(t, conn, cipher, auth.AuthKey, destroyAuthKeyOkTypeID)
mustHave(t, replies, destroyAuthKeyOkTypeID, "destroy_auth_key_ok")
if _, found, err := srv.authKeys.Get(context.Background(), auth.AuthKey.ID); err != nil || found {
t.Fatalf("auth key after destroy: found=%v err=%v", found, err)
}
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
defer cancel()
var frame bin.Buffer
if err := conn.Recv(ctx, &frame); err == nil {
t.Fatal("destroy_auth_key requester remained readable after required ok")
}
}
// TestBadServerSalt 验证客户端带错 server_salt 时 server 返回 bad_server_salt，

24
internal/mtprotoedge/inbound_preflight.go
View file

@ -68,7 +68,6 @@ type inboundPlan struct {
logicalMin int64
releases []func()
rpcPrepared bool
rpcReservation *inboundRPCBatchReservation
rpcTasks []inboundRPC
rpcOwners []*rpcResultOwnerLease
@ -101,7 +100,7 @@ func (p *inboundPlan) commitRPCBatch() error {
// batch is runnable immediately; using the old deferred scheduler token here
// was not a real barrier on a busy Conn because an existing ready token could
// dequeue newly appended tasks before activateRPCBatch ran.
_, err := p.rpcReservation.commit(p.rpcTasks, false)
err := p.rpcReservation.commit(p.rpcTasks)
if err != nil {
return err
}
@ -619,7 +618,6 @@ func preflightInboundItem(msgID int64, seqNo int32, typeID uint32, content bool,
// into one consistent terminal FLOOD_WAIT result per uncached RPC; no business
// handler from the batch is allowed to start in that case.
func (s *Server) prepareInboundRPCBatch(ctx context.Context, c *Conn, plan *inboundPlan) error {
plan.rpcPrepared = true
// Keep service-only frames (ping/ack/http_wait) allocation-free here. These
// collections are needed only after the first real API RPC acquires ownership.
var indices []int
@ -815,18 +813,22 @@ func (s *Server) executeInboundPlan(ctx context.Context, cs *connState, c *Conn,
s.log.Debug("Received destroy_auth_key", zap.String("auth_key_id", c.authKeyHex))
if err := s.authKeys.Delete(ctx, c.authKeyID); err != nil {
s.log.Warn("Delete auth key failed", zap.String("auth_key_id", c.authKeyHex), zap.Error(err))
return c.SendAsync(ctx, proto.MessageServerResponse, &destroyAuthKeyFail{})
return c.SendRequiredControl(ctx, proto.MessageServerResponse, &destroyAuthKeyFail{})
}
c.keyDestroyed.Store(true)
// Fence every other active/claiming generation before acknowledging the
// deletion. The exact requester remains writable only long enough to put the
// required destroy_auth_key_ok frame on the wire.
s.conns.CloseSessionsForRawAuthKeyExceptConn(c.authKeyID, c)
return c.SendAsync(ctx, proto.MessageServerResponse, &destroyAuthKeyOk{})
case inboundItemRPC:
if plan.rpcPrepared {
continue
}
if err := s.enqueueRPC(ctx, c, item.msgID, item.typeID, &bin.Buffer{Buf: item.body}); err != nil {
if err := c.SendRequiredControl(ctx, proto.MessageServerResponse, &destroyAuthKeyOk{}); err != nil {
return err
}
c.beginTerminalShutdown()
c.closeTransport()
return nil
case inboundItemRPC:
// prepareInboundRPCBatch owns every fresh RPC before synchronous service
// execution begins; commitRPCBatch publishes them after all protocol barriers.
continue
case inboundItemCapacityError:
if err := s.sendResult(ctx, c, item.msgID, &mt.RPCError{
ErrorCode: 420,

315
internal/mtprotoedge/inbound_rpc.go
View file

@ -39,14 +39,7 @@ type inboundRPC struct {
ticket *inboundRPCTicket
}
const (
inboundRPCTicketQueued int32 = iota
inboundRPCTicketRunning
inboundRPCTicketDone
)
type inboundRPCTicket struct {
state atomic.Int32
onTimeout func()
}
@ -84,19 +77,6 @@ type inboundRPCGlobalReservation struct {
released atomic.Bool
}
// inboundRPCReservation 同时持有全局和单连接的“Copy 前”预算。commit/abort 只能成功一次；
// 无论 Copy 后连接关闭、入队成功还是调用方提前返回，预算都有唯一归还路径。
type inboundRPCReservation struct {
conn *Conn
global *inboundRPCGlobalReservation
ctx context.Context
method string
size int
enqueuedAt time.Time
deadline time.Time
once sync.Once
}
// inboundRPCSpec 是 container preflight 与 RPC scheduler 之间的有界 admission 描述。
// method 仅用于 metrics，size 是在 Copy 之前必须预留的 request body 字节数。
type inboundRPCSpec struct {
@ -202,31 +182,6 @@ func (s *inboundRPCScheduler) stop(timeout time.Duration) {
}
}
func (s *inboundRPCScheduler) reserveGlobal(size int) (*inboundRPCGlobalReservation, string, error) {
if size < 0 {
size = 0
}
size64 := int64(size)
s.budgetMu.Lock()
defer s.budgetMu.Unlock()
select {
case <-s.stopCh:
return nil, "scheduler_closed", ErrConnClosed
default:
}
if s.tasks >= s.maxTasks {
return nil, "global_task_budget", ErrInboundRPCQueueFull
}
// 用减法比较避免 s.bytes+size64 溢出。
if size64 > s.maxBytes-s.bytes {
return nil, "global_byte_budget", ErrInboundRPCQueueFull
}
s.tasks++
s.bytes += size64
return &inboundRPCGlobalReservation{scheduler: s, size: size64}, "", nil
}
// reserveGlobalBatch 在一次 budgetMu 临界区内检查并预留整批条数/字节。
// 返回的每个 reservation 仍由对应 task 单独归还，避免一个慢 RPC 持有
// 整个 container 已完成任务的预算。
@ -312,12 +267,6 @@ func releaseInboundRPCGlobalBatch(reservations []*inboundRPCGlobalReservation) {
scheduler.budgetMu.Unlock()
}
func (s *inboundRPCScheduler) budgetSnapshot() (tasks int, bytes int64) {
s.budgetMu.Lock()
defer s.budgetMu.Unlock()
return s.tasks, s.bytes
}
func (s *inboundRPCScheduler) schedule(c *Conn) {
if s == nil || c == nil {
return
@ -454,86 +403,6 @@ func (c *Conn) startInboundRPCScheduler(scheduler *inboundRPCScheduler, maxInfli
// rpcQueue 保持 nil；首个成功 commit 才由 append 分配，静默连接零队列内存。
}
// reserveInboundRPC 必须在 request body Copy 前调用。它先拿进程级条数/字节预算，
// 再预占单连接队列槽和字节预算；commit 或 abort 负责唯一释放。
func (c *Conn) reserveInboundRPC(ctx context.Context, method string, size int) (*inboundRPCReservation, error) {
if ctx == nil {
ctx = context.Background()
}
select {
case <-ctx.Done():
c.metrics.InboundRPCDropped(method, "context_done")
return nil, ctx.Err()
default:
}
if c.terminal.Load() {
c.metrics.InboundRPCDropped(method, "scheduler_closed")
return nil, ErrConnClosed
}
if c.rpcScheduler == nil {
c.metrics.InboundRPCDropped(method, "scheduler_closed")
return nil, ErrConnClosed
}
global, reason, err := c.rpcScheduler.reserveGlobal(size)
if err != nil {
c.metrics.InboundRPCDropped(method, reason)
return nil, err
}
now := time.Now()
deadline := time.Time{}
if c.rpcTimeout > 0 {
deadline = now.Add(c.rpcTimeout)
}
if ctxDeadline, ok := ctx.Deadline(); ok && (deadline.IsZero() || ctxDeadline.Before(deadline)) {
deadline = ctxDeadline
}
if size < 0 {
size = 0
}
c.rpcMu.Lock()
if err := ctx.Err(); err != nil {
c.rpcMu.Unlock()
global.release()
c.metrics.InboundRPCDropped(method, "context_done")
return nil, err
}
if c.rpcClosed || c.terminal.Load() {
c.rpcMu.Unlock()
global.release()
c.metrics.InboundRPCDropped(method, "scheduler_closed")
return nil, ErrConnClosed
}
if c.rpcReserved+len(c.rpcQueue) >= c.rpcQueueSize {
c.rpcMu.Unlock()
global.release()
c.metrics.InboundRPCDropped(method, "queue_full")
return nil, ErrInboundRPCQueueFull
}
if int64(size) > maxInflightRPCBytes-c.inflightRPCBytes.Load() {
c.rpcMu.Unlock()
global.release()
c.metrics.InboundRPCDropped(method, "byte_budget")
return nil, ErrInboundRPCQueueFull
}
c.rpcReserved++
c.inflightRPCBytes.Add(int64(size))
// Add 与 close 的 Wait 由 rpcMu 排序：close 置 rpcClosed 后不会再发生 Add。
c.rpcReservationWG.Add(1)
c.rpcMu.Unlock()
return &inboundRPCReservation{
conn: c,
global: global,
ctx: ctx,
method: method,
size: size,
enqueuedAt: now,
deadline: deadline,
}, nil
}
// reserveInboundRPCBatch 必须在 container 内任何 request body Copy 前调用。
// 全局预算只锁一次，单连接预算也只锁一次；任一限制不满足时
// 整批失败，不会留下部分 task/字节 reservation。
@ -547,7 +416,7 @@ func (c *Conn) reserveInboundRPCBatch(ctx context.Context, specs []inboundRPCSpe
return nil, ctx.Err()
default:
}
if c.terminal.Load() {
if c.isRetired() {
c.dropInboundRPCSpecs(specs, "scheduler_closed")
return nil, ErrConnClosed
}
@ -589,7 +458,7 @@ func (c *Conn) reserveInboundRPCBatch(ctx context.Context, specs []inboundRPCSpe
c.dropInboundRPCSpecs(normalized, "context_done")
return nil, err
}
if c.rpcClosed || c.terminal.Load() {
if c.rpcClosed || c.isRetired() {
c.rpcMu.Unlock()
releaseInboundRPCGlobalBatch(globals)
c.dropInboundRPCSpecs(normalized, "scheduler_closed")
@ -630,113 +499,12 @@ func (c *Conn) dropInboundRPCSpecs(specs []inboundRPCSpec, reason string) {
}
}
// enqueueInboundRPC 是测试和已持有独立 body 的便捷入口。生产收包路径使用
// reserveInboundRPC -> Copy -> commit，保证真正的 Copy 前预算。
func (c *Conn) enqueueInboundRPC(ctx context.Context, task inboundRPC) error {
reservation, err := c.reserveInboundRPC(ctx, task.method, task.size)
if err != nil {
return err
}
defer reservation.abort()
return reservation.commit(task)
}
func (r *inboundRPCReservation) commit(task inboundRPC) error {
result := ErrConnClosed
var (
committed bool
reschedule bool
queueLen int
queueCap int
)
r.once.Do(func() {
c := r.conn
c.rpcMu.Lock()
c.rpcReserved--
if c.rpcClosed || c.terminal.Load() {
c.inflightRPCBytes.Add(-int64(r.size))
} else {
// The request deadline starts when admission succeeds, not when a worker
// eventually dequeues the request. This bounds total queue + execution
// latency and lets a queued request emit its explicit timeout on time.
if r.deadline.IsZero() {
task.ctx, task.cancel = context.WithCancel(r.ctx)
} else {
task.ctx, task.cancel = context.WithDeadline(r.ctx, r.deadline)
}
task.stopRoot = context.AfterFunc(c.rpcRootCtx, task.cancel)
task.method = r.method
task.enqueuedAt = r.enqueuedAt
task.deadline = r.deadline
task.size = r.size
task.budget = r.global
ticket := &inboundRPCTicket{}
if task.onTimeout != nil {
onTimeout := task.onTimeout
var timeoutOnce sync.Once
ticket.onTimeout = func() {
timeoutOnce.Do(onTimeout)
}
task.onTimeout = ticket.onTimeout
}
task.ticket = ticket
if task.onTimeout != nil && !task.deadline.IsZero() {
taskCtx := task.ctx
task.stopTimeout = context.AfterFunc(taskCtx, func() {
if errors.Is(taskCtx.Err(), context.DeadlineExceeded) {
c.expireInboundRPCTicket(ticket)
}
})
}
c.rpcQueue = append(c.rpcQueue, task)
queueLen = len(c.rpcQueue)
queueCap = c.rpcQueueSize
if c.rpcRunning < c.rpcMaxInflight && !c.rpcReady {
c.rpcReady = true
reschedule = true
}
committed = true
result = nil
}
c.rpcMu.Unlock()
c.rpcReservationWG.Done()
if !committed {
r.global.release()
}
})
if committed {
r.conn.metrics.InboundRPCQueued(r.method, queueLen, queueCap)
if reschedule {
r.conn.rpcScheduler.schedule(r.conn)
}
}
return result
}
func (r *inboundRPCReservation) abort() {
// commit 在一次 rpcMu 临界区内把整批 task append 到队列并立即发布 ready token。
// 协议 barrier 必须在调用 commit 前完成；延迟发布 token 无法阻止已有 worker
// 从同一连接队列取走新任务，因此不提供虚假的 deferred-schedule 模式。
func (r *inboundRPCBatchReservation) commit(tasks []inboundRPC) (result error) {
if r == nil {
return
}
r.once.Do(func() {
c := r.conn
c.rpcMu.Lock()
c.rpcReserved--
c.inflightRPCBytes.Add(-int64(r.size))
c.rpcMu.Unlock()
c.rpcReservationWG.Done()
r.global.release()
})
}
// commit 在一次 rpcMu 临界区内把整批 task append 到队列。
// deferSchedule=false 保持旧的立即调度语义；true 则返回一个幂等 activate
// 函数，让调用方先完成 new_session_created 等协议 barrier 再启动 worker。
//
// 延迟调度只能延迟本次 commit 新产生的 ready token；调用方应在连接的
// 首个 admission batch 使用它，不得把它当作已有 worker 的全局暂停锁。
func (r *inboundRPCBatchReservation) commit(tasks []inboundRPC, deferSchedule bool) (activate func(), result error) {
if r == nil {
return nil, ErrConnClosed
return ErrConnClosed
}
result = ErrConnClosed
var (
@ -769,7 +537,7 @@ func (r *inboundRPCBatchReservation) commit(tasks []inboundRPC, deferSchedule bo
if len(tasks) != len(r.entries) {
c.inflightRPCBytes.Add(-r.totalSize)
result = errInboundRPCBatchTaskCount
} else if c.rpcClosed || c.terminal.Load() {
} else if c.rpcClosed || c.isRetired() {
c.inflightRPCBytes.Add(-r.totalSize)
} else {
prepared := make([]inboundRPC, len(tasks))
@ -831,19 +599,10 @@ func (r *inboundRPCBatchReservation) commit(tasks []inboundRPC, deferSchedule bo
r.conn.metrics.InboundRPCQueued(entry.method, firstQueueLen+i, queueCap)
}
if reschedule {
var once sync.Once
activate = func() {
once.Do(func() {
r.conn.rpcScheduler.schedule(r.conn)
})
}
if !deferSchedule {
activate()
activate = nil
}
r.conn.rpcScheduler.schedule(r.conn)
}
}
return activate, result
return result
}
func (r *inboundRPCBatchReservation) abort() {
@ -884,9 +643,6 @@ func (c *Conn) takeInboundRPC() (task inboundRPC, ok, reschedule bool) {
c.rpcQueue = nil
}
c.rpcRunning++
if task.ticket != nil {
task.ticket.state.Store(inboundRPCTicketRunning)
}
c.rpcWG.Add(1)
if len(c.rpcQueue) > 0 && c.rpcRunning < c.rpcMaxInflight {
c.rpcReady = true
@ -920,10 +676,7 @@ func (c *Conn) runInboundRPC(task inboundRPC) {
}
func (c *Conn) finishInboundRPC(task inboundRPC) {
if task.ticket != nil {
task.ticket.state.Store(inboundRPCTicketDone)
}
timeoutHandoff := stopInboundRPCTask(task)
stopInboundRPCTask(task)
var reschedule bool
c.rpcMu.Lock()
c.rpcRunning--
@ -940,12 +693,6 @@ func (c *Conn) finishInboundRPC(task inboundRPC) {
// with a newly admitted body under the same byte accounting.
task = inboundRPC{}
reservation.release()
if timeoutHandoff != nil {
// stopTimeout(false) means the deadline callback may already have read
// Running but not yet entered ticket.onTimeout. Calling the sync.Once wrapper
// here either performs or joins that response before owner release/Abort.
timeoutHandoff()
}
if release != nil {
release()
}
@ -957,8 +704,8 @@ func (c *Conn) finishInboundRPC(task inboundRPC) {
// expireInboundRPCTicket removes a request that is still queued and returns its
// memory/task reservations immediately. If the worker won the dequeue race, the
// same callback only signals the running request's response gate; its body remains
// owned until the handler exits.
// callback does nothing: the running handler owns the only terminal response and
// its deadline is represented solely by context cancellation.
func (c *Conn) expireInboundRPCTicket(ticket *inboundRPCTicket) {
if ticket == nil {
return
@ -986,7 +733,6 @@ func (c *Conn) expireInboundRPCTicket(ticket *inboundRPCTicket) {
}
}
c.inflightRPCBytes.Add(-int64(task.size))
ticket.state.Store(inboundRPCTicketDone)
found = true
break
}
@ -999,7 +745,7 @@ func (c *Conn) expireInboundRPCTicket(ticket *inboundRPCTicket) {
method := task.method
reservation := task.budget
release := task.release
_ = stopInboundRPCTask(task)
stopInboundRPCTask(task)
// Drop the run/context closures before returning the byte reservation. Otherwise an
// onTimeout callback that blocks or performs a slow write can keep the copied request body
// reachable after the global scheduler has advertised those bytes as available again.
@ -1014,24 +760,14 @@ func (c *Conn) expireInboundRPCTicket(ticket *inboundRPCTicket) {
}
return
}
if ticket.state.Load() == inboundRPCTicketRunning && ticket.onTimeout != nil {
ticket.onTimeout()
}
}
// stopInboundRPCTask disarms callbacks before canceling the context so a normal
// completion or connection close cannot manufacture an RPC_TIMEOUT response.
// If the runtime already started a deadline callback, the returned sync.Once
// wrapper is a mandatory handoff: callers invoke it before owner release so a
// callback paused between ticket-state inspection and response-gate claim cannot
// publish into a later flight generation.
func stopInboundRPCTask(task inboundRPC) (timeoutHandoff func()) {
// stopInboundRPCTask disarms queue-expiration cleanup before canceling the
// context. Once a worker dequeues the task, the deadline only cancels the
// handler; it never races the handler with an early RPC_TIMEOUT response.
func stopInboundRPCTask(task inboundRPC) {
if task.stopTimeout != nil {
stopped := task.stopTimeout()
if !stopped && task.ctx != nil && errors.Is(task.ctx.Err(), context.DeadlineExceeded) &&
task.ticket != nil && task.ticket.onTimeout != nil {
timeoutHandoff = task.ticket.onTimeout
}
task.stopTimeout()
}
if task.stopRoot != nil {
task.stopRoot()
@ -1039,7 +775,6 @@ func stopInboundRPCTask(task inboundRPC) (timeoutHandoff func()) {
if task.cancel != nil {
task.cancel()
}
return timeoutHandoff
}
func (c *Conn) closeInboundRPCScheduler() {
@ -1051,8 +786,8 @@ func (c *Conn) closeInboundRPCScheduler() {
}
// beginCloseInboundRPCScheduler publishes closure, cancels running work and releases queued
// requests without waiting for handlers. ForceClose uses this phase before transport.Close so a
// pathological/blocking transport implementation cannot leave the RPC admission gate open.
// requests without waiting for handlers. Shutdown publishes this phase before transport.Close so
// a pathological/blocking transport implementation cannot leave the RPC admission gate open.
func (c *Conn) beginCloseInboundRPCScheduler() {
if c.rpcScheduler == nil {
return
@ -1078,18 +813,12 @@ func (c *Conn) beginCloseInboundRPCScheduler() {
for i := range queued {
task := queued[i]
queued[i] = inboundRPC{}
if task.ticket != nil {
task.ticket.state.Store(inboundRPCTicketDone)
}
method := task.method
reservation := task.budget
release := task.release
timeoutHandoff := stopInboundRPCTask(task)
stopInboundRPCTask(task)
task = inboundRPC{}
reservation.release()
if timeoutHandoff != nil {
timeoutHandoff()
}
if release != nil {
release()
}

26
internal/mtprotoedge/inbound_rpc_batch_test.go
View file

@ -130,9 +130,8 @@ func TestInboundRPCBatchAbortReturnsEveryReservationExactlyOnce(t *testing.T) {
}
}
func TestInboundRPCBatchCommitAppendsAllAndDefersSchedule(t *testing.T) {
func TestInboundRPCBatchCommitAppendsAllAndSchedulesAtomically(t *testing.T) {
scheduler := newInboundRPCScheduler(1, 8, 1<<20)
scheduler.start()
c := newInboundTestConn(scheduler, 1, 4, time.Second)
defer func() {
c.closeInboundRPCScheduler()
@ -159,13 +158,9 @@ func TestInboundRPCBatchCommitAppendsAllAndDefersSchedule(t *testing.T) {
return nil
}
}
activate, err := reservation.commit(tasks, true)
if err != nil {
if err := reservation.commit(tasks); err != nil {
t.Fatalf("commit batch: %v", err)
}
if activate == nil {
t.Fatal("deferred commit did not return an activation function")
}
c.rpcMu.Lock()
queued := len(c.rpcQueue)
ready := c.rpcReady
@ -173,17 +168,16 @@ func TestInboundRPCBatchCommitAppendsAllAndDefersSchedule(t *testing.T) {
if queued != len(specs) || !ready {
t.Fatalf("atomic queue state after commit = queued %d ready %v, want %d/true", queued, ready, len(specs))
}
if got := scheduler.readyLen(); got != 0 {
t.Fatalf("scheduler ready tokens before activation = %d, want zero", got)
if got := scheduler.readyLen(); got != 1 {
t.Fatalf("scheduler ready tokens after commit = %d, want one", got)
}
select {
case method := <-runs:
t.Fatalf("RPC %q ran before deferred activation", method)
t.Fatalf("RPC %q ran before scheduler start", method)
default:
}
activate()
activate() // activation is idempotent.
scheduler.start()
for _, want := range []string{"one", "two", "three"} {
select {
case got := <-runs:
@ -212,7 +206,7 @@ func TestInboundRPCBatchCommitMismatchReleasesAllWithoutEnqueue(t *testing.T) {
if err != nil {
t.Fatalf("reserve batch: %v", err)
}
if _, err := reservation.commit([]inboundRPC{{}}, false); !errors.Is(err, errInboundRPCBatchTaskCount) {
if err := reservation.commit([]inboundRPC{{}}); !errors.Is(err, errInboundRPCBatchTaskCount) {
t.Fatalf("commit task mismatch err = %v, want %v", err, errInboundRPCBatchTaskCount)
}
if tasks, bytes := scheduler.budgetSnapshot(); tasks != 0 || bytes != 0 {
@ -250,7 +244,7 @@ func TestInboundRPCBatchCommitRacingCloseNeverPartiallyEnqueues(t *testing.T) {
}()
waitInboundRPCBatchConnClosed(t, c)
if _, err := reservation.commit(make([]inboundRPC, 3), false); !errors.Is(err, ErrConnClosed) {
if err := reservation.commit(make([]inboundRPC, 3)); !errors.Is(err, ErrConnClosed) {
t.Fatalf("commit after close err = %v, want ErrConnClosed", err)
}
select {
@ -290,8 +284,8 @@ func TestInboundRPCBatchCommitAfterTerminalFenceRejectsAll(t *testing.T) {
// Session replacement and revocation publish terminal before the slower
// physical-close path. A reservation held across that fence must not be able
// to append even one stale task.
c.terminal.Store(true)
if _, err := reservation.commit(make([]inboundRPC, 2), false); !errors.Is(err, ErrConnClosed) {
c.retire()
if err := reservation.commit(make([]inboundRPC, 2)); !errors.Is(err, ErrConnClosed) {
t.Fatalf("commit after terminal fence err = %v, want ErrConnClosed", err)
}

54
internal/mtprotoedge/inbound_rpc_test.go
View file

@ -3,58 +3,11 @@ package mtprotoedge
import (
"context"
"errors"
"sync"
"sync/atomic"
"testing"
"time"
)
func TestStopInboundRPCTaskJoinsStartedDeadlineCallback(t *testing.T) {
ctx, cancel := context.WithDeadline(context.Background(), time.Now().Add(-time.Second))
defer cancel()
started := make(chan struct{})
release := make(chan struct{})
var once sync.Once
ticket := &inboundRPCTicket{}
ticket.onTimeout = func() {
once.Do(func() {
close(started)
<-release
})
}
task := inboundRPC{
ctx: ctx,
ticket: ticket,
stopTimeout: func() bool { return false }, // runtime callback already scheduled
}
handoff := stopInboundRPCTask(task)
if handoff == nil {
t.Fatal("started deadline callback did not produce handoff")
}
go ticket.onTimeout()
select {
case <-started:
case <-time.After(time.Second):
t.Fatal("deadline callback did not start")
}
joined := make(chan struct{})
go func() {
handoff()
close(joined)
}()
select {
case <-joined:
t.Fatal("handoff returned before in-flight callback completed")
case <-time.After(20 * time.Millisecond):
}
close(release)
select {
case <-joined:
case <-time.After(time.Second):
t.Fatal("handoff did not join completed callback")
}
}
func newInboundTestConn(s *inboundRPCScheduler, maxInflight, queueSize int, timeout time.Duration) *Conn {
c := &Conn{metrics: NopMetrics{}}
c.startInboundRPCScheduler(s, maxInflight, queueSize, timeout)
@ -438,7 +391,7 @@ func TestInboundRPCCloseDisarmsQueuedTimeout(t *testing.T) {
}
}
func TestInboundRPCRunningTimeoutSignalsWithoutReleasingBodyEarly(t *testing.T) {
func TestInboundRPCRunningDeadlineCancelsWithoutEarlyTimeout(t *testing.T) {
scheduler := newInboundRPCScheduler(1, 8, 1<<20)
scheduler.start()
c := newInboundTestConn(scheduler, 1, 4, 30*time.Millisecond)
@ -465,10 +418,11 @@ func TestInboundRPCRunningTimeoutSignalsWithoutReleasingBodyEarly(t *testing.T)
t.Fatalf("enqueue running task: %v", err)
}
<-started
time.Sleep(80 * time.Millisecond)
select {
case <-timedOut:
case <-time.After(time.Second):
t.Fatal("running task did not signal timeout while handler ignored cancellation")
t.Fatal("running task emitted an early timeout before handler convergence")
default:
}
if tasks, bytes := scheduler.budgetSnapshot(); tasks != 1 || bytes != 7 {
t.Fatalf("running body budget after timeout = (%d, %d), want retained (1, 7)", tasks, bytes)

47
internal/mtprotoedge/outbound.go
View file

@ -336,10 +336,7 @@ func (c *Conn) Close() {
// close. It closes both producer gates and cancels RPC work before any potentially blocking
// transport.Close call, so a timed-out batch close cannot keep accepting memory/work.
func (c *Conn) beginTerminalShutdown() {
c.terminal.Store(true)
// Retirement is irreversible. SessionManager activation only uses CAS from
// provisional/claiming, so a stale goroutine cannot publish this Conn again.
c.lifecycle.Store(uint32(connLifecycleRetired))
c.retire()
c.signalOutboundStop()
c.beginCloseInboundRPCScheduler()
}
@ -367,16 +364,6 @@ func (c *Conn) waitOutboundShutdownUntil(timeout time.Duration) bool {
}
}
// ForceClose 停止连接并关闭底层 transport。
// 仅用于授权撤销 / destroy_auth_key 这类“必须让对端立即断线”的路径；普通生命周期仍由
// serveConn 统一关闭 transport，避免正常 push/索引清理把长连接误伤成硬断。
func (c *Conn) ForceClose() {
c.beginTerminalShutdown()
c.closeTransport()
c.closeInboundRPCScheduler()
c.waitOutboundShutdown()
}
// closeTransport 只关闭物理 transport，不等待 outbound actor。写失败路径运行在
// actor 自身 goroutine 中，若在这里调用 Close 会等待 outboundDone 而自锁。
func (c *Conn) closeTransport() {
@ -395,7 +382,7 @@ func (c *Conn) closeTransport() {
}
// failTransport 把不可恢复的写错误提升为连接级 terminal failure。它只负责
// 标记 terminal + 关闭 socket；handleOutboundOp 返回后，actor 自己发停止信号并退出，
// 把 lifecycle 推进到 retired 并关闭 socket；handleOutboundOp 返回后，actor 自己发停止信号并退出，
// serveConn 被 Close 解开 Recv 后负责注销索引。
func (c *Conn) failTransport() {
// Publish both producer gates before Close: a custom/broken transport may block in
@ -417,10 +404,9 @@ func (c *Conn) fenceUndeliveredRPCResult() {
// A replacement/shutdown that already published terminal owns physical
// lifecycle cleanup (and may intentionally transfer the lease). Only the
// resultless task that wins false->true is allowed to close this generation.
if !c.terminal.CompareAndSwap(false, true) {
if !c.retire() {
return
}
c.lifecycle.Store(uint32(connLifecycleRetired))
c.signalOutboundStop()
if c.transportLease != nil {
c.transportLease.startCloseAlreadyFenced()
@ -461,11 +447,6 @@ func (c *Conn) Send(ctx context.Context, t proto.MessageType, msg bin.Encoder) e
return c.send(ctx, t, msg, false)
}
// SendPriority 加密并优先发送一条 server 控制消息。
func (c *Conn) SendPriority(ctx context.Context, t proto.MessageType, msg bin.Encoder) error {
return c.send(ctx, t, msg, true)
}
// SendRequiredControl writes a protocol-critical control message before the caller commits
// the state transition guarded by that message. One absolute deadline covers encode admission,
// body-budget reservation, control-queue admission and the physical transport write. A failure
@ -515,7 +496,7 @@ func (c *Conn) sendBestEffort(ctx context.Context, t proto.MessageType, msg bin.
if c.outbound == nil || c.outboundControl == nil {
return ErrConnClosed
}
if c.terminal.Load() {
if c.isRetired() {
return ErrConnClosed
}
writeCtx := context.Background()
@ -645,7 +626,7 @@ func (c *Conn) SendAsync(ctx context.Context, t proto.MessageType, msg bin.Encod
if c.outbound == nil || c.outboundControl == nil {
return ErrConnClosed
}
if c.terminal.Load() {
if c.isRetired() {
return ErrConnClosed
}
op, err := c.newOutboundSendOp(ctx, t, msg, nil, true)
@ -677,7 +658,7 @@ func (c *Conn) SendAsync(ctx context.Context, t proto.MessageType, msg bin.Encod
// AckServerMessages 接收客户端 msgs_ack，释放已确认的 server 出站消息。
func (c *Conn) AckServerMessages(ids []int64) {
if len(ids) == 0 || c.outbound == nil || c.outboundControl == nil || c.terminal.Load() {
if len(ids) == 0 || c.outbound == nil || c.outboundControl == nil || c.isRetired() {
return
}
op, err := c.newOutboundVectorOp(outboundAck, ids)
@ -790,7 +771,7 @@ func (c *Conn) enqueueOutboundRegistered(ctx context.Context, op outboundOp) err
if ctx == nil {
ctx = context.Background()
}
if c.terminal.Load() {
if c.isRetired() {
return ErrConnClosed
}
q := c.outbound
@ -820,7 +801,7 @@ func (c *Conn) enqueueOutboundRegistered(ctx context.Context, op outboundOp) err
func (c *Conn) beginOutboundEnqueue() bool {
c.outboundEnqueueMu.Lock()
defer c.outboundEnqueueMu.Unlock()
if c.outboundClosing || c.terminal.Load() {
if c.outboundClosing || c.isRetired() {
return false
}
c.outboundEnqueueWG.Add(1)
@ -840,14 +821,14 @@ func (c *Conn) outboundLoop() {
close(c.outboundDone)
}()
for {
if c.terminal.Load() {
if c.isRetired() {
c.signalOutboundStop()
c.drainOutbound()
return
}
select {
case op := <-c.outboundControl:
if c.terminal.Load() {
if c.isRetired() {
op.releaseReservation(c.outboundTrackedBudget)
op.finish(outboundResult{err: ErrConnClosed})
c.signalOutboundStop()
@ -855,7 +836,7 @@ func (c *Conn) outboundLoop() {
return
}
c.handleOutboundOp(state, op)
if c.terminal.Load() {
if c.isRetired() {
c.signalOutboundStop()
c.drainOutbound()
return
@ -868,7 +849,7 @@ func (c *Conn) outboundLoop() {
c.drainOutbound()
return
case op := <-c.outboundControl:
if c.terminal.Load() {
if c.isRetired() {
op.releaseReservation(c.outboundTrackedBudget)
op.finish(outboundResult{err: ErrConnClosed})
c.signalOutboundStop()
@ -877,7 +858,7 @@ func (c *Conn) outboundLoop() {
}
c.handleOutboundOp(state, op)
case op := <-c.outbound:
if c.terminal.Load() {
if c.isRetired() {
op.releaseReservation(c.outboundTrackedBudget)
op.finish(outboundResult{err: ErrConnClosed})
c.signalOutboundStop()
@ -886,7 +867,7 @@ func (c *Conn) outboundLoop() {
}
c.handleOutboundOp(state, op)
}
if c.terminal.Load() {
if c.isRetired() {
c.signalOutboundStop()
c.drainOutbound()
return

10
internal/mtprotoedge/outbound_required_control_test.go
View file

@ -89,7 +89,7 @@ func TestSendRequiredControlWaitsForPhysicalWriteAndReturnsBudget(t *testing.T)
case <-time.After(time.Second):
t.Fatal("SendRequiredControl did not return after physical write")
}
if c.terminal.Load() {
if c.isRetired() {
t.Fatal("successful required control terminally closed the connection")
}
if got := controlBudget.snapshot(); got != 0 {
@ -163,7 +163,7 @@ func TestSendRequiredControlQueueDeadlineTerminatesAndReturnsBudget(t *testing.T
if elapsed := time.Since(started); elapsed > 500*time.Millisecond {
t.Fatalf("full control queue waited %v, want parent-deadline-bounded admission", elapsed)
}
if !c.terminal.Load() {
if !c.isRetired() {
t.Fatal("required control queue failure did not terminally close the connection")
}
if got := tr.sends.Load(); got != 0 {
@ -197,7 +197,7 @@ func TestSendRequiredControlBlockedWriteUsesWholeOperationDeadline(t *testing.T)
case <-time.After(time.Second):
t.Fatal("outbound actor did not stop after required-control write timeout")
}
if !c.terminal.Load() {
if !c.isRetired() {
t.Fatal("blocked required control did not terminally close the connection")
}
if got := tr.closes.Load(); got != 1 {
@ -224,7 +224,7 @@ func TestSendRequiredControlWriteFailureTerminatesAndReturnsBudget(t *testing.T)
case <-time.After(time.Second):
t.Fatal("outbound actor did not stop after required-control write failure")
}
if !c.terminal.Load() {
if !c.isRetired() {
t.Fatal("write-failed required control did not terminally close the connection")
}
if got := tr.closes.Load(); got != 1 {
@ -250,7 +250,7 @@ func TestSendRequiredControlBudgetFailureIsTerminal(t *testing.T) {
case <-time.After(time.Second):
t.Fatal("outbound actor did not stop after required-control budget failure")
}
if !c.terminal.Load() {
if !c.isRetired() {
t.Fatal("required-control budget failure did not terminally close the connection")
}
if got := tr.sends.Load(); got != 0 {

2
internal/mtprotoedge/outbound_test.go
View file

@ -286,7 +286,7 @@ func TestOutboundScratchAdmissionUsesWriteTimeoutWithoutClosingHealthyConnection
if got := tr.sends.Load(); got != 0 {
t.Fatalf("writer called %d times without scratch, want 0", got)
}
if c.terminal.Load() {
if c.isRetired() {
t.Fatal("scratch admission timeout terminally closed a healthy connection")
}
select {

283
internal/mtprotoedge/production_compat_test.go Normal file
View file

@ -0,0 +1,283 @@
package mtprotoedge
import (
"context"
"errors"
"sort"
"time"
"go.uber.org/zap"
"github.com/gotd/td/bin"
"github.com/gotd/td/proto"
)
var ErrSessionAmbiguous = errors.New("session id is shared by multiple auth keys")
// ForceClose is intentionally test-only. Production shutdown paths use the
// narrower lifecycle primitives so callers cannot bypass ownership rules.
func (c *Conn) ForceClose() {
c.beginTerminalShutdown()
c.closeTransport()
c.closeInboundRPCScheduler()
c.waitOutboundShutdown()
}
// newRPCResultCache keeps older focused cache tests concise without exposing a
// second production constructor.
func newRPCResultCache(now func() time.Time) *rpcResultCache {
return newRPCResultCacheWithFlightLimit(now, rpcResultFlightDefaultMaxPending)
}
// Conns is a white-box test accessor. Production wires the shared manager
// explicitly and does not need a second access path through Server.
func (s *Server) Conns() *SessionManager {
return s.conns
}
// Register is a test fixture shortcut for tests that do not exercise the
// wire-level required-control barrier.
func (m *SessionManager) Register(c *Conn) error {
if c == nil {
return ErrSessionActivationSuperseded
}
if c.isActive() {
m.mu.RLock()
current := m.bySession[connSessionKey(c)]
m.mu.RUnlock()
if current == c {
return nil
}
return ErrSessionActivationSuperseded
}
if err := m.BeginActivation(c); err != nil {
return err
}
if err := m.PublishActivation(c); err != nil {
m.AbortActivation(c)
return err
}
return nil
}
func (m *SessionManager) uniqueSessionForTestLocked(sessionID int64) (*Conn, sessionKey, bool, bool) {
var (
found *Conn
foundKey sessionKey
)
for key, c := range m.bySession {
if key.sessionID != sessionID {
continue
}
if found != nil {
return nil, sessionKey{}, false, true
}
found, foundKey = c, key
}
return found, foundKey, found != nil, false
}
// The session-id-only helpers below preserve focused legacy tests without
// carrying an ambiguous global index or API in production.
func (m *SessionManager) BindUser(sessionID, userID int64) {
m.mu.Lock()
defer m.mu.Unlock()
c, key, ok, ambiguous := m.uniqueSessionForTestLocked(sessionID)
if !ambiguous && ok {
m.bindUserLocked(c, key, userID)
}
}
func (m *SessionManager) UserID(sessionID int64) (int64, bool) {
m.mu.RLock()
c, _, ok, ambiguous := m.uniqueSessionForTestLocked(sessionID)
m.mu.RUnlock()
if ambiguous || !ok {
return 0, false
}
userID := c.userID.Load()
return userID, userID != 0
}
func (m *SessionManager) UserIDResolved(sessionID int64) (int64, bool) {
m.mu.RLock()
c, _, ok, ambiguous := m.uniqueSessionForTestLocked(sessionID)
m.mu.RUnlock()
if ambiguous || !ok {
return 0, false
}
return c.UserIDResolved()
}
func (m *SessionManager) UserIDForAuthKey(authKeyID [8]byte, sessionID int64) (int64, bool) {
m.mu.RLock()
c, ok := m.bySession[sessionKey{authKeyID: authKeyID, sessionID: sessionID}]
m.mu.RUnlock()
if !ok {
return 0, false
}
userID := c.userID.Load()
return userID, userID != 0
}
func (m *SessionManager) BindAuthKey(sessionID int64, authKeyID [8]byte) {
m.mu.Lock()
defer m.mu.Unlock()
c, key, ok, ambiguous := m.uniqueSessionForTestLocked(sessionID)
if !ambiguous && ok {
m.bindAuthKeyLocked(c, key, authKeyID)
}
}
func (m *SessionManager) AuthKeyID(sessionID int64) ([8]byte, bool) {
m.mu.RLock()
c, _, ok, ambiguous := m.uniqueSessionForTestLocked(sessionID)
m.mu.RUnlock()
if ambiguous || !ok {
return [8]byte{}, false
}
return c.BusinessAuthKeyID()
}
func (m *SessionManager) SetReceivesUpdates(sessionID int64, receives bool) {
m.mu.Lock()
c, key, ok, ambiguous := m.uniqueSessionForTestLocked(sessionID)
if ambiguous || !ok {
m.mu.Unlock()
return
}
owner, start := m.setReceivesUpdatesLocked(c, key, receives)
m.mu.Unlock()
if start {
go m.runFlush(c, key, owner, 0)
}
}
func (m *SessionManager) PushToSession(ctx context.Context, sessionID int64, t proto.MessageType, msg bin.Encoder) error {
m.mu.RLock()
c, key, ok, ambiguous := m.uniqueSessionForTestLocked(sessionID)
if ambiguous {
m.mu.RUnlock()
return ErrSessionAmbiguous
}
if !ok {
m.mu.RUnlock()
return ErrSessionNotFound
}
ready := c.receivesUpdates.Load()
m.mu.RUnlock()
if ready {
return c.Send(ctx, t, msg)
}
return m.queueOrSendPrepared(ctx, key, t, msg)
}
func (m *SessionManager) PushToUser(ctx context.Context, userID int64, t proto.MessageType, msg bin.Encoder) (int, error) {
return m.PushToUserExceptAuthKeySession(ctx, userID, [8]byte{}, 0, t, msg)
}
func (m *SessionManager) PushToUserExceptSession(ctx context.Context, userID, excludeSessionID int64, t proto.MessageType, msg bin.Encoder) (int, error) {
return m.pushToUser(ctx, userID, nil, excludeSessionID, t, msg)
}
func (m *SessionManager) PushToUserExceptSessionBestEffort(ctx context.Context, userID, excludeSessionID int64, t proto.MessageType, msg bin.Encoder, timeout time.Duration) (int, error) {
return m.pushToUserBestEffort(ctx, userID, nil, excludeSessionID, t, msg, timeout)
}
func (m *SessionManager) Online() int {
m.mu.RLock()
defer m.mu.RUnlock()
return len(m.bySession)
}
func (m *SessionManager) OnlineChannelIDsAfter(afterChannelID int64, limit int) []int64 {
if limit <= 0 {
return nil
}
const maxRecoveryPage = 4096
if limit > maxRecoveryPage {
limit = maxRecoveryPage
}
all := m.OnlineChannelIDsSnapshot()
start := sort.Search(len(all), func(i int) bool { return all[i] > afterChannelID })
end := start + limit
if end > len(all) {
end = len(all)
}
return all[start:end]
}
func (m *SessionManager) queueLocked(key sessionKey, t proto.MessageType, msg bin.Encoder) bool {
encoded, reservation, err := m.preparePendingPush(context.Background(), msg)
if err != nil {
m.log.Debug("Drop pending push outside byte budget",
zap.String("auth_key_id", sessionKeyLog(key.authKeyID)),
zap.Int64("session_id", key.sessionID),
zap.Error(err),
)
return false
}
defer reservation.release()
return m.queuePreparedLocked(key, t, encoded, reservation)
}
func (cs *connState) track(msgID int64, seqNo int32, content bool, state byte) {
cs.trackInbound(msgID, seqNo, content, false, state)
}
func (s *Server) dispatch(ctx context.Context, cs *connState, c *Conn, msgID int64, seqNo int32, b *bin.Buffer, acks *[]int64) error {
plan, err := s.preflightInbound(cs, msgID, seqNo, b.Buf)
if err != nil {
var bad *dispatchBadMsgError
if errors.As(err, &bad) && c != nil {
return s.sendBadMsg(ctx, c, bad.msgID, bad.seqNo, bad.code)
}
return err
}
defer plan.close()
plan.commitState(cs)
*acks = append(*acks, plan.ackIDs...)
return s.executeInboundPlan(ctx, cs, c, plan)
}
// inboundRPCReservation adapts legacy focused tests to the sole production
// reservation state machine: a batch with exactly one entry.
type inboundRPCReservation struct {
batch *inboundRPCBatchReservation
}
func (c *Conn) reserveInboundRPC(ctx context.Context, method string, size int) (*inboundRPCReservation, error) {
batch, err := c.reserveInboundRPCBatch(ctx, []inboundRPCSpec{{method: method, size: size}})
if err != nil {
return nil, err
}
return &inboundRPCReservation{batch: batch}, nil
}
func (r *inboundRPCReservation) commit(task inboundRPC) error {
if r == nil || r.batch == nil {
return ErrConnClosed
}
return r.batch.commit([]inboundRPC{task})
}
func (r *inboundRPCReservation) abort() {
if r != nil && r.batch != nil {
r.batch.abort()
}
}
func (c *Conn) enqueueInboundRPC(ctx context.Context, task inboundRPC) error {
reservation, err := c.reserveInboundRPC(ctx, task.method, task.size)
if err != nil {
return err
}
defer reservation.abort()
return reservation.commit(task)
}
func (s *inboundRPCScheduler) budgetSnapshot() (tasks int, bytes int64) {
s.budgetMu.Lock()
defer s.budgetMu.Unlock()
return s.tasks, s.bytes
}

4
internal/mtprotoedge/rpc_result_cache.go
View file

@ -55,10 +55,6 @@ type rpcResultCacheShard struct {
pending map[rpcResultCacheKey]*rpcResultFlight
}
func newRPCResultCache(now func() time.Time) *rpcResultCache {
return newRPCResultCacheWithFlightLimit(now, rpcResultFlightDefaultMaxPending)
}
func newRPCResultCacheWithFlightLimit(now func() time.Time, maxPending int) *rpcResultCache {
if now == nil {
now = time.Now

4
internal/mtprotoedge/rpc_result_delivery_test.go
View file

@ -102,8 +102,8 @@ func TestRPCResultPrewriteFailureFencesConnBeforeCachePublication(t *testing.T)
for !tr.closed.Load() && time.Now().Before(closeDeadline) {
time.Sleep(time.Millisecond)
}
if !c.terminal.Load() || !tr.closed.Load() || c.isPhysicalTransportCurrentOpen() {
t.Fatalf("failed delivery did not fence Conn: terminal=%v closed=%v current_open=%v", c.terminal.Load(), tr.closed.Load(), c.isPhysicalTransportCurrentOpen())
if !c.isRetired() || !tr.closed.Load() || c.isPhysicalTransportCurrentOpen() {
t.Fatalf("failed delivery did not fence Conn: retired=%v closed=%v current_open=%v", c.isRetired(), tr.closed.Load(), c.isPhysicalTransportCurrentOpen())
}
completed, acquireErr := s.rpcResults.Acquire(key.ID, c.sessionID, reqMsgID)
if acquireErr != nil || completed.state != rpcResultAcquireCompleted || completed.encoded == nil {

51
internal/mtprotoedge/rpc_test.go
View file

@ -151,7 +151,7 @@ func TestInboundRPCQueuedDeadlineReturnsRPCTimeout(t *testing.T) {
}
}
func TestInboundRPCRunningDeadlineReturnsExactlyOneTimeout(t *testing.T) {
func TestInboundRPCRunningDeadlineWaitsForHandlerTerminalResult(t *testing.T) {
for _, tc := range []struct {
name string
honorContext bool
@ -185,40 +185,33 @@ func TestInboundRPCRunningDeadlineReturnsExactlyOneTimeout(t *testing.T) {
t.Fatal("timed out waiting for running rpc")
}
// In the ignore-context case this result must arrive before release is closed: the
// scheduler deadline, not eventual handler return, owns the timeout response.
if !tc.honorContext {
// Let the deadline expire while Dispatch is still running. No early timeout
// may win; after the handler reports committed success, that success is the
// sole terminal result.
time.Sleep(120 * time.Millisecond)
close(handler.release)
}
result := readRPCResultForRequest(t, conn, cipher, auth.AuthKey, reqID)
var rpcErr mt.RPCError
if err := rpcErr.Decode(&bin.Buffer{Buf: result.Result}); err != nil {
t.Fatalf("decode running rpc timeout: %v", err)
if tc.honorContext {
var rpcErr mt.RPCError
if err := rpcErr.Decode(&bin.Buffer{Buf: result.Result}); err != nil {
t.Fatalf("decode converged rpc timeout: %v", err)
}
if rpcErr.ErrorCode != 500 || rpcErr.ErrorMessage != "RPC_TIMEOUT" {
t.Fatalf("rpc_error = %d %q, want 500 RPC_TIMEOUT", rpcErr.ErrorCode, rpcErr.ErrorMessage)
}
close(handler.release)
} else {
var config tg.Config
if err := config.Decode(&bin.Buffer{Buf: result.Result}); err != nil {
t.Fatalf("decode committed success after deadline: %v", err)
}
}
if rpcErr.ErrorCode != 500 || rpcErr.ErrorMessage != "RPC_TIMEOUT" {
t.Fatalf("rpc_error = %d %q, want 500 RPC_TIMEOUT", rpcErr.ErrorCode, rpcErr.ErrorMessage)
}
close(handler.release)
})
}
}
func TestRPCResponseGateExactlyOnce(t *testing.T) {
for i := 0; i < 100; i++ {
gate := &rpcResponseGate{}
results := make(chan bool, 2)
go func() { results <- gate.tryNormal() }()
go func() { results <- gate.tryTimeout() }()
wins := 0
if <-results {
wins++
}
if <-results {
wins++
}
if wins != 1 {
t.Fatalf("iteration %d response gate winners = %d, want 1", i, wins)
}
}
}
func TestDuplicateRPCResultAcrossReconnectUsesSessionCache(t *testing.T) {
const dc = 2
handler := &countingConfigRPC{}

19
internal/mtprotoedge/server.go
View file

@ -120,8 +120,6 @@ type Options struct {
RSAKey *rsa.PrivateKey
// AuthKeys 持久化 auth key。默认内存实现。
AuthKeys store.AuthKeyStore
// Sessions 记录在线 MTProto session（持久化数据）。默认内存实现。
Sessions store.SessionStore
// ActiveSessions 管理活跃连接。默认新建；传入时可让 RPC 层共享同一注册表。
ActiveSessions *SessionManager
// RPC 是 typed RPC 路由。nil 时加密 RPC 被丢弃并记录。
@ -198,9 +196,6 @@ func (o *Options) setDefaults() {
if o.AuthKeys == nil {
o.AuthKeys = memory.NewAuthKeyStore()
}
if o.Sessions == nil {
o.Sessions = memory.NewSessionStore()
}
if o.Metrics == nil {
o.Metrics = NopMetrics{}
}
@ -241,7 +236,6 @@ type Server struct {
dc int
key exchange.PrivateKey
authKeys store.AuthKeyStore
sessions store.SessionStore
conns *SessionManager
rpc RPCHandler
metrics Metrics
@ -287,7 +281,6 @@ func New(opts Options) *Server {
dc: opts.DC,
key: exchange.PrivateKey{RSA: opts.RSAKey},
authKeys: opts.AuthKeys,
sessions: opts.Sessions,
conns: conns,
rpc: opts.RPC,
metrics: opts.Metrics,
@ -300,11 +293,6 @@ func New(opts Options) *Server {
}
}
// Conns 返回活跃连接注册表，供业务层主动推送（updates 等）。
func (s *Server) Conns() *SessionManager {
return s.conns
}
// newConn 基于一次解密结果创建一个可发送的连接对象。
func (s *Server) newConn(tc transport.Conn, key crypto.AuthKey, sessionID, salt int64) *Conn {
if lease, ok := tc.(*physicalTransportLease); ok {
@ -682,12 +670,11 @@ func (s *Server) serveConn(ctx context.Context, raw transport.Conn) (err error)
// 已建立连接复用缓存密钥走快路径（fetchedKey=nil）：避开每帧回查 AuthKeyStore——
// 这是 mtprotoedge 层最热的库访问点。密钥材料创建后不可变；销毁(destroy_auth_key)/
// 撤销由 SessionManager 主动 Close 连接保证失效，不依赖此被动回查。仅 destroy_auth_key
// 的发起连接置 keyDestroyed，使其下一帧回落到 Get→AuthKeyNotFound。尚未进入
// SessionManager 的 bad-salt provisional 会在 handleEncrypted 建立 activation claim
// 撤销由 SessionManager 主动 Close 连接保证失效，不依赖被动的“下一帧 -404”。
// 尚未进入 SessionManager 的 bad-salt provisional 会在 handleEncrypted 建立 activation claim
// 后精确复查一次，既把撤销与激活线性化，也不把 salt storm 放大成 PG 写风暴。
var fetchedKey *store.AuthKeyData
if current == nil || current.authKeyID != authKeyID || current.keyDestroyed.Load() {
if current == nil || current.authKeyID != authKeyID {
d, found, err := s.authKeys.Get(ctx, authKeyID)
if err != nil {
return fmt.Errorf("lookup auth key: %w", err)

12
internal/mtprotoedge/session_activation_test.go
View file

@ -46,8 +46,8 @@ func TestSessionActivationGatesReplacementBeforePublishing(t *testing.T) {
if !newConn.isActive() {
t.Fatal("replacement was not activated")
}
if oldConn.lifecycleState() != connLifecycleRetired || !oldConn.terminal.Load() {
t.Fatalf("old connection gates = lifecycle:%v terminal:%v", oldConn.lifecycleState(), oldConn.terminal.Load())
if !oldConn.isRetired() {
t.Fatalf("old connection lifecycle=%v", oldConn.lifecycleState())
}
if err := oldConn.SendAsync(context.Background(), proto.MessageFromServer, &mt.MsgsAck{}); !errors.Is(err, ErrConnClosed) {
t.Fatalf("stale old connection send error = %v, want ErrConnClosed", err)
@ -80,8 +80,8 @@ func TestSessionActivationClaimPreemptionCannotReversePublish(t *testing.T) {
if err := manager.BeginActivation(second); err != nil {
t.Fatalf("begin superseding activation: %v", err)
}
if first.lifecycleState() != connLifecycleRetired || !first.terminal.Load() {
t.Fatalf("superseded first lifecycle=%v terminal=%v", first.lifecycleState(), first.terminal.Load())
if !first.isRetired() {
t.Fatalf("superseded first lifecycle=%v", first.lifecycleState())
}
if err := manager.PublishActivation(first); !errors.Is(err, ErrSessionActivationSuperseded) {
t.Fatalf("stale publish error = %v, want superseded", err)
@ -205,8 +205,8 @@ func TestRawAuthKeyCloseExactConnDoesNotExcludeSameSessionReplacement(t *testing
manager.mu.RLock()
active, claim := manager.bySession[sessionKey{authKeyID: key, sessionID: sessionID}], manager.claims[sessionKey{authKeyID: key, sessionID: sessionID}]
manager.mu.RUnlock()
if active != nil || claim != nil || !replacement.terminal.Load() {
t.Fatalf("same-session replacement escaped exact exclusion: active=%p claim=%p terminal=%v", active, claim, replacement.terminal.Load())
if active != nil || claim != nil || !replacement.isRetired() {
t.Fatalf("same-session replacement escaped exact exclusion: active=%p claim=%p lifecycle=%v", active, claim, replacement.lifecycleState())
}
}

16
internal/mtprotoedge/session_barrier_integration_test.go
View file

@ -81,8 +81,8 @@ func TestBadServerSaltRetainsOneProvisionalConnUntilCorrected(t *testing.T) {
if err != nil {
t.Fatalf("first bad salt: %v", err)
}
if firstConn == nil || firstConn.lifecycleState() != connLifecycleProvisional || firstConn.terminal.Load() {
t.Fatalf("first correction lifecycle conn=%p state=%v terminal=%v", firstConn, firstConn.lifecycleState(), firstConn.terminal.Load())
if firstConn == nil || firstConn.lifecycleState() != connLifecycleProvisional {
t.Fatalf("first correction lifecycle conn=%p state=%v", firstConn, firstConn.lifecycleState())
}
if cs.createdFloor != 0 || len(cs.seen) != 0 || handler.calls.Load() != 0 {
t.Fatalf("bad salt admitted state: floor=%d seen=%d calls=%d", cs.createdFloor, len(cs.seen), handler.calls.Load())
@ -192,8 +192,8 @@ func TestWrongSaltSessionChangeTransfersPhysicalOwnership(t *testing.T) {
if newConn == nil || newConn == oldConn || newConn.lifecycleState() != connLifecycleProvisional {
t.Fatalf("replacement conn old=%p new=%p state=%v", oldConn, newConn, newConn.lifecycleState())
}
if !oldConn.terminal.Load() || tr.closed.Load() {
t.Fatalf("transfer state old_terminal=%v raw_closed=%v", oldConn.terminal.Load(), tr.closed.Load())
if !oldConn.isRetired() || tr.closed.Load() {
t.Fatalf("transfer state old_lifecycle=%v raw_closed=%v", oldConn.lifecycleState(), tr.closed.Load())
}
// A delayed stale close must not tear down the generation already transferred
// to the new logical session.
@ -483,10 +483,10 @@ func TestCrossConnectionInflightRPCHasOneBusinessOwnerAndReplaysResult(t *testin
}()
deadline := time.Now().Add(2 * time.Second)
for !firstConn.terminal.Load() && time.Now().Before(deadline) {
for !firstConn.isRetired() && time.Now().Before(deadline) {
time.Sleep(time.Millisecond)
}
if !firstConn.terminal.Load() {
if !firstConn.isRetired() {
t.Fatal("replacement did not fence the first physical connection")
}
if got := handler.calls.Load(); got != 1 {
@ -612,8 +612,8 @@ func TestCrossConnectionInflightAbortRetriesOnlyAfterOldOwnerStops(t *testing.T)
if resultCount != 1 {
t.Fatalf("sequential retry result count = %d, want 1", resultCount)
}
if !firstConn.terminal.Load() || secondConn == nil || !secondConn.isActive() {
t.Fatalf("replacement lifecycle = old terminal:%v new:%p active:%v", firstConn.terminal.Load(), secondConn, secondConn != nil && secondConn.isActive())
if !firstConn.isRetired() || secondConn == nil || !secondConn.isActive() {
t.Fatalf("replacement lifecycle = old:%v new:%p active:%v", firstConn.lifecycleState(), secondConn, secondConn != nil && secondConn.isActive())
}
secondConn.ForceClose()
}

298
internal/mtprotoedge/session_manager.go
View file

@ -19,9 +19,6 @@ import (
// ErrSessionNotFound 表示目标 session 当前无活跃连接。
var ErrSessionNotFound = errors.New("session not found")
// ErrSessionAmbiguous 表示仅用 session_id 无法唯一定位连接。
var ErrSessionAmbiguous = errors.New("session id is shared by multiple auth keys")
var (
ErrSessionActivationSuperseded = errors.New("session activation superseded")
ErrSessionActivationFence = errors.New("session activation could not fence previous writer")
@ -119,8 +116,8 @@ type SessionLifecycleObserver interface {
// SessionManager 是活跃连接注册表，支持按 session / auth-key / user 查找并主动 push。
//
// 它管理运行态的在线连接，与持久化的 store.SessionStore 互补：后者记录 session 数据，
// 前者持有可发送的活跃连接。所有方法并发安全。
// 它只管理进程内运行态，持有可发送的活跃连接；协议可恢复事实由 auth key、客户端重连
// 和 durable updates/difference 链路承担。所有方法并发安全。
type SessionManager struct {
mu sync.RWMutex
bySession map[sessionKey]*Conn
@ -129,7 +126,6 @@ type SessionManager struct {
// is on the wire and PublishActivation validates the same owner.
claims map[sessionKey]*Conn
claimsByAuth map[[8]byte]map[int64]*Conn // raw authKeyID -> sessionID -> provisional claim
bySessionID map[int64]map[[8]byte]*Conn // sessionID → raw authKeyID → Conn，用于兼容旧 API 的唯一性检查
byAuthKey map[[8]byte]map[int64]*Conn // raw authKeyID → sessionID → Conn
byBusinessAuthKey map[[8]byte]map[sessionKey]*Conn
byUser map[int64]map[sessionKey]*Conn
@ -154,7 +150,6 @@ func NewSessionManager(log *zap.Logger) *SessionManager {
bySession: make(map[sessionKey]*Conn),
claims: make(map[sessionKey]*Conn),
claimsByAuth: make(map[[8]byte]map[int64]*Conn),
bySessionID: make(map[int64]map[[8]byte]*Conn),
byAuthKey: make(map[[8]byte]map[int64]*Conn),
byBusinessAuthKey: make(map[[8]byte]map[sessionKey]*Conn),
byUser: make(map[int64]map[sessionKey]*Conn),
@ -189,7 +184,7 @@ func (m *SessionManager) BeginActivation(c *Conn) error {
key := connSessionKey(c)
retired := make([]*Conn, 0, 2)
m.mu.Lock()
if c.terminal.Load() || !c.isPhysicalTransportCurrentOpen() || c.lifecycleState() != connLifecycleClaiming {
if !c.isPhysicalTransportCurrentOpen() || c.lifecycleState() != connLifecycleClaiming {
c.beginTerminalShutdown()
m.mu.Unlock()
return ErrConnClosed
@ -248,7 +243,7 @@ func (m *SessionManager) PublishActivation(c *Conn) error {
if m.claims[key] != c {
return ErrSessionActivationSuperseded
}
if c.terminal.Load() || c.lifecycleState() != connLifecycleClaiming {
if c.lifecycleState() != connLifecycleClaiming {
m.removeClaimLocked(key, c)
return ErrConnClosed
}
@ -265,7 +260,6 @@ func (m *SessionManager) PublishActivation(c *Conn) error {
}
m.removeClaimLocked(key, c)
m.bySession[key] = c
addSessionIDIndex(m.bySessionID, c.sessionID, c.authKeyID, c)
addConnIndex(m.byAuthKey, c.authKeyID, c.sessionID, c)
if businessAuthKeyID, resolved := c.BusinessAuthKeyID(); resolved {
addBusinessAuthKeyIndex(m.byBusinessAuthKey, businessAuthKeyID, key, c)
@ -306,32 +300,6 @@ func (m *SessionManager) AbortActivation(c *Conn) {
}
}
// Register is retained for tests and embedders that do not have a wire-level
// required-control barrier. Production encrypted traffic uses the explicit
// BeginActivation -> SendRequiredControl -> PublishActivation sequence.
func (m *SessionManager) Register(c *Conn) error {
if c == nil {
return ErrSessionActivationSuperseded
}
if c.isActive() {
m.mu.RLock()
current := m.bySession[connSessionKey(c)]
m.mu.RUnlock()
if current == c {
return nil
}
return ErrSessionActivationSuperseded
}
if err := m.BeginActivation(c); err != nil {
return err
}
if err := m.PublishActivation(c); err != nil {
m.AbortActivation(c)
return err
}
return nil
}
// Unregister 注销一个连接（仅当它仍是当前注册的同一对象，避免误删重连后的新连接）。
// 观察者对未登录连接（userID=0）也回调：业务层据此清理按 session 维度的缓存条目，
// 否则未登录连接的元数据只能等容量上限驱逐。
@ -373,38 +341,6 @@ func (m *SessionManager) Unregister(c *Conn) {
}
}
// DestroySession 移除指定 session 的运行态索引，供 MTProto destroy_session 使用。
func (m *SessionManager) DestroySession(sessionID int64) bool {
m.mu.Lock()
c, key, ok, ambiguous := m.uniqueSessionLocked(sessionID)
if ambiguous || !ok {
if !ambiguous {
m.dropPendingBySessionLocked(sessionID)
}
m.mu.Unlock()
return false
}
offlineUser := m.retireConnLocked(c, true)
lastForUser := offlineUser != 0 && len(m.byUser[offlineUser]) == 0
observer := m.lifecycle
m.log.Debug("Session destroyed",
zap.String("auth_key_id", sessionKeyLog(key.authKeyID)),
zap.Int64("session_id", sessionID),
zap.Int("online", len(m.bySession)),
)
m.mu.Unlock()
if !forceCloseConnBatch([]*Conn{c}, forceCloseBatchTimeout) {
m.log.Warn("Destroyed session close exceeded shared deadline",
zap.String("auth_key_id", sessionKeyLog(key.authKeyID)),
zap.Int64("session_id", sessionID),
)
}
if observer != nil && offlineUser != 0 {
observer.SessionOffline(key.authKeyID, sessionID, offlineUser, lastForUser)
}
return true
}
// DestroySessionForAuthKey 精确移除某个 raw auth_key_id 下的 session。
func (m *SessionManager) DestroySessionForAuthKey(authKeyID [8]byte, sessionID int64) bool {
m.mu.Lock()
@ -447,22 +383,6 @@ func (m *SessionManager) DestroySessionForAuthKey(authKeyID [8]byte, sessionID i
return true
}
// BindUser 缓存 session 的授权用户。userID=0 表示当前 auth_key 已确认未登录。
// 登录后绑定非 0 userID，使其可经 PushToUser 收到推送。
func (m *SessionManager) BindUser(sessionID, userID int64) {
m.mu.Lock()
c, key, ok, ambiguous := m.uniqueSessionLocked(sessionID)
if ambiguous || !ok {
if ambiguous {
m.log.Warn("Skip BindUser for ambiguous session_id", zap.Int64("session_id", sessionID))
}
m.mu.Unlock()
return
}
m.bindUserLocked(c, key, userID)
m.mu.Unlock()
}
// BindUserForAuthKey 缓存指定 raw auth_key_id + session_id 的授权用户。
func (m *SessionManager) BindUserForAuthKey(authKeyID [8]byte, sessionID, userID int64) {
m.mu.Lock()
@ -500,48 +420,6 @@ func (m *SessionManager) bindUserLocked(c *Conn, key sessionKey, userID int64) {
}
}
// UserID 返回 session 当前缓存的登录用户 id。未绑定或离线时 ok=false。
func (m *SessionManager) UserID(sessionID int64) (int64, bool) {
m.mu.RLock()
c, _, ok, ambiguous := m.uniqueSessionLocked(sessionID)
m.mu.RUnlock()
if ambiguous || !ok {
return 0, false
}
userID := c.userID.Load()
if userID == 0 {
return 0, false
}
return userID, true
}
// UserIDForAuthKey 返回指定 raw auth_key_id + session_id 当前缓存的登录用户 id。
func (m *SessionManager) UserIDForAuthKey(authKeyID [8]byte, sessionID int64) (int64, bool) {
m.mu.RLock()
c, ok := m.bySession[sessionKey{authKeyID: authKeyID, sessionID: sessionID}]
m.mu.RUnlock()
if !ok {
return 0, false
}
userID := c.userID.Load()
if userID == 0 {
return 0, false
}
return userID, true
}
// UserIDResolved 返回 session 的 user_id 授权状态是否已经查过。
// resolved=true 且 userID=0 表示该 session 当前未登录。
func (m *SessionManager) UserIDResolved(sessionID int64) (int64, bool) {
m.mu.RLock()
c, _, ok, ambiguous := m.uniqueSessionLocked(sessionID)
m.mu.RUnlock()
if ambiguous || !ok {
return 0, false
}
return c.UserIDResolved()
}
// UserIDResolvedForAuthKey 返回指定 raw auth_key_id + session_id 的 user_id 缓存状态。
func (m *SessionManager) UserIDResolvedForAuthKey(authKeyID [8]byte, sessionID int64) (int64, bool) {
m.mu.RLock()
@ -553,21 +431,6 @@ func (m *SessionManager) UserIDResolvedForAuthKey(authKeyID [8]byte, sessionID i
return c.UserIDResolved()
}
// BindAuthKey 缓存业务视角 auth_key_id（temp auth_key 解析后的 perm auth_key）。
func (m *SessionManager) BindAuthKey(sessionID int64, authKeyID [8]byte) {
m.mu.Lock()
c, key, ok, ambiguous := m.uniqueSessionLocked(sessionID)
if ambiguous || !ok {
if ambiguous {
m.log.Warn("Skip BindAuthKey for ambiguous session_id", zap.Int64("session_id", sessionID))
}
m.mu.Unlock()
return
}
m.bindAuthKeyLocked(c, key, authKeyID)
m.mu.Unlock()
}
// BindAuthKeyForSession 缓存指定 raw auth_key_id + session_id 的业务 auth_key_id。
func (m *SessionManager) BindAuthKeyForSession(rawAuthKeyID [8]byte, sessionID int64, authKeyID [8]byte) {
m.mu.Lock()
@ -603,18 +466,6 @@ func (m *SessionManager) bindAuthKeyLocked(c *Conn, key sessionKey, authKeyID [8
}
}
// AuthKeyID 返回 session 缓存的业务视角 auth_key_id。
// ok=false 表示该连接尚未完成 temp→perm 解析。
func (m *SessionManager) AuthKeyID(sessionID int64) ([8]byte, bool) {
m.mu.RLock()
c, _, ok, ambiguous := m.uniqueSessionLocked(sessionID)
m.mu.RUnlock()
if ambiguous || !ok {
return [8]byte{}, false
}
return c.BusinessAuthKeyID()
}
// AuthKeyIDForSession 返回指定 raw auth_key_id + session_id 缓存的业务 auth_key_id。
func (m *SessionManager) AuthKeyIDForSession(rawAuthKeyID [8]byte, sessionID int64) ([8]byte, bool) {
m.mu.RLock()
@ -863,28 +714,6 @@ func (m *SessionManager) UnbindAuthKey(authKeyID [8]byte) int {
return count
}
// SetReceivesUpdates 标记 session 是否已完成 updates 同步入口。
//
// TDesktop 登录后会先调用 updates.getState/getDifference 建立本地同步基线。
// 在此之前收到的主动 updates 先暂存，待 session 可接收后再异步下发。
func (m *SessionManager) SetReceivesUpdates(sessionID int64, receives bool) {
m.mu.Lock()
c, key, ok, ambiguous := m.uniqueSessionLocked(sessionID)
if ambiguous || !ok {
if ambiguous {
m.log.Warn("Skip SetReceivesUpdates for ambiguous session_id", zap.Int64("session_id", sessionID))
}
m.mu.Unlock()
return
}
owner, start := m.setReceivesUpdatesLocked(c, key, receives)
m.mu.Unlock()
if start {
go m.runFlush(c, key, owner, 0)
}
}
// setReceivesUpdatesLocked 是置位/复位的共同内核，调用方须持有 m.mu。
// 置位且有暂存时不立即置 receivesUpdates：标记 flushing 并返回该批暂存所属的 userID，
// 交由 runFlush 排空后原子置位，期间新到推送继续进 pending，保证暂存与实时推送的
@ -1059,26 +888,6 @@ func (m *SessionManager) SetReceivesUpdatesForAuthKey(authKeyID [8]byte, session
}
}
// PushToSession 向指定 session 推送一条消息。
func (m *SessionManager) PushToSession(ctx context.Context, sessionID int64, t proto.MessageType, msg bin.Encoder) error {
m.mu.RLock()
c, key, ok, ambiguous := m.uniqueSessionLocked(sessionID)
if ambiguous {
m.mu.RUnlock()
return ErrSessionAmbiguous
}
if !ok {
m.mu.RUnlock()
return ErrSessionNotFound
}
ready := c.receivesUpdates.Load()
m.mu.RUnlock()
if ready {
return c.Send(ctx, t, msg)
}
return m.queueOrSendPrepared(ctx, key, t, msg)
}
// PushToSessionForAuthKey 向指定 raw auth_key_id + session_id 推送一条消息。
func (m *SessionManager) PushToSessionForAuthKey(ctx context.Context, authKeyID [8]byte, sessionID int64, t proto.MessageType, msg bin.Encoder) error {
m.mu.RLock()
@ -1134,18 +943,6 @@ func (m *SessionManager) PushToSessionForAuthKeyImmediate(ctx context.Context, a
return c.SendBestEffort(ctx, t, msg, 2*time.Second)
}
// PushToUser 向某 user 所有活跃连接推送，返回已发送或已暂存的连接数。
// 发送在释放锁后进行，避免持锁阻塞于网络 IO。
func (m *SessionManager) PushToUser(ctx context.Context, userID int64, t proto.MessageType, msg bin.Encoder) (int, error) {
return m.PushToUserExceptAuthKeySession(ctx, userID, [8]byte{}, 0, t, msg)
}
// PushToUserExceptSession 向某 user 所有活跃连接推送，但跳过指定 session。
// 未完成 updates 同步入口的 session 会先暂存，等 SetReceivesUpdates(true) 后再发。
func (m *SessionManager) PushToUserExceptSession(ctx context.Context, userID, excludeSessionID int64, t proto.MessageType, msg bin.Encoder) (int, error) {
return m.pushToUser(ctx, userID, nil, excludeSessionID, t, msg)
}
// PushToUserExceptAuthKeySession 向某 user 所有活跃连接推送，跳过指定 raw auth_key + session。
func (m *SessionManager) PushToUserExceptAuthKeySession(ctx context.Context, userID int64, excludeAuthKeyID [8]byte, excludeSessionID int64, t proto.MessageType, msg bin.Encoder) (int, error) {
return m.pushToUser(ctx, userID, &excludeAuthKeyID, excludeSessionID, t, msg)
@ -1275,10 +1072,6 @@ func (m *SessionManager) PushToUserTransientExceptAuthKeySession(ctx context.Con
})
}
func (m *SessionManager) PushToUserExceptSessionBestEffort(ctx context.Context, userID, excludeSessionID int64, t proto.MessageType, msg bin.Encoder, timeout time.Duration) (int, error) {
return m.pushToUserBestEffort(ctx, userID, nil, excludeSessionID, t, msg, timeout)
}
func (m *SessionManager) PushToUserExceptAuthKeySessionBestEffort(ctx context.Context, userID int64, excludeAuthKeyID [8]byte, excludeSessionID int64, t proto.MessageType, msg bin.Encoder, timeout time.Duration) (int, error) {
return m.pushToUserBestEffort(ctx, userID, &excludeAuthKeyID, excludeSessionID, t, msg, timeout)
}
@ -1490,13 +1283,6 @@ func (m *SessionManager) pushToUserWithSender(ctx context.Context, userID int64,
return sent + queued, firstErr
}
// Online 返回当前活跃连接数。
func (m *SessionManager) Online() int {
m.mu.RLock()
defer m.mu.RUnlock()
return len(m.bySession)
}
// ActiveRawAuthKeyIDs 返回当前物理连接实际使用的 raw auth_key_id 去重快照。
// maintenance 用它保护“已建 key 但尚未登录”的长连接不被 orphan GC 删除；不能用
// business/temp→perm key 替代，否则活跃 temp 连接仍可能误删。
@ -1757,25 +1543,6 @@ func (m *SessionManager) OnlineChannelIDsSnapshot() []int64 {
return out
}
// OnlineChannelIDsAfter is retained for bounded diagnostics/tests. Production recovery takes one
// OnlineChannelIDsSnapshot per generation and slices it into pages, avoiding repeated full scans.
func (m *SessionManager) OnlineChannelIDsAfter(afterChannelID int64, limit int) []int64 {
if limit <= 0 {
return nil
}
const maxRecoveryPage = 4096
if limit > maxRecoveryPage {
limit = maxRecoveryPage
}
all := m.OnlineChannelIDsSnapshot()
start := sort.Search(len(all), func(i int) bool { return all[i] > afterChannelID })
end := start + limit
if end > len(all) {
end = len(all)
}
return all[start:end]
}
func (m *SessionManager) onlineChannelUsers(index map[int64]map[sessionKey]int64, channelID int64, limit int) []int64 {
if channelID == 0 {
return nil
@ -1813,7 +1580,6 @@ func (m *SessionManager) removeLocked(c *Conn, dropPending bool) int64 {
return 0
}
delete(m.bySession, key)
removeSessionIDIndex(m.bySessionID, c.sessionID, c.authKeyID)
removeConnIndex(m.byAuthKey, c.authKeyID, c.sessionID)
if businessAuthKeyID, resolved := c.BusinessAuthKeyID(); resolved {
removeBusinessAuthKeyIndex(m.byBusinessAuthKey, businessAuthKeyID, key)
@ -2100,44 +1866,6 @@ func (m *SessionManager) queuePreparedLocked(key sessionKey, t proto.MessageType
return true
}
// queueLocked remains as a test/internal single-target convenience. Production fan-out prepares
// outside m.mu and calls queuePreparedLocked so TL encoding never serializes the session registry.
func (m *SessionManager) queueLocked(key sessionKey, t proto.MessageType, msg bin.Encoder) bool {
encoded, reservation, err := m.preparePendingPush(context.Background(), msg)
if err != nil {
m.log.Debug("Drop pending push outside byte budget",
zap.String("auth_key_id", sessionKeyLog(key.authKeyID)),
zap.Int64("session_id", key.sessionID),
zap.Error(err),
)
return false
}
defer reservation.release()
return m.queuePreparedLocked(key, t, encoded, reservation)
}
func (m *SessionManager) uniqueSessionLocked(sessionID int64) (*Conn, sessionKey, bool, bool) {
set := m.bySessionID[sessionID]
if len(set) == 0 {
return nil, sessionKey{}, false, false
}
if len(set) > 1 {
return nil, sessionKey{}, false, true
}
for authKeyID, c := range set {
return c, sessionKey{authKeyID: authKeyID, sessionID: sessionID}, true, false
}
return nil, sessionKey{}, false, false
}
func (m *SessionManager) dropPendingBySessionLocked(sessionID int64) {
for key := range m.pending {
if key.sessionID == sessionID {
m.deletePendingLocked(key)
}
}
}
// RunPendingSweeper 周期回收长期滞留的 pending 暂存：被动老化（queueLocked/takePendingLocked）
// 只在「有新推送」或「就绪后取出」时触发，对「已注册但迟迟不调 getState、又恰好没有新推送、
// 也不断连（持续 ping 保活）」的连接无法回收其超龄 pending。本 sweeper 给出一个主动兜底，
@ -2236,24 +1964,6 @@ func removeBusinessAuthKeyIndex(idx map[[8]byte]map[sessionKey]*Conn, authKeyID
}
}
func addSessionIDIndex(idx map[int64]map[[8]byte]*Conn, sessionID int64, authKeyID [8]byte, c *Conn) {
set := idx[sessionID]
if set == nil {
set = make(map[[8]byte]*Conn)
idx[sessionID] = set
}
set[authKeyID] = c
}
func removeSessionIDIndex(idx map[int64]map[[8]byte]*Conn, sessionID int64, authKeyID [8]byte) {
if set := idx[sessionID]; set != nil {
delete(set, authKeyID)
if len(set) == 0 {
delete(idx, sessionID)
}
}
}
func addUserIndex(idx map[int64]map[sessionKey]*Conn, userID int64, key sessionKey, c *Conn) {
set := idx[userID]
if set == nil {

25
internal/mtprotoedge/session_manager_test.go
View file

@ -282,11 +282,11 @@ func TestSessionManagerBestEffortFanoutUsesOneBudgetAndDropsOnlySlowConsumers(t
if got := len(healthy.outbound); got != 1 {
t.Fatalf("healthy queued ops = %d, want 1", got)
}
if healthy.terminal.Load() {
if healthy.isRetired() {
t.Fatal("healthy session was terminalized")
}
for i, c := range slow {
if !c.terminal.Load() {
if !c.isRetired() {
t.Fatalf("slow session %d was not terminalized", i)
}
if tr := c.transport.(*closeCountingTransport); tr.closes != 1 {
@ -344,6 +344,15 @@ func TestSessionManagerScopesSameSessionIDByAuthKey(t *testing.T) {
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) {
@ -513,7 +522,7 @@ func TestForceCloseBatchTimeoutStillClosesProducerAndRPCGates(t *testing.T) {
t.Fatalf("timed batch close blocked for %v", elapsed)
}
for i, c := range conns {
if !c.terminal.Load() {
if !c.isRetired() {
t.Fatalf("connection %d producer gate remains open after batch timeout", i)
}
select {
@ -609,12 +618,12 @@ func TestPushToUserAuthKeyUsesOneDeadlineAndDropsOnlySlowPFSConnections(t *testi
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 {
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.terminal.Load(), slowOneTransport.closes, slowTwo.terminal.Load(), slowTwoTransport.closes)
slowOne.isRetired(), slowOneTransport.closes, slowTwo.isRetired(), slowTwoTransport.closes)
}
if healthy.terminal.Load() || healthyTransport.closes != 0 {
t.Fatalf("healthy connection was dropped: terminal=%v closes=%d", healthy.terminal.Load(), healthyTransport.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:
@ -834,7 +843,7 @@ func TestPendingFlushGlobalBodyPressureDoesNotTerminateHealthyConnection(t *test
// 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() {
if c.isRetired() {
t.Fatal("shared body pressure terminated a healthy pending-flush connection")
}
if !c.receivesUpdates.Load() {

2
internal/mtprotoedge/shutdown_gate_test.go
View file

@ -40,7 +40,7 @@ func TestTerminalFailurePathsCloseGatesBeforeBlockingTransportClose(t *testing.T
if tr.closes.Load() == 0 {
t.Fatal("terminal path did not enter transport.Close")
}
if !c.terminal.Load() {
if !c.isRetired() {
t.Fatal("producer terminal gate was not published before blocking Close")
}
select {

12
internal/mtprotoedge/transport_ownership_test.go
View file

@ -148,8 +148,8 @@ func TestPhysicalCloseFencesActivationPublication(t *testing.T) {
if c.publishActivation() {
t.Fatal("closed physical transport published active Conn")
}
if !c.terminal.Load() || c.lifecycleState() != connLifecycleRetired {
t.Fatalf("closed Conn terminal=%v lifecycle=%v", c.terminal.Load(), c.lifecycleState())
if !c.isRetired() {
t.Fatalf("closed Conn lifecycle=%v", c.lifecycleState())
}
c.Close()
}
@ -161,7 +161,7 @@ func TestPhysicalCloseBitPreventsActivationClaimBeforeLogicalFence(t *testing.T)
c := s.newConnWithLease(lease, newTestAuthKey(t), 73002, 1)
// Hold the binding lock so CloseAny can linearize the physical closed bit but
// cannot yet publish c.terminal. beginActivationClaim must inspect the lease
// cannot yet retire the logical Conn. beginActivationClaim must inspect the lease
// itself and refuse this otherwise-dangerous window.
owner.bindingMu.Lock()
closeDone := make(chan error, 1)
@ -174,7 +174,7 @@ func TestPhysicalCloseBitPreventsActivationClaimBeforeLogicalFence(t *testing.T)
owner.bindingMu.Unlock()
t.Fatal("CloseAny did not publish closed bit")
}
if c.terminal.Load() {
if c.isRetired() {
owner.bindingMu.Unlock()
t.Fatal("logical fence escaped held binding lock")
}
@ -191,8 +191,8 @@ func TestPhysicalCloseBitPreventsActivationClaimBeforeLogicalFence(t *testing.T)
case <-time.After(time.Second):
t.Fatal("CloseAny did not finish")
}
if !c.terminal.Load() || c.lifecycleState() != connLifecycleRetired {
t.Fatalf("logical fence terminal=%v lifecycle=%v", c.terminal.Load(), c.lifecycleState())
if !c.isRetired() {
t.Fatalf("logical fence lifecycle=%v", c.lifecycleState())
}
c.Close()
}