notify/transport_write.go

package notify

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

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

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
	}
	for len(data) > 0 {
		n, err := conn.Write(data)
		if n > 0 {
			data = data[n:]
		}
		if err != nil {
			return err
		}
		if n == 0 {
			return io.ErrNoProgress
		}
	}
	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 {
	if conn == nil {
		return net.ErrClosed
	}
	lock := rawConnWriteLock(conn)
	lock.Lock()
	defer lock.Unlock()
	if !deadline.IsZero() {
		if err := conn.SetWriteDeadline(deadline); err != nil {
			return err
		}
		defer func() {
			_ = conn.SetWriteDeadline(time.Time{})
		}()
	}
	return fn(conn)
}

func rawConnWriteLock(conn net.Conn) *sync.Mutex {
	if conn == nil {
		return &sync.Mutex{}
	}
	if lock, ok := transportConnWriteLocks.Load(conn); ok {
		return lock.(*sync.Mutex)
	}
	lock := &sync.Mutex{}
	actual, _ := transportConnWriteLocks.LoadOrStore(conn, lock)
	return actual.(*sync.Mutex)
}

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
	}
}