package mtprotoedge import ( "container/list" "context" "errors" "fmt" "sync" "sync/atomic" "time" "telesrv/internal/transport" ) // ErrInboundRPCQueueFull 表示 inbound RPC 已触达单连接或进程级预算。 var ErrInboundRPCQueueFull = errors.New("inbound rpc queue full") // maxInflightRPCBytes 是单连接所有已预留、排队和执行中 RPC 的内存 charge 上限。 // legacy 路径的 charge 等于 copied body;exact 路径在 typed decode 前按 wire 大小、 // 生成对象/向量/interface/string/bytes 放大保守计费。它不是可接收 wire bytes 上限。 // 进程级预算先兜底;这里再隔离单个连接,避免一个客户端独占全局内存。 const maxInflightRPCBytes = 32 << 20 // 32 MiB // rpcCloseWaitTimeout 是连接/Server 关闭时等待在途 RPC 或共享 worker 退出的上限。 const rpcCloseWaitTimeout = 5 * time.Second type inboundRPC struct { ctx context.Context cancel context.CancelFunc stopRoot func() bool stopTimeout func() bool method string enqueuedAt time.Time deadline time.Time size int run func(context.Context) error onTimeout func() // release drops request-scoped ownership which is independent from the body // budget (for example a cross-connection in-flight RPC claim). It runs once // after the request has either published a terminal result or become // impossible to run; callers should still make the callback idempotent. release func() budget *inboundRPCGlobalReservation ticket *inboundRPCTicket gate *inboundRPCGate } // inboundRPCGate keeps dependency waits out of the shared worker pool. A // gated task remains within the ordinary queue/task/byte budgets, but workers // skip it until every prerequisite publishes its terminal execution outcome. type inboundRPCGate struct { remaining atomic.Int32 failed atomic.Bool ready atomic.Bool wake func() } func newInboundRPCGate(prerequisites int, wake func()) *inboundRPCGate { if prerequisites < 0 { prerequisites = 0 } g := &inboundRPCGate{wake: wake} // One sentinel keeps the gate closed while subscribers are installed; the // caller resolves it after registration is complete. g.remaining.Store(int32(prerequisites + 1)) return g } func (g *inboundRPCGate) resolve(success bool) { if g == nil { return } if !success { g.failed.Store(true) } remaining := g.remaining.Add(-1) if remaining < 0 { panic("mtproto inbound RPC gate counter underflow") } if remaining == 0 && g.ready.CompareAndSwap(false, true) && g.wake != nil { g.wake() } } func (g *inboundRPCGate) runnable() bool { return g == nil || g.ready.Load() } func (g *inboundRPCGate) success() bool { return g == nil || (g.ready.Load() && !g.failed.Load()) } type inboundRPCTicket struct { onTimeout func() } // inboundRPCScheduler 是 Server 级共享调度器。ready 中每个 Conn 最多只有一个有效令牌; // worker 每次只从该连接取一条,再把仍可运行的连接放回队尾,因此单个热点连接不能长期 // 占住共享池。worker 在首条任务到达后才创建,空闲 Server 不预起 256 个 goroutine。 type inboundRPCScheduler struct { workers int maxTasks int maxBytes int64 // ready is an intrusive scheduler-owned queue rather than a bounded channel. A connection // has at most one element, and close removes that element in O(1). This prevents closed-Conn // stale tokens from filling a channel and making every worker block while trying to reschedule. readyMu sync.Mutex ready *list.List readyIndex map[*Conn]*list.Element readyWake chan struct{} stopCh chan struct{} lifecycleMu sync.Mutex started bool stopped bool workersStarted bool workerWG sync.WaitGroup budgetMu sync.Mutex tasks int bytes int64 } type inboundRPCGlobalReservation struct { scheduler *inboundRPCScheduler size int64 released atomic.Bool } // inboundRPCSpec 是 container preflight 与 RPC scheduler 之间的有界 admission 描述。 // method 仅用于 metrics;size 是 materialization charge。legacy 路径等于 Copy 前 // request body 字节数,exact 路径是 typed decode 前计算的保守内存上界。 type inboundRPCSpec struct { method string size int } type inboundRPCBatchEntry struct { global *inboundRPCGlobalReservation method string size int } // inboundRPCBatchReservation 把一个 container 内的 RPC 视为一个 admission 单元。 // reserve 一次预留整批的全局/单连接条数和字节预算;commit 一次 append // 全部任务;abort 一次归还全部预算。这防止 container 只执行前半批。 type inboundRPCBatchReservation struct { conn *Conn ctx context.Context entries []inboundRPCBatchEntry totalSize int64 once sync.Once } var errInboundRPCBatchTaskCount = errors.New("inbound rpc batch task count mismatch") var errInboundRPCBatchSelection = errors.New("inbound rpc batch selection is invalid") func newInboundRPCScheduler(workers, maxTasks int, maxBytes int64) *inboundRPCScheduler { if workers <= 0 { workers = 1 } if maxTasks <= 0 { maxTasks = 1 } if maxBytes <= 0 { maxBytes = 1 } return &inboundRPCScheduler{ workers: workers, maxTasks: maxTasks, maxBytes: maxBytes, ready: list.New(), readyIndex: make(map[*Conn]*list.Element), readyWake: make(chan struct{}, 1), stopCh: make(chan struct{}), } } // start 允许共享池开始消费。已在 start 前进入 ready 的任务会保留顺序,便于启动突发, // 也使测试能够确定性验证轮转公平性。 func (s *inboundRPCScheduler) start() { s.lifecycleMu.Lock() if s.stopped { s.lifecycleMu.Unlock() return } s.started = true shouldStart := s.readyLen() > 0 s.lifecycleMu.Unlock() if shouldStart { s.ensureWorkers() } } func (s *inboundRPCScheduler) ensureWorkers() { s.lifecycleMu.Lock() defer s.lifecycleMu.Unlock() if !s.started || s.stopped || s.workersStarted { return } s.workersStarted = true s.workerWG.Add(s.workers) for i := 0; i < s.workers; i++ { go s.worker() } } func (s *inboundRPCScheduler) stop(timeout time.Duration) { s.lifecycleMu.Lock() if !s.stopped { s.stopped = true s.budgetMu.Lock() // 与 reserveGlobal 在同一把锁下切断新任务;已持有 reservation 的任务仍由 // 对应 Conn 的 commit/abort/close 路径精确归还。 close(s.stopCh) s.budgetMu.Unlock() } s.lifecycleMu.Unlock() done := make(chan struct{}) go func() { s.workerWG.Wait() close(done) }() if timeout <= 0 { <-done return } timer := time.NewTimer(timeout) defer timer.Stop() select { case <-done: case <-timer.C: } } // reserveGlobalBatch 在一次 budgetMu 临界区内检查并预留整批条数/字节。 // 返回的每个 reservation 仍由对应 task 单独归还,避免一个慢 RPC 持有 // 整个 container 已完成任务的预算。 func (s *inboundRPCScheduler) reserveGlobalBatch(sizes []int) ([]*inboundRPCGlobalReservation, string, error) { reservations := make([]*inboundRPCGlobalReservation, len(sizes)) s.budgetMu.Lock() defer s.budgetMu.Unlock() select { case <-s.stopCh: return nil, "scheduler_closed", ErrConnClosed default: } if len(sizes) > s.maxTasks-s.tasks { return nil, "global_task_budget", ErrInboundRPCQueueFull } // 逐项从剩余预算中减,避免 total 和 s.bytes+total 溢出。 remaining := s.maxBytes - s.bytes var total int64 for i, size := range sizes { if size < 0 { size = 0 } size64 := int64(size) if size64 > remaining-total { return nil, "global_byte_budget", ErrInboundRPCQueueFull } total += size64 reservations[i] = &inboundRPCGlobalReservation{scheduler: s, size: size64} } s.tasks += len(sizes) s.bytes += total return reservations, "", nil } func (r *inboundRPCGlobalReservation) release() { if r == nil || r.scheduler == nil { return } if !r.released.CompareAndSwap(false, true) { return } s := r.scheduler s.budgetMu.Lock() s.tasks-- s.bytes -= r.size s.budgetMu.Unlock() } // releaseInboundRPCGlobalBatch 使 batch abort/commit-failure 在一次全局锁内 // 归还它仍持有的全部预算。每张 ticket 的 CAS 保证与任何并发释放幂等。 func releaseInboundRPCGlobalBatch(reservations []*inboundRPCGlobalReservation) { var ( scheduler *inboundRPCScheduler tasks int bytes int64 ) for _, reservation := range reservations { if reservation == nil || reservation.scheduler == nil || !reservation.released.CompareAndSwap(false, true) { continue } if scheduler == nil { scheduler = reservation.scheduler } // 一个 batch 只能由同一 scheduler 创建。若未来出现混合调用, // 仍通过单张 release 正确归还,而不会扣错 scheduler。 if reservation.scheduler != scheduler { reservation.scheduler.budgetMu.Lock() reservation.scheduler.tasks-- reservation.scheduler.bytes -= reservation.size reservation.scheduler.budgetMu.Unlock() continue } tasks++ bytes += reservation.size } if scheduler == nil || tasks == 0 { return } scheduler.budgetMu.Lock() scheduler.tasks -= tasks scheduler.bytes -= bytes scheduler.budgetMu.Unlock() } func (s *inboundRPCScheduler) schedule(c *Conn) { if s == nil || c == nil { return } // rpcReady/rpcClosed and queue membership must be tested/installed while holding rpcMu. // Otherwise close can remove the old token between the test and enqueue, leaving a new stale // token behind after the connection is already terminal. c.rpcMu.Lock() eligible := c.rpcReady && !c.rpcClosed added := false if eligible { added = s.enqueueReady(c) } c.rpcMu.Unlock() if !added { return } s.signalReady() s.ensureWorkers() } func (s *inboundRPCScheduler) worker() { defer s.workerWG.Done() for { select { case <-s.stopCh: return default: } if c := s.popReady(); c != nil { task, ok, reschedule := c.takeInboundRPC() if reschedule { s.schedule(c) } if ok { c.runInboundRPC(task) } continue } select { case <-s.readyWake: case <-s.stopCh: return } } } func (s *inboundRPCScheduler) enqueueReady(c *Conn) bool { select { case <-s.stopCh: return false default: } s.readyMu.Lock() defer s.readyMu.Unlock() select { case <-s.stopCh: return false default: } if _, exists := s.readyIndex[c]; exists { return false } s.readyIndex[c] = s.ready.PushBack(c) return true } func (s *inboundRPCScheduler) popReady() *Conn { s.readyMu.Lock() front := s.ready.Front() if front == nil { s.readyMu.Unlock() return nil } c, _ := front.Value.(*Conn) s.ready.Remove(front) delete(s.readyIndex, c) hasMore := s.ready.Len() > 0 s.readyMu.Unlock() if hasMore { // Wake another worker while this worker begins the task. A capacity-one wake channel is // sufficient: every pop cascades another wake until the queue is drained. s.signalReady() } return c } func (s *inboundRPCScheduler) unschedule(c *Conn) { if s == nil || c == nil { return } s.readyMu.Lock() if el := s.readyIndex[c]; el != nil { s.ready.Remove(el) delete(s.readyIndex, c) } hasMore := s.ready.Len() > 0 s.readyMu.Unlock() if hasMore { s.signalReady() } } func (s *inboundRPCScheduler) readyLen() int { s.readyMu.Lock() defer s.readyMu.Unlock() return s.ready.Len() } func (s *inboundRPCScheduler) signalReady() { select { case s.readyWake <- struct{}{}: default: } } func (c *Conn) startInboundRPCScheduler(scheduler *inboundRPCScheduler, maxInflight, queueSize int, timeout time.Duration) { if c.metrics == nil { c.metrics = NopMetrics{} } if maxInflight <= 0 { maxInflight = 1 } if queueSize <= 0 { queueSize = 1 } rootCtx, cancel := context.WithCancel(context.Background()) c.rpcScheduler = scheduler c.rpcCancel = cancel c.rpcTimeout = timeout c.rpcRootCtx = rootCtx c.rpcMaxInflight = maxInflight c.rpcQueueSize = queueSize // rpcQueue 保持 nil;首个成功 commit 才由 append 分配,静默连接零队列内存。 } // reserveInboundRPCBatch 必须在 container 内任何 request body Copy 前调用。 // 全局预算只锁一次,单连接预算也只锁一次;任一限制不满足时 // 整批失败,不会留下部分 task/字节 reservation。 func (c *Conn) reserveInboundRPCBatch(ctx context.Context, specs []inboundRPCSpec) (*inboundRPCBatchReservation, error) { if ctx == nil { ctx = context.Background() } select { case <-ctx.Done(): c.dropInboundRPCSpecs(specs, "context_done") return nil, ctx.Err() default: } if c.isRetired() { c.dropInboundRPCSpecs(specs, "scheduler_closed") return nil, ErrConnClosed } if c.rpcScheduler == nil { c.dropInboundRPCSpecs(specs, "scheduler_closed") return nil, ErrConnClosed } normalized := make([]inboundRPCSpec, len(specs)) sizes := make([]int, len(specs)) for i, spec := range specs { if spec.size < 0 { spec.size = 0 } normalized[i] = spec sizes[i] = spec.size } globals, reason, err := c.rpcScheduler.reserveGlobalBatch(sizes) if err != nil { c.dropInboundRPCSpecs(normalized, reason) return nil, err } entries := make([]inboundRPCBatchEntry, len(normalized)) var totalSize int64 for i, spec := range normalized { entries[i] = inboundRPCBatchEntry{ global: globals[i], method: spec.method, size: spec.size, } totalSize += int64(spec.size) } c.rpcMu.Lock() if err := ctx.Err(); err != nil { c.rpcMu.Unlock() releaseInboundRPCGlobalBatch(globals) c.dropInboundRPCSpecs(normalized, "context_done") return nil, err } if c.rpcClosed || c.isRetired() { c.rpcMu.Unlock() releaseInboundRPCGlobalBatch(globals) c.dropInboundRPCSpecs(normalized, "scheduler_closed") return nil, ErrConnClosed } // 用减法比较,避免对抗性 batch 的 len 加法溢出。 availableSlots := c.rpcQueueSize - c.rpcReserved - len(c.rpcQueue) if len(entries) > availableSlots { c.rpcMu.Unlock() releaseInboundRPCGlobalBatch(globals) c.dropInboundRPCSpecs(normalized, "queue_full") return nil, ErrInboundRPCQueueFull } if totalSize > maxInflightRPCBytes-c.inflightRPCBytes.Load() { c.rpcMu.Unlock() releaseInboundRPCGlobalBatch(globals) c.dropInboundRPCSpecs(normalized, "byte_budget") return nil, ErrInboundRPCQueueFull } c.rpcReserved += len(entries) c.inflightRPCBytes.Add(totalSize) // Add 与 close 的 Wait 由 rpcMu 排序:close 置 rpcClosed 后不会再发生 Add。 // 整批 reservation 只需一个 waiter;commit/abort 也只会完成一次。 c.rpcReservationWG.Add(1) c.rpcMu.Unlock() return &inboundRPCBatchReservation{ conn: c, ctx: ctx, entries: entries, totalSize: totalSize, }, nil } func (c *Conn) dropInboundRPCSpecs(specs []inboundRPCSpec, reason string) { for _, spec := range specs { c.metrics.InboundRPCDropped(spec.method, reason) } } // growEntry raises one provisional task's materialization charge without a // release/reacquire window. Exact gzip admission calls this after the expanded // size is known but before the generated typed decoder can allocate the request // graph. A failure leaves every existing reservation unchanged so the caller // can reject and abort the whole container atomically. func (r *inboundRPCBatchReservation) growEntry(index, targetSize int) error { if r == nil || index < 0 || index >= len(r.entries) || targetSize < 0 { return errInboundRPCBatchSelection } entry := &r.entries[index] if targetSize <= entry.size { return nil } if entry.global == nil || entry.global.scheduler == nil || entry.global.released.Load() { return ErrConnClosed } delta := int64(targetSize) - int64(entry.size) scheduler := entry.global.scheduler conn := r.conn // Match initial admission's lock order: global scheduler budget, then the // connection queue budget. Release paths never hold rpcMu while acquiring // budgetMu, so this cannot invert task completion or close. scheduler.budgetMu.Lock() defer scheduler.budgetMu.Unlock() select { case <-scheduler.stopCh: return ErrConnClosed default: } if delta > scheduler.maxBytes-scheduler.bytes { return fmt.Errorf("%w: grow exact admission global byte budget by %d", ErrInboundRPCQueueFull, delta) } conn.rpcMu.Lock() defer conn.rpcMu.Unlock() if err := r.ctx.Err(); err != nil { return err } if conn.rpcClosed || conn.isRetired() { return ErrConnClosed } if delta > int64(maxInflightRPCBytes)-conn.inflightRPCBytes.Load() { return fmt.Errorf("%w: grow exact admission connection byte budget by %d", ErrInboundRPCQueueFull, delta) } scheduler.bytes += delta entry.global.size += delta entry.size = targetSize r.totalSize += delta conn.inflightRPCBytes.Add(delta) return nil } // retain keeps a subset of a provisional batch on the original connection and // global reservations. Exact-layer admission uses this after typed decode has // classified completed replays, pending joins, admission errors, and fresh // owners. A retained fresh owner therefore never passes through a release then // reacquire window where another connection could consume its memory/task // budget. Entries not retained are returned immediately as one batch. // // indices are reservation-entry indices, must be strictly increasing, and have // a one-to-one correspondence with specs. The original conservative byte // charge is intentionally retained; a typed request remains reachable from its // scheduler task until completion, so shrinking it to the wire size would make // the materialized graph unaccounted. func (r *inboundRPCBatchReservation) retain(indices []int, specs []inboundRPCSpec) error { if r == nil || len(indices) != len(specs) || len(indices) > len(r.entries) { return errInboundRPCBatchSelection } retained := make([]inboundRPCBatchEntry, len(indices)) removed := make([]*inboundRPCGlobalReservation, 0, len(r.entries)-len(indices)) var ( previous = -1 retainedSize int64 removedSize int64 ) selected := 0 for output, index := range indices { if index <= previous || index < 0 || index >= len(r.entries) { return errInboundRPCBatchSelection } previous = index for selected < index { entry := r.entries[selected] removed = append(removed, entry.global) removedSize += int64(entry.size) selected++ } entry := r.entries[index] // The caller may improve the provisional metric label but may not increase // or replace its byte charge after materialization has started. if specs[output].size > entry.size { return errInboundRPCBatchSelection } entry.method = specs[output].method retained[output] = entry retainedSize += int64(entry.size) selected = index + 1 } for selected < len(r.entries) { entry := r.entries[selected] removed = append(removed, entry.global) removedSize += int64(entry.size) selected++ } c := r.conn c.rpcMu.Lock() c.rpcReserved -= len(removed) c.inflightRPCBytes.Add(-removedSize) r.entries = retained r.totalSize = retainedSize c.rpcMu.Unlock() releaseInboundRPCGlobalBatch(removed) if len(retained) == 0 { // Finish the one reservation waiter as part of the same ownership // transition. abort sees an empty batch and is therefore accounting-only. r.abort() } return nil } // commit 在一次 rpcMu 临界区内把整批 task append 到队列并立即发布 ready token。 // 协议 barrier 必须在调用 commit 前完成;延迟发布 token 无法阻止已有 worker // 从同一连接队列取走新任务,因此不提供虚假的 deferred-schedule 模式。 func (r *inboundRPCBatchReservation) commit(tasks []inboundRPC) (result error) { if r == nil { return ErrConnClosed } result = ErrConnClosed var ( committed bool reschedule bool firstQueueLen int queueCap int globals []*inboundRPCGlobalReservation ) r.once.Do(func() { c := r.conn // A container reservation may deliberately span protocol-critical barriers // (session ownership claim and new_session_created's physical write). Queue // and execution latency starts only when the fully admitted batch becomes // runnable, not while it is waiting behind those independent barriers. enqueuedAt := time.Now() deadline := time.Time{} if c.rpcTimeout > 0 { deadline = enqueuedAt.Add(c.rpcTimeout) } if ctxDeadline, ok := r.ctx.Deadline(); ok && (deadline.IsZero() || ctxDeadline.Before(deadline)) { deadline = ctxDeadline } globals = make([]*inboundRPCGlobalReservation, len(r.entries)) for i := range r.entries { globals[i] = r.entries[i].global } c.rpcMu.Lock() c.rpcReserved -= len(r.entries) if len(tasks) != len(r.entries) { c.inflightRPCBytes.Add(-r.totalSize) result = errInboundRPCBatchTaskCount } else if c.rpcClosed || c.isRetired() { c.inflightRPCBytes.Add(-r.totalSize) } else { prepared := make([]inboundRPC, len(tasks)) // The request deadline starts when admission succeeds, not when a worker // eventually dequeues the request. This bounds total queue + execution // latency and lets a queued request emit its explicit timeout on time. for i := range tasks { task := tasks[i] entry := r.entries[i] if deadline.IsZero() { task.ctx, task.cancel = context.WithCancel(r.ctx) } else { task.ctx, task.cancel = context.WithDeadline(r.ctx, deadline) } task.stopRoot = context.AfterFunc(c.rpcRootCtx, task.cancel) task.method = entry.method task.enqueuedAt = enqueuedAt task.deadline = deadline task.size = entry.size task.budget = entry.global ticket := &inboundRPCTicket{} if task.onTimeout != nil { onTimeout := task.onTimeout var timeoutOnce sync.Once ticket.onTimeout = func() { timeoutOnce.Do(onTimeout) } task.onTimeout = ticket.onTimeout } task.ticket = ticket if task.onTimeout != nil && !task.deadline.IsZero() { taskCtx := task.ctx task.stopTimeout = context.AfterFunc(taskCtx, func() { if errors.Is(taskCtx.Err(), context.DeadlineExceeded) { c.expireInboundRPCTicket(ticket) } }) } prepared[i] = task } firstQueueLen = len(c.rpcQueue) + 1 c.rpcQueue = append(c.rpcQueue, prepared...) queueCap = c.rpcQueueSize if c.rpcRunning < c.rpcMaxInflight && !c.rpcReady && c.hasRunnableInboundRPCLocked() { c.rpcReady = true reschedule = true } committed = true result = nil } c.rpcMu.Unlock() c.rpcReservationWG.Done() if !committed { releaseInboundRPCGlobalBatch(globals) } }) if committed { for i, entry := range r.entries { r.conn.metrics.InboundRPCQueued(entry.method, firstQueueLen+i, queueCap) } if reschedule { r.conn.rpcScheduler.schedule(r.conn) } } return result } func (r *inboundRPCBatchReservation) abort() { if r == nil { return } r.once.Do(func() { c := r.conn c.rpcMu.Lock() c.rpcReserved -= len(r.entries) c.inflightRPCBytes.Add(-r.totalSize) c.rpcMu.Unlock() c.rpcReservationWG.Done() globals := make([]*inboundRPCGlobalReservation, len(r.entries)) for i := range r.entries { globals[i] = r.entries[i].global } releaseInboundRPCGlobalBatch(globals) }) } func (c *Conn) takeInboundRPC() (task inboundRPC, ok, reschedule bool) { c.rpcMu.Lock() defer c.rpcMu.Unlock() // ready token 是可替代的:收到一个 token 就消费当前“已调度”状态。关闭后或 // 已被另一 token 抢先处理时,这只是一个无害 stale token。 if !c.rpcReady { return inboundRPC{}, false, false } c.rpcReady = false if c.rpcClosed || len(c.rpcQueue) == 0 || c.rpcRunning >= c.rpcMaxInflight { return inboundRPC{}, false, false } index := -1 for i := range c.rpcQueue { if c.rpcQueue[i].gate.runnable() { index = i break } } if index < 0 { return inboundRPC{}, false, false } task = c.rpcQueue[index] if index == 0 { // The overwhelmingly common FIFO path advances the slice head in O(1). // Clear the departed element first so its request graph and closures are // not retained by the shared backing array. c.rpcQueue[0] = inboundRPC{} c.rpcQueue = c.rpcQueue[1:] } else { copy(c.rpcQueue[index:], c.rpcQueue[index+1:]) last := len(c.rpcQueue) - 1 c.rpcQueue[last] = inboundRPC{} c.rpcQueue = c.rpcQueue[:last] } if len(c.rpcQueue) == 0 { c.rpcQueue = nil } c.rpcRunning++ c.rpcWG.Add(1) if c.rpcRunning < c.rpcMaxInflight && c.hasRunnableInboundRPCLocked() { c.rpcReady = true reschedule = true } return task, true, reschedule } func (c *Conn) hasRunnableInboundRPCLocked() bool { if c.rpcReplayRestores > 0 { return false } for i := range c.rpcQueue { if c.rpcQueue[i].gate.runnable() { return true } } return false } // beginRPCReplayRestore installs a scheduler-level barrier without occupying a // global RPC worker. The returned idempotent completion function wakes queued // work only after the last overlapping restore has finished. func (c *Conn) beginRPCReplayRestore() func() { if c == nil { return func() {} } active := false unschedule := false c.rpcMu.Lock() if !c.rpcClosed && !c.isRetired() { c.rpcReplayRestores++ active = true if c.rpcReady { c.rpcReady = false unschedule = true } } c.rpcMu.Unlock() if unschedule && c.rpcScheduler != nil { c.rpcScheduler.unschedule(c) } var once sync.Once return func() { once.Do(func() { if !active { return } reschedule := false c.rpcMu.Lock() c.rpcReplayRestores-- if c.rpcReplayRestores < 0 { c.rpcMu.Unlock() panic("mtproto inbound RPC replay restore barrier underflow") } if c.rpcReplayRestores == 0 && !c.rpcClosed && c.rpcRunning < c.rpcMaxInflight && !c.rpcReady && c.hasRunnableInboundRPCLocked() { c.rpcReady = true reschedule = true } c.rpcMu.Unlock() if reschedule && c.rpcScheduler != nil { c.rpcScheduler.schedule(c) } }) } } // wakeInboundRPC is safe from a flight completion callback. It does no work // while the gate's task has not been committed yet; commit performs the same // runnable scan and publishes the initial scheduler token. func (c *Conn) wakeInboundRPC() { if c == nil || c.rpcScheduler == nil { return } reschedule := false c.rpcMu.Lock() if !c.rpcClosed && c.rpcRunning < c.rpcMaxInflight && !c.rpcReady && c.hasRunnableInboundRPCLocked() { c.rpcReady = true reschedule = true } c.rpcMu.Unlock() if reschedule { c.rpcScheduler.schedule(c) } } func (c *Conn) runInboundRPC(task inboundRPC) { defer c.finishInboundRPC(task) now := time.Now() ctxErr := task.ctx.Err() if (!task.deadline.IsZero() && !now.Before(task.deadline)) || errors.Is(ctxErr, context.DeadlineExceeded) { c.metrics.InboundRPCDropped(task.method, "queue_timeout") if task.onTimeout != nil { task.onTimeout() } return } if ctxErr != nil { c.metrics.InboundRPCDropped(task.method, "context_done") return } c.metrics.InboundRPCStarted(task.method, now.Sub(task.enqueuedAt)) ctx := task.ctx // 把连接对端 IP 作为中立传输元数据注入,绑定设备授权时写 authorizations.ip。 // Edge 只产生传输事实;由 RPC 层的 transport.ClientIPFrom 消费,避免反向依赖。 if ip := c.clientIP(); ip != "" { ctx = transport.WithClientIP(ctx, ip) } if task.run != nil { _ = task.run(ctx) } } func (c *Conn) finishInboundRPC(task inboundRPC) { stopInboundRPCTask(task) var reschedule bool c.rpcMu.Lock() c.rpcRunning-- c.inflightRPCBytes.Add(-int64(task.size)) if !c.rpcClosed && c.rpcRunning < c.rpcMaxInflight && !c.rpcReady && c.hasRunnableInboundRPCLocked() { c.rpcReady = true reschedule = true } c.rpcMu.Unlock() reservation := task.budget release := task.release // The scheduler budget may be reused immediately after release. Clear request-owned // closures/context references first so slow metrics/rescheduling cannot overlap the old body // with a newly admitted body under the same byte accounting. task = inboundRPC{} reservation.release() if release != nil { release() } c.rpcWG.Done() if reschedule { c.rpcScheduler.schedule(c) } } // expireInboundRPCTicket removes a request that is still queued and returns its // memory/task reservations immediately. If the worker won the dequeue race, the // callback does nothing: the running handler owns the only terminal response and // its deadline is represented solely by context cancellation. func (c *Conn) expireInboundRPCTicket(ticket *inboundRPCTicket) { if ticket == nil { return } var ( task inboundRPC found bool unschedule bool ) c.rpcMu.Lock() for i := range c.rpcQueue { if c.rpcQueue[i].ticket != ticket { continue } task = c.rpcQueue[i] copy(c.rpcQueue[i:], c.rpcQueue[i+1:]) last := len(c.rpcQueue) - 1 c.rpcQueue[last] = inboundRPC{} c.rpcQueue = c.rpcQueue[:last] if len(c.rpcQueue) == 0 { c.rpcQueue = nil if c.rpcReady { c.rpcReady = false unschedule = true } } c.inflightRPCBytes.Add(-int64(task.size)) found = true break } c.rpcMu.Unlock() if unschedule { c.rpcScheduler.unschedule(c) } if found { method := task.method reservation := task.budget release := task.release stopInboundRPCTask(task) // Drop the run/context closures before returning the byte reservation. Otherwise an // onTimeout callback that blocks or performs a slow write can keep the copied request body // reachable after the global scheduler has advertised those bytes as available again. task = inboundRPC{} reservation.release() c.metrics.InboundRPCDropped(method, "queue_timeout") if ticket.onTimeout != nil { ticket.onTimeout() } if release != nil { release() } return } } // stopInboundRPCTask disarms queue-expiration cleanup before canceling the // context. Once a worker dequeues the task, the deadline only cancels the // handler; it never races the handler with an early RPC_TIMEOUT response. func stopInboundRPCTask(task inboundRPC) { if task.stopTimeout != nil { task.stopTimeout() } if task.stopRoot != nil { task.stopRoot() } if task.cancel != nil { task.cancel() } } func (c *Conn) closeInboundRPCScheduler() { c.beginCloseInboundRPCScheduler() if c.rpcScheduler == nil { return } c.waitInboundShutdown(rpcCloseWaitTimeout) } // beginCloseInboundRPCScheduler publishes closure, cancels running work and releases queued // requests without waiting for handlers. Shutdown publishes this phase before transport.Close so // a pathological/blocking transport implementation cannot leave the RPC admission gate open. func (c *Conn) beginCloseInboundRPCScheduler() { if c.rpcScheduler == nil { return } c.rpcClose.Do(func() { c.rpcMu.Lock() c.rpcClosed = true c.rpcReady = false queued := c.rpcQueue c.rpcQueue = nil for i := range queued { c.inflightRPCBytes.Add(-int64(queued[i].size)) } c.rpcMu.Unlock() // Remove the scheduler-owned token after rpcClosed/rpcReady become visible. schedule() // takes rpcMu while installing a token, so either it finishes first and is removed here, // or it observes the closed state and cannot enqueue a new stale token afterward. c.rpcScheduler.unschedule(c) if c.rpcCancel != nil { c.rpcCancel() } for i := range queued { task := queued[i] queued[i] = inboundRPC{} method := task.method reservation := task.budget release := task.release stopInboundRPCTask(task) task = inboundRPC{} reservation.release() if release != nil { release() } c.metrics.InboundRPCDropped(method, "connection_closed") } }) } // waitInboundShutdown 等 Copy 前 reservation 完成 commit/abort,以及本连接已经出队的 RPC // 完成,二者共用一个 timeout。超时后 reservation/共享 worker 会在底层调用最终返回时自行 // 收敛;连接 root context 已取消。 func (c *Conn) waitInboundShutdown(timeout time.Duration) bool { done := make(chan struct{}) go func() { c.rpcReservationWG.Wait() c.rpcWG.Wait() close(done) }() if timeout <= 0 { return false } timer := time.NewTimer(timeout) defer timer.Stop() select { case <-done: return true case <-timer.C: return false } }