owpengram-server/internal/mtprotoedge/transport_ownership.go

package mtprotoedge

import (
	"context"
	"sync"
	"sync/atomic"
	"time"

	"github.com/iamxvbaba/td/bin"
	"github.com/iamxvbaba/td/transport"
)

const (
	physicalTransportClosedBit     = uint64(1) << 63
	physicalTransportGenerationMax = physicalTransportClosedBit - 1
)

// physicalTransportOwner separates the lifetime of one physical socket from
// the logical Conn generations that successively use it. state is either an
// open generation in [1, physicalTransportGenerationMax], or that generation
// with physicalTransportClosedBit set.
//
// Transfer and owner-close race through one atomic state transition. This is
// the critical property: when transfer wins, a stale generation can no longer
// close the socket; when close wins, no later generation can be published.
type physicalTransportOwner struct {
	raw transport.Conn

	// writeMu orders a completed send before generation transfer. CloseAny and
	// owner-close deliberately do not wait for it: raw.Close must be able to
	// interrupt a transport implementation blocked inside Send.
	writeMu sync.Mutex
	state   atomic.Uint64

	// binding identifies the logical Conn currently attached to the open
	// generation. Physical close fences that exact Conn before touching raw;
	// generation matching prevents a stale close from retiring a replacement.
	bindingMu        sync.Mutex
	boundGeneration  uint64
	boundLogicalConn *Conn

	closeDone chan struct{}
	closeErr  error
}

// physicalTransportLease is the capability owned by one logical Conn
// generation. It implements transport.Conn so exchange/protocol-error paths
// can use the initial lease from the moment serveConn starts.
type physicalTransportLease struct {
	owner      *physicalTransportOwner
	generation uint64
}

var _ transport.Conn = (*physicalTransportLease)(nil)

func newPhysicalTransportOwner(raw transport.Conn) (*physicalTransportOwner, *physicalTransportLease) {
	owner := &physicalTransportOwner{
		raw:       raw,
		closeDone: make(chan struct{}),
	}
	owner.state.Store(1)
	return owner, &physicalTransportLease{owner: owner, generation: 1}
}

// Transfer atomically hands the physical transport to the next logical Conn
// generation. Holding writeMu ensures every send admitted by the old lease has
// returned before the generation changes. CloseAny may still interrupt such a
// send and make the CAS fail, which is the required shutdown ordering.
func (l *physicalTransportLease) Transfer() (*physicalTransportLease, bool) {
	if l == nil || l.owner == nil || l.generation == 0 {
		return nil, false
	}
	owner := l.owner
	owner.writeMu.Lock()
	defer owner.writeMu.Unlock()

	if l.generation >= physicalTransportGenerationMax {
		return nil, false
	}
	if !owner.state.CompareAndSwap(l.generation, l.generation+1) {
		return nil, false
	}
	owner.bindingMu.Lock()
	if owner.boundGeneration == l.generation {
		owner.boundGeneration = l.generation + 1
		owner.boundLogicalConn = nil
	}
	owner.bindingMu.Unlock()
	return &physicalTransportLease{owner: owner, generation: l.generation + 1}, true
}

// IsCurrentOpen reports whether this lease is still the unique open owner.
// Callers use it at protocol publication barriers after a required write.
func (l *physicalTransportLease) IsCurrentOpen() bool {
	return l != nil && l.owner != nil && l.generation != 0 &&
		l.owner.state.Load() == l.generation
}

// bindLogicalConn publishes the logical Conn for this exact generation. If
// physical close already won, the new Conn is terminally fenced immediately
// and can never pass SessionManager publication checks.
func (l *physicalTransportLease) bindLogicalConn(c *Conn) bool {
	if l == nil || l.owner == nil || c == nil {
		if c != nil {
			c.beginTerminalShutdown()
		}
		return false
	}
	owner := l.owner
	owner.bindingMu.Lock()
	open := owner.state.Load() == l.generation
	if open {
		owner.boundGeneration = l.generation
		owner.boundLogicalConn = c
	}
	owner.bindingMu.Unlock()
	if !open {
		c.beginTerminalShutdown()
	}
	return open
}

// Send admits a write only while this lease is the current open generation.
// The lock also serializes quick ACK/protocol writes with the outbound actor.
func (l *physicalTransportLease) Send(ctx context.Context, b *bin.Buffer) error {
	return l.withCurrentWriter(func(raw transport.Conn) error {
		return raw.Send(ctx, b)
	})
}

// SendDeadline preserves the allocation-free fast path of compatTransportConn.
func (l *physicalTransportLease) SendDeadline(deadline time.Time, b *bin.Buffer) error {
	return l.withCurrentWriter(func(raw transport.Conn) error {
		if writer, ok := raw.(deadlineOutboundWriter); ok {
			return writer.SendDeadline(deadline, b)
		}
		ctx := context.Background()
		cancel := func() {}
		if !deadline.IsZero() {
			ctx, cancel = context.WithDeadline(ctx, deadline)
		}
		defer cancel()
		return raw.Send(ctx, b)
	})
}

// SendDeadlineWithScratch forwards the globally budgeted codec scratch through the generation
// lease while preserving the same write-ownership barrier as SendDeadline.
func (l *physicalTransportLease) SendDeadlineWithScratch(deadline time.Time, b *bin.Buffer, scratch *[]byte) error {
	return l.withCurrentWriter(func(raw transport.Conn) error {
		if writer, ok := raw.(deadlineOutboundScratchWriter); ok {
			return writer.SendDeadlineWithScratch(deadline, b, scratch)
		}
		if writer, ok := raw.(deadlineOutboundWriter); ok {
			return writer.SendDeadline(deadline, b)
		}
		ctx := context.Background()
		cancel := func() {}
		if !deadline.IsZero() {
			ctx, cancel = context.WithDeadline(ctx, deadline)
		}
		defer cancel()
		return raw.Send(ctx, b)
	})
}

// SendDeadlineWithScratchGuarded evaluates guard while holding the physical
// write-ownership lock, immediately before entering the raw writer. Conn-level
// checks performed before this call are insufficient: a quick ACK or protocol
// write may hold writeMu across a temporary-key expiry boundary. The guard must
// not close/fence the connection itself because that would re-enter transport
// shutdown while writeMu is held; callers handle its error after the lock drops.
func (l *physicalTransportLease) SendDeadlineWithScratchGuarded(deadline time.Time, b *bin.Buffer, scratch *[]byte, guard func() error) error {
	return l.withCurrentWriterGuarded(guard, func(raw transport.Conn) error {
		if writer, ok := raw.(deadlineOutboundScratchWriter); ok {
			return writer.SendDeadlineWithScratch(deadline, b, scratch)
		}
		if writer, ok := raw.(deadlineOutboundWriter); ok {
			return writer.SendDeadline(deadline, b)
		}
		ctx := context.Background()
		cancel := func() {}
		if !deadline.IsZero() {
			ctx, cancel = context.WithDeadline(ctx, deadline)
		}
		defer cancel()
		return raw.Send(ctx, b)
	})
}

func (l *physicalTransportLease) withCurrentWriter(send func(transport.Conn) error) error {
	return l.withCurrentWriterGuarded(nil, send)
}

func (l *physicalTransportLease) withCurrentWriterGuarded(guard func() error, send func(transport.Conn) error) error {
	if l == nil || l.owner == nil || l.owner.raw == nil {
		return ErrConnClosed
	}
	owner := l.owner
	owner.writeMu.Lock()
	defer owner.writeMu.Unlock()
	if owner.state.Load() != l.generation {
		return ErrConnClosed
	}
	if guard != nil {
		if err := guard(); err != nil {
			return err
		}
	}
	return send(owner.raw)
}

// Recv is owned by serveConn rather than a logical Conn generation. It is a
// direct proxy so the read loop remains valid after ownership transfers.
func (l *physicalTransportLease) Recv(ctx context.Context, b *bin.Buffer) error {
	if l == nil || l.owner == nil || l.owner.raw == nil {
		return ErrConnClosed
	}
	return l.owner.raw.Recv(ctx, b)
}

// RecvDeadline preserves serveConn's direct-deadline fast path.
func (l *physicalTransportLease) RecvDeadline(deadline time.Time, b *bin.Buffer) error {
	if l == nil || l.owner == nil || l.owner.raw == nil {
		return ErrConnClosed
	}
	if receiver, ok := l.owner.raw.(deadlineReceiver); ok {
		return receiver.RecvDeadline(deadline, b)
	}
	ctx := context.Background()
	cancel := func() {}
	if !deadline.IsZero() {
		ctx, cancel = context.WithDeadline(ctx, deadline)
	}
	defer cancel()
	return l.owner.raw.Recv(ctx, b)
}

// Close closes the physical transport only if this lease still owns the
// current generation. A stale logical Conn therefore cannot close a socket
// already transferred to its replacement.
func (l *physicalTransportLease) Close() error {
	if l == nil || l.owner == nil || l.generation == 0 {
		return nil
	}
	owner := l.owner
	for {
		state := owner.state.Load()
		if state&physicalTransportClosedBit != 0 {
			return owner.waitClosed()
		}
		if state != l.generation {
			return nil
		}
		if owner.state.CompareAndSwap(state, state|physicalTransportClosedBit) {
			owner.fenceBoundGeneration(l.generation)
			return owner.closeRaw()
		}
	}
}

// startCloseAlreadyFenced is the non-reentrant close capability for the logical
// Conn that has already published terminal/lifecycle gates itself. It
// synchronously publishes the physical closed bit (so Transfer cannot win), then
// runs the potentially pathological raw.Close outside the RPC worker/flight
// handoff. It neither calls fenceBoundGeneration nor waits for a CloseAny that
// already won, avoiding both lifecycle reentry and close cycles.
func (l *physicalTransportLease) startCloseAlreadyFenced() {
	if l == nil || l.owner == nil || l.generation == 0 {
		return
	}
	owner := l.owner
	for {
		state := owner.state.Load()
		if state&physicalTransportClosedBit != 0 || state != l.generation {
			return
		}
		if owner.state.CompareAndSwap(state, state|physicalTransportClosedBit) {
			owner.bindingMu.Lock()
			if owner.boundGeneration == l.generation {
				owner.boundLogicalConn = nil
			}
			owner.bindingMu.Unlock()
			go func() { _ = owner.closeRaw() }()
			return
		}
	}
}

// CloseAny is the unconditional physical-socket capability retained by
// serveConn. It marks the owner closed before calling raw.Close and never waits
// for writeMu, so it can break a blocked Send or Recv.
func (o *physicalTransportOwner) CloseAny() error {
	if o == nil {
		return nil
	}
	for {
		state := o.state.Load()
		if state&physicalTransportClosedBit != 0 {
			return o.waitClosed()
		}
		if o.state.CompareAndSwap(state, state|physicalTransportClosedBit) {
			o.fenceBoundGeneration(state)
			return o.closeRaw()
		}
	}
}

func (o *physicalTransportOwner) fenceBoundGeneration(generation uint64) {
	o.bindingMu.Lock()
	var c *Conn
	if o.boundGeneration == generation {
		c = o.boundLogicalConn
		o.boundLogicalConn = nil
	}
	o.bindingMu.Unlock()
	if c != nil {
		c.beginTerminalShutdown()
	}
}

func (o *physicalTransportOwner) closeRaw() error {
	if o.raw != nil {
		o.closeErr = o.raw.Close()
	}
	close(o.closeDone)
	return o.closeErr
}

func (o *physicalTransportOwner) waitClosed() error {
	<-o.closeDone
	return o.closeErr
}

// Forward the optional compat-transport capabilities hidden by the lease.
// These keep frame-budget ownership and quick-ack semantics unchanged while
// serveConn operates on the initial lease instead of the raw transport.
func (l *physicalTransportLease) releaseInboundFrame() {
	if l == nil || l.owner == nil {
		return
	}
	if releaser, ok := l.owner.raw.(inboundFrameOwnershipReleaser); ok {
		releaser.releaseInboundFrame()
	}
}

func (l *physicalTransportLease) retainInboundFrameBytes(n int64) bool {
	if l == nil || l.owner == nil {
		return true
	}
	if retainer, ok := l.owner.raw.(inboundFrameBackingRetainer); ok {
		return retainer.retainInboundFrameBytes(n)
	}
	return true
}

func (l *physicalTransportLease) ConsumeQuickAckRequested() bool {
	if l == nil || l.owner == nil {
		return false
	}
	if quick, ok := l.owner.raw.(quickAckTransport); ok {
		return quick.ConsumeQuickAckRequested()
	}
	return false
}

func (l *physicalTransportLease) SendQuickAck(ctx context.Context, token uint32) error {
	return l.withCurrentWriter(func(raw transport.Conn) error {
		if quick, ok := raw.(quickAckTransport); ok {
			return quick.SendQuickAck(ctx, token)
		}
		return nil
	})
}

func (l *physicalTransportLease) SendQuickAckDeadline(deadline time.Time, token uint32) error {
	return l.withCurrentWriter(func(raw transport.Conn) error {
		if quick, ok := raw.(deadlineQuickAckTransport); ok {
			return quick.SendQuickAckDeadline(deadline, token)
		}
		if quick, ok := raw.(quickAckTransport); ok {
			ctx := context.Background()
			cancel := func() {}
			if !deadline.IsZero() {
				ctx, cancel = context.WithDeadline(ctx, deadline)
			}
			defer cancel()
			return quick.SendQuickAck(ctx, token)
		}
		return nil
	})
}