notify/transport_binding.go

package notify

import (
	"b612.me/stario"
	"context"
	"net"
	"sync"
	"sync/atomic"
	"time"
)

const transportWorkerStopWait = time.Second

// transportBinding models the currently attached physical transport for a
// logical session. The binding can be swapped later without forcing callers to
// reach into raw conn fields directly.
type transportBinding struct {
	conn  net.Conn
	queue *stario.StarQueue

	writeGateOnce sync.Once
	writeGate     chan struct{}
	writeGateRef  *connWriteGateRef
	writeGateDone sync.Once

	writeActive        atomic.Int64
	writeStopping      atomic.Bool
	writeDrainOnce     sync.Once
	writeDrain         chan struct{}
	writeDrainDoneOnce sync.Once

	adaptiveTx adaptiveTxState

	controlMu     sync.Mutex
	controlSender *controlBatchSender

	streamMu     sync.Mutex
	streamSender *streamBatchSender

	bulkMu     sync.Mutex
	bulkSender *bulkBatchSender
}

func newTransportBinding(conn net.Conn, queue *stario.StarQueue) *transportBinding {
	if conn == nil && queue == nil {
		return nil
	}
	return &transportBinding{conn: conn, queue: queue}
}

func (b *transportBinding) connSnapshot() net.Conn {
	if b == nil {
		return nil
	}
	return b.conn
}

func (b *transportBinding) queueSnapshot() *stario.StarQueue {
	if b == nil {
		return nil
	}
	return b.queue
}

func (b *transportBinding) withConnWriteLock(fn func(net.Conn) error) error {
	return b.withConnWriteLockDeadline(time.Time{}, fn)
}

func (b *transportBinding) withConnWriteLockDeadline(deadline time.Time, fn func(net.Conn) error) error {
	if b == nil {
		return net.ErrClosed
	}
	if err := b.beginConnWrite(); err != nil {
		return err
	}
	<-b.writeGateSnapshot()
	defer b.unlockConnWrite()
	conn := b.connSnapshot()
	if conn == nil {
		return net.ErrClosed
	}
	if !deadline.IsZero() {
		if err := conn.SetWriteDeadline(deadline); err != nil {
			return err
		}
		defer func() {
			_ = conn.SetWriteDeadline(time.Time{})
		}()
	}
	return fn(conn)
}

func (b *transportBinding) withConnWriteLockContextTimeout(ctx context.Context, timeout time.Duration, fn func(net.Conn) error) (bool, error) {
	return b.withConnWriteLockContextStopTimeout(ctx, nil, timeout, fn)
}

func (b *transportBinding) withConnWriteLockContextStopTimeout(ctx context.Context, stop <-chan struct{}, timeout time.Duration, fn func(net.Conn) error) (bool, error) {
	return b.withConnWriteLockContextStopDeadline(ctx, stop, writeDeadlineFromTimeout(timeout), fn)
}

// withConnWriteLockContextStopDeadline carries the caller's context deadline
// into the physical socket write. A context only used for queue admission is
// otherwise unable to interrupt a net.Conn.Write once the write has started.
func (b *transportBinding) withConnWriteLockContextStopDeadline(ctx context.Context, stop <-chan struct{}, deadline time.Time, fn func(net.Conn) error) (bool, error) {
	return b.withConnWriteLockContextStopDeadlineMode(ctx, stop, deadline, true, fn)
}

// Sender-owned writes are drained by the sender's stop/flush lifecycle. They
// only need the shutdown admission check, avoiding activity-counter work on
// the bulk, stream, and control hot paths.
func (b *transportBinding) withConnWriteLockContextStopDeadlineManaged(ctx context.Context, stop <-chan struct{}, deadline time.Time, fn func(net.Conn) error) (bool, error) {
	return b.withConnWriteLockContextStopDeadlineMode(ctx, stop, deadline, false, fn)
}

func (b *transportBinding) withConnWriteLockContextStopDeadlineMode(ctx context.Context, stop <-chan struct{}, deadline time.Time, trackActivity bool, fn func(net.Conn) error) (bool, error) {
	if b == nil {
		return false, net.ErrClosed
	}
	if ctx == nil {
		ctx = context.Background()
	}
	if trackActivity {
		if err := b.beginConnWrite(); err != nil {
			return false, err
		}
	} else if b.writeStopping.Load() {
		return false, net.ErrClosed
	}
	deadline = earlierWriteDeadline(deadline, contextDeadline(ctx))
	if err := lockWriteGateContextDeadline(ctx, stop, b.writeGateSnapshot(), deadline); err != nil {
		if trackActivity {
			b.finishConnWrite()
		}
		return false, err
	}
	if trackActivity {
		defer b.unlockConnWrite()
	} else {
		defer b.unlockConnWriteManaged()
	}
	if err := ctx.Err(); err != nil {
		return false, err
	}
	conn := b.connSnapshot()
	if conn == nil {
		return false, net.ErrClosed
	}
	if !deadline.IsZero() {
		if err := conn.SetWriteDeadline(deadline); err != nil {
			return true, err
		}
		defer func() {
			_ = conn.SetWriteDeadline(time.Time{})
		}()
	}
	return true, fn(conn)
}

func contextDeadline(ctx context.Context) time.Time {
	if ctx == nil {
		return time.Time{}
	}
	deadline, ok := ctx.Deadline()
	if !ok {
		return time.Time{}
	}
	return deadline
}

func earlierWriteDeadline(left time.Time, right time.Time) time.Time {
	if left.IsZero() {
		return right
	}
	if right.IsZero() || left.Before(right) {
		return left
	}
	return right
}

func (b *transportBinding) lockConnWriteContext(ctx context.Context) error {
	return b.lockConnWriteContextStop(ctx, nil)
}

func (b *transportBinding) lockConnWriteContextStop(ctx context.Context, stop <-chan struct{}) error {
	if b == nil {
		return net.ErrClosed
	}
	if ctx == nil {
		ctx = context.Background()
	}
	if err := b.beginConnWrite(); err != nil {
		return err
	}
	if err := lockWriteGateContextDeadline(ctx, stop, b.writeGateSnapshot(), time.Time{}); err != nil {
		b.finishConnWrite()
		return err
	}
	return nil
}

func (b *transportBinding) unlockConnWrite() {
	if b == nil {
		return
	}
	b.writeGateSnapshot() <- struct{}{}
	b.finishConnWrite()
}

func (b *transportBinding) unlockConnWriteManaged() {
	if b == nil {
		return
	}
	b.writeGateSnapshot() <- struct{}{}
}

func (b *transportBinding) beginConnWrite() error {
	if b == nil || b.writeStopping.Load() {
		return net.ErrClosed
	}
	b.writeActive.Add(1)
	if b.writeStopping.Load() {
		b.finishConnWrite()
		return net.ErrClosed
	}
	return nil
}

func (b *transportBinding) finishConnWrite() {
	if b == nil {
		return
	}
	if b.writeActive.Add(-1) == 0 && b.writeStopping.Load() {
		b.signalConnWritesDrained()
	}
}

func (b *transportBinding) stopConnWrites() <-chan struct{} {
	if b == nil {
		done := make(chan struct{})
		close(done)
		return done
	}
	b.beginConnWriteShutdown()
	done := b.writeDrainSnapshot()
	return done
}

func (b *transportBinding) beginConnWriteShutdown() {
	if b == nil {
		return
	}
	b.writeStopping.Store(true)
	if b.writeActive.Load() == 0 {
		b.signalConnWritesDrained()
	}
}

func (b *transportBinding) writeDrainSnapshot() chan struct{} {
	b.writeDrainOnce.Do(func() {
		b.writeDrain = make(chan struct{})
	})
	return b.writeDrain
}

func (b *transportBinding) signalConnWritesDrained() {
	done := b.writeDrainSnapshot()
	b.writeDrainDoneOnce.Do(func() {
		close(done)
	})
}

func (b *transportBinding) writeGateSnapshot() chan struct{} {
	if b == nil {
		return nil
	}
	b.writeGateOnce.Do(func() {
		if conn := b.connSnapshot(); conn != nil {
			b.writeGateRef = retainRawConnWriteGate(conn)
			b.writeGate = b.writeGateRef.gate
		} else {
			b.writeGate = newConnWriteGate()
		}
	})
	return b.writeGate
}

func (b *transportBinding) bulkBatchSenderSnapshotWithCodec(codec bulkBatchCodec, writeTimeout func() time.Duration) *bulkBatchSender {
	if b == nil {
		return nil
	}
	b.bulkMu.Lock()
	defer b.bulkMu.Unlock()
	if b.bulkSender != nil {
		return b.bulkSender
	}
	b.bulkSender = newBulkBatchSender(b, codec, writeTimeout)
	return b.bulkSender
}

func (b *transportBinding) clientBulkBatchSenderSnapshot(c *ClientCommon) *bulkBatchSender {
	if b == nil || c == nil {
		return nil
	}
	return b.bulkBatchSenderSnapshotWithCodec(bulkBatchCodec{
		encodeSingle: c.encodeBulkFastPayloadPooled,
		encodeBatch:  c.encodeBulkFastBatchPayloadPooled,
	}, c.maxWriteTimeoutSnapshot)
}

func (b *transportBinding) serverBulkBatchSenderSnapshot(logical *LogicalConn) *bulkBatchSender {
	if b == nil || logical == nil {
		return nil
	}
	server := logical.Server()
	common, ok := server.(*ServerCommon)
	if !ok || common == nil {
		return nil
	}
	return b.bulkBatchSenderSnapshotWithCodec(bulkBatchCodec{
		encodeSingle: func(frame bulkFastFrame) ([]byte, func(), error) {
			return common.encodeBulkFastPayloadLogicalPooled(logical, frame)
		},
		encodeBatch: func(frames []bulkFastFrame) ([]byte, func(), error) {
			return common.encodeBulkFastBatchPayloadLogicalPooled(logical, frames)
		},
	}, logical.maxWriteTimeoutSnapshot)
}

func (b *transportBinding) controlBatchSenderSnapshot() *controlBatchSender {
	if b == nil {
		return nil
	}
	b.controlMu.Lock()
	defer b.controlMu.Unlock()
	if b.controlSender != nil {
		return b.controlSender
	}
	b.controlSender = newControlBatchSender(b)
	return b.controlSender
}

func (b *transportBinding) streamBatchSenderSnapshotWithCodec(codec streamBatchCodec, writeTimeout func() time.Duration) *streamBatchSender {
	if b == nil {
		return nil
	}
	b.streamMu.Lock()
	defer b.streamMu.Unlock()
	if b.streamSender != nil {
		return b.streamSender
	}
	b.streamSender = newStreamBatchSender(b, codec, writeTimeout)
	return b.streamSender
}

func (b *transportBinding) clientStreamBatchSenderSnapshot(c *ClientCommon) *streamBatchSender {
	if b == nil || c == nil {
		return nil
	}
	return b.streamBatchSenderSnapshotWithCodec(streamBatchCodec{
		encodeSingle: c.encodeFastStreamPayload,
		encodeBatch:  c.encodeFastStreamBatchPayload,
	}, c.maxWriteTimeoutSnapshot)
}

func (b *transportBinding) serverStreamBatchSenderSnapshot(logical *LogicalConn) *streamBatchSender {
	if b == nil || logical == nil {
		return nil
	}
	server := logical.Server()
	common, ok := server.(*ServerCommon)
	if !ok || common == nil {
		return nil
	}
	return b.streamBatchSenderSnapshotWithCodec(streamBatchCodec{
		encodeSingle: func(frame streamFastDataFrame) ([]byte, error) {
			return common.encodeFastStreamPayloadLogical(logical, frame)
		},
		encodeBatch: func(frames []streamFastDataFrame) ([]byte, error) {
			return common.encodeFastStreamBatchPayloadLogical(logical, frames)
		},
	}, logical.maxWriteTimeoutSnapshot)
}

func (b *transportBinding) stopBackgroundWorkers() {
	b.stopBackgroundWorkersWithClose(false)
}

// stopReplacedTransportBinding interrupts an in-flight physical write before
// waiting for the old binding's workers. A binding may be replaced while
// retaining the same socket (for example, when only its queue changes), so
// that case keeps the connection open.
func stopReplacedTransportBinding(oldBinding *transportBinding, nextBinding *transportBinding, closeConn bool) {
	if oldBinding == nil {
		return
	}
	if closeConn && (nextBinding == nil || oldBinding.connSnapshot() != nextBinding.connSnapshot()) {
		oldBinding.stopBackgroundWorkersWithClose(true)
		return
	}
	oldBinding.stopBackgroundWorkers()
}

func (b *transportBinding) stopBackgroundWorkersWithClose(closeConn bool) {
	if b == nil {
		return
	}
	b.beginConnWriteShutdown()
	b.controlMu.Lock()
	controlSender := b.controlSender
	b.controlMu.Unlock()
	b.streamMu.Lock()
	streamSender := b.streamSender
	b.streamMu.Unlock()
	b.bulkMu.Lock()
	bulkSender := b.bulkSender
	b.bulkMu.Unlock()
	// Closing first is required to interrupt an in-flight physical write during
	// actual transport shutdown. Handoff callers normally keep the old socket
	// alive, but a bounded wait below closes it if a sender is already inside
	// Conn.Write. Reusing a socket after that point could otherwise preserve a
	// partial frame and block the handoff forever.
	if closeConn {
		b.closeConn()
	}
	workersDone := make(chan struct{})
	go func() {
		if controlSender != nil {
			controlSender.stop()
		}
		if streamSender != nil {
			streamSender.stop()
		}
		if bulkSender != nil {
			bulkSender.stop()
		}
		b.releaseWriteGate()
		close(workersDone)
	}()
	timer := time.NewTimer(transportWorkerStopWait)
	defer timer.Stop()
	select {
	case <-workersDone:
		return
	case <-timer.C:
		if !closeConn {
			// A same-socket handoff may retain the connection only while no old
			// sender is active. Once the bounded wait expires the socket is
			// unsafe to reuse, so force it closed and let the stopper finish
			// asynchronously.
			b.closeConn()
		}
	}
}

func (b *transportBinding) releaseWriteGate() {
	if b == nil {
		return
	}
	b.writeGateSnapshot()
	<-b.stopConnWrites()
	b.writeGateDone.Do(func() {
		releaseRawConnWriteGate(b.connSnapshot(), b.writeGateRef)
	})
}

func (b *transportBinding) closeConn() {
	if b == nil {
		return
	}
	if conn := b.connSnapshot(); conn != nil {
		_ = conn.Close()
	}
}