From 305e8a0008e9541755f9d34ed095fe918370886c Mon Sep 17 00:00:00 2001 From: A Date: Mon, 13 Jul 2026 19:05:41 +0800 Subject: [PATCH] fix: sync MTProto startup and egress fixes --- cmd/telesrv/main.go | 27 +- internal/mtprotoedge/conn.go | 16 +- internal/mtprotoedge/connection_intake.go | 95 ++++ .../mtprotoedge/connection_intake_test.go | 170 +++++++ internal/mtprotoedge/encrypted.go | 292 ++++++----- internal/mtprotoedge/inbound_preflight.go | 179 +++++-- internal/mtprotoedge/metrics.go | 18 + internal/mtprotoedge/outbound.go | 422 +++++++++++++--- internal/mtprotoedge/outbound_test.go | 14 +- internal/mtprotoedge/rpc_flight.go | 158 +++++- internal/mtprotoedge/rpc_result_cache.go | 7 +- internal/mtprotoedge/rpc_result_egress.go | 110 +++++ .../mtprotoedge/rpc_result_egress_test.go | 444 +++++++++++++++++ internal/mtprotoedge/rpc_rewrap.go | 455 ++++++++++++++++++ internal/mtprotoedge/rpc_rewrap_test.go | 299 ++++++++++++ internal/mtprotoedge/same_port_mux.go | 27 +- internal/mtprotoedge/server.go | 129 ++++- .../session_barrier_integration_test.go | 58 ++- internal/mtprotoedge/transport_compat.go | 76 ++- internal/mtprotoedge/transport_compat_test.go | 71 ++- internal/postresponse/callbacks.go | 25 +- internal/postresponse/callbacks_test.go | 28 ++ internal/rpc/router.go | 72 ++- internal/rpc/router_dispatch_test.go | 5 +- 24 files changed, 2880 insertions(+), 317 deletions(-) create mode 100644 internal/mtprotoedge/connection_intake.go create mode 100644 internal/mtprotoedge/connection_intake_test.go create mode 100644 internal/mtprotoedge/rpc_result_egress.go create mode 100644 internal/mtprotoedge/rpc_result_egress_test.go create mode 100644 internal/mtprotoedge/rpc_rewrap.go create mode 100644 internal/mtprotoedge/rpc_rewrap_test.go create mode 100644 internal/postresponse/callbacks_test.go diff --git a/cmd/telesrv/main.go b/cmd/telesrv/main.go index 0bb03121..a0fe0ee9 100644 --- a/cmd/telesrv/main.go +++ b/cmd/telesrv/main.go @@ -347,11 +347,6 @@ func run(logger *zap.Logger) error { defer func() { _ = rdb.Close() }() logger.Info("持久化依赖就绪", zap.String("redis", cfg.RedisAddr)) - ln, err := net.Listen("tcp", cfg.ListenAddr) - if err != nil { - return fmt.Errorf("listen %q: %w", cfg.ListenAddr, err) - } - authKeyStore := postgres.NewAuthKeyStore(pool) userStore := postgres.NewUserStore(pool) authzStore := postgres.NewAuthorizationStore(pool) @@ -858,14 +853,18 @@ func run(logger *zap.Logger) error { OutboundControlQueueSize: cfg.MTProtoOutboundControlQueueSize, OutboundTrackedGlobalMaxBytes: cfg.MTProtoOutboundTrackedGlobalMaxBytes, OutboundWriteGlobalMaxBytes: cfg.MTProtoOutboundWriteGlobalMaxBytes, + OnServing: func(_ net.Addr) { + logger.Info("telesrv 服务就绪", + zap.String("listen", cfg.ListenAddr), + zap.String("advertise", net.JoinHostPort(cfg.AdvertiseIP, portStr)), + zap.Int("pid", os.Getpid()), + zap.String("git_commit", buildMeta.Commit), + zap.Uint("schema_version", migrationStatus.Version), + zap.String("blob_backend", "localfs"), + ) + }, }) - logger.Info("telesrv 服务就绪", - zap.String("listen", cfg.ListenAddr), - zap.String("advertise", net.JoinHostPort(cfg.AdvertiseIP, portStr)), - zap.Int("pid", os.Getpid()), - zap.String("git_commit", buildMeta.Commit), - zap.Uint("schema_version", migrationStatus.Version), - zap.String("blob_backend", "localfs"), - ) - return srv.Serve(ctx, ln) + // This is intentionally the final startup operation. ListenAndServe owns the + // public listener so no seed/prewarm work can run after port 2398 is exposed. + return srv.ListenAndServe(ctx, cfg.ListenAddr) } diff --git a/internal/mtprotoedge/conn.go b/internal/mtprotoedge/conn.go index d1d65979..cc9eddca 100644 --- a/internal/mtprotoedge/conn.go +++ b/internal/mtprotoedge/conn.go @@ -61,9 +61,13 @@ type Conn struct { outbound chan outboundOp outboundControl chan outboundOp - outboundStop chan struct{} - outboundDone chan struct{} - outboundClose sync.Once + // Critical RPC results (session/difference convergence) and large bulk + // responses have independent bounded lanes. The actor remains the sole writer. + outboundCritical chan outboundOp + outboundBulk chan outboundOp + outboundStop chan struct{} + outboundDone chan struct{} + outboundClose sync.Once // outboundEnqueueMu orders producer registration against terminal close. Close // flips closing under this lock before waiting, so no WaitGroup Add can race Wait. outboundEnqueueMu sync.Mutex @@ -103,6 +107,12 @@ type Conn struct { rpcRunning int rpcReady bool rpcClosed bool + // Rewrap aliasing never delays execution. initialized stops collecting + // candidates after the first valid init wrapper on this physical generation. + rpcRewrapInitialized atomic.Bool + // rpcResultAcked is invoked by the sole outbound actor after it resolves an + // acknowledged server frame back to the rpc_result request msg_id. + rpcResultAcked func(*Conn, int64) // inflightRPCBytes 跟踪已入队未完成的 inbound RPC body 总字节，配合 maxInflightRPCBytes // 给 RPC 队列设字节预算（不止限条数），防对抗客户端发大请求撑内存。 inflightRPCBytes atomic.Int64 diff --git a/internal/mtprotoedge/connection_intake.go b/internal/mtprotoedge/connection_intake.go new file mode 100644 index 00000000..e77add6e --- /dev/null +++ b/internal/mtprotoedge/connection_intake.go @@ -0,0 +1,95 @@ +package mtprotoedge + +import ( + "net" + "time" + + "go.uber.org/zap" +) + +type connectionIntakeEvent struct { + stage string + outcome string + transport string + remote string + local string + duration time.Duration + bytes int + err error +} + +func (s *Server) signalServing(addr net.Addr) { + if s.onServing != nil { + s.onServing(addr) + } +} + +func (s *Server) observeConnectionIntake(event connectionIntakeEvent) { + fields := []zap.Field{ + zap.String("phase", event.stage), + zap.String("outcome", event.outcome), + } + if event.transport != "" { + fields = append(fields, zap.String("transport", event.transport)) + } + if event.remote != "" { + fields = append(fields, zap.String("remote_addr", event.remote)) + } + if event.local != "" { + fields = append(fields, zap.String("local_addr", event.local)) + } + if event.duration > 0 { + fields = append(fields, zap.Duration("duration", event.duration)) + } + if event.bytes > 0 { + fields = append(fields, zap.Int("bytes", event.bytes)) + } + if event.err != nil { + fields = append(fields, zap.Error(event.err)) + } + s.log.Debug("Connection intake", fields...) + if metrics, ok := s.metrics.(ConnectionIntakeMetrics); ok { + metrics.ConnectionIntake(event.stage, event.outcome, event.duration) + } +} + +func (s *Server) observeRawAccepts(ln net.Listener) net.Listener { + return &connectionObservedListener{Listener: ln, observe: s.observeConnectionIntake} +} + +type connectionObservedListener struct { + net.Listener + observe func(connectionIntakeEvent) +} + +func (l *connectionObservedListener) Accept() (net.Conn, error) { + conn, err := l.Listener.Accept() + if err != nil { + return nil, err + } + l.observe(connectionIntakeEvent{ + stage: "raw_accept", outcome: "accepted", remote: connRemote(conn), local: connLocal(conn), + }) + return conn, nil +} + +func connRemote(conn net.Conn) string { + if conn == nil || conn.RemoteAddr() == nil { + return "" + } + return conn.RemoteAddr().String() +} + +func connLocal(conn net.Conn) string { + if conn == nil || conn.LocalAddr() == nil { + return "" + } + return conn.LocalAddr().String() +} + +func intakeTransport(obfuscated bool) string { + if obfuscated { + return "obfuscated_tcp" + } + return "tcp" +} diff --git a/internal/mtprotoedge/connection_intake_test.go b/internal/mtprotoedge/connection_intake_test.go new file mode 100644 index 00000000..919c14ea --- /dev/null +++ b/internal/mtprotoedge/connection_intake_test.go @@ -0,0 +1,170 @@ +package mtprotoedge + +import ( + "context" + "net" + "strings" + "sync" + "testing" + "time" +) + +func TestListenAndServeCallsServingHookOnBoundListener(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + metrics := newIntakeCaptureMetrics() + hookEntered := make(chan net.Addr, 1) + releaseHook := make(chan struct{}) + hookReleased := false + defer func() { + if !hookReleased { + close(releaseHook) + } + }() + srv := New(Options{Metrics: metrics, OnServing: func(addr net.Addr) { + hookEntered <- addr + <-releaseHook + }}) + done := make(chan error, 1) + go func() { done <- srv.ListenAndServe(ctx, "127.0.0.1:0") }() + + var addr net.Addr + select { + case addr = <-hookEntered: + case <-time.After(3 * time.Second): + t.Fatal("serving callback was not published") + } + conn, err := net.DialTimeout("tcp", addr.String(), time.Second) + if err != nil { + t.Fatalf("dial after serving callback: %v", err) + } + _ = conn.Close() + + // The observation hook is still blocked: raw_accept can only advance here + // when the intake loop was installed before OnServing was called. + deadline := time.After(3 * time.Second) + for !hasIntakeEvent(metrics.snapshot(), "raw_accept", "accepted") { + select { + case <-metrics.wake: + case <-deadline: + t.Fatal("raw accept loop did not run while serving hook was blocked") + } + } + close(releaseHook) + hookReleased = true + + cancel() + select { + case err := <-done: + if err != nil { + t.Fatalf("ListenAndServe shutdown: %v", err) + } + case <-time.After(3 * time.Second): + t.Fatal("ListenAndServe did not stop") + } +} + +type intakeCaptureMetrics struct { + NopMetrics + + mu sync.Mutex + events []intakeMetricEvent + wake chan struct{} +} + +type intakeMetricEvent struct { + stage string + outcome string +} + +func newIntakeCaptureMetrics() *intakeCaptureMetrics { + return &intakeCaptureMetrics{wake: make(chan struct{}, 16)} +} + +func (m *intakeCaptureMetrics) ConnectionIntake(stage, outcome string, _ time.Duration) { + m.mu.Lock() + m.events = append(m.events, intakeMetricEvent{stage: stage, outcome: outcome}) + m.mu.Unlock() + select { + case m.wake <- struct{}{}: + default: + } +} + +func (m *intakeCaptureMetrics) snapshot() []intakeMetricEvent { + m.mu.Lock() + defer m.mu.Unlock() + return append([]intakeMetricEvent(nil), m.events...) +} + +func TestConnectionIntakeStagesIncludePrePromotionDisconnect(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + metrics := newIntakeCaptureMetrics() + ready := make(chan net.Addr, 1) + srv := New(Options{ + Metrics: metrics, + ObfuscatedTCP: true, + WebSocket: true, + OnServing: func(addr net.Addr) { + ready <- addr + }, + }) + done := make(chan error, 1) + go func() { done <- srv.ListenAndServe(ctx, "127.0.0.1:0") }() + + var addr net.Addr + select { + case addr = <-ready: + case <-time.After(3 * time.Second): + t.Fatal("serving callback was not published") + } + conn, err := net.DialTimeout("tcp", addr.String(), time.Second) + if err != nil { + t.Fatalf("dial: %v", err) + } + // Four non-HTTP bytes let the mux classify the connection as raw MTProto; + // closing before the remaining obfuscated2 header arrives must still expose + // the exact transport_promote failure phase. + if _, err := conn.Write([]byte{1, 2, 3, 4}); err != nil { + t.Fatalf("write prefix: %v", err) + } + _ = conn.Close() + + deadline := time.After(3 * time.Second) + for { + events := metrics.snapshot() + if hasIntakeEvent(events, "raw_accept", "accepted") && + hasIntakeEvent(events, "mux_sniff", "ready") && + hasIntakeEvent(events, "transport_dispatch", "accepted") && + hasIntakeEvent(events, "transport_promote", "client_disconnect") { + break + } + select { + case <-metrics.wake: + case <-deadline: + t.Fatalf("intake events did not converge: %v", events) + } + } + + cancel() + select { + case err := <-done: + if err != nil && !strings.Contains(err.Error(), "closed") { + t.Fatalf("server shutdown: %v", err) + } + case <-time.After(3 * time.Second): + t.Fatal("server did not stop") + } +} + +func hasIntakeEvent(events []intakeMetricEvent, stage, outcome string) bool { + for _, event := range events { + if event.stage == stage && event.outcome == outcome { + return true + } + } + return false +} diff --git a/internal/mtprotoedge/encrypted.go b/internal/mtprotoedge/encrypted.go index a31c6589..c3e05fe5 100644 --- a/internal/mtprotoedge/encrypted.go +++ b/internal/mtprotoedge/encrypted.go @@ -260,6 +260,9 @@ func (s *Server) handleEncrypted(ctx context.Context, tc transport.Conn, cs *con if err := s.executeInboundPlan(ctx, cs, current, plan); err != nil { return current, err } + if err := plan.commitRewrapAliases(s); err != nil { + return current, err + } if err := plan.commitRPCBatch(); err != nil { return current, err } @@ -268,12 +271,6 @@ func (s *Server) handleEncrypted(ctx context.Context, tc transport.Conn, cs *con return current, err } } - // An overlapping cross-connection owner may run until RPCTimeout. ACK the - // accepted duplicate before joining it so the new client does not build a - // retransmit storm while its read loop waits for the shared result. - if err := s.executePendingRPCReplays(ctx, current, plan); err != nil { - return current, err - } return current, nil } @@ -577,58 +574,6 @@ func mergeStateInfo(primary, fallback []byte) []byte { return info } -// 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) - if err != nil { - if errors.Is(err, ErrRPCResultFlightCapacity) { - c.metrics.InboundRPCDropped(method, "flight_capacity") - return s.handleInboundRPCAdmissionError(ctx, c, msgID, method, ErrInboundRPCQueueFull) - } - return err - } - switch claim.state { - case rpcResultAcquireCompleted: - s.log.Info("RPC duplicate replay from session cache", - zap.String("method", method), - zap.Int64("msg_id", msgID), - zap.String("auth_key_id", c.authKeyHex), - zap.Int64("session_id", c.sessionID), - ) - return s.sendCachedRPCResult(ctx, c, claim.encoded) - case rpcResultAcquirePending: - encoded, ok, waitErr := claim.waiter.Wait(ctx) - if waitErr != nil || !ok || encoded == nil { - return waitErr - } - return s.sendCachedRPCResult(ctx, c, encoded) - case rpcResultAcquireOwner: - // Ownership transfers to the queued task only after commit succeeds. - default: - return ErrRPCResultFlightInvalid - } - owner := claim.owner - transferred := false - defer func() { - if !transferred { - owner.Abort() - } - }() - // 两级条数/字节预算必须先于 Copy：对抗客户端不能用大量满尺寸请求在“判断队列满” - // 之前制造一轮无上限的临时 body 分配。reservation 在 commit/abort 间唯一持有预算。 - 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([]inboundRPC{s.newInboundRPCTask(c, msgID, method, body, owner)}) - transferred = err == nil - return s.handleInboundRPCAdmissionError(ctx, c, msgID, method, err) -} - // newInboundRPCTask builds the exactly-once timeout/result gate shared by the // single-message and atomic container-batch admission paths. body must already // be an independently owned, budgeted copy. @@ -675,7 +620,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}); err != nil { + if err := s.handleRPC(taskCtx, c, msgID, method, &bin.Buffer{Buf: body}, owner); err != nil { fields := []zap.Field{ zap.Int64("msg_id", msgID), zap.String("auth_key_id", c.authKeyHex), @@ -711,24 +656,37 @@ 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) error { +func (s *Server) handleRPC(ctx context.Context, c *Conn, msgID int64, method string, b *bin.Buffer, owner *rpcResultOwnerLease) error { if s.rpc == nil { if err := ctx.Err(); err != nil { return err } s.log.Warn("No RPC handler configured", zap.String("method", method)) - return s.sendResult(ctx, c, msgID, &mt.RPCError{ + return s.publishRPCResult(c, msgID, method, owner, &mt.RPCError{ ErrorCode: 500, ErrorMessage: "NOT_IMPLEMENTED", - }) + }, nil) } ctx = postresponse.WithCallbacks(ctx) ctx, dbStats := dbtrace.WithStats(ctx) start := s.clock.Now() - result, err := s.rpc.Dispatch(ctx, c.authKeyID, c.sessionID, b) + effectiveMethod := method + var ( + result bin.Encoder + err error + ) + if detailed, ok := s.rpc.(RPCHandlerWithMethod); ok { + var innerMethod string + result, innerMethod, err = detailed.DispatchWithMethod(ctx, c.authKeyID, c.sessionID, b) + if innerMethod != "" { + effectiveMethod = innerMethod + } + } else { + result, err = s.rpc.Dispatch(ctx, c.authKeyID, c.sessionID, b) + } dur := s.clock.Now().Sub(start) - s.metrics.RPCHandled(method, dur, err) + s.metrics.RPCHandled(effectiveMethod, dur, err) // 刷新本连接协商 layer（invokeWithLayer/initConnection 已被 Dispatch 处理并登记）， // 供 rpc_result 与后续 push 出站降级使用。仅在确实观测到 layer 时更新——缓存被驱逐 // 时 NegotiatedLayer 返回 ok=false，此时必须保留连接已记住的 layer，绝不覆盖成默认值， @@ -739,12 +697,15 @@ func (s *Server) handleRPC(ctx context.Context, c *Conn, msgID int64, method str fields := make([]zap.Field, 0, 12) fields = append(fields, - zap.String("method", method), + zap.String("method", effectiveMethod), zap.String("auth_key_id", c.authKeyHex), zap.Int64("session_id", c.sessionID), zap.Int64("msg_id", msgID), zap.Duration("dur", dur), ) + if effectiveMethod != method { + fields = append(fields, zap.String("outer_method", method)) + } if businessAuthKeyHex, ok := c.BusinessAuthKeyHex(); ok { fields = append(fields, zap.String("business_auth_key_id", businessAuthKeyHex)) } @@ -767,33 +728,12 @@ func (s *Server) handleRPC(ctx context.Context, c *Conn, msgID int64, method str terminal = &mt.RPCError{ErrorCode: 500, ErrorMessage: "RPC_TIMEOUT"} } 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) - if runPostResponse { - postresponse.Run(context.WithoutCancel(ctx)) - } - } - } else { - // An individual RPC deadline can expire while the physical connection is - // still healthy. Use a fresh bounded delivery context; publishing cache-only - // here would strand same-Conn duplicates behind an ACK with no result. - writeTimeout := c.writeTimeout - if writeTimeout <= 0 || writeTimeout > 5*time.Second { - writeTimeout = 5 * time.Second - } - responseCtx, cancel := context.WithTimeout(context.Background(), writeTimeout) - sendErr := s.sendResult(responseCtx, c, msgID, terminal) - cancel() - if sendErr != nil { - s.log.Debug("Send canceled RPC result failed", append(fields, zap.Error(sendErr))...) - } else if runPostResponse { - postresponse.Run(context.WithoutCancel(ctx)) - } + var after func() + if runPostResponse { + after = postresponse.Take(context.WithoutCancel(ctx)) + } + if sendErr := s.publishRPCResult(c, msgID, effectiveMethod, owner, terminal, after); sendErr != nil { + s.log.Debug("Publish canceled RPC result failed", append(fields, zap.Error(sendErr))...) } } cancelFields := append(fields, zap.NamedError("context_error", ctxErr)) @@ -808,23 +748,126 @@ func (s *Server) handleRPC(ctx context.Context, c *Conn, msgID int64, method str var rpcErr *tgerr.Error if errors.As(err, &rpcErr) { s.log.Info("RPC error", append(fields, zap.Int("code", rpcErr.Code), zap.String("error", rpcErr.Message))...) - return s.sendResult(ctx, c, msgID, &mt.RPCError{ + return s.publishRPCResult(c, msgID, effectiveMethod, owner, &mt.RPCError{ ErrorCode: rpcErr.Code, ErrorMessage: rpcErr.Message, - }) + }, nil) } s.log.Info("RPC internal error", append(fields, zap.Error(err))...) - return s.sendResult(ctx, c, msgID, &mt.RPCError{ + return s.publishRPCResult(c, msgID, effectiveMethod, owner, &mt.RPCError{ ErrorCode: 500, ErrorMessage: "INTERNAL", - }) + }, nil) } s.log.Info("RPC handled", fields...) - if err := s.sendResult(ctx, c, msgID, result); err != nil { + return s.publishRPCResult(c, msgID, effectiveMethod, owner, result, postresponse.Take(ctx)) +} + +// publishRPCResult ends the inbound worker's ownership at bounded egress +// admission. Physical delivery, fencing, completed-cache publication and the +// post-response hook are thereafter owned by the single outbound actor. +func (s *Server) publishRPCResult( + c *Conn, + reqMsgID int64, + method string, + owner *rpcResultOwnerLease, + result bin.Encoder, + afterDelivered func(), +) error { + if result == nil { + result = &mt.RPCError{ErrorCode: 500, ErrorMessage: "INTERNAL"} + } + prepareTimeout := c.writeTimeout + if prepareTimeout <= 0 || prepareTimeout > 5*time.Second { + prepareTimeout = 5 * time.Second + } + prepareCtx, cancel := context.WithTimeout(context.Background(), prepareTimeout) + defer cancel() + encoded, err := s.encodeRPCResultContext(prepareCtx, c, reqMsgID, result) + if err != nil { + s.log.Warn("Encode RPC result failed; publishing INTERNAL", + zap.String("method", method), zap.Int64("req_msg_id", reqMsgID), zap.Error(err)) + afterDelivered = nil + encoded, err = s.encodeRPCResultContext(prepareCtx, c, reqMsgID, &mt.RPCError{ + ErrorCode: 500, ErrorMessage: "INTERNAL", + }) + if err != nil { + c.fenceUndeliveredRPCResult() + return err + } + } + if owner != nil && owner.Delivery() != nil { + // The owner-level delivery coordinator exists before the handler starts, so + // an initConnection rewrap can retarget even while result encoding is still + // pending. The encoded body itself remains immutable; the actor clones only + // the 12-byte rpc_result prefix when it snapshots the physical target. + encoded.delivery = owner.Delivery() + } + if afterDelivered != nil { + encoded.delivery.fn = afterDelivered + } + if owner != nil && !owner.HandOff() { + return ErrRPCResultFlightInvalid + } + + priority := rpcResultPriority(method, encoded) + encoded.priority = priority + resultLogLevel := zap.DebugLevel + if encoded.compressed || priority == outboundPriorityCritical || priority == outboundPriorityBulk { + // Keep ordinary small RPCs at debug, but make convergence and bulk/gzip + // delivery visible in default service logs. These are the + // responses whose queueing and write latency diagnose startup Updating. + resultLogLevel = zap.InfoLevel + } + if metrics, ok := s.metrics.(RPCResultMetrics); ok { + metrics.RPCResultPrepared(method, priority.String(), encoded.uncompressedBytes, len(encoded.body), encoded.compressed) + } + egressStarted := time.Now() + terminal := func(deliveryErr error) { + latency := time.Since(egressStarted) + deliveredReqMsgID := encoded.writtenRequestID() + if metrics, ok := s.metrics.(RPCResultMetrics); ok { + metrics.RPCResultDelivered(method, latency, len(encoded.body), deliveryErr) + } + if deliveryErr != nil { + encoded.markReplayable() + c.fenceUndeliveredRPCResult() + s.storeRPCResult(c, reqMsgID, encoded) + if checked := s.log.Check(resultLogLevel, "RPC result delivery fenced for replay"); checked != nil { + checked.Write( + zap.String("method", method), zap.Int64("req_msg_id", reqMsgID), + zap.Int64("delivered_req_msg_id", deliveredReqMsgID), + zap.String("auth_key_id", c.authKeyHex), zap.Int64("session_id", c.sessionID), + zap.Int("wire_bytes", len(encoded.body)), zap.Bool("gzip", encoded.compressed), + zap.Error(deliveryErr)) + } + return + } + s.storeRPCResult(c, reqMsgID, encoded) + encoded.markDelivered() + if checked := s.log.Check(resultLogLevel, "RPC result delivered"); checked != nil { + checked.Write( + zap.String("method", method), zap.Int64("req_msg_id", reqMsgID), + zap.Int64("delivered_req_msg_id", deliveredReqMsgID), + zap.String("auth_key_id", c.authKeyHex), zap.Int64("session_id", c.sessionID), + zap.Int("wire_bytes", len(encoded.body)), zap.Bool("gzip", encoded.compressed), + zap.Duration("egress_latency", latency)) + } + } + encoded.markQueued() + if err := c.enqueueEncodedDelivery(prepareCtx, proto.MessageServerResponse, encoded, priority, terminal); err != nil { + // HandOff already made the egress path the terminal owner. No bytes were + // admitted, so fence this generation before publishing a replayable result. + terminal(err) return err } - postresponse.Run(ctx) + if checked := s.log.Check(resultLogLevel, "RPC result admitted"); checked != nil { + checked.Write( + zap.String("method", method), zap.Int64("req_msg_id", reqMsgID), + zap.Int("wire_bytes", len(encoded.body)), zap.Int("inner_bytes", encoded.uncompressedBytes), + zap.Bool("gzip", encoded.compressed), zap.String("priority", priority.String())) + } return nil } @@ -833,13 +876,13 @@ func (s *Server) sendResult(ctx context.Context, c *Conn, reqMsgID int64, result if result == nil { result = &mt.RPCError{ErrorCode: 500, ErrorMessage: "INTERNAL"} } - encoded, err := s.encodeRPCResult(c, reqMsgID, result) + encoded, err := s.encodeRPCResultContext(ctx, c, reqMsgID, result) if err != nil { // The business operation has already crossed atomic admission. Convert an // invalid result encoder into one deterministic terminal RPC error instead of // aborting the flight and allowing a reconnect to execute the operation again. s.log.Warn("Encode RPC result failed; sending INTERNAL", zap.Int64("req_msg_id", reqMsgID), zap.Error(err)) - encoded, err = s.encodeRPCResult(c, reqMsgID, &mt.RPCError{ + encoded, err = s.encodeRPCResultContext(ctx, c, reqMsgID, &mt.RPCError{ ErrorCode: 500, ErrorMessage: "INTERNAL", }) @@ -854,12 +897,14 @@ func (s *Server) sendResult(ctx context.Context, c *Conn, reqMsgID int64, result // (queue/context/scratch deadline); without this terminal barrier a later // same-Conn duplicate would be ACKed while no result can ever arrive. c.fenceUndeliveredRPCResult() + encoded.markReplayable() s.storeRPCResult(c, reqMsgID, encoded) return err } // On a live Conn, completed means the rpc_result has reached the reliable byte // stream. Same-physical duplicates can therefore be ACK-only without data loss. s.storeRPCResult(c, reqMsgID, encoded) + encoded.markDelivered() return nil } @@ -873,8 +918,10 @@ func (s *Server) sendCachedRPCResult(ctx context.Context, c *Conn, encoded *enco } if err := c.SendEncoded(ctx, proto.MessageServerResponse, encoded); err != nil { c.fenceUndeliveredRPCResult() + encoded.markReplayable() return err } + encoded.markDelivered() return nil } @@ -884,30 +931,43 @@ func (s *Server) sendCachedRPCResult(ctx context.Context, c *Conn, encoded *enco // 降级改写字节时才重建整条消息。降级失败 fail-safe：记日志并发送 canonical 字节—— // 宁可老客户端对个别长尾对象渲染异常，也不让连接/流崩。 func (s *Server) encodeRPCResult(c *Conn, reqMsgID int64, result bin.Encoder) (*encodedOutboundMessage, error) { - const headerLen = 4 + 8 // rpc_result#f35c6d01 type_id + req_msg_id - var buf bin.Buffer - buf.PutID(proto.ResultTypeID) - buf.PutLong(reqMsgID) - if err := result.Encode(&buf); err != nil { + return s.encodeRPCResultContext(context.Background(), c, reqMsgID, result) +} + +func (s *Server) encodeRPCResultContext(ctx context.Context, c *Conn, reqMsgID int64, result bin.Encoder) (*encodedOutboundMessage, error) { + if ctx == nil { + ctx = context.Background() + } + var inner bin.Buffer + // Terminal result preparation must survive physical-generation retirement: + // an overlapping replacement may already be waiting to replay this owner's + // result. Only the bounded preparation context, not the old socket stop, owns it. + if err := withOutboundEncodeSlot(ctx, nil, func() error { + return result.Encode(&inner) + }); err != nil { return nil, fmt.Errorf("encode rpc result: %w", err) } + innerBody := inner.Raw() if layer := c.ClientLayer(); layer < layerwire.CanonicalLayer { - inner := buf.Buf[headerLen:] - if down, err := layerwire.Transcode(inner, layer); err != nil { + if down, err := layerwire.Transcode(innerBody, layer); err != nil { s.log.Warn("layerwire downgrade failed; sending canonical rpc_result", zap.Int("layer", layer), zap.Int64("req_msg_id", reqMsgID), zap.Error(err)) - } else if !sameBacking(down, inner) { - var rebuilt bin.Buffer - rebuilt.PutID(proto.ResultTypeID) - rebuilt.PutLong(reqMsgID) - rebuilt.Put(down) - buf = rebuilt + } else { + innerBody = down } } + + wireInner, compressed, err := encodeAdaptiveRPCResultInner(ctx, nil, innerBody) + if err != nil { + return nil, fmt.Errorf("compress rpc result: %w", err) + } + var out bin.Buffer + out.PutID(proto.ResultTypeID) + out.PutLong(reqMsgID) + out.Put(wireInner) return &encodedOutboundMessage{ - typeID: proto.ResultTypeID, - body: buf.Raw(), - reqMsgID: reqMsgID, + typeID: proto.ResultTypeID, body: out.Raw(), reqMsgID: reqMsgID, + compressed: compressed, uncompressedBytes: len(innerBody), delivery: newRPCResultDelivery(), }, nil } diff --git a/internal/mtprotoedge/inbound_preflight.go b/internal/mtprotoedge/inbound_preflight.go index d38995ec..66188a0d 100644 --- a/internal/mtprotoedge/inbound_preflight.go +++ b/internal/mtprotoedge/inbound_preflight.go @@ -34,7 +34,10 @@ const ( inboundItemDestroyAuthKey inboundItemRPC inboundItemCapacityError - inboundItemPendingRPC + // inboundItemRewrappedRPC is an initConnection retry whose exact inner TL + // request is already executing (or completed) under the client's old msg_id. + // It never dispatches business code a second time. + inboundItemRewrappedRPC // inboundItemReplayRPC is a request first observed by this physical Conn whose // terminal result already exists in the cross-connection cache. It is distinct // from inboundItemDuplicate: a duplicate already present in this Conn's seen @@ -71,6 +74,8 @@ type inboundPlan struct { rpcReservation *inboundRPCBatchReservation rpcTasks []inboundRPC rpcOwners []*rpcResultOwnerLease + rewrapAliases []*rpcRewrapAlias + rewrapIndices []int } func (p *inboundPlan) close() { @@ -85,12 +90,42 @@ func (p *inboundPlan) close() { owner.Abort() } p.rpcOwners = nil + for _, alias := range p.rewrapAliases { + alias.releaseCandidate() + if alias != nil && alias.newOwner != nil { + alias.newOwner.Abort() + } + } + p.rewrapAliases = nil + p.rewrapIndices = nil for i := len(p.releases) - 1; i >= 0; i-- { p.releases[i]() } p.releases = nil } +func (p *inboundPlan) commitRewrapAliases(s *Server) error { + if p == nil || len(p.rewrapAliases) == 0 { + return nil + } + for i, alias := range p.rewrapAliases { + if err := alias.activate(s); err != nil { + for _, pending := range p.rewrapAliases[i:] { + pending.releaseCandidate() + if pending != nil && pending.newOwner != nil { + pending.newOwner.Abort() + } + } + p.rewrapAliases = nil + p.rewrapIndices = nil + return err + } + } + p.rewrapAliases = nil + p.rewrapIndices = nil + return nil +} + func (p *inboundPlan) commitRPCBatch() error { if p == nil || p.rpcReservation == nil { return nil @@ -624,6 +659,7 @@ func (s *Server) prepareInboundRPCBatch(ctx context.Context, c *Conn, plan *inbo var specs []inboundRPCSpec var ownersInPlan map[int64]*rpcResultOwnerLease flightCapacity := false + clearedPostInitCandidates := false for i := range plan.items { item := &plan.items[i] if item.kind != inboundItemRPC && item.kind != inboundItemDuplicate { @@ -638,6 +674,71 @@ func (s *Server) prepareInboundRPCBatch(ctx context.Context, c *Conn, plan *inbo continue } method := s.typeName(item.typeID) + init, isInitRewrap := decodeRPCRewrapInit(item.body) + if isInitRewrap { + firstInit := !c.rpcRewrapInitialized.Swap(true) + c.SetClientLayer(init.layer) + if candidate := s.rpcRewrap.claim(c, init.inner); candidate != nil { + claim, err := s.rpcResults.Acquire(c.authKeyID, c.sessionID, item.msgID) + if errors.Is(err, ErrRPCResultFlightCapacity) { + s.rpcRewrap.release(candidate) + c.metrics.InboundRPCDropped(candidate.method, "flight_capacity") + flightCapacity = true + item.kind = inboundItemCapacityError + continue + } + if err != nil { + s.rpcRewrap.release(candidate) + return err + } + switch claim.state { + case rpcResultAcquireCompleted: + s.rpcRewrap.commit(candidate) + item.kind = inboundItemReplayRPC + item.payload = claim.encoded + case rpcResultAcquirePending: + s.rpcRewrap.commit(candidate) + item.kind = inboundItemRewrappedRPC + item.payload = claim.waiter + plan.rewrapAliases = append(plan.rewrapAliases, &rpcRewrapAlias{ + conn: c, newReqID: item.msgID, method: candidate.method, + oldWaiter: claim.waiter, observeInit: firstInit, init: init, + }) + plan.rewrapIndices = append(plan.rewrapIndices, i) + case rpcResultAcquireOwner: + if ownersInPlan == nil { + ownersInPlan = make(map[int64]*rpcResultOwnerLease) + } + ownersInPlan[item.msgID] = claim.owner + item.kind = inboundItemRewrappedRPC + item.payload = claim.owner + plan.rewrapAliases = append(plan.rewrapAliases, &rpcRewrapAlias{ + conn: c, newReqID: item.msgID, method: candidate.method, + oldWaiter: candidate.waiter, newOwner: claim.owner, + sourceConn: candidate.source, sourceOwner: candidate.owner, + observeInit: firstInit, init: init, + candidate: candidate, registry: s.rpcRewrap, + }) + plan.rewrapIndices = append(plan.rewrapIndices, i) + default: + s.rpcRewrap.release(candidate) + return ErrRPCResultFlightInvalid + } + s.log.Info("RPC init rewrap matched", + zap.String("method", candidate.method), + zap.Int64("old_req_msg_id", candidate.reqMsgID), + zap.Int64("new_req_msg_id", item.msgID), + zap.Bool("same_connection", candidate.source == c), + zap.String("auth_key_id", c.authKeyHex), zap.Int64("session_id", c.sessionID)) + continue + } + } else if c.rpcRewrapInitialized.Load() && !clearedPostInitCandidates { + // A naked request after this connection has observed initConnection is + // event-level proof that the client finished moving its old running set. + // Retire any unmatched candidates without a timer. + s.rpcRewrap.clearSession(c) + clearedPostInitCandidates = true + } claim, err := s.rpcResults.Acquire(c.authKeyID, c.sessionID, item.msgID) if errors.Is(err, ErrRPCResultFlightCapacity) { if item.kind == inboundItemRPC { @@ -669,8 +770,11 @@ func (s *Server) prepareInboundRPCBatch(ctx context.Context, c *Conn, plan *inbo item.kind = inboundItemDuplicate item.payload = nil } else { - item.kind = inboundItemPendingRPC + item.kind = inboundItemRewrappedRPC item.payload = claim.waiter + plan.rewrapAliases = append(plan.rewrapAliases, &rpcRewrapAlias{ + conn: c, newReqID: item.msgID, method: method, oldWaiter: claim.waiter, + }) } case rpcResultAcquireOwner: if item.kind == inboundItemDuplicate { @@ -689,6 +793,9 @@ func (s *Server) prepareInboundRPCBatch(ctx context.Context, c *Conn, plan *inbo item.payload = claim.owner indices = append(indices, i) specs = append(specs, inboundRPCSpec{method: method, size: len(item.body)}) + if !c.rpcRewrapInitialized.Load() && !isInitRewrap { + s.rpcRewrap.register(c, item.body, item.msgID, method, claim.owner) + } default: return ErrRPCResultFlightInvalid } @@ -700,6 +807,17 @@ func (s *Server) prepareInboundRPCBatch(ctx context.Context, c *Conn, plan *inbo for _, index := range indices { plan.items[index].kind = inboundItemCapacityError } + for _, index := range plan.rewrapIndices { + plan.items[index].kind = inboundItemCapacityError + } + for _, alias := range plan.rewrapAliases { + alias.releaseCandidate() + if alias.newOwner != nil { + alias.newOwner.Abort() + } + } + plan.rewrapAliases = nil + plan.rewrapIndices = nil return nil } if len(specs) == 0 { @@ -712,6 +830,17 @@ func (s *Server) prepareInboundRPCBatch(ctx context.Context, c *Conn, plan *inbo for _, index := range indices { plan.items[index].kind = inboundItemCapacityError } + for _, index := range plan.rewrapIndices { + plan.items[index].kind = inboundItemCapacityError + } + for _, alias := range plan.rewrapAliases { + alias.releaseCandidate() + if alias.newOwner != nil { + alias.newOwner.Abort() + } + } + plan.rewrapAliases = nil + plan.rewrapIndices = nil return nil } return err @@ -751,10 +880,10 @@ func (s *Server) executeInboundPlan(ctx context.Context, cs *connState, c *Conn, } else if err := s.replayRPCResultByRequest(ctx, c, item.msgID); err != nil { return err } - case inboundItemPendingRPC: - // Wait only after this plan's fresh RPC batch has become runnable; see - // executePendingRPCReplays. Blocking here could otherwise deadlock on - // an owner appended by the same container. + case inboundItemRewrappedRPC: + // Activation is deferred until every session/control barrier has + // committed. It subscribes to the original result event and never waits + // or dispatches the business handler again. continue case inboundItemPing: if err := s.sendPong(ctx, c, item.msgID, item.payload.(mt.PingRequest).PingID); err != nil { @@ -842,41 +971,3 @@ func (s *Server) executeInboundPlan(ctx context.Context, cs *connState, c *Conn, } return nil } - -// executePendingRPCReplays joins owners already running on another physical -// connection for the same MTProto session. It never dispatches business code: -// owner Put publishes the immutable result to every waiter; owner Abort leaves -// the client free to retry after the old execution has definitely stopped. -func (s *Server) executePendingRPCReplays(ctx context.Context, c *Conn, plan *inboundPlan) error { - for _, item := range plan.items { - if item.kind != inboundItemPendingRPC { - continue - } - waiter, _ := item.payload.(*rpcResultWaiter) - if waiter == nil { - continue - } - encoded, ok, err := waiter.Wait(ctx) - if err != nil { - return err - } - if !ok || encoded == nil { - // The old owner stopped without publishing a result (normally because - // replacement cancellation reached the handler before it committed). - // A fresh-connection item still owns the decoded request body, so reacquire - // only after the prior flight is definitively gone. This is sequential - // retry, never concurrent business execution. Same-Conn seen duplicates - // have no body here and rely on the client's ordinary resend path. - if len(item.body) > 0 { - if err := s.enqueueRPC(ctx, c, item.msgID, item.typeID, &bin.Buffer{Buf: item.body}); err != nil { - return err - } - } - continue - } - if err := s.sendCachedRPCResult(ctx, c, encoded); err != nil { - return err - } - } - return nil -} diff --git a/internal/mtprotoedge/metrics.go b/internal/mtprotoedge/metrics.go index 74e1e4f3..622827a3 100644 --- a/internal/mtprotoedge/metrics.go +++ b/internal/mtprotoedge/metrics.go @@ -30,6 +30,24 @@ type Metrics interface { OutboundQueueWait(len, cap int) } +// RPCResultMetrics is an optional extension for the detached response pipeline. +// Keeping it separate preserves lightweight embedders while production exporters +// can observe preparation/compression and end-to-end delivery independently from +// business handler latency. +type RPCResultMetrics interface { + RPCResultPrepared(method, priority string, innerBytes, wireBytes int, compressed bool) + RPCResultDelivered(method string, egressLatency time.Duration, wireBytes int, err error) +} + +// ConnectionIntakeMetrics is an optional extension for the pre-session +// connection pipeline. stage is one of raw_accept, mux_sniff, mux_delivery, +// transport_dispatch, transport_promote, or first_frame; outcome is a bounded +// value suitable for metrics labels. Remote addresses are deliberately kept in +// structured logs only so exporters cannot create unbounded cardinality. +type ConnectionIntakeMetrics interface { + ConnectionIntake(stage, outcome string, duration time.Duration) +} + // NopMetrics 是 Metrics 的空实现。 type NopMetrics struct{} diff --git a/internal/mtprotoedge/outbound.go b/internal/mtprotoedge/outbound.go index 95778775..7f397742 100644 --- a/internal/mtprotoedge/outbound.go +++ b/internal/mtprotoedge/outbound.go @@ -36,10 +36,12 @@ const ( // 原实现为每条 Conn eager 分配 1024 + 256 个 outboundOp 槽；在大连接数下仅空队列 // backing 就占据大量常驻内存。默认缩至 128 + 32，仍覆盖 TDesktop 启动/推送突发， // 慢消费者继续由 best-effort timeout + durable difference 降级。 - defaultOutboundQueueSize = 128 - defaultOutboundControlQueueSize = 32 - defaultOutboundTrackedMaxBytes = int64(512 << 20) // 512 MiB / Server - defaultOutboundControlMaxBytes = int64(64 << 20) // ack/state/resend vectors / Server + defaultOutboundQueueSize = 128 + defaultOutboundControlQueueSize = 32 + defaultOutboundCriticalQueueSize = 16 + defaultOutboundBulkQueueSize = 16 + defaultOutboundTrackedMaxBytes = int64(512 << 20) // 512 MiB / Server + defaultOutboundControlMaxBytes = int64(64 << 20) // ack/state/resend vectors / Server // requiredControlMaxWait bounds protocol barriers such as new_session_created from // the beginning of encoding through the completed physical write. These frames gate // subsequent session state transitions, so timing out must close the connection instead @@ -67,13 +69,31 @@ const ( outboundResendByRequest ) +type outboundPriority uint8 + +const ( + outboundPriorityNormal outboundPriority = iota + outboundPriorityCritical + outboundPriorityBulk + outboundPriorityControl +) + +const ( + // Large responses are scheduled separately so a startup prefetch cannot sit + // ahead of an already-prepared session convergence result. The threshold is + // applied after layer conversion and adaptive gzip. + bulkOutboundThreshold = 64 << 10 + maxOrdinaryBeforeBulk = 16 +) + type outboundOp struct { - kind outboundOpKind - control bool - ctx context.Context - msgType proto.MessageType - msg bin.Encoder - encoded *encodedOutboundMessage + kind outboundOpKind + control bool + priority outboundPriority + ctx context.Context + msgType proto.MessageType + msg bin.Encoder + encoded *encodedOutboundMessage // reservedBytes accounts for the encoded body while it is queued. For a // content frame the reservation is transferred to resend tracking after a // successful write; every other terminal path releases it exactly once. @@ -83,12 +103,160 @@ type outboundOp struct { reqMsgID int64 enqueuedAt time.Time done chan outboundResult + // terminal is owned by the outbound actor after successful queue admission. + // It resolves detached RPC-result ownership on every physical terminal path. + terminal func(error) } type encodedOutboundMessage struct { - body []byte - typeID uint32 - reqMsgID int64 + body []byte + typeID uint32 + reqMsgID int64 + priority outboundPriority + delivery *rpcResultDelivery + compressed bool + uncompressedBytes int +} + +type rpcResultDeliveryState uint32 + +const ( + rpcResultDeliveryPrepared rpcResultDeliveryState = iota + 1 + rpcResultDeliveryQueued + rpcResultDeliveryWriting + rpcResultDeliveryReplayable + rpcResultDeliveryDelivered +) + +type rpcResultDelivery struct { + state atomic.Uint32 + mu sync.Mutex + // targetReqMsgID may change only before the outbound actor enters writing. + // writtenReqMsgID is the actor's immutable snapshot for the physical frame. + targetReqMsgID int64 + writtenReqMsgID int64 + once sync.Once + fn func() +} + +func newRPCResultDelivery(reqMsgID ...int64) *rpcResultDelivery { + d := &rpcResultDelivery{} + if len(reqMsgID) > 0 { + d.targetReqMsgID = reqMsgID[0] + } + d.state.Store(uint32(rpcResultDeliveryPrepared)) + return d +} + +func (m *encodedOutboundMessage) deliveryState() rpcResultDeliveryState { + if m == nil || m.delivery == nil { + return 0 + } + return rpcResultDeliveryState(m.delivery.state.Load()) +} + +func (m *encodedOutboundMessage) markQueued() { + if m == nil || m.delivery == nil { + return + } + m.delivery.mu.Lock() + if m.deliveryState() == rpcResultDeliveryPrepared { + m.delivery.state.Store(uint32(rpcResultDeliveryQueued)) + } + m.delivery.mu.Unlock() +} + +// beginWriting linearizes retargeting against the sole outbound actor. The +// returned req_msg_id is immutable for this physical write. +func (m *encodedOutboundMessage) beginWriting() int64 { + if m == nil || m.delivery == nil { + if m == nil { + return 0 + } + return m.reqMsgID + } + m.delivery.mu.Lock() + if m.delivery.targetReqMsgID == 0 { + m.delivery.targetReqMsgID = m.reqMsgID + } + m.delivery.writtenReqMsgID = m.delivery.targetReqMsgID + m.delivery.state.Store(uint32(rpcResultDeliveryWriting)) + target := m.delivery.writtenReqMsgID + m.delivery.mu.Unlock() + return target +} + +func (m *encodedOutboundMessage) tryRetarget(reqMsgID int64) bool { + if m == nil || m.delivery == nil || reqMsgID == 0 { + return false + } + m.delivery.mu.Lock() + defer m.delivery.mu.Unlock() + state := m.deliveryState() + if state != rpcResultDeliveryPrepared && state != rpcResultDeliveryQueued { + return false + } + m.delivery.targetReqMsgID = reqMsgID + return true +} + +func (m *encodedOutboundMessage) writtenRequestID() int64 { + if m == nil || m.delivery == nil { + if m == nil { + return 0 + } + return m.reqMsgID + } + m.delivery.mu.Lock() + id := m.delivery.writtenReqMsgID + if id == 0 { + id = m.delivery.targetReqMsgID + } + m.delivery.mu.Unlock() + if id == 0 { + return m.reqMsgID + } + return id +} + +func (m *encodedOutboundMessage) markReplayable() { + if m == nil || m.delivery == nil { + return + } + m.delivery.mu.Lock() + if m.deliveryState() != rpcResultDeliveryDelivered { + m.delivery.state.Store(uint32(rpcResultDeliveryReplayable)) + } + m.delivery.mu.Unlock() +} + +func (m *encodedOutboundMessage) markDelivered() { + if m == nil || m.delivery == nil { + return + } + m.delivery.mu.Lock() + m.delivery.state.Store(uint32(rpcResultDeliveryDelivered)) + m.delivery.mu.Unlock() + if m.delivery.fn != nil { + m.delivery.once.Do(func() { scheduleRPCDeliveryHook(m.delivery.fn) }) + } +} + +func cloneRPCResultForRequest(encoded *encodedOutboundMessage, reqMsgID int64, shareDelivery bool) (*encodedOutboundMessage, error) { + if encoded == nil || encoded.typeID != proto.ResultTypeID || len(encoded.body) < 12 || reqMsgID == 0 { + return nil, errors.New("invalid rpc_result retarget") + } + body := append([]byte(nil), encoded.body...) + binary.LittleEndian.PutUint64(body[4:12], uint64(reqMsgID)) + delivery := newRPCResultDelivery(reqMsgID) + if shareDelivery { + delivery = encoded.delivery + } + return &encodedOutboundMessage{ + body: body, typeID: encoded.typeID, reqMsgID: reqMsgID, + priority: encoded.priority, delivery: delivery, compressed: encoded.compressed, + uncompressedBytes: encoded.uncompressedBytes, + }, nil } type outboundResult struct { @@ -314,12 +482,20 @@ func (c *Conn) startOutbound() { if c.outboundQueueSize <= 0 { c.outboundQueueSize = defaultOutboundQueueSize } + if c.outboundQueueSize < 3 { + c.outboundQueueSize = 3 + } if c.outboundControlQueueSize <= 0 { c.outboundControlQueueSize = defaultOutboundControlQueueSize } c.ensureOutboundTrackedBudget() - c.outbound = make(chan outboundOp, c.outboundQueueSize) + criticalSize := min(defaultOutboundCriticalQueueSize, max(1, c.outboundQueueSize/8)) + bulkSize := min(defaultOutboundBulkQueueSize, max(1, c.outboundQueueSize/8)) + normalSize := c.outboundQueueSize - criticalSize - bulkSize + c.outbound = make(chan outboundOp, normalSize) c.outboundControl = make(chan outboundOp, c.outboundControlQueueSize) + c.outboundCritical = make(chan outboundOp, criticalSize) + c.outboundBulk = make(chan outboundOp, bulkSize) c.outboundStop = make(chan struct{}) c.outboundDone = make(chan struct{}) go c.outboundLoop() @@ -525,8 +701,9 @@ func (c *Conn) sendBestEffort(ctx context.Context, t proto.MessageType, msg bin. op.enqueuedAt = time.Now() // 快路径：非阻塞入队。fan-out 每 (conn × push) 都走这里，队列有空位时不为 // 本次推送分配任何 timer（此前 timeout>0 无条件 WithTimeout，稳态白建 timer）。 + q := c.outboundQueue(op) select { - case c.outbound <- op: + case q <- op: return nil case <-c.outboundStop: op.releaseReservation(c.outboundTrackedBudget) @@ -538,7 +715,7 @@ func (c *Conn) sendBestEffort(ctx context.Context, t proto.MessageType, msg bin. c.metrics.OutboundDropped("push_queue_full") return ErrOutboundQueueFull } - c.metrics.OutboundQueueWait(len(c.outbound), cap(c.outbound)) + c.metrics.OutboundQueueWait(len(q), cap(q)) if ctx == nil { ctx = context.Background() } @@ -549,7 +726,7 @@ func (c *Conn) sendBestEffort(ctx context.Context, t proto.MessageType, msg bin. timeoutC = timer.C } select { - case c.outbound <- op: + case q <- op: return nil case <-timeoutC: op.releaseReservation(c.outboundTrackedBudget) @@ -572,6 +749,41 @@ func (c *Conn) SendEncoded(ctx context.Context, t proto.MessageType, encoded *en return c.sendOutbound(ctx, t, nil, encoded, false) } +// enqueueEncodedDelivery transfers an immutable body to the bounded egress actor +// and returns after queue admission, not after socket I/O. terminal becomes actor- +// owned only on success and is invoked for write success, write failure, or drain. +func (c *Conn) enqueueEncodedDelivery( + ctx context.Context, + t proto.MessageType, + encoded *encodedOutboundMessage, + priority outboundPriority, + terminal func(error), +) error { + if c.outbound == nil || c.outboundControl == nil || c.outboundCritical == nil || c.outboundBulk == nil { + return ErrConnClosed + } + op, err := c.newOutboundSendOp(ctx, t, nil, encoded, false) + if err != nil { + c.failOutboundBudget(err) + return err + } + if !c.beginOutboundEnqueue() { + op.releaseReservation(c.outboundTrackedBudget) + return ErrConnClosed + } + op.priority = priority + op.ctx = context.Background() + op.enqueuedAt = time.Now() + op.terminal = terminal + if err := c.enqueueOutboundRegistered(ctx, op); err != nil { + op.releaseReservation(c.outboundTrackedBudget) + c.endOutboundEnqueue() + return err + } + c.endOutboundEnqueue() + return nil +} + func (c *Conn) sendOutbound(ctx context.Context, t proto.MessageType, msg bin.Encoder, encoded *encodedOutboundMessage, control bool) error { if c.outbound == nil || c.outboundControl == nil { return ErrConnClosed @@ -774,10 +986,7 @@ func (c *Conn) enqueueOutboundRegistered(ctx context.Context, op outboundOp) err if c.isRetired() { return ErrConnClosed } - q := c.outbound - if op.control { - q = c.outboundControl - } + q := c.outboundQueue(op) select { case q <- op: return nil @@ -798,6 +1007,20 @@ func (c *Conn) enqueueOutboundRegistered(ctx context.Context, op outboundOp) err } } +func (c *Conn) outboundQueue(op outboundOp) chan outboundOp { + if op.control || op.priority == outboundPriorityControl { + return c.outboundControl + } + switch op.priority { + case outboundPriorityCritical: + return c.outboundCritical + case outboundPriorityBulk: + return c.outboundBulk + default: + return c.outbound + } +} + func (c *Conn) beginOutboundEnqueue() bool { c.outboundEnqueueMu.Lock() defer c.outboundEnqueueMu.Unlock() @@ -814,6 +1037,7 @@ func (c *Conn) endOutboundEnqueue() { func (c *Conn) outboundLoop() { state := newOutboundState(c.outboundTrackedBudget) + ordinarySinceBulk := 0 defer func() { // pending frames belong exclusively to this actor. Releasing after drain ensures no // resend path can race the final budget return and no Conn body survives actor exit. @@ -826,47 +1050,19 @@ func (c *Conn) outboundLoop() { c.drainOutbound() return } - select { - case op := <-c.outboundControl: - if c.isRetired() { - op.releaseReservation(c.outboundTrackedBudget) - op.finish(outboundResult{err: ErrConnClosed}) - c.signalOutboundStop() - c.drainOutbound() - return - } - c.handleOutboundOp(state, op) - if c.isRetired() { - c.signalOutboundStop() - c.drainOutbound() - return - } - continue - default: - } - select { - case <-c.outboundStop: + op, ok := c.nextOutboundOp(&ordinarySinceBulk) + if !ok { c.drainOutbound() return - case op := <-c.outboundControl: - if c.isRetired() { - op.releaseReservation(c.outboundTrackedBudget) - op.finish(outboundResult{err: ErrConnClosed}) - c.signalOutboundStop() - c.drainOutbound() - return - } - c.handleOutboundOp(state, op) - case op := <-c.outbound: - if c.isRetired() { - op.releaseReservation(c.outboundTrackedBudget) - op.finish(outboundResult{err: ErrConnClosed}) - c.signalOutboundStop() - c.drainOutbound() - return - } - c.handleOutboundOp(state, op) } + if c.isRetired() { + op.releaseReservation(c.outboundTrackedBudget) + op.finish(outboundResult{err: ErrConnClosed}) + c.signalOutboundStop() + c.drainOutbound() + return + } + c.handleOutboundOp(state, op) if c.isRetired() { c.signalOutboundStop() c.drainOutbound() @@ -875,6 +1071,55 @@ func (c *Conn) outboundLoop() { } } +// nextOutboundOp applies the connection-wide egress policy without introducing +// another writer. Required protocol controls stay strict, convergence RPCs pass +// ordinary/bulk work, and a bounded ordinary burst guarantees bulk progress. +func (c *Conn) nextOutboundOp(ordinarySinceBulk *int) (outboundOp, bool) { + try := func(q <-chan outboundOp) (outboundOp, bool) { + select { + case op := <-q: + return op, true + default: + return outboundOp{}, false + } + } + if op, ok := try(c.outboundControl); ok { + return op, true + } + if op, ok := try(c.outboundCritical); ok { + return op, true + } + if *ordinarySinceBulk >= maxOrdinaryBeforeBulk { + if op, ok := try(c.outboundBulk); ok { + *ordinarySinceBulk = 0 + return op, true + } + } + if op, ok := try(c.outbound); ok { + *ordinarySinceBulk++ + return op, true + } + if op, ok := try(c.outboundBulk); ok { + *ordinarySinceBulk = 0 + return op, true + } + + select { + case <-c.outboundStop: + return outboundOp{}, false + case op := <-c.outboundControl: + return op, true + case op := <-c.outboundCritical: + return op, true + case op := <-c.outbound: + *ordinarySinceBulk++ + return op, true + case op := <-c.outboundBulk: + *ordinarySinceBulk = 0 + return op, true + } +} + func (c *Conn) drainOutbound() { // signalOutboundStop closes the producer gate before the actor gets here. // Waiting first guarantees every producer either enqueued an owned op or @@ -885,9 +1130,15 @@ func (c *Conn) drainOutbound() { case op := <-c.outboundControl: op.releaseReservation(c.outboundTrackedBudget) op.finish(outboundResult{err: ErrConnClosed}) + case op := <-c.outboundCritical: + op.releaseReservation(c.outboundTrackedBudget) + op.finish(outboundResult{err: ErrConnClosed}) case op := <-c.outbound: op.releaseReservation(c.outboundTrackedBudget) op.finish(outboundResult{err: ErrConnClosed}) + case op := <-c.outboundBulk: + op.releaseReservation(c.outboundTrackedBudget) + op.finish(outboundResult{err: ErrConnClosed}) default: return } @@ -902,7 +1153,11 @@ func (c *Conn) handleOutboundOp(state *outboundState, op outboundOp) { case outboundSend: c.handleOutboundSend(state, op) case outboundAck: - state.ack(op.ids) + for _, reqMsgID := range state.ack(op.ids) { + if c.rpcResultAcked != nil { + c.rpcResultAcked(c, reqMsgID) + } + } case outboundQueryState: op.finish(outboundResult{info: state.stateInfo(op.ids)}) case outboundResend: @@ -917,7 +1172,17 @@ func (c *Conn) handleOutboundOp(state *outboundState, op outboundOp) { } func (c *Conn) handleOutboundSend(state *outboundState, op outboundOp) { - frame, err := c.buildFrame(op.ctx, op.msgType, op.msg, op.encoded) + var err error + if op.encoded != nil { + targetReqMsgID := op.encoded.beginWriting() + if targetReqMsgID != 0 && targetReqMsgID != op.encoded.reqMsgID { + op.encoded, err = cloneRPCResultForRequest(op.encoded, targetReqMsgID, true) + } + } + var frame *outboundFrame + if err == nil { + frame, err = c.buildFrame(op.ctx, op.msgType, op.msg, op.encoded) + } reserved := op.reservedBytes reservationBudget := op.reservationBudget if reservationBudget == nil { @@ -1018,6 +1283,9 @@ func (c *Conn) handleOutboundResendByRequest(state *outboundState, ctx context.C } func (op outboundOp) finish(res outboundResult) { + if op.terminal != nil { + op.terminal(res.err) + } if op.done == nil { return } @@ -1093,6 +1361,7 @@ func (c *Conn) newOutboundSendOp(ctx context.Context, t proto.MessageType, msg b kind: outboundSend, msgType: t, encoded: encoded, + priority: classifyOutboundPriority(encoded, priorityControl), reservedBytes: bytes, reservationBudget: budget, }, nil @@ -1112,11 +1381,25 @@ func (c *Conn) newOutboundSendOp(ctx context.Context, t proto.MessageType, msg b kind: outboundSend, msgType: t, encoded: encoded, + priority: classifyOutboundPriority(encoded, priorityControl), reservedBytes: bytes, reservationBudget: budget, }, nil } +func classifyOutboundPriority(encoded *encodedOutboundMessage, control bool) outboundPriority { + if control { + return outboundPriorityControl + } + if encoded != nil && encoded.priority != outboundPriorityNormal { + return encoded.priority + } + if encoded != nil && len(encoded.body) >= bulkOutboundThreshold { + return outboundPriorityBulk + } + return outboundPriorityNormal +} + func (c *Conn) outboundMessageBudget(typeID uint32, priorityControl bool) *outboundTrackedBudget { if priorityControl || encodedControlFrame(typeID) { return c.ensureOutboundControlTrackedBudget() @@ -1210,7 +1493,11 @@ func (c *Conn) downgradedCloneContext(ctx context.Context, encoded *encodedOutbo if sameBacking(down, encoded.body) { return encoded // 直通：未变(mt.*/顶层未知)，无需拷贝或重算 typeID } - out := &encodedOutboundMessage{body: down, typeID: encoded.typeID, reqMsgID: encoded.reqMsgID} + out := &encodedOutboundMessage{ + body: down, typeID: encoded.typeID, reqMsgID: encoded.reqMsgID, + priority: encoded.priority, delivery: encoded.delivery, compressed: encoded.compressed, + uncompressedBytes: encoded.uncompressedBytes, + } if id, e := (&bin.Buffer{Buf: down}).PeekID(); e == nil { out.typeID = id } @@ -1535,8 +1822,16 @@ func (s *outboundState) addReserved(frame *outboundFrame) int { return s.shrinkPending() } -func (s *outboundState) ack(ids []int64) { +func (s *outboundState) ack(ids []int64) []int64 { + var requestIDs []int64 for _, id := range ids { + frame, ok := s.pending[id] + if !ok { + continue + } + if frame.reqMsgID != 0 { + requestIDs = append(requestIDs, frame.reqMsgID) + } if !s.removePending(id) { continue } @@ -1545,6 +1840,7 @@ func (s *outboundState) ack(ids []int64) { if len(s.order) > s.maxMessages*2 { s.compactOrder() } + return requestIDs } func (s *outboundState) stateInfo(ids []int64) []byte { diff --git a/internal/mtprotoedge/outbound_test.go b/internal/mtprotoedge/outbound_test.go index f54de3fc..82057f25 100644 --- a/internal/mtprotoedge/outbound_test.go +++ b/internal/mtprotoedge/outbound_test.go @@ -397,8 +397,8 @@ func TestOutboundQueueBackingUsesSmallConfigurableBounds(t *testing.T) { c := &Conn{metrics: NopMetrics{}} c.startOutbound() defer c.Close() - if got := cap(c.outbound); got != defaultOutboundQueueSize { - t.Fatalf("normal queue cap = %d, want %d", got, defaultOutboundQueueSize) + if got := cap(c.outbound) + cap(c.outboundCritical) + cap(c.outboundBulk); got != defaultOutboundQueueSize { + t.Fatalf("ordinary lane total cap = %d, want %d", got, defaultOutboundQueueSize) } if got := cap(c.outboundControl); got != defaultOutboundControlQueueSize { t.Fatalf("control queue cap = %d, want %d", got, defaultOutboundControlQueueSize) @@ -413,8 +413,8 @@ func TestOutboundQueueBackingUsesSmallConfigurableBounds(t *testing.T) { } c.startOutbound() defer c.Close() - if got := cap(c.outbound); got != 7 { - t.Fatalf("normal queue cap = %d, want 7", got) + if got := cap(c.outbound) + cap(c.outboundCritical) + cap(c.outboundBulk); got != 7 { + t.Fatalf("ordinary lane total cap = %d, want 7", got) } if got := cap(c.outboundControl); got != 3 { t.Fatalf("control queue cap = %d, want 3", got) @@ -446,9 +446,11 @@ func TestServerNewConnectionsShareOutboundBudgetAndQueueLimits(t *testing.T) { defer c1.Close() defer c2.Close() - if cap(c1.outbound) != 7 || cap(c1.outboundControl) != 3 || cap(c2.outbound) != 7 || cap(c2.outboundControl) != 3 { + c1Ordinary := cap(c1.outbound) + cap(c1.outboundCritical) + cap(c1.outboundBulk) + c2Ordinary := cap(c2.outbound) + cap(c2.outboundCritical) + cap(c2.outboundBulk) + if c1Ordinary != 7 || cap(c1.outboundControl) != 3 || c2Ordinary != 7 || cap(c2.outboundControl) != 3 { t.Fatalf("server queue caps = %d/%d and %d/%d, want 7/3", - cap(c1.outbound), cap(c1.outboundControl), cap(c2.outbound), cap(c2.outboundControl)) + c1Ordinary, cap(c1.outboundControl), c2Ordinary, cap(c2.outboundControl)) } if c1.outboundTrackedBudget != srv.outboundTrackedBudget || c2.outboundTrackedBudget != srv.outboundTrackedBudget { t.Fatal("server connections did not receive the shared outbound tracking budget") diff --git a/internal/mtprotoedge/rpc_flight.go b/internal/mtprotoedge/rpc_flight.go index 3c790c65..8aaebff7 100644 --- a/internal/mtprotoedge/rpc_flight.go +++ b/internal/mtprotoedge/rpc_flight.go @@ -3,6 +3,7 @@ package mtprotoedge import ( "context" "errors" + "sync" "sync/atomic" ) @@ -41,12 +42,15 @@ type rpcResultAcquire struct { // done channel is closed exactly once while holding the owning cache shard lock; // channel close publishes encoded/ok to all waiters without a waiter goroutine. type rpcResultFlight struct { - done chan struct{} - encoded *encodedOutboundMessage - ok bool + done chan struct{} + encoded *encodedOutboundMessage + ok bool + subscribers []func(*encodedOutboundMessage, bool) } type rpcResultWaiter struct { + cache *rpcResultCache + key rpcResultCacheKey flight *rpcResultFlight } @@ -78,10 +82,115 @@ func (w *rpcResultWaiter) Wait(ctx context.Context) (encoded *encodedOutboundMes } } +// Subscribe registers an event callback without creating a goroutine or +// occupying an RPC worker. The callback is invoked after the cache shard lock is +// released; it must remain non-blocking. +func (w *rpcResultWaiter) Subscribe(fn func(*encodedOutboundMessage, bool)) error { + if w == nil || w.cache == nil || w.flight == nil || fn == nil { + return ErrRPCResultFlightInvalid + } + s := w.cache.shard(w.key) + var ( + encoded *encodedOutboundMessage + ok bool + ready bool + ) + s.mu.Lock() + if flight, exists := s.pending[w.key]; exists && flight == w.flight { + flight.subscribers = append(flight.subscribers, fn) + s.mu.Unlock() + return nil + } + select { + case <-w.flight.done: + encoded, ok, ready = w.flight.encoded, w.flight.ok, true + default: + } + s.mu.Unlock() + if !ready { + return ErrRPCResultFlightInvalid + } + fn(encoded, ok) + return nil +} + type rpcResultOwnerLease struct { - cache *rpcResultCache - key rpcResultCacheKey - flight *rpcResultFlight + cache *rpcResultCache + key rpcResultCacheKey + flight *rpcResultFlight + delivery *rpcResultDelivery + hookMu sync.Mutex + abortHook func() + // handedOff means the inbound worker transferred terminal-result ownership to + // the bounded egress pipeline. Its ordinary release callback must no longer + // abort the flight merely because the socket write is still pending. + handedOff atomic.Bool +} + +func (l *rpcResultOwnerLease) SetAbortHook(fn func()) { + if l == nil { + return + } + l.hookMu.Lock() + l.abortHook = fn + l.hookMu.Unlock() +} + +// InstallAbortHook installs fn only while this lease still owns the pending +// flight. The shard lock linearizes installation with Abort/Put so a registry +// cannot publish a candidate after its owner has already disappeared. +func (l *rpcResultOwnerLease) InstallAbortHook(fn func()) bool { + if l == nil || l.cache == nil || l.flight == nil || fn == nil { + return false + } + s := l.cache.shard(l.key) + s.mu.Lock() + flight, ok := s.pending[l.key] + if !ok || flight != l.flight { + s.mu.Unlock() + return false + } + l.hookMu.Lock() + l.abortHook = fn + l.hookMu.Unlock() + s.mu.Unlock() + return true +} + +func (l *rpcResultOwnerLease) Waiter() *rpcResultWaiter { + if l == nil || l.cache == nil || l.flight == nil { + return nil + } + return &rpcResultWaiter{cache: l.cache, key: l.key, flight: l.flight} +} + +func (l *rpcResultOwnerLease) TryRetarget(reqMsgID int64) bool { + return l != nil && l.delivery != nil && (&encodedOutboundMessage{delivery: l.delivery}).tryRetarget(reqMsgID) +} + +func (l *rpcResultOwnerLease) Delivery() *rpcResultDelivery { + if l == nil { + return nil + } + return l.delivery +} + +// HandOff transfers completion responsibility from the inbound RPC task to an +// already-admitted egress operation. The egress terminal callback must resolve +// the flight through Put on both successful delivery and fenced failure. +func (l *rpcResultOwnerLease) HandOff() bool { + if l == nil || l.cache == nil || l.flight == nil { + return false + } + s := l.cache.shard(l.key) + s.mu.Lock() + defer s.mu.Unlock() + flight, ok := s.pending[l.key] + if !ok || flight != l.flight { + return false + } + l.handedOff.Store(true) + return true } // Abort releases an unfinished owner claim and wakes every waiter with no @@ -91,17 +200,37 @@ func (l *rpcResultOwnerLease) Abort() bool { if l == nil || l.cache == nil || l.flight == nil { return false } + if l.handedOff.Load() { + return false + } s := l.cache.shard(l.key) s.mu.Lock() - defer s.mu.Unlock() + if l.handedOff.Load() { + s.mu.Unlock() + return false + } flight, ok := s.pending[l.key] if !ok || flight != l.flight { + s.mu.Unlock() return false } delete(s.pending, l.key) l.cache.flightLimit.release() + subscribers := append([]func(*encodedOutboundMessage, bool){}, flight.subscribers...) + flight.subscribers = nil close(flight.done) + s.mu.Unlock() + l.hookMu.Lock() + abortHook := l.abortHook + l.abortHook = nil + l.hookMu.Unlock() + if abortHook != nil { + abortHook() + } + for _, subscriber := range subscribers { + subscriber(nil, false) + } return true } @@ -168,7 +297,7 @@ func (c *rpcResultCache) Acquire(authKeyID [8]byte, sessionID, reqMsgID int64) ( if flight, ok := s.pending[key]; ok { return rpcResultAcquire{ state: rpcResultAcquirePending, - waiter: &rpcResultWaiter{flight: flight}, + waiter: &rpcResultWaiter{cache: c, key: key, flight: flight}, }, nil } if !c.flightLimit.reserve() { @@ -181,7 +310,9 @@ func (c *rpcResultCache) Acquire(authKeyID [8]byte, sessionID, reqMsgID int64) ( s.pending[key] = flight return rpcResultAcquire{ state: rpcResultAcquireOwner, - owner: &rpcResultOwnerLease{cache: c, key: key, flight: flight}, + owner: &rpcResultOwnerLease{ + cache: c, key: key, flight: flight, delivery: newRPCResultDelivery(reqMsgID), + }, }, nil } @@ -191,17 +322,20 @@ func (c *rpcResultCache) completeRPCResultFlightLocked( s *rpcResultCacheShard, key rpcResultCacheKey, encoded *encodedOutboundMessage, -) { +) []func(*encodedOutboundMessage, bool) { if c == nil || s == nil || encoded == nil { - return + return nil } flight, ok := s.pending[key] if !ok { - return + return nil } delete(s.pending, key) flight.encoded = encoded flight.ok = true c.flightLimit.release() + subscribers := append([]func(*encodedOutboundMessage, bool){}, flight.subscribers...) + flight.subscribers = nil close(flight.done) + return subscribers } diff --git a/internal/mtprotoedge/rpc_result_cache.go b/internal/mtprotoedge/rpc_result_cache.go index 11eb7d82..231996ef 100644 --- a/internal/mtprotoedge/rpc_result_cache.go +++ b/internal/mtprotoedge/rpc_result_cache.go @@ -118,7 +118,6 @@ func (c *rpcResultCache) Put(authKeyID [8]byte, sessionID, reqMsgID int64, encod now := s.now() s.mu.Lock() - defer s.mu.Unlock() if cacheable { s.expireLocked(now) @@ -139,7 +138,11 @@ func (c *rpcResultCache) Put(authKeyID [8]byte, sessionID, reqMsgID int64, encod // Resolve the independent in-flight entry only after the completed cache has // been published. Waiters awakened by this close can therefore immediately // observe either the shared encoded result or the completed Get entry. - c.completeRPCResultFlightLocked(s, key, encoded) + subscribers := c.completeRPCResultFlightLocked(s, key, encoded) + s.mu.Unlock() + for _, subscriber := range subscribers { + subscriber(encoded, true) + } } func (s *rpcResultCacheShard) expireLocked(now time.Time) { diff --git a/internal/mtprotoedge/rpc_result_egress.go b/internal/mtprotoedge/rpc_result_egress.go new file mode 100644 index 00000000..6e89c919 --- /dev/null +++ b/internal/mtprotoedge/rpc_result_egress.go @@ -0,0 +1,110 @@ +package mtprotoedge + +import ( + "context" + "strings" + "sync" + + "github.com/gotd/td/bin" + "github.com/gotd/td/proto" +) + +const ( + rpcResultGZIPMinBytes = 4 << 10 + rpcResultGZIPMaxInputBytes = (10 << 20) - 1 // gotd client decompression hard limit. + rpcResultGZIPMinSavedBytes = 1 << 10 + rpcResultGZIPMinSavedDivisor = 12 // Require roughly 8.3% reduction. + rpcResultGZIPConcurrency = 8 + rpcDeliveryHookConcurrency = 8 + rpcDeliveryHookQueueSize = 1024 +) + +var rpcResultGZIPSlots = make(chan struct{}, rpcResultGZIPConcurrency) + +var ( + rpcDeliveryHooksOnce sync.Once + rpcDeliveryHooks chan func() +) + +// scheduleRPCDeliveryHook keeps database/update follow-up work off the sole +// socket writer. Hooks are internal, bounded and timeout-aware at registration +// sites; the fixed worker set prevents one delivered result from stalling every +// subsequent frame on that connection. +func scheduleRPCDeliveryHook(fn func()) { + if fn == nil { + return + } + rpcDeliveryHooksOnce.Do(func() { + rpcDeliveryHooks = make(chan func(), rpcDeliveryHookQueueSize) + for range rpcDeliveryHookConcurrency { + go func() { + for hook := range rpcDeliveryHooks { + hook() + } + }() + } + }) + rpcDeliveryHooks <- fn +} + +// encodeAdaptiveRPCResultInner returns either the original layer-specific TL +// object or one complete gzip_packed object. Compression is CPU bounded and is +// retained only when it materially reduces the non-preemptible transport frame. +func encodeAdaptiveRPCResultInner(ctx context.Context, stop <-chan struct{}, inner []byte) ([]byte, bool, error) { + if len(inner) < rpcResultGZIPMinBytes || len(inner) > rpcResultGZIPMaxInputBytes { + return inner, false, nil + } + if ctx == nil { + ctx = context.Background() + } + select { + case rpcResultGZIPSlots <- struct{}{}: + case <-ctx.Done(): + return nil, false, ctx.Err() + case <-stop: + return nil, false, ErrConnClosed + } + defer func() { <-rpcResultGZIPSlots }() + + var packed bin.Buffer + if err := (proto.GZIP{Data: inner}).Encode(&packed); err != nil { + return nil, false, err + } + saved := len(inner) - packed.Len() + required := max(rpcResultGZIPMinSavedBytes, len(inner)/rpcResultGZIPMinSavedDivisor) + if saved < required { + return inner, false, nil + } + return packed.Raw(), true, nil +} + +// rpcResultPriority is protocol scheduling metadata, not handler business +// behavior. Difference/state responses converge the update state, while the +// dialogs+pinned pair converges the initial chat list in both TDesktop and +// Android. These bootstrap barriers must pass background prefetch regardless of +// platform or their own encoded size. +func rpcResultPriority(method string, encoded *encodedOutboundMessage) outboundPriority { + base := method + if i := strings.IndexByte(base, '#'); i >= 0 { + base = base[:i] + } + switch base { + case "updates.getDifference", "updates.getChannelDifference", "updates.getState", + "messages.getDialogs", "messages.getPinnedDialogs": + return outboundPriorityCritical + } + return classifyOutboundPriority(encoded, false) +} + +func (p outboundPriority) String() string { + switch p { + case outboundPriorityCritical: + return "convergence" + case outboundPriorityBulk: + return "bulk" + case outboundPriorityControl: + return "control" + default: + return "normal" + } +} diff --git a/internal/mtprotoedge/rpc_result_egress_test.go b/internal/mtprotoedge/rpc_result_egress_test.go new file mode 100644 index 00000000..cb9e355b --- /dev/null +++ b/internal/mtprotoedge/rpc_result_egress_test.go @@ -0,0 +1,444 @@ +package mtprotoedge + +import ( + "bytes" + "context" + "crypto/rand" + "io" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/gotd/td/bin" + "github.com/gotd/td/crypto" + "github.com/gotd/td/proto" + "github.com/gotd/td/tg" +) + +type opaqueRPCResult struct{ body []byte } + +func (o opaqueRPCResult) Encode(b *bin.Buffer) error { + b.PutID(0x10203040) + b.Put(o.body) + return nil +} + +func TestEncodeRPCResultUsesAdaptiveGZIP(t *testing.T) { + s := New(Options{}) + c := &Conn{metrics: NopMetrics{}} + large := &tg.DataJSON{Data: string(bytes.Repeat([]byte("sticker-metadata-"), 16<<10))} + encoded, err := s.encodeRPCResult(c, 123, large) + if err != nil { + t.Fatalf("encode compressed rpc_result: %v", err) + } + if !encoded.compressed { + t.Fatal("compressible large rpc_result was not gzip_packed") + } + if encoded.uncompressedBytes <= len(encoded.body) { + t.Fatalf("compressed wire=%d is not smaller than inner=%d", len(encoded.body), encoded.uncompressedBytes) + } + var result proto.Result + if err := result.Decode(&bin.Buffer{Buf: encoded.body}); err != nil { + t.Fatalf("decode rpc_result: %v", err) + } + var packed proto.GZIP + if err := packed.Decode(&bin.Buffer{Buf: result.Result}); err != nil { + t.Fatalf("decode gzip_packed: %v", err) + } + var decoded tg.DataJSON + if err := decoded.Decode(&bin.Buffer{Buf: packed.Data}); err != nil { + t.Fatalf("decode compressed inner result: %v", err) + } + if decoded.Data != large.Data { + t.Fatal("gzip round trip changed rpc_result") + } +} + +func TestEncodeRPCResultKeepsIncompressibleBodyRaw(t *testing.T) { + s := New(Options{}) + c := &Conn{metrics: NopMetrics{}} + raw := make([]byte, 96<<10) + if _, err := rand.Read(raw); err != nil { + t.Fatalf("random body: %v", err) + } + encoded, err := s.encodeRPCResult(c, 456, opaqueRPCResult{body: raw}) + if err != nil { + t.Fatalf("encode incompressible rpc_result: %v", err) + } + if encoded.compressed { + t.Fatal("incompressible rpc_result retained a larger gzip envelope") + } + var result proto.Result + if err := result.Decode(&bin.Buffer{Buf: encoded.body}); err != nil { + t.Fatalf("decode raw rpc_result: %v", err) + } + id, err := (&bin.Buffer{Buf: result.Result}).PeekID() + if err != nil || id != 0x10203040 { + t.Fatalf("raw result type = %#x err=%v", id, err) + } +} + +func TestBootstrapBarriersAlwaysUseConvergenceLane(t *testing.T) { + large := &encodedOutboundMessage{body: make([]byte, bulkOutboundThreshold)} + for _, method := range []string{ + "updates.getDifference", "updates.getDifference#25939651", + "updates.getChannelDifference#03173d78", "updates.getState", + "messages.getDialogs", "messages.getDialogs#a0f4cb4f", + "messages.getPinnedDialogs", "messages.getPinnedDialogs#d6b94df2", + } { + if got := rpcResultPriority(method, large); got != outboundPriorityCritical { + t.Fatalf("priority(%q) = %s, want convergence", method, got.String()) + } + } + if got := rpcResultPriority("messages.getStickerSet", large); got != outboundPriorityBulk { + t.Fatalf("sticker-set priority = %s, want bulk", got.String()) + } +} + +type gatedRecordingTransport struct { + started chan struct{} + release chan struct{} + once sync.Once + sends atomic.Int32 + mu sync.Mutex + frames [][]byte +} + +func newGatedRecordingTransport() *gatedRecordingTransport { + return &gatedRecordingTransport{started: make(chan struct{}), release: make(chan struct{})} +} + +func (t *gatedRecordingTransport) Send(_ context.Context, b *bin.Buffer) error { + if t.sends.Add(1) == 1 { + close(t.started) + <-t.release + } + t.mu.Lock() + t.frames = append(t.frames, append([]byte(nil), b.Raw()...)) + t.mu.Unlock() + return nil +} + +func (*gatedRecordingTransport) Recv(context.Context, *bin.Buffer) error { return io.EOF } +func (t *gatedRecordingTransport) Close() error { + t.once.Do(func() { close(t.release) }) + return nil +} + +func (t *gatedRecordingTransport) snapshot() [][]byte { + t.mu.Lock() + defer t.mu.Unlock() + out := make([][]byte, len(t.frames)) + for i := range t.frames { + out[i] = append([]byte(nil), t.frames[i]...) + } + return out +} + +func encodedRPCResultForPriorityTest(reqMsgID int64, payloadBytes int) *encodedOutboundMessage { + var b bin.Buffer + b.PutID(proto.ResultTypeID) + b.PutLong(reqMsgID) + b.PutID(tg.BoolTrueTypeID) + if payloadBytes > 0 { + b.Put(make([]byte, payloadBytes)) + } + return &encodedOutboundMessage{typeID: proto.ResultTypeID, reqMsgID: reqMsgID, body: b.Raw()} +} + +func TestConvergenceResultPassesQueuedBulkAfterBlockedWrite(t *testing.T) { + tr := newGatedRecordingTransport() + c := newOutboundTestConn(t, tr, newOutboundTrackedBudget(2<<20)) + gate := &encodedOutboundMessage{typeID: tg.UpdatesTooLongTypeID, body: []byte{0x0b, 0xa1, 0x01, 0xe3}} + if err := c.SendBestEffortEncoded(context.Background(), proto.MessageFromServer, gate, 0); err != nil { + t.Fatalf("enqueue gate: %v", err) + } + select { + case <-tr.started: + case <-time.After(time.Second): + t.Fatal("first write did not block") + } + + const ordinaryResults = 17 + for i := 0; i < ordinaryResults; i++ { + ordinary := encodedRPCResultForPriorityTest(2000+int64(i), 0) + if err := c.enqueueEncodedDelivery(context.Background(), proto.MessageServerResponse, ordinary, outboundPriorityNormal, nil); err != nil { + t.Fatalf("enqueue ordinary result %d: %v", i, err) + } + } + bulk := encodedRPCResultForPriorityTest(1001, bulkOutboundThreshold) + critical := encodedRPCResultForPriorityTest(1002, 0) + if err := c.enqueueEncodedDelivery(context.Background(), proto.MessageServerResponse, bulk, outboundPriorityBulk, nil); err != nil { + t.Fatalf("enqueue bulk: %v", err) + } + if err := c.enqueueEncodedDelivery(context.Background(), proto.MessageServerResponse, critical, outboundPriorityCritical, nil); err != nil { + t.Fatalf("enqueue convergence result: %v", err) + } + tr.once.Do(func() { close(tr.release) }) + wantSends := int32(1 + ordinaryResults + 2) + deadline := time.Now().Add(2 * time.Second) + for tr.sends.Load() < wantSends && time.Now().Before(deadline) { + time.Sleep(time.Millisecond) + } + if got := tr.sends.Load(); got != wantSends { + t.Fatalf("physical sends = %d, want %d", got, wantSends) + } + + var resultOrder []int64 + clientCipher := crypto.NewClientCipher(rand.Reader) + for _, frame := range tr.snapshot() { + data, err := clientCipher.DecryptFromBuffer(c.key, &bin.Buffer{Buf: frame}) + if err != nil { + t.Fatalf("decrypt frame: %v", err) + } + plain := &bin.Buffer{Buf: append([]byte(nil), data.Data()...)} + id, err := plain.PeekID() + if err != nil || id != proto.ResultTypeID { + continue + } + var result proto.Result + if err := result.Decode(plain); err != nil { + t.Fatalf("decode result: %v", err) + } + resultOrder = append(resultOrder, result.RequestMessageID) + } + if len(resultOrder) != ordinaryResults+2 || resultOrder[0] != 1002 { + t.Fatalf("rpc_result order = %v, want convergence first", resultOrder) + } + bulkIndex := -1 + for i, reqMsgID := range resultOrder { + if reqMsgID == 1001 { + bulkIndex = i + break + } + } + if bulkIndex < 0 || bulkIndex > maxOrdinaryBeforeBulk+1 { + t.Fatalf("bulk result index = %d in %v, want bounded ordinary burst", bulkIndex, resultOrder) + } +} + +type immediateLargeRPC struct{} + +func (immediateLargeRPC) Dispatch(context.Context, [8]byte, int64, *bin.Buffer) (bin.Encoder, error) { + return &tg.DataJSON{Data: string(bytes.Repeat([]byte("large-sticker-set"), 16<<10))}, nil +} + +func (immediateLargeRPC) NegotiatedLayer([8]byte, int64) (int, bool) { return 227, true } + +type wrappedConvergenceRPC struct{ immediateLargeRPC } + +func (w wrappedConvergenceRPC) DispatchWithMethod( + ctx context.Context, + authKeyID [8]byte, + sessionID int64, + b *bin.Buffer, +) (bin.Encoder, string, error) { + result, err := w.Dispatch(ctx, authKeyID, sessionID, b) + return result, "updates.getDifference", err +} + +type captureRPCResultMetrics struct { + NopMetrics + mu sync.Mutex + preparedMethod string + priority string + innerBytes int + wireBytes int + compressed bool + delivered chan error +} + +func (m *captureRPCResultMetrics) RPCResultPrepared(method, priority string, innerBytes, wireBytes int, compressed bool) { + m.mu.Lock() + m.preparedMethod, m.priority = method, priority + m.innerBytes, m.wireBytes, m.compressed = innerBytes, wireBytes, compressed + m.mu.Unlock() +} + +func (m *captureRPCResultMetrics) RPCResultDelivered(_ string, _ time.Duration, _ int, err error) { + m.delivered <- err +} + +func TestRPCResultPipelineExportsPreparationAndDeliveryMetrics(t *testing.T) { + metrics := &captureRPCResultMetrics{delivered: make(chan error, 1)} + s := New(Options{Metrics: metrics}) + c := newOutboundTestConn(t, &failAfterTransport{}, newOutboundTrackedBudget(1<<20)) + const reqMsgID = int64(9050) + claim, err := s.rpcResults.Acquire(c.authKeyID, c.sessionID, reqMsgID) + if err != nil || claim.state != rpcResultAcquireOwner { + t.Fatalf("acquire flight = %+v err=%v", claim, err) + } + result := &tg.DataJSON{Data: string(bytes.Repeat([]byte("sticker-data"), 12<<10))} + if err := s.publishRPCResult(c, reqMsgID, "updates.getDifference#25939651", claim.owner, result, nil); err != nil { + t.Fatalf("publish result: %v", err) + } + select { + case err := <-metrics.delivered: + if err != nil { + t.Fatalf("delivery metric error: %v", err) + } + case <-time.After(time.Second): + t.Fatal("missing delivery metric") + } + metrics.mu.Lock() + defer metrics.mu.Unlock() + if metrics.preparedMethod != "updates.getDifference#25939651" || metrics.priority != "convergence" { + t.Fatalf("prepared metric = %q/%q", metrics.preparedMethod, metrics.priority) + } + if !metrics.compressed || metrics.innerBytes <= metrics.wireBytes { + t.Fatalf("compression metric = compressed:%v inner:%d wire:%d", metrics.compressed, metrics.innerBytes, metrics.wireBytes) + } +} + +func TestWrappedConvergenceMethodDrivesEgressAndReplayPriority(t *testing.T) { + metrics := &captureRPCResultMetrics{delivered: make(chan error, 1)} + s := New(Options{RPC: wrappedConvergenceRPC{}, Metrics: metrics}) + c := newOutboundTestConn(t, &failAfterTransport{}, newOutboundTrackedBudget(1<<20)) + const reqMsgID = int64(9051) + claim, err := s.rpcResults.Acquire(c.authKeyID, c.sessionID, reqMsgID) + if err != nil || claim.state != rpcResultAcquireOwner { + t.Fatalf("acquire flight = %+v err=%v", claim, err) + } + body := mustEncodeTL(t, &tg.HelpGetConfigRequest{}) + if err := s.handleRPC(context.Background(), c, reqMsgID, "invokeWithLayer#da9b0d0d", &bin.Buffer{Buf: body}, claim.owner); err != nil { + t.Fatalf("handle wrapped convergence RPC: %v", err) + } + select { + case err := <-metrics.delivered: + if err != nil { + t.Fatalf("delivery metric error: %v", err) + } + case <-time.After(time.Second): + t.Fatal("missing delivery metric") + } + metrics.mu.Lock() + method, priority := metrics.preparedMethod, metrics.priority + metrics.mu.Unlock() + if method != "updates.getDifference" || priority != "convergence" { + t.Fatalf("wrapped prepared metric = %q/%q, want updates.getDifference/convergence", method, priority) + } + var cached *encodedOutboundMessage + deadline := time.Now().Add(time.Second) + for time.Now().Before(deadline) { + if got, ok := s.rpcResults.Get(c.authKeyID, c.sessionID, reqMsgID); ok { + cached = got + break + } + time.Sleep(time.Millisecond) + } + if cached == nil { + t.Fatal("wrapped convergence result missing from replay cache") + } + if got := classifyOutboundPriority(cached, false); got != outboundPriorityCritical { + t.Fatalf("cached convergence priority = %s, want convergence", got.String()) + } +} + +func TestRPCWorkerReleasesAfterEgressAdmissionWhileWriteBlocked(t *testing.T) { + s := New(Options{RPC: immediateLargeRPC{}, WriteTimeout: time.Second}) + tr := newGatedRecordingTransport() + c := newOutboundTestConn(t, tr, newOutboundTrackedBudget(2<<20)) + const reqMsgID = int64(9001) + claim, err := s.rpcResults.Acquire(c.authKeyID, c.sessionID, reqMsgID) + if err != nil || claim.state != rpcResultAcquireOwner { + t.Fatalf("acquire flight = %+v err=%v", claim, err) + } + body := mustEncodeTL(t, &tg.HelpGetConfigRequest{}) + task := s.newInboundRPCTask(c, reqMsgID, "updates.getDifference#25939651", body, claim.owner) + done := make(chan error, 1) + go func() { done <- task.run(context.Background()) }() + select { + case <-tr.started: + case <-time.After(time.Second): + t.Fatal("rpc_result write did not start") + } + select { + case err := <-done: + if err != nil { + t.Fatalf("RPC worker result: %v", err) + } + case <-time.After(250 * time.Millisecond): + t.Fatal("RPC worker remained coupled to blocked physical write") + } + acquired, err := s.rpcResults.Acquire(c.authKeyID, c.sessionID, reqMsgID) + if err != nil || acquired.state != rpcResultAcquirePending { + t.Fatalf("blocked delivery flight = %+v err=%v, want pending", acquired, err) + } + if task.release != nil { + task.release() + } + if claim.owner.Abort() { + t.Fatal("detached egress flight was aborted by inbound release") + } + tr.once.Do(func() { close(tr.release) }) + deadline := time.Now().Add(time.Second) + for { + if _, ok := s.rpcResults.Get(c.authKeyID, c.sessionID, reqMsgID); ok { + break + } + if time.Now().After(deadline) { + t.Fatal("delivered result was not published to replay cache") + } + time.Sleep(time.Millisecond) + } +} + +func TestDeliveryHookRunsOnceAfterReplayNotFailedWrite(t *testing.T) { + s := New(Options{}) + failing := &failAfterTransport{} + failing.failAt.Store(1) + oldConn := newOutboundTestConn(t, failing, newOutboundTrackedBudget(1<<20)) + const reqMsgID = int64(9101) + claim, err := s.rpcResults.Acquire(oldConn.authKeyID, oldConn.sessionID, reqMsgID) + if err != nil || claim.state != rpcResultAcquireOwner { + t.Fatalf("acquire flight = %+v err=%v", claim, err) + } + var hooks atomic.Int32 + if err := s.publishRPCResult(oldConn, reqMsgID, "updates.getDifference", claim.owner, + &tg.DataJSON{Data: "difference"}, func() { hooks.Add(1) }); err != nil { + // Admission succeeds; the asynchronous physical failure is observed below. + t.Fatalf("publish result: %v", err) + } + deadline := time.Now().Add(time.Second) + var cached *encodedOutboundMessage + for time.Now().Before(deadline) { + if got, ok := s.rpcResults.Get(oldConn.authKeyID, oldConn.sessionID, reqMsgID); ok { + cached = got + break + } + time.Sleep(time.Millisecond) + } + if cached == nil { + t.Fatal("failed write was not fenced and published for replay") + } + if got := cached.deliveryState(); got != rpcResultDeliveryReplayable { + t.Fatalf("failed delivery state = %d, want replayable", got) + } + if got := hooks.Load(); got != 0 { + t.Fatalf("delivery hooks after failed write = %d, want 0", got) + } + + replayTransport := &failAfterTransport{} + replayConn := newOutboundTestConn(t, replayTransport, newOutboundTrackedBudget(1<<20)) + if err := s.sendCachedRPCResult(context.Background(), replayConn, cached); err != nil { + t.Fatalf("replay result: %v", err) + } + deadline = time.Now().Add(time.Second) + for hooks.Load() != 1 && time.Now().Before(deadline) { + time.Sleep(time.Millisecond) + } + if got := hooks.Load(); got != 1 { + t.Fatalf("delivery hooks after replay = %d, want 1", got) + } + if got := cached.deliveryState(); got != rpcResultDeliveryDelivered { + t.Fatalf("replayed delivery state = %d, want delivered", got) + } + if err := s.sendCachedRPCResult(context.Background(), replayConn, cached); err != nil { + t.Fatalf("second replay: %v", err) + } + time.Sleep(20 * time.Millisecond) + if got := hooks.Load(); got != 1 { + t.Fatalf("delivery hooks after duplicate replay = %d, want 1", got) + } +} diff --git a/internal/mtprotoedge/rpc_rewrap.go b/internal/mtprotoedge/rpc_rewrap.go new file mode 100644 index 00000000..2b5d89c2 --- /dev/null +++ b/internal/mtprotoedge/rpc_rewrap.go @@ -0,0 +1,455 @@ +package mtprotoedge + +import ( + "context" + "crypto/sha256" + "sync" + "sync/atomic" + "time" + + "go.uber.org/zap" + + "github.com/gotd/td/bin" + "github.com/gotd/td/proto" + "github.com/gotd/td/tg" +) + +// rpcRewrapRegistry links only an explicit official-client transition: +// outstanding naked request -> invokeWithLayer(initConnection(the exact same +// request)). It is not a general content-dedup cache. Entries are hard-bounded +// and are retired by protocol events (client ACK, alias consumption, owner +// abort, or the first post-init naked request), never by client identity or by +// delaying request execution. +type rpcRewrapRegistry struct { + mu sync.Mutex + max int + total int + byKey map[rpcRewrapKey][]*rpcRewrapCandidate + bySession map[rpcRewrapSessionKey]map[*rpcRewrapCandidate]struct{} + byRequest map[rpcRewrapRequestKey]*rpcRewrapCandidate +} + +type rpcRewrapSessionKey struct { + authKeyID [8]byte + sessionID int64 +} + +type rpcRewrapKey struct { + rpcRewrapSessionKey + fingerprint [sha256.Size]byte +} + +type rpcRewrapRequestKey struct { + rpcRewrapSessionKey + reqMsgID int64 +} + +type rpcRewrapCandidate struct { + active bool + claimed bool + key rpcRewrapKey + source *Conn + reqMsgID int64 + method string + owner *rpcResultOwnerLease + waiter *rpcResultWaiter +} + +func newRPCRewrapRegistry(max int) *rpcRewrapRegistry { + if max <= 0 { + max = rpcResultFlightDefaultMaxPending + } + return &rpcRewrapRegistry{ + max: max, + byKey: make(map[rpcRewrapKey][]*rpcRewrapCandidate), + bySession: make(map[rpcRewrapSessionKey]map[*rpcRewrapCandidate]struct{}), + byRequest: make(map[rpcRewrapRequestKey]*rpcRewrapCandidate), + } +} + +func (r *rpcRewrapRegistry) register(c *Conn, body []byte, reqMsgID int64, method string, owner *rpcResultOwnerLease) bool { + if r == nil || c == nil || c.rpcRewrapInitialized.Load() || owner == nil { + return false + } + session := rpcRewrapSessionKey{authKeyID: c.authKeyID, sessionID: c.sessionID} + key := rpcRewrapKey{rpcRewrapSessionKey: session, fingerprint: sha256.Sum256(body)} + candidate := &rpcRewrapCandidate{ + active: true, key: key, source: c, reqMsgID: reqMsgID, method: method, + owner: owner, waiter: owner.Waiter(), + } + if candidate.waiter == nil { + return false + } + r.mu.Lock() + if r.total >= r.max { + r.mu.Unlock() + return false + } + r.byKey[key] = append(r.byKey[key], candidate) + set := r.bySession[session] + if set == nil { + set = make(map[*rpcRewrapCandidate]struct{}) + r.bySession[session] = set + } + set[candidate] = struct{}{} + r.byRequest[rpcRewrapRequestKey{rpcRewrapSessionKey: session, reqMsgID: reqMsgID}] = candidate + r.total++ + r.mu.Unlock() + if !owner.InstallAbortHook(func() { r.remove(candidate) }) { + r.remove(candidate) + return false + } + return true +} + +func (r *rpcRewrapRegistry) claim(c *Conn, inner []byte) *rpcRewrapCandidate { + if r == nil || c == nil { + return nil + } + session := rpcRewrapSessionKey{authKeyID: c.authKeyID, sessionID: c.sessionID} + key := rpcRewrapKey{rpcRewrapSessionKey: session, fingerprint: sha256.Sum256(inner)} + r.mu.Lock() + queue := r.byKey[key] + for _, candidate := range queue { + if !candidate.active || candidate.claimed { + continue + } + candidate.claimed = true + r.mu.Unlock() + return candidate + } + r.mu.Unlock() + return nil +} + +func (r *rpcRewrapRegistry) commit(candidate *rpcRewrapCandidate) { + if r == nil || candidate == nil { + return + } + r.mu.Lock() + r.removeLocked(candidate) + r.mu.Unlock() + candidate.owner.SetAbortHook(nil) +} + +func (r *rpcRewrapRegistry) release(candidate *rpcRewrapCandidate) { + if r == nil || candidate == nil { + return + } + r.mu.Lock() + if candidate.active { + candidate.claimed = false + } + r.mu.Unlock() +} + +func (r *rpcRewrapRegistry) remove(candidate *rpcRewrapCandidate) { + if r == nil || candidate == nil { + return + } + r.mu.Lock() + r.removeLocked(candidate) + r.mu.Unlock() +} + +// acknowledge retires a candidate only after the client explicitly ACKs the +// physical rpc_result. A successful socket write alone is insufficient proof: +// the client may already have reassigned the request to a new msg_id without +// parsing that old result. +func (r *rpcRewrapRegistry) acknowledge(c *Conn, reqMsgID int64) { + if r == nil || c == nil || reqMsgID == 0 { + return + } + request := rpcRewrapRequestKey{ + rpcRewrapSessionKey: rpcRewrapSessionKey{authKeyID: c.authKeyID, sessionID: c.sessionID}, + reqMsgID: reqMsgID, + } + r.mu.Lock() + candidate := r.byRequest[request] + r.removeLocked(candidate) + r.mu.Unlock() + if candidate != nil { + candidate.owner.SetAbortHook(nil) + } +} + +func (r *rpcRewrapRegistry) removeLocked(candidate *rpcRewrapCandidate) { + if candidate == nil || !candidate.active { + return + } + candidate.active = false + candidate.claimed = false + r.total-- + session := candidate.key.rpcRewrapSessionKey + delete(r.byRequest, rpcRewrapRequestKey{rpcRewrapSessionKey: session, reqMsgID: candidate.reqMsgID}) + if set := r.bySession[session]; set != nil { + delete(set, candidate) + if len(set) == 0 { + delete(r.bySession, session) + } + } + queue := r.byKey[candidate.key] + for i, existing := range queue { + if existing != candidate { + continue + } + copy(queue[i:], queue[i+1:]) + queue[len(queue)-1] = nil + queue = queue[:len(queue)-1] + break + } + if len(queue) == 0 { + delete(r.byKey, candidate.key) + } else { + r.byKey[candidate.key] = queue + } +} + +func (r *rpcRewrapRegistry) clearSession(c *Conn) { + if r == nil || c == nil { + return + } + session := rpcRewrapSessionKey{authKeyID: c.authKeyID, sessionID: c.sessionID} + r.mu.Lock() + set := r.bySession[session] + owners := make([]*rpcResultOwnerLease, 0, len(set)) + for candidate := range set { + owners = append(owners, candidate.owner) + r.removeLocked(candidate) + } + r.mu.Unlock() + for _, owner := range owners { + owner.SetAbortHook(nil) + } +} + +type rpcRewrapInit struct { + layer int + apiID int + deviceModel string + system string + appVersion string + systemLang string + langPack string + langCode string + inner []byte +} + +type rpcRewrapRawObject struct { + data []byte +} + +func (o *rpcRewrapRawObject) Decode(b *bin.Buffer) error { + if _, err := b.PeekID(); err != nil { + return err + } + o.data = b.Buf + b.Skip(len(b.Buf)) + return nil +} + +func (o *rpcRewrapRawObject) Encode(b *bin.Buffer) error { + b.Put(o.data) + return nil +} + +func decodeRPCRewrapInit(body []byte) (rpcRewrapInit, bool) { + b := &bin.Buffer{Buf: body} + if err := b.ConsumeID(tg.InvokeWithLayerRequestTypeID); err != nil { + return rpcRewrapInit{}, false + } + layer, err := b.Int() + if err != nil || layer <= 0 { + return rpcRewrapInit{}, false + } + raw := &rpcRewrapRawObject{} + req := tg.InitConnectionRequest{Query: raw} + if err := req.Decode(b); err != nil || b.Len() != 0 || len(raw.data) < bin.Word { + return rpcRewrapInit{}, false + } + return rpcRewrapInit{ + layer: layer, apiID: req.APIID, deviceModel: req.DeviceModel, + system: req.SystemVersion, appVersion: req.AppVersion, + systemLang: req.SystemLangCode, langPack: req.LangPack, langCode: req.LangCode, + inner: raw.data, + }, true +} + +type rpcRewrapAlias struct { + conn *Conn + newReqID int64 + method string + oldWaiter *rpcResultWaiter + newOwner *rpcResultOwnerLease + sourceConn *Conn + sourceOwner *rpcResultOwnerLease + retargeted atomic.Bool + observeInit bool + init rpcRewrapInit + candidate *rpcRewrapCandidate + registry *rpcRewrapRegistry +} + +var ( + rpcRewrapDeliveryOnce sync.Once + rpcRewrapDeliveryJobs chan func() +) + +const ( + rpcRewrapDeliveryWorkers = 4 + rpcRewrapDeliveryQueue = 256 +) + +func scheduleRPCRewrapDelivery(fn func()) bool { + if fn == nil { + return false + } + rpcRewrapDeliveryOnce.Do(func() { + rpcRewrapDeliveryJobs = make(chan func(), rpcRewrapDeliveryQueue) + for range rpcRewrapDeliveryWorkers { + go func() { + for job := range rpcRewrapDeliveryJobs { + job() + } + }() + } + }) + select { + case rpcRewrapDeliveryJobs <- fn: + return true + default: + return false + } +} + +func (a *rpcRewrapAlias) activate(s *Server) error { + if a == nil || s == nil || a.conn == nil || a.oldWaiter == nil { + return ErrRPCResultFlightInvalid + } + err := a.oldWaiter.Subscribe(func(encoded *encodedOutboundMessage, ok bool) { + if !ok || encoded == nil { + if a.newOwner != nil { + a.newOwner.Abort() + } + a.conn.fenceUndeliveredRPCResult() + return + } + if a.newOwner == nil { + if !scheduleRPCRewrapDelivery(func() { + ctx, cancel := context.WithTimeout(context.Background(), min(5*time.Second, max(time.Second, a.conn.writeTimeout))) + defer cancel() + if err := s.sendCachedRPCResult(ctx, a.conn, encoded); err != nil && !isClientDisconnect(err) { + s.log.Debug("RPC init rewrap pending replay failed", zap.Error(err)) + } + }) { + a.conn.fenceUndeliveredRPCResult() + } + return + } + clone, err := cloneRPCResultForRequest(encoded, a.newReqID, false) + if err != nil { + a.newOwner.Abort() + a.conn.fenceUndeliveredRPCResult() + return + } + if a.retargeted.Load() { + if !a.newOwner.HandOff() { + a.conn.fenceUndeliveredRPCResult() + return + } + clone.markDelivered() + s.storeRPCResult(a.conn, a.newReqID, clone) + s.log.Info("RPC init rewrap result retargeted", + zap.String("method", a.method), zap.Int64("new_req_msg_id", a.newReqID), + zap.String("auth_key_id", a.conn.authKeyHex), zap.Int64("session_id", a.conn.sessionID)) + return + } + if !scheduleRPCRewrapDelivery(func() { + s.publishRewrappedRPCResult(a.conn, a.newReqID, a.method, a.newOwner, clone) + }) { + // The completed result is durable in memory. Fence before publishing it + // under the new msg_id so a replacement can replay without re-executing. + a.conn.fenceUndeliveredRPCResult() + if a.newOwner.HandOff() { + clone.markReplayable() + s.storeRPCResult(a.conn, a.newReqID, clone) + } + } + }) + if err != nil { + s.rpcRewrap.release(a.candidate) + return err + } + // Subscribe first so every terminal owner event has a consumer. If completion + // wins this race the callback replays under the new ID; if retarget wins, the + // sole outbound actor snapshots the new ID before writing. + if a.newOwner != nil && a.sourceConn == a.conn && a.sourceOwner != nil { + a.retargeted.Store(a.sourceOwner.TryRetarget(a.newReqID)) + } + if a.observeInit { + s.scheduleRewrappedInitObservation(a.conn, a.init) + } + s.rpcRewrap.commit(a.candidate) + a.candidate = nil + return nil +} + +func (a *rpcRewrapAlias) releaseCandidate() { + if a == nil || a.candidate == nil { + return + } + a.registry.release(a.candidate) + a.candidate = nil +} + +func (s *Server) scheduleRewrappedInitObservation(c *Conn, init rpcRewrapInit) { + observer, ok := s.rpc.(RPCInitConnectionObserver) + if !ok || c == nil { + return + } + if !scheduleRPCRewrapDelivery(func() { + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + if err := observer.ObserveInitConnection( + ctx, c.authKeyID, c.sessionID, init.layer, init.apiID, + init.deviceModel, init.system, init.appVersion, init.systemLang, + init.langPack, init.langCode, + ); err != nil { + s.log.Debug("Observe rewrapped initConnection failed", zap.Error(err)) + } + }) { + s.log.Debug("Observe rewrapped initConnection dropped", zap.String("auth_key_id", c.authKeyHex)) + } +} + +func (s *Server) publishRewrappedRPCResult(c *Conn, reqMsgID int64, method string, owner *rpcResultOwnerLease, encoded *encodedOutboundMessage) { + if s == nil || c == nil || owner == nil || encoded == nil { + return + } + if !owner.HandOff() { + c.fenceUndeliveredRPCResult() + return + } + priority := rpcResultPriority(method, encoded) + encoded.priority = priority + terminal := func(deliveryErr error) { + if deliveryErr != nil { + encoded.markReplayable() + c.fenceUndeliveredRPCResult() + } else { + encoded.markDelivered() + } + s.storeRPCResult(c, reqMsgID, encoded) + } + encoded.markQueued() + ctx, cancel := context.WithTimeout(context.Background(), min(5*time.Second, max(time.Second, c.writeTimeout))) + defer cancel() + if err := c.enqueueEncodedDelivery(ctx, proto.MessageServerResponse, encoded, priority, terminal); err != nil { + terminal(err) + return + } + s.log.Info("RPC init rewrap result replay admitted", + zap.String("method", method), zap.Int64("req_msg_id", reqMsgID), + zap.String("auth_key_id", c.authKeyHex), zap.Int64("session_id", c.sessionID), + zap.Int("wire_bytes", len(encoded.body))) +} diff --git a/internal/mtprotoedge/rpc_rewrap_test.go b/internal/mtprotoedge/rpc_rewrap_test.go new file mode 100644 index 00000000..7582a478 --- /dev/null +++ b/internal/mtprotoedge/rpc_rewrap_test.go @@ -0,0 +1,299 @@ +package mtprotoedge + +import ( + "context" + "encoding/binary" + "sync/atomic" + "testing" + "time" + + "github.com/gotd/td/bin" + "github.com/gotd/td/crypto" + "github.com/gotd/td/mt" + "github.com/gotd/td/tg" +) + +func encodeRewrapTestRequest(t *testing.T) ([]byte, []byte) { + t.Helper() + var inner bin.Buffer + if err := (&tg.HelpGetConfigRequest{}).Encode(&inner); err != nil { + t.Fatalf("encode inner request: %v", err) + } + var wrapped bin.Buffer + if err := (&tg.InvokeWithLayerRequest{ + Layer: 227, + Query: &tg.InitConnectionRequest{ + APIID: 6, DeviceModel: "Pixel", SystemVersion: "SDK 36", + AppVersion: "12.8.1", SystemLangCode: "en", LangPack: "android", LangCode: "en", + Query: &tg.HelpGetConfigRequest{}, + }, + }).Encode(&wrapped); err != nil { + t.Fatalf("encode wrapped request: %v", err) + } + return append([]byte(nil), inner.Raw()...), append([]byte(nil), wrapped.Raw()...) +} + +func TestDecodeRPCRewrapInitExtractsExactInnerQuery(t *testing.T) { + inner, wrapped := encodeRewrapTestRequest(t) + init, ok := decodeRPCRewrapInit(wrapped) + if !ok { + t.Fatal("valid invokeWithLayer(initConnection(query)) was not recognized") + } + if init.layer != 227 || init.apiID != 6 || init.langPack != "android" { + t.Fatalf("metadata = layer:%d api:%d lang_pack:%q", init.layer, init.apiID, init.langPack) + } + if string(init.inner) != string(inner) { + t.Fatalf("inner = %x, want %x", init.inner, inner) + } + if _, ok := decodeRPCRewrapInit(inner); ok { + t.Fatal("naked request must not be classified as an init rewrap") + } +} + +func TestRPCResultDeliveryRetargetHasExactWritingBarrier(t *testing.T) { + const oldReqID, newReqID, tooLateReqID = int64(101), int64(202), int64(303) + encoded := &encodedOutboundMessage{ + typeID: mt.RPCResultTypeID, reqMsgID: oldReqID, + body: make([]byte, 16), delivery: newRPCResultDelivery(oldReqID), + } + if !encoded.tryRetarget(newReqID) { + t.Fatal("prepared result should be retargetable") + } + if got := encoded.beginWriting(); got != newReqID { + t.Fatalf("writing target = %d, want %d", got, newReqID) + } + if encoded.tryRetarget(tooLateReqID) { + t.Fatal("writing result must not be mutated") + } +} + +func TestRPCResultWaiterSubscribeIsEventDriven(t *testing.T) { + cache := newRPCResultCacheWithFlightLimit(time.Now, 8) + claim, err := cache.Acquire([8]byte{1}, 2, 3) + if err != nil || claim.state != rpcResultAcquireOwner { + t.Fatalf("Acquire owner = %+v, %v", claim, err) + } + var called atomic.Bool + if err := claim.owner.Waiter().Subscribe(func(encoded *encodedOutboundMessage, ok bool) { + if !ok || encoded == nil { + t.Errorf("subscriber result = %#v, %v", encoded, ok) + } + called.Store(true) + }); err != nil { + t.Fatalf("Subscribe: %v", err) + } + if called.Load() { + t.Fatal("Subscribe waited for or fabricated a result") + } + encoded := &encodedOutboundMessage{typeID: mt.RPCResultTypeID, body: make([]byte, 16), reqMsgID: 3} + cache.Put([8]byte{1}, 2, 3, encoded) + if !called.Load() { + t.Fatal("completion event did not invoke subscriber") + } +} + +func TestRPCResultOwnerAbortHookInstallationIsFlightBound(t *testing.T) { + cache := newRPCResultCacheWithFlightLimit(time.Now, 8) + claim, err := cache.Acquire([8]byte{2}, 3, 4) + if err != nil || claim.state != rpcResultAcquireOwner { + t.Fatalf("Acquire owner = %+v, %v", claim, err) + } + var called atomic.Bool + if !claim.owner.InstallAbortHook(func() { called.Store(true) }) { + t.Fatal("live owner rejected abort hook") + } + if !claim.owner.Abort() || !called.Load() { + t.Fatal("owner abort did not invoke installed hook") + } + if claim.owner.InstallAbortHook(func() {}) { + t.Fatal("completed flight accepted a new abort hook") + } +} + +func TestRPCRewrapRegistryIsPlatformAgnosticAndAckBound(t *testing.T) { + cache := newRPCResultCacheWithFlightLimit(time.Now, 8) + claim, err := cache.Acquire([8]byte{4}, 5, 6) + if err != nil || claim.state != rpcResultAcquireOwner { + t.Fatalf("Acquire owner = %+v, %v", claim, err) + } + c := &Conn{authKeyID: [8]byte{4}, sessionID: 5} + r := newRPCRewrapRegistry(8) + body := []byte{1, 2, 3, 4} + if !r.register(c, body, 6, "test.method", claim.owner) { + t.Fatal("protocol candidate was rejected without platform metadata") + } + candidate := r.claim(c, body) + if candidate == nil { + t.Fatal("exact protocol fingerprint did not match") + } + if got := r.claim(c, body); got != nil { + t.Fatal("one candidate was claimed by two rewrapped requests") + } + r.release(candidate) + if got := r.claim(&Conn{authKeyID: [8]byte{7}, sessionID: 5}, body); got != nil { + t.Fatal("fingerprint crossed the auth-key boundary") + } + if got := r.claim(c, []byte{4, 3, 2, 1}); got != nil { + t.Fatal("mismatched TL bytes claimed a candidate") + } + if got := r.claim(c, body); got != candidate { + t.Fatal("released candidate was lost from the fingerprint index") + } + r.release(candidate) + if r.total != 1 { + t.Fatalf("released claim retired candidate: total=%d", r.total) + } + r.acknowledge(c, 6) + if r.total != 0 || len(r.byKey) != 0 || len(r.bySession) != 0 || len(r.byRequest) != 0 { + t.Fatalf("ACK did not retire every candidate index: total=%d key=%d session=%d request=%d", + r.total, len(r.byKey), len(r.bySession), len(r.byRequest)) + } +} + +func TestInitRewrapAfterWritingReplaysWithoutBusinessExecution(t *testing.T) { + inner, wrapped := encodeRewrapTestRequest(t) + s := New(Options{RPCGlobalWorkers: 1, RPCGlobalMaxTasks: 16, RPCGlobalMaxBytes: 1 << 20}) + transport := &collectingSessionTransport{} + key := newTestAuthKey(t) + c := s.newConn(transport, key, 88, 99) + defer c.ForceClose() + + const oldReqID, newReqID = int64(3001), int64(4001) + oldPlan := &inboundPlan{items: []inboundItem{{ + kind: inboundItemRPC, msgID: oldReqID, typeID: tg.HelpGetConfigRequestTypeID, body: inner, + }}} + defer oldPlan.close() + if err := s.prepareInboundRPCBatch(context.Background(), c, oldPlan); err != nil { + t.Fatalf("prepare old request: %v", err) + } + oldOwner := oldPlan.rpcOwners[0] + if got := (&encodedOutboundMessage{delivery: oldOwner.Delivery()}).beginWriting(); got != oldReqID { + t.Fatalf("old writing target = %d, want %d", got, oldReqID) + } + + newPlan := &inboundPlan{items: []inboundItem{{ + kind: inboundItemRPC, msgID: newReqID, typeID: tg.InvokeWithLayerRequestTypeID, body: wrapped, + }}} + defer newPlan.close() + if err := s.prepareInboundRPCBatch(context.Background(), c, newPlan); err != nil { + t.Fatalf("prepare rewrapped request: %v", err) + } + if len(newPlan.rpcTasks) != 0 || len(newPlan.rewrapAliases) != 1 { + t.Fatalf("late rewrap dispatched business: tasks=%d aliases=%d", len(newPlan.rpcTasks), len(newPlan.rewrapAliases)) + } + if err := newPlan.commitRewrapAliases(s); err != nil { + t.Fatalf("activate late alias: %v", err) + } + + encoded, err := s.encodeRPCResultContext(context.Background(), c, oldReqID, &mt.RPCError{ + ErrorCode: 400, ErrorMessage: "TEST", + }) + if err != nil { + t.Fatalf("encode old result: %v", err) + } + encoded.delivery = oldOwner.Delivery() + if !oldOwner.HandOff() { + t.Fatal("old owner handoff failed") + } + s.rpcResults.Put(c.authKeyID, c.sessionID, oldReqID, encoded) + + deadline := time.Now().Add(2 * time.Second) + var replayed *encodedOutboundMessage + for time.Now().Before(deadline) { + if got, ok := s.rpcResults.Get(c.authKeyID, c.sessionID, newReqID); ok { + replayed = got + break + } + time.Sleep(time.Millisecond) + } + if replayed == nil { + t.Fatal("late alias did not publish a result under the new request ID") + } + if got := int64(binary.LittleEndian.Uint64(replayed.body[4:12])); got != newReqID { + t.Fatalf("replayed req_msg_id = %d, want %d", got, newReqID) + } + if len(transport.snapshot()) != 1 { + t.Fatalf("late alias physical result count = %d, want 1", len(transport.snapshot())) + } +} + +func TestOutboundAckReturnsRPCRequestIDs(t *testing.T) { + state := outboundState{ + pending: map[int64]*outboundFrame{ + 10: {msgID: 10, reqMsgID: 101}, + 20: {msgID: 20}, + }, + byRequest: map[int64]int64{101: 10}, + maxMessages: 8, + } + got := state.ack([]int64{10, 20, 30}) + if len(got) != 1 || got[0] != 101 { + t.Fatalf("acked request IDs = %v, want [101]", got) + } +} + +func TestInitRewrapAliasesExecutionAndRetargetsQueuedResult(t *testing.T) { + inner, wrapped := encodeRewrapTestRequest(t) + s := New(Options{RPCGlobalWorkers: 1, RPCGlobalMaxTasks: 16, RPCGlobalMaxBytes: 1 << 20}) + c := &Conn{ + metrics: NopMetrics{}, authKeyID: [8]byte{9}, authKeyHex: "09", + sessionID: 77, key: crypto.AuthKey{ID: [8]byte{9}}, + } + c.startInboundRPCScheduler(s.rpcScheduler, 1, 8, time.Second) + defer c.closeInboundRPCScheduler() + + const oldReqID, newReqID = int64(1001), int64(2001) + oldPlan := &inboundPlan{items: []inboundItem{{ + kind: inboundItemRPC, msgID: oldReqID, typeID: tg.HelpGetConfigRequestTypeID, body: inner, + }}} + defer oldPlan.close() + if err := s.prepareInboundRPCBatch(context.Background(), c, oldPlan); err != nil { + t.Fatalf("prepare old request: %v", err) + } + if len(oldPlan.rpcTasks) != 1 || len(oldPlan.rpcOwners) != 1 { + t.Fatalf("old admission tasks=%d owners=%d", len(oldPlan.rpcTasks), len(oldPlan.rpcOwners)) + } + oldOwner := oldPlan.rpcOwners[0] + + newPlan := &inboundPlan{items: []inboundItem{{ + kind: inboundItemRPC, msgID: newReqID, typeID: tg.InvokeWithLayerRequestTypeID, body: wrapped, + }}} + defer newPlan.close() + if err := s.prepareInboundRPCBatch(context.Background(), c, newPlan); err != nil { + t.Fatalf("prepare rewrapped request: %v", err) + } + if len(newPlan.rpcTasks) != 0 || len(newPlan.rewrapAliases) != 1 || newPlan.items[0].kind != inboundItemRewrappedRPC { + t.Fatalf("rewrap admission tasks=%d aliases=%d kind=%d", len(newPlan.rpcTasks), len(newPlan.rewrapAliases), newPlan.items[0].kind) + } + if err := newPlan.commitRewrapAliases(s); err != nil { + t.Fatalf("activate alias: %v", err) + } + + encoded, err := s.encodeRPCResultContext(context.Background(), c, oldReqID, &mt.RPCError{ + ErrorCode: 400, ErrorMessage: "TEST", + }) + if err != nil { + t.Fatalf("encode old result: %v", err) + } + encoded.delivery = oldOwner.Delivery() + encoded.markQueued() + if got := encoded.beginWriting(); got != newReqID { + t.Fatalf("physical result target = %d, want new req %d", got, newReqID) + } + if !oldOwner.HandOff() { + t.Fatal("old owner handoff failed") + } + encoded.markDelivered() + s.rpcResults.Put(c.authKeyID, c.sessionID, oldReqID, encoded) + + aliased, ok := s.rpcResults.Get(c.authKeyID, c.sessionID, newReqID) + if !ok { + t.Fatal("new req_msg_id result was not completed") + } + if got := int64(binary.LittleEndian.Uint64(aliased.body[4:12])); got != newReqID { + t.Fatalf("cached aliased req_msg_id = %d, want %d", got, newReqID) + } + if s.rpcRewrap.total != 0 { + t.Fatalf("rewrap registry retained %d consumed candidates", s.rpcRewrap.total) + } +} diff --git a/internal/mtprotoedge/same_port_mux.go b/internal/mtprotoedge/same_port_mux.go index 72eff5b3..4862ebf4 100644 --- a/internal/mtprotoedge/same_port_mux.go +++ b/internal/mtprotoedge/same_port_mux.go @@ -34,6 +34,7 @@ type samePortMux struct { base net.Listener addr net.Addr sniffTimeout time.Duration + observe func(connectionIntakeEvent) tcp *samePortMuxListener http *samePortMuxListener @@ -49,7 +50,7 @@ type samePortMux struct { sniffing map[net.Conn]struct{} } -func newSamePortMux(base net.Listener, sniffTimeout time.Duration) *samePortMux { +func newSamePortMux(base net.Listener, sniffTimeout time.Duration, observers ...func(connectionIntakeEvent)) *samePortMux { if sniffTimeout <= 0 { sniffTimeout = 5 * time.Second } @@ -60,6 +61,9 @@ func newSamePortMux(base net.Listener, sniffTimeout time.Duration) *samePortMux closed: make(chan struct{}), sniffing: make(map[net.Conn]struct{}), } + if len(observers) > 0 { + m.observe = observers[0] + } m.tcp = newSamePortMuxListener(m.addr, m.closed) m.http = newSamePortMuxListener(m.addr, m.closed) return m @@ -144,10 +148,13 @@ func (m *samePortMux) Close() error { // dispatch 窥探单条连接的前 4 字节并把它交给 tcp 或 http 子 listener。窥探带 sniffTimeout // 读上界，慢/半开连接最多占用本 goroutine sniffTimeout 后即被回收。 func (m *samePortMux) dispatch(ctx context.Context, conn net.Conn) { + started := time.Now() + remote, local := connRemote(conn), connLocal(conn) // SetReadDeadline bounds an otherwise healthy slow-loris connection, but Close only owns // the base listener, not sockets Accept has already returned. Register temporary ownership so // mux shutdown can close this read immediately. finishSniff removes the socket before hand-off. if !m.beginSniff(conn) { + m.observeEvent(connectionIntakeEvent{stage: "mux_sniff", outcome: "closed", remote: remote, local: local, duration: time.Since(started)}) _ = conn.Close() return } @@ -160,10 +167,16 @@ func (m *samePortMux) dispatch(ctx context.Context, conn net.Conn) { var header [4]byte if err := conn.SetReadDeadline(time.Now().Add(m.sniffTimeout)); err != nil { + m.observeEvent(connectionIntakeEvent{stage: "mux_sniff", outcome: "error", remote: remote, local: local, duration: time.Since(started), err: err}) _ = conn.Close() return } if _, err := io.ReadFull(conn, header[:]); err != nil { + outcome := "error" + if isClientDisconnect(err) { + outcome = "client_disconnect" + } + m.observeEvent(connectionIntakeEvent{stage: "mux_sniff", outcome: outcome, remote: remote, local: local, duration: time.Since(started), err: err}) _ = conn.Close() return } @@ -171,11 +184,13 @@ func (m *samePortMux) dispatch(ctx context.Context, conn net.Conn) { // registry entry under sniffMu is the hand-off barrier: Close either captured and closed this // socket, or it can no longer find it. A concurrently closed mux refuses delivery. if !m.finishSniff(conn) { + m.observeEvent(connectionIntakeEvent{stage: "mux_sniff", outcome: "closed", remote: remote, local: local, duration: time.Since(started)}) _ = conn.Close() return } finishedSniff = true if err := conn.SetReadDeadline(time.Time{}); err != nil { + m.observeEvent(connectionIntakeEvent{stage: "mux_sniff", outcome: "error", remote: remote, local: local, duration: time.Since(started), err: err}) _ = conn.Close() return } @@ -186,14 +201,24 @@ func (m *samePortMux) dispatch(ctx context.Context, conn net.Conn) { } target := m.tcp + transport := "tcp" if isHTTPHeaderPrefix(header) { target = m.http + transport = "websocket" } + m.observeEvent(connectionIntakeEvent{stage: "mux_sniff", outcome: "ready", transport: transport, remote: remote, local: local, duration: time.Since(started), bytes: len(header)}) if !target.deliver(ctx, wrapped) { + m.observeEvent(connectionIntakeEvent{stage: "mux_delivery", outcome: "closed", transport: transport, remote: remote, local: local, duration: time.Since(started)}) _ = conn.Close() } } +func (m *samePortMux) observeEvent(event connectionIntakeEvent) { + if m.observe != nil { + m.observe(event) + } +} + func (m *samePortMux) beginSniff(conn net.Conn) bool { m.sniffMu.Lock() defer m.sniffMu.Unlock() diff --git a/internal/mtprotoedge/server.go b/internal/mtprotoedge/server.go index 718b3c08..7c057382 100644 --- a/internal/mtprotoedge/server.go +++ b/internal/mtprotoedge/server.go @@ -43,6 +43,27 @@ type RPCHandler interface { NegotiatedLayer(authKeyID [8]byte, sessionID int64) (int, bool) } +// RPCHandlerWithMethod returns the canonical innermost RPC method after the +// router has peeled invokeWithLayer/initConnection/invokeAfter wrappers. Egress +// scheduling must use this identity: the outer wrapper is not a useful signal +// for prioritizing updates convergence over catalog/media responses. +type RPCHandlerWithMethod interface { + DispatchWithMethod(ctx context.Context, authKeyID [8]byte, sessionID int64, b *bin.Buffer) (bin.Encoder, string, error) +} + +// RPCInitConnectionObserver records wrapper metadata when the edge aliases an +// initConnection reissue to an already-running request and therefore correctly +// skips a second business Dispatch. +type RPCInitConnectionObserver interface { + ObserveInitConnection( + ctx context.Context, + authKeyID [8]byte, + sessionID int64, + layer, apiID int, + deviceModel, systemVersion, appVersion, systemLangCode, langPack, langCode string, + ) error +} + // Options 配置 Server。 type Options struct { // Logger 日志器。默认 zap.NewNop()。 @@ -126,6 +147,10 @@ type Options struct { RPC RPCHandler // Metrics 接收连接层指标。默认 NopMetrics。 Metrics Metrics + // OnServing is called after the connection intake loops have been installed. + // It is a platform-neutral observation hook; all slow initialization must + // finish before ListenAndServe is entered. + OnServing func(net.Addr) // Clock 用于消息 ID 与时间戳。默认 clock.System。 Clock clock.Clock // Rand 随机源。默认 crypto.DefaultRand()。 @@ -239,6 +264,7 @@ type Server struct { conns *SessionManager rpc RPCHandler metrics Metrics + onServing func(net.Addr) cipher crypto.Cipher clock clock.Clock rand io.Reader @@ -246,6 +272,7 @@ type Server struct { admission *admissionController rpcResults *rpcResultCache + rpcRewrap *rpcRewrapRegistry // onFrame 是测试钩子：收到一帧时回调其字节数；生产为 nil。 onFrame func(n int) @@ -284,15 +311,30 @@ func New(opts Options) *Server { conns: conns, rpc: opts.RPC, metrics: opts.Metrics, + onServing: opts.OnServing, cipher: crypto.NewServerCipher(opts.Rand), clock: opts.Clock, rand: opts.Rand, types: tmap.New(tg.TypesMap(), mt.TypesMap(), proto.TypesMap()), rpcResults: newRPCResultCacheWithFlightLimit(opts.Clock.Now, opts.RPCGlobalMaxTasks), + rpcRewrap: newRPCRewrapRegistry(opts.RPCGlobalMaxTasks), admission: newAdmissionController(opts.MaxConnections, opts.MaxConnectionsPerIP, opts.MaxConcurrentHandshakes), } } +// ListenAndServe binds the public MTProto socket and immediately enters Serve. +// Keeping listener ownership at the connection edge prevents callers from +// exposing a TCP port and then performing slow seed/cache initialization while +// clients are already completing handshakes into an unread accept backlog. +func (s *Server) ListenAndServe(ctx context.Context, addr string) error { + var lc net.ListenConfig + ln, err := lc.Listen(ctx, "tcp", addr) + if err != nil { + return fmt.Errorf("listen %q: %w", addr, err) + } + return s.Serve(ctx, ln) +} + // newConn 基于一次解密结果创建一个可发送的连接对象。 func (s *Server) newConn(tc transport.Conn, key crypto.AuthKey, sessionID, salt int64) *Conn { if lease, ok := tc.(*physicalTransportLease); ok { @@ -338,6 +380,7 @@ func (s *Server) buildConn(tc transport.Conn, lease *physicalTransportLease, key outboundTrackedBudget: s.outboundTrackedBudget, outboundControlTrackedBudget: s.outboundControlBudget, outboundScratchPool: s.outboundScratchPool, + rpcResultAcked: s.rpcRewrap.acknowledge, } c.startOutbound() c.startInboundRPCScheduler(s.rpcScheduler, s.rpcInflight, s.rpcQueueSize, s.rpcTimeout) @@ -354,7 +397,7 @@ func (s *Server) Serve(ctx context.Context, ln net.Listener) error { defer s.rpcScheduler.stop(rpcCloseWaitTimeout) // 只在最外层 listener 包一次，确保 same-port mux 的 sniff/HTTP upgrade 也计入 // raw admission，而不是等连接已经分流后才计数。 - ln = s.admission.wrapListener(ln) + ln = s.observeRawAccepts(s.admission.wrapListener(ln)) if s.websocket { return s.serveMixed(ctx, ln) } @@ -367,8 +410,12 @@ func (s *Server) serveTCP(ctx context.Context, ln net.Listener) error { s.log.Info("Serving", zap.String("addr", ln.Addr().String()), zap.Int("dc", s.dc), zap.Bool("obfuscated_tcp", s.obfuscated)) defer s.log.Info("Stopped") - - return s.acceptLoop(ctx, ln, s.obfuscated) + errCh := make(chan error, 1) + go func() { + errCh <- s.acceptLoop(ctx, ln, s.obfuscated) + }() + s.signalServing(ln.Addr()) + return <-errCh } func (s *Server) serveMixed(ctx context.Context, ln net.Listener) error { @@ -382,7 +429,7 @@ func (s *Server) serveMixed(ctx context.Context, ln net.Listener) error { // minDuration(5s,...) 把嗅探压到 5s，比非 mux 路径激进 12 倍，会把这些暖连接在 5s 误杀， // 触发客户端 6s 重连风暴并误判「后端不健康」回退到外部 DNS。per-conn goroutine 模型已消解 // slow-loris 接入饥饿，故嗅探用满 handshakeTimeout 是安全的。 - mux := newSamePortMux(ln, s.handshakeTimeout) + mux := newSamePortMux(ln, s.handshakeTimeout, s.observeConnectionIntake) wsRawLn, wsHandler := transport.WebsocketListener(ln.Addr()) wsLn := newTransportPacketMessageListener(wsRawLn) @@ -430,7 +477,7 @@ func (s *Server) serveMixed(ctx context.Context, ln net.Listener) error { // WebSocket：gotd 升级处理器已剥离 obfuscated2 并补回 codec tag，这里只需探测 codec。 go func() { defer wg.Done() - errCh <- s.acceptLoop(ctx, wsLn, false) + errCh <- s.acceptLoopTransport(ctx, wsLn, false, "websocket") }() go func() { defer wg.Done() @@ -444,6 +491,7 @@ func (s *Server) serveMixed(ctx context.Context, ln net.Listener) error { } errCh <- nil }() + s.signalServing(ln.Addr()) // The four services form one lifecycle: even a clean/closed-listener return from any one // component means the remaining three can no longer make forward progress as a complete @@ -468,6 +516,10 @@ func (s *Server) serveMixed(ctx context.Context, ln net.Listener) error { // 整个监听循环。obfuscated 为 true 时先走 obfuscated2 去混淆（裸 MTProto TCP）；WebSocket // 连接传 false（gotd 升级处理器已完成去混淆）。 func (s *Server) acceptLoop(ctx context.Context, ln net.Listener, obfuscated bool) error { + return s.acceptLoopTransport(ctx, ln, obfuscated, intakeTransport(obfuscated)) +} + +func (s *Server) acceptLoopTransport(ctx context.Context, ln net.Listener, obfuscated bool, transportName string) error { ctx, cancel := context.WithCancel(ctx) var wg sync.WaitGroup defer func() { @@ -505,7 +557,14 @@ func (s *Server) acceptLoop(ctx context.Context, ln net.Listener, obfuscated boo wg.Add(1) go func() { defer wg.Done() - s.serveDetectedConn(ctx, raw, obfuscated) + s.observeConnectionIntake(connectionIntakeEvent{ + stage: "transport_dispatch", + outcome: "accepted", + transport: transportName, + remote: connRemote(raw), + local: connLocal(raw), + }) + s.serveDetectedConn(ctx, raw, obfuscated, transportName) }() } } @@ -513,7 +572,9 @@ func (s *Server) acceptLoop(ctx context.Context, ln net.Listener, obfuscated boo // serveDetectedConn 把一条裸连接提升为 transport.Conn（去混淆 + codec 探测）后运行 MTProto // 连接循环。提升过程的读取放在本 goroutine、且受握手读超时约束，而非塞在 accept 循环里， // 这样慢连接不会阻塞其他连接接入，去混淆/codec 握手本身也有时间上界。 -func (s *Server) serveDetectedConn(ctx context.Context, raw net.Conn, obfuscated bool) { +func (s *Server) serveDetectedConn(ctx context.Context, raw net.Conn, obfuscated bool, transportName string) { + started := time.Now() + remote, local := connRemote(raw), connLocal(raw) // 握手读超时只覆盖去混淆 + codec 探测这一小段；用真实墙钟时间（SetReadDeadline 语义）， // 不走可能被测试注入的逻辑 clock。 if err := raw.SetReadDeadline(time.Now().Add(s.handshakeTimeout)); err != nil { @@ -536,21 +597,37 @@ func (s *Server) serveDetectedConn(ctx context.Context, raw net.Conn, obfuscated conn, err := s.promoteConn(raw, obfuscated) close(promoted) if err != nil { - // 去混淆/codec 探测失败（读超时、客户端中途断开、坏 init 等）只影响这一条连接， - // 记 debug 即可。 - if !isClientDisconnect(err) { - s.log.Debug("Transport handshake failed", zap.Error(err)) + outcome := "error" + if isClientDisconnect(err) { + outcome = "client_disconnect" } + s.observeConnectionIntake(connectionIntakeEvent{ + stage: "transport_promote", + outcome: outcome, + transport: transportName, + remote: remote, + local: local, + duration: time.Since(started), + err: err, + }) _ = raw.Close() return } + s.observeConnectionIntake(connectionIntakeEvent{ + stage: "transport_promote", + outcome: "ready", + transport: transportName, + remote: remote, + local: local, + duration: time.Since(started), + }) // 探测完成，撤掉握手读超时；后续每帧读写由 serveConn / 传输层各自管理超时。 if err := raw.SetReadDeadline(time.Time{}); err != nil { _ = conn.Close() return } - if err := s.serveConn(ctx, conn); err != nil && !isClientDisconnect(err) { - s.log.Info("Connection closed with error", zap.Error(err)) + if err := s.serveConn(ctx, conn, remote, local); err != nil && !isClientDisconnect(err) { + s.log.Info("Connection closed with error", zap.String("remote_addr", remote), zap.String("local_addr", local), zap.Error(err)) } } @@ -572,12 +649,14 @@ func (s *Server) promoteConn(raw net.Conn, obfuscated bool) (transport.Conn, err // - auth_key_id 未注册：回 AuthKeyNotFound，促使客户端重新握手。 // // 连接建立 session 后注册到 SessionManager，结束时注销。 -func (s *Server) serveConn(ctx context.Context, raw transport.Conn) (err error) { +func (s *Server) serveConn(ctx context.Context, raw transport.Conn, remote, local string) (err error) { transportOwner, conn := newPhysicalTransportOwner(raw) s.metrics.ConnOpened() - s.log.Debug("Connection accepted") + s.log.Debug("MTProto connection loop started", zap.String("remote_addr", remote), zap.String("local_addr", local)) var current *Conn + firstFrameStarted := time.Now() + firstFrameSeen := false defer func() { // A successful Recv transfers the frame reservation to serveConn. Release it only after // this stack has stopped using b/plain; transport.Close may have raced us earlier and must @@ -595,7 +674,7 @@ func (s *Server) serveConn(ctx context.Context, raw transport.Conn) (err error) current.Close() } s.metrics.ConnClosed() - s.log.Debug("Connection closed", zap.Error(err)) + s.log.Debug("Connection closed", zap.String("remote_addr", remote), zap.String("local_addr", local), zap.Error(err)) }() // ctx 取消或处理结束时关闭连接，解除 Recv 阻塞。 @@ -624,8 +703,25 @@ func (s *Server) serveConn(ctx context.Context, raw transport.Conn) (err error) timeout = s.handshakeTimeout } if err := s.recv(ctx, conn, &b, timeout); err != nil { + if !firstFrameSeen { + outcome := "error" + if isClientDisconnect(err) { + outcome = "client_disconnect" + } + s.observeConnectionIntake(connectionIntakeEvent{ + stage: "first_frame", outcome: outcome, remote: remote, local: local, + duration: time.Since(firstFrameStarted), err: err, + }) + } return err } + if !firstFrameSeen { + firstFrameSeen = true + s.observeConnectionIntake(connectionIntakeEvent{ + stage: "first_frame", outcome: "ready", remote: remote, local: local, + duration: time.Since(firstFrameStarted), bytes: b.Len(), + }) + } if s.onFrame != nil { s.onFrame(b.Len()) } @@ -742,6 +838,7 @@ func (s *Server) recv(ctx context.Context, conn transport.Conn, b *bin.Buffer, t func isClientDisconnect(err error) bool { switch { case errors.Is(err, io.EOF), + errors.Is(err, io.ErrUnexpectedEOF), errors.Is(err, net.ErrClosed), errors.Is(err, context.Canceled), errors.Is(err, context.DeadlineExceeded): diff --git a/internal/mtprotoedge/session_barrier_integration_test.go b/internal/mtprotoedge/session_barrier_integration_test.go index c8142c92..755694fc 100644 --- a/internal/mtprotoedge/session_barrier_integration_test.go +++ b/internal/mtprotoedge/session_barrier_integration_test.go @@ -492,29 +492,28 @@ func TestCrossConnectionInflightRPCHasOneBusinessOwnerAndReplaysResult(t *testin if got := handler.calls.Load(); got != 1 { t.Fatalf("overlapping reconnect executed %d business handlers, want 1", got) } - select { - case got := <-secondDone: - t.Fatalf("duplicate reconnect returned before owner completion: %v", got.err) - default: - } - - close(handler.release) var second handleResult select { case second = <-secondDone: - case <-time.After(3 * time.Second): - t.Fatal("duplicate reconnect did not receive owner result") + case <-time.After(time.Second): + t.Fatal("event-driven duplicate admission blocked the replacement read loop") } if second.err != nil { t.Fatalf("second handleEncrypted: %v", second.err) } if second.conn == nil || !second.conn.isActive() { - t.Fatalf("second connection was not active after replay: %p", second.conn) + t.Fatalf("second connection was not active after non-blocking admission: %p", second.conn) } + + close(handler.release) if got := handler.calls.Load(); got != 1 { t.Fatalf("cross-connection duplicate business calls = %d, want 1", got) } + deadline = time.Now().Add(3 * time.Second) + for len(secondTransport.snapshot()) < 3 && time.Now().Before(deadline) { + time.Sleep(time.Millisecond) + } resultCount := 0 for _, wire := range secondTransport.snapshot() { data, decryptErr := crypto.NewClientCipher(rand.Reader).DecryptFromBuffer(key, &bin.Buffer{Buf: wire}) @@ -582,9 +581,34 @@ func TestCrossConnectionInflightAbortRetriesOnlyAfterOldOwnerStops(t *testing.T) secondState := newConnState() var secondPlain bin.Buffer secondConn, err := s.handleEncrypted(context.Background(), secondTransport, secondState, nil, &stored, secondFrame, &secondPlain) - if err != nil { + if err != nil && !errors.Is(err, ErrConnClosed) { t.Fatalf("second handleEncrypted: %v", err) } + if secondConn == nil { + t.Fatal("second handleEncrypted returned no logical connection") + } + deadline := time.Now().Add(2 * time.Second) + for (handler.active.Load() != 0 || !secondConn.isRetired()) && time.Now().Before(deadline) { + time.Sleep(time.Millisecond) + } + if got := handler.calls.Load(); got != 1 { + t.Fatalf("event subscription re-executed aborted owner: calls=%d", got) + } + if handler.active.Load() != 0 || secondConn == nil || !secondConn.isRetired() { + t.Fatalf("abort convergence = active:%d second:%p state:%v", handler.active.Load(), secondConn, secondConn.lifecycleState()) + } + + // The aborted owner has now definitively stopped and the subscribed + // replacement generation is fenced. A fresh physical retry may acquire the + // same msg_id and execute, still with max concurrency one. + thirdFrame, _ := encryptedRPCFrameForBarrierTest(t, key, salt, sessionID, msgID) + thirdTransport := &collectingSessionTransport{} + thirdState := newConnState() + var thirdPlain bin.Buffer + thirdConn, err := s.handleEncrypted(context.Background(), thirdTransport, thirdState, nil, &stored, thirdFrame, &thirdPlain) + if err != nil { + t.Fatalf("third handleEncrypted: %v", err) + } waitForAtomicCalls(t, &handler.calls, 2) waitInboundRPCBatchBudget(t, s.rpcScheduler, 0, 0) if got := handler.max.Load(); got != 1 { @@ -595,7 +619,11 @@ func TestCrossConnectionInflightAbortRetriesOnlyAfterOldOwnerStops(t *testing.T) } resultCount := 0 - for _, wire := range secondTransport.snapshot() { + deadline = time.Now().Add(2 * time.Second) + for len(thirdTransport.snapshot()) < 3 && time.Now().Before(deadline) { + time.Sleep(time.Millisecond) + } + for _, wire := range thirdTransport.snapshot() { data, decryptErr := crypto.NewClientCipher(rand.Reader).DecryptFromBuffer(key, &bin.Buffer{Buf: wire}) if decryptErr != nil { t.Fatalf("decrypt sequential-retry reply: %v", decryptErr) @@ -612,10 +640,10 @@ func TestCrossConnectionInflightAbortRetriesOnlyAfterOldOwnerStops(t *testing.T) if resultCount != 1 { t.Fatalf("sequential retry result count = %d, want 1", resultCount) } - if !firstConn.isRetired() || secondConn == nil || !secondConn.isActive() { - t.Fatalf("replacement lifecycle = old:%v new:%p active:%v", firstConn.lifecycleState(), secondConn, secondConn != nil && secondConn.isActive()) + if !firstConn.isRetired() || !secondConn.isRetired() || thirdConn == nil || !thirdConn.isActive() { + t.Fatalf("replacement lifecycle = first:%v second:%v third:%p active:%v", firstConn.lifecycleState(), secondConn.lifecycleState(), thirdConn, thirdConn != nil && thirdConn.isActive()) } - secondConn.ForceClose() + thirdConn.ForceClose() } func waitForAtomicCalls(t *testing.T, calls interface{ Load() int32 }, want int32) { diff --git a/internal/mtprotoedge/transport_compat.go b/internal/mtprotoedge/transport_compat.go index 122aa329..83fee282 100644 --- a/internal/mtprotoedge/transport_compat.go +++ b/internal/mtprotoedge/transport_compat.go @@ -26,6 +26,9 @@ const quickAckResponseFlag = uint32(1 << 31) const ( maxCompatPacketOverhead = 7 // 4-byte header + up to 3 bytes padded-intermediate padding. maxRetainedDirectMessageScratch = 64 << 10 + progressiveWriteMinBytes = 64 << 10 + progressiveWriteChunkBytes = 32 << 10 + progressiveWriteIdleTimeout = 10 * time.Second ) type transportListener interface { @@ -237,12 +240,83 @@ func (c *compatTransportConn) sendDeadline(deadline time.Time, b *bin.Buffer, sc } return nil } - if err := c.codec.Write(c.conn, b); err != nil { + writer := io.Writer(c.conn) + // A transport packet is still physically serialized as one uninterrupted + // byte sequence. Only large raw-TCP payloads use bounded chunks so progress + // refreshes the idle deadline and a stalled link reports how far it got. + if b.Len() >= progressiveWriteMinBytes { + writer = &progressiveDeadlineWriter{ + conn: c.conn, hardDeadline: deadline, + idleTimeout: progressiveWriteIdleTimeout, chunkBytes: progressiveWriteChunkBytes, + } + } + if err := c.codec.Write(writer, b); err != nil { return errors.Wrap(err, "write") } return nil } +type progressiveWriteError struct { + Written int + Chunks int + Err error +} + +func (e *progressiveWriteError) Error() string { + return fmt.Sprintf("progressive transport write after %d bytes/%d chunks: %v", e.Written, e.Chunks, e.Err) +} + +func (e *progressiveWriteError) Unwrap() error { return e.Err } + +type progressiveDeadlineWriter struct { + conn net.Conn + hardDeadline time.Time + idleTimeout time.Duration + chunkBytes int + written int + chunks int +} + +func (w *progressiveDeadlineWriter) Write(p []byte) (int, error) { + if w == nil || w.conn == nil { + return 0, io.ErrClosedPipe + } + startWritten := w.written + chunkBytes := w.chunkBytes + if chunkBytes <= 0 { + chunkBytes = progressiveWriteChunkBytes + } + for len(p) > 0 { + deadline := w.hardDeadline + if w.idleTimeout > 0 { + idle := time.Now().Add(w.idleTimeout) + if deadline.IsZero() || idle.Before(deadline) { + deadline = idle + } + } + if err := w.conn.SetWriteDeadline(deadline); err != nil { + return w.written - startWritten, &progressiveWriteError{Written: w.written, Chunks: w.chunks, Err: err} + } + chunk := p + if len(chunk) > chunkBytes { + chunk = chunk[:chunkBytes] + } + n, err := w.conn.Write(chunk) + w.written += n + if n > 0 { + w.chunks++ + p = p[n:] + } + if err != nil { + return w.written - startWritten, &progressiveWriteError{Written: w.written, Chunks: w.chunks, Err: err} + } + if n == 0 { + return w.written - startWritten, &progressiveWriteError{Written: w.written, Chunks: w.chunks, Err: io.ErrShortWrite} + } + } + return w.written - startWritten, nil +} + func (c *compatTransportConn) writeTransportPacketMessage(b *bin.Buffer, scratch *[]byte) error { required := b.Len() + maxCompatPacketOverhead if cap(*scratch) < required { diff --git a/internal/mtprotoedge/transport_compat_test.go b/internal/mtprotoedge/transport_compat_test.go index 34e9944c..6bfd926a 100644 --- a/internal/mtprotoedge/transport_compat_test.go +++ b/internal/mtprotoedge/transport_compat_test.go @@ -13,8 +13,9 @@ import ( type messageWriteTestConn struct { bytes.Buffer - writes int - maxWrite int + writes int + maxWrite int + deadlines []time.Time } func (c *messageWriteTestConn) Read([]byte) (int, error) { return 0, io.EOF } @@ -25,12 +26,15 @@ func (c *messageWriteTestConn) Write(p []byte) (int, error) { } return c.Buffer.Write(p) } -func (*messageWriteTestConn) Close() error { return nil } -func (*messageWriteTestConn) LocalAddr() net.Addr { return messageWriteTestAddr("local") } -func (*messageWriteTestConn) RemoteAddr() net.Addr { return messageWriteTestAddr("remote") } -func (*messageWriteTestConn) SetDeadline(time.Time) error { return nil } -func (*messageWriteTestConn) SetReadDeadline(time.Time) error { return nil } -func (*messageWriteTestConn) SetWriteDeadline(time.Time) error { return nil } +func (*messageWriteTestConn) Close() error { return nil } +func (*messageWriteTestConn) LocalAddr() net.Addr { return messageWriteTestAddr("local") } +func (*messageWriteTestConn) RemoteAddr() net.Addr { return messageWriteTestAddr("remote") } +func (*messageWriteTestConn) SetDeadline(time.Time) error { return nil } +func (*messageWriteTestConn) SetReadDeadline(time.Time) error { return nil } +func (c *messageWriteTestConn) SetWriteDeadline(d time.Time) error { + c.deadlines = append(c.deadlines, d) + return nil +} type messageWriteTestAddr string @@ -42,6 +46,33 @@ type countWriteBuffer struct { writes int } +func TestProgressiveDeadlineWriterChunksLargeRawPacket(t *testing.T) { + raw := &messageWriteTestConn{maxWrite: 8 << 10} + hard := time.Now().Add(time.Minute) + w := &progressiveDeadlineWriter{ + conn: raw, hardDeadline: hard, idleTimeout: time.Second, chunkBytes: 32 << 10, + } + payload := bytes.Repeat([]byte{0x5a}, 192<<10) + n, err := w.Write(payload) + if err != nil { + t.Fatalf("progressive write: %v", err) + } + if n != len(payload) || !bytes.Equal(raw.Bytes(), payload) { + t.Fatalf("progressive write bytes = %d/%d", n, len(raw.Bytes())) + } + if raw.writes < len(payload)/(8<<10) { + t.Fatalf("physical chunks = %d, want progress across partial writes", raw.writes) + } + if len(raw.deadlines) != raw.writes { + t.Fatalf("deadline refreshes = %d, writes = %d", len(raw.deadlines), raw.writes) + } + for _, deadline := range raw.deadlines { + if deadline.After(hard) { + t.Fatalf("idle deadline %v exceeded hard deadline %v", deadline, hard) + } + } +} + func (w *countWriteBuffer) Write(p []byte) (int, error) { w.writes++ return w.Buffer.Write(p) @@ -158,6 +189,30 @@ func TestCompatTransportCodecsWriteSegmentedPacketWithoutFullCopy(t *testing.T) }) } +func TestCompatRawTransportUsesProgressiveWriterForLargePacket(t *testing.T) { + var payload bin.Buffer + payload.Put(bytes.Repeat([]byte{0x6b}, 192<<10)) + raw := &messageWriteTestConn{maxWrite: 8 << 10} + conn := &compatTransportConn{conn: raw, codec: &quickAckAbridgedCodec{}} + if err := conn.SendDeadline(time.Now().Add(time.Minute), &payload); err != nil { + t.Fatalf("send large raw packet: %v", err) + } + if raw.writes < 20 { + t.Fatalf("large raw packet writes = %d, want observable progressive chunks", raw.writes) + } + if len(raw.deadlines) != raw.writes+1 { // initial hard deadline + one per progress write. + t.Fatalf("large packet deadline updates = %d, writes = %d", len(raw.deadlines), raw.writes) + } + var decoded bin.Buffer + requested, err := readQuickAckAbridged(bytes.NewReader(raw.Bytes()), &decoded) + if err != nil { + t.Fatalf("read large packet: %v", err) + } + if requested || !bytes.Equal(decoded.Raw(), payload.Raw()) { + t.Fatal("progressive raw write changed transport packet") + } +} + func TestCompatTransportWebSocketWritesOneMessagePerPacket(t *testing.T) { var payload bin.Buffer payload.PutInt32(0x01020304) diff --git a/internal/postresponse/callbacks.go b/internal/postresponse/callbacks.go index 95015b18..29dea0c0 100644 --- a/internal/postresponse/callbacks.go +++ b/internal/postresponse/callbacks.go @@ -39,15 +39,34 @@ func Register(ctx context.Context, cb func()) bool { } func Run(ctx context.Context) { + run := Take(ctx) + if run != nil { + run() + } +} + +// Take transfers ownership of every currently registered callback to the caller. +// The returned function is idempotent and may safely outlive the request context. +// MTProto uses this to release a business worker after admitting rpc_result while +// still delaying follow-up updates until the result reaches the reliable stream. +func Take(ctx context.Context) func() { cbs, ok := ctx.Value(callbacksKey{}).(*callbacks) if !ok || cbs == nil { - return + return nil } cbs.mu.Lock() list := append([]callback(nil), cbs.list...) cbs.list = nil cbs.mu.Unlock() - for _, cb := range list { - cb() + if len(list) == 0 { + return nil + } + var once sync.Once + return func() { + once.Do(func() { + for _, cb := range list { + cb() + } + }) } } diff --git a/internal/postresponse/callbacks_test.go b/internal/postresponse/callbacks_test.go new file mode 100644 index 00000000..94c52b37 --- /dev/null +++ b/internal/postresponse/callbacks_test.go @@ -0,0 +1,28 @@ +package postresponse + +import ( + "context" + "sync/atomic" + "testing" +) + +func TestTakeTransfersCallbacksAndRunsOnce(t *testing.T) { + ctx := WithCallbacks(context.Background()) + var calls atomic.Int32 + if !Register(ctx, func() { calls.Add(1) }) || !Register(ctx, func() { calls.Add(1) }) { + t.Fatal("register callbacks") + } + run := Take(ctx) + if run == nil { + t.Fatal("Take returned no callback") + } + Run(ctx) + if got := calls.Load(); got != 0 { + t.Fatalf("callbacks remained attached after Take: %d", got) + } + run() + run() + if got := calls.Load(); got != 2 { + t.Fatalf("transferred callbacks ran %d times, want 2 total", got) + } +} diff --git a/internal/rpc/router.go b/internal/rpc/router.go index ef3d4dad..76c8debb 100644 --- a/internal/rpc/router.go +++ b/internal/rpc/router.go @@ -218,23 +218,31 @@ func New(cfg Config, deps Deps, log *zap.Logger, clk clock.Clock) *Router { return r } -// Dispatch 路由一条 RPC 请求：先剥离 invokeWithLayer / initConnection / -// invokeWithoutUpdates / invokeAfter* 等 wrapper（注入 layer / 客户端信息到 ctx）， -// 再按 TypeID 路由到 typed handler。满足 mtprotoedge.RPCHandler。 +// Dispatch routes one RPC and preserves the historical two-value API used by +// domain/RPC tests. The MTProto edge uses DispatchWithMethod so outbound +// scheduling sees the canonical inner method rather than an invoke wrapper. func (r *Router) Dispatch(ctx context.Context, authKeyID [8]byte, sessionID int64, b *bin.Buffer) (bin.Encoder, error) { + enc, _, err := r.DispatchWithMethod(ctx, authKeyID, sessionID, b) + return enc, err +} + +// DispatchWithMethod 路由一条 RPC 请求：先剥离 invokeWithLayer / initConnection / +// invokeWithoutUpdates / invokeAfter* 等 wrapper（注入 layer / 客户端信息到 ctx）， +// 再按 TypeID 路由到 typed handler，并返回 canonical innermost method。 +func (r *Router) DispatchWithMethod(ctx context.Context, authKeyID [8]byte, sessionID int64, b *bin.Buffer) (bin.Encoder, string, error) { preStart := r.clock.Now() ctx = withInboundRPCBytes(ctx, b.Len()) ctx = WithRawAuthKeyID(ctx, authKeyID) effectiveAuthKeyID, err := r.effectiveAuthKeyID(ctx, authKeyID, sessionID) if err != nil { - return nil, internalErr() + return nil, "", internalErr() } tAuth := r.clock.Now() ctx = WithAuthKeyID(ctx, effectiveAuthKeyID) ctx = WithSessionID(ctx, sessionID) userID, hasUserID, err := r.effectiveUserID(ctx, authKeyID, effectiveAuthKeyID, sessionID) if err != nil { - return nil, internalErr() + return nil, "", internalErr() } if hasUserID { ctx = WithUserID(ctx, userID) @@ -299,7 +307,9 @@ func (r *Router) Dispatch(ctx context.Context, authKeyID [8]byte, sessionID int6 ctx = WithClientInfo(ctx, info.clientInfo) } } - return r.dispatch(ctx, b, 0) + meta := rpcDispatchMetadata{} + enc, err := r.dispatch(ctx, b, 0, &meta) + return enc, meta.method, err } func (r *Router) effectiveAuthKeyID(ctx context.Context, rawAuthKeyID [8]byte, sessionID int64) ([8]byte, error) { @@ -494,7 +504,11 @@ func (r *Router) invalidateAuthUserCache(authKeyID [8]byte) { r.authUserSF.Forget(authKeyClientInfoSingleflightPrefix + key) } -func (r *Router) dispatch(ctx context.Context, b *bin.Buffer, depth int) (bin.Encoder, error) { +type rpcDispatchMetadata struct { + method string +} + +func (r *Router) dispatch(ctx context.Context, b *bin.Buffer, depth int, meta *rpcDispatchMetadata) (bin.Encoder, error) { if depth > maxWrapperDepth { return nil, wrapperTooDeepErr() } @@ -516,13 +530,13 @@ func (r *Router) dispatch(ctx context.Context, b *bin.Buffer, depth int) (bin.En // query 紧跟 layer，buffer 剩余即内层请求。 ctx = WithLayer(ctx, layer) r.rememberClientLayer(ctx, layer) - return r.dispatch(ctx, b, depth+1) + return r.dispatch(ctx, b, depth+1, meta) case tg.InvokeWithoutUpdatesRequestTypeID: if err := b.ConsumeID(id); err != nil { return nil, err } - return r.dispatch(withInvokeWithoutUpdates(ctx), b, depth+1) + return r.dispatch(withInvokeWithoutUpdates(ctx), b, depth+1, meta) case tg.InvokeAfterMsgRequestTypeID: if err := b.ConsumeID(id); err != nil { @@ -531,7 +545,7 @@ func (r *Router) dispatch(ctx context.Context, b *bin.Buffer, depth int) (bin.En if _, err := b.Long(); err != nil { return nil, fmt.Errorf("decode invokeAfterMsg msg_id: %w", err) } - return r.dispatch(ctx, b, depth+1) + return r.dispatch(ctx, b, depth+1, meta) case tg.InvokeAfterMsgsRequestTypeID: if err := b.ConsumeID(id); err != nil { @@ -549,7 +563,7 @@ func (r *Router) dispatch(ctx context.Context, b *bin.Buffer, depth int) (bin.En return nil, fmt.Errorf("decode invokeAfterMsgs msg_ids[%d]: %w", i, err) } } - return r.dispatch(ctx, b, depth+1) + return r.dispatch(ctx, b, depth+1, meta) case tg.InitConnectionRequestTypeID: req := &tg.InitConnectionRequest{Query: &rawObject{}} @@ -578,7 +592,7 @@ func (r *Router) dispatch(ctx context.Context, b *bin.Buffer, depth int) (bin.En if !ok { return nil, fmt.Errorf("initConnection query: unexpected type %T", req.Query) } - return r.dispatch(ctx, &bin.Buffer{Buf: inner.data}, depth+1) + return r.dispatch(ctx, &bin.Buffer{Buf: inner.data}, depth+1, meta) default: // 入站兼容统一入口（先于鉴权门/dispatcher）：layerwire 把老客户端请求升级为 @@ -606,6 +620,9 @@ func (r *Router) dispatch(ctx context.Context, b *bin.Buffer, depth int) (bin.En } } } + if meta != nil { + meta.method = tlTypeName(id) + } knownRequest, structuralErr := layerwire.ValidateRoutableRequest(b.Buf) if !knownRequest { if structuralErr != nil { @@ -834,6 +851,37 @@ func (r *Router) NegotiatedLayer(authKeyID [8]byte, sessionID int64) (int, bool) return currentClientLayer, false } +// ObserveInitConnection records the protocol metadata of an initConnection +// whose inner request was aliased by mtprotoedge to an already-running naked +// request. It deliberately performs no handler dispatch and therefore cannot +// repeat business side effects. +func (r *Router) ObserveInitConnection( + ctx context.Context, + rawAuthKeyID [8]byte, + sessionID int64, + layer, apiID int, + deviceModel, systemVersion, appVersion, systemLangCode, langPack, langCode string, +) error { + if ctx == nil { + ctx = context.Background() + } + ctx = WithRawAuthKeyID(ctx, rawAuthKeyID) + effectiveAuthKeyID, err := r.effectiveAuthKeyID(ctx, rawAuthKeyID, sessionID) + if err != nil { + return err + } + ctx = WithAuthKeyID(ctx, effectiveAuthKeyID) + ctx = WithSessionID(ctx, sessionID) + ctx = WithLayer(ctx, layer) + r.rememberClientLayer(ctx, layer) + r.rememberClientInfo(ctx, ClientInfo{ + APIID: apiID, DeviceModel: deviceModel, SystemVersion: systemVersion, + AppVersion: appVersion, SystemLangCode: systemLangCode, + LangPack: langPack, LangCode: langCode, + }) + return nil +} + // mutateClientSessionInfo 在单个临界区内完成「读旧值-修改-写回」，避免 // RLock 读出与 Lock 写回之间被并发写覆盖的窗口。 func (r *Router) mutateClientSessionInfo(ctx context.Context, mutate func(*clientSessionInfo)) { diff --git a/internal/rpc/router_dispatch_test.go b/internal/rpc/router_dispatch_test.go index c07bfb57..ff3b1cef 100644 --- a/internal/rpc/router_dispatch_test.go +++ b/internal/rpc/router_dispatch_test.go @@ -49,10 +49,13 @@ func TestDispatchUnwrapsWrappers(t *testing.T) { t.Fatalf("encode wrapped request: %v", err) } - enc, err := r.Dispatch(context.Background(), [8]byte{}, 0, &b) + enc, method, err := r.DispatchWithMethod(context.Background(), [8]byte{}, 0, &b) if err != nil { t.Fatalf("dispatch: %v", err) } + if method != "help.getConfig" { + t.Fatalf("effective method = %q, want help.getConfig", method) + } cfg, ok := enc.(*tg.Config) if !ok { t.Fatalf("result type = %T, want *tg.Config", enc)