feat: sync multilayer td integration

This commit is contained in:
A 2026-07-15 13:32:06 +08:00
parent 20a310f6ca
commit 766c5db992
491 changed files with 26235 additions and 35340 deletions

View file

@ -3,15 +3,18 @@ package mtprotoedge
import (
"context"
"crypto/sha256"
"fmt"
"log"
"runtime/debug"
"sync"
"sync/atomic"
"time"
"go.uber.org/zap"
"github.com/gotd/td/bin"
"github.com/gotd/td/proto"
"github.com/gotd/td/tg"
"github.com/iamxvbaba/td/bin"
"github.com/iamxvbaba/td/proto"
"github.com/iamxvbaba/td/tg"
)
// rpcRewrapRegistry links only an explicit official-client transition:
@ -37,6 +40,9 @@ type rpcRewrapSessionKey struct {
type rpcRewrapKey struct {
rpcRewrapSessionKey
fingerprint [sha256.Size]byte
semantic tg.LayerSemanticRequestIdentity
call tg.LayerCallIdentity
exact bool
}
type rpcRewrapRequestKey struct {
@ -102,6 +108,59 @@ func (r *rpcRewrapRegistry) register(c *Conn, body []byte, reqMsgID int64, metho
return true
}
func (r *rpcRewrapRegistry) registerSemantic(
c *Conn,
identity tg.LayerSemanticRequestIdentity,
call tg.LayerCallIdentity,
reqMsgID int64,
method string,
owner *rpcResultOwnerLease,
) bool {
if identity.Method() == 0 || identity.CanonicalSize() <= 0 {
return false
}
return r.registerKey(c, rpcRewrapKey{
rpcRewrapSessionKey: rpcRewrapSessionKey{authKeyID: c.authKeyID, sessionID: c.sessionID},
semantic: identity,
call: call,
exact: true,
}, reqMsgID, method, owner)
}
func (r *rpcRewrapRegistry) registerKey(c *Conn, key rpcRewrapKey, reqMsgID int64, method string, owner *rpcResultOwnerLease) bool {
if r == nil || c == nil || c.rpcRewrapInitialized.Load() || owner == nil {
return false
}
candidate := &rpcRewrapCandidate{
active: true, key: key, source: c, reqMsgID: reqMsgID, method: method,
owner: owner, waiter: owner.Waiter(),
}
if candidate.waiter == nil {
return false
}
session := key.rpcRewrapSessionKey
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
@ -122,6 +181,37 @@ func (r *rpcRewrapRegistry) claim(c *Conn, inner []byte) *rpcRewrapCandidate {
return nil
}
func (r *rpcRewrapRegistry) claimSemantic(
c *Conn,
identity tg.LayerSemanticRequestIdentity,
call tg.LayerCallIdentity,
) *rpcRewrapCandidate {
if r == nil || c == nil || identity.Method() == 0 || identity.CanonicalSize() <= 0 {
return nil
}
return r.claimKey(rpcRewrapKey{
rpcRewrapSessionKey: rpcRewrapSessionKey{authKeyID: c.authKeyID, sessionID: c.sessionID},
semantic: identity,
call: call,
exact: true,
})
}
func (r *rpcRewrapRegistry) claimKey(key rpcRewrapKey) *rpcRewrapCandidate {
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
@ -276,107 +366,636 @@ func decodeRPCRewrapInit(body []byte) (rpcRewrapInit, bool) {
}
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
conn *Conn
itemIndex int
newReqID int64
method string
oldWaiter *rpcResultWaiter
newOwner *rpcResultOwnerLease
sourceConn *Conn
sourceOwner *rpcResultOwnerLease
retargeted atomic.Bool
observeInit bool
init rpcRewrapInit
candidate *rpcRewrapCandidate
registry *rpcRewrapRegistry
afterSuccessfulDelivery func() error
finishReplayRestore func()
afterOnce sync.Once
deliveredFinalizeOnce sync.Once
deliveredFinalizeErr error
resultStoreClaimed atomic.Bool
executionOK atomic.Bool
// bodyReservation pins the replay/retarget clone from before allocation
// through queue residence. It is concurrency-safe because the watchdog and
// outbound actor race to release or take the same one-shot ownership token.
bodyReservation *outboundBodyReservation
}
func (a *rpcRewrapAlias) beginReplayRestore() {
if a == nil || a.conn == nil || a.finishReplayRestore != nil {
return
}
a.finishReplayRestore = a.conn.beginRPCReplayRestore()
}
func (a *rpcRewrapAlias) runAfterSuccessfulDelivery() (err error) {
if a == nil || a.afterSuccessfulDelivery == nil {
return nil
}
a.afterOnce.Do(func() {
if a.executionOK.Load() {
err = a.afterSuccessfulDelivery()
}
})
return err
}
func (a *rpcRewrapAlias) finishReplayRestoreWithoutDelivery() {
if a == nil {
return
}
a.releaseBodyReservation()
// Win or wait for any concurrent callback before dropping the barrier.
a.afterOnce.Do(func() {})
a.releaseReplayRestoreBarrier()
}
func (a *rpcRewrapAlias) releaseBodyReservation() {
if a != nil && a.bodyReservation != nil {
a.bodyReservation.release()
}
}
func (a *rpcRewrapAlias) releaseReplayRestoreBarrier() {
if a != nil && a.finishReplayRestore != nil {
a.finishReplayRestore()
}
}
func (a *rpcRewrapAlias) releaseDeferredLogicalHook() {
if a == nil || a.sourceOwner == nil || a.sourceOwner.Delivery() == nil ||
a.sourceOwner.Delivery().coordinator == nil {
return
}
a.sourceOwner.Delivery().coordinator.releaseDeferredHook()
}
func (a *rpcRewrapAlias) storeResultOnce(s *Server, encoded *encodedOutboundMessage) {
if a == nil || s == nil || encoded == nil || !a.resultStoreClaimed.CompareAndSwap(false, true) {
return
}
s.storeRPCResult(a.conn, a.newReqID, encoded)
}
func claimRPCRewrapLogicalHook(
ctx context.Context,
encoded *encodedOutboundMessage,
) (*rpcResultDeliveryHookClaim, error) {
if encoded == nil {
return nil, nil
}
// Only this alias is allowed to consume a sticky TryRetarget deferral. If a
// late physical success races another replacement replay, the coordinator
// waits for its Claimed/InProgress hook to publish Done.
return encoded.claimLogicalDeliveryHook(ctx, true)
}
// completeDeliveredRPCRewrapResult is safe after a watchdog has already fenced
// this physical generation. The caller has independent proof that the
// retargeted bytes reached the stream; deliveredFinalizeOnce, the shared hook
// coordinator and cache publication make late/concurrent invocations converge
// while preserving replacement -> logical -> cache -> barrier order.
func (s *Server) completeDeliveredRPCRewrapResult(
ctx context.Context,
a *rpcRewrapAlias,
encoded *encodedOutboundMessage,
source string,
) error {
if s == nil || a == nil || encoded == nil {
return ErrRPCResultFlightInvalid
}
a.deliveredFinalizeOnce.Do(func() {
defer a.releaseBodyReservation()
defer a.releaseReplayRestoreBarrier()
restoreCtx, cancel := boundedRPCReplayRestoreContext(ctx)
defer cancel()
logical, claimErr := claimRPCRewrapLogicalHook(restoreCtx, encoded)
encoded.markDelivered()
if claimErr != nil {
a.conn.fenceUndeliveredRPCResult()
a.deliveredFinalizeErr = fmt.Errorf("wait for rewrapped rpc_result logical restore: %w", claimErr)
a.storeResultOnce(s, encoded)
return
}
a.deliveredFinalizeErr = s.runBoundedRPCReplayRestore(
restoreCtx, a.conn, source, logical, a.runAfterSuccessfulDelivery,
)
a.storeResultOnce(s, encoded)
})
return a.deliveredFinalizeErr
}
var (
rpcRewrapDeliveryOnce sync.Once
rpcRewrapDeliveryJobs chan func()
rpcRewrapDeliveryOnce sync.Once
rpcRewrapDeliveryJobs chan rpcRewrapDeliveryJob
rpcRewrapObservationOnce sync.Once
rpcRewrapObservationJobs chan rpcRewrapDeliveryJob
)
const (
rpcRewrapDeliveryWorkers = 4
rpcRewrapDeliveryQueue = 256
rpcRewrapObserverWorkers = 1
rpcRewrapObserverQueue = 64
// Queue residence, physical delivery and ordered restore share one absolute
// deadline. An admitted alias must never retain a Conn scheduler barrier for
// minutes behind older slow jobs.
rpcRewrapDeliveryQueueTimeout = 5 * time.Second
)
func scheduleRPCRewrapDelivery(fn func()) bool {
if fn == nil {
type rpcRewrapDeliveryJob struct {
run func(*rpcRewrapDeliveryControl, time.Time)
fail func(error)
deadline time.Time
control *rpcRewrapDeliveryControl
}
type rpcRewrapDeliveryJobState uint32
const (
rpcRewrapJobPending rpcRewrapDeliveryJobState = iota
rpcRewrapJobRunning
rpcRewrapJobCommitted
rpcRewrapJobComplete
rpcRewrapJobFailed
)
// rpcRewrapDeliveryControl lets an independent deadline timer retire queued,
// running and physically committed jobs. A late worker cannot enter run after
// the timer wins; a committed non-cooperative restore is fenced by fail, and
// its eventual return cannot report failure or finish the barrier a second time.
type rpcRewrapDeliveryControl struct {
state atomic.Uint32
timerMu sync.Mutex
timer *time.Timer
}
type rpcRewrapPhysicalOutcome struct {
err error
owned bool
}
// waitRPCRewrapPhysicalTerminal deliberately keeps one of the four bounded
// workers attached to an in-progress actor write even after the watchdog fences
// the Conn. A broken transport may therefore strand at most four workers, while
// queued jobs still time out independently. If that transport later reports
// success, the worker cannot lose the logical hook merely because timeout won
// before its goroutine resumed.
func waitRPCRewrapPhysicalTerminal(
c *Conn,
ctx context.Context,
encoded *encodedOutboundMessage,
reserved *outboundBodyReservation,
control *rpcRewrapDeliveryControl,
) rpcRewrapPhysicalOutcome {
terminal := make(chan rpcRewrapPhysicalOutcome, 1)
_ = c.sendOutboundWithTerminalReserved(
ctx, proto.MessageServerResponse, nil, encoded, false,
func(err error) {
terminal <- rpcRewrapPhysicalOutcome{err: err, owned: control.commit()}
},
reserved,
)
return <-terminal
}
func newRPCRewrapDeliveryControl() *rpcRewrapDeliveryControl {
c := &rpcRewrapDeliveryControl{}
c.state.Store(uint32(rpcRewrapJobPending))
return c
}
func (c *rpcRewrapDeliveryControl) transition(from, to rpcRewrapDeliveryJobState) bool {
return c != nil && c.state.CompareAndSwap(uint32(from), uint32(to))
}
func (c *rpcRewrapDeliveryControl) fail() bool {
if c == nil {
return true
}
for {
state := rpcRewrapDeliveryJobState(c.state.Load())
if state == rpcRewrapJobComplete || state == rpcRewrapJobFailed {
return false
}
if c.transition(state, rpcRewrapJobFailed) {
c.stopTimer()
return true
}
}
}
func (c *rpcRewrapDeliveryControl) timeout() bool {
if c == nil {
return true
}
for {
state := rpcRewrapDeliveryJobState(c.state.Load())
if state != rpcRewrapJobPending && state != rpcRewrapJobRunning &&
state != rpcRewrapJobCommitted {
return false
}
if c.transition(state, rpcRewrapJobFailed) {
c.stopTimer()
return true
}
}
}
// commit records successful physical delivery (or an already-proven retarget)
// without disarming the watchdog. The same absolute deadline covers the
// replacement/logical restore and cache/barrier terminal path; complete is the
// only successful transition that stops the timer.
func (c *rpcRewrapDeliveryControl) commit() bool {
if c == nil || !c.transition(rpcRewrapJobRunning, rpcRewrapJobCommitted) {
return false
}
rpcRewrapDeliveryOnce.Do(func() {
rpcRewrapDeliveryJobs = make(chan func(), rpcRewrapDeliveryQueue)
for range rpcRewrapDeliveryWorkers {
return true
}
func (c *rpcRewrapDeliveryControl) running() bool {
return c != nil && rpcRewrapDeliveryJobState(c.state.Load()) == rpcRewrapJobRunning
}
func (c *rpcRewrapDeliveryControl) complete() {
if c == nil {
return
}
for {
state := rpcRewrapDeliveryJobState(c.state.Load())
if state != rpcRewrapJobRunning && state != rpcRewrapJobCommitted {
return
}
if c.transition(state, rpcRewrapJobComplete) {
c.stopTimer()
return
}
}
}
func (c *rpcRewrapDeliveryControl) installTimer(timer *time.Timer) {
if c == nil || timer == nil {
return
}
c.timerMu.Lock()
c.timer = timer
state := rpcRewrapDeliveryJobState(c.state.Load())
terminal := state == rpcRewrapJobComplete || state == rpcRewrapJobFailed
c.timerMu.Unlock()
if terminal {
timer.Stop()
}
}
func (c *rpcRewrapDeliveryControl) stopTimer() {
if c == nil {
return
}
c.timerMu.Lock()
timer := c.timer
c.timer = nil
c.timerMu.Unlock()
if timer != nil {
timer.Stop()
}
}
func (j rpcRewrapDeliveryJob) reportFailure(err error) {
if err == nil {
err = fmt.Errorf("rpc rewrap delivery job failed")
}
defer func() {
if recovered := recover(); recovered != nil {
log.Printf("mtprotoedge: rpc rewrap delivery failure callback panicked: %v\n%s", recovered, debug.Stack())
}
}()
if j.fail != nil {
j.fail(err)
return
}
log.Printf("mtprotoedge: rpc rewrap delivery job failed: %v", err)
}
func runRPCRewrapDeliveryJob(j rpcRewrapDeliveryJob) {
if j.run == nil {
j.reportFailure(fmt.Errorf("nil rpc rewrap delivery job"))
return
}
control := j.control
if control == nil {
control = newRPCRewrapDeliveryControl()
}
if !control.transition(rpcRewrapJobPending, rpcRewrapJobRunning) {
return
}
if !j.deadline.IsZero() && !time.Now().Before(j.deadline) {
if control.fail() {
j.reportFailure(context.DeadlineExceeded)
}
return
}
defer func() {
if recovered := recover(); recovered != nil {
if control.fail() {
j.reportFailure(fmt.Errorf("rpc rewrap delivery panic: %v", recovered))
}
log.Printf("mtprotoedge: rpc rewrap delivery job panicked: %v\n%s", recovered, debug.Stack())
return
}
control.complete()
}()
j.run(control, j.deadline)
}
func scheduleRPCRewrapJob(
job rpcRewrapDeliveryJob,
once *sync.Once,
jobs *chan rpcRewrapDeliveryJob,
workers, queue int,
) bool {
if job.run == nil {
return false
}
if job.deadline.IsZero() {
job.deadline = time.Now().Add(rpcRewrapDeliveryQueueTimeout)
}
job.control = newRPCRewrapDeliveryControl()
once.Do(func() {
*jobs = make(chan rpcRewrapDeliveryJob, queue)
for range workers {
go func() {
for job := range rpcRewrapDeliveryJobs {
job()
for job := range *jobs {
runRPCRewrapDeliveryJob(job)
}
}()
}
})
delay := time.Until(job.deadline)
if delay < 0 {
delay = 0
}
timer := time.AfterFunc(delay, func() {
if job.control.timeout() {
job.reportFailure(context.DeadlineExceeded)
}
})
job.control.installTimer(timer)
select {
case rpcRewrapDeliveryJobs <- fn:
case *jobs <- job:
return true
default:
// If the independent timer already won, it owns the fail callback and the
// caller must not report queue failure a second time.
if job.control.transition(rpcRewrapJobPending, rpcRewrapJobComplete) {
job.control.stopTimer()
return false
}
return true
}
}
func scheduleRPCRewrapDeliveryJob(job rpcRewrapDeliveryJob) bool {
return scheduleRPCRewrapJob(job, &rpcRewrapDeliveryOnce, &rpcRewrapDeliveryJobs,
rpcRewrapDeliveryWorkers, rpcRewrapDeliveryQueue)
}
func scheduleRPCRewrapObservation(fn func()) bool {
if fn == nil {
return false
}
return scheduleRPCRewrapJob(rpcRewrapDeliveryJob{
deadline: time.Now().Add(rpcRewrapDeliveryQueueTimeout),
run: func(*rpcRewrapDeliveryControl, time.Time) { fn() },
}, &rpcRewrapObservationOnce, &rpcRewrapObservationJobs,
rpcRewrapObserverWorkers, rpcRewrapObserverQueue)
}
func (s *Server) rpcRewrapRestoreJob(
a *rpcRewrapAlias,
source string,
run func(*rpcRewrapDeliveryControl, time.Time),
) rpcRewrapDeliveryJob {
return rpcRewrapDeliveryJob{
deadline: time.Now().Add(rpcRewrapDeliveryQueueTimeout),
run: run,
fail: func(err error) {
if a != nil && a.conn != nil {
a.conn.fenceUndeliveredRPCResult()
}
if a != nil {
a.releaseBodyReservation()
a.releaseReplayRestoreBarrier()
}
if s != nil && s.log != nil {
s.log.Warn("RPC rewrap delivery job failed",
zap.String("source", source), zap.Error(err))
}
},
}
}
func (s *Server) failRPCRewrapResultJob(
a *rpcRewrapAlias,
encoded *encodedOutboundMessage,
err error,
) {
if a != nil {
defer a.releaseBodyReservation()
}
if a == nil || a.conn == nil || a.newOwner == nil || encoded == nil {
return
}
a.conn.fenceUndeliveredRPCResult()
publish := a.newOwner.HandOff()
encoded.markReplayable()
encoded.releaseDeferredLogicalDeliveryHook()
// Release the connection-local scheduler before any defensive cache panic;
// the physical generation is already fenced, so no following task can run.
a.releaseReplayRestoreBarrier()
if publish {
a.storeResultOnce(s, encoded)
}
if s != nil && s.log != nil {
s.log.Warn("RPC rewrap result job failed; exact result retained",
zap.String("method", a.method), zap.Int64("req_msg_id", a.newReqID), zap.Error(err))
}
}
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) {
// Install the scheduler barrier synchronously, before this plan publishes
// any following naked RPC tasks. The asynchronous physical replay below is
// then free to use a bounded rewrap worker without an ordering race.
a.beginReplayRestore()
var executionSubscriber func(bool)
if a.newOwner != nil || a.afterSuccessfulDelivery != nil {
executionSubscriber = func(success bool) {
a.executionOK.Store(success)
if a.newOwner != nil {
a.newOwner.CompleteExecution(success)
}
}
}
resultSubscriber := func(encoded *encodedOutboundMessage, ok bool) {
if !ok || encoded == nil {
if a.newOwner != nil {
a.newOwner.Abort()
}
a.conn.fenceUndeliveredRPCResult()
a.releaseDeferredLogicalHook()
a.finishReplayRestoreWithoutDelivery()
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() {
attempt, reserved, cloneErr := a.conn.cloneRPCResultForRequestReserved(encoded, encoded.reqMsgID, false)
if cloneErr != nil {
a.conn.failOutboundBudget(cloneErr)
a.conn.fenceUndeliveredRPCResult()
a.finishReplayRestoreWithoutDelivery()
return
}
clone.markDelivered()
s.storeRPCResult(a.conn, a.newReqID, clone)
a.bodyReservation = reserved
job := s.rpcRewrapRestoreJob(a, "pending init rewrap replay", func(control *rpcRewrapDeliveryControl, deadline time.Time) {
ctx, cancel := context.WithDeadline(context.Background(), deadline)
defer cancel()
outcome := waitRPCRewrapPhysicalTerminal(a.conn, ctx, attempt, a.bodyReservation, control)
if outcome.err != nil {
if !outcome.owned {
return
}
a.conn.fenceUndeliveredRPCResult()
attempt.markReplayable()
a.releaseBodyReservation()
a.releaseReplayRestoreBarrier()
if !isClientDisconnect(outcome.err) {
s.log.Debug("RPC init rewrap pending replay failed", zap.Error(outcome.err))
}
return
}
// A watchdog may win after the transport has already returned physical
// success but before this goroutine resumes. Success is irrevocable: run
// the once-only restore with a fresh bounded lifetime if timeout failure
// already fenced/released this physical generation.
restoreParent := ctx
if !outcome.owned {
restoreParent = context.Background()
}
restoreCtx, cancelRestore := boundedRPCReplayRestoreContext(restoreParent)
defer cancelRestore()
logical, claimErr := attempt.claimLogicalDeliveryHook(restoreCtx, false)
attempt.markDelivered()
if claimErr != nil {
a.conn.fenceUndeliveredRPCResult()
a.releaseBodyReservation()
a.releaseReplayRestoreBarrier()
return
}
restoreErr := s.runBoundedRPCReplayRestore(
restoreCtx, a.conn, "pending init rewrap replay", logical, a.runAfterSuccessfulDelivery,
)
a.releaseBodyReservation()
a.releaseReplayRestoreBarrier()
if restoreErr != nil && !isClientDisconnect(restoreErr) {
s.log.Debug("RPC init rewrap pending replay failed", zap.Error(restoreErr))
}
})
if !scheduleRPCRewrapDeliveryJob(job) {
a.conn.fenceUndeliveredRPCResult()
a.finishReplayRestoreWithoutDelivery()
}
return
}
// A successful physical write under the retargeted req_msg_id is the only
// proof that lets the alias reuse that attempt. A mere TryRetarget success
// is not proof: the original socket may have failed before any bytes landed.
retargetDelivered := a.retargeted.Load() &&
encoded.deliveryState() == rpcResultDeliveryDelivered &&
encoded.writtenRequestID() == a.newReqID
clone, reserved, err := a.conn.cloneRPCResultForRequestReserved(encoded, a.newReqID, retargetDelivered)
if err != nil {
a.conn.failOutboundBudget(err)
a.newOwner.Abort()
a.conn.fenceUndeliveredRPCResult()
a.finishReplayRestoreWithoutDelivery()
return
}
a.bodyReservation = reserved
if retargetDelivered {
if !a.newOwner.HandOff() {
a.conn.fenceUndeliveredRPCResult()
a.finishReplayRestoreWithoutDelivery()
return
}
job := s.rpcRewrapRestoreJob(a, "retargeted init rewrap result", func(control *rpcRewrapDeliveryControl, deadline time.Time) {
if !control.commit() {
return
}
restoreCtx, cancelRestore := context.WithDeadline(context.Background(), deadline)
defer cancelRestore()
restoreErr := s.completeDeliveredRPCRewrapResult(
restoreCtx, a, clone, "retargeted init rewrap result",
)
if restoreErr != nil && !isClientDisconnect(restoreErr) {
s.log.Debug("Retargeted RPC restore failed", zap.Error(restoreErr))
}
})
job.fail = func(err error) {
// Never enter deliveredFinalizeOnce from the timer goroutine: the worker
// may already own a non-cooperative restore. Fence and release its Conn
// barrier first, then retain the immutable delivered result for a later
// replacement replay, which will wait on coordinator Claimed/InProgress.
a.conn.fenceUndeliveredRPCResult()
clone.markReplayable()
clone.releaseDeferredLogicalDeliveryHook()
a.releaseReplayRestoreBarrier()
a.storeResultOnce(s, clone)
a.releaseBodyReservation()
s.log.Warn("Retargeted RPC restore watchdog expired",
zap.String("method", a.method), zap.Int64("req_msg_id", a.newReqID), zap.Error(err))
}
if !scheduleRPCRewrapDeliveryJob(job) {
job.fail(ErrOutboundQueueFull)
}
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)
}) {
job := s.rpcRewrapRestoreJob(a, "pending init rewrap result", func(control *rpcRewrapDeliveryControl, deadline time.Time) {
s.publishRewrappedRPCResult(a.conn, a.newReqID, a.method, a.newOwner, clone, a, control, deadline)
})
// Once this alias consumed the source candidate, expiration or panic of
// the admitted worker job must still publish the immutable result under
// the new msg_id. Otherwise the alias owner would remain pending forever
// (or a reconnect could execute the business request a second time).
job.fail = func(err error) { s.failRPCRewrapResultJob(a, clone, err) }
if !scheduleRPCRewrapDeliveryJob(job) {
// 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)
}
s.failRPCRewrapResultJob(a, clone, ErrOutboundQueueFull)
}
})
}
var err error
if executionSubscriber != nil {
err = a.oldWaiter.SubscribeResultAndExecution(resultSubscriber, executionSubscriber)
} else {
err = a.oldWaiter.Subscribe(resultSubscriber)
}
if err != nil {
a.finishReplayRestoreWithoutDelivery()
s.rpcRewrap.release(a.candidate)
return err
}
@ -407,7 +1026,7 @@ func (s *Server) scheduleRewrappedInitObservation(c *Conn, init rpcRewrapInit) {
if !ok || c == nil {
return
}
if !scheduleRPCRewrapDelivery(func() {
if !scheduleRPCRewrapObservation(func() {
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
if err := observer.ObserveInitConnection(
@ -422,33 +1041,75 @@ func (s *Server) scheduleRewrappedInitObservation(c *Conn, init rpcRewrapInit) {
}
}
func (s *Server) publishRewrappedRPCResult(c *Conn, reqMsgID int64, method string, owner *rpcResultOwnerLease, encoded *encodedOutboundMessage) {
func (s *Server) publishRewrappedRPCResult(
c *Conn,
reqMsgID int64,
method string,
owner *rpcResultOwnerLease,
encoded *encodedOutboundMessage,
alias *rpcRewrapAlias,
control *rpcRewrapDeliveryControl,
deadline time.Time,
) {
if s == nil || c == nil || owner == nil || encoded == nil {
if alias != nil {
alias.finishReplayRestoreWithoutDelivery()
}
return
}
if !owner.HandOff() {
c.fenceUndeliveredRPCResult()
alias.finishReplayRestoreWithoutDelivery()
return
}
if control == nil || !control.running() {
alias.releaseBodyReservation()
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)))
if deadline.IsZero() {
deadline = time.Now().Add(rpcRewrapDeliveryQueueTimeout)
}
ctx, cancel := context.WithDeadline(context.Background(), deadline)
defer cancel()
if err := c.enqueueEncodedDelivery(ctx, proto.MessageServerResponse, encoded, priority, terminal); err != nil {
terminal(err)
// Rewrap delivery is synchronous on this small bounded worker pool. This
// makes the queue deadline cover the physical write and lets the pending
// logical hook join the same per-Conn ordered restore, without touching the
// process-wide asynchronous hook executor.
outcome := waitRPCRewrapPhysicalTerminal(c, ctx, encoded, alias.bodyReservation, control)
if outcome.err != nil {
if !outcome.owned {
return
}
encoded.markReplayable()
encoded.releaseDeferredLogicalDeliveryHook()
c.fenceUndeliveredRPCResult()
alias.storeResultOnce(s, encoded)
alias.releaseBodyReservation()
alias.releaseReplayRestoreBarrier()
return
}
s.log.Info("RPC init rewrap result replay admitted",
// Physical success outranks an already-fired watchdog. The timeout path may
// have fenced and cached a replayable clone, but it cannot revoke bytes; the
// shared once/coordinator below still completes logical state exactly once.
// Run replacement metadata then the original logical hook before publishing
// the alias cache entry. Whole-finalization once also covers a watchdog racing
// a late physical terminal, so completed metadata cannot be overwritten.
restoreParent := ctx
if !outcome.owned {
restoreParent = context.Background()
}
restoreCtx, cancelRestore := boundedRPCReplayRestoreContext(restoreParent)
defer cancelRestore()
restoreErr := s.completeDeliveredRPCRewrapResult(
restoreCtx, alias, encoded, "physically delivered init rewrap result",
)
if restoreErr != nil && !isClientDisconnect(restoreErr) {
s.log.Debug("RPC init rewrap delivered-state restore failed", zap.Error(restoreErr))
}
s.log.Info("RPC init rewrap result replay delivered",
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)))