perf: sync protocol and core hardening updates
This commit is contained in:
A
parent
152fed3b87
commit
4390ebf5a9
283 changed files with 29231 additions and 2295 deletions
|
|
@ -1,30 +1,337 @@
|
|||
package mtprotoedge
|
||||
|
||||
import (
|
||||
"container/list"
|
||||
"context"
|
||||
"errors"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
)
|
||||
|
||||
// ErrInboundRPCQueueFull 表示单连接 RPC 队列已满。
|
||||
// ErrInboundRPCQueueFull 表示 inbound RPC 已触达单连接或进程级预算。
|
||||
var ErrInboundRPCQueueFull = errors.New("inbound rpc queue full")
|
||||
|
||||
// maxInflightRPCBytes 是单连接已入队未完成 inbound RPC body 的总字节上限。
|
||||
// 队列除按条数(queueSize)限制外，再按字节预算兜底：对抗客户端发满大请求时按字节先拒绝。
|
||||
// maxInflightRPCBytes 是单连接所有已预留、排队和执行中 RPC body 的总字节上限。
|
||||
// 进程级预算在 Copy 前先兜底；这里再隔离单个连接，避免一个客户端独占全局内存。
|
||||
const maxInflightRPCBytes = 32 << 20 // 32 MiB
|
||||
|
||||
// rpcCloseWaitTimeout 是连接关闭时等待 inbound RPC worker 退出的上限。
|
||||
// rpcCloseWaitTimeout 是连接/Server 关闭时等待在途 RPC 或共享 worker 退出的上限。
|
||||
const rpcCloseWaitTimeout = 5 * time.Second
|
||||
|
||||
type inboundRPC struct {
|
||||
ctx context.Context
|
||||
method string
|
||||
enqueuedAt time.Time
|
||||
size int
|
||||
run func(context.Context) error
|
||||
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()
|
||||
budget *inboundRPCGlobalReservation
|
||||
ticket *inboundRPCTicket
|
||||
}
|
||||
|
||||
func (c *Conn) startInboundRPCScheduler(maxInflight, queueSize int, timeout time.Duration) {
|
||||
const (
|
||||
inboundRPCTicketQueued int32 = iota
|
||||
inboundRPCTicketRunning
|
||||
inboundRPCTicketDone
|
||||
)
|
||||
|
||||
type inboundRPCTicket struct {
|
||||
state atomic.Int32
|
||||
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
|
||||
once sync.Once
|
||||
}
|
||||
|
||||
// inboundRPCReservation 同时持有全局和单连接的“Copy 前”预算。commit/abort 只能成功一次；
|
||||
// 无论 Copy 后连接关闭、入队成功还是调用方提前返回，预算都有唯一归还路径。
|
||||
type inboundRPCReservation struct {
|
||||
conn *Conn
|
||||
global *inboundRPCGlobalReservation
|
||||
ctx context.Context
|
||||
method string
|
||||
size int
|
||||
enqueuedAt time.Time
|
||||
deadline time.Time
|
||||
once sync.Once
|
||||
}
|
||||
|
||||
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:
|
||||
}
|
||||
}
|
||||
|
||||
func (s *inboundRPCScheduler) reserveGlobal(size int) (*inboundRPCGlobalReservation, string, error) {
|
||||
if size < 0 {
|
||||
size = 0
|
||||
}
|
||||
size64 := int64(size)
|
||||
s.budgetMu.Lock()
|
||||
defer s.budgetMu.Unlock()
|
||||
|
||||
select {
|
||||
case <-s.stopCh:
|
||||
return nil, "scheduler_closed", ErrConnClosed
|
||||
default:
|
||||
}
|
||||
if s.tasks >= s.maxTasks {
|
||||
return nil, "global_task_budget", ErrInboundRPCQueueFull
|
||||
}
|
||||
// 用减法比较避免 s.bytes+size64 溢出。
|
||||
if size64 > s.maxBytes-s.bytes {
|
||||
return nil, "global_byte_budget", ErrInboundRPCQueueFull
|
||||
}
|
||||
s.tasks++
|
||||
s.bytes += size64
|
||||
return &inboundRPCGlobalReservation{scheduler: s, size: size64}, "", nil
|
||||
}
|
||||
|
||||
func (r *inboundRPCGlobalReservation) release() {
|
||||
if r == nil || r.scheduler == nil {
|
||||
return
|
||||
}
|
||||
r.once.Do(func() {
|
||||
s := r.scheduler
|
||||
s.budgetMu.Lock()
|
||||
s.tasks--
|
||||
s.bytes -= r.size
|
||||
s.budgetMu.Unlock()
|
||||
})
|
||||
}
|
||||
|
||||
func (s *inboundRPCScheduler) budgetSnapshot() (tasks int, bytes int64) {
|
||||
s.budgetMu.Lock()
|
||||
defer s.budgetMu.Unlock()
|
||||
return s.tasks, s.bytes
|
||||
}
|
||||
|
||||
func (s *inboundRPCScheduler) schedule(c *Conn) {
|
||||
if s == nil || c == nil {
|
||||
return
|
||||
}
|
||||
// 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{}
|
||||
}
|
||||
|
|
@ -35,163 +342,413 @@ func (c *Conn) startInboundRPCScheduler(maxInflight, queueSize int, timeout time
|
|||
queueSize = 1
|
||||
}
|
||||
rootCtx, cancel := context.WithCancel(context.Background())
|
||||
c.rpcQueue = make(chan inboundRPC, queueSize)
|
||||
c.rpcStop = make(chan struct{})
|
||||
c.rpcScheduler = scheduler
|
||||
c.rpcCancel = cancel
|
||||
c.rpcTimeout = timeout
|
||||
c.rpcRootCtx = rootCtx
|
||||
c.rpcMaxInflight = maxInflight
|
||||
// worker 懒启动：不在此处起 worker；首个 RPC 入队时由 ensureInboundRPCWorkers 起，
|
||||
// 避免握手后静默 / 纯推送目标连接白白钉住 maxInflight 个 goroutine。
|
||||
c.rpcQueueSize = queueSize
|
||||
// rpcQueue 保持 nil；首个成功 commit 才由 append 分配，静默连接零队列内存。
|
||||
}
|
||||
|
||||
// ensureInboundRPCWorkers 懒启动 maxInflight 个 RPC worker（仅一次），在 enqueueInboundRPC
|
||||
// 入队成功后调用。从不发 RPC 的连接（半开 / 纯推送）由此完全不起 worker。
|
||||
func (c *Conn) ensureInboundRPCWorkers() {
|
||||
c.rpcWorkersOnce.Do(func() {
|
||||
c.rpcWG.Add(c.rpcMaxInflight)
|
||||
for i := 0; i < c.rpcMaxInflight; i++ {
|
||||
go c.inboundRPCWorker(c.rpcRootCtx)
|
||||
// reserveInboundRPC 必须在 request body Copy 前调用。它先拿进程级条数/字节预算，
|
||||
// 再预占单连接队列槽和字节预算；commit 或 abort 负责唯一释放。
|
||||
func (c *Conn) reserveInboundRPC(ctx context.Context, method string, size int) (*inboundRPCReservation, error) {
|
||||
if ctx == nil {
|
||||
ctx = context.Background()
|
||||
}
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
c.metrics.InboundRPCDropped(method, "context_done")
|
||||
return nil, ctx.Err()
|
||||
default:
|
||||
}
|
||||
if c.rpcScheduler == nil {
|
||||
c.metrics.InboundRPCDropped(method, "scheduler_closed")
|
||||
return nil, ErrConnClosed
|
||||
}
|
||||
global, reason, err := c.rpcScheduler.reserveGlobal(size)
|
||||
if err != nil {
|
||||
c.metrics.InboundRPCDropped(method, reason)
|
||||
return nil, err
|
||||
}
|
||||
|
||||
now := time.Now()
|
||||
deadline := time.Time{}
|
||||
if c.rpcTimeout > 0 {
|
||||
deadline = now.Add(c.rpcTimeout)
|
||||
}
|
||||
if ctxDeadline, ok := ctx.Deadline(); ok && (deadline.IsZero() || ctxDeadline.Before(deadline)) {
|
||||
deadline = ctxDeadline
|
||||
}
|
||||
if size < 0 {
|
||||
size = 0
|
||||
}
|
||||
|
||||
c.rpcMu.Lock()
|
||||
if err := ctx.Err(); err != nil {
|
||||
c.rpcMu.Unlock()
|
||||
global.release()
|
||||
c.metrics.InboundRPCDropped(method, "context_done")
|
||||
return nil, err
|
||||
}
|
||||
if c.rpcClosed {
|
||||
c.rpcMu.Unlock()
|
||||
global.release()
|
||||
c.metrics.InboundRPCDropped(method, "scheduler_closed")
|
||||
return nil, ErrConnClosed
|
||||
}
|
||||
if c.rpcReserved+len(c.rpcQueue) >= c.rpcQueueSize {
|
||||
c.rpcMu.Unlock()
|
||||
global.release()
|
||||
c.metrics.InboundRPCDropped(method, "queue_full")
|
||||
return nil, ErrInboundRPCQueueFull
|
||||
}
|
||||
if int64(size) > maxInflightRPCBytes-c.inflightRPCBytes.Load() {
|
||||
c.rpcMu.Unlock()
|
||||
global.release()
|
||||
c.metrics.InboundRPCDropped(method, "byte_budget")
|
||||
return nil, ErrInboundRPCQueueFull
|
||||
}
|
||||
c.rpcReserved++
|
||||
c.inflightRPCBytes.Add(int64(size))
|
||||
// Add 与 close 的 Wait 由 rpcMu 排序：close 置 rpcClosed 后不会再发生 Add。
|
||||
c.rpcReservationWG.Add(1)
|
||||
c.rpcMu.Unlock()
|
||||
|
||||
return &inboundRPCReservation{
|
||||
conn: c,
|
||||
global: global,
|
||||
ctx: ctx,
|
||||
method: method,
|
||||
size: size,
|
||||
enqueuedAt: now,
|
||||
deadline: deadline,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// enqueueInboundRPC 是测试和已持有独立 body 的便捷入口。生产收包路径使用
|
||||
// reserveInboundRPC -> Copy -> commit，保证真正的 Copy 前预算。
|
||||
func (c *Conn) enqueueInboundRPC(ctx context.Context, task inboundRPC) error {
|
||||
reservation, err := c.reserveInboundRPC(ctx, task.method, task.size)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer reservation.abort()
|
||||
return reservation.commit(task)
|
||||
}
|
||||
|
||||
func (r *inboundRPCReservation) commit(task inboundRPC) error {
|
||||
result := ErrConnClosed
|
||||
var (
|
||||
committed bool
|
||||
reschedule bool
|
||||
queueLen int
|
||||
queueCap int
|
||||
)
|
||||
r.once.Do(func() {
|
||||
c := r.conn
|
||||
c.rpcMu.Lock()
|
||||
c.rpcReserved--
|
||||
if c.rpcClosed {
|
||||
c.inflightRPCBytes.Add(-int64(r.size))
|
||||
} else {
|
||||
// The request deadline starts when admission succeeds, not when a worker
|
||||
// eventually dequeues the request. This bounds total queue + execution
|
||||
// latency and lets a queued request emit its explicit timeout on time.
|
||||
if r.deadline.IsZero() {
|
||||
task.ctx, task.cancel = context.WithCancel(r.ctx)
|
||||
} else {
|
||||
task.ctx, task.cancel = context.WithDeadline(r.ctx, r.deadline)
|
||||
}
|
||||
task.stopRoot = context.AfterFunc(c.rpcRootCtx, task.cancel)
|
||||
task.method = r.method
|
||||
task.enqueuedAt = r.enqueuedAt
|
||||
task.deadline = r.deadline
|
||||
task.size = r.size
|
||||
task.budget = r.global
|
||||
ticket := &inboundRPCTicket{}
|
||||
if task.onTimeout != nil {
|
||||
onTimeout := task.onTimeout
|
||||
var timeoutOnce sync.Once
|
||||
ticket.onTimeout = func() {
|
||||
timeoutOnce.Do(onTimeout)
|
||||
}
|
||||
task.onTimeout = ticket.onTimeout
|
||||
}
|
||||
task.ticket = ticket
|
||||
if task.onTimeout != nil && !task.deadline.IsZero() {
|
||||
taskCtx := task.ctx
|
||||
task.stopTimeout = context.AfterFunc(taskCtx, func() {
|
||||
if errors.Is(taskCtx.Err(), context.DeadlineExceeded) {
|
||||
c.expireInboundRPCTicket(ticket)
|
||||
}
|
||||
})
|
||||
}
|
||||
c.rpcQueue = append(c.rpcQueue, task)
|
||||
queueLen = len(c.rpcQueue)
|
||||
queueCap = c.rpcQueueSize
|
||||
if c.rpcRunning < c.rpcMaxInflight && !c.rpcReady {
|
||||
c.rpcReady = true
|
||||
reschedule = true
|
||||
}
|
||||
committed = true
|
||||
result = nil
|
||||
}
|
||||
c.rpcMu.Unlock()
|
||||
c.rpcReservationWG.Done()
|
||||
if !committed {
|
||||
r.global.release()
|
||||
}
|
||||
})
|
||||
if committed {
|
||||
r.conn.metrics.InboundRPCQueued(r.method, queueLen, queueCap)
|
||||
if reschedule {
|
||||
r.conn.rpcScheduler.schedule(r.conn)
|
||||
}
|
||||
}
|
||||
return result
|
||||
}
|
||||
|
||||
func (r *inboundRPCReservation) abort() {
|
||||
if r == nil {
|
||||
return
|
||||
}
|
||||
r.once.Do(func() {
|
||||
c := r.conn
|
||||
c.rpcMu.Lock()
|
||||
c.rpcReserved--
|
||||
c.inflightRPCBytes.Add(-int64(r.size))
|
||||
c.rpcMu.Unlock()
|
||||
c.rpcReservationWG.Done()
|
||||
r.global.release()
|
||||
})
|
||||
}
|
||||
|
||||
func (c *Conn) enqueueInboundRPC(ctx context.Context, task inboundRPC) error {
|
||||
if ctx == nil {
|
||||
ctx = context.Background()
|
||||
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
|
||||
}
|
||||
if c.rpcQueue == nil || c.rpcStop == nil {
|
||||
c.metrics.InboundRPCDropped(task.method, "scheduler_closed")
|
||||
return ErrConnClosed
|
||||
c.rpcReady = false
|
||||
if c.rpcClosed || len(c.rpcQueue) == 0 || c.rpcRunning >= c.rpcMaxInflight {
|
||||
return inboundRPC{}, false, false
|
||||
}
|
||||
task.ctx = ctx
|
||||
task.enqueuedAt = time.Now()
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
task = c.rpcQueue[0]
|
||||
c.rpcQueue[0] = inboundRPC{}
|
||||
c.rpcQueue = c.rpcQueue[1:]
|
||||
if len(c.rpcQueue) == 0 {
|
||||
c.rpcQueue = nil
|
||||
}
|
||||
c.rpcRunning++
|
||||
if task.ticket != nil {
|
||||
task.ticket.state.Store(inboundRPCTicketRunning)
|
||||
}
|
||||
c.rpcWG.Add(1)
|
||||
if len(c.rpcQueue) > 0 && c.rpcRunning < c.rpcMaxInflight {
|
||||
c.rpcReady = true
|
||||
reschedule = true
|
||||
}
|
||||
return task, true, reschedule
|
||||
}
|
||||
|
||||
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 ctx.Err()
|
||||
case <-c.rpcStop:
|
||||
c.metrics.InboundRPCDropped(task.method, "scheduler_closed")
|
||||
return ErrConnClosed
|
||||
default:
|
||||
return
|
||||
}
|
||||
// 字节预算：先预扣 size，超 maxInflightRPCBytes 则回滚并拒绝（与条数上限并列的第二道闸）。
|
||||
if task.size > 0 {
|
||||
if c.inflightRPCBytes.Add(int64(task.size)) > maxInflightRPCBytes {
|
||||
c.inflightRPCBytes.Add(-int64(task.size))
|
||||
c.metrics.InboundRPCDropped(task.method, "byte_budget")
|
||||
return ErrInboundRPCQueueFull
|
||||
}
|
||||
}
|
||||
select {
|
||||
case c.rpcQueue <- task:
|
||||
c.ensureInboundRPCWorkers()
|
||||
c.metrics.InboundRPCQueued(task.method, len(c.rpcQueue), cap(c.rpcQueue))
|
||||
return nil
|
||||
case <-ctx.Done():
|
||||
c.releaseInflightRPCBytes(task.size)
|
||||
c.metrics.InboundRPCDropped(task.method, "context_done")
|
||||
return ctx.Err()
|
||||
case <-c.rpcStop:
|
||||
c.releaseInflightRPCBytes(task.size)
|
||||
c.metrics.InboundRPCDropped(task.method, "scheduler_closed")
|
||||
return ErrConnClosed
|
||||
default:
|
||||
c.releaseInflightRPCBytes(task.size)
|
||||
c.metrics.InboundRPCDropped(task.method, "queue_full")
|
||||
return ErrInboundRPCQueueFull
|
||||
}
|
||||
}
|
||||
|
||||
// releaseInflightRPCBytes 归还字节预算。与 enqueueInboundRPC 的预扣严格配对：
|
||||
// 入队失败时回滚、worker 执行完(runInboundRPC)或排空丢弃(drainInboundRPCQueue)时释放。
|
||||
func (c *Conn) releaseInflightRPCBytes(size int) {
|
||||
if size > 0 {
|
||||
c.inflightRPCBytes.Add(-int64(size))
|
||||
}
|
||||
}
|
||||
|
||||
func (c *Conn) inboundRPCWorker(rootCtx context.Context) {
|
||||
defer c.rpcWG.Done()
|
||||
for {
|
||||
select {
|
||||
case <-c.rpcStop:
|
||||
return
|
||||
default:
|
||||
}
|
||||
select {
|
||||
case task := <-c.rpcQueue:
|
||||
c.runInboundRPC(rootCtx, task)
|
||||
case <-c.rpcStop:
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (c *Conn) runInboundRPC(rootCtx context.Context, task inboundRPC) {
|
||||
defer c.releaseInflightRPCBytes(task.size)
|
||||
queueWait := time.Since(task.enqueuedAt)
|
||||
c.metrics.InboundRPCStarted(task.method, queueWait)
|
||||
c.metrics.InboundRPCStarted(task.method, now.Sub(task.enqueuedAt))
|
||||
ctx := task.ctx
|
||||
if ctx == nil {
|
||||
ctx = context.Background()
|
||||
if task.run != nil {
|
||||
_ = task.run(ctx)
|
||||
}
|
||||
// 合并两个取消源（task.ctx 与 rootCtx）+ 超时为最少的 context 层数：
|
||||
// WithTimeout/WithCancel 的 cancel 直接作为 AfterFunc 回调，省掉单独的中间层。
|
||||
var cancel context.CancelFunc
|
||||
if c.rpcTimeout > 0 {
|
||||
ctx, cancel = context.WithTimeout(ctx, c.rpcTimeout)
|
||||
} else {
|
||||
ctx, cancel = context.WithCancel(ctx)
|
||||
}
|
||||
|
||||
func (c *Conn) finishInboundRPC(task inboundRPC) {
|
||||
if task.ticket != nil {
|
||||
task.ticket.state.Store(inboundRPCTicketDone)
|
||||
}
|
||||
stopInboundRPCTask(task)
|
||||
var reschedule bool
|
||||
c.rpcMu.Lock()
|
||||
c.rpcRunning--
|
||||
c.inflightRPCBytes.Add(-int64(task.size))
|
||||
if !c.rpcClosed && len(c.rpcQueue) > 0 && c.rpcRunning < c.rpcMaxInflight && !c.rpcReady {
|
||||
c.rpcReady = true
|
||||
reschedule = true
|
||||
}
|
||||
c.rpcMu.Unlock()
|
||||
reservation := task.budget
|
||||
// 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()
|
||||
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
|
||||
// same callback only signals the running request's response gate; its body remains
|
||||
// owned until the handler exits.
|
||||
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))
|
||||
ticket.state.Store(inboundRPCTicketDone)
|
||||
found = true
|
||||
break
|
||||
}
|
||||
c.rpcMu.Unlock()
|
||||
|
||||
if unschedule {
|
||||
c.rpcScheduler.unschedule(c)
|
||||
}
|
||||
if found {
|
||||
method := task.method
|
||||
reservation := task.budget
|
||||
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()
|
||||
}
|
||||
return
|
||||
}
|
||||
if ticket.state.Load() == inboundRPCTicketRunning && ticket.onTimeout != nil {
|
||||
ticket.onTimeout()
|
||||
}
|
||||
}
|
||||
|
||||
// stopInboundRPCTask disarms callbacks before canceling the context so a normal
|
||||
// completion or connection close cannot manufacture an RPC_TIMEOUT response.
|
||||
// A deadline callback already in flight is harmless because enqueueRPC's response
|
||||
// gate makes timeout and normal rpc_result mutually exclusive.
|
||||
func stopInboundRPCTask(task inboundRPC) {
|
||||
if task.stopTimeout != nil {
|
||||
task.stopTimeout()
|
||||
}
|
||||
if task.stopRoot != nil {
|
||||
task.stopRoot()
|
||||
}
|
||||
if task.cancel != nil {
|
||||
task.cancel()
|
||||
}
|
||||
defer cancel()
|
||||
stopRoot := context.AfterFunc(rootCtx, cancel)
|
||||
defer stopRoot()
|
||||
_ = task.run(ctx)
|
||||
}
|
||||
|
||||
func (c *Conn) closeInboundRPCScheduler() {
|
||||
if c.rpcStop == nil {
|
||||
c.beginCloseInboundRPCScheduler()
|
||||
if c.rpcScheduler == nil {
|
||||
return
|
||||
}
|
||||
c.waitInboundShutdown(rpcCloseWaitTimeout)
|
||||
}
|
||||
|
||||
// beginCloseInboundRPCScheduler publishes closure, cancels running work and releases queued
|
||||
// requests without waiting for handlers. ForceClose uses this phase before transport.Close so a
|
||||
// pathological/blocking transport implementation cannot leave the RPC admission gate open.
|
||||
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()
|
||||
}
|
||||
close(c.rpcStop)
|
||||
// 抢占懒启动 Once：若 worker 尚未起，封住其启动，避免后续 ensureInboundRPCWorkers 的
|
||||
// rpcWG.Add 与下面的 rpcWG.Wait 并发（WaitGroup 误用）。Once 互斥保证 Add happens-before Wait。
|
||||
c.rpcWorkersOnce.Do(func() {})
|
||||
c.drainInboundRPCQueue()
|
||||
// 等 worker 退出，使关闭对 inbound 与 outbound（<-outboundDone）收敛对称；带超时防慢 handler 卡死。
|
||||
c.waitInboundWorkers(rpcCloseWaitTimeout)
|
||||
for i := range queued {
|
||||
task := queued[i]
|
||||
queued[i] = inboundRPC{}
|
||||
if task.ticket != nil {
|
||||
task.ticket.state.Store(inboundRPCTicketDone)
|
||||
}
|
||||
method := task.method
|
||||
reservation := task.budget
|
||||
stopInboundRPCTask(task)
|
||||
task = inboundRPC{}
|
||||
reservation.release()
|
||||
c.metrics.InboundRPCDropped(method, "connection_closed")
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// waitInboundWorkers 等所有 inbound RPC worker 退出，最长 timeout。超时则放弃等待，
|
||||
// worker 在其阻塞的底层调用返回后自行退出（rpcCancel 已发，最终收敛）。
|
||||
func (c *Conn) waitInboundWorkers(timeout time.Duration) {
|
||||
// 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:
|
||||
}
|
||||
}
|
||||
|
||||
func (c *Conn) drainInboundRPCQueue() {
|
||||
for {
|
||||
select {
|
||||
case task := <-c.rpcQueue:
|
||||
c.releaseInflightRPCBytes(task.size)
|
||||
c.metrics.InboundRPCDropped(task.method, "connection_closed")
|
||||
default:
|
||||
return
|
||||
}
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue