fix: sync MTProto startup and egress fixes

This commit is contained in:
A 2026-07-13 19:05:41 +08:00
parent 50803a604c
commit 305e8a0008
24 changed files with 2880 additions and 317 deletions

View file

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

View file

@ -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

View file

@ -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"
}

View file

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

View file

@ -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)
// 刷新本连接协商 layerinvokeWithLayer/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
}

View file

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

View file

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

View file

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

View file

@ -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")

View file

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

View file

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

View file

@ -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"
}
}

View file

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

View file

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

View file

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

View file

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

View file

@ -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 {
// WebSocketgotd 升级处理器已剥离 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 TCPWebSocket
// 连接传 falsegotd 升级处理器已完成去混淆)。
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):

View file

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

View file

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

View file

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

View file

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

View file

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

View file

@ -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 紧跟 layerbuffer 剩余即内层请求。
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:
// 入站兼容统一入口(先于鉴权门/dispatcherlayerwire 把老客户端请求升级为
@ -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)) {

View file

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