fix: sync MTProto startup and egress fixes
This commit is contained in:
parent
50803a604c
commit
305e8a0008
24 changed files with 2880 additions and 317 deletions
|
|
@ -347,11 +347,6 @@ func run(logger *zap.Logger) error {
|
||||||
defer func() { _ = rdb.Close() }()
|
defer func() { _ = rdb.Close() }()
|
||||||
logger.Info("持久化依赖就绪", zap.String("redis", cfg.RedisAddr))
|
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)
|
authKeyStore := postgres.NewAuthKeyStore(pool)
|
||||||
userStore := postgres.NewUserStore(pool)
|
userStore := postgres.NewUserStore(pool)
|
||||||
authzStore := postgres.NewAuthorizationStore(pool)
|
authzStore := postgres.NewAuthorizationStore(pool)
|
||||||
|
|
@ -858,7 +853,7 @@ func run(logger *zap.Logger) error {
|
||||||
OutboundControlQueueSize: cfg.MTProtoOutboundControlQueueSize,
|
OutboundControlQueueSize: cfg.MTProtoOutboundControlQueueSize,
|
||||||
OutboundTrackedGlobalMaxBytes: cfg.MTProtoOutboundTrackedGlobalMaxBytes,
|
OutboundTrackedGlobalMaxBytes: cfg.MTProtoOutboundTrackedGlobalMaxBytes,
|
||||||
OutboundWriteGlobalMaxBytes: cfg.MTProtoOutboundWriteGlobalMaxBytes,
|
OutboundWriteGlobalMaxBytes: cfg.MTProtoOutboundWriteGlobalMaxBytes,
|
||||||
})
|
OnServing: func(_ net.Addr) {
|
||||||
logger.Info("telesrv 服务就绪",
|
logger.Info("telesrv 服务就绪",
|
||||||
zap.String("listen", cfg.ListenAddr),
|
zap.String("listen", cfg.ListenAddr),
|
||||||
zap.String("advertise", net.JoinHostPort(cfg.AdvertiseIP, portStr)),
|
zap.String("advertise", net.JoinHostPort(cfg.AdvertiseIP, portStr)),
|
||||||
|
|
@ -867,5 +862,9 @@ func run(logger *zap.Logger) error {
|
||||||
zap.Uint("schema_version", migrationStatus.Version),
|
zap.Uint("schema_version", migrationStatus.Version),
|
||||||
zap.String("blob_backend", "localfs"),
|
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)
|
||||||
}
|
}
|
||||||
|
|
|
||||||
|
|
@ -61,6 +61,10 @@ type Conn struct {
|
||||||
|
|
||||||
outbound chan outboundOp
|
outbound chan outboundOp
|
||||||
outboundControl chan outboundOp
|
outboundControl chan outboundOp
|
||||||
|
// 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{}
|
outboundStop chan struct{}
|
||||||
outboundDone chan struct{}
|
outboundDone chan struct{}
|
||||||
outboundClose sync.Once
|
outboundClose sync.Once
|
||||||
|
|
@ -103,6 +107,12 @@ type Conn struct {
|
||||||
rpcRunning int
|
rpcRunning int
|
||||||
rpcReady bool
|
rpcReady bool
|
||||||
rpcClosed 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
|
// inflightRPCBytes 跟踪已入队未完成的 inbound RPC body 总字节,配合 maxInflightRPCBytes
|
||||||
// 给 RPC 队列设字节预算(不止限条数),防对抗客户端发大请求撑内存。
|
// 给 RPC 队列设字节预算(不止限条数),防对抗客户端发大请求撑内存。
|
||||||
inflightRPCBytes atomic.Int64
|
inflightRPCBytes atomic.Int64
|
||||||
|
|
|
||||||
95
internal/mtprotoedge/connection_intake.go
Normal file
95
internal/mtprotoedge/connection_intake.go
Normal 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"
|
||||||
|
}
|
||||||
170
internal/mtprotoedge/connection_intake_test.go
Normal file
170
internal/mtprotoedge/connection_intake_test.go
Normal 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
|
||||||
|
}
|
||||||
|
|
@ -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 {
|
if err := s.executeInboundPlan(ctx, cs, current, plan); err != nil {
|
||||||
return current, err
|
return current, err
|
||||||
}
|
}
|
||||||
|
if err := plan.commitRewrapAliases(s); err != nil {
|
||||||
|
return current, err
|
||||||
|
}
|
||||||
if err := plan.commitRPCBatch(); err != nil {
|
if err := plan.commitRPCBatch(); err != nil {
|
||||||
return current, err
|
return current, err
|
||||||
}
|
}
|
||||||
|
|
@ -268,12 +271,6 @@ func (s *Server) handleEncrypted(ctx context.Context, tc transport.Conn, cs *con
|
||||||
return current, err
|
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
|
return current, nil
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
@ -577,58 +574,6 @@ func mergeStateInfo(primary, fallback []byte) []byte {
|
||||||
return info
|
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
|
// newInboundRPCTask builds the exactly-once timeout/result gate shared by the
|
||||||
// single-message and atomic container-batch admission paths. body must already
|
// single-message and atomic container-batch admission paths. body must already
|
||||||
// be an independently owned, budgeted copy.
|
// 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 {
|
run: func(taskCtx context.Context) error {
|
||||||
// body 是预算成功后生成的独立副本,且每个任务只 run 一次,
|
// body 是预算成功后生成的独立副本,且每个任务只 run 一次,
|
||||||
// 无需再 append 拷贝;直接复用,省掉一份 inbound 在途内存。
|
// 无需再 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{
|
fields := []zap.Field{
|
||||||
zap.Int64("msg_id", msgID),
|
zap.Int64("msg_id", msgID),
|
||||||
zap.String("auth_key_id", c.authKeyHex),
|
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 回发。
|
// 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 s.rpc == nil {
|
||||||
if err := ctx.Err(); err != nil {
|
if err := ctx.Err(); err != nil {
|
||||||
return err
|
return err
|
||||||
}
|
}
|
||||||
s.log.Warn("No RPC handler configured", zap.String("method", method))
|
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,
|
ErrorCode: 500,
|
||||||
ErrorMessage: "NOT_IMPLEMENTED",
|
ErrorMessage: "NOT_IMPLEMENTED",
|
||||||
})
|
}, nil)
|
||||||
}
|
}
|
||||||
|
|
||||||
ctx = postresponse.WithCallbacks(ctx)
|
ctx = postresponse.WithCallbacks(ctx)
|
||||||
ctx, dbStats := dbtrace.WithStats(ctx)
|
ctx, dbStats := dbtrace.WithStats(ctx)
|
||||||
start := s.clock.Now()
|
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)
|
dur := s.clock.Now().Sub(start)
|
||||||
s.metrics.RPCHandled(method, dur, err)
|
s.metrics.RPCHandled(effectiveMethod, dur, err)
|
||||||
// 刷新本连接协商 layer(invokeWithLayer/initConnection 已被 Dispatch 处理并登记),
|
// 刷新本连接协商 layer(invokeWithLayer/initConnection 已被 Dispatch 处理并登记),
|
||||||
// 供 rpc_result 与后续 push 出站降级使用。仅在确实观测到 layer 时更新——缓存被驱逐
|
// 供 rpc_result 与后续 push 出站降级使用。仅在确实观测到 layer 时更新——缓存被驱逐
|
||||||
// 时 NegotiatedLayer 返回 ok=false,此时必须保留连接已记住的 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 := make([]zap.Field, 0, 12)
|
||||||
fields = append(fields,
|
fields = append(fields,
|
||||||
zap.String("method", method),
|
zap.String("method", effectiveMethod),
|
||||||
zap.String("auth_key_id", c.authKeyHex),
|
zap.String("auth_key_id", c.authKeyHex),
|
||||||
zap.Int64("session_id", c.sessionID),
|
zap.Int64("session_id", c.sessionID),
|
||||||
zap.Int64("msg_id", msgID),
|
zap.Int64("msg_id", msgID),
|
||||||
zap.Duration("dur", dur),
|
zap.Duration("dur", dur),
|
||||||
)
|
)
|
||||||
|
if effectiveMethod != method {
|
||||||
|
fields = append(fields, zap.String("outer_method", method))
|
||||||
|
}
|
||||||
if businessAuthKeyHex, ok := c.BusinessAuthKeyHex(); ok {
|
if businessAuthKeyHex, ok := c.BusinessAuthKeyHex(); ok {
|
||||||
fields = append(fields, zap.String("business_auth_key_id", businessAuthKeyHex))
|
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"}
|
terminal = &mt.RPCError{ErrorCode: 500, ErrorMessage: "RPC_TIMEOUT"}
|
||||||
}
|
}
|
||||||
if terminal != nil {
|
if terminal != nil {
|
||||||
if c.isRetired() || !c.isPhysicalTransportCurrentOpen() {
|
var after func()
|
||||||
// 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 {
|
if runPostResponse {
|
||||||
postresponse.Run(context.WithoutCancel(ctx))
|
after = postresponse.Take(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))
|
|
||||||
}
|
}
|
||||||
|
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))
|
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
|
var rpcErr *tgerr.Error
|
||||||
if errors.As(err, &rpcErr) {
|
if errors.As(err, &rpcErr) {
|
||||||
s.log.Info("RPC error", append(fields, zap.Int("code", rpcErr.Code), zap.String("error", rpcErr.Message))...)
|
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,
|
ErrorCode: rpcErr.Code,
|
||||||
ErrorMessage: rpcErr.Message,
|
ErrorMessage: rpcErr.Message,
|
||||||
})
|
}, nil)
|
||||||
}
|
}
|
||||||
s.log.Info("RPC internal error", append(fields, zap.Error(err))...)
|
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,
|
ErrorCode: 500,
|
||||||
ErrorMessage: "INTERNAL",
|
ErrorMessage: "INTERNAL",
|
||||||
})
|
}, nil)
|
||||||
}
|
}
|
||||||
|
|
||||||
s.log.Info("RPC handled", fields...)
|
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
|
return err
|
||||||
}
|
}
|
||||||
postresponse.Run(ctx)
|
}
|
||||||
|
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
|
||||||
|
}
|
||||||
|
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
|
return nil
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
@ -833,13 +876,13 @@ func (s *Server) sendResult(ctx context.Context, c *Conn, reqMsgID int64, result
|
||||||
if result == nil {
|
if result == nil {
|
||||||
result = &mt.RPCError{ErrorCode: 500, ErrorMessage: "INTERNAL"}
|
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 {
|
if err != nil {
|
||||||
// The business operation has already crossed atomic admission. Convert an
|
// The business operation has already crossed atomic admission. Convert an
|
||||||
// invalid result encoder into one deterministic terminal RPC error instead of
|
// invalid result encoder into one deterministic terminal RPC error instead of
|
||||||
// aborting the flight and allowing a reconnect to execute the operation again.
|
// 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))
|
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,
|
ErrorCode: 500,
|
||||||
ErrorMessage: "INTERNAL",
|
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
|
// (queue/context/scratch deadline); without this terminal barrier a later
|
||||||
// same-Conn duplicate would be ACKed while no result can ever arrive.
|
// same-Conn duplicate would be ACKed while no result can ever arrive.
|
||||||
c.fenceUndeliveredRPCResult()
|
c.fenceUndeliveredRPCResult()
|
||||||
|
encoded.markReplayable()
|
||||||
s.storeRPCResult(c, reqMsgID, encoded)
|
s.storeRPCResult(c, reqMsgID, encoded)
|
||||||
return err
|
return err
|
||||||
}
|
}
|
||||||
// On a live Conn, completed means the rpc_result has reached the reliable byte
|
// 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.
|
// stream. Same-physical duplicates can therefore be ACK-only without data loss.
|
||||||
s.storeRPCResult(c, reqMsgID, encoded)
|
s.storeRPCResult(c, reqMsgID, encoded)
|
||||||
|
encoded.markDelivered()
|
||||||
return nil
|
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 {
|
if err := c.SendEncoded(ctx, proto.MessageServerResponse, encoded); err != nil {
|
||||||
c.fenceUndeliveredRPCResult()
|
c.fenceUndeliveredRPCResult()
|
||||||
|
encoded.markReplayable()
|
||||||
return err
|
return err
|
||||||
}
|
}
|
||||||
|
encoded.markDelivered()
|
||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
@ -884,30 +931,43 @@ func (s *Server) sendCachedRPCResult(ctx context.Context, c *Conn, encoded *enco
|
||||||
// 降级改写字节时才重建整条消息。降级失败 fail-safe:记日志并发送 canonical 字节——
|
// 降级改写字节时才重建整条消息。降级失败 fail-safe:记日志并发送 canonical 字节——
|
||||||
// 宁可老客户端对个别长尾对象渲染异常,也不让连接/流崩。
|
// 宁可老客户端对个别长尾对象渲染异常,也不让连接/流崩。
|
||||||
func (s *Server) encodeRPCResult(c *Conn, reqMsgID int64, result bin.Encoder) (*encodedOutboundMessage, error) {
|
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
|
return s.encodeRPCResultContext(context.Background(), c, reqMsgID, result)
|
||||||
var buf bin.Buffer
|
}
|
||||||
buf.PutID(proto.ResultTypeID)
|
|
||||||
buf.PutLong(reqMsgID)
|
func (s *Server) encodeRPCResultContext(ctx context.Context, c *Conn, reqMsgID int64, result bin.Encoder) (*encodedOutboundMessage, error) {
|
||||||
if err := result.Encode(&buf); err != nil {
|
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)
|
return nil, fmt.Errorf("encode rpc result: %w", err)
|
||||||
}
|
}
|
||||||
|
innerBody := inner.Raw()
|
||||||
if layer := c.ClientLayer(); layer < layerwire.CanonicalLayer {
|
if layer := c.ClientLayer(); layer < layerwire.CanonicalLayer {
|
||||||
inner := buf.Buf[headerLen:]
|
if down, err := layerwire.Transcode(innerBody, layer); err != nil {
|
||||||
if down, err := layerwire.Transcode(inner, layer); err != nil {
|
|
||||||
s.log.Warn("layerwire downgrade failed; sending canonical rpc_result",
|
s.log.Warn("layerwire downgrade failed; sending canonical rpc_result",
|
||||||
zap.Int("layer", layer), zap.Int64("req_msg_id", reqMsgID), zap.Error(err))
|
zap.Int("layer", layer), zap.Int64("req_msg_id", reqMsgID), zap.Error(err))
|
||||||
} else if !sameBacking(down, inner) {
|
} else {
|
||||||
var rebuilt bin.Buffer
|
innerBody = down
|
||||||
rebuilt.PutID(proto.ResultTypeID)
|
|
||||||
rebuilt.PutLong(reqMsgID)
|
|
||||||
rebuilt.Put(down)
|
|
||||||
buf = rebuilt
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
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{
|
return &encodedOutboundMessage{
|
||||||
typeID: proto.ResultTypeID,
|
typeID: proto.ResultTypeID, body: out.Raw(), reqMsgID: reqMsgID,
|
||||||
body: buf.Raw(),
|
compressed: compressed, uncompressedBytes: len(innerBody), delivery: newRPCResultDelivery(),
|
||||||
reqMsgID: reqMsgID,
|
|
||||||
}, nil
|
}, nil
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -34,7 +34,10 @@ const (
|
||||||
inboundItemDestroyAuthKey
|
inboundItemDestroyAuthKey
|
||||||
inboundItemRPC
|
inboundItemRPC
|
||||||
inboundItemCapacityError
|
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
|
// inboundItemReplayRPC is a request first observed by this physical Conn whose
|
||||||
// terminal result already exists in the cross-connection cache. It is distinct
|
// terminal result already exists in the cross-connection cache. It is distinct
|
||||||
// from inboundItemDuplicate: a duplicate already present in this Conn's seen
|
// from inboundItemDuplicate: a duplicate already present in this Conn's seen
|
||||||
|
|
@ -71,6 +74,8 @@ type inboundPlan struct {
|
||||||
rpcReservation *inboundRPCBatchReservation
|
rpcReservation *inboundRPCBatchReservation
|
||||||
rpcTasks []inboundRPC
|
rpcTasks []inboundRPC
|
||||||
rpcOwners []*rpcResultOwnerLease
|
rpcOwners []*rpcResultOwnerLease
|
||||||
|
rewrapAliases []*rpcRewrapAlias
|
||||||
|
rewrapIndices []int
|
||||||
}
|
}
|
||||||
|
|
||||||
func (p *inboundPlan) close() {
|
func (p *inboundPlan) close() {
|
||||||
|
|
@ -85,12 +90,42 @@ func (p *inboundPlan) close() {
|
||||||
owner.Abort()
|
owner.Abort()
|
||||||
}
|
}
|
||||||
p.rpcOwners = nil
|
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-- {
|
for i := len(p.releases) - 1; i >= 0; i-- {
|
||||||
p.releases[i]()
|
p.releases[i]()
|
||||||
}
|
}
|
||||||
p.releases = nil
|
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 {
|
func (p *inboundPlan) commitRPCBatch() error {
|
||||||
if p == nil || p.rpcReservation == nil {
|
if p == nil || p.rpcReservation == nil {
|
||||||
return nil
|
return nil
|
||||||
|
|
@ -624,6 +659,7 @@ func (s *Server) prepareInboundRPCBatch(ctx context.Context, c *Conn, plan *inbo
|
||||||
var specs []inboundRPCSpec
|
var specs []inboundRPCSpec
|
||||||
var ownersInPlan map[int64]*rpcResultOwnerLease
|
var ownersInPlan map[int64]*rpcResultOwnerLease
|
||||||
flightCapacity := false
|
flightCapacity := false
|
||||||
|
clearedPostInitCandidates := false
|
||||||
for i := range plan.items {
|
for i := range plan.items {
|
||||||
item := &plan.items[i]
|
item := &plan.items[i]
|
||||||
if item.kind != inboundItemRPC && item.kind != inboundItemDuplicate {
|
if item.kind != inboundItemRPC && item.kind != inboundItemDuplicate {
|
||||||
|
|
@ -638,6 +674,71 @@ func (s *Server) prepareInboundRPCBatch(ctx context.Context, c *Conn, plan *inbo
|
||||||
continue
|
continue
|
||||||
}
|
}
|
||||||
method := s.typeName(item.typeID)
|
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)
|
claim, err := s.rpcResults.Acquire(c.authKeyID, c.sessionID, item.msgID)
|
||||||
if errors.Is(err, ErrRPCResultFlightCapacity) {
|
if errors.Is(err, ErrRPCResultFlightCapacity) {
|
||||||
if item.kind == inboundItemRPC {
|
if item.kind == inboundItemRPC {
|
||||||
|
|
@ -669,8 +770,11 @@ func (s *Server) prepareInboundRPCBatch(ctx context.Context, c *Conn, plan *inbo
|
||||||
item.kind = inboundItemDuplicate
|
item.kind = inboundItemDuplicate
|
||||||
item.payload = nil
|
item.payload = nil
|
||||||
} else {
|
} else {
|
||||||
item.kind = inboundItemPendingRPC
|
item.kind = inboundItemRewrappedRPC
|
||||||
item.payload = claim.waiter
|
item.payload = claim.waiter
|
||||||
|
plan.rewrapAliases = append(plan.rewrapAliases, &rpcRewrapAlias{
|
||||||
|
conn: c, newReqID: item.msgID, method: method, oldWaiter: claim.waiter,
|
||||||
|
})
|
||||||
}
|
}
|
||||||
case rpcResultAcquireOwner:
|
case rpcResultAcquireOwner:
|
||||||
if item.kind == inboundItemDuplicate {
|
if item.kind == inboundItemDuplicate {
|
||||||
|
|
@ -689,6 +793,9 @@ func (s *Server) prepareInboundRPCBatch(ctx context.Context, c *Conn, plan *inbo
|
||||||
item.payload = claim.owner
|
item.payload = claim.owner
|
||||||
indices = append(indices, i)
|
indices = append(indices, i)
|
||||||
specs = append(specs, inboundRPCSpec{method: method, size: len(item.body)})
|
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:
|
default:
|
||||||
return ErrRPCResultFlightInvalid
|
return ErrRPCResultFlightInvalid
|
||||||
}
|
}
|
||||||
|
|
@ -700,6 +807,17 @@ func (s *Server) prepareInboundRPCBatch(ctx context.Context, c *Conn, plan *inbo
|
||||||
for _, index := range indices {
|
for _, index := range indices {
|
||||||
plan.items[index].kind = inboundItemCapacityError
|
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 nil
|
||||||
}
|
}
|
||||||
if len(specs) == 0 {
|
if len(specs) == 0 {
|
||||||
|
|
@ -712,6 +830,17 @@ func (s *Server) prepareInboundRPCBatch(ctx context.Context, c *Conn, plan *inbo
|
||||||
for _, index := range indices {
|
for _, index := range indices {
|
||||||
plan.items[index].kind = inboundItemCapacityError
|
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 nil
|
||||||
}
|
}
|
||||||
return err
|
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 {
|
} else if err := s.replayRPCResultByRequest(ctx, c, item.msgID); err != nil {
|
||||||
return err
|
return err
|
||||||
}
|
}
|
||||||
case inboundItemPendingRPC:
|
case inboundItemRewrappedRPC:
|
||||||
// Wait only after this plan's fresh RPC batch has become runnable; see
|
// Activation is deferred until every session/control barrier has
|
||||||
// executePendingRPCReplays. Blocking here could otherwise deadlock on
|
// committed. It subscribes to the original result event and never waits
|
||||||
// an owner appended by the same container.
|
// or dispatches the business handler again.
|
||||||
continue
|
continue
|
||||||
case inboundItemPing:
|
case inboundItemPing:
|
||||||
if err := s.sendPong(ctx, c, item.msgID, item.payload.(mt.PingRequest).PingID); err != nil {
|
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
|
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
|
|
||||||
}
|
|
||||||
|
|
|
||||||
|
|
@ -30,6 +30,24 @@ type Metrics interface {
|
||||||
OutboundQueueWait(len, cap int)
|
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 的空实现。
|
// NopMetrics 是 Metrics 的空实现。
|
||||||
type NopMetrics struct{}
|
type NopMetrics struct{}
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -38,6 +38,8 @@ const (
|
||||||
// 慢消费者继续由 best-effort timeout + durable difference 降级。
|
// 慢消费者继续由 best-effort timeout + durable difference 降级。
|
||||||
defaultOutboundQueueSize = 128
|
defaultOutboundQueueSize = 128
|
||||||
defaultOutboundControlQueueSize = 32
|
defaultOutboundControlQueueSize = 32
|
||||||
|
defaultOutboundCriticalQueueSize = 16
|
||||||
|
defaultOutboundBulkQueueSize = 16
|
||||||
defaultOutboundTrackedMaxBytes = int64(512 << 20) // 512 MiB / Server
|
defaultOutboundTrackedMaxBytes = int64(512 << 20) // 512 MiB / Server
|
||||||
defaultOutboundControlMaxBytes = int64(64 << 20) // ack/state/resend vectors / Server
|
defaultOutboundControlMaxBytes = int64(64 << 20) // ack/state/resend vectors / Server
|
||||||
// requiredControlMaxWait bounds protocol barriers such as new_session_created from
|
// requiredControlMaxWait bounds protocol barriers such as new_session_created from
|
||||||
|
|
@ -67,9 +69,27 @@ const (
|
||||||
outboundResendByRequest
|
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 {
|
type outboundOp struct {
|
||||||
kind outboundOpKind
|
kind outboundOpKind
|
||||||
control bool
|
control bool
|
||||||
|
priority outboundPriority
|
||||||
ctx context.Context
|
ctx context.Context
|
||||||
msgType proto.MessageType
|
msgType proto.MessageType
|
||||||
msg bin.Encoder
|
msg bin.Encoder
|
||||||
|
|
@ -83,12 +103,160 @@ type outboundOp struct {
|
||||||
reqMsgID int64
|
reqMsgID int64
|
||||||
enqueuedAt time.Time
|
enqueuedAt time.Time
|
||||||
done chan outboundResult
|
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 {
|
type encodedOutboundMessage struct {
|
||||||
body []byte
|
body []byte
|
||||||
typeID uint32
|
typeID uint32
|
||||||
reqMsgID int64
|
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 {
|
type outboundResult struct {
|
||||||
|
|
@ -314,12 +482,20 @@ func (c *Conn) startOutbound() {
|
||||||
if c.outboundQueueSize <= 0 {
|
if c.outboundQueueSize <= 0 {
|
||||||
c.outboundQueueSize = defaultOutboundQueueSize
|
c.outboundQueueSize = defaultOutboundQueueSize
|
||||||
}
|
}
|
||||||
|
if c.outboundQueueSize < 3 {
|
||||||
|
c.outboundQueueSize = 3
|
||||||
|
}
|
||||||
if c.outboundControlQueueSize <= 0 {
|
if c.outboundControlQueueSize <= 0 {
|
||||||
c.outboundControlQueueSize = defaultOutboundControlQueueSize
|
c.outboundControlQueueSize = defaultOutboundControlQueueSize
|
||||||
}
|
}
|
||||||
c.ensureOutboundTrackedBudget()
|
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.outboundControl = make(chan outboundOp, c.outboundControlQueueSize)
|
||||||
|
c.outboundCritical = make(chan outboundOp, criticalSize)
|
||||||
|
c.outboundBulk = make(chan outboundOp, bulkSize)
|
||||||
c.outboundStop = make(chan struct{})
|
c.outboundStop = make(chan struct{})
|
||||||
c.outboundDone = make(chan struct{})
|
c.outboundDone = make(chan struct{})
|
||||||
go c.outboundLoop()
|
go c.outboundLoop()
|
||||||
|
|
@ -525,8 +701,9 @@ func (c *Conn) sendBestEffort(ctx context.Context, t proto.MessageType, msg bin.
|
||||||
op.enqueuedAt = time.Now()
|
op.enqueuedAt = time.Now()
|
||||||
// 快路径:非阻塞入队。fan-out 每 (conn × push) 都走这里,队列有空位时不为
|
// 快路径:非阻塞入队。fan-out 每 (conn × push) 都走这里,队列有空位时不为
|
||||||
// 本次推送分配任何 timer(此前 timeout>0 无条件 WithTimeout,稳态白建 timer)。
|
// 本次推送分配任何 timer(此前 timeout>0 无条件 WithTimeout,稳态白建 timer)。
|
||||||
|
q := c.outboundQueue(op)
|
||||||
select {
|
select {
|
||||||
case c.outbound <- op:
|
case q <- op:
|
||||||
return nil
|
return nil
|
||||||
case <-c.outboundStop:
|
case <-c.outboundStop:
|
||||||
op.releaseReservation(c.outboundTrackedBudget)
|
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")
|
c.metrics.OutboundDropped("push_queue_full")
|
||||||
return ErrOutboundQueueFull
|
return ErrOutboundQueueFull
|
||||||
}
|
}
|
||||||
c.metrics.OutboundQueueWait(len(c.outbound), cap(c.outbound))
|
c.metrics.OutboundQueueWait(len(q), cap(q))
|
||||||
if ctx == nil {
|
if ctx == nil {
|
||||||
ctx = context.Background()
|
ctx = context.Background()
|
||||||
}
|
}
|
||||||
|
|
@ -549,7 +726,7 @@ func (c *Conn) sendBestEffort(ctx context.Context, t proto.MessageType, msg bin.
|
||||||
timeoutC = timer.C
|
timeoutC = timer.C
|
||||||
}
|
}
|
||||||
select {
|
select {
|
||||||
case c.outbound <- op:
|
case q <- op:
|
||||||
return nil
|
return nil
|
||||||
case <-timeoutC:
|
case <-timeoutC:
|
||||||
op.releaseReservation(c.outboundTrackedBudget)
|
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)
|
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 {
|
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 {
|
if c.outbound == nil || c.outboundControl == nil {
|
||||||
return ErrConnClosed
|
return ErrConnClosed
|
||||||
|
|
@ -774,10 +986,7 @@ func (c *Conn) enqueueOutboundRegistered(ctx context.Context, op outboundOp) err
|
||||||
if c.isRetired() {
|
if c.isRetired() {
|
||||||
return ErrConnClosed
|
return ErrConnClosed
|
||||||
}
|
}
|
||||||
q := c.outbound
|
q := c.outboundQueue(op)
|
||||||
if op.control {
|
|
||||||
q = c.outboundControl
|
|
||||||
}
|
|
||||||
select {
|
select {
|
||||||
case q <- op:
|
case q <- op:
|
||||||
return nil
|
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 {
|
func (c *Conn) beginOutboundEnqueue() bool {
|
||||||
c.outboundEnqueueMu.Lock()
|
c.outboundEnqueueMu.Lock()
|
||||||
defer c.outboundEnqueueMu.Unlock()
|
defer c.outboundEnqueueMu.Unlock()
|
||||||
|
|
@ -814,6 +1037,7 @@ func (c *Conn) endOutboundEnqueue() {
|
||||||
|
|
||||||
func (c *Conn) outboundLoop() {
|
func (c *Conn) outboundLoop() {
|
||||||
state := newOutboundState(c.outboundTrackedBudget)
|
state := newOutboundState(c.outboundTrackedBudget)
|
||||||
|
ordinarySinceBulk := 0
|
||||||
defer func() {
|
defer func() {
|
||||||
// pending frames belong exclusively to this actor. Releasing after drain ensures no
|
// 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.
|
// resend path can race the final budget return and no Conn body survives actor exit.
|
||||||
|
|
@ -826,8 +1050,11 @@ func (c *Conn) outboundLoop() {
|
||||||
c.drainOutbound()
|
c.drainOutbound()
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
select {
|
op, ok := c.nextOutboundOp(&ordinarySinceBulk)
|
||||||
case op := <-c.outboundControl:
|
if !ok {
|
||||||
|
c.drainOutbound()
|
||||||
|
return
|
||||||
|
}
|
||||||
if c.isRetired() {
|
if c.isRetired() {
|
||||||
op.releaseReservation(c.outboundTrackedBudget)
|
op.releaseReservation(c.outboundTrackedBudget)
|
||||||
op.finish(outboundResult{err: ErrConnClosed})
|
op.finish(outboundResult{err: ErrConnClosed})
|
||||||
|
|
@ -841,37 +1068,55 @@ func (c *Conn) outboundLoop() {
|
||||||
c.drainOutbound()
|
c.drainOutbound()
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
continue
|
|
||||||
default:
|
|
||||||
}
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// 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 {
|
select {
|
||||||
case <-c.outboundStop:
|
case <-c.outboundStop:
|
||||||
c.drainOutbound()
|
return outboundOp{}, false
|
||||||
return
|
|
||||||
case op := <-c.outboundControl:
|
case op := <-c.outboundControl:
|
||||||
if c.isRetired() {
|
return op, true
|
||||||
op.releaseReservation(c.outboundTrackedBudget)
|
case op := <-c.outboundCritical:
|
||||||
op.finish(outboundResult{err: ErrConnClosed})
|
return op, true
|
||||||
c.signalOutboundStop()
|
|
||||||
c.drainOutbound()
|
|
||||||
return
|
|
||||||
}
|
|
||||||
c.handleOutboundOp(state, op)
|
|
||||||
case op := <-c.outbound:
|
case op := <-c.outbound:
|
||||||
if c.isRetired() {
|
*ordinarySinceBulk++
|
||||||
op.releaseReservation(c.outboundTrackedBudget)
|
return op, true
|
||||||
op.finish(outboundResult{err: ErrConnClosed})
|
case op := <-c.outboundBulk:
|
||||||
c.signalOutboundStop()
|
*ordinarySinceBulk = 0
|
||||||
c.drainOutbound()
|
return op, true
|
||||||
return
|
|
||||||
}
|
|
||||||
c.handleOutboundOp(state, op)
|
|
||||||
}
|
|
||||||
if c.isRetired() {
|
|
||||||
c.signalOutboundStop()
|
|
||||||
c.drainOutbound()
|
|
||||||
return
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
@ -885,9 +1130,15 @@ func (c *Conn) drainOutbound() {
|
||||||
case op := <-c.outboundControl:
|
case op := <-c.outboundControl:
|
||||||
op.releaseReservation(c.outboundTrackedBudget)
|
op.releaseReservation(c.outboundTrackedBudget)
|
||||||
op.finish(outboundResult{err: ErrConnClosed})
|
op.finish(outboundResult{err: ErrConnClosed})
|
||||||
|
case op := <-c.outboundCritical:
|
||||||
|
op.releaseReservation(c.outboundTrackedBudget)
|
||||||
|
op.finish(outboundResult{err: ErrConnClosed})
|
||||||
case op := <-c.outbound:
|
case op := <-c.outbound:
|
||||||
op.releaseReservation(c.outboundTrackedBudget)
|
op.releaseReservation(c.outboundTrackedBudget)
|
||||||
op.finish(outboundResult{err: ErrConnClosed})
|
op.finish(outboundResult{err: ErrConnClosed})
|
||||||
|
case op := <-c.outboundBulk:
|
||||||
|
op.releaseReservation(c.outboundTrackedBudget)
|
||||||
|
op.finish(outboundResult{err: ErrConnClosed})
|
||||||
default:
|
default:
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
|
@ -902,7 +1153,11 @@ func (c *Conn) handleOutboundOp(state *outboundState, op outboundOp) {
|
||||||
case outboundSend:
|
case outboundSend:
|
||||||
c.handleOutboundSend(state, op)
|
c.handleOutboundSend(state, op)
|
||||||
case outboundAck:
|
case outboundAck:
|
||||||
state.ack(op.ids)
|
for _, reqMsgID := range state.ack(op.ids) {
|
||||||
|
if c.rpcResultAcked != nil {
|
||||||
|
c.rpcResultAcked(c, reqMsgID)
|
||||||
|
}
|
||||||
|
}
|
||||||
case outboundQueryState:
|
case outboundQueryState:
|
||||||
op.finish(outboundResult{info: state.stateInfo(op.ids)})
|
op.finish(outboundResult{info: state.stateInfo(op.ids)})
|
||||||
case outboundResend:
|
case outboundResend:
|
||||||
|
|
@ -917,7 +1172,17 @@ func (c *Conn) handleOutboundOp(state *outboundState, op outboundOp) {
|
||||||
}
|
}
|
||||||
|
|
||||||
func (c *Conn) handleOutboundSend(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
|
reserved := op.reservedBytes
|
||||||
reservationBudget := op.reservationBudget
|
reservationBudget := op.reservationBudget
|
||||||
if reservationBudget == nil {
|
if reservationBudget == nil {
|
||||||
|
|
@ -1018,6 +1283,9 @@ func (c *Conn) handleOutboundResendByRequest(state *outboundState, ctx context.C
|
||||||
}
|
}
|
||||||
|
|
||||||
func (op outboundOp) finish(res outboundResult) {
|
func (op outboundOp) finish(res outboundResult) {
|
||||||
|
if op.terminal != nil {
|
||||||
|
op.terminal(res.err)
|
||||||
|
}
|
||||||
if op.done == nil {
|
if op.done == nil {
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
|
@ -1093,6 +1361,7 @@ func (c *Conn) newOutboundSendOp(ctx context.Context, t proto.MessageType, msg b
|
||||||
kind: outboundSend,
|
kind: outboundSend,
|
||||||
msgType: t,
|
msgType: t,
|
||||||
encoded: encoded,
|
encoded: encoded,
|
||||||
|
priority: classifyOutboundPriority(encoded, priorityControl),
|
||||||
reservedBytes: bytes,
|
reservedBytes: bytes,
|
||||||
reservationBudget: budget,
|
reservationBudget: budget,
|
||||||
}, nil
|
}, nil
|
||||||
|
|
@ -1112,11 +1381,25 @@ func (c *Conn) newOutboundSendOp(ctx context.Context, t proto.MessageType, msg b
|
||||||
kind: outboundSend,
|
kind: outboundSend,
|
||||||
msgType: t,
|
msgType: t,
|
||||||
encoded: encoded,
|
encoded: encoded,
|
||||||
|
priority: classifyOutboundPriority(encoded, priorityControl),
|
||||||
reservedBytes: bytes,
|
reservedBytes: bytes,
|
||||||
reservationBudget: budget,
|
reservationBudget: budget,
|
||||||
}, nil
|
}, 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 {
|
func (c *Conn) outboundMessageBudget(typeID uint32, priorityControl bool) *outboundTrackedBudget {
|
||||||
if priorityControl || encodedControlFrame(typeID) {
|
if priorityControl || encodedControlFrame(typeID) {
|
||||||
return c.ensureOutboundControlTrackedBudget()
|
return c.ensureOutboundControlTrackedBudget()
|
||||||
|
|
@ -1210,7 +1493,11 @@ func (c *Conn) downgradedCloneContext(ctx context.Context, encoded *encodedOutbo
|
||||||
if sameBacking(down, encoded.body) {
|
if sameBacking(down, encoded.body) {
|
||||||
return encoded // 直通:未变(mt.*/顶层未知),无需拷贝或重算 typeID
|
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 {
|
if id, e := (&bin.Buffer{Buf: down}).PeekID(); e == nil {
|
||||||
out.typeID = id
|
out.typeID = id
|
||||||
}
|
}
|
||||||
|
|
@ -1535,8 +1822,16 @@ func (s *outboundState) addReserved(frame *outboundFrame) int {
|
||||||
return s.shrinkPending()
|
return s.shrinkPending()
|
||||||
}
|
}
|
||||||
|
|
||||||
func (s *outboundState) ack(ids []int64) {
|
func (s *outboundState) ack(ids []int64) []int64 {
|
||||||
|
var requestIDs []int64
|
||||||
for _, id := range ids {
|
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) {
|
if !s.removePending(id) {
|
||||||
continue
|
continue
|
||||||
}
|
}
|
||||||
|
|
@ -1545,6 +1840,7 @@ func (s *outboundState) ack(ids []int64) {
|
||||||
if len(s.order) > s.maxMessages*2 {
|
if len(s.order) > s.maxMessages*2 {
|
||||||
s.compactOrder()
|
s.compactOrder()
|
||||||
}
|
}
|
||||||
|
return requestIDs
|
||||||
}
|
}
|
||||||
|
|
||||||
func (s *outboundState) stateInfo(ids []int64) []byte {
|
func (s *outboundState) stateInfo(ids []int64) []byte {
|
||||||
|
|
|
||||||
|
|
@ -397,8 +397,8 @@ func TestOutboundQueueBackingUsesSmallConfigurableBounds(t *testing.T) {
|
||||||
c := &Conn{metrics: NopMetrics{}}
|
c := &Conn{metrics: NopMetrics{}}
|
||||||
c.startOutbound()
|
c.startOutbound()
|
||||||
defer c.Close()
|
defer c.Close()
|
||||||
if got := cap(c.outbound); got != defaultOutboundQueueSize {
|
if got := cap(c.outbound) + cap(c.outboundCritical) + cap(c.outboundBulk); got != defaultOutboundQueueSize {
|
||||||
t.Fatalf("normal queue cap = %d, want %d", got, defaultOutboundQueueSize)
|
t.Fatalf("ordinary lane total cap = %d, want %d", got, defaultOutboundQueueSize)
|
||||||
}
|
}
|
||||||
if got := cap(c.outboundControl); got != defaultOutboundControlQueueSize {
|
if got := cap(c.outboundControl); got != defaultOutboundControlQueueSize {
|
||||||
t.Fatalf("control queue cap = %d, want %d", got, defaultOutboundControlQueueSize)
|
t.Fatalf("control queue cap = %d, want %d", got, defaultOutboundControlQueueSize)
|
||||||
|
|
@ -413,8 +413,8 @@ func TestOutboundQueueBackingUsesSmallConfigurableBounds(t *testing.T) {
|
||||||
}
|
}
|
||||||
c.startOutbound()
|
c.startOutbound()
|
||||||
defer c.Close()
|
defer c.Close()
|
||||||
if got := cap(c.outbound); got != 7 {
|
if got := cap(c.outbound) + cap(c.outboundCritical) + cap(c.outboundBulk); got != 7 {
|
||||||
t.Fatalf("normal queue cap = %d, want 7", got)
|
t.Fatalf("ordinary lane total cap = %d, want 7", got)
|
||||||
}
|
}
|
||||||
if got := cap(c.outboundControl); got != 3 {
|
if got := cap(c.outboundControl); got != 3 {
|
||||||
t.Fatalf("control queue cap = %d, want 3", got)
|
t.Fatalf("control queue cap = %d, want 3", got)
|
||||||
|
|
@ -446,9 +446,11 @@ func TestServerNewConnectionsShareOutboundBudgetAndQueueLimits(t *testing.T) {
|
||||||
defer c1.Close()
|
defer c1.Close()
|
||||||
defer c2.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",
|
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 {
|
if c1.outboundTrackedBudget != srv.outboundTrackedBudget || c2.outboundTrackedBudget != srv.outboundTrackedBudget {
|
||||||
t.Fatal("server connections did not receive the shared outbound tracking budget")
|
t.Fatal("server connections did not receive the shared outbound tracking budget")
|
||||||
|
|
|
||||||
|
|
@ -3,6 +3,7 @@ package mtprotoedge
|
||||||
import (
|
import (
|
||||||
"context"
|
"context"
|
||||||
"errors"
|
"errors"
|
||||||
|
"sync"
|
||||||
"sync/atomic"
|
"sync/atomic"
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
@ -44,9 +45,12 @@ type rpcResultFlight struct {
|
||||||
done chan struct{}
|
done chan struct{}
|
||||||
encoded *encodedOutboundMessage
|
encoded *encodedOutboundMessage
|
||||||
ok bool
|
ok bool
|
||||||
|
subscribers []func(*encodedOutboundMessage, bool)
|
||||||
}
|
}
|
||||||
|
|
||||||
type rpcResultWaiter struct {
|
type rpcResultWaiter struct {
|
||||||
|
cache *rpcResultCache
|
||||||
|
key rpcResultCacheKey
|
||||||
flight *rpcResultFlight
|
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 {
|
type rpcResultOwnerLease struct {
|
||||||
cache *rpcResultCache
|
cache *rpcResultCache
|
||||||
key rpcResultCacheKey
|
key rpcResultCacheKey
|
||||||
flight *rpcResultFlight
|
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
|
// 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 {
|
if l == nil || l.cache == nil || l.flight == nil {
|
||||||
return false
|
return false
|
||||||
}
|
}
|
||||||
|
if l.handedOff.Load() {
|
||||||
|
return false
|
||||||
|
}
|
||||||
s := l.cache.shard(l.key)
|
s := l.cache.shard(l.key)
|
||||||
s.mu.Lock()
|
s.mu.Lock()
|
||||||
defer s.mu.Unlock()
|
if l.handedOff.Load() {
|
||||||
|
s.mu.Unlock()
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
|
||||||
flight, ok := s.pending[l.key]
|
flight, ok := s.pending[l.key]
|
||||||
if !ok || flight != l.flight {
|
if !ok || flight != l.flight {
|
||||||
|
s.mu.Unlock()
|
||||||
return false
|
return false
|
||||||
}
|
}
|
||||||
delete(s.pending, l.key)
|
delete(s.pending, l.key)
|
||||||
l.cache.flightLimit.release()
|
l.cache.flightLimit.release()
|
||||||
|
subscribers := append([]func(*encodedOutboundMessage, bool){}, flight.subscribers...)
|
||||||
|
flight.subscribers = nil
|
||||||
close(flight.done)
|
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
|
return true
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
@ -168,7 +297,7 @@ func (c *rpcResultCache) Acquire(authKeyID [8]byte, sessionID, reqMsgID int64) (
|
||||||
if flight, ok := s.pending[key]; ok {
|
if flight, ok := s.pending[key]; ok {
|
||||||
return rpcResultAcquire{
|
return rpcResultAcquire{
|
||||||
state: rpcResultAcquirePending,
|
state: rpcResultAcquirePending,
|
||||||
waiter: &rpcResultWaiter{flight: flight},
|
waiter: &rpcResultWaiter{cache: c, key: key, flight: flight},
|
||||||
}, nil
|
}, nil
|
||||||
}
|
}
|
||||||
if !c.flightLimit.reserve() {
|
if !c.flightLimit.reserve() {
|
||||||
|
|
@ -181,7 +310,9 @@ func (c *rpcResultCache) Acquire(authKeyID [8]byte, sessionID, reqMsgID int64) (
|
||||||
s.pending[key] = flight
|
s.pending[key] = flight
|
||||||
return rpcResultAcquire{
|
return rpcResultAcquire{
|
||||||
state: rpcResultAcquireOwner,
|
state: rpcResultAcquireOwner,
|
||||||
owner: &rpcResultOwnerLease{cache: c, key: key, flight: flight},
|
owner: &rpcResultOwnerLease{
|
||||||
|
cache: c, key: key, flight: flight, delivery: newRPCResultDelivery(reqMsgID),
|
||||||
|
},
|
||||||
}, nil
|
}, nil
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
@ -191,17 +322,20 @@ func (c *rpcResultCache) completeRPCResultFlightLocked(
|
||||||
s *rpcResultCacheShard,
|
s *rpcResultCacheShard,
|
||||||
key rpcResultCacheKey,
|
key rpcResultCacheKey,
|
||||||
encoded *encodedOutboundMessage,
|
encoded *encodedOutboundMessage,
|
||||||
) {
|
) []func(*encodedOutboundMessage, bool) {
|
||||||
if c == nil || s == nil || encoded == nil {
|
if c == nil || s == nil || encoded == nil {
|
||||||
return
|
return nil
|
||||||
}
|
}
|
||||||
flight, ok := s.pending[key]
|
flight, ok := s.pending[key]
|
||||||
if !ok {
|
if !ok {
|
||||||
return
|
return nil
|
||||||
}
|
}
|
||||||
delete(s.pending, key)
|
delete(s.pending, key)
|
||||||
flight.encoded = encoded
|
flight.encoded = encoded
|
||||||
flight.ok = true
|
flight.ok = true
|
||||||
c.flightLimit.release()
|
c.flightLimit.release()
|
||||||
|
subscribers := append([]func(*encodedOutboundMessage, bool){}, flight.subscribers...)
|
||||||
|
flight.subscribers = nil
|
||||||
close(flight.done)
|
close(flight.done)
|
||||||
|
return subscribers
|
||||||
}
|
}
|
||||||
|
|
|
||||||
|
|
@ -118,7 +118,6 @@ func (c *rpcResultCache) Put(authKeyID [8]byte, sessionID, reqMsgID int64, encod
|
||||||
now := s.now()
|
now := s.now()
|
||||||
|
|
||||||
s.mu.Lock()
|
s.mu.Lock()
|
||||||
defer s.mu.Unlock()
|
|
||||||
|
|
||||||
if cacheable {
|
if cacheable {
|
||||||
s.expireLocked(now)
|
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
|
// Resolve the independent in-flight entry only after the completed cache has
|
||||||
// been published. Waiters awakened by this close can therefore immediately
|
// been published. Waiters awakened by this close can therefore immediately
|
||||||
// observe either the shared encoded result or the completed Get entry.
|
// 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) {
|
func (s *rpcResultCacheShard) expireLocked(now time.Time) {
|
||||||
|
|
|
||||||
110
internal/mtprotoedge/rpc_result_egress.go
Normal file
110
internal/mtprotoedge/rpc_result_egress.go
Normal 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"
|
||||||
|
}
|
||||||
|
}
|
||||||
444
internal/mtprotoedge/rpc_result_egress_test.go
Normal file
444
internal/mtprotoedge/rpc_result_egress_test.go
Normal 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)
|
||||||
|
}
|
||||||
|
}
|
||||||
455
internal/mtprotoedge/rpc_rewrap.go
Normal file
455
internal/mtprotoedge/rpc_rewrap.go
Normal 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)))
|
||||||
|
}
|
||||||
299
internal/mtprotoedge/rpc_rewrap_test.go
Normal file
299
internal/mtprotoedge/rpc_rewrap_test.go
Normal 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)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
@ -34,6 +34,7 @@ type samePortMux struct {
|
||||||
base net.Listener
|
base net.Listener
|
||||||
addr net.Addr
|
addr net.Addr
|
||||||
sniffTimeout time.Duration
|
sniffTimeout time.Duration
|
||||||
|
observe func(connectionIntakeEvent)
|
||||||
|
|
||||||
tcp *samePortMuxListener
|
tcp *samePortMuxListener
|
||||||
http *samePortMuxListener
|
http *samePortMuxListener
|
||||||
|
|
@ -49,7 +50,7 @@ type samePortMux struct {
|
||||||
sniffing map[net.Conn]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 {
|
if sniffTimeout <= 0 {
|
||||||
sniffTimeout = 5 * time.Second
|
sniffTimeout = 5 * time.Second
|
||||||
}
|
}
|
||||||
|
|
@ -60,6 +61,9 @@ func newSamePortMux(base net.Listener, sniffTimeout time.Duration) *samePortMux
|
||||||
closed: make(chan struct{}),
|
closed: make(chan struct{}),
|
||||||
sniffing: make(map[net.Conn]struct{}),
|
sniffing: make(map[net.Conn]struct{}),
|
||||||
}
|
}
|
||||||
|
if len(observers) > 0 {
|
||||||
|
m.observe = observers[0]
|
||||||
|
}
|
||||||
m.tcp = newSamePortMuxListener(m.addr, m.closed)
|
m.tcp = newSamePortMuxListener(m.addr, m.closed)
|
||||||
m.http = newSamePortMuxListener(m.addr, m.closed)
|
m.http = newSamePortMuxListener(m.addr, m.closed)
|
||||||
return m
|
return m
|
||||||
|
|
@ -144,10 +148,13 @@ func (m *samePortMux) Close() error {
|
||||||
// dispatch 窥探单条连接的前 4 字节并把它交给 tcp 或 http 子 listener。窥探带 sniffTimeout
|
// dispatch 窥探单条连接的前 4 字节并把它交给 tcp 或 http 子 listener。窥探带 sniffTimeout
|
||||||
// 读上界,慢/半开连接最多占用本 goroutine sniffTimeout 后即被回收。
|
// 读上界,慢/半开连接最多占用本 goroutine sniffTimeout 后即被回收。
|
||||||
func (m *samePortMux) dispatch(ctx context.Context, conn net.Conn) {
|
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
|
// 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
|
// 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.
|
// mux shutdown can close this read immediately. finishSniff removes the socket before hand-off.
|
||||||
if !m.beginSniff(conn) {
|
if !m.beginSniff(conn) {
|
||||||
|
m.observeEvent(connectionIntakeEvent{stage: "mux_sniff", outcome: "closed", remote: remote, local: local, duration: time.Since(started)})
|
||||||
_ = conn.Close()
|
_ = conn.Close()
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
|
@ -160,10 +167,16 @@ func (m *samePortMux) dispatch(ctx context.Context, conn net.Conn) {
|
||||||
|
|
||||||
var header [4]byte
|
var header [4]byte
|
||||||
if err := conn.SetReadDeadline(time.Now().Add(m.sniffTimeout)); err != nil {
|
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()
|
_ = conn.Close()
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
if _, err := io.ReadFull(conn, header[:]); err != nil {
|
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()
|
_ = conn.Close()
|
||||||
return
|
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
|
// 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.
|
// socket, or it can no longer find it. A concurrently closed mux refuses delivery.
|
||||||
if !m.finishSniff(conn) {
|
if !m.finishSniff(conn) {
|
||||||
|
m.observeEvent(connectionIntakeEvent{stage: "mux_sniff", outcome: "closed", remote: remote, local: local, duration: time.Since(started)})
|
||||||
_ = conn.Close()
|
_ = conn.Close()
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
finishedSniff = true
|
finishedSniff = true
|
||||||
if err := conn.SetReadDeadline(time.Time{}); err != nil {
|
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()
|
_ = conn.Close()
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
|
@ -186,14 +201,24 @@ func (m *samePortMux) dispatch(ctx context.Context, conn net.Conn) {
|
||||||
}
|
}
|
||||||
|
|
||||||
target := m.tcp
|
target := m.tcp
|
||||||
|
transport := "tcp"
|
||||||
if isHTTPHeaderPrefix(header) {
|
if isHTTPHeaderPrefix(header) {
|
||||||
target = m.http
|
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) {
|
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()
|
_ = conn.Close()
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func (m *samePortMux) observeEvent(event connectionIntakeEvent) {
|
||||||
|
if m.observe != nil {
|
||||||
|
m.observe(event)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
func (m *samePortMux) beginSniff(conn net.Conn) bool {
|
func (m *samePortMux) beginSniff(conn net.Conn) bool {
|
||||||
m.sniffMu.Lock()
|
m.sniffMu.Lock()
|
||||||
defer m.sniffMu.Unlock()
|
defer m.sniffMu.Unlock()
|
||||||
|
|
|
||||||
|
|
@ -43,6 +43,27 @@ type RPCHandler interface {
|
||||||
NegotiatedLayer(authKeyID [8]byte, sessionID int64) (int, bool)
|
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。
|
// Options 配置 Server。
|
||||||
type Options struct {
|
type Options struct {
|
||||||
// Logger 日志器。默认 zap.NewNop()。
|
// Logger 日志器。默认 zap.NewNop()。
|
||||||
|
|
@ -126,6 +147,10 @@ type Options struct {
|
||||||
RPC RPCHandler
|
RPC RPCHandler
|
||||||
// Metrics 接收连接层指标。默认 NopMetrics。
|
// Metrics 接收连接层指标。默认 NopMetrics。
|
||||||
Metrics Metrics
|
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 用于消息 ID 与时间戳。默认 clock.System。
|
||||||
Clock clock.Clock
|
Clock clock.Clock
|
||||||
// Rand 随机源。默认 crypto.DefaultRand()。
|
// Rand 随机源。默认 crypto.DefaultRand()。
|
||||||
|
|
@ -239,6 +264,7 @@ type Server struct {
|
||||||
conns *SessionManager
|
conns *SessionManager
|
||||||
rpc RPCHandler
|
rpc RPCHandler
|
||||||
metrics Metrics
|
metrics Metrics
|
||||||
|
onServing func(net.Addr)
|
||||||
cipher crypto.Cipher
|
cipher crypto.Cipher
|
||||||
clock clock.Clock
|
clock clock.Clock
|
||||||
rand io.Reader
|
rand io.Reader
|
||||||
|
|
@ -246,6 +272,7 @@ type Server struct {
|
||||||
admission *admissionController
|
admission *admissionController
|
||||||
|
|
||||||
rpcResults *rpcResultCache
|
rpcResults *rpcResultCache
|
||||||
|
rpcRewrap *rpcRewrapRegistry
|
||||||
|
|
||||||
// onFrame 是测试钩子:收到一帧时回调其字节数;生产为 nil。
|
// onFrame 是测试钩子:收到一帧时回调其字节数;生产为 nil。
|
||||||
onFrame func(n int)
|
onFrame func(n int)
|
||||||
|
|
@ -284,15 +311,30 @@ func New(opts Options) *Server {
|
||||||
conns: conns,
|
conns: conns,
|
||||||
rpc: opts.RPC,
|
rpc: opts.RPC,
|
||||||
metrics: opts.Metrics,
|
metrics: opts.Metrics,
|
||||||
|
onServing: opts.OnServing,
|
||||||
cipher: crypto.NewServerCipher(opts.Rand),
|
cipher: crypto.NewServerCipher(opts.Rand),
|
||||||
clock: opts.Clock,
|
clock: opts.Clock,
|
||||||
rand: opts.Rand,
|
rand: opts.Rand,
|
||||||
types: tmap.New(tg.TypesMap(), mt.TypesMap(), proto.TypesMap()),
|
types: tmap.New(tg.TypesMap(), mt.TypesMap(), proto.TypesMap()),
|
||||||
rpcResults: newRPCResultCacheWithFlightLimit(opts.Clock.Now, opts.RPCGlobalMaxTasks),
|
rpcResults: newRPCResultCacheWithFlightLimit(opts.Clock.Now, opts.RPCGlobalMaxTasks),
|
||||||
|
rpcRewrap: newRPCRewrapRegistry(opts.RPCGlobalMaxTasks),
|
||||||
admission: newAdmissionController(opts.MaxConnections, opts.MaxConnectionsPerIP, opts.MaxConcurrentHandshakes),
|
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 基于一次解密结果创建一个可发送的连接对象。
|
// newConn 基于一次解密结果创建一个可发送的连接对象。
|
||||||
func (s *Server) newConn(tc transport.Conn, key crypto.AuthKey, sessionID, salt int64) *Conn {
|
func (s *Server) newConn(tc transport.Conn, key crypto.AuthKey, sessionID, salt int64) *Conn {
|
||||||
if lease, ok := tc.(*physicalTransportLease); ok {
|
if lease, ok := tc.(*physicalTransportLease); ok {
|
||||||
|
|
@ -338,6 +380,7 @@ func (s *Server) buildConn(tc transport.Conn, lease *physicalTransportLease, key
|
||||||
outboundTrackedBudget: s.outboundTrackedBudget,
|
outboundTrackedBudget: s.outboundTrackedBudget,
|
||||||
outboundControlTrackedBudget: s.outboundControlBudget,
|
outboundControlTrackedBudget: s.outboundControlBudget,
|
||||||
outboundScratchPool: s.outboundScratchPool,
|
outboundScratchPool: s.outboundScratchPool,
|
||||||
|
rpcResultAcked: s.rpcRewrap.acknowledge,
|
||||||
}
|
}
|
||||||
c.startOutbound()
|
c.startOutbound()
|
||||||
c.startInboundRPCScheduler(s.rpcScheduler, s.rpcInflight, s.rpcQueueSize, s.rpcTimeout)
|
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)
|
defer s.rpcScheduler.stop(rpcCloseWaitTimeout)
|
||||||
// 只在最外层 listener 包一次,确保 same-port mux 的 sniff/HTTP upgrade 也计入
|
// 只在最外层 listener 包一次,确保 same-port mux 的 sniff/HTTP upgrade 也计入
|
||||||
// raw admission,而不是等连接已经分流后才计数。
|
// raw admission,而不是等连接已经分流后才计数。
|
||||||
ln = s.admission.wrapListener(ln)
|
ln = s.observeRawAccepts(s.admission.wrapListener(ln))
|
||||||
if s.websocket {
|
if s.websocket {
|
||||||
return s.serveMixed(ctx, ln)
|
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))
|
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")
|
defer s.log.Info("Stopped")
|
||||||
|
errCh := make(chan error, 1)
|
||||||
return s.acceptLoop(ctx, ln, s.obfuscated)
|
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 {
|
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 误杀,
|
// minDuration(5s,...) 把嗅探压到 5s,比非 mux 路径激进 12 倍,会把这些暖连接在 5s 误杀,
|
||||||
// 触发客户端 6s 重连风暴并误判「后端不健康」回退到外部 DNS。per-conn goroutine 模型已消解
|
// 触发客户端 6s 重连风暴并误判「后端不健康」回退到外部 DNS。per-conn goroutine 模型已消解
|
||||||
// slow-loris 接入饥饿,故嗅探用满 handshakeTimeout 是安全的。
|
// slow-loris 接入饥饿,故嗅探用满 handshakeTimeout 是安全的。
|
||||||
mux := newSamePortMux(ln, s.handshakeTimeout)
|
mux := newSamePortMux(ln, s.handshakeTimeout, s.observeConnectionIntake)
|
||||||
wsRawLn, wsHandler := transport.WebsocketListener(ln.Addr())
|
wsRawLn, wsHandler := transport.WebsocketListener(ln.Addr())
|
||||||
wsLn := newTransportPacketMessageListener(wsRawLn)
|
wsLn := newTransportPacketMessageListener(wsRawLn)
|
||||||
|
|
||||||
|
|
@ -430,7 +477,7 @@ func (s *Server) serveMixed(ctx context.Context, ln net.Listener) error {
|
||||||
// WebSocket:gotd 升级处理器已剥离 obfuscated2 并补回 codec tag,这里只需探测 codec。
|
// WebSocket:gotd 升级处理器已剥离 obfuscated2 并补回 codec tag,这里只需探测 codec。
|
||||||
go func() {
|
go func() {
|
||||||
defer wg.Done()
|
defer wg.Done()
|
||||||
errCh <- s.acceptLoop(ctx, wsLn, false)
|
errCh <- s.acceptLoopTransport(ctx, wsLn, false, "websocket")
|
||||||
}()
|
}()
|
||||||
go func() {
|
go func() {
|
||||||
defer wg.Done()
|
defer wg.Done()
|
||||||
|
|
@ -444,6 +491,7 @@ func (s *Server) serveMixed(ctx context.Context, ln net.Listener) error {
|
||||||
}
|
}
|
||||||
errCh <- nil
|
errCh <- nil
|
||||||
}()
|
}()
|
||||||
|
s.signalServing(ln.Addr())
|
||||||
|
|
||||||
// The four services form one lifecycle: even a clean/closed-listener return from any one
|
// 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
|
// component means the remaining three can no longer make forward progress as a complete
|
||||||
|
|
@ -468,6 +516,10 @@ func (s *Server) serveMixed(ctx context.Context, ln net.Listener) error {
|
||||||
// 整个监听循环。obfuscated 为 true 时先走 obfuscated2 去混淆(裸 MTProto TCP);WebSocket
|
// 整个监听循环。obfuscated 为 true 时先走 obfuscated2 去混淆(裸 MTProto TCP);WebSocket
|
||||||
// 连接传 false(gotd 升级处理器已完成去混淆)。
|
// 连接传 false(gotd 升级处理器已完成去混淆)。
|
||||||
func (s *Server) acceptLoop(ctx context.Context, ln net.Listener, obfuscated bool) error {
|
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)
|
ctx, cancel := context.WithCancel(ctx)
|
||||||
var wg sync.WaitGroup
|
var wg sync.WaitGroup
|
||||||
defer func() {
|
defer func() {
|
||||||
|
|
@ -505,7 +557,14 @@ func (s *Server) acceptLoop(ctx context.Context, ln net.Listener, obfuscated boo
|
||||||
wg.Add(1)
|
wg.Add(1)
|
||||||
go func() {
|
go func() {
|
||||||
defer wg.Done()
|
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
|
// serveDetectedConn 把一条裸连接提升为 transport.Conn(去混淆 + codec 探测)后运行 MTProto
|
||||||
// 连接循环。提升过程的读取放在本 goroutine、且受握手读超时约束,而非塞在 accept 循环里,
|
// 连接循环。提升过程的读取放在本 goroutine、且受握手读超时约束,而非塞在 accept 循环里,
|
||||||
// 这样慢连接不会阻塞其他连接接入,去混淆/codec 握手本身也有时间上界。
|
// 这样慢连接不会阻塞其他连接接入,去混淆/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 语义),
|
// 握手读超时只覆盖去混淆 + codec 探测这一小段;用真实墙钟时间(SetReadDeadline 语义),
|
||||||
// 不走可能被测试注入的逻辑 clock。
|
// 不走可能被测试注入的逻辑 clock。
|
||||||
if err := raw.SetReadDeadline(time.Now().Add(s.handshakeTimeout)); err != nil {
|
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)
|
conn, err := s.promoteConn(raw, obfuscated)
|
||||||
close(promoted)
|
close(promoted)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
// 去混淆/codec 探测失败(读超时、客户端中途断开、坏 init 等)只影响这一条连接,
|
outcome := "error"
|
||||||
// 记 debug 即可。
|
if isClientDisconnect(err) {
|
||||||
if !isClientDisconnect(err) {
|
outcome = "client_disconnect"
|
||||||
s.log.Debug("Transport handshake failed", zap.Error(err))
|
|
||||||
}
|
}
|
||||||
|
s.observeConnectionIntake(connectionIntakeEvent{
|
||||||
|
stage: "transport_promote",
|
||||||
|
outcome: outcome,
|
||||||
|
transport: transportName,
|
||||||
|
remote: remote,
|
||||||
|
local: local,
|
||||||
|
duration: time.Since(started),
|
||||||
|
err: err,
|
||||||
|
})
|
||||||
_ = raw.Close()
|
_ = raw.Close()
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
s.observeConnectionIntake(connectionIntakeEvent{
|
||||||
|
stage: "transport_promote",
|
||||||
|
outcome: "ready",
|
||||||
|
transport: transportName,
|
||||||
|
remote: remote,
|
||||||
|
local: local,
|
||||||
|
duration: time.Since(started),
|
||||||
|
})
|
||||||
// 探测完成,撤掉握手读超时;后续每帧读写由 serveConn / 传输层各自管理超时。
|
// 探测完成,撤掉握手读超时;后续每帧读写由 serveConn / 传输层各自管理超时。
|
||||||
if err := raw.SetReadDeadline(time.Time{}); err != nil {
|
if err := raw.SetReadDeadline(time.Time{}); err != nil {
|
||||||
_ = conn.Close()
|
_ = conn.Close()
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
if err := s.serveConn(ctx, conn); err != nil && !isClientDisconnect(err) {
|
if err := s.serveConn(ctx, conn, remote, local); err != nil && !isClientDisconnect(err) {
|
||||||
s.log.Info("Connection closed with error", zap.Error(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,促使客户端重新握手。
|
// - auth_key_id 未注册:回 AuthKeyNotFound,促使客户端重新握手。
|
||||||
//
|
//
|
||||||
// 连接建立 session 后注册到 SessionManager,结束时注销。
|
// 连接建立 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)
|
transportOwner, conn := newPhysicalTransportOwner(raw)
|
||||||
s.metrics.ConnOpened()
|
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
|
var current *Conn
|
||||||
|
firstFrameStarted := time.Now()
|
||||||
|
firstFrameSeen := false
|
||||||
defer func() {
|
defer func() {
|
||||||
// A successful Recv transfers the frame reservation to serveConn. Release it only after
|
// 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
|
// 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()
|
current.Close()
|
||||||
}
|
}
|
||||||
s.metrics.ConnClosed()
|
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 阻塞。
|
// ctx 取消或处理结束时关闭连接,解除 Recv 阻塞。
|
||||||
|
|
@ -624,8 +703,25 @@ func (s *Server) serveConn(ctx context.Context, raw transport.Conn) (err error)
|
||||||
timeout = s.handshakeTimeout
|
timeout = s.handshakeTimeout
|
||||||
}
|
}
|
||||||
if err := s.recv(ctx, conn, &b, timeout); err != nil {
|
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
|
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 {
|
if s.onFrame != nil {
|
||||||
s.onFrame(b.Len())
|
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 {
|
func isClientDisconnect(err error) bool {
|
||||||
switch {
|
switch {
|
||||||
case errors.Is(err, io.EOF),
|
case errors.Is(err, io.EOF),
|
||||||
|
errors.Is(err, io.ErrUnexpectedEOF),
|
||||||
errors.Is(err, net.ErrClosed),
|
errors.Is(err, net.ErrClosed),
|
||||||
errors.Is(err, context.Canceled),
|
errors.Is(err, context.Canceled),
|
||||||
errors.Is(err, context.DeadlineExceeded):
|
errors.Is(err, context.DeadlineExceeded):
|
||||||
|
|
|
||||||
|
|
@ -492,29 +492,28 @@ func TestCrossConnectionInflightRPCHasOneBusinessOwnerAndReplaysResult(t *testin
|
||||||
if got := handler.calls.Load(); got != 1 {
|
if got := handler.calls.Load(); got != 1 {
|
||||||
t.Fatalf("overlapping reconnect executed %d business handlers, want 1", got)
|
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
|
var second handleResult
|
||||||
select {
|
select {
|
||||||
case second = <-secondDone:
|
case second = <-secondDone:
|
||||||
case <-time.After(3 * time.Second):
|
case <-time.After(time.Second):
|
||||||
t.Fatal("duplicate reconnect did not receive owner result")
|
t.Fatal("event-driven duplicate admission blocked the replacement read loop")
|
||||||
}
|
}
|
||||||
if second.err != nil {
|
if second.err != nil {
|
||||||
t.Fatalf("second handleEncrypted: %v", second.err)
|
t.Fatalf("second handleEncrypted: %v", second.err)
|
||||||
}
|
}
|
||||||
if second.conn == nil || !second.conn.isActive() {
|
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 {
|
if got := handler.calls.Load(); got != 1 {
|
||||||
t.Fatalf("cross-connection duplicate business calls = %d, want 1", got)
|
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
|
resultCount := 0
|
||||||
for _, wire := range secondTransport.snapshot() {
|
for _, wire := range secondTransport.snapshot() {
|
||||||
data, decryptErr := crypto.NewClientCipher(rand.Reader).DecryptFromBuffer(key, &bin.Buffer{Buf: wire})
|
data, decryptErr := crypto.NewClientCipher(rand.Reader).DecryptFromBuffer(key, &bin.Buffer{Buf: wire})
|
||||||
|
|
@ -582,9 +581,34 @@ func TestCrossConnectionInflightAbortRetriesOnlyAfterOldOwnerStops(t *testing.T)
|
||||||
secondState := newConnState()
|
secondState := newConnState()
|
||||||
var secondPlain bin.Buffer
|
var secondPlain bin.Buffer
|
||||||
secondConn, err := s.handleEncrypted(context.Background(), secondTransport, secondState, nil, &stored, secondFrame, &secondPlain)
|
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)
|
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)
|
waitForAtomicCalls(t, &handler.calls, 2)
|
||||||
waitInboundRPCBatchBudget(t, s.rpcScheduler, 0, 0)
|
waitInboundRPCBatchBudget(t, s.rpcScheduler, 0, 0)
|
||||||
if got := handler.max.Load(); got != 1 {
|
if got := handler.max.Load(); got != 1 {
|
||||||
|
|
@ -595,7 +619,11 @@ func TestCrossConnectionInflightAbortRetriesOnlyAfterOldOwnerStops(t *testing.T)
|
||||||
}
|
}
|
||||||
|
|
||||||
resultCount := 0
|
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})
|
data, decryptErr := crypto.NewClientCipher(rand.Reader).DecryptFromBuffer(key, &bin.Buffer{Buf: wire})
|
||||||
if decryptErr != nil {
|
if decryptErr != nil {
|
||||||
t.Fatalf("decrypt sequential-retry reply: %v", decryptErr)
|
t.Fatalf("decrypt sequential-retry reply: %v", decryptErr)
|
||||||
|
|
@ -612,10 +640,10 @@ func TestCrossConnectionInflightAbortRetriesOnlyAfterOldOwnerStops(t *testing.T)
|
||||||
if resultCount != 1 {
|
if resultCount != 1 {
|
||||||
t.Fatalf("sequential retry result count = %d, want 1", resultCount)
|
t.Fatalf("sequential retry result count = %d, want 1", resultCount)
|
||||||
}
|
}
|
||||||
if !firstConn.isRetired() || secondConn == nil || !secondConn.isActive() {
|
if !firstConn.isRetired() || !secondConn.isRetired() || thirdConn == nil || !thirdConn.isActive() {
|
||||||
t.Fatalf("replacement lifecycle = old:%v new:%p active:%v", firstConn.lifecycleState(), secondConn, secondConn != nil && secondConn.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) {
|
func waitForAtomicCalls(t *testing.T, calls interface{ Load() int32 }, want int32) {
|
||||||
|
|
|
||||||
|
|
@ -26,6 +26,9 @@ const quickAckResponseFlag = uint32(1 << 31)
|
||||||
const (
|
const (
|
||||||
maxCompatPacketOverhead = 7 // 4-byte header + up to 3 bytes padded-intermediate padding.
|
maxCompatPacketOverhead = 7 // 4-byte header + up to 3 bytes padded-intermediate padding.
|
||||||
maxRetainedDirectMessageScratch = 64 << 10
|
maxRetainedDirectMessageScratch = 64 << 10
|
||||||
|
progressiveWriteMinBytes = 64 << 10
|
||||||
|
progressiveWriteChunkBytes = 32 << 10
|
||||||
|
progressiveWriteIdleTimeout = 10 * time.Second
|
||||||
)
|
)
|
||||||
|
|
||||||
type transportListener interface {
|
type transportListener interface {
|
||||||
|
|
@ -237,12 +240,83 @@ func (c *compatTransportConn) sendDeadline(deadline time.Time, b *bin.Buffer, sc
|
||||||
}
|
}
|
||||||
return nil
|
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 errors.Wrap(err, "write")
|
||||||
}
|
}
|
||||||
return nil
|
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 {
|
func (c *compatTransportConn) writeTransportPacketMessage(b *bin.Buffer, scratch *[]byte) error {
|
||||||
required := b.Len() + maxCompatPacketOverhead
|
required := b.Len() + maxCompatPacketOverhead
|
||||||
if cap(*scratch) < required {
|
if cap(*scratch) < required {
|
||||||
|
|
|
||||||
|
|
@ -15,6 +15,7 @@ type messageWriteTestConn struct {
|
||||||
bytes.Buffer
|
bytes.Buffer
|
||||||
writes int
|
writes int
|
||||||
maxWrite int
|
maxWrite int
|
||||||
|
deadlines []time.Time
|
||||||
}
|
}
|
||||||
|
|
||||||
func (c *messageWriteTestConn) Read([]byte) (int, error) { return 0, io.EOF }
|
func (c *messageWriteTestConn) Read([]byte) (int, error) { return 0, io.EOF }
|
||||||
|
|
@ -30,7 +31,10 @@ func (*messageWriteTestConn) LocalAddr() net.Addr { return messageW
|
||||||
func (*messageWriteTestConn) RemoteAddr() net.Addr { return messageWriteTestAddr("remote") }
|
func (*messageWriteTestConn) RemoteAddr() net.Addr { return messageWriteTestAddr("remote") }
|
||||||
func (*messageWriteTestConn) SetDeadline(time.Time) error { return nil }
|
func (*messageWriteTestConn) SetDeadline(time.Time) error { return nil }
|
||||||
func (*messageWriteTestConn) SetReadDeadline(time.Time) error { return nil }
|
func (*messageWriteTestConn) SetReadDeadline(time.Time) error { return nil }
|
||||||
func (*messageWriteTestConn) SetWriteDeadline(time.Time) error { return nil }
|
func (c *messageWriteTestConn) SetWriteDeadline(d time.Time) error {
|
||||||
|
c.deadlines = append(c.deadlines, d)
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
type messageWriteTestAddr string
|
type messageWriteTestAddr string
|
||||||
|
|
||||||
|
|
@ -42,6 +46,33 @@ type countWriteBuffer struct {
|
||||||
writes int
|
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) {
|
func (w *countWriteBuffer) Write(p []byte) (int, error) {
|
||||||
w.writes++
|
w.writes++
|
||||||
return w.Buffer.Write(p)
|
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) {
|
func TestCompatTransportWebSocketWritesOneMessagePerPacket(t *testing.T) {
|
||||||
var payload bin.Buffer
|
var payload bin.Buffer
|
||||||
payload.PutInt32(0x01020304)
|
payload.PutInt32(0x01020304)
|
||||||
|
|
|
||||||
|
|
@ -39,15 +39,34 @@ func Register(ctx context.Context, cb func()) bool {
|
||||||
}
|
}
|
||||||
|
|
||||||
func Run(ctx context.Context) {
|
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)
|
cbs, ok := ctx.Value(callbacksKey{}).(*callbacks)
|
||||||
if !ok || cbs == nil {
|
if !ok || cbs == nil {
|
||||||
return
|
return nil
|
||||||
}
|
}
|
||||||
cbs.mu.Lock()
|
cbs.mu.Lock()
|
||||||
list := append([]callback(nil), cbs.list...)
|
list := append([]callback(nil), cbs.list...)
|
||||||
cbs.list = nil
|
cbs.list = nil
|
||||||
cbs.mu.Unlock()
|
cbs.mu.Unlock()
|
||||||
|
if len(list) == 0 {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
var once sync.Once
|
||||||
|
return func() {
|
||||||
|
once.Do(func() {
|
||||||
for _, cb := range list {
|
for _, cb := range list {
|
||||||
cb()
|
cb()
|
||||||
}
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
|
||||||
28
internal/postresponse/callbacks_test.go
Normal file
28
internal/postresponse/callbacks_test.go
Normal 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)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
@ -218,23 +218,31 @@ func New(cfg Config, deps Deps, log *zap.Logger, clk clock.Clock) *Router {
|
||||||
return r
|
return r
|
||||||
}
|
}
|
||||||
|
|
||||||
// Dispatch 路由一条 RPC 请求:先剥离 invokeWithLayer / initConnection /
|
// Dispatch routes one RPC and preserves the historical two-value API used by
|
||||||
// invokeWithoutUpdates / invokeAfter* 等 wrapper(注入 layer / 客户端信息到 ctx),
|
// domain/RPC tests. The MTProto edge uses DispatchWithMethod so outbound
|
||||||
// 再按 TypeID 路由到 typed handler。满足 mtprotoedge.RPCHandler。
|
// 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) {
|
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()
|
preStart := r.clock.Now()
|
||||||
ctx = withInboundRPCBytes(ctx, b.Len())
|
ctx = withInboundRPCBytes(ctx, b.Len())
|
||||||
ctx = WithRawAuthKeyID(ctx, authKeyID)
|
ctx = WithRawAuthKeyID(ctx, authKeyID)
|
||||||
effectiveAuthKeyID, err := r.effectiveAuthKeyID(ctx, authKeyID, sessionID)
|
effectiveAuthKeyID, err := r.effectiveAuthKeyID(ctx, authKeyID, sessionID)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return nil, internalErr()
|
return nil, "", internalErr()
|
||||||
}
|
}
|
||||||
tAuth := r.clock.Now()
|
tAuth := r.clock.Now()
|
||||||
ctx = WithAuthKeyID(ctx, effectiveAuthKeyID)
|
ctx = WithAuthKeyID(ctx, effectiveAuthKeyID)
|
||||||
ctx = WithSessionID(ctx, sessionID)
|
ctx = WithSessionID(ctx, sessionID)
|
||||||
userID, hasUserID, err := r.effectiveUserID(ctx, authKeyID, effectiveAuthKeyID, sessionID)
|
userID, hasUserID, err := r.effectiveUserID(ctx, authKeyID, effectiveAuthKeyID, sessionID)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return nil, internalErr()
|
return nil, "", internalErr()
|
||||||
}
|
}
|
||||||
if hasUserID {
|
if hasUserID {
|
||||||
ctx = WithUserID(ctx, userID)
|
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)
|
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) {
|
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)
|
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 {
|
if depth > maxWrapperDepth {
|
||||||
return nil, wrapperTooDeepErr()
|
return nil, wrapperTooDeepErr()
|
||||||
}
|
}
|
||||||
|
|
@ -516,13 +530,13 @@ func (r *Router) dispatch(ctx context.Context, b *bin.Buffer, depth int) (bin.En
|
||||||
// query 紧跟 layer,buffer 剩余即内层请求。
|
// query 紧跟 layer,buffer 剩余即内层请求。
|
||||||
ctx = WithLayer(ctx, layer)
|
ctx = WithLayer(ctx, layer)
|
||||||
r.rememberClientLayer(ctx, layer)
|
r.rememberClientLayer(ctx, layer)
|
||||||
return r.dispatch(ctx, b, depth+1)
|
return r.dispatch(ctx, b, depth+1, meta)
|
||||||
|
|
||||||
case tg.InvokeWithoutUpdatesRequestTypeID:
|
case tg.InvokeWithoutUpdatesRequestTypeID:
|
||||||
if err := b.ConsumeID(id); err != nil {
|
if err := b.ConsumeID(id); err != nil {
|
||||||
return nil, err
|
return nil, err
|
||||||
}
|
}
|
||||||
return r.dispatch(withInvokeWithoutUpdates(ctx), b, depth+1)
|
return r.dispatch(withInvokeWithoutUpdates(ctx), b, depth+1, meta)
|
||||||
|
|
||||||
case tg.InvokeAfterMsgRequestTypeID:
|
case tg.InvokeAfterMsgRequestTypeID:
|
||||||
if err := b.ConsumeID(id); err != nil {
|
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 {
|
if _, err := b.Long(); err != nil {
|
||||||
return nil, fmt.Errorf("decode invokeAfterMsg msg_id: %w", err)
|
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:
|
case tg.InvokeAfterMsgsRequestTypeID:
|
||||||
if err := b.ConsumeID(id); err != nil {
|
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 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:
|
case tg.InitConnectionRequestTypeID:
|
||||||
req := &tg.InitConnectionRequest{Query: &rawObject{}}
|
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 {
|
if !ok {
|
||||||
return nil, fmt.Errorf("initConnection query: unexpected type %T", req.Query)
|
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:
|
default:
|
||||||
// 入站兼容统一入口(先于鉴权门/dispatcher):layerwire 把老客户端请求升级为
|
// 入站兼容统一入口(先于鉴权门/dispatcher):layerwire 把老客户端请求升级为
|
||||||
|
|
@ -606,6 +620,9 @@ func (r *Router) dispatch(ctx context.Context, b *bin.Buffer, depth int) (bin.En
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
if meta != nil {
|
||||||
|
meta.method = tlTypeName(id)
|
||||||
|
}
|
||||||
knownRequest, structuralErr := layerwire.ValidateRoutableRequest(b.Buf)
|
knownRequest, structuralErr := layerwire.ValidateRoutableRequest(b.Buf)
|
||||||
if !knownRequest {
|
if !knownRequest {
|
||||||
if structuralErr != nil {
|
if structuralErr != nil {
|
||||||
|
|
@ -834,6 +851,37 @@ func (r *Router) NegotiatedLayer(authKeyID [8]byte, sessionID int64) (int, bool)
|
||||||
return currentClientLayer, false
|
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 在单个临界区内完成「读旧值-修改-写回」,避免
|
// mutateClientSessionInfo 在单个临界区内完成「读旧值-修改-写回」,避免
|
||||||
// RLock 读出与 Lock 写回之间被并发写覆盖的窗口。
|
// RLock 读出与 Lock 写回之间被并发写覆盖的窗口。
|
||||||
func (r *Router) mutateClientSessionInfo(ctx context.Context, mutate func(*clientSessionInfo)) {
|
func (r *Router) mutateClientSessionInfo(ctx context.Context, mutate func(*clientSessionInfo)) {
|
||||||
|
|
|
||||||
|
|
@ -49,10 +49,13 @@ func TestDispatchUnwrapsWrappers(t *testing.T) {
|
||||||
t.Fatalf("encode wrapped request: %v", err)
|
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 {
|
if err != nil {
|
||||||
t.Fatalf("dispatch: %v", err)
|
t.Fatalf("dispatch: %v", err)
|
||||||
}
|
}
|
||||||
|
if method != "help.getConfig" {
|
||||||
|
t.Fatalf("effective method = %q, want help.getConfig", method)
|
||||||
|
}
|
||||||
cfg, ok := enc.(*tg.Config)
|
cfg, ok := enc.(*tg.Config)
|
||||||
if !ok {
|
if !ok {
|
||||||
t.Fatalf("result type = %T, want *tg.Config", enc)
|
t.Fatalf("result type = %T, want *tg.Config", enc)
|
||||||
|
|
|
||||||
Loading…
Add table
Add a link
Reference in a new issue