diff --git a/cmd/telesrv/main.go b/cmd/telesrv/main.go index e2d9e627..e7d801fb 100644 --- a/cmd/telesrv/main.go +++ b/cmd/telesrv/main.go @@ -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, diff --git a/deploy/migrations/0001_init.up.sql b/deploy/migrations/0001_init.up.sql index da9643a3..7fedd390 100644 --- a/deploy/migrations/0001_init.up.sql +++ b/deploy/migrations/0001_init.up.sql @@ -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)))) ); diff --git a/deploy/migrations/0082_dispatch_outbox_exclusion_pair.down.sql b/deploy/migrations/0082_dispatch_outbox_exclusion_pair.down.sql new file mode 100644 index 00000000..f78dc9c5 --- /dev/null +++ b/deploy/migrations/0082_dispatch_outbox_exclusion_pair.down.sql @@ -0,0 +1,2 @@ +ALTER TABLE dispatch_outbox + DROP CONSTRAINT IF EXISTS dispatch_outbox_exclusion_pair_check; diff --git a/deploy/migrations/0082_dispatch_outbox_exclusion_pair.up.sql b/deploy/migrations/0082_dispatch_outbox_exclusion_pair.up.sql new file mode 100644 index 00000000..179cc234 --- /dev/null +++ b/deploy/migrations/0082_dispatch_outbox_exclusion_pair.up.sql @@ -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; diff --git a/internal/mtprotoedge/conn.go b/internal/mtprotoedge/conn.go index 575d9afb..d1d65979 100644 --- a/internal/mtprotoedge/conn.go +++ b/internal/mtprotoedge/conn.go @@ -82,12 +82,9 @@ type Conn struct { outboundControlBudgetOnce sync.Once outboundScratchPool *outboundScratchPool outboundScratchOnce sync.Once - // terminal 表示该 logical Conn 已停止接受新的出站操作。写失败时由 - // outbound actor 置位并只发停止信号，不能在 actor 内等待自身退出。 - terminal atomic.Bool - // lifecycle is a monotonic activation state machine. In particular, retired - // never transitions back to claiming/active; this closes the stale-read-loop - // ABA where an evicted Conn observed "not active" and registered itself again. + // lifecycle is the sole monotonic activation/retirement state machine. + // retired never transitions back to claiming/active; one atomic state avoids + // contradictory activation and shutdown observations. lifecycle atomic.Uint32 transportClose sync.Once @@ -138,14 +135,6 @@ type Conn struct { membershipGen atomic.Int64 // createdAt 是连接建立时刻，供同 auth_key session 数触顶时驱逐真正最旧的连接。 createdAt time.Time - // keyDestroyed 标记本连接的 auth_key 已被 destroy_auth_key 删除。serveConn 对已建立 - // 连接复用缓存密钥跳过每帧 AuthKeyStore 回查；置位后强制回落到 Get→AuthKeyNotFound， - // 维持「destroy_auth_key 发起连接下一帧自然失效」契约。只由 destroy_auth_key 处理器置位。 - keyDestroyed atomic.Bool - // lastSessionSaveUnix 是上次把本连接 session 持久化到 SessionStore 的 unix 秒，用于把 - // 每帧 Save 去抖到固定间隔——session 持久化是软状态（生产无热读路径，仅观测/未来用）。 - // 只由单连接的读循环 goroutine 访问。 - lastSessionSaveUnix atomic.Int64 // clientLayer 是本连接协商的 TL layer（invokeWithLayer/initConnection），由 handleRPC // 在每次 Dispatch 后从 RPC 注册表刷新。出站(rpc_result/push)按此把 227 对象降级给老客户端； // 0 表示尚未协商，按 canonical(227) 处理=不降级。 @@ -159,8 +148,29 @@ func (c *Conn) lifecycleState() connLifecycle { return connLifecycle(c.lifecycle.Load()) } +func (c *Conn) isRetired() bool { + return c == nil || c.lifecycleState() == connLifecycleRetired +} + +// retire irreversibly fences the logical connection. The caller that wins the +// transition may additionally own one-shot physical cleanup. +func (c *Conn) retire() bool { + if c == nil { + return false + } + for { + state := c.lifecycle.Load() + if connLifecycle(state) == connLifecycleRetired { + return false + } + if c.lifecycle.CompareAndSwap(state, uint32(connLifecycleRetired)) { + return true + } + } +} + func (c *Conn) beginActivationClaim() bool { - if c == nil || c.terminal.Load() || !c.isPhysicalTransportCurrentOpen() { + if c == nil || !c.isPhysicalTransportCurrentOpen() { return false } if !c.lifecycle.CompareAndSwap(uint32(connLifecycleProvisional), uint32(connLifecycleClaiming)) { @@ -169,31 +179,31 @@ func (c *Conn) beginActivationClaim() bool { // Physical close can win after the pre-check but before the lifecycle CAS. // Do not let a doomed claimant enter SessionManager and retire a healthy old // owner for the same logical session. - if c.terminal.Load() || !c.isPhysicalTransportCurrentOpen() { - c.lifecycle.Store(uint32(connLifecycleRetired)) + if c.lifecycleState() != connLifecycleClaiming || !c.isPhysicalTransportCurrentOpen() { + c.retire() return false } return true } func (c *Conn) publishActivation() bool { - if c == nil || c.terminal.Load() || !c.isPhysicalTransportCurrentOpen() { + if c == nil || !c.isPhysicalTransportCurrentOpen() { return false } if !c.lifecycle.CompareAndSwap(uint32(connLifecycleClaiming), uint32(connLifecycleActive)) { return false } // A concurrent transport failure can retire the Conn between the first - // terminal check and the CAS. Never let that intermediate active value escape. - if c.terminal.Load() || !c.isPhysicalTransportCurrentOpen() { - c.lifecycle.Store(uint32(connLifecycleRetired)) + // physical check and the CAS. Never let that intermediate active value escape. + if c.lifecycleState() != connLifecycleActive || !c.isPhysicalTransportCurrentOpen() { + c.retire() return false } return true } func (c *Conn) isActive() bool { - return c != nil && !c.terminal.Load() && c.lifecycleState() == connLifecycleActive + return c != nil && c.lifecycleState() == connLifecycleActive && c.isPhysicalTransportCurrentOpen() } // transferTransportOwnership hands this Conn's physical socket to the next diff --git a/internal/mtprotoedge/encrypted.go b/internal/mtprotoedge/encrypted.go index 76bbcaba..a31c6589 100644 --- a/internal/mtprotoedge/encrypted.go +++ b/internal/mtprotoedge/encrypted.go @@ -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, diff --git a/internal/mtprotoedge/encrypted_test.go b/internal/mtprotoedge/encrypted_test.go index 9c27c2d9..ab48a05d 100644 --- a/internal/mtprotoedge/encrypted_test.go +++ b/internal/mtprotoedge/encrypted_test.go @@ -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， diff --git a/internal/mtprotoedge/inbound_preflight.go b/internal/mtprotoedge/inbound_preflight.go index c6aeccf5..d38995ec 100644 --- a/internal/mtprotoedge/inbound_preflight.go +++ b/internal/mtprotoedge/inbound_preflight.go @@ -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, diff --git a/internal/mtprotoedge/inbound_rpc.go b/internal/mtprotoedge/inbound_rpc.go index 2d33f8ea..d7cb4272 100644 --- a/internal/mtprotoedge/inbound_rpc.go +++ b/internal/mtprotoedge/inbound_rpc.go @@ -39,14 +39,7 @@ type inboundRPC struct { ticket *inboundRPCTicket } -const ( - inboundRPCTicketQueued int32 = iota - inboundRPCTicketRunning - inboundRPCTicketDone -) - type inboundRPCTicket struct { - state atomic.Int32 onTimeout func() } @@ -84,19 +77,6 @@ type inboundRPCGlobalReservation struct { released atomic.Bool } -// inboundRPCReservation 同时持有全局和单连接的“Copy 前”预算。commit/abort 只能成功一次； -// 无论 Copy 后连接关闭、入队成功还是调用方提前返回，预算都有唯一归还路径。 -type inboundRPCReservation struct { - conn *Conn - global *inboundRPCGlobalReservation - ctx context.Context - method string - size int - enqueuedAt time.Time - deadline time.Time - once sync.Once -} - // inboundRPCSpec 是 container preflight 与 RPC scheduler 之间的有界 admission 描述。 // method 仅用于 metrics，size 是在 Copy 之前必须预留的 request body 字节数。 type inboundRPCSpec struct { @@ -202,31 +182,6 @@ func (s *inboundRPCScheduler) stop(timeout time.Duration) { } } -func (s *inboundRPCScheduler) reserveGlobal(size int) (*inboundRPCGlobalReservation, string, error) { - if size < 0 { - size = 0 - } - size64 := int64(size) - s.budgetMu.Lock() - defer s.budgetMu.Unlock() - - select { - case <-s.stopCh: - return nil, "scheduler_closed", ErrConnClosed - default: - } - if s.tasks >= s.maxTasks { - return nil, "global_task_budget", ErrInboundRPCQueueFull - } - // 用减法比较避免 s.bytes+size64 溢出。 - if size64 > s.maxBytes-s.bytes { - return nil, "global_byte_budget", ErrInboundRPCQueueFull - } - s.tasks++ - s.bytes += size64 - return &inboundRPCGlobalReservation{scheduler: s, size: size64}, "", nil -} - // reserveGlobalBatch 在一次 budgetMu 临界区内检查并预留整批条数/字节。 // 返回的每个 reservation 仍由对应 task 单独归还，避免一个慢 RPC 持有 // 整个 container 已完成任务的预算。 @@ -312,12 +267,6 @@ func releaseInboundRPCGlobalBatch(reservations []*inboundRPCGlobalReservation) { scheduler.budgetMu.Unlock() } -func (s *inboundRPCScheduler) budgetSnapshot() (tasks int, bytes int64) { - s.budgetMu.Lock() - defer s.budgetMu.Unlock() - return s.tasks, s.bytes -} - func (s *inboundRPCScheduler) schedule(c *Conn) { if s == nil || c == nil { return @@ -454,86 +403,6 @@ func (c *Conn) startInboundRPCScheduler(scheduler *inboundRPCScheduler, maxInfli // rpcQueue 保持 nil；首个成功 commit 才由 append 分配，静默连接零队列内存。 } -// reserveInboundRPC 必须在 request body Copy 前调用。它先拿进程级条数/字节预算， -// 再预占单连接队列槽和字节预算；commit 或 abort 负责唯一释放。 -func (c *Conn) reserveInboundRPC(ctx context.Context, method string, size int) (*inboundRPCReservation, error) { - if ctx == nil { - ctx = context.Background() - } - select { - case <-ctx.Done(): - c.metrics.InboundRPCDropped(method, "context_done") - return nil, ctx.Err() - default: - } - if c.terminal.Load() { - c.metrics.InboundRPCDropped(method, "scheduler_closed") - return nil, ErrConnClosed - } - if c.rpcScheduler == nil { - c.metrics.InboundRPCDropped(method, "scheduler_closed") - return nil, ErrConnClosed - } - global, reason, err := c.rpcScheduler.reserveGlobal(size) - if err != nil { - c.metrics.InboundRPCDropped(method, reason) - return nil, err - } - - now := time.Now() - deadline := time.Time{} - if c.rpcTimeout > 0 { - deadline = now.Add(c.rpcTimeout) - } - if ctxDeadline, ok := ctx.Deadline(); ok && (deadline.IsZero() || ctxDeadline.Before(deadline)) { - deadline = ctxDeadline - } - if size < 0 { - size = 0 - } - - c.rpcMu.Lock() - if err := ctx.Err(); err != nil { - c.rpcMu.Unlock() - global.release() - c.metrics.InboundRPCDropped(method, "context_done") - return nil, err - } - if c.rpcClosed || c.terminal.Load() { - c.rpcMu.Unlock() - global.release() - c.metrics.InboundRPCDropped(method, "scheduler_closed") - return nil, ErrConnClosed - } - if c.rpcReserved+len(c.rpcQueue) >= c.rpcQueueSize { - c.rpcMu.Unlock() - global.release() - c.metrics.InboundRPCDropped(method, "queue_full") - return nil, ErrInboundRPCQueueFull - } - if int64(size) > maxInflightRPCBytes-c.inflightRPCBytes.Load() { - c.rpcMu.Unlock() - global.release() - c.metrics.InboundRPCDropped(method, "byte_budget") - return nil, ErrInboundRPCQueueFull - } - c.rpcReserved++ - c.inflightRPCBytes.Add(int64(size)) - // Add 与 close 的 Wait 由 rpcMu 排序：close 置 rpcClosed 后不会再发生 Add。 - c.rpcReservationWG.Add(1) - c.rpcMu.Unlock() - - return &inboundRPCReservation{ - conn: c, - global: global, - ctx: ctx, - method: method, - size: size, - enqueuedAt: now, - deadline: deadline, - }, nil -} - // reserveInboundRPCBatch 必须在 container 内任何 request body Copy 前调用。 // 全局预算只锁一次，单连接预算也只锁一次；任一限制不满足时 // 整批失败，不会留下部分 task/字节 reservation。 @@ -547,7 +416,7 @@ func (c *Conn) reserveInboundRPCBatch(ctx context.Context, specs []inboundRPCSpe return nil, ctx.Err() default: } - if c.terminal.Load() { + if c.isRetired() { c.dropInboundRPCSpecs(specs, "scheduler_closed") return nil, ErrConnClosed } @@ -589,7 +458,7 @@ func (c *Conn) reserveInboundRPCBatch(ctx context.Context, specs []inboundRPCSpe c.dropInboundRPCSpecs(normalized, "context_done") return nil, err } - if c.rpcClosed || c.terminal.Load() { + if c.rpcClosed || c.isRetired() { c.rpcMu.Unlock() releaseInboundRPCGlobalBatch(globals) c.dropInboundRPCSpecs(normalized, "scheduler_closed") @@ -630,113 +499,12 @@ func (c *Conn) dropInboundRPCSpecs(specs []inboundRPCSpec, reason string) { } } -// enqueueInboundRPC 是测试和已持有独立 body 的便捷入口。生产收包路径使用 -// reserveInboundRPC -> Copy -> commit，保证真正的 Copy 前预算。 -func (c *Conn) enqueueInboundRPC(ctx context.Context, task inboundRPC) error { - reservation, err := c.reserveInboundRPC(ctx, task.method, task.size) - if err != nil { - return err - } - defer reservation.abort() - return reservation.commit(task) -} - -func (r *inboundRPCReservation) commit(task inboundRPC) error { - result := ErrConnClosed - var ( - committed bool - reschedule bool - queueLen int - queueCap int - ) - r.once.Do(func() { - c := r.conn - c.rpcMu.Lock() - c.rpcReserved-- - if c.rpcClosed || c.terminal.Load() { - c.inflightRPCBytes.Add(-int64(r.size)) - } else { - // The request deadline starts when admission succeeds, not when a worker - // eventually dequeues the request. This bounds total queue + execution - // latency and lets a queued request emit its explicit timeout on time. - if r.deadline.IsZero() { - task.ctx, task.cancel = context.WithCancel(r.ctx) - } else { - task.ctx, task.cancel = context.WithDeadline(r.ctx, r.deadline) - } - task.stopRoot = context.AfterFunc(c.rpcRootCtx, task.cancel) - task.method = r.method - task.enqueuedAt = r.enqueuedAt - task.deadline = r.deadline - task.size = r.size - task.budget = r.global - ticket := &inboundRPCTicket{} - if task.onTimeout != nil { - onTimeout := task.onTimeout - var timeoutOnce sync.Once - ticket.onTimeout = func() { - timeoutOnce.Do(onTimeout) - } - task.onTimeout = ticket.onTimeout - } - task.ticket = ticket - if task.onTimeout != nil && !task.deadline.IsZero() { - taskCtx := task.ctx - task.stopTimeout = context.AfterFunc(taskCtx, func() { - if errors.Is(taskCtx.Err(), context.DeadlineExceeded) { - c.expireInboundRPCTicket(ticket) - } - }) - } - c.rpcQueue = append(c.rpcQueue, task) - queueLen = len(c.rpcQueue) - queueCap = c.rpcQueueSize - if c.rpcRunning < c.rpcMaxInflight && !c.rpcReady { - c.rpcReady = true - reschedule = true - } - committed = true - result = nil - } - c.rpcMu.Unlock() - c.rpcReservationWG.Done() - if !committed { - r.global.release() - } - }) - if committed { - r.conn.metrics.InboundRPCQueued(r.method, queueLen, queueCap) - if reschedule { - r.conn.rpcScheduler.schedule(r.conn) - } - } - return result -} - -func (r *inboundRPCReservation) abort() { +// commit 在一次 rpcMu 临界区内把整批 task append 到队列并立即发布 ready token。 +// 协议 barrier 必须在调用 commit 前完成；延迟发布 token 无法阻止已有 worker +// 从同一连接队列取走新任务，因此不提供虚假的 deferred-schedule 模式。 +func (r *inboundRPCBatchReservation) commit(tasks []inboundRPC) (result error) { if r == nil { - return - } - r.once.Do(func() { - c := r.conn - c.rpcMu.Lock() - c.rpcReserved-- - c.inflightRPCBytes.Add(-int64(r.size)) - c.rpcMu.Unlock() - c.rpcReservationWG.Done() - r.global.release() - }) -} - -// commit 在一次 rpcMu 临界区内把整批 task append 到队列。 -// deferSchedule=false 保持旧的立即调度语义；true 则返回一个幂等 activate -// 函数，让调用方先完成 new_session_created 等协议 barrier 再启动 worker。 -// -// 延迟调度只能延迟本次 commit 新产生的 ready token；调用方应在连接的 -// 首个 admission batch 使用它，不得把它当作已有 worker 的全局暂停锁。 -func (r *inboundRPCBatchReservation) commit(tasks []inboundRPC, deferSchedule bool) (activate func(), result error) { - if r == nil { - return nil, ErrConnClosed + return ErrConnClosed } result = ErrConnClosed var ( @@ -769,7 +537,7 @@ func (r *inboundRPCBatchReservation) commit(tasks []inboundRPC, deferSchedule bo if len(tasks) != len(r.entries) { c.inflightRPCBytes.Add(-r.totalSize) result = errInboundRPCBatchTaskCount - } else if c.rpcClosed || c.terminal.Load() { + } else if c.rpcClosed || c.isRetired() { c.inflightRPCBytes.Add(-r.totalSize) } else { prepared := make([]inboundRPC, len(tasks)) @@ -831,19 +599,10 @@ func (r *inboundRPCBatchReservation) commit(tasks []inboundRPC, deferSchedule bo r.conn.metrics.InboundRPCQueued(entry.method, firstQueueLen+i, queueCap) } if reschedule { - var once sync.Once - activate = func() { - once.Do(func() { - r.conn.rpcScheduler.schedule(r.conn) - }) - } - if !deferSchedule { - activate() - activate = nil - } + r.conn.rpcScheduler.schedule(r.conn) } } - return activate, result + return result } func (r *inboundRPCBatchReservation) abort() { @@ -884,9 +643,6 @@ func (c *Conn) takeInboundRPC() (task inboundRPC, ok, reschedule bool) { c.rpcQueue = nil } c.rpcRunning++ - if task.ticket != nil { - task.ticket.state.Store(inboundRPCTicketRunning) - } c.rpcWG.Add(1) if len(c.rpcQueue) > 0 && c.rpcRunning < c.rpcMaxInflight { c.rpcReady = true @@ -920,10 +676,7 @@ func (c *Conn) runInboundRPC(task inboundRPC) { } func (c *Conn) finishInboundRPC(task inboundRPC) { - if task.ticket != nil { - task.ticket.state.Store(inboundRPCTicketDone) - } - timeoutHandoff := stopInboundRPCTask(task) + stopInboundRPCTask(task) var reschedule bool c.rpcMu.Lock() c.rpcRunning-- @@ -940,12 +693,6 @@ func (c *Conn) finishInboundRPC(task inboundRPC) { // with a newly admitted body under the same byte accounting. task = inboundRPC{} reservation.release() - if timeoutHandoff != nil { - // stopTimeout(false) means the deadline callback may already have read - // Running but not yet entered ticket.onTimeout. Calling the sync.Once wrapper - // here either performs or joins that response before owner release/Abort. - timeoutHandoff() - } if release != nil { release() } @@ -957,8 +704,8 @@ func (c *Conn) finishInboundRPC(task inboundRPC) { // expireInboundRPCTicket removes a request that is still queued and returns its // memory/task reservations immediately. If the worker won the dequeue race, the -// same callback only signals the running request's response gate; its body remains -// owned until the handler exits. +// callback does nothing: the running handler owns the only terminal response and +// its deadline is represented solely by context cancellation. func (c *Conn) expireInboundRPCTicket(ticket *inboundRPCTicket) { if ticket == nil { return @@ -986,7 +733,6 @@ func (c *Conn) expireInboundRPCTicket(ticket *inboundRPCTicket) { } } c.inflightRPCBytes.Add(-int64(task.size)) - ticket.state.Store(inboundRPCTicketDone) found = true break } @@ -999,7 +745,7 @@ func (c *Conn) expireInboundRPCTicket(ticket *inboundRPCTicket) { method := task.method reservation := task.budget release := task.release - _ = stopInboundRPCTask(task) + stopInboundRPCTask(task) // Drop the run/context closures before returning the byte reservation. Otherwise an // onTimeout callback that blocks or performs a slow write can keep the copied request body // reachable after the global scheduler has advertised those bytes as available again. @@ -1014,24 +760,14 @@ func (c *Conn) expireInboundRPCTicket(ticket *inboundRPCTicket) { } return } - if ticket.state.Load() == inboundRPCTicketRunning && ticket.onTimeout != nil { - ticket.onTimeout() - } } -// stopInboundRPCTask disarms callbacks before canceling the context so a normal -// completion or connection close cannot manufacture an RPC_TIMEOUT response. -// If the runtime already started a deadline callback, the returned sync.Once -// wrapper is a mandatory handoff: callers invoke it before owner release so a -// callback paused between ticket-state inspection and response-gate claim cannot -// publish into a later flight generation. -func stopInboundRPCTask(task inboundRPC) (timeoutHandoff func()) { +// stopInboundRPCTask disarms queue-expiration cleanup before canceling the +// context. Once a worker dequeues the task, the deadline only cancels the +// handler; it never races the handler with an early RPC_TIMEOUT response. +func stopInboundRPCTask(task inboundRPC) { if task.stopTimeout != nil { - stopped := task.stopTimeout() - if !stopped && task.ctx != nil && errors.Is(task.ctx.Err(), context.DeadlineExceeded) && - task.ticket != nil && task.ticket.onTimeout != nil { - timeoutHandoff = task.ticket.onTimeout - } + task.stopTimeout() } if task.stopRoot != nil { task.stopRoot() @@ -1039,7 +775,6 @@ func stopInboundRPCTask(task inboundRPC) (timeoutHandoff func()) { if task.cancel != nil { task.cancel() } - return timeoutHandoff } func (c *Conn) closeInboundRPCScheduler() { @@ -1051,8 +786,8 @@ func (c *Conn) closeInboundRPCScheduler() { } // beginCloseInboundRPCScheduler publishes closure, cancels running work and releases queued -// requests without waiting for handlers. ForceClose uses this phase before transport.Close so a -// pathological/blocking transport implementation cannot leave the RPC admission gate open. +// requests without waiting for handlers. Shutdown publishes this phase before transport.Close so +// a pathological/blocking transport implementation cannot leave the RPC admission gate open. func (c *Conn) beginCloseInboundRPCScheduler() { if c.rpcScheduler == nil { return @@ -1078,18 +813,12 @@ func (c *Conn) beginCloseInboundRPCScheduler() { for i := range queued { task := queued[i] queued[i] = inboundRPC{} - if task.ticket != nil { - task.ticket.state.Store(inboundRPCTicketDone) - } method := task.method reservation := task.budget release := task.release - timeoutHandoff := stopInboundRPCTask(task) + stopInboundRPCTask(task) task = inboundRPC{} reservation.release() - if timeoutHandoff != nil { - timeoutHandoff() - } if release != nil { release() } diff --git a/internal/mtprotoedge/inbound_rpc_batch_test.go b/internal/mtprotoedge/inbound_rpc_batch_test.go index 52553152..423b1ab2 100644 --- a/internal/mtprotoedge/inbound_rpc_batch_test.go +++ b/internal/mtprotoedge/inbound_rpc_batch_test.go @@ -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) } diff --git a/internal/mtprotoedge/inbound_rpc_test.go b/internal/mtprotoedge/inbound_rpc_test.go index 1f16fe20..e58b8f1e 100644 --- a/internal/mtprotoedge/inbound_rpc_test.go +++ b/internal/mtprotoedge/inbound_rpc_test.go @@ -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) diff --git a/internal/mtprotoedge/outbound.go b/internal/mtprotoedge/outbound.go index 4f8b2412..1bd21b19 100644 --- a/internal/mtprotoedge/outbound.go +++ b/internal/mtprotoedge/outbound.go @@ -336,10 +336,7 @@ func (c *Conn) Close() { // close. It closes both producer gates and cancels RPC work before any potentially blocking // transport.Close call, so a timed-out batch close cannot keep accepting memory/work. func (c *Conn) beginTerminalShutdown() { - c.terminal.Store(true) - // Retirement is irreversible. SessionManager activation only uses CAS from - // provisional/claiming, so a stale goroutine cannot publish this Conn again. - c.lifecycle.Store(uint32(connLifecycleRetired)) + c.retire() c.signalOutboundStop() c.beginCloseInboundRPCScheduler() } @@ -367,16 +364,6 @@ func (c *Conn) waitOutboundShutdownUntil(timeout time.Duration) bool { } } -// ForceClose 停止连接并关闭底层 transport。 -// 仅用于授权撤销 / destroy_auth_key 这类“必须让对端立即断线”的路径；普通生命周期仍由 -// serveConn 统一关闭 transport，避免正常 push/索引清理把长连接误伤成硬断。 -func (c *Conn) ForceClose() { - c.beginTerminalShutdown() - c.closeTransport() - c.closeInboundRPCScheduler() - c.waitOutboundShutdown() -} - // closeTransport 只关闭物理 transport，不等待 outbound actor。写失败路径运行在 // actor 自身 goroutine 中，若在这里调用 Close 会等待 outboundDone 而自锁。 func (c *Conn) closeTransport() { @@ -395,7 +382,7 @@ func (c *Conn) closeTransport() { } // failTransport 把不可恢复的写错误提升为连接级 terminal failure。它只负责 -// 标记 terminal + 关闭 socket；handleOutboundOp 返回后，actor 自己发停止信号并退出， +// 把 lifecycle 推进到 retired 并关闭 socket；handleOutboundOp 返回后，actor 自己发停止信号并退出， // serveConn 被 Close 解开 Recv 后负责注销索引。 func (c *Conn) failTransport() { // Publish both producer gates before Close: a custom/broken transport may block in @@ -417,10 +404,9 @@ func (c *Conn) fenceUndeliveredRPCResult() { // A replacement/shutdown that already published terminal owns physical // lifecycle cleanup (and may intentionally transfer the lease). Only the // resultless task that wins false->true is allowed to close this generation. - if !c.terminal.CompareAndSwap(false, true) { + if !c.retire() { return } - c.lifecycle.Store(uint32(connLifecycleRetired)) c.signalOutboundStop() if c.transportLease != nil { c.transportLease.startCloseAlreadyFenced() @@ -461,11 +447,6 @@ func (c *Conn) Send(ctx context.Context, t proto.MessageType, msg bin.Encoder) e return c.send(ctx, t, msg, false) } -// SendPriority 加密并优先发送一条 server 控制消息。 -func (c *Conn) SendPriority(ctx context.Context, t proto.MessageType, msg bin.Encoder) error { - return c.send(ctx, t, msg, true) -} - // SendRequiredControl writes a protocol-critical control message before the caller commits // the state transition guarded by that message. One absolute deadline covers encode admission, // body-budget reservation, control-queue admission and the physical transport write. A failure @@ -515,7 +496,7 @@ func (c *Conn) sendBestEffort(ctx context.Context, t proto.MessageType, msg bin. if c.outbound == nil || c.outboundControl == nil { return ErrConnClosed } - if c.terminal.Load() { + if c.isRetired() { return ErrConnClosed } writeCtx := context.Background() @@ -645,7 +626,7 @@ func (c *Conn) SendAsync(ctx context.Context, t proto.MessageType, msg bin.Encod if c.outbound == nil || c.outboundControl == nil { return ErrConnClosed } - if c.terminal.Load() { + if c.isRetired() { return ErrConnClosed } op, err := c.newOutboundSendOp(ctx, t, msg, nil, true) @@ -677,7 +658,7 @@ func (c *Conn) SendAsync(ctx context.Context, t proto.MessageType, msg bin.Encod // AckServerMessages 接收客户端 msgs_ack，释放已确认的 server 出站消息。 func (c *Conn) AckServerMessages(ids []int64) { - if len(ids) == 0 || c.outbound == nil || c.outboundControl == nil || c.terminal.Load() { + if len(ids) == 0 || c.outbound == nil || c.outboundControl == nil || c.isRetired() { return } op, err := c.newOutboundVectorOp(outboundAck, ids) @@ -790,7 +771,7 @@ func (c *Conn) enqueueOutboundRegistered(ctx context.Context, op outboundOp) err if ctx == nil { ctx = context.Background() } - if c.terminal.Load() { + if c.isRetired() { return ErrConnClosed } q := c.outbound @@ -820,7 +801,7 @@ func (c *Conn) enqueueOutboundRegistered(ctx context.Context, op outboundOp) err func (c *Conn) beginOutboundEnqueue() bool { c.outboundEnqueueMu.Lock() defer c.outboundEnqueueMu.Unlock() - if c.outboundClosing || c.terminal.Load() { + if c.outboundClosing || c.isRetired() { return false } c.outboundEnqueueWG.Add(1) @@ -840,14 +821,14 @@ func (c *Conn) outboundLoop() { close(c.outboundDone) }() for { - if c.terminal.Load() { + if c.isRetired() { c.signalOutboundStop() c.drainOutbound() return } select { case op := <-c.outboundControl: - if c.terminal.Load() { + if c.isRetired() { op.releaseReservation(c.outboundTrackedBudget) op.finish(outboundResult{err: ErrConnClosed}) c.signalOutboundStop() @@ -855,7 +836,7 @@ func (c *Conn) outboundLoop() { return } c.handleOutboundOp(state, op) - if c.terminal.Load() { + if c.isRetired() { c.signalOutboundStop() c.drainOutbound() return @@ -868,7 +849,7 @@ func (c *Conn) outboundLoop() { c.drainOutbound() return case op := <-c.outboundControl: - if c.terminal.Load() { + if c.isRetired() { op.releaseReservation(c.outboundTrackedBudget) op.finish(outboundResult{err: ErrConnClosed}) c.signalOutboundStop() @@ -877,7 +858,7 @@ func (c *Conn) outboundLoop() { } c.handleOutboundOp(state, op) case op := <-c.outbound: - if c.terminal.Load() { + if c.isRetired() { op.releaseReservation(c.outboundTrackedBudget) op.finish(outboundResult{err: ErrConnClosed}) c.signalOutboundStop() @@ -886,7 +867,7 @@ func (c *Conn) outboundLoop() { } c.handleOutboundOp(state, op) } - if c.terminal.Load() { + if c.isRetired() { c.signalOutboundStop() c.drainOutbound() return diff --git a/internal/mtprotoedge/outbound_required_control_test.go b/internal/mtprotoedge/outbound_required_control_test.go index 5d4fca21..abe405bd 100644 --- a/internal/mtprotoedge/outbound_required_control_test.go +++ b/internal/mtprotoedge/outbound_required_control_test.go @@ -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 { diff --git a/internal/mtprotoedge/outbound_test.go b/internal/mtprotoedge/outbound_test.go index 798b4eba..82b3cf12 100644 --- a/internal/mtprotoedge/outbound_test.go +++ b/internal/mtprotoedge/outbound_test.go @@ -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 { diff --git a/internal/mtprotoedge/production_compat_test.go b/internal/mtprotoedge/production_compat_test.go new file mode 100644 index 00000000..fecfbcff --- /dev/null +++ b/internal/mtprotoedge/production_compat_test.go @@ -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 +} diff --git a/internal/mtprotoedge/rpc_result_cache.go b/internal/mtprotoedge/rpc_result_cache.go index 98fc5827..11eb7d82 100644 --- a/internal/mtprotoedge/rpc_result_cache.go +++ b/internal/mtprotoedge/rpc_result_cache.go @@ -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 diff --git a/internal/mtprotoedge/rpc_result_delivery_test.go b/internal/mtprotoedge/rpc_result_delivery_test.go index 30905727..7ba45954 100644 --- a/internal/mtprotoedge/rpc_result_delivery_test.go +++ b/internal/mtprotoedge/rpc_result_delivery_test.go @@ -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 { diff --git a/internal/mtprotoedge/rpc_test.go b/internal/mtprotoedge/rpc_test.go index 8ec2e191..9ed7c4fb 100644 --- a/internal/mtprotoedge/rpc_test.go +++ b/internal/mtprotoedge/rpc_test.go @@ -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{} diff --git a/internal/mtprotoedge/server.go b/internal/mtprotoedge/server.go index c8d0752f..b141f4dd 100644 --- a/internal/mtprotoedge/server.go +++ b/internal/mtprotoedge/server.go @@ -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) diff --git a/internal/mtprotoedge/session_activation_test.go b/internal/mtprotoedge/session_activation_test.go index 7838e642..c6fbe638 100644 --- a/internal/mtprotoedge/session_activation_test.go +++ b/internal/mtprotoedge/session_activation_test.go @@ -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()) } } diff --git a/internal/mtprotoedge/session_barrier_integration_test.go b/internal/mtprotoedge/session_barrier_integration_test.go index e3f98e8e..c8142c92 100644 --- a/internal/mtprotoedge/session_barrier_integration_test.go +++ b/internal/mtprotoedge/session_barrier_integration_test.go @@ -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() } diff --git a/internal/mtprotoedge/session_manager.go b/internal/mtprotoedge/session_manager.go index fc15e676..6185271f 100644 --- a/internal/mtprotoedge/session_manager.go +++ b/internal/mtprotoedge/session_manager.go @@ -19,9 +19,6 @@ import ( // ErrSessionNotFound 表示目标 session 当前无活跃连接。 var ErrSessionNotFound = errors.New("session not found") -// ErrSessionAmbiguous 表示仅用 session_id 无法唯一定位连接。 -var ErrSessionAmbiguous = errors.New("session id is shared by multiple auth keys") - var ( ErrSessionActivationSuperseded = errors.New("session activation superseded") ErrSessionActivationFence = errors.New("session activation could not fence previous writer") @@ -119,8 +116,8 @@ type SessionLifecycleObserver interface { // SessionManager 是活跃连接注册表，支持按 session / auth-key / user 查找并主动 push。 // -// 它管理运行态的在线连接，与持久化的 store.SessionStore 互补：后者记录 session 数据， -// 前者持有可发送的活跃连接。所有方法并发安全。 +// 它只管理进程内运行态，持有可发送的活跃连接；协议可恢复事实由 auth key、客户端重连 +// 和 durable updates/difference 链路承担。所有方法并发安全。 type SessionManager struct { mu sync.RWMutex bySession map[sessionKey]*Conn @@ -129,7 +126,6 @@ type SessionManager struct { // is on the wire and PublishActivation validates the same owner. claims map[sessionKey]*Conn claimsByAuth map[[8]byte]map[int64]*Conn // raw authKeyID -> sessionID -> provisional claim - bySessionID map[int64]map[[8]byte]*Conn // sessionID → raw authKeyID → Conn，用于兼容旧 API 的唯一性检查 byAuthKey map[[8]byte]map[int64]*Conn // raw authKeyID → sessionID → Conn byBusinessAuthKey map[[8]byte]map[sessionKey]*Conn byUser map[int64]map[sessionKey]*Conn @@ -154,7 +150,6 @@ func NewSessionManager(log *zap.Logger) *SessionManager { bySession: make(map[sessionKey]*Conn), claims: make(map[sessionKey]*Conn), claimsByAuth: make(map[[8]byte]map[int64]*Conn), - bySessionID: make(map[int64]map[[8]byte]*Conn), byAuthKey: make(map[[8]byte]map[int64]*Conn), byBusinessAuthKey: make(map[[8]byte]map[sessionKey]*Conn), byUser: make(map[int64]map[sessionKey]*Conn), @@ -189,7 +184,7 @@ func (m *SessionManager) BeginActivation(c *Conn) error { key := connSessionKey(c) retired := make([]*Conn, 0, 2) m.mu.Lock() - if c.terminal.Load() || !c.isPhysicalTransportCurrentOpen() || c.lifecycleState() != connLifecycleClaiming { + if !c.isPhysicalTransportCurrentOpen() || c.lifecycleState() != connLifecycleClaiming { c.beginTerminalShutdown() m.mu.Unlock() return ErrConnClosed @@ -248,7 +243,7 @@ func (m *SessionManager) PublishActivation(c *Conn) error { if m.claims[key] != c { return ErrSessionActivationSuperseded } - if c.terminal.Load() || c.lifecycleState() != connLifecycleClaiming { + if c.lifecycleState() != connLifecycleClaiming { m.removeClaimLocked(key, c) return ErrConnClosed } @@ -265,7 +260,6 @@ func (m *SessionManager) PublishActivation(c *Conn) error { } m.removeClaimLocked(key, c) m.bySession[key] = c - addSessionIDIndex(m.bySessionID, c.sessionID, c.authKeyID, c) addConnIndex(m.byAuthKey, c.authKeyID, c.sessionID, c) if businessAuthKeyID, resolved := c.BusinessAuthKeyID(); resolved { addBusinessAuthKeyIndex(m.byBusinessAuthKey, businessAuthKeyID, key, c) @@ -306,32 +300,6 @@ func (m *SessionManager) AbortActivation(c *Conn) { } } -// Register is retained for tests and embedders that do not have a wire-level -// required-control barrier. Production encrypted traffic uses the explicit -// BeginActivation -> SendRequiredControl -> PublishActivation sequence. -func (m *SessionManager) Register(c *Conn) error { - if c == nil { - return ErrSessionActivationSuperseded - } - if c.isActive() { - m.mu.RLock() - current := m.bySession[connSessionKey(c)] - m.mu.RUnlock() - if current == c { - return nil - } - return ErrSessionActivationSuperseded - } - if err := m.BeginActivation(c); err != nil { - return err - } - if err := m.PublishActivation(c); err != nil { - m.AbortActivation(c) - return err - } - return nil -} - // Unregister 注销一个连接（仅当它仍是当前注册的同一对象，避免误删重连后的新连接）。 // 观察者对未登录连接（userID=0）也回调：业务层据此清理按 session 维度的缓存条目， // 否则未登录连接的元数据只能等容量上限驱逐。 @@ -373,38 +341,6 @@ func (m *SessionManager) Unregister(c *Conn) { } } -// DestroySession 移除指定 session 的运行态索引，供 MTProto destroy_session 使用。 -func (m *SessionManager) DestroySession(sessionID int64) bool { - m.mu.Lock() - c, key, ok, ambiguous := m.uniqueSessionLocked(sessionID) - if ambiguous || !ok { - if !ambiguous { - m.dropPendingBySessionLocked(sessionID) - } - m.mu.Unlock() - return false - } - offlineUser := m.retireConnLocked(c, true) - lastForUser := offlineUser != 0 && len(m.byUser[offlineUser]) == 0 - observer := m.lifecycle - m.log.Debug("Session destroyed", - zap.String("auth_key_id", sessionKeyLog(key.authKeyID)), - zap.Int64("session_id", sessionID), - zap.Int("online", len(m.bySession)), - ) - m.mu.Unlock() - if !forceCloseConnBatch([]*Conn{c}, forceCloseBatchTimeout) { - m.log.Warn("Destroyed session close exceeded shared deadline", - zap.String("auth_key_id", sessionKeyLog(key.authKeyID)), - zap.Int64("session_id", sessionID), - ) - } - if observer != nil && offlineUser != 0 { - observer.SessionOffline(key.authKeyID, sessionID, offlineUser, lastForUser) - } - return true -} - // DestroySessionForAuthKey 精确移除某个 raw auth_key_id 下的 session。 func (m *SessionManager) DestroySessionForAuthKey(authKeyID [8]byte, sessionID int64) bool { m.mu.Lock() @@ -447,22 +383,6 @@ func (m *SessionManager) DestroySessionForAuthKey(authKeyID [8]byte, sessionID i return true } -// BindUser 缓存 session 的授权用户。userID=0 表示当前 auth_key 已确认未登录。 -// 登录后绑定非 0 userID，使其可经 PushToUser 收到推送。 -func (m *SessionManager) BindUser(sessionID, userID int64) { - m.mu.Lock() - c, key, ok, ambiguous := m.uniqueSessionLocked(sessionID) - if ambiguous || !ok { - if ambiguous { - m.log.Warn("Skip BindUser for ambiguous session_id", zap.Int64("session_id", sessionID)) - } - m.mu.Unlock() - return - } - m.bindUserLocked(c, key, userID) - m.mu.Unlock() -} - // BindUserForAuthKey 缓存指定 raw auth_key_id + session_id 的授权用户。 func (m *SessionManager) BindUserForAuthKey(authKeyID [8]byte, sessionID, userID int64) { m.mu.Lock() @@ -500,48 +420,6 @@ func (m *SessionManager) bindUserLocked(c *Conn, key sessionKey, userID int64) { } } -// UserID 返回 session 当前缓存的登录用户 id。未绑定或离线时 ok=false。 -func (m *SessionManager) UserID(sessionID int64) (int64, bool) { - m.mu.RLock() - c, _, ok, ambiguous := m.uniqueSessionLocked(sessionID) - m.mu.RUnlock() - if ambiguous || !ok { - return 0, false - } - userID := c.userID.Load() - if userID == 0 { - return 0, false - } - return userID, true -} - -// UserIDForAuthKey 返回指定 raw auth_key_id + session_id 当前缓存的登录用户 id。 -func (m *SessionManager) UserIDForAuthKey(authKeyID [8]byte, sessionID int64) (int64, bool) { - m.mu.RLock() - c, ok := m.bySession[sessionKey{authKeyID: authKeyID, sessionID: sessionID}] - m.mu.RUnlock() - if !ok { - return 0, false - } - userID := c.userID.Load() - if userID == 0 { - return 0, false - } - return userID, true -} - -// UserIDResolved 返回 session 的 user_id 授权状态是否已经查过。 -// resolved=true 且 userID=0 表示该 session 当前未登录。 -func (m *SessionManager) UserIDResolved(sessionID int64) (int64, bool) { - m.mu.RLock() - c, _, ok, ambiguous := m.uniqueSessionLocked(sessionID) - m.mu.RUnlock() - if ambiguous || !ok { - return 0, false - } - return c.UserIDResolved() -} - // UserIDResolvedForAuthKey 返回指定 raw auth_key_id + session_id 的 user_id 缓存状态。 func (m *SessionManager) UserIDResolvedForAuthKey(authKeyID [8]byte, sessionID int64) (int64, bool) { m.mu.RLock() @@ -553,21 +431,6 @@ func (m *SessionManager) UserIDResolvedForAuthKey(authKeyID [8]byte, sessionID i return c.UserIDResolved() } -// BindAuthKey 缓存业务视角 auth_key_id（temp auth_key 解析后的 perm auth_key）。 -func (m *SessionManager) BindAuthKey(sessionID int64, authKeyID [8]byte) { - m.mu.Lock() - c, key, ok, ambiguous := m.uniqueSessionLocked(sessionID) - if ambiguous || !ok { - if ambiguous { - m.log.Warn("Skip BindAuthKey for ambiguous session_id", zap.Int64("session_id", sessionID)) - } - m.mu.Unlock() - return - } - m.bindAuthKeyLocked(c, key, authKeyID) - m.mu.Unlock() -} - // BindAuthKeyForSession 缓存指定 raw auth_key_id + session_id 的业务 auth_key_id。 func (m *SessionManager) BindAuthKeyForSession(rawAuthKeyID [8]byte, sessionID int64, authKeyID [8]byte) { m.mu.Lock() @@ -603,18 +466,6 @@ func (m *SessionManager) bindAuthKeyLocked(c *Conn, key sessionKey, authKeyID [8 } } -// AuthKeyID 返回 session 缓存的业务视角 auth_key_id。 -// ok=false 表示该连接尚未完成 temp→perm 解析。 -func (m *SessionManager) AuthKeyID(sessionID int64) ([8]byte, bool) { - m.mu.RLock() - c, _, ok, ambiguous := m.uniqueSessionLocked(sessionID) - m.mu.RUnlock() - if ambiguous || !ok { - return [8]byte{}, false - } - return c.BusinessAuthKeyID() -} - // AuthKeyIDForSession 返回指定 raw auth_key_id + session_id 缓存的业务 auth_key_id。 func (m *SessionManager) AuthKeyIDForSession(rawAuthKeyID [8]byte, sessionID int64) ([8]byte, bool) { m.mu.RLock() @@ -863,28 +714,6 @@ func (m *SessionManager) UnbindAuthKey(authKeyID [8]byte) int { return count } -// SetReceivesUpdates 标记 session 是否已完成 updates 同步入口。 -// -// TDesktop 登录后会先调用 updates.getState/getDifference 建立本地同步基线。 -// 在此之前收到的主动 updates 先暂存，待 session 可接收后再异步下发。 -func (m *SessionManager) SetReceivesUpdates(sessionID int64, receives bool) { - m.mu.Lock() - c, key, ok, ambiguous := m.uniqueSessionLocked(sessionID) - if ambiguous || !ok { - if ambiguous { - m.log.Warn("Skip SetReceivesUpdates for ambiguous session_id", zap.Int64("session_id", sessionID)) - } - m.mu.Unlock() - return - } - owner, start := m.setReceivesUpdatesLocked(c, key, receives) - m.mu.Unlock() - - if start { - go m.runFlush(c, key, owner, 0) - } -} - // setReceivesUpdatesLocked 是置位/复位的共同内核，调用方须持有 m.mu。 // 置位且有暂存时不立即置 receivesUpdates：标记 flushing 并返回该批暂存所属的 userID， // 交由 runFlush 排空后原子置位，期间新到推送继续进 pending，保证暂存与实时推送的 @@ -1059,26 +888,6 @@ func (m *SessionManager) SetReceivesUpdatesForAuthKey(authKeyID [8]byte, session } } -// PushToSession 向指定 session 推送一条消息。 -func (m *SessionManager) PushToSession(ctx context.Context, sessionID int64, t proto.MessageType, msg bin.Encoder) error { - m.mu.RLock() - c, key, ok, ambiguous := m.uniqueSessionLocked(sessionID) - if ambiguous { - m.mu.RUnlock() - return ErrSessionAmbiguous - } - if !ok { - m.mu.RUnlock() - return ErrSessionNotFound - } - ready := c.receivesUpdates.Load() - m.mu.RUnlock() - if ready { - return c.Send(ctx, t, msg) - } - return m.queueOrSendPrepared(ctx, key, t, msg) -} - // PushToSessionForAuthKey 向指定 raw auth_key_id + session_id 推送一条消息。 func (m *SessionManager) PushToSessionForAuthKey(ctx context.Context, authKeyID [8]byte, sessionID int64, t proto.MessageType, msg bin.Encoder) error { m.mu.RLock() @@ -1134,18 +943,6 @@ func (m *SessionManager) PushToSessionForAuthKeyImmediate(ctx context.Context, a return c.SendBestEffort(ctx, t, msg, 2*time.Second) } -// PushToUser 向某 user 所有活跃连接推送，返回已发送或已暂存的连接数。 -// 发送在释放锁后进行，避免持锁阻塞于网络 IO。 -func (m *SessionManager) PushToUser(ctx context.Context, userID int64, t proto.MessageType, msg bin.Encoder) (int, error) { - return m.PushToUserExceptAuthKeySession(ctx, userID, [8]byte{}, 0, t, msg) -} - -// PushToUserExceptSession 向某 user 所有活跃连接推送，但跳过指定 session。 -// 未完成 updates 同步入口的 session 会先暂存，等 SetReceivesUpdates(true) 后再发。 -func (m *SessionManager) PushToUserExceptSession(ctx context.Context, userID, excludeSessionID int64, t proto.MessageType, msg bin.Encoder) (int, error) { - return m.pushToUser(ctx, userID, nil, excludeSessionID, t, msg) -} - // PushToUserExceptAuthKeySession 向某 user 所有活跃连接推送，跳过指定 raw auth_key + session。 func (m *SessionManager) PushToUserExceptAuthKeySession(ctx context.Context, userID int64, excludeAuthKeyID [8]byte, excludeSessionID int64, t proto.MessageType, msg bin.Encoder) (int, error) { return m.pushToUser(ctx, userID, &excludeAuthKeyID, excludeSessionID, t, msg) @@ -1275,10 +1072,6 @@ func (m *SessionManager) PushToUserTransientExceptAuthKeySession(ctx context.Con }) } -func (m *SessionManager) PushToUserExceptSessionBestEffort(ctx context.Context, userID, excludeSessionID int64, t proto.MessageType, msg bin.Encoder, timeout time.Duration) (int, error) { - return m.pushToUserBestEffort(ctx, userID, nil, excludeSessionID, t, msg, timeout) -} - func (m *SessionManager) PushToUserExceptAuthKeySessionBestEffort(ctx context.Context, userID int64, excludeAuthKeyID [8]byte, excludeSessionID int64, t proto.MessageType, msg bin.Encoder, timeout time.Duration) (int, error) { return m.pushToUserBestEffort(ctx, userID, &excludeAuthKeyID, excludeSessionID, t, msg, timeout) } @@ -1490,13 +1283,6 @@ func (m *SessionManager) pushToUserWithSender(ctx context.Context, userID int64, return sent + queued, firstErr } -// Online 返回当前活跃连接数。 -func (m *SessionManager) Online() int { - m.mu.RLock() - defer m.mu.RUnlock() - return len(m.bySession) -} - // ActiveRawAuthKeyIDs 返回当前物理连接实际使用的 raw auth_key_id 去重快照。 // maintenance 用它保护“已建 key 但尚未登录”的长连接不被 orphan GC 删除；不能用 // business/temp→perm key 替代，否则活跃 temp 连接仍可能误删。 @@ -1757,25 +1543,6 @@ func (m *SessionManager) OnlineChannelIDsSnapshot() []int64 { return out } -// OnlineChannelIDsAfter is retained for bounded diagnostics/tests. Production recovery takes one -// OnlineChannelIDsSnapshot per generation and slices it into pages, avoiding repeated full scans. -func (m *SessionManager) OnlineChannelIDsAfter(afterChannelID int64, limit int) []int64 { - if limit <= 0 { - return nil - } - const maxRecoveryPage = 4096 - if limit > maxRecoveryPage { - limit = maxRecoveryPage - } - all := m.OnlineChannelIDsSnapshot() - start := sort.Search(len(all), func(i int) bool { return all[i] > afterChannelID }) - end := start + limit - if end > len(all) { - end = len(all) - } - return all[start:end] -} - func (m *SessionManager) onlineChannelUsers(index map[int64]map[sessionKey]int64, channelID int64, limit int) []int64 { if channelID == 0 { return nil @@ -1813,7 +1580,6 @@ func (m *SessionManager) removeLocked(c *Conn, dropPending bool) int64 { return 0 } delete(m.bySession, key) - removeSessionIDIndex(m.bySessionID, c.sessionID, c.authKeyID) removeConnIndex(m.byAuthKey, c.authKeyID, c.sessionID) if businessAuthKeyID, resolved := c.BusinessAuthKeyID(); resolved { removeBusinessAuthKeyIndex(m.byBusinessAuthKey, businessAuthKeyID, key) @@ -2100,44 +1866,6 @@ func (m *SessionManager) queuePreparedLocked(key sessionKey, t proto.MessageType return true } -// queueLocked remains as a test/internal single-target convenience. Production fan-out prepares -// outside m.mu and calls queuePreparedLocked so TL encoding never serializes the session registry. -func (m *SessionManager) queueLocked(key sessionKey, t proto.MessageType, msg bin.Encoder) bool { - encoded, reservation, err := m.preparePendingPush(context.Background(), msg) - if err != nil { - m.log.Debug("Drop pending push outside byte budget", - zap.String("auth_key_id", sessionKeyLog(key.authKeyID)), - zap.Int64("session_id", key.sessionID), - zap.Error(err), - ) - return false - } - defer reservation.release() - return m.queuePreparedLocked(key, t, encoded, reservation) -} - -func (m *SessionManager) uniqueSessionLocked(sessionID int64) (*Conn, sessionKey, bool, bool) { - set := m.bySessionID[sessionID] - if len(set) == 0 { - return nil, sessionKey{}, false, false - } - if len(set) > 1 { - return nil, sessionKey{}, false, true - } - for authKeyID, c := range set { - return c, sessionKey{authKeyID: authKeyID, sessionID: sessionID}, true, false - } - return nil, sessionKey{}, false, false -} - -func (m *SessionManager) dropPendingBySessionLocked(sessionID int64) { - for key := range m.pending { - if key.sessionID == sessionID { - m.deletePendingLocked(key) - } - } -} - // RunPendingSweeper 周期回收长期滞留的 pending 暂存：被动老化（queueLocked/takePendingLocked） // 只在「有新推送」或「就绪后取出」时触发，对「已注册但迟迟不调 getState、又恰好没有新推送、 // 也不断连（持续 ping 保活）」的连接无法回收其超龄 pending。本 sweeper 给出一个主动兜底， @@ -2236,24 +1964,6 @@ func removeBusinessAuthKeyIndex(idx map[[8]byte]map[sessionKey]*Conn, authKeyID } } -func addSessionIDIndex(idx map[int64]map[[8]byte]*Conn, sessionID int64, authKeyID [8]byte, c *Conn) { - set := idx[sessionID] - if set == nil { - set = make(map[[8]byte]*Conn) - idx[sessionID] = set - } - set[authKeyID] = c -} - -func removeSessionIDIndex(idx map[int64]map[[8]byte]*Conn, sessionID int64, authKeyID [8]byte) { - if set := idx[sessionID]; set != nil { - delete(set, authKeyID) - if len(set) == 0 { - delete(idx, sessionID) - } - } -} - func addUserIndex(idx map[int64]map[sessionKey]*Conn, userID int64, key sessionKey, c *Conn) { set := idx[userID] if set == nil { diff --git a/internal/mtprotoedge/session_manager_test.go b/internal/mtprotoedge/session_manager_test.go index 3d3d5bbc..0f3829ba 100644 --- a/internal/mtprotoedge/session_manager_test.go +++ b/internal/mtprotoedge/session_manager_test.go @@ -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() { diff --git a/internal/mtprotoedge/shutdown_gate_test.go b/internal/mtprotoedge/shutdown_gate_test.go index f953ee3a..fae24f9c 100644 --- a/internal/mtprotoedge/shutdown_gate_test.go +++ b/internal/mtprotoedge/shutdown_gate_test.go @@ -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 { diff --git a/internal/mtprotoedge/transport_ownership_test.go b/internal/mtprotoedge/transport_ownership_test.go index 18e1f90e..ee80403d 100644 --- a/internal/mtprotoedge/transport_ownership_test.go +++ b/internal/mtprotoedge/transport_ownership_test.go @@ -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() } diff --git a/internal/rpc/auth.go b/internal/rpc/auth.go index fbda59ac..efd3b971 100644 --- a/internal/rpc/auth.go +++ b/internal/rpc/auth.go @@ -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)) } }() diff --git a/internal/rpc/channel_fanout_dispatcher_test.go b/internal/rpc/channel_fanout_dispatcher_test.go index ec4f8b1a..11f84e96 100644 --- a/internal/rpc/channel_fanout_dispatcher_test.go +++ b/internal/rpc/channel_fanout_dispatcher_test.go @@ -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 { diff --git a/internal/rpc/convert_auth.go b/internal/rpc/convert_auth.go index 0c77e3a7..ea3b0f55 100644 --- a/internal/rpc/convert_auth.go +++ b/internal/rpc/convert_auth.go @@ -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 } diff --git a/internal/rpc/deps.go b/internal/rpc/deps.go index 78769c05..dbc8b705 100644 --- a/internal/rpc/deps.go +++ b/internal/rpc/deps.go @@ -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) } diff --git a/internal/rpc/outbox_dispatcher.go b/internal/rpc/outbox_dispatcher.go index 9d9dd57c..57e10389 100644 --- a/internal/rpc/outbox_dispatcher.go +++ b/internal/rpc/outbox_dispatcher.go @@ -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 并发 claim：ClaimPending 用 FOR UPDATE SKIP LOCKED，worker 间认领不重叠。 @@ -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 { diff --git a/internal/rpc/outbox_dispatcher_test.go b/internal/rpc/outbox_dispatcher_test.go index 22f244dc..00e4829b 100644 --- a/internal/rpc/outbox_dispatcher_test.go +++ b/internal/rpc/outbox_dispatcher_test.go @@ -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{}} diff --git a/internal/rpc/phone_push.go b/internal/rpc/phone_push.go index 92d9a0ce..b0a670a1 100644 --- a/internal/rpc/phone_push.go +++ b/internal/rpc/phone_push.go @@ -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) diff --git a/internal/rpc/phone_rpc_test.go b/internal/rpc/phone_rpc_test.go index 34f96569..5184bab0 100644 --- a/internal/rpc/phone_rpc_test.go +++ b/internal/rpc/phone_rpc_test.go @@ -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) diff --git a/internal/rpc/photos.go b/internal/rpc/photos.go index 48066848..35986342 100644 --- a/internal/rpc/photos.go +++ b/internal/rpc/photos.go @@ -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)) } } diff --git a/internal/rpc/push.go b/internal/rpc/push.go index 7a875225..e4ed8838 100644 --- a/internal/rpc/push.go +++ b/internal/rpc/push.go @@ -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)) } } diff --git a/internal/rpc/router.go b/internal/rpc/router.go index 3bf6ac61..f2dde113 100644 --- a/internal/rpc/router.go +++ b/internal/rpc/router.go @@ -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() diff --git a/internal/rpc/router_auth_cache_test.go b/internal/rpc/router_auth_cache_test.go index 1ec10459..3e0321fe 100644 --- a/internal/rpc/router_auth_cache_test.go +++ b/internal/rpc/router_auth_cache_test.go @@ -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, diff --git a/internal/rpc/rpc_testkit_sessions_test.go b/internal/rpc/rpc_testkit_sessions_test.go index 3ee01d96..f88baf74 100644 --- a/internal/rpc/rpc_testkit_sessions_test.go +++ b/internal/rpc/rpc_testkit_sessions_test.go @@ -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 diff --git a/internal/rpc/updates.go b/internal/rpc/updates.go index a5068f8d..f28e81d0 100644 --- a/internal/rpc/updates.go +++ b/internal/rpc/updates.go @@ -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 } diff --git a/internal/store/doc.go b/internal/store/doc.go index dbb9c1fe..344620e2 100644 --- a/internal/store/doc.go +++ b/internal/store/doc.go @@ -1,13 +1,13 @@ // Package store 定义存储接口与协议层 DTO，不含具体实现。 // -// 布局：主包只放接口（AuthKeyStore / SessionStore / UserStore / AuthorizationStore / +// 布局：主包只放接口（AuthKeyStore / UserStore / AuthorizationStore / // CodeStore / UpdateStateStore / UpdateEventStore 等）与协议 DTO；三种后端实现各自独立成对称子包： // - store/memory —— 内存实现，测试替身与本地兜底 // - store/postgres —— PostgreSQL（pgx + sqlc 生成查询 + golang-migrate 迁移） // - store/redisstore —— Redis（go-redis） // // 类型边界：接口签名分两类—— -// - 协议产物用 store 自有 DTO：AuthKeyData、SessionData、PhoneCode（不依赖 tg.*，也非业务实体）； +// - 协议产物用 store 自有 DTO：AuthKeyData、PhoneCode（不依赖 tg.*，也非业务实体）； // - 业务实体直接用 domain：UserStore / AuthorizationStore / MessageStore / UpdateEventStore // 收发 domain.User / domain.Authorization / domain.Message / domain.UpdateEvent。 package store diff --git a/internal/store/memory/auth.go b/internal/store/memory/auth.go index e0fc4ec6..63fd7d80 100644 --- a/internal/store/memory/auth.go +++ b/internal/store/memory/auth.go @@ -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 diff --git a/internal/store/postgres/contiguous_pts_integration_test.go b/internal/store/postgres/contiguous_pts_integration_test.go index fd6e747e..05bef6c1 100644 --- a/internal/store/postgres/contiguous_pts_integration_test.go +++ b/internal/store/postgres/contiguous_pts_integration_test.go @@ -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, ` diff --git a/internal/store/postgres/dispatch_outbox.go b/internal/store/postgres/dispatch_outbox.go index aeb9bad5..a1a862ef 100644 --- a/internal/store/postgres/dispatch_outbox.go +++ b/internal/store/postgres/dispatch_outbox.go @@ -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 diff --git a/internal/store/postgres/dispatch_outbox_exclusion_integration_test.go b/internal/store/postgres/dispatch_outbox_exclusion_integration_test.go new file mode 100644 index 00000000..fc385787 --- /dev/null +++ b/internal/store/postgres/dispatch_outbox_exclusion_integration_test.go @@ -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) + } +} diff --git a/internal/store/postgres/dispatch_outbox_exclusion_test.go b/internal/store/postgres/dispatch_outbox_exclusion_test.go new file mode 100644 index 00000000..3675ee4a --- /dev/null +++ b/internal/store/postgres/dispatch_outbox_exclusion_test.go @@ -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) + } + }) + } +} diff --git a/internal/store/postgres/login_code_delivery.go b/internal/store/postgres/login_code_delivery.go index 827373a1..0453ce06 100644 --- a/internal/store/postgres/login_code_delivery.go +++ b/internal/store/postgres/login_code_delivery.go @@ -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), diff --git a/internal/store/postgres/message_delete.go b/internal/store/postgres/message_delete.go index ab69c730..862b2da4 100644 --- a/internal/store/postgres/message_delete.go +++ b/internal/store/postgres/message_delete.go @@ -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), diff --git a/internal/store/postgres/message_edit.go b/internal/store/postgres/message_edit.go index dbcce1c0..3fbc1649 100644 --- a/internal/store/postgres/message_edit.go +++ b/internal/store/postgres/message_edit.go @@ -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), diff --git a/internal/store/postgres/message_history.go b/internal/store/postgres/message_history.go index c04422f9..b9497606 100644 --- a/internal/store/postgres/message_history.go +++ b/internal/store/postgres/message_history.go @@ -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), diff --git a/internal/store/postgres/message_pin.go b/internal/store/postgres/message_pin.go index 9e11d05b..7ed596ef 100644 --- a/internal/store/postgres/message_pin.go +++ b/internal/store/postgres/message_pin.go @@ -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), diff --git a/internal/store/postgres/message_read.go b/internal/store/postgres/message_read.go index 6f89e6ee..3390b2a4 100644 --- a/internal/store/postgres/message_read.go +++ b/internal/store/postgres/message_read.go @@ -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), diff --git a/internal/store/postgres/message_send.go b/internal/store/postgres/message_send.go index 306e5cce..500f74f7 100644 --- a/internal/store/postgres/message_send.go +++ b/internal/store/postgres/message_send.go @@ -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), diff --git a/internal/store/postgres/phone_change.go b/internal/store/postgres/phone_change.go index d5eff6f0..8dd26719 100644 --- a/internal/store/postgres/phone_change.go +++ b/internal/store/postgres/phone_change.go @@ -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), diff --git a/internal/store/postgres/update_event.go b/internal/store/postgres/update_event.go index 43619f44..f937528e 100644 --- a/internal/store/postgres/update_event.go +++ b/internal/store/postgres/update_event.go @@ -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), diff --git a/internal/store/redisstore/redis.go b/internal/store/redisstore/redis.go index f7491907..cfb9d0a9 100644 --- a/internal/store/redisstore/redis.go +++ b/internal/store/redisstore/redis.go @@ -1,4 +1,4 @@ -// Package redisstore 用 Redis 实现高频易失态的存储接口（第一阶段：SessionStore）。 +// Package redisstore 用 Redis 实现高频、易失且可重建的短状态、缓存、计数器与限流。 // // 职责边界见 docs/persistence-layer.md §1：Redis 存「态与计数」，丢失可由 PG/协议恢复。 package redisstore diff --git a/internal/store/redisstore/session.go b/internal/store/redisstore/session.go deleted file mode 100644 index 2e5507ff..00000000 --- a/internal/store/redisstore/session.go +++ /dev/null @@ -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 中的序列化形态（不含 ID，ID 即 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 -} diff --git a/internal/store/redisstore/session_integration_test.go b/internal/store/redisstore/session_integration_test.go deleted file mode 100644 index aff049e5..00000000 --- a/internal/store/redisstore/session_integration_test.go +++ /dev/null @@ -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") - } -} diff --git a/internal/store/session.go b/internal/store/session.go deleted file mode 100644 index e42401b9..00000000 --- a/internal/store/session.go +++ /dev/null @@ -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 -}