feat: sync multilayer td integration

internal/mtprotoedge/rpc_flight.go

@@ -5,6 +5,8 @@ import (
 	"errors"
 	"sync"
 	"sync/atomic"
+
+	"github.com/iamxvbaba/td/tg"
 )
 
 const rpcResultFlightDefaultMaxPending = 8192
@@ -12,10 +14,49 @@ const rpcResultFlightDefaultMaxPending = 8192
 var (
 	// ErrRPCResultFlightCapacity is returned before installing a new owner when
 	// the process-wide in-flight claim table has reached its hard bound.
-	ErrRPCResultFlightCapacity = errors.New("mtproto rpc result in-flight capacity exhausted")
-	ErrRPCResultFlightInvalid  = errors.New("mtproto rpc result in-flight claim is invalid")
+	ErrRPCResultFlightCapacity     = errors.New("mtproto rpc result in-flight capacity exhausted")
+	ErrRPCResultSubscriberCapacity = errors.New("mtproto rpc result subscriber capacity exhausted")
+	ErrRPCResultFlightInvalid      = errors.New("mtproto rpc result in-flight claim is invalid")
+	ErrRPCResultIdentityMismatch   = errors.New("mtproto rpc result request identity mismatch")
+	ErrRPCAdmissionSeqExhausted    = errors.New("mtproto rpc admission sequence exhausted")
 )
 
+// rpcResultIdentityMismatchError carries the winner's immutable admission
+// profile from inside the cache shard critical section. A replacement Conn can
+// re-decode the same naked body under that grammar even if the winner aborts
+// immediately after the mismatch is returned.
+type rpcResultIdentityMismatchError struct {
+	profile    tg.LayerProfile
+	hasProfile bool
+}
+
+func (e *rpcResultIdentityMismatchError) Error() string { return ErrRPCResultIdentityMismatch.Error() }
+func (e *rpcResultIdentityMismatchError) Is(target error) bool {
+	return target == ErrRPCResultIdentityMismatch
+}
+
+func identityMismatch(identity rpcResultRequestIdentity) error {
+	return &rpcResultIdentityMismatchError{profile: identity.profile, hasProfile: identity.valid && identity.profile != 0}
+}
+
+type rpcResultRequestIdentity struct {
+	exact tg.LayerPreparedCallIdentity
+	// profile is retained separately because LayerPreparedCallIdentity is opaque.
+	// It lets a same-msg_id replay be re-admitted with the original request
+	// grammar after the session default has moved to another Layer.
+	profile tg.LayerProfile
+	valid   bool
+}
+
+func (i rpcResultRequestIdentity) matches(requested rpcResultRequestIdentity) bool {
+	if !requested.valid {
+		// Legacy service/test callers carry no API request identity and preserve
+		// the historical cache lookup behavior. Exact callers must always match.
+		return true
+	}
+	return i.valid && i.exact == requested.exact
+}
+
 type rpcResultAcquireState uint8
 
 const (
@@ -32,20 +73,35 @@ const (
 //   - pending: waiter joins the already-running owner;
 //   - owner: owner must eventually complete through rpcResultCache.Put or Abort.
 type rpcResultAcquire struct {
-	state   rpcResultAcquireState
-	encoded *encodedOutboundMessage
-	waiter  *rpcResultWaiter
-	owner   *rpcResultOwnerLease
+	state          rpcResultAcquireState
+	admissionSeq   uint64
+	encoded        *encodedOutboundMessage
+	waiter         *rpcResultWaiter
+	owner          *rpcResultOwnerLease
+	executionKnown bool
+	executionOK    bool
 }
 
-// rpcResultFlight is not part of the completed cache LRU/TTL lifecycle. Its
+// rpcResultFlight is not part of the completed cache TTL lifecycle. Its
 // done channel is closed exactly once while holding the owning cache shard lock;
 // channel close publishes encoded/ok to all waiters without a waiter goroutine.
 type rpcResultFlight struct {
-	done        chan struct{}
-	encoded     *encodedOutboundMessage
-	ok          bool
-	subscribers []func(*encodedOutboundMessage, bool)
+	done                 chan struct{}
+	encoded              *encodedOutboundMessage
+	ok                   bool
+	subscribers          []func(*encodedOutboundMessage, bool)
+	executionDone        bool
+	executionOK          bool
+	executionSubscribers []func(bool)
+	// subscriberSlots counts callbacks retained by this pending flight. Result
+	// and execution callbacks are charged independently; a replay alias installs
+	// both atomically so a capacity failure cannot leave half an alias behind.
+	subscriberSlots int
+	identity        rpcResultRequestIdentity
+	admissionSeq    uint64
+	// reservation owns one entry and at least one byte at global, raw-auth and
+	// session scopes. Put transfers it to a result/tombstone; Abort releases it.
+	reservation *rpcResultBudgetReservation
 }
 
 type rpcResultWaiter struct {
@@ -86,31 +142,107 @@ func (w *rpcResultWaiter) Wait(ctx context.Context) (encoded *encodedOutboundMes
 // 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 {
+	if fn == nil {
+		return ErrRPCResultFlightInvalid
+	}
+	return w.subscribe(fn, nil)
+}
+
+// SubscribeExecution registers a non-blocking callback for the terminal
+// business outcome. This is deliberately separate from physical rpc_result
+// delivery: invokeAfter* ordering depends on handler completion and must not
+// occupy a shared RPC worker while a socket write is pending.
+func (w *rpcResultWaiter) SubscribeExecution(fn func(bool)) error {
+	if fn == nil {
+		return ErrRPCResultFlightInvalid
+	}
+	return w.subscribe(nil, fn)
+}
+
+// SubscribeResultAndExecution atomically installs both halves of one pending
+// replay alias. Either both callbacks are retained, or neither is; this avoids
+// an execution-only closure surviving when the result subscriber hits a hard
+// capacity limit.
+func (w *rpcResultWaiter) SubscribeResultAndExecution(
+	resultFn func(*encodedOutboundMessage, bool),
+	executionFn func(bool),
+) error {
+	if resultFn == nil || executionFn == nil {
+		return ErrRPCResultFlightInvalid
+	}
+	return w.subscribe(resultFn, executionFn)
+}
+
+func (w *rpcResultWaiter) subscribe(
+	resultFn func(*encodedOutboundMessage, bool),
+	executionFn func(bool),
+) error {
+	if w == nil || w.cache == nil || w.flight == nil || (resultFn == nil && executionFn == nil) {
 		return ErrRPCResultFlightInvalid
 	}
 	s := w.cache.shard(w.key)
 	var (
-		encoded *encodedOutboundMessage
-		ok      bool
-		ready   bool
+		encoded        *encodedOutboundMessage
+		resultOK       bool
+		resultReady    = resultFn == nil
+		executionOK    bool
+		executionReady = executionFn == nil
 	)
 	s.mu.Lock()
 	if flight, exists := s.pending[w.key]; exists && flight == w.flight {
-		flight.subscribers = append(flight.subscribers, fn)
+		slots := 0
+		if resultFn != nil {
+			slots++
+		}
+		if executionFn != nil && !flight.executionDone {
+			slots++
+		}
+		if slots > 0 {
+			if flight.subscriberSlots > w.cache.subscriberPerFlight-slots ||
+				!w.cache.subscriberBudget.reserve(w.key, slots) {
+				s.mu.Unlock()
+				return ErrRPCResultSubscriberCapacity
+			}
+			flight.subscriberSlots += slots
+		}
+		if resultFn != nil {
+			flight.subscribers = append(flight.subscribers, resultFn)
+		}
+		if executionFn != nil {
+			if flight.executionDone {
+				executionOK, executionReady = flight.executionOK, true
+			} else {
+				flight.executionSubscribers = append(flight.executionSubscribers, executionFn)
+			}
+		}
 		s.mu.Unlock()
+		if executionReady && executionFn != nil {
+			executionFn(executionOK)
+		}
 		return nil
 	}
-	select {
-	case <-w.flight.done:
-		encoded, ok, ready = w.flight.encoded, w.flight.ok, true
-	default:
+	if resultFn != nil {
+		select {
+		case <-w.flight.done:
+			encoded, resultOK, resultReady = w.flight.encoded, w.flight.ok, true
+		default:
+		}
+	}
+	if executionFn != nil && w.flight.executionDone {
+		executionOK, executionReady = w.flight.executionOK, true
 	}
 	s.mu.Unlock()
-	if !ready {
+	// Atomic late subscription also means no callback runs unless every requested
+	// terminal state is available.
+	if !resultReady || !executionReady {
 		return ErrRPCResultFlightInvalid
 	}
-	fn(encoded, ok)
+	if executionFn != nil {
+		executionFn(executionOK)
+	}
+	if resultFn != nil {
+		resultFn(encoded, resultOK)
+	}
 	return nil
 }
 
@@ -175,6 +307,33 @@ func (l *rpcResultOwnerLease) Delivery() *rpcResultDelivery {
 	return l.delivery
 }
 
+// CompleteExecution publishes the handler outcome exactly once while this
+// lease still owns the flight. success=false includes RPC errors, internal
+// failures and dependency failures. Delivery/cache completion remains a
+// separate later transition.
+func (l *rpcResultOwnerLease) CompleteExecution(success bool) bool {
+	if l == nil || l.cache == nil || l.flight == nil {
+		return false
+	}
+	s := l.cache.shard(l.key)
+	s.mu.Lock()
+	flight, ok := s.pending[l.key]
+	if !ok || flight != l.flight || flight.executionDone {
+		s.mu.Unlock()
+		return false
+	}
+	flight.executionDone = true
+	flight.executionOK = success
+	subscribers := append([]func(bool){}, flight.executionSubscribers...)
+	flight.executionSubscribers = nil
+	l.cache.releaseFlightSubscriberSlotsLocked(l.key, flight, len(subscribers))
+	s.mu.Unlock()
+	for _, subscriber := range subscribers {
+		subscriber(success)
+	}
+	return true
+}
+
 // 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.
@@ -216,9 +375,23 @@ func (l *rpcResultOwnerLease) Abort() bool {
 		return false
 	}
 	delete(s.pending, l.key)
+	if flight.reservation != nil {
+		flight.reservation.release()
+		flight.reservation = nil
+	}
 	l.cache.flightLimit.release()
+	l.cache.activeAdmissions.retire(flight.admissionSeq)
 	subscribers := append([]func(*encodedOutboundMessage, bool){}, flight.subscribers...)
 	flight.subscribers = nil
+	executionSubscribers := append([]func(bool){}, flight.executionSubscribers...)
+	flight.executionSubscribers = nil
+	l.cache.releaseFlightSubscriberSlotsLocked(
+		l.key, flight, len(subscribers)+len(executionSubscribers),
+	)
+	if !flight.executionDone {
+		flight.executionDone = true
+		flight.executionOK = false
+	}
 	close(flight.done)
 	s.mu.Unlock()
 	l.hookMu.Lock()
@@ -231,6 +404,9 @@ func (l *rpcResultOwnerLease) Abort() bool {
 	for _, subscriber := range subscribers {
 		subscriber(nil, false)
 	}
+	for _, subscriber := range executionSubscribers {
+		subscriber(false)
+	}
 	return true
 }
 
@@ -240,25 +416,39 @@ type rpcResultFlightLimit struct {
 }
 
 func (l *rpcResultFlightLimit) reserve() bool {
+	return l.reserveN(1)
+}
+
+func (l *rpcResultFlightLimit) reserveN(delta int64) bool {
 	if l == nil || l.max <= 0 {
 		return false
 	}
+	if delta <= 0 {
+		return false
+	}
 	for {
 		used := l.used.Load()
-		if used >= l.max {
+		if used < 0 || used > l.max-delta {
 			return false
 		}
-		if l.used.CompareAndSwap(used, used+1) {
+		if l.used.CompareAndSwap(used, used+delta) {
 			return true
 		}
 	}
 }
 
 func (l *rpcResultFlightLimit) release() {
+	l.releaseN(1)
+}
+
+func (l *rpcResultFlightLimit) releaseN(delta int64) {
 	if l == nil {
 		return
 	}
-	if remaining := l.used.Add(-1); remaining < 0 {
+	if delta <= 0 {
+		panic("mtproto rpc result counter release must be positive")
+	}
+	if remaining := l.used.Add(-delta); remaining < 0 {
 		// Put/Abort use map removal and lease identity to make double release
 		// impossible. Fail fast instead of masking a capacity-accounting bug that
 		// could otherwise admit more owners than the configured hard limit.
@@ -275,45 +465,159 @@ func (l *rpcResultFlightLimit) snapshot() int64 {
 
 // Acquire atomically returns a completed result, joins the existing in-flight
 // owner, or installs the unique owner lease. Pending entries have a separate
-// lifecycle from completed cache trim/TTL and consume one process-wide slot.
+// lifecycle from completed cache TTL but reserve the same global/auth/session
+// ownership that Put later transfers to a completed result or tombstone.
 func (c *rpcResultCache) Acquire(authKeyID [8]byte, sessionID, reqMsgID int64) (rpcResultAcquire, error) {
+	return c.acquire(authKeyID, sessionID, reqMsgID, rpcResultRequestIdentity{})
+}
+
+func (c *rpcResultCache) AcquireIdentified(
+	authKeyID [8]byte,
+	sessionID, reqMsgID int64,
+	identity tg.LayerPreparedCallIdentity,
+) (rpcResultAcquire, error) {
+	return c.acquire(authKeyID, sessionID, reqMsgID, rpcResultRequestIdentity{exact: identity, valid: true})
+}
+
+// AcquireLayerIdentified is the production exact-RPC claim. In addition to the
+// immutable full request identity it retains the admission profile required to
+// decode a later same-msg_id naked replay under its original grammar.
+func (c *rpcResultCache) AcquireLayerIdentified(
+	authKeyID [8]byte,
+	sessionID, reqMsgID int64,
+	profile tg.LayerProfile,
+	identity tg.LayerPreparedCallIdentity,
+) (rpcResultAcquire, error) {
+	return c.acquire(authKeyID, sessionID, reqMsgID, rpcResultRequestIdentity{
+		exact: identity, profile: profile, valid: true,
+	})
+}
+
+// ExactAdmissionProfile returns the immutable profile of an existing exact
+// owner/result. It does not create or join a flight. Callers still perform
+// AcquireLayerIdentified after decode, which atomically rejects a same-msg_id
+// body change by comparing the full prepared identity.
+func (c *rpcResultCache) ExactAdmissionProfile(authKeyID [8]byte, sessionID, reqMsgID int64) (tg.LayerProfile, bool) {
+	if c == nil || reqMsgID == 0 {
+		return 0, false
+	}
+	key := rpcResultCacheKey{authKeyID: authKeyID, sessionID: sessionID, reqMsgID: reqMsgID}
+	s := c.shard(key)
+	now := s.now()
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	if elem := s.byKey[key]; elem != nil {
+		entry := elem.Value.(*rpcResultCacheEntry)
+		if entry.expiresAt.After(now) {
+			if entry.identity.valid && entry.identity.profile != 0 {
+				return entry.identity.profile, true
+			}
+			return 0, false
+		}
+		s.removeElement(elem)
+	}
+	if flight := s.pending[key]; flight != nil && flight.identity.valid && flight.identity.profile != 0 {
+		return flight.identity.profile, true
+	}
+	return 0, false
+}
+
+func (c *rpcResultCache) acquire(
+	authKeyID [8]byte,
+	sessionID, reqMsgID int64,
+	identity rpcResultRequestIdentity,
+) (rpcResultAcquire, error) {
 	if c == nil || reqMsgID == 0 {
 		return rpcResultAcquire{}, ErrRPCResultFlightInvalid
 	}
 	key := rpcResultCacheKey{authKeyID: authKeyID, sessionID: sessionID, reqMsgID: reqMsgID}
 	s := c.shard(key)
-	now := s.now()
+	reclaimedExpired := false
+	for {
+		now := s.now()
+		s.mu.Lock()
+		s.expireLocked(now)
 
-	s.mu.Lock()
-	defer s.mu.Unlock()
-
-	if elem, ok := s.byKey[key]; ok {
-		entry := elem.Value.(*rpcResultCacheEntry)
-		if entry.expiresAt.After(now) {
-			return rpcResultAcquire{state: rpcResultAcquireCompleted, encoded: entry.encoded}, nil
+		if elem, ok := s.byKey[key]; ok {
+			entry := elem.Value.(*rpcResultCacheEntry)
+			if !entry.identity.matches(identity) {
+				s.mu.Unlock()
+				return rpcResultAcquire{}, identityMismatch(entry.identity)
+			}
+			if entry.capacity || entry.encoded == nil {
+				s.mu.Unlock()
+				return rpcResultAcquire{}, ErrRPCResultFlightCapacity
+			}
+			result := rpcResultAcquire{
+				state: rpcResultAcquireCompleted, admissionSeq: entry.admissionSeq, encoded: entry.encoded,
+				executionKnown: entry.executionKnown, executionOK: entry.executionOK,
+			}
+			s.mu.Unlock()
+			return result, nil
 		}
-		s.removeElement(elem)
-	}
-	if flight, ok := s.pending[key]; ok {
+		if flight, ok := s.pending[key]; ok {
+			if !flight.identity.matches(identity) {
+				s.mu.Unlock()
+				return rpcResultAcquire{}, identityMismatch(flight.identity)
+			}
+			result := rpcResultAcquire{
+				state:        rpcResultAcquirePending,
+				admissionSeq: flight.admissionSeq,
+				waiter:       &rpcResultWaiter{cache: c, key: key, flight: flight},
+			}
+			s.mu.Unlock()
+			return result, nil
+		}
+		if s.maxEntries > 0 && len(s.byKey)+len(s.pending) >= s.maxEntries {
+			s.mu.Unlock()
+			return rpcResultAcquire{}, ErrRPCResultFlightCapacity
+		}
+		if !c.flightLimit.reserve() {
+			s.mu.Unlock()
+			return rpcResultAcquire{}, ErrRPCResultFlightCapacity
+		}
+		reservation := c.fairBudget.reserveOwner(key)
+		if reservation == nil {
+			c.flightLimit.release()
+			s.mu.Unlock()
+			if reclaimedExpired {
+				return rpcResultAcquire{}, ErrRPCResultFlightCapacity
+			}
+			// Expired rows in another full-key shard may be the only consumers at
+			// the global, auth or session scope. Reap once, then retry every identity
+			// and capacity check because another goroutine may have won this key.
+			c.expireCompletedResults()
+			reclaimedExpired = true
+			continue
+		}
+		var admissionSeq uint64
+		if identity.valid {
+			var err error
+			admissionSeq, err = c.activeAdmissions.allocateAndRegister(&c.nextAdmissionSeq)
+			if err != nil {
+				reservation.release()
+				c.flightLimit.release()
+				s.mu.Unlock()
+				return rpcResultAcquire{}, err
+			}
+		}
+		flight := &rpcResultFlight{
+			done: make(chan struct{}), identity: identity, admissionSeq: admissionSeq,
+			reservation: reservation,
+		}
+		if s.pending == nil {
+			s.pending = make(map[rpcResultCacheKey]*rpcResultFlight)
+		}
+		s.pending[key] = flight
+		s.mu.Unlock()
 		return rpcResultAcquire{
-			state:  rpcResultAcquirePending,
-			waiter: &rpcResultWaiter{cache: c, key: key, flight: flight},
+			state:        rpcResultAcquireOwner,
+			admissionSeq: admissionSeq,
+			owner: &rpcResultOwnerLease{
+				cache: c, key: key, flight: flight, delivery: newRPCResultDelivery(reqMsgID),
+			},
 		}, nil
 	}
-	if !c.flightLimit.reserve() {
-		return rpcResultAcquire{}, ErrRPCResultFlightCapacity
-	}
-	flight := &rpcResultFlight{done: make(chan struct{})}
-	if s.pending == nil {
-		s.pending = make(map[rpcResultCacheKey]*rpcResultFlight)
-	}
-	s.pending[key] = flight
-	return rpcResultAcquire{
-		state: rpcResultAcquireOwner,
-		owner: &rpcResultOwnerLease{
-			cache: c, key: key, flight: flight, delivery: newRPCResultDelivery(reqMsgID),
-		},
-	}, nil
 }
 
 // completeRPCResultFlightLocked publishes encoded to the current owner claim.
@@ -322,20 +626,62 @@ func (c *rpcResultCache) completeRPCResultFlightLocked(
 	s *rpcResultCacheShard,
 	key rpcResultCacheKey,
 	encoded *encodedOutboundMessage,
-) []func(*encodedOutboundMessage, bool) {
+) (
+	[]func(*encodedOutboundMessage, bool),
+	[]func(bool),
+	bool,
+) {
 	if c == nil || s == nil || encoded == nil {
-		return nil
+		return nil, nil, false
 	}
 	flight, ok := s.pending[key]
 	if !ok {
-		return nil
+		return nil, nil, false
 	}
 	delete(s.pending, key)
+	if flight.reservation != nil {
+		flight.reservation.releasePending()
+		// The completed entry now owns the same reservation.
+		flight.reservation = nil
+	}
 	flight.encoded = encoded
 	flight.ok = true
 	c.flightLimit.release()
+	c.activeAdmissions.retire(flight.admissionSeq)
 	subscribers := append([]func(*encodedOutboundMessage, bool){}, flight.subscribers...)
 	flight.subscribers = nil
+	executionSubscribers := append([]func(bool){}, flight.executionSubscribers...)
+	flight.executionSubscribers = nil
+	executionOK := flight.executionOK
+	if !flight.executionDone {
+		// A result without an explicit handler-completion proof cannot satisfy an
+		// invokeAfter dependency. Production completes execution before Put; this
+		// branch is the conservative terminal cleanup for defensive callers.
+		flight.executionDone = true
+		flight.executionOK = false
+		executionOK = false
+	}
+	c.releaseFlightSubscriberSlotsLocked(
+		key, flight, len(subscribers)+len(executionSubscribers),
+	)
 	close(flight.done)
-	return subscribers
+	return subscribers, executionSubscribers, executionOK
+}
+
+// releaseFlightSubscriberSlotsLocked releases callbacks detached from a flight.
+// The caller holds that flight's cache shard lock, preserving the only lock
+// order used by subscription: shard -> subscriber budget.
+func (c *rpcResultCache) releaseFlightSubscriberSlotsLocked(
+	key rpcResultCacheKey,
+	flight *rpcResultFlight,
+	slots int,
+) {
+	if slots == 0 {
+		return
+	}
+	if c == nil || flight == nil || slots < 0 || flight.subscriberSlots < slots {
+		panic("mtproto rpc result subscriber slot underflow")
+	}
+	flight.subscriberSlots -= slots
+	c.subscriberBudget.release(key, slots)
 }