merged from gramsrv upstream

internal/mtprotoedge/rpc_result_egress.go

@@ -8,24 +8,31 @@ import (
 	"strings"
 	"sync"
 	"sync/atomic"
+	"time"
 
 	"github.com/iamxvbaba/td/bin"
 	"github.com/iamxvbaba/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
+	rpcResultGZIPMinBytes            = 4 << 10
+	rpcResultGZIPMaxInputBytes       = (10 << 20) - 1 // gotd client decompression hard limit.
+	rpcResultGZIPMinSavedBytes       = 1 << 10
+	rpcResultGZIPMinSavedDivisor     = 12 // Require roughly 8.3% reduction.
+	rpcResultGZIPConcurrency         = 8
+	defaultRPCDeliveryHookWorkers    = 32
+	defaultRPCDeliveryHookMaxPending = 16_384
 )
 
 var rpcResultGZIPSlots = make(chan struct{}, rpcResultGZIPConcurrency)
 
-var defaultRPCDeliveryHookExecutor = newRPCDeliveryHookExecutor(rpcDeliveryHookConcurrency, rpcDeliveryHookQueueSize)
+// This executor is only a compatibility boundary for directly constructed
+// Conn/encoded-message tests. Production Conns always use their owning Server's
+// isolated executor.
+var defaultRPCDeliveryHookExecutor = newRPCDeliveryHookExecutor(
+	defaultRPCDeliveryHookWorkers,
+	defaultRPCDeliveryHookMaxPending,
+)
 
 // ErrRPCDeliveryHookCapacity means an RPC result with a delivery-dependent
 // transition cannot reserve reliable executor capacity. The result must not be
@@ -53,20 +60,43 @@ type rpcDeliveryHookJob struct {
 	fn     func()
 }
 
-// rpcDeliveryHookExecutor has process lifetime. Capacity bounds queued plus
-// running hooks; every physical write reserves a ticket before admission. A
-// successful writer therefore performs only one short O(1) queue append and
-// never waits for capacity or hook work. Failed writes release their ticket,
-// while the shared logical coordinator remains eligible for a later replay.
+// rpcDeliveryHookExecutor is owned by one Server. Capacity bounds reserved plus
+// queued plus running hooks; every physical write reserves a ticket before
+// admission. A successful writer therefore performs only one short O(1) queue
+// append and never waits for capacity or hook work. Failed writes release their
+// ticket, while the shared logical coordinator remains eligible for a later
+// replay.
 type rpcDeliveryHookExecutor struct {
-	slots chan struct{}
+	workers  int
+	capacity int
+	slots    chan struct{}
+	start    sync.Once
+	wg       sync.WaitGroup
 
-	mu   sync.Mutex
-	cond *sync.Cond
-	head *rpcDeliveryHookJob
-	tail *rpcDeliveryHookJob
+	mu       sync.Mutex
+	cond     *sync.Cond
+	head     *rpcDeliveryHookJob
+	tail     *rpcDeliveryHookJob
+	stopping bool
 
-	panics atomic.Uint64
+	queued        atomic.Int64
+	running       atomic.Int64
+	completed     atomic.Uint64
+	rejected      atomic.Uint64
+	panics        atomic.Uint64
+	durationNanos atomic.Uint64
+}
+
+type rpcDeliveryHookRuntimeSnapshot struct {
+	workers         int64
+	capacity        int64
+	reserved        int64
+	queued          int64
+	running         int64
+	completed       uint64
+	rejected        uint64
+	panics          uint64
+	durationSeconds float64
 }
 
 func newRPCDeliveryHookExecutor(workers, capacity int) *rpcDeliveryHookExecutor {
@@ -76,24 +106,50 @@ func newRPCDeliveryHookExecutor(workers, capacity int) *rpcDeliveryHookExecutor
 	if capacity < workers {
 		capacity = workers
 	}
-	e := &rpcDeliveryHookExecutor{slots: make(chan struct{}, capacity)}
-	e.cond = sync.NewCond(&e.mu)
-	for range workers {
-		go e.run()
+	e := &rpcDeliveryHookExecutor{
+		workers:  workers,
+		capacity: capacity,
+		slots:    make(chan struct{}, capacity),
 	}
+	e.cond = sync.NewCond(&e.mu)
 	return e
 }
 
+func (e *rpcDeliveryHookExecutor) startWorkers() {
+	if e == nil {
+		return
+	}
+	e.start.Do(func() {
+		e.mu.Lock()
+		defer e.mu.Unlock()
+		if e.stopping {
+			return
+		}
+		e.wg.Add(e.workers)
+		for range e.workers {
+			go e.run()
+		}
+	})
+}
+
 func (e *rpcDeliveryHookExecutor) reserve() (*rpcDeliveryHookTicket, bool) {
 	if e == nil {
 		return nil, false
 	}
+	e.startWorkers()
+	e.mu.Lock()
+	defer e.mu.Unlock()
+	if e.stopping {
+		e.rejected.Add(1)
+		return nil, false
+	}
 	select {
 	case e.slots <- struct{}{}:
 		ticket := &rpcDeliveryHookTicket{executor: e}
 		ticket.state.Store(uint32(rpcDeliveryHookTicketReserved))
 		return ticket, true
 	default:
+		e.rejected.Add(1)
 		return nil, false
 	}
 }
@@ -105,6 +161,7 @@ func (t *rpcDeliveryHookTicket) release() {
 		return
 	}
 	<-t.executor.slots
+	t.executor.signalStateChange()
 }
 
 func (t *rpcDeliveryHookTicket) submit(fn func()) bool {
@@ -127,14 +184,20 @@ func (e *rpcDeliveryHookExecutor) enqueue(job *rpcDeliveryHookJob) {
 		e.tail.next = job
 	}
 	e.tail = job
+	e.queued.Add(1)
 	e.cond.Signal()
 	e.mu.Unlock()
 }
 
 func (e *rpcDeliveryHookExecutor) run() {
+	defer e.wg.Done()
 	for {
 		e.mu.Lock()
 		for e.head == nil {
+			if e.stopping && len(e.slots) == 0 {
+				e.mu.Unlock()
+				return
+			}
 			e.cond.Wait()
 		}
 		job := e.head
@@ -143,27 +206,88 @@ func (e *rpcDeliveryHookExecutor) run() {
 			e.tail = nil
 		}
 		job.next = nil
+		e.queued.Add(-1)
+		e.running.Add(1)
 		e.mu.Unlock()
 		e.runOne(job)
 	}
 }
 
 func (e *rpcDeliveryHookExecutor) runOne(job *rpcDeliveryHookJob) {
+	started := time.Now()
 	defer func() {
 		if recovered := recover(); recovered != nil {
 			e.panics.Add(1)
 			log.Printf("mtprotoedge: rpc delivery hook panic: %v\n%s", recovered, debug.Stack())
 		}
+		e.durationNanos.Add(uint64(time.Since(started)))
+		e.completed.Add(1)
+		e.running.Add(-1)
 		if job != nil && job.ticket != nil {
 			job.ticket.state.Store(uint32(rpcDeliveryHookTicketDone))
 			<-e.slots
 		}
+		e.signalStateChange()
 	}()
 	if job != nil && job.fn != nil {
 		job.fn()
 	}
 }
 
+func (e *rpcDeliveryHookExecutor) signalStateChange() {
+	if e == nil {
+		return
+	}
+	e.mu.Lock()
+	e.cond.Broadcast()
+	e.mu.Unlock()
+}
+
+// stop rejects new reservations and lets every already-reserved ticket either
+// be released or submitted and executed. Timing out never abandons jobs: the
+// existing workers continue draining under their Server-owned executor.
+func (e *rpcDeliveryHookExecutor) stop(timeout time.Duration) bool {
+	if e == nil {
+		return true
+	}
+	e.mu.Lock()
+	e.stopping = true
+	e.cond.Broadcast()
+	e.mu.Unlock()
+	done := make(chan struct{})
+	go func() {
+		e.wg.Wait()
+		close(done)
+	}()
+	if timeout <= 0 {
+		<-done
+		return true
+	}
+	select {
+	case <-done:
+		return true
+	case <-time.After(timeout):
+		return false
+	}
+}
+
+func (e *rpcDeliveryHookExecutor) runtimeSnapshot() rpcDeliveryHookRuntimeSnapshot {
+	if e == nil {
+		return rpcDeliveryHookRuntimeSnapshot{}
+	}
+	return rpcDeliveryHookRuntimeSnapshot{
+		workers:         int64(e.workers),
+		capacity:        int64(e.capacity),
+		reserved:        int64(len(e.slots)),
+		queued:          e.queued.Load(),
+		running:         e.running.Load(),
+		completed:       e.completed.Load(),
+		rejected:        e.rejected.Load(),
+		panics:          e.panics.Load(),
+		durationSeconds: float64(e.durationNanos.Load()) / float64(time.Second),
+	}
+}
+
 // 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.