fix: sync scoped connection and outbox exclusion updates

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

View file

@ -277,7 +277,8 @@ func run(logger *zap.Logger) error {
// goroutine/锁竞争的定位全靠此端点。早于重负载初始化启动,连 seed/预热阶段也可剖析。
startDebugServer(ctx, cfg.DebugAddr, logger)
// 持久化依赖:先迁移 schema再建立连接。auth_key 落 PostgreSQL、session 落 Redis。
// 持久化依赖:先迁移 schema再建立连接。auth key 与业务事实落 PostgreSQL
// Redis 只承载可重建的短 TTL 状态、缓存、计数器和限流。
// 依赖由 deploy/docker-compose.yml 启动;连不上则启动失败(开发期须先 docker compose up
migrationStatus, err := postgres.MigrateAndStatus(cfg.PostgresDSN)
if err != nil {
@ -421,7 +422,6 @@ func run(logger *zap.Logger) error {
helpStore := postgres.NewHelpStore(pool)
aiComposeStore := postgres.NewAIComposeStore(pool)
tempAuthKeyStore := postgres.NewTempAuthKeyBindingStore(pool)
sessionStore := redisstore.NewSessionStore(rdb, redisstore.DefaultSessionTTL)
inlineRegistryStore := redisstore.NewInlineRegistryStore(rdb)
codeStore := redisstore.NewCodeStore(rdb)
rateLimiter := redisstore.NewRateLimiter(rdb)
@ -791,7 +791,6 @@ func run(logger *zap.Logger) error {
RSAKey: rsaKey,
RPC: router,
AuthKeys: authKeyStore,
Sessions: sessionStore,
ActiveSessions: activeSessions,
ObfuscatedTCP: true,
WebSocket: cfg.WebSocketEnable,

View file

@ -2061,7 +2061,8 @@ CREATE TABLE public.dispatch_outbox (
last_error text DEFAULT ''::text NOT NULL,
created_at timestamp with time zone DEFAULT now() NOT NULL,
updated_at timestamp with time zone DEFAULT now() NOT NULL,
CONSTRAINT dispatch_outbox_status_check CHECK (((status)::text = ANY (ARRAY[('pending'::character varying)::text, ('dispatching'::character varying)::text, ('failed'::character varying)::text])))
CONSTRAINT dispatch_outbox_status_check CHECK (((status)::text = ANY (ARRAY[('pending'::character varying)::text, ('dispatching'::character varying)::text, ('failed'::character varying)::text]))),
CONSTRAINT dispatch_outbox_exclusion_pair_check CHECK ((((exclude_auth_key_id = 0) AND (exclude_session_id = 0)) OR ((exclude_auth_key_id <> 0) AND (exclude_session_id <> 0))))
);

View file

@ -0,0 +1,2 @@
ALTER TABLE dispatch_outbox
DROP CONSTRAINT IF EXISTS dispatch_outbox_exclusion_pair_check;

View file

@ -0,0 +1,33 @@
-- Excluding the originating device requires the exact physical
-- (raw auth_key_id, session_id) tuple. Install the guard first so no new
-- half-pair can race the explicit cleanup of legacy invalid online tasks.
DO $$
BEGIN
-- Fresh databases already receive this constraint from 0001_init; upgraded
-- databases do not. Keep one migration stream valid for both shapes.
IF NOT EXISTS (
SELECT 1
FROM pg_constraint
WHERE conrelid = 'dispatch_outbox'::regclass
AND conname = 'dispatch_outbox_exclusion_pair_check'
) THEN
ALTER TABLE dispatch_outbox
ADD CONSTRAINT dispatch_outbox_exclusion_pair_check
CHECK (
(exclude_auth_key_id = 0 AND exclude_session_id = 0)
OR
(exclude_auth_key_id <> 0 AND exclude_session_id <> 0)
) NOT VALID;
END IF;
END
$$;
-- dispatch_outbox is only an online delivery task queue. Its durable source
-- remains user_update_events through the existing (target_user_id, pts) FK,
-- and the delete trigger promotes the next per-user head. Missing tuple parts
-- cannot be reconstructed safely, so remove rather than normalize bad rows.
DELETE FROM dispatch_outbox
WHERE (exclude_auth_key_id = 0) <> (exclude_session_id = 0);
ALTER TABLE dispatch_outbox
VALIDATE CONSTRAINT dispatch_outbox_exclusion_pair_check;

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 layerinvokeWithLayer/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

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,

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

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,

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 仅用于 metricssize 是在 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()
}

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)
}

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)

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 + 关闭 sockethandleOutboundOp 返回后actor 自己发停止信号并退出,
// 把 lifecycle 推进到 retired 并关闭 sockethandleOutboundOp 返回后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

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 {

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 {

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
}

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

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 {

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{}

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)

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())
}
}

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()
}

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_idtemp 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 {

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() {

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 {

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()
}

View file

@ -77,12 +77,7 @@ func (r *Router) onAuthBindTempAuthKey(ctx context.Context, req *tg.AuthBindTemp
r.tempKeyResolveCache.Delete(id)
}
if r.deps.Sessions != nil {
if scoped, ok := r.scopedSessions(); ok {
rawAuthKeyID, _ := RawAuthKeyIDFrom(ctx)
scoped.BindAuthKeyForSession(rawAuthKeyID, sessionID, authKeyIDFromInt64(req.PermAuthKeyID))
} else {
r.deps.Sessions.BindAuthKey(sessionID, authKeyIDFromInt64(req.PermAuthKeyID))
}
r.deps.Sessions.BindAuthKeyForSession(id, sessionID, authKeyIDFromInt64(req.PermAuthKeyID))
}
r.invalidateAuthUserCache(id)
return true, nil
@ -192,14 +187,8 @@ func (r *Router) bindLoginTokenTarget(target loginTokenTarget, userID int64) {
if r.deps.Sessions == nil || target.sessionID == 0 {
return
}
if scoped, ok := r.scopedSessions(); ok && target.rawAuthKeyID != ([8]byte{}) {
scoped.BindAuthKeyForSession(target.rawAuthKeyID, target.sessionID, target.authKeyID)
scoped.BindUserForAuthKey(target.rawAuthKeyID, target.sessionID, userID)
r.announceSessionOnline(loginTokenTargetContext(target, userID), userID)
return
}
r.deps.Sessions.BindAuthKey(target.sessionID, target.authKeyID)
r.deps.Sessions.BindUser(target.sessionID, userID)
r.deps.Sessions.BindAuthKeyForSession(target.rawAuthKeyID, target.sessionID, target.authKeyID)
r.deps.Sessions.BindUserForAuthKey(target.rawAuthKeyID, target.sessionID, userID)
r.announceSessionOnline(loginTokenTargetContext(target, userID), userID)
}
@ -220,19 +209,13 @@ func (r *Router) pushLoginTokenAccepted(ctx context.Context, target loginTokenTa
Update: &tg.UpdateLoginToken{},
Date: int(r.clock.Now().Unix()),
}
if immediate, ok := r.deps.Sessions.(ScopedImmediateSessionPusher); ok && target.rawAuthKeyID != ([8]byte{}) {
if immediate, ok := r.deps.Sessions.(ImmediateSessionPusher); ok {
if err := immediate.PushToSessionForAuthKeyImmediate(ctx, target.rawAuthKeyID, target.sessionID, proto.MessageFromServer, updates); err != nil {
r.log.Debug("push login token accepted immediate", zap.Int64("session_id", target.sessionID), zap.Error(err))
}
return
}
if scoped, ok := r.scopedSessions(); ok && target.rawAuthKeyID != ([8]byte{}) {
if err := scoped.PushToSessionForAuthKey(ctx, target.rawAuthKeyID, target.sessionID, proto.MessageFromServer, updates); err != nil {
r.log.Debug("push login token accepted", zap.Int64("session_id", target.sessionID), zap.Error(err))
}
return
}
if err := r.deps.Sessions.PushToSession(ctx, target.sessionID, proto.MessageFromServer, updates); err != nil {
if err := r.deps.Sessions.PushToSessionForAuthKey(ctx, target.rawAuthKeyID, target.sessionID, proto.MessageFromServer, updates); err != nil {
r.log.Debug("push login token accepted", zap.Int64("session_id", target.sessionID), zap.Error(err))
}
}
@ -713,13 +696,7 @@ func (r *Router) bindSessionUser(ctx context.Context, userID int64) {
if !ok {
return
}
if scoped, ok := r.scopedSessions(); ok {
rawAuthKeyID, _ := RawAuthKeyIDFrom(ctx)
scoped.BindUserForAuthKey(rawAuthKeyID, sessionID, userID)
r.announceSessionOnline(ctx, userID)
return
}
r.deps.Sessions.BindUser(sessionID, userID)
r.deps.Sessions.BindUserForAuthKey(rawAuthKeyIDForOrigin(ctx), sessionID, userID)
r.announceSessionOnline(ctx, userID)
}
@ -776,13 +753,7 @@ func (r *Router) pushSignInServiceNotificationToOthers(ctx context.Context, u do
go func() {
pushCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
if scoped, ok := r.scopedSessions(); ok {
if sent, err := scoped.PushToUserExceptAuthKeySession(pushCtx, u.ID, rawAuthKeyID, sessionID, proto.MessageFromServer, notification); err != nil {
r.log.Debug("push sign-in service notification", zap.Int64("user_id", u.ID), zap.Int("sent", sent), zap.Error(err))
}
return
}
if sent, err := r.deps.Sessions.PushToUserExceptSession(pushCtx, u.ID, sessionID, proto.MessageFromServer, notification); err != nil {
if sent, err := r.deps.Sessions.PushToUserExceptAuthKeySession(pushCtx, u.ID, rawAuthKeyID, sessionID, proto.MessageFromServer, notification); err != nil {
r.log.Debug("push sign-in service notification", zap.Int64("user_id", u.ID), zap.Int("sent", sent), zap.Error(err))
}
}()

View file

@ -65,7 +65,7 @@ func newRecoveryFanoutSessions(onlineChannels []int64, release <-chan struct{})
}
}
func (s *recoveryFanoutSessions) PushToUserExceptSession(ctx context.Context, _ int64, _ int64, _ proto.MessageType, msg bin.Encoder) (int, error) {
func (s *recoveryFanoutSessions) PushToUserExceptAuthKeySession(ctx context.Context, _ int64, _ [8]byte, _ int64, _ proto.MessageType, msg bin.Encoder) (int, error) {
s.startOnce.Do(func() { close(s.pushStarted) })
if s.pushRelease != nil {
select {
@ -397,7 +397,7 @@ func newOverflowNudgeSessions(onlineByChannel map[int64][]int64) *overflowNudgeS
}
}
func (s *overflowNudgeSessions) PushToUserExceptSession(ctx context.Context, userID, excludeSessionID int64, typ proto.MessageType, msg bin.Encoder) (int, error) {
func (s *overflowNudgeSessions) PushToUserExceptAuthKeySession(ctx context.Context, userID int64, excludeAuthKeyID [8]byte, excludeSessionID int64, typ proto.MessageType, msg bin.Encoder) (int, error) {
if updates, ok := msg.(*tg.Updates); ok && len(updates.Updates) == 1 {
if nudge, ok := updates.Updates[0].(*tg.UpdateChannelTooLong); ok {
pts, _ := nudge.GetPts()
@ -407,7 +407,7 @@ func (s *overflowNudgeSessions) PushToUserExceptSession(ctx context.Context, use
s.mu.Unlock()
}
}
return s.captureSessions.PushToUserExceptSession(ctx, userID, excludeSessionID, typ, msg)
return s.captureSessions.PushToUserExceptAuthKeySession(ctx, userID, excludeAuthKeyID, excludeSessionID, typ, msg)
}
func (s *overflowNudgeSessions) OnlineChannelMemberUserIDsExcluding(channelID int64, exclude map[int64]struct{}, limit int) []int64 {
@ -445,11 +445,11 @@ func newNudgeSessions(online []int64) *nudgeSessions {
return &nudgeSessions{captureSessions: &captureSessions{}, online: online, byUser: map[int64]bin.Encoder{}}
}
func (s *nudgeSessions) PushToUserExceptSession(ctx context.Context, userID, excludeSessionID int64, t proto.MessageType, msg bin.Encoder) (int, error) {
func (s *nudgeSessions) PushToUserExceptAuthKeySession(ctx context.Context, userID int64, excludeAuthKeyID [8]byte, excludeSessionID int64, t proto.MessageType, msg bin.Encoder) (int, error) {
s.mu.Lock()
s.byUser[userID] = msg
s.mu.Unlock()
return s.captureSessions.PushToUserExceptSession(ctx, userID, excludeSessionID, t, msg)
return s.captureSessions.PushToUserExceptAuthKeySession(ctx, userID, excludeAuthKeyID, excludeSessionID, t, msg)
}
func (s *nudgeSessions) OnlineChannelMemberUserIDsExcluding(_ int64, exclude map[int64]struct{}, limit int) []int64 {

View file

@ -29,26 +29,12 @@ func (r *Router) currentUserID(ctx context.Context) (int64, bool, error) {
}
if r.deps.Sessions != nil {
if sessionID, ok := SessionIDFrom(ctx); ok {
if scoped, ok := r.scopedSessions(); ok {
if rawAuthKeyID, ok := RawAuthKeyIDFrom(ctx); ok {
if userID, resolved := scoped.UserIDResolvedForAuthKey(rawAuthKeyID, sessionID); resolved {
if userID == 0 {
if authKeyID, ok := AuthKeyIDFrom(ctx); ok {
if cachedUserID, ok := r.positiveCachedAuthUser(authKeyID); ok {
scoped.BindUserForAuthKey(rawAuthKeyID, sessionID, cachedUserID)
r.announceSessionOnline(ctx, cachedUserID)
return cachedUserID, true, nil
}
}
}
return userID, userID != 0, nil
}
}
} else if userID, resolved := r.deps.Sessions.UserIDResolved(sessionID); resolved {
rawAuthKeyID := rawAuthKeyIDForOrigin(ctx)
if userID, resolved := r.deps.Sessions.UserIDResolvedForAuthKey(rawAuthKeyID, sessionID); resolved {
if userID == 0 {
if authKeyID, ok := AuthKeyIDFrom(ctx); ok {
if cachedUserID, ok := r.positiveCachedAuthUser(authKeyID); ok {
r.deps.Sessions.BindUser(sessionID, cachedUserID)
r.deps.Sessions.BindUserForAuthKey(rawAuthKeyID, sessionID, cachedUserID)
r.announceSessionOnline(ctx, cachedUserID)
return cachedUserID, true, nil
}

View file

@ -49,36 +49,26 @@ type AuthService interface {
}
// SessionBinder 抽象登录后 session 与 user 的在线绑定。
//
// MTProto session 的完整身份是 raw auth_key_id + session_id。所有定位单个 session
// 的方法都必须携带这两个值;禁止退回只按 session_id 查询,否则不同 auth key 复用
// 同一随机 session_id 时会产生跨账号绑定、排除或推送歧义。
type SessionBinder interface {
BindAuthKey(sessionID int64, authKeyID [8]byte)
AuthKeyID(sessionID int64) ([8]byte, bool)
BindUser(sessionID, userID int64)
UserID(sessionID int64) (int64, bool)
UserIDResolved(sessionID int64) (userID int64, resolved bool)
UnbindAuthKey(authKeyID [8]byte) int
SetReceivesUpdates(sessionID int64, receives bool)
PushToSession(ctx context.Context, sessionID int64, t proto.MessageType, msg bin.Encoder) error
PushToUserExceptSession(ctx context.Context, userID, excludeSessionID int64, t proto.MessageType, msg bin.Encoder) (int, error)
}
// ScopedSessionBinder 是 SessionBinder 的精确版本:所有定位都带 raw auth_key_id + session_id。
// 生产 mtprotoedge.SessionManager 实现它;测试替身和旧实现可以只实现 SessionBinder。
type ScopedSessionBinder interface {
BindAuthKeyForSession(rawAuthKeyID [8]byte, sessionID int64, authKeyID [8]byte)
AuthKeyIDForSession(rawAuthKeyID [8]byte, sessionID int64) ([8]byte, bool)
BindUserForAuthKey(rawAuthKeyID [8]byte, sessionID, userID int64)
UserIDForAuthKey(rawAuthKeyID [8]byte, sessionID int64) (int64, bool)
UserIDResolvedForAuthKey(rawAuthKeyID [8]byte, sessionID int64) (userID int64, resolved bool)
UnbindAuthKey(authKeyID [8]byte) int
SetReceivesUpdatesForAuthKey(rawAuthKeyID [8]byte, sessionID int64, receives bool)
PushToSessionForAuthKey(ctx context.Context, rawAuthKeyID [8]byte, sessionID int64, t proto.MessageType, msg bin.Encoder) error
// excludeAuthKeyID is the physical/raw auth key, paired with session_id.
// excludeAuthKeyID/excludeSessionID 必须同时为零(不排除)或同时非零(精确排除)。
PushToUserExceptAuthKeySession(ctx context.Context, userID int64, excludeAuthKeyID [8]byte, excludeSessionID int64, t proto.MessageType, msg bin.Encoder) (int, error)
}
// ScopedImmediateSessionPusher 是可选的登录前信号直推能力。
// ImmediateSessionPusher 是可选的登录前信号直推能力。
// 它绕过登录后 updates-ready 队列,只能用于会解锁登录流程本身的握手消息,
// 例如 updateLoginToken。
type ScopedImmediateSessionPusher interface {
type ImmediateSessionPusher interface {
PushToSessionForAuthKeyImmediate(ctx context.Context, rawAuthKeyID [8]byte, sessionID int64, t proto.MessageType, msg bin.Encoder) error
}
@ -110,13 +100,9 @@ type RawSessionTerminator interface {
CloseSessionsForRawAuthKeyExcept(authKeyID [8]byte, exceptSessionID int64) int
}
// BestEffortSessionBinder 是 updates fanout 的短超时推送接口;不用于 RPC result/ack。
// BestEffortSessionBinder 是带 raw auth_key_id 精确排除当前设备的短超时推送接口;
// 不用于 RPC result/ack。
type BestEffortSessionBinder interface {
PushToUserExceptSessionBestEffort(ctx context.Context, userID, excludeSessionID int64, t proto.MessageType, msg bin.Encoder, timeout time.Duration) (int, error)
}
// ScopedBestEffortSessionBinder 是带 raw auth_key_id 精确排除当前设备的 best-effort 版本。
type ScopedBestEffortSessionBinder interface {
PushToUserExceptAuthKeySessionBestEffort(ctx context.Context, userID int64, excludeAuthKeyID [8]byte, excludeSessionID int64, t proto.MessageType, msg bin.Encoder, timeout time.Duration) (int, error)
}

View file

@ -29,7 +29,10 @@ const (
defaultOutboxMaxIdleInterval = 1 * time.Second
)
var errMissingOutboxEvent = errors.New("missing outbox update event")
var (
errMissingOutboxEvent = errors.New("missing outbox update event")
errInvalidOutboxExclusionPair = errors.New("outbox exclusion requires both raw auth key and session id")
)
// OutboxDispatcher 把 PG transactional outbox 中的 update 批量推给在线 session。
// 多 worker 并发 claimClaimPending 用 FOR UPDATE SKIP LOCKEDworker 间认领不重叠。
@ -273,6 +276,23 @@ type outboxEventKey struct {
// dispatchBatch 批量加载已 claim 事件、逐条 push、批量标记 delivered失败项单独退避重试。
func (d *OutboxDispatcher) dispatchBatch(ctx context.Context, items []store.DispatchOutboxItem, loader batchEventLoader, marker batchOutboxMarker) {
valid := make([]store.DispatchOutboxItem, 0, len(items))
blockedUsers := make(map[int64]struct{})
for _, item := range items {
if _, blocked := blockedUsers[item.TargetUserID]; blocked {
continue
}
if err := validateOutboxExclusionPair(item); err != nil {
d.markDispatchFailed(ctx, item, err)
blockedUsers[item.TargetUserID] = struct{}{}
continue
}
valid = append(valid, item)
}
if len(valid) == 0 {
return
}
items = valid
cursors := make([]store.EventCursor, len(items))
for i, item := range items {
cursors[i] = store.EventCursor{UserID: item.TargetUserID, Pts: item.Pts}
@ -299,7 +319,7 @@ func (d *OutboxDispatcher) dispatchBatch(ctx context.Context, items []store.Disp
start := time.Now()
ready := make([]outboxDispatchReady, 0, len(items))
requests := make([]OutboxUpdateRequest, 0, len(items))
blockedUsers := make(map[int64]struct{})
clear(blockedUsers)
for _, item := range items {
if _, blocked := blockedUsers[item.TargetUserID]; blocked {
continue
@ -362,6 +382,10 @@ type outboxDispatchReady struct {
func (d *OutboxDispatcher) dispatchItem(ctx context.Context, item store.DispatchOutboxItem) bool {
start := time.Now()
if err := validateOutboxExclusionPair(item); err != nil {
d.markDispatchFailed(ctx, item, err)
return false
}
events, err := d.events.ListAfter(ctx, item.TargetUserID, item.Pts-1, 1)
if err != nil {
d.markDispatchFailed(ctx, item, err)
@ -443,23 +467,29 @@ func (d *OutboxDispatcher) buildOutboxUpdates(ctx context.Context, requests []Ou
// 接口剩余的非 context 错误通常是确定性的编码/构造错误,必须进入 failed不能永久占着
// dispatching head 靠租约空转。只有 dispatcher shutdown/deadline 属于可重试中断。
func (d *OutboxDispatcher) pushOutboxUpdate(ctx context.Context, item store.DispatchOutboxItem, update *tg.Updates) (sent int, retriable bool, err error) {
var zeroAuthKeyID [8]byte
if err := validateOutboxExclusionPair(item); err != nil {
return 0, false, errInvalidOutboxExclusionPair
}
if d.pushTimeout > 0 {
if scoped, ok := d.sessions.(ScopedBestEffortSessionBinder); ok && item.ExcludeAuthKeyID != zeroAuthKeyID {
sent, err = scoped.PushToUserExceptAuthKeySessionBestEffort(ctx, item.TargetUserID, item.ExcludeAuthKeyID, item.ExcludeSessionID, proto.MessageFromServer, update, d.pushTimeout)
return sent, outboxPushInterrupted(err), err
}
if bestEffort, ok := d.sessions.(BestEffortSessionBinder); ok {
sent, err = bestEffort.PushToUserExceptSessionBestEffort(ctx, item.TargetUserID, item.ExcludeSessionID, proto.MessageFromServer, update, d.pushTimeout)
sent, err = bestEffort.PushToUserExceptAuthKeySessionBestEffort(ctx, item.TargetUserID, item.ExcludeAuthKeyID, item.ExcludeSessionID, proto.MessageFromServer, update, d.pushTimeout)
return sent, outboxPushInterrupted(err), err
}
}
if scoped, ok := d.sessions.(ScopedSessionBinder); ok && item.ExcludeAuthKeyID != zeroAuthKeyID {
sent, err = scoped.PushToUserExceptAuthKeySession(ctx, item.TargetUserID, item.ExcludeAuthKeyID, item.ExcludeSessionID, proto.MessageFromServer, update)
return sent, false, err
sent, err = d.sessions.PushToUserExceptAuthKeySession(ctx, item.TargetUserID, item.ExcludeAuthKeyID, item.ExcludeSessionID, proto.MessageFromServer, update)
if err != nil {
return sent, outboxPushInterrupted(err), err
}
sent, err = d.sessions.PushToUserExceptSession(ctx, item.TargetUserID, item.ExcludeSessionID, proto.MessageFromServer, update)
return sent, false, err
return sent, false, nil
}
func validateOutboxExclusionPair(item store.DispatchOutboxItem) error {
var zeroAuthKeyID [8]byte
if (item.ExcludeAuthKeyID != zeroAuthKeyID) != (item.ExcludeSessionID != 0) {
return errInvalidOutboxExclusionPair
}
return nil
}
func outboxPushInterrupted(err error) bool {

View file

@ -33,6 +33,7 @@ func TestOutboxDispatcherPushesNewMessageAndMarksDelivered(t *testing.T) {
TargetUserID: msg.OwnerUserID,
Pts: msg.Pts,
EventType: domain.UpdateEventNewMessage,
ExcludeAuthKeyID: [8]byte{1},
ExcludeSessionID: 99,
}}}
events := &captureUpdateEventStore{events: []domain.UpdateEvent{{
@ -103,6 +104,78 @@ func TestOutboxDispatcherUsesScopedAuthKeyExclusion(t *testing.T) {
}
}
func TestOutboxDispatcherRejectsPartialSessionExclusion(t *testing.T) {
tests := []struct {
name string
authKeyID [8]byte
sessionID int64
}{
{name: "auth key only", authKeyID: [8]byte{1}},
{name: "session only", sessionID: 99},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
const userID = int64(1000000002)
outbox := &captureDispatchOutbox{items: []store.DispatchOutboxItem{{
ID: 58,
TargetUserID: userID,
Pts: 10,
EventType: domain.UpdateEventPeerSettings,
ExcludeAuthKeyID: tt.authKeyID,
ExcludeSessionID: tt.sessionID,
}}}
// No event exists: exclusion shape must win before event loading so a bad
// durable row is not mislabeled as merely missing its payload.
events := &captureUpdateEventStore{}
sessions := &captureSessions{}
metrics := &captureOutboxMetrics{}
dispatcher := NewOutboxDispatcher(events, outbox, sessions, zaptest.NewLogger(t), WithOutboxMetrics(metrics))
dispatcher.DispatchOnce(context.Background())
if !outbox.failed || outbox.delivered {
t.Fatalf("outbox failed=%v delivered=%v, want failed without delivery", outbox.failed, outbox.delivered)
}
if outbox.failedError != errInvalidOutboxExclusionPair.Error() {
t.Fatalf("failed error = %q, want %q", outbox.failedError, errInvalidOutboxExclusionPair)
}
if sessions.message != nil {
t.Fatalf("invalid exclusion unexpectedly pushed %T", sessions.message)
}
if metrics.failed != 1 || metrics.delivered != 0 {
t.Fatalf("metrics failed=%d delivered=%d, want 1/0", metrics.failed, metrics.delivered)
}
})
}
}
func TestOutboxDispatcherBatchRejectsPartialExclusionBeforeNoop(t *testing.T) {
const userID = int64(1000000002)
events := &batchEventStore{captureUpdateEventStore: &captureUpdateEventStore{events: []domain.UpdateEvent{{
UserID: userID,
Type: domain.UpdateEventNoop,
Pts: 10,
}}}}
outbox := &batchDispatchOutbox{captureDispatchOutbox: &captureDispatchOutbox{items: []store.DispatchOutboxItem{{
ID: 59,
TargetUserID: userID,
Pts: 10,
EventType: domain.UpdateEventNoop,
ExcludeAuthKeyID: [8]byte{1},
}}}}
dispatcher := NewOutboxDispatcher(events, outbox, &captureSessions{}, zaptest.NewLogger(t))
dispatcher.DispatchOnce(context.Background())
if !outbox.failed || outbox.delivered || len(outbox.deliveredBatch) != 0 {
t.Fatalf("batch invalid pair failed=%v delivered=%v batch=%v", outbox.failed, outbox.delivered, outbox.deliveredBatch)
}
if outbox.failedError != errInvalidOutboxExclusionPair.Error() {
t.Fatalf("failed error = %q, want %q", outbox.failedError, errInvalidOutboxExclusionPair)
}
if len(events.batchCursors) != 0 {
t.Fatalf("invalid pair reached batch event loader: %+v", events.batchCursors)
}
}
// TestOutboxDispatcherBatchPath 覆盖生产批量路径store 同时具备 BatchByCursor + MarkDeliveredBatch
// 时DispatchOnce 一次批量取事件、推送、再批量标记 delivered而非逐条。
func TestOutboxDispatcherBatchPath(t *testing.T) {
@ -129,6 +202,7 @@ func TestOutboxDispatcherBatchPath(t *testing.T) {
TargetUserID: msg.OwnerUserID,
Pts: msg.Pts,
EventType: domain.UpdateEventNewMessage,
ExcludeAuthKeyID: [8]byte{1},
ExcludeSessionID: 99,
}}}}
sessions := &captureSessions{}
@ -661,22 +735,34 @@ func TestOutboxDispatcherUsesBestEffortPush(t *testing.T) {
Message: msg,
Users: []domain.User{{ID: msg.From.ID, FirstName: "Sender"}},
}}}
outbox := &captureDispatchOutbox{items: []store.DispatchOutboxItem{{
ID: 55,
TargetUserID: msg.OwnerUserID,
Pts: msg.Pts,
EventType: domain.UpdateEventNewMessage,
ExcludeSessionID: 99,
}}}
sessions := &captureBestEffortSessions{captureSessions: &captureSessions{}}
dispatcher := NewOutboxDispatcher(events, outbox, sessions, zaptest.NewLogger(t), WithOutboxPushTimeout(50*time.Millisecond))
dispatcher.DispatchOnce(context.Background())
for _, tt := range []struct {
name string
authKeyID [8]byte
sessionID int64
}{
{name: "exclude origin", authKeyID: [8]byte{1}, sessionID: 99},
{name: "exclude none"},
} {
t.Run(tt.name, func(t *testing.T) {
outbox := &captureDispatchOutbox{items: []store.DispatchOutboxItem{{
ID: 55,
TargetUserID: msg.OwnerUserID,
Pts: msg.Pts,
EventType: domain.UpdateEventNewMessage,
ExcludeAuthKeyID: tt.authKeyID,
ExcludeSessionID: tt.sessionID,
}}}
sessions := &captureBestEffortSessions{captureSessions: &captureSessions{}}
dispatcher := NewOutboxDispatcher(events, outbox, sessions, zaptest.NewLogger(t), WithOutboxPushTimeout(50*time.Millisecond))
dispatcher.DispatchOnce(context.Background())
if !sessions.bestEffort || sessions.timeout != 50*time.Millisecond {
t.Fatalf("best-effort push = %v timeout %v, want true/50ms", sessions.bestEffort, sessions.timeout)
}
if !outbox.delivered || outbox.failed {
t.Fatalf("outbox delivered=%v failed=%v, want delivered after accepted best-effort push", outbox.delivered, outbox.failed)
if !sessions.bestEffort || sessions.timeout != 50*time.Millisecond {
t.Fatalf("best-effort push = %v timeout %v, want true/50ms", sessions.bestEffort, sessions.timeout)
}
if !outbox.delivered || outbox.failed {
t.Fatalf("outbox delivered=%v failed=%v, want delivered after accepted best-effort push", outbox.delivered, outbox.failed)
}
})
}
}
@ -686,10 +772,10 @@ type captureBestEffortSessions struct {
timeout time.Duration
}
func (s *captureBestEffortSessions) PushToUserExceptSessionBestEffort(ctx context.Context, userID, excludeSessionID int64, t proto.MessageType, msg bin.Encoder, timeout time.Duration) (int, error) {
func (s *captureBestEffortSessions) PushToUserExceptAuthKeySessionBestEffort(ctx context.Context, userID int64, excludeAuthKeyID [8]byte, excludeSessionID int64, t proto.MessageType, msg bin.Encoder, timeout time.Duration) (int, error) {
s.bestEffort = true
s.timeout = timeout
return s.PushToUserExceptSession(ctx, userID, excludeSessionID, t, msg)
return s.PushToUserExceptAuthKeySession(ctx, userID, excludeAuthKeyID, excludeSessionID, t, msg)
}
type orderedOutboxCaptureSessions struct {
@ -709,7 +795,7 @@ type selectiveFailOutboxSessions struct {
attempts []outboxPushAttempt
}
func (s *selectiveFailOutboxSessions) PushToUserExceptSession(_ context.Context, userID, excludeSessionID int64, t proto.MessageType, msg bin.Encoder) (int, error) {
func (s *selectiveFailOutboxSessions) PushToUserExceptAuthKeySession(_ context.Context, userID int64, excludeAuthKeyID [8]byte, excludeSessionID int64, t proto.MessageType, msg bin.Encoder) (int, error) {
pts := 0
if updates, ok := msg.(*tg.Updates); ok {
pts = firstOutboxUpdatePts(updates)
@ -718,18 +804,18 @@ func (s *selectiveFailOutboxSessions) PushToUserExceptSession(_ context.Context,
if userID == s.failUserID && pts == s.failPts {
return 0, errors.New("injected outbox push failure")
}
return s.captureSessions.PushToUserExceptSession(context.Background(), userID, excludeSessionID, t, msg)
return s.captureSessions.PushToUserExceptAuthKeySession(context.Background(), userID, excludeAuthKeyID, excludeSessionID, t, msg)
}
func (s *selectiveFailOutboxSessions) pushAttempts() []outboxPushAttempt {
return append([]outboxPushAttempt(nil), s.attempts...)
}
func (s *orderedOutboxCaptureSessions) PushToUserExceptSession(_ context.Context, userID, excludeSessionID int64, t proto.MessageType, msg bin.Encoder) (int, error) {
func (s *orderedOutboxCaptureSessions) PushToUserExceptAuthKeySession(_ context.Context, userID int64, excludeAuthKeyID [8]byte, excludeSessionID int64, t proto.MessageType, msg bin.Encoder) (int, error) {
if updates, ok := msg.(*tg.Updates); ok {
s.pushed = append(s.pushed, firstOutboxUpdatePts(updates))
}
return s.captureSessions.PushToUserExceptSession(context.Background(), userID, excludeSessionID, t, msg)
return s.captureSessions.PushToUserExceptAuthKeySession(context.Background(), userID, excludeAuthKeyID, excludeSessionID, t, msg)
}
func (s *orderedOutboxCaptureSessions) pushedPts() []int {
@ -916,32 +1002,30 @@ func (s *captureScopedSessions) immediatePushSnapshot() (proto.MessageType, bin.
}
func (s *captureScopedSessions) BindAuthKeyForSession(rawAuthKeyID [8]byte, sessionID int64, authKeyID [8]byte) {
s.BindAuthKey(sessionID, authKeyID)
s.captureSessions.BindAuthKeyForSession(rawAuthKeyID, sessionID, authKeyID)
s.setScopedAuthKeyID(rawAuthKeyID)
}
func (s *captureScopedSessions) AuthKeyIDForSession([8]byte, int64) ([8]byte, bool) {
return s.AuthKeyID(0)
func (s *captureScopedSessions) AuthKeyIDForSession(rawAuthKeyID [8]byte, sessionID int64) ([8]byte, bool) {
return s.captureSessions.AuthKeyIDForSession(rawAuthKeyID, sessionID)
}
func (s *captureScopedSessions) BindUserForAuthKey(rawAuthKeyID [8]byte, sessionID, userID int64) {
s.BindUser(sessionID, userID)
s.captureSessions.BindUserForAuthKey(rawAuthKeyID, sessionID, userID)
s.setScopedAuthKeyID(rawAuthKeyID)
}
func (s *captureScopedSessions) UserIDForAuthKey([8]byte, int64) (int64, bool) {
return s.UserID(0)
func (s *captureScopedSessions) UserIDResolvedForAuthKey(rawAuthKeyID [8]byte, sessionID int64) (int64, bool) {
return s.captureSessions.UserIDResolvedForAuthKey(rawAuthKeyID, sessionID)
}
func (s *captureScopedSessions) UserIDResolvedForAuthKey([8]byte, int64) (int64, bool) {
return s.UserIDResolved(0)
func (s *captureScopedSessions) SetReceivesUpdatesForAuthKey(rawAuthKeyID [8]byte, sessionID int64, receives bool) {
s.captureSessions.SetReceivesUpdatesForAuthKey(rawAuthKeyID, sessionID, receives)
}
func (s *captureScopedSessions) SetReceivesUpdatesForAuthKey([8]byte, int64, bool) {}
func (s *captureScopedSessions) PushToSessionForAuthKey(_ context.Context, rawAuthKeyID [8]byte, sessionID int64, t proto.MessageType, msg bin.Encoder) error {
s.setScopedAuthKeyID(rawAuthKeyID)
return s.PushToSession(context.Background(), sessionID, t, msg)
return s.captureSessions.PushToSessionForAuthKey(context.Background(), rawAuthKeyID, sessionID, t, msg)
}
func (s *captureScopedSessions) PushToSessionForAuthKeyImmediate(_ context.Context, rawAuthKeyID [8]byte, sessionID int64, t proto.MessageType, msg bin.Encoder) error {
@ -951,12 +1035,12 @@ func (s *captureScopedSessions) PushToSessionForAuthKeyImmediate(_ context.Conte
s.immediateType = t
s.immediateMsg = msg
s.scopedMu.Unlock()
return s.PushToSession(context.Background(), sessionID, t, msg)
return s.captureSessions.PushToSessionForAuthKey(context.Background(), rawAuthKeyID, sessionID, t, msg)
}
func (s *captureScopedSessions) PushToUserExceptAuthKeySession(_ context.Context, userID int64, excludeAuthKeyID [8]byte, excludeSessionID int64, t proto.MessageType, msg bin.Encoder) (int, error) {
s.setScopedAuthKeyID(excludeAuthKeyID)
return s.PushToUserExceptSession(context.Background(), userID, excludeSessionID, t, msg)
return s.captureSessions.PushToUserExceptAuthKeySession(context.Background(), userID, excludeAuthKeyID, excludeSessionID, t, msg)
}
func (s *captureDispatchOutbox) ClaimPending(context.Context, int) ([]store.DispatchOutboxItem, error) {
@ -1035,7 +1119,7 @@ type interruptedBestEffortSessions struct {
attempts int
}
func (s *interruptedBestEffortSessions) PushToUserExceptSessionBestEffort(_ context.Context, _ int64, _ int64, _ proto.MessageType, _ bin.Encoder, _ time.Duration) (int, error) {
func (s *interruptedBestEffortSessions) PushToUserExceptAuthKeySessionBestEffort(_ context.Context, _ int64, _ [8]byte, _ int64, _ proto.MessageType, _ bin.Encoder, _ time.Duration) (int, error) {
s.attempts++
return 0, context.DeadlineExceeded
}
@ -1065,6 +1149,7 @@ func TestOutboxDispatcherDefersOnPushInterruption(t *testing.T) {
TargetUserID: msg.OwnerUserID,
Pts: msg.Pts,
EventType: domain.UpdateEventNewMessage,
ExcludeAuthKeyID: [8]byte{1},
ExcludeSessionID: 99,
}}}
sessions := &interruptedBestEffortSessions{captureSessions: &captureSessions{}}

View file

@ -63,11 +63,9 @@ func (r *Router) pushPhoneSignalingData(ctx context.Context, targetUserID int64,
Date: int(r.clock.Now().Unix()),
Seq: 0,
}
if !device.Zero() {
if scoped, ok := r.scopedSessions(); ok {
if err := scoped.PushToSessionForAuthKey(ctx, device.RawAuthKeyID, device.SessionID, proto.MessageFromServer, upd); err == nil {
return
}
if !device.Zero() && r.deps.Sessions != nil {
if err := r.deps.Sessions.PushToSessionForAuthKey(ctx, device.RawAuthKeyID, device.SessionID, proto.MessageFromServer, upd); err == nil {
return
}
}
r.pushUserMessage(ctx, targetUserID, "phone call signaling", upd)

View file

@ -25,6 +25,7 @@ import (
// phonePushRecord 记录一次定向推送(目标用户、被排除的 session、载荷
type phonePushRecord struct {
userID int64
targetSession int64
excludeSession int64
msg bin.Encoder
}
@ -35,18 +36,25 @@ type phoneCaptureSessions struct {
log []phonePushRecord
}
func (s *phoneCaptureSessions) BindAuthKey(int64, [8]byte) {}
func (s *phoneCaptureSessions) AuthKeyID(int64) ([8]byte, bool) { return [8]byte{}, false }
func (s *phoneCaptureSessions) BindUser(int64, int64) {}
func (s *phoneCaptureSessions) UserID(int64) (int64, bool) { return 0, false }
func (s *phoneCaptureSessions) UserIDResolved(int64) (int64, bool) { return 0, false }
func (s *phoneCaptureSessions) UnbindAuthKey([8]byte) int { return 0 }
func (s *phoneCaptureSessions) SetReceivesUpdates(int64, bool) {}
func (s *phoneCaptureSessions) PushToSession(context.Context, int64, proto.MessageType, bin.Encoder) error {
func (s *phoneCaptureSessions) BindAuthKeyForSession([8]byte, int64, [8]byte) {}
func (s *phoneCaptureSessions) AuthKeyIDForSession([8]byte, int64) ([8]byte, bool) {
return [8]byte{}, false
}
func (s *phoneCaptureSessions) BindUserForAuthKey([8]byte, int64, int64) {}
func (s *phoneCaptureSessions) UserIDResolvedForAuthKey([8]byte, int64) (int64, bool) {
return 0, false
}
func (s *phoneCaptureSessions) UnbindAuthKey([8]byte) int { return 0 }
func (s *phoneCaptureSessions) SetReceivesUpdatesForAuthKey([8]byte, int64, bool) {}
func (s *phoneCaptureSessions) PushToSessionForAuthKey(_ context.Context, _ [8]byte, sessionID int64, _ proto.MessageType, msg bin.Encoder) error {
s.mu.Lock()
defer s.mu.Unlock()
s.log = append(s.log, phonePushRecord{targetSession: sessionID, msg: msg})
return nil
}
func (s *phoneCaptureSessions) PushToUserExceptSession(_ context.Context, userID, excludeSessionID int64, _ proto.MessageType, msg bin.Encoder) (int, error) {
func (s *phoneCaptureSessions) PushToUserExceptAuthKeySession(_ context.Context, userID int64, _ [8]byte, excludeSessionID int64, _ proto.MessageType, msg bin.Encoder) (int, error) {
s.mu.Lock()
defer s.mu.Unlock()
s.log = append(s.log, phonePushRecord{userID: userID, excludeSession: excludeSessionID, msg: msg})
@ -308,8 +316,8 @@ func TestPhoneCallRPCHappyPath(t *testing.T) {
t.Fatalf("sendSignalingData = %v err=%v", okSig, err)
}
pushes = f.sessions.records()
if len(pushes) != 1 || pushes[0].userID != f.callee.ID {
t.Fatalf("signaling pushes = %+v, want one to callee", pushes)
if len(pushes) != 1 || pushes[0].targetSession != phoneCalleeSession {
t.Fatalf("signaling pushes = %+v, want one to callee session", pushes)
}
sigUpdates := pushes[0].msg.(*tg.Updates)
sig, ok := sigUpdates.Updates[0].(*tg.UpdatePhoneCallSignalingData)

View file

@ -567,20 +567,11 @@ func (r *Router) pushSelfPhotoUpdateToCurrentSession(ctx context.Context, update
if !ok {
return
}
rawAuthKeyID, hasRawAuthKeyID := RawAuthKeyIDFrom(ctx)
rawAuthKeyID := rawAuthKeyIDForOrigin(ctx)
push := func() {
pushCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
if scoped, ok := r.scopedSessions(); ok {
if !hasRawAuthKeyID {
return
}
if err := scoped.PushToSessionForAuthKey(pushCtx, rawAuthKeyID, sessionID, proto.MessageFromServer, updates); err != nil {
r.log.Debug("push self photo update to current session", zap.Int64("session_id", sessionID), zap.Error(err))
}
return
}
if err := r.deps.Sessions.PushToSession(pushCtx, sessionID, proto.MessageFromServer, updates); err != nil {
if err := r.deps.Sessions.PushToSessionForAuthKey(pushCtx, rawAuthKeyID, sessionID, proto.MessageFromServer, updates); err != nil {
r.log.Debug("push self photo update to current session", zap.Int64("session_id", sessionID), zap.Error(err))
}
}

View file

@ -15,16 +15,8 @@ func (r *Router) pushUserMessage(ctx context.Context, userID int64, logMessage s
sessionID, _ := SessionIDFrom(ctx)
if timeout := r.cfg.OutboundPushTimeout; timeout > 0 {
authKeyID := rawAuthKeyIDForOrigin(ctx)
if scoped, ok := r.deps.Sessions.(ScopedBestEffortSessionBinder); ok {
if sent, err := scoped.PushToUserExceptAuthKeySessionBestEffort(ctx, userID, authKeyID, sessionID, proto.MessageFromServer, msg, timeout); err != nil {
r.log.Debug(logMessage, zap.Int64("user_id", userID), zap.Int("sent", sent), zap.Duration("timeout", timeout), zap.Error(err))
return sent
} else {
return sent
}
}
if bestEffort, ok := r.deps.Sessions.(BestEffortSessionBinder); ok {
if sent, err := bestEffort.PushToUserExceptSessionBestEffort(ctx, userID, sessionID, proto.MessageFromServer, msg, timeout); err != nil {
if sent, err := bestEffort.PushToUserExceptAuthKeySessionBestEffort(ctx, userID, authKeyID, sessionID, proto.MessageFromServer, msg, timeout); err != nil {
r.log.Debug(logMessage, zap.Int64("user_id", userID), zap.Int("sent", sent), zap.Duration("timeout", timeout), zap.Error(err))
return sent
} else {
@ -32,16 +24,8 @@ func (r *Router) pushUserMessage(ctx context.Context, userID int64, logMessage s
}
}
}
if scoped, ok := r.scopedSessions(); ok {
authKeyID := rawAuthKeyIDForOrigin(ctx)
if sent, err := scoped.PushToUserExceptAuthKeySession(ctx, userID, authKeyID, sessionID, proto.MessageFromServer, msg); err != nil {
r.log.Debug(logMessage, zap.Int64("user_id", userID), zap.Int("sent", sent), zap.Error(err))
return sent
} else {
return sent
}
}
if sent, err := r.deps.Sessions.PushToUserExceptSession(ctx, userID, sessionID, proto.MessageFromServer, msg); err != nil {
authKeyID := rawAuthKeyIDForOrigin(ctx)
if sent, err := r.deps.Sessions.PushToUserExceptAuthKeySession(ctx, userID, authKeyID, sessionID, proto.MessageFromServer, msg); err != nil {
r.log.Debug(logMessage, zap.Int64("user_id", userID), zap.Int("sent", sent), zap.Error(err))
return sent
} else {
@ -76,17 +60,7 @@ func (r *Router) pushCurrentSessionMessage(ctx context.Context, logMessage strin
if !ok {
return
}
if scoped, ok := r.scopedSessions(); ok {
rawAuthKeyID, ok := RawAuthKeyIDFrom(ctx)
if !ok {
return
}
if err := scoped.PushToSessionForAuthKey(ctx, rawAuthKeyID, sessionID, proto.MessageFromServer, msg); err != nil {
r.log.Debug(logMessage, zap.Int64("session_id", sessionID), zap.Error(err))
}
return
}
if err := r.deps.Sessions.PushToSession(ctx, sessionID, proto.MessageFromServer, msg); err != nil {
if err := r.deps.Sessions.PushToSessionForAuthKey(ctx, rawAuthKeyIDForOrigin(ctx), sessionID, proto.MessageFromServer, msg); err != nil {
r.log.Debug(logMessage, zap.Int64("session_id", sessionID), zap.Error(err))
}
}

View file

@ -308,12 +308,7 @@ func (r *Router) effectiveAuthKeyID(ctx context.Context, rawAuthKeyID [8]byte, s
hasCached bool
)
if r.deps.Sessions != nil {
if scoped, ok := r.deps.Sessions.(ScopedSessionBinder); ok {
if id, ok := scoped.AuthKeyIDForSession(rawAuthKeyID, sessionID); ok {
cached = id
hasCached = true
}
} else if id, ok := r.deps.Sessions.AuthKeyID(sessionID); ok {
if id, ok := r.deps.Sessions.AuthKeyIDForSession(rawAuthKeyID, sessionID); ok {
cached = id
hasCached = true
}
@ -381,39 +376,22 @@ func (r *Router) effectiveAuthKeyID(ctx context.Context, rawAuthKeyID [8]byte, s
func (r *Router) bindEffectiveAuthKey(rawAuthKeyID [8]byte, sessionID int64, effective [8]byte) {
if r.deps.Sessions != nil {
if scoped, ok := r.deps.Sessions.(ScopedSessionBinder); ok {
scoped.BindAuthKeyForSession(rawAuthKeyID, sessionID, effective)
} else {
r.deps.Sessions.BindAuthKey(sessionID, effective)
}
r.deps.Sessions.BindAuthKeyForSession(rawAuthKeyID, sessionID, effective)
}
}
func (r *Router) effectiveUserID(ctx context.Context, rawAuthKeyID, authKeyID [8]byte, sessionID int64) (int64, bool, error) {
if userID, ok := UserIDFrom(ctx); ok {
if scoped, ok := r.scopedSessions(); ok {
scoped.BindUserForAuthKey(rawAuthKeyID, sessionID, userID)
} else if r.deps.Sessions != nil {
r.deps.Sessions.BindUser(sessionID, userID)
if r.deps.Sessions != nil {
r.deps.Sessions.BindUserForAuthKey(rawAuthKeyID, sessionID, userID)
}
return userID, true, nil
}
if r.deps.Sessions != nil {
if scoped, ok := r.deps.Sessions.(ScopedSessionBinder); ok {
if userID, resolved := scoped.UserIDResolvedForAuthKey(rawAuthKeyID, sessionID); resolved {
if userID == 0 {
if cachedUserID, ok := r.positiveCachedAuthUser(authKeyID); ok {
scoped.BindUserForAuthKey(rawAuthKeyID, sessionID, cachedUserID)
r.announceSessionOnline(ctx, cachedUserID)
return cachedUserID, true, nil
}
}
return userID, userID != 0, nil
}
} else if userID, resolved := r.deps.Sessions.UserIDResolved(sessionID); resolved {
if userID, resolved := r.deps.Sessions.UserIDResolvedForAuthKey(rawAuthKeyID, sessionID); resolved {
if userID == 0 {
if cachedUserID, ok := r.positiveCachedAuthUser(authKeyID); ok {
r.deps.Sessions.BindUser(sessionID, cachedUserID)
r.deps.Sessions.BindUserForAuthKey(rawAuthKeyID, sessionID, cachedUserID)
r.announceSessionOnline(ctx, cachedUserID)
return cachedUserID, true, nil
}
@ -434,32 +412,16 @@ func (r *Router) effectiveUserID(ctx context.Context, rawAuthKeyID, authKeyID [8
return 0, false, err
}
if r.deps.Sessions != nil {
if scoped, ok := r.deps.Sessions.(ScopedSessionBinder); ok {
if cachedUserID, resolved := scoped.UserIDResolvedForAuthKey(rawAuthKeyID, sessionID); resolved {
if cachedUserID != 0 || !found {
return cachedUserID, cachedUserID != 0, nil
}
}
} else {
if cachedUserID, resolved := r.deps.Sessions.UserIDResolved(sessionID); resolved {
if cachedUserID != 0 || !found {
return cachedUserID, cachedUserID != 0, nil
}
if cachedUserID, resolved := r.deps.Sessions.UserIDResolvedForAuthKey(rawAuthKeyID, sessionID); resolved {
if cachedUserID != 0 || !found {
return cachedUserID, cachedUserID != 0, nil
}
}
if found {
if scoped, ok := r.deps.Sessions.(ScopedSessionBinder); ok {
scoped.BindUserForAuthKey(rawAuthKeyID, sessionID, userID)
} else {
r.deps.Sessions.BindUser(sessionID, userID)
}
r.deps.Sessions.BindUserForAuthKey(rawAuthKeyID, sessionID, userID)
r.announceSessionOnline(ctx, userID)
} else {
if scoped, ok := r.deps.Sessions.(ScopedSessionBinder); ok {
scoped.BindUserForAuthKey(rawAuthKeyID, sessionID, 0)
} else {
r.deps.Sessions.BindUser(sessionID, 0)
}
r.deps.Sessions.BindUserForAuthKey(rawAuthKeyID, sessionID, 0)
}
}
return userID, found, nil
@ -532,14 +494,6 @@ func (r *Router) invalidateAuthUserCache(authKeyID [8]byte) {
r.authUserSF.Forget(authKeyClientInfoSingleflightPrefix + key)
}
func (r *Router) scopedSessions() (ScopedSessionBinder, bool) {
if r.deps.Sessions == nil {
return nil, false
}
scoped, ok := r.deps.Sessions.(ScopedSessionBinder)
return scoped, ok
}
func (r *Router) dispatch(ctx context.Context, b *bin.Buffer, depth int) (bin.Encoder, error) {
if depth > maxWrapperDepth {
return nil, wrapperTooDeepErr()

View file

@ -26,7 +26,7 @@ func newAuthBindingCaptureSessions() *authBindingCaptureSessions {
return &authBindingCaptureSessions{captureSessions: &captureSessions{}}
}
func (s *authBindingCaptureSessions) PushToUserExceptSession(_ context.Context, userID, _ int64, t proto.MessageType, msg bin.Encoder) (int, error) {
func (s *authBindingCaptureSessions) PushToUserExceptAuthKeySession(_ context.Context, userID int64, _ [8]byte, _ int64, t proto.MessageType, msg bin.Encoder) (int, error) {
s.mu.Lock()
defer s.mu.Unlock()
s.messageType = t
@ -43,8 +43,8 @@ func TestDispatchPromotesNegativeSessionCacheFromPositiveAuthCache(t *testing.T)
userID = int64(1000000001)
)
sessions := newAuthBindingCaptureSessions()
sessions.BindAuthKey(sessionID, authKeyID)
sessions.BindUser(sessionID, 0)
sessions.BindAuthKeyForSession(authKeyID, sessionID, authKeyID)
sessions.BindUserForAuthKey(authKeyID, sessionID, 0)
auth := &captureAuthService{}
r := New(Config{}, Deps{
Auth: auth,

View file

@ -9,6 +9,7 @@ import (
type captureSessions struct {
mu sync.Mutex
rawAuthKeyID [8]byte
sessionID int64
userID int64
userResolved bool
@ -81,39 +82,35 @@ func (s *captureSessions) clearMessages() {
s.pushUserIDs = nil
}
func (s *captureSessions) BindAuthKey(sessionID int64, authKeyID [8]byte) {
func (s *captureSessions) BindAuthKeyForSession(rawAuthKeyID [8]byte, sessionID int64, authKeyID [8]byte) {
s.mu.Lock()
defer s.mu.Unlock()
if s.authKeyResolved && s.authKeyID != authKeyID {
s.userID = 0
s.userResolved = false
}
s.rawAuthKeyID = rawAuthKeyID
s.sessionID = sessionID
s.authKeyID = authKeyID
s.authKeyResolved = true
}
func (s *captureSessions) AuthKeyID(int64) ([8]byte, bool) {
func (s *captureSessions) AuthKeyIDForSession([8]byte, int64) ([8]byte, bool) {
s.mu.Lock()
defer s.mu.Unlock()
return s.authKeyID, s.authKeyResolved
}
func (s *captureSessions) BindUser(sessionID, userID int64) {
func (s *captureSessions) BindUserForAuthKey(rawAuthKeyID [8]byte, sessionID, userID int64) {
s.mu.Lock()
defer s.mu.Unlock()
s.rawAuthKeyID = rawAuthKeyID
s.sessionID = sessionID
s.userID = userID
s.userResolved = true
}
func (s *captureSessions) UserID(int64) (int64, bool) {
s.mu.Lock()
defer s.mu.Unlock()
return s.userID, s.userID != 0
}
func (s *captureSessions) UserIDResolved(int64) (int64, bool) {
func (s *captureSessions) UserIDResolvedForAuthKey([8]byte, int64) (int64, bool) {
s.mu.Lock()
defer s.mu.Unlock()
return s.userID, s.userResolved
@ -130,26 +127,29 @@ func (s *captureSessions) UnbindAuthKey(authKeyID [8]byte) int {
return 0
}
func (s *captureSessions) SetReceivesUpdates(sessionID int64, receives bool) {
func (s *captureSessions) SetReceivesUpdatesForAuthKey(rawAuthKeyID [8]byte, sessionID int64, receives bool) {
s.mu.Lock()
defer s.mu.Unlock()
s.rawAuthKeyID = rawAuthKeyID
s.sessionID = sessionID
s.receives = receives
}
func (s *captureSessions) PushToSession(_ context.Context, sessionID int64, t proto.MessageType, msg bin.Encoder) error {
func (s *captureSessions) PushToSessionForAuthKey(_ context.Context, rawAuthKeyID [8]byte, sessionID int64, t proto.MessageType, msg bin.Encoder) error {
s.mu.Lock()
defer s.mu.Unlock()
s.rawAuthKeyID = rawAuthKeyID
s.sessionID = sessionID
s.messageType = t
s.message = msg
return nil
}
func (s *captureSessions) PushToUserExceptSession(_ context.Context, userID, excludeSessionID int64, t proto.MessageType, msg bin.Encoder) (int, error) {
func (s *captureSessions) PushToUserExceptAuthKeySession(_ context.Context, userID int64, excludeAuthKeyID [8]byte, excludeSessionID int64, t proto.MessageType, msg bin.Encoder) (int, error) {
s.mu.Lock()
defer s.mu.Unlock()
s.userID = userID
s.rawAuthKeyID = excludeAuthKeyID
s.sessionID = excludeSessionID
s.messageType = t
s.message = msg

View file

@ -192,13 +192,7 @@ func (r *Router) markSessionReceivesUpdates(ctx context.Context, userID int64) {
if !ok {
return
}
if scoped, ok := r.scopedSessions(); ok {
if rawAuthKeyID, ok := RawAuthKeyIDFrom(ctx); ok {
scoped.SetReceivesUpdatesForAuthKey(rawAuthKeyID, sessionID, true)
return
}
}
r.deps.Sessions.SetReceivesUpdates(sessionID, true)
r.deps.Sessions.SetReceivesUpdatesForAuthKey(rawAuthKeyIDForOrigin(ctx), sessionID, true)
}
func ptr[T any](v T) *T { return &v }

View file

@ -1,13 +1,13 @@
// Package store 定义存储接口与协议层 DTO不含具体实现。
//
// 布局主包只放接口AuthKeyStore / SessionStore / UserStore / AuthorizationStore /
// 布局主包只放接口AuthKeyStore / UserStore / AuthorizationStore /
// CodeStore / UpdateStateStore / UpdateEventStore 等)与协议 DTO三种后端实现各自独立成对称子包
// - store/memory —— 内存实现,测试替身与本地兜底
// - store/postgres —— PostgreSQLpgx + sqlc 生成查询 + golang-migrate 迁移)
// - store/redisstore —— Redisgo-redis
//
// 类型边界:接口签名分两类——
// - 协议产物用 store 自有 DTOAuthKeyData、SessionData、PhoneCode不依赖 tg.*,也非业务实体);
// - 协议产物用 store 自有 DTOAuthKeyData、PhoneCode不依赖 tg.*,也非业务实体);
// - 业务实体直接用 domainUserStore / AuthorizationStore / MessageStore / UpdateEventStore
// 收发 domain.User / domain.Authorization / domain.Message / domain.UpdateEvent。
package store

View file

@ -73,38 +73,6 @@ func (s *AuthKeyStore) Delete(_ context.Context, id [8]byte) error {
return nil
}
// SessionStore 是 store.SessionStore 的内存实现。
type SessionStore struct {
mu sync.RWMutex
sessions map[int64]store.SessionData
}
// NewSessionStore 创建内存 SessionStore。
func NewSessionStore() *SessionStore {
return &SessionStore{sessions: make(map[int64]store.SessionData)}
}
func (s *SessionStore) Save(_ context.Context, d store.SessionData) error {
s.mu.Lock()
s.sessions[d.ID] = d
s.mu.Unlock()
return nil
}
func (s *SessionStore) Get(_ context.Context, id int64) (store.SessionData, bool, error) {
s.mu.RLock()
d, ok := s.sessions[id]
s.mu.RUnlock()
return d, ok, nil
}
func (s *SessionStore) Delete(_ context.Context, id int64) error {
s.mu.Lock()
delete(s.sessions, id)
s.mu.Unlock()
return nil
}
// TempAuthKeyBindingStore 是 store.TempAuthKeyBindingStore 的内存实现。
type TempAuthKeyBindingStore struct {
mu sync.RWMutex

View file

@ -169,7 +169,7 @@ func TestDispatchOutboxLifecycleKeepsDurableEvents(t *testing.T) {
Peer: domain.Peer{Type: domain.PeerTypeUser, ID: owner.ID + int64(pts)},
Bool: pts%2 == 0,
}
if _, err := events.AppendAllocatedWithDispatch(ctx, owner.ID, event, [8]byte{}, sessionID); err != nil {
if _, err := events.AppendAllocatedWithDispatch(ctx, owner.ID, event, [8]byte{1}, sessionID); err != nil {
t.Fatalf("AppendAllocatedWithDispatch pts=%d: %v", pts, err)
}
if _, err := tx.Exec(ctx, `

View file

@ -2,6 +2,7 @@ package postgres
import (
"context"
"errors"
"fmt"
"time"
@ -18,6 +19,21 @@ const (
maxDispatchPoisonCleanupBatch = 1000
)
var errInvalidDispatchOutboxExclusionPair = errors.New("dispatch outbox exclusion requires both raw auth key and session id")
// enqueueDispatch is the only production write boundary for dispatch_outbox.
// A zero pair means no originating session is excluded; a non-zero pair identifies
// one exact physical raw-auth/session tuple. A half pair is never meaningful because
// session IDs are not globally unique and must fail the surrounding transaction.
func enqueueDispatch(ctx context.Context, q *sqlcgen.Queries, arg sqlcgen.EnqueueDispatchParams) error {
hasAuthKey := arg.ExcludeAuthKeyID != 0
hasSession := arg.ExcludeSessionID != 0
if hasAuthKey != hasSession {
return errInvalidDispatchOutboxExclusionPair
}
return q.EnqueueDispatch(ctx, arg)
}
// DispatchOutboxStore 用 PostgreSQL 实现 transactional outbox。
type DispatchOutboxStore struct {
q *sqlcgen.Queries

View file

@ -0,0 +1,82 @@
package postgres
import (
"context"
"errors"
"testing"
"github.com/jackc/pgx/v5/pgconn"
"telesrv/internal/domain"
)
func TestDispatchOutboxExclusionPairInvariantPostgres(t *testing.T) {
pool := testPool(t)
ctx := context.Background()
suffix := randomSuffix(t)
owner := createTestUser(t, ctx, NewUserStore(pool), "+1887"+suffix+"01", "OutboxPair", "")
t.Cleanup(func() { _, _ = pool.Exec(ctx, "DELETE FROM users WHERE id = $1", owner.ID) })
event := domain.UpdateEvent{
Type: domain.UpdateEventDialogPinned,
PtsCount: 1,
Date: 1700002300,
Peer: domain.Peer{Type: domain.PeerTypeUser, ID: owner.ID},
Bool: true,
}
for _, test := range []struct {
name string
authKeyID [8]byte
sessionID int64
}{
{name: "auth key only", authKeyID: [8]byte{1}},
{name: "session only", sessionID: 77},
} {
t.Run("write boundary "+test.name, func(t *testing.T) {
_, err := NewUpdateEventStore(pool).AppendAllocatedWithDispatch(ctx, owner.ID, event, test.authKeyID, test.sessionID)
if !errors.Is(err, errInvalidDispatchOutboxExclusionPair) {
t.Fatalf("AppendAllocatedWithDispatch error = %v, want %v", err, errInvalidDispatchOutboxExclusionPair)
}
})
}
var eventCount int
if err := pool.QueryRow(ctx, "SELECT count(*)::int FROM user_update_events WHERE user_id = $1", owner.ID).Scan(&eventCount); err != nil {
t.Fatalf("count events after rejected writes: %v", err)
}
if eventCount != 0 {
t.Fatalf("events after rejected writes = %d, want 0 (transaction rollback)", eventCount)
}
stored, err := NewUpdateEventStore(pool).AppendAllocated(ctx, owner.ID, event)
if err != nil {
t.Fatalf("append durable event for constraint test: %v", err)
}
for _, test := range []struct {
name string
authKeyID int64
sessionID int64
}{
{name: "auth key only", authKeyID: 1},
{name: "session only", sessionID: 77},
} {
t.Run("database constraint "+test.name, func(t *testing.T) {
_, err := pool.Exec(ctx, `
INSERT INTO dispatch_outbox (
target_user_id, pts, event_type, exclude_auth_key_id, exclude_session_id
) VALUES ($1, $2, $3, $4, $5)`, owner.ID, stored.Pts, string(stored.Type), test.authKeyID, test.sessionID)
var pgErr *pgconn.PgError
if !errors.As(err, &pgErr) || pgErr.Code != "23514" || pgErr.ConstraintName != "dispatch_outbox_exclusion_pair_check" {
t.Fatalf("direct insert error = %v, want check violation from dispatch_outbox_exclusion_pair_check", err)
}
})
}
var outboxCount int
if err := pool.QueryRow(ctx, "SELECT count(*)::int FROM dispatch_outbox WHERE target_user_id = $1", owner.ID).Scan(&outboxCount); err != nil {
t.Fatalf("count outbox after rejected inserts: %v", err)
}
if outboxCount != 0 {
t.Fatalf("outbox rows after rejected inserts = %d, want 0", outboxCount)
}
}

View file

@ -0,0 +1,31 @@
package postgres
import (
"context"
"errors"
"testing"
"telesrv/internal/store/postgres/sqlcgen"
)
func TestEnqueueDispatchRejectsHalfExclusionPair(t *testing.T) {
tests := []struct {
name string
authKeyID int64
sessionID int64
}{
{name: "auth key only", authKeyID: 1},
{name: "session only", sessionID: 1},
}
for _, test := range tests {
t.Run(test.name, func(t *testing.T) {
err := enqueueDispatch(context.Background(), nil, sqlcgen.EnqueueDispatchParams{
ExcludeAuthKeyID: test.authKeyID,
ExcludeSessionID: test.sessionID,
})
if !errors.Is(err, errInvalidDispatchOutboxExclusionPair) {
t.Fatalf("enqueueDispatch error = %v, want %v", err, errInvalidDispatchOutboxExclusionPair)
}
})
}
}

View file

@ -186,7 +186,7 @@ func (s *MessageStore) DeliverLoginCodeMessage(ctx context.Context, req domain.L
if err := appendNewMessageEvent(ctx, qtx, msg); err != nil {
return domain.LoginCodeDeliveryResult{}, err
}
if err := qtx.EnqueueDispatch(ctx, sqlcgen.EnqueueDispatchParams{
if err := enqueueDispatch(ctx, qtx, sqlcgen.EnqueueDispatchParams{
TargetUserID: req.UserID,
Pts: int32(msg.Pts),
EventType: string(domain.UpdateEventNewMessage),

View file

@ -227,7 +227,7 @@ WHERE sender_user_id = $1
dispatchAuthKeyID = excludeAuthKeyID
dispatchSessionID = excludeSessionID
}
if err := q.EnqueueDispatch(ctx, sqlcgen.EnqueueDispatchParams{
if err := enqueueDispatch(ctx, q, sqlcgen.EnqueueDispatchParams{
TargetUserID: userID,
Pts: int32(deletePts),
EventType: string(domain.UpdateEventDeleteMessages),
@ -256,7 +256,7 @@ WHERE sender_user_id = $1
}); err != nil {
return res, fmt.Errorf("advance dialog read inbox after delete correction: %w", err)
}
if err := q.EnqueueDispatch(ctx, sqlcgen.EnqueueDispatchParams{
if err := enqueueDispatch(ctx, q, sqlcgen.EnqueueDispatchParams{
TargetUserID: userID,
Pts: int32(correction.Pts),
EventType: string(domain.UpdateEventReadHistoryInbox),

View file

@ -161,7 +161,7 @@ WHERE owner_user_id = $1 AND box_id = $2`, box.OwnerUserID, box.BoxID, int32(pts
if err := appendUserUpdateEvent(ctx, tx, qtx, msg.OwnerUserID, event); err != nil {
return res, fmt.Errorf("append web page event: %w", err)
}
if err := qtx.EnqueueDispatch(ctx, sqlcgen.EnqueueDispatchParams{
if err := enqueueDispatch(ctx, qtx, sqlcgen.EnqueueDispatchParams{
TargetUserID: msg.OwnerUserID,
Pts: int32(pts),
EventType: string(domain.UpdateEventWebPage),
@ -257,7 +257,7 @@ WHERE message_sender_id = $1 AND private_message_id = $2`, messageSenderID, targ
dispatchAuthKeyID = req.OriginAuthKeyID
dispatchSessionID = req.OriginSessionID
}
if err := qtx.EnqueueDispatch(ctx, sqlcgen.EnqueueDispatchParams{
if err := enqueueDispatch(ctx, qtx, sqlcgen.EnqueueDispatchParams{
TargetUserID: msg.OwnerUserID,
Pts: int32(pts),
EventType: string(domain.UpdateEventEditMessage),

View file

@ -379,7 +379,7 @@ func (s *MessageStore) ReadHistory(ctx context.Context, req domain.ReadHistoryRe
if err := appendUserUpdateEvent(ctx, tx, qtx, req.OwnerUserID, res.InboxEvent); err != nil {
return res, fmt.Errorf("append read inbox event: %w", err)
}
if err := qtx.EnqueueDispatch(ctx, sqlcgen.EnqueueDispatchParams{
if err := enqueueDispatch(ctx, qtx, sqlcgen.EnqueueDispatchParams{
TargetUserID: req.OwnerUserID,
Pts: int32(readerPts),
EventType: string(domain.UpdateEventReadHistoryInbox),
@ -414,7 +414,7 @@ func (s *MessageStore) ReadHistory(ctx context.Context, req domain.ReadHistoryRe
if err := appendUserUpdateEvent(ctx, tx, qtx, candidate.SenderOwnerUserID, res.OutboxEvent); err != nil {
return res, fmt.Errorf("append read outbox event: %w", err)
}
if err := qtx.EnqueueDispatch(ctx, sqlcgen.EnqueueDispatchParams{
if err := enqueueDispatch(ctx, qtx, sqlcgen.EnqueueDispatchParams{
TargetUserID: candidate.SenderOwnerUserID,
Pts: int32(senderPts),
EventType: string(domain.UpdateEventReadHistoryOutbox),

View file

@ -141,7 +141,7 @@ func (s *MessageStore) PinPrivateMessage(ctx context.Context, req domain.PinPriv
dispatchAuthKeyID = req.OriginAuthKeyID
dispatchSessionID = req.OriginSessionID
}
if err := qtx.EnqueueDispatch(ctx, sqlcgen.EnqueueDispatchParams{
if err := enqueueDispatch(ctx, qtx, sqlcgen.EnqueueDispatchParams{
TargetUserID: side.userID,
Pts: int32(pts),
EventType: string(domain.UpdateEventPinnedMessages),
@ -286,7 +286,7 @@ func (s *MessageStore) UnpinAllPrivateMessages(ctx context.Context, req domain.U
dispatchAuthKeyID = req.OriginAuthKeyID
dispatchSessionID = req.OriginSessionID
}
if err := qtx.EnqueueDispatch(ctx, sqlcgen.EnqueueDispatchParams{
if err := enqueueDispatch(ctx, qtx, sqlcgen.EnqueueDispatchParams{
TargetUserID: side.userID,
Pts: int32(pts),
EventType: string(domain.UpdateEventPinnedMessages),

View file

@ -182,7 +182,7 @@ WHERE d.user_id = $1
if err := appendUserUpdateEvent(ctx, tx, qtx, req.OwnerUserID, res.Event); err != nil {
return res, fmt.Errorf("append read message contents event: %w", err)
}
if err := qtx.EnqueueDispatch(ctx, sqlcgen.EnqueueDispatchParams{
if err := enqueueDispatch(ctx, qtx, sqlcgen.EnqueueDispatchParams{
TargetUserID: req.OwnerUserID,
Pts: int32(pts),
EventType: string(domain.UpdateEventReadMessageContents),
@ -242,7 +242,7 @@ RETURNING box_id`, senderID, senderPrivateMessageIDs[senderID])
if err := appendUserUpdateEvent(ctx, tx, qtx, senderID, event); err != nil {
return res, fmt.Errorf("append sender content read event: %w", err)
}
if err := qtx.EnqueueDispatch(ctx, sqlcgen.EnqueueDispatchParams{
if err := enqueueDispatch(ctx, qtx, sqlcgen.EnqueueDispatchParams{
TargetUserID: senderID,
Pts: int32(senderPts),
EventType: string(domain.UpdateEventReadMessageContents),

View file

@ -298,7 +298,7 @@ func (s *MessageStore) sendPrivateTextOnce(ctx context.Context, req domain.SendP
if err := appendNewMessageEvent(ctx, qtx, sender); err != nil {
return domain.SendPrivateTextResult{}, err
}
if err := qtx.EnqueueDispatch(ctx, sqlcgen.EnqueueDispatchParams{
if err := enqueueDispatch(ctx, qtx, sqlcgen.EnqueueDispatchParams{
TargetUserID: req.SenderUserID,
Pts: int32(senderPts),
EventType: string(domain.UpdateEventNewMessage),
@ -360,7 +360,7 @@ func (s *MessageStore) sendPrivateTextOnce(ctx context.Context, req domain.SendP
if err := appendNewMessageEvent(ctx, qtx, recipient); err != nil {
return domain.SendPrivateTextResult{}, err
}
if err := qtx.EnqueueDispatch(ctx, sqlcgen.EnqueueDispatchParams{
if err := enqueueDispatch(ctx, qtx, sqlcgen.EnqueueDispatchParams{
TargetUserID: req.RecipientUserID,
Pts: int32(recipientPts),
EventType: string(domain.UpdateEventNewMessage),

View file

@ -92,7 +92,7 @@ func (s *PhoneChangeStore) ChangePhone(ctx context.Context, req domain.PhoneChan
if err := appendUserUpdateEvent(ctx, tx, qtx, req.UserID, event); err != nil {
return domain.PhoneChangeResult{}, fmt.Errorf("append phone change event: %w", err)
}
if err := qtx.EnqueueDispatch(ctx, sqlcgen.EnqueueDispatchParams{
if err := enqueueDispatch(ctx, qtx, sqlcgen.EnqueueDispatchParams{
TargetUserID: req.UserID,
Pts: int32(event.Pts),
EventType: string(event.Type),

View file

@ -122,7 +122,7 @@ func (s *UpdateEventStore) appendInTx(ctx context.Context, db sqlcgen.DBTX, q *s
return domain.UpdateEvent{}, fmt.Errorf("append update event: %w", err)
}
if dispatch {
if err := q.EnqueueDispatch(ctx, sqlcgen.EnqueueDispatchParams{
if err := enqueueDispatch(ctx, q, sqlcgen.EnqueueDispatchParams{
TargetUserID: userID,
Pts: int32(event.Pts),
EventType: string(event.Type),

View file

@ -1,4 +1,4 @@
// Package redisstore 用 Redis 实现高频易失态的存储接口第一阶段SessionStore
// Package redisstore 用 Redis 实现高频、易失且可重建的短状态、缓存、计数器与限流
//
// 职责边界见 docs/persistence-layer.md §1Redis 存「态与计数」,丢失可由 PG/协议恢复。
package redisstore

View file

@ -1,76 +0,0 @@
package redisstore
import (
"context"
"encoding/json"
"errors"
"fmt"
"time"
"github.com/redis/go-redis/v9"
"telesrv/internal/store"
)
// DefaultSessionTTL 是 session 记录的默认过期时间。
// session 是连接态:过期或丢失后,客户端重连会触发 new_session_created / bad_server_salt 重建,
// 因此 TTL 不必很长。每个随机 session_id 都落一条记录且断连不删,过长的 TTL
// 只会堆积死 session移动端每次重连一条。7 天足够覆盖常规离线窗口。
const DefaultSessionTTL = 7 * 24 * time.Hour
// SessionStore 用 Redis 实现 store.SessionStore。
type SessionStore struct {
c *redis.Client
ttl time.Duration
}
// NewSessionStore 创建 Redis SessionStore。ttl<=0 表示永不过期。
func NewSessionStore(c *redis.Client, ttl time.Duration) *SessionStore {
return &SessionStore{c: c, ttl: ttl}
}
func sessionKey(id int64) string {
return fmt.Sprintf("session:%d", id)
}
// sessionValue 是 SessionData 在 Redis 中的序列化形态(不含 IDID 即 key
type sessionValue struct {
AuthKeyID [8]byte `json:"auth_key_id"`
Salt int64 `json:"salt"`
LastSeen int64 `json:"last_seen"`
}
// Save 实现 store.SessionStore。
func (s *SessionStore) Save(ctx context.Context, d store.SessionData) error {
v, err := json.Marshal(sessionValue{AuthKeyID: d.AuthKeyID, Salt: d.Salt, LastSeen: d.LastSeen})
if err != nil {
return fmt.Errorf("marshal session: %w", err)
}
if err := s.c.Set(ctx, sessionKey(d.ID), v, s.ttl).Err(); err != nil {
return fmt.Errorf("redis set session: %w", err)
}
return nil
}
// Get 实现 store.SessionStore。不存在时 found=false。
func (s *SessionStore) Get(ctx context.Context, id int64) (store.SessionData, bool, error) {
raw, err := s.c.Get(ctx, sessionKey(id)).Bytes()
if err != nil {
if errors.Is(err, redis.Nil) {
return store.SessionData{}, false, nil
}
return store.SessionData{}, false, fmt.Errorf("redis get session: %w", err)
}
var v sessionValue
if err := json.Unmarshal(raw, &v); err != nil {
return store.SessionData{}, false, fmt.Errorf("unmarshal session: %w", err)
}
return store.SessionData{ID: id, AuthKeyID: v.AuthKeyID, Salt: v.Salt, LastSeen: v.LastSeen}, true, nil
}
func (s *SessionStore) Delete(ctx context.Context, id int64) error {
if err := s.c.Del(ctx, sessionKey(id)).Err(); err != nil {
return fmt.Errorf("redis delete session: %w", err)
}
return nil
}

View file

@ -1,52 +0,0 @@
package redisstore
import (
"context"
"os"
"testing"
"time"
"telesrv/internal/store"
)
// TestSessionStoreRoundTrip 验证 session 落 Redis 后能用全新 store 实例原样读回。
// 未设 TELESRV_TEST_REDIS_ADDR 则跳过。
func TestSessionStoreRoundTrip(t *testing.T) {
addr := os.Getenv("TELESRV_TEST_REDIS_ADDR")
if addr == "" {
t.Skip("set TELESRV_TEST_REDIS_ADDR to run redis integration test")
}
ctx := context.Background()
c, err := Open(ctx, addr, "", 0)
if err != nil {
t.Fatalf("open: %v", err)
}
t.Cleanup(func() { _ = c.Close() })
want := store.SessionData{
ID: 0x1234beef,
AuthKeyID: [8]byte{1, 2, 3, 4, 5, 6, 7, 8},
Salt: 42,
LastSeen: 1000,
}
t.Cleanup(func() { _ = c.Del(ctx, sessionKey(want.ID)).Err() })
if err := NewSessionStore(c, time.Minute).Save(ctx, want); err != nil {
t.Fatalf("save: %v", err)
}
got, found, err := NewSessionStore(c, time.Minute).Get(ctx, want.ID)
if err != nil {
t.Fatalf("get: %v", err)
}
if !found {
t.Fatal("session not found after save")
}
if got != want {
t.Fatalf("round trip mismatch: got %+v want %+v", got, want)
}
if _, found, _ := NewSessionStore(c, time.Minute).Get(ctx, 999999); found {
t.Fatal("unexpected found for missing session")
}
}

View file

@ -1,23 +0,0 @@
package store
import "context"
// SessionData 是一条 MTProto session 记录client 生成的 session_id
//
// 后续里程碑会扩展 device / layer 等字段。
type SessionData struct {
ID int64 // session_id客户端生成
AuthKeyID [8]byte // 绑定的 auth key
Salt int64 // 当前 server salt
LastSeen int64 // unix 秒
}
// SessionStore 记录在线 MTProto session。实现见 store/memory测试替身、store/redisstore。
type SessionStore interface {
// Save 保存或更新一条 session 记录。
Save(ctx context.Context, s SessionData) error
// Get 按 session_id 查询;不存在时 found=false。
Get(ctx context.Context, id int64) (data SessionData, found bool, err error)
// Delete 删除一条 session 记录;不存在时不报错。
Delete(ctx context.Context, id int64) error
}