notify/transport_write.go

package notify

import (
	"b612.me/stario"
	"context"
	"errors"
	"io"
	"net"
	"strings"
	"sync"
	"time"
)

var transportConnWriteGates sync.Map
var errTransportFrameQueueUnavailable = errors.New("transport frame queue is unavailable")

type connWriteGateRef struct {
	mu   sync.Mutex
	gate chan struct{}
	refs int
}

type vectoredBuffersWriter interface {
	WriteBuffers(*net.Buffers) (int64, error)
}

type vectoredConnUnwrapper interface {
	UnwrapConn() net.Conn
}

func writeFullToConn(conn net.Conn, data []byte) error {
	if conn == nil {
		return net.ErrClosed
	}
	return withRawConnWriteLock(conn, func(conn net.Conn) error {
		return writeFullToConnUnlocked(conn, data)
	})
}

func writeFullToConnUnlocked(conn net.Conn, data []byte) error {
	if conn == nil {
		return net.ErrClosed
	}
	return writeFullToWriterUnlocked(conn, data)
}

func writeFullToWriterUnlocked(writer io.Writer, data []byte) error {
	if writer == nil {
		return io.ErrClosedPipe
	}
	for len(data) > 0 {
		n, err := writer.Write(data)
		if n > 0 {
			data = data[n:]
		}
		if err != nil {
			return err
		}
		if n == 0 {
			return io.ErrNoProgress
		}
	}
	return nil
}

func writeNetBuffersFullUnlocked(conn net.Conn, buffers net.Buffers) error {
	if conn == nil {
		return net.ErrClosed
	}
	writer, writeFn := vectoredWriteStrategy(conn)
	if writeFn == nil {
		return writeRemainingBuffersUnlocked(conn, buffers)
	}
	n, err := writeFn(&buffers)
	if err != nil {
		return err
	}
	if len(buffers) == 0 {
		return nil
	}
	if n == 0 {
		return io.ErrNoProgress
	}
	return writeRemainingBuffersUnlocked(writer, buffers)
}

func vectoredWriteStrategy(conn net.Conn) (io.Writer, func(*net.Buffers) (int64, error)) {
	current := conn
	for depth := 0; depth < 8 && current != nil; depth++ {
		if writer, ok := current.(vectoredBuffersWriter); ok {
			target := current
			return target, writer.WriteBuffers
		}
		switch target := current.(type) {
		case *net.TCPConn:
			return target, func(bufs *net.Buffers) (int64, error) {
				return bufs.WriteTo(target)
			}
		case *net.UnixConn:
			return target, func(bufs *net.Buffers) (int64, error) {
				return bufs.WriteTo(target)
			}
		}
		unwrapper, ok := current.(vectoredConnUnwrapper)
		if !ok {
			break
		}
		next := unwrapper.UnwrapConn()
		if next == nil || next == current {
			break
		}
		current = next
	}
	return nil, nil
}

func writeRemainingBuffersUnlocked(writer io.Writer, buffers net.Buffers) error {
	for _, part := range buffers {
		if len(part) == 0 {
			continue
		}
		if err := writeFullToWriterUnlocked(writer, part); err != nil {
			return err
		}
	}
	return nil
}

func withRawConnWriteLock(conn net.Conn, fn func(net.Conn) error) error {
	return withRawConnWriteLockDeadline(conn, time.Time{}, fn)
}

func withRawConnWriteLockDeadline(conn net.Conn, deadline time.Time, fn func(net.Conn) error) error {
	_, err := withRawConnWriteLockContextDeadline(context.Background(), conn, deadline, fn)
	return err
}

func withRawConnWriteLockContextDeadline(ctx context.Context, conn net.Conn, deadline time.Time, fn func(net.Conn) error) (bool, error) {
	if conn == nil {
		return false, net.ErrClosed
	}
	if ctx == nil {
		ctx = context.Background()
	}
	gateRef := retainRawConnWriteGate(conn)
	defer releaseRawConnWriteGate(conn, gateRef)
	gate := gateRef.gate
	if err := lockWriteGateContextDeadline(ctx, nil, gate, deadline); err != nil {
		return false, err
	}
	defer func() { gate <- struct{}{} }()
	if err := ctx.Err(); err != nil {
		return false, err
	}
	deadline = earlierWriteDeadline(deadline, contextDeadline(ctx))
	if !deadline.IsZero() {
		if err := conn.SetWriteDeadline(deadline); err != nil {
			return true, err
		}
		defer func() {
			_ = conn.SetWriteDeadline(time.Time{})
		}()
	}
	return true, fn(conn)
}

func retainRawConnWriteGate(conn net.Conn) *connWriteGateRef {
	if conn == nil {
		return &connWriteGateRef{gate: newConnWriteGate(), refs: 1}
	}
	for {
		candidate := &connWriteGateRef{gate: newConnWriteGate(), refs: 1}
		actual, loaded := transportConnWriteGates.LoadOrStore(conn, candidate)
		if !loaded {
			return candidate
		}
		ref := actual.(*connWriteGateRef)
		ref.mu.Lock()
		if ref.refs > 0 {
			ref.refs++
			ref.mu.Unlock()
			return ref
		}
		ref.mu.Unlock()
		transportConnWriteGates.CompareAndDelete(conn, ref)
	}
}

func releaseRawConnWriteGate(conn net.Conn, ref *connWriteGateRef) {
	if ref == nil {
		return
	}
	ref.mu.Lock()
	if ref.refs > 0 {
		ref.refs--
	}
	remove := ref.refs == 0
	ref.mu.Unlock()
	if remove && conn != nil {
		transportConnWriteGates.CompareAndDelete(conn, ref)
	}
}

func newConnWriteGate() chan struct{} {
	gate := make(chan struct{}, 1)
	gate <- struct{}{}
	return gate
}

func lockWriteGateContextDeadline(ctx context.Context, stop <-chan struct{}, gate chan struct{}, deadline time.Time) error {
	if ctx == nil {
		ctx = context.Background()
	}
	select {
	case <-ctx.Done():
		return ctx.Err()
	case <-stop:
		return net.ErrClosed
	default:
	}
	select {
	case <-gate:
		return nil
	default:
	}
	if deadline.IsZero() {
		select {
		case <-ctx.Done():
			return ctx.Err()
		case <-stop:
			return net.ErrClosed
		case <-gate:
			return nil
		}
	}
	wait := time.Until(deadline)
	if wait <= 0 {
		return context.DeadlineExceeded
	}
	timer := time.NewTimer(wait)
	defer timer.Stop()
	select {
	case <-ctx.Done():
		return ctx.Err()
	case <-stop:
		return net.ErrClosed
	case <-timer.C:
		return context.DeadlineExceeded
	case <-gate:
		return nil
	}
}

func shorterPositiveDuration(left time.Duration, right time.Duration) time.Duration {
	left = maxDuration(0, left)
	right = maxDuration(0, right)
	if left == 0 {
		return right
	}
	if right == 0 || left < right {
		return left
	}
	return right
}

func writeFramedPayloadUnlocked(conn net.Conn, queue *stario.StarQueue, payload []byte) error {
	if conn == nil {
		return net.ErrClosed
	}
	if queue == nil {
		return errTransportFrameQueueUnavailable
	}
	if isPacketTransportConn(conn) {
		return writeFullToConnUnlocked(conn, queue.BuildMessage(payload))
	}
	return queue.WriteFrameBuffers(conn, payload)
}

func writeFramedPayloadBatchUnlocked(conn net.Conn, queue *stario.StarQueue, payloads [][]byte) error {
	if conn == nil {
		return net.ErrClosed
	}
	if queue == nil {
		return errTransportFrameQueueUnavailable
	}
	if len(payloads) == 0 {
		return nil
	}
	if isPacketTransportConn(conn) {
		for _, payload := range payloads {
			if err := writeFullToConnUnlocked(conn, queue.BuildMessage(payload)); err != nil {
				return err
			}
		}
		return nil
	}
	return queue.WriteFramesBuffers(conn, payloads...)
}

func isPacketTransportConn(conn net.Conn) bool {
	if conn == nil {
		return false
	}
	if _, ok := conn.(*net.UDPConn); ok {
		return true
	}
	return isPacketNetwork(addrNetwork(conn.LocalAddr())) || isPacketNetwork(addrNetwork(conn.RemoteAddr()))
}

func addrNetwork(addr net.Addr) string {
	if addr == nil {
		return ""
	}
	return addr.Network()
}

func isPacketNetwork(network string) bool {
	switch strings.ToLower(network) {
	case "udp", "udp4", "udp6":
		return true
	default:
		return false
	}
}