golib/gnet/net.go

package gnet

import (
	"crypto/tls"
	"errors"
	"fmt"
	"math/rand/v2"
	"net"
	"sync"
	"time"
)

const (
	ClientReadTimout  = 10 * time.Second
	ClientWriteTimout = 5 * time.Second
)

const (
	ServerReadTimout   = 60 * time.Second
	ServerWriteTimeout = 5 * time.Second
)

const (
	IdleTime = 1 * time.Second
)

const (
	DialTimout = 10 * time.Second
)

const (
	MaxBuffSize = 4096
)

var (
	// ErrConnNotFound 连接不存在
	ErrConnNotFound = errors.New("network: connection not found")
)

type Timeout struct {
	Msg string
}

func (t *Timeout) Timeout() bool { return true }

func (t *Timeout) Error() string {
	if t.Msg == "" {
		return "network: timeout"
	}
	return fmt.Sprintf("network: timeout -> %s", t.Msg)
}

// ReadMultiplexer 读取复用
type ReadMultiplexer interface {
	// ReadMux 将读取的数据存储至内部切片中, b 则是内部切片的指针引用. ReadMux 被调用时, 总是会清除上一次保存的数据. 即你需要将 b 使用完毕
	// 以后再调用, 否则数据将会被覆盖.
	ReadMux() (b []byte, err error)
}

// Config 连接配置
// 当任意Timeout未设定时则表示无超时
type Config struct {
	ReadTimeout  time.Duration
	WriteTimeout time.Duration
	Timeout      time.Duration // Read and Write
	DialTimeout  time.Duration

	Reconnect bool // Client Only
	MuxBuff   int  // ReadMultiplexer.ReadMux Only
}

func (c *Config) Client() *Config {
	c.ReadTimeout = ClientReadTimout
	c.WriteTimeout = ClientWriteTimout
	c.DialTimeout = DialTimout
	return c
}

func (c *Config) Server() *Config {
	c.ReadTimeout = ServerReadTimout
	c.WriteTimeout = ServerWriteTimeout
	return c
}

// TCPConn 基于 net.Conn 增加在调用 Read 和 Write 时补充超时设置
type TCPConn struct {
	net.Conn
	mu sync.Mutex

	buf    []byte
	Config *Config
}

func (t *TCPConn) setReadTimeout() (err error) {
	if t.Config == nil {
		return
	}
	if t.Config.Timeout > 0 {
		return t.Conn.SetDeadline(time.Now().Add(t.Config.Timeout))
	}
	if t.Config.ReadTimeout > 0 {
		return t.Conn.SetReadDeadline(time.Now().Add(t.Config.ReadTimeout))
	}
	return
}

func (t *TCPConn) setWriteTimout() (err error) {
	if t.Config == nil {
		return
	}
	if t.Config.Timeout > 0 {
		return t.Conn.SetDeadline(time.Now().Add(t.Config.Timeout))
	}
	if t.Config.WriteTimeout > 0 {
		return t.Conn.SetWriteDeadline(time.Now().Add(t.Config.WriteTimeout))
	}
	return
}

func (t *TCPConn) Read(b []byte) (n int, err error) {
	t.mu.Lock()
	defer t.mu.Unlock()
	if err = t.setReadTimeout(); err != nil {
		return
	}
	return t.Conn.Read(b)
}

func (t *TCPConn) ReadMux() (b []byte, err error) {
	if len(t.buf) == 0 {
		bufSize := t.Config.MuxBuff
		if bufSize <= 0 {
			bufSize = MaxBuffSize
		}
		t.buf = make([]byte, bufSize)
	}
	n, err := t.Read(t.buf)
	if err != nil {
		return nil, err
	}
	return t.buf[:n], nil
}

func (t *TCPConn) Write(b []byte) (n int, err error) {
	t.mu.Lock()
	defer t.mu.Unlock()
	if err = t.setReadTimeout(); err != nil {
		return
	}
	return t.Conn.Write(b)
}

func (t *TCPConn) Close() error {
	err := t.Conn.Close()
	t.buf = nil
	return err
}

type Connection interface {
	IsConnected() bool
	IsClosed() bool
	Reconnecting() bool
}

type tcpAliveConn struct {
	net.Conn
	DialTimeout time.Duration
	mu          *sync.Mutex

	handing bool
	closed  bool
}

func (t *tcpAliveConn) IsConnected() bool {
	if t.Conn == nil {
		return false
	}
	if t.handing || t.closed {
		return false
	}
	return true
}

func (t *tcpAliveConn) IsClosed() bool {
	return t.closed
}

func (t *tcpAliveConn) Reconnecting() bool {
	if t.Conn == nil {
		return false
	}
	return t.handing && !t.closed
}

// hasAvailableNetFace
// 检查当前操作系统中是否存在可用的网卡, 无可用的网卡时挂起重连操作
// 修复部分操作系统(Windows)休眠后网卡状态异常导致 net.DialTimeout 锥栈溢出(然后panic)的问题
func (t *tcpAliveConn) hasAvailableNetFace() bool {
	ift, err := net.Interfaces()
	if err != nil {
		return false
	}
	i := 0
	for _, ifi := range ift {
		// FlagUp 网线插入, FlagLoopback 本机循环网卡 FlagRunning 活动的网卡
		if ifi.Flags&net.FlagUp != 0 && ifi.Flags&net.FlagLoopback == 0 && ifi.Flags&net.FlagRunning != 0 {
			i++
		}
	}
	return i > 0
}

func (t *tcpAliveConn) Dial(addr net.Addr) (net.Conn, error) {
	tcpConn, err := net.DialTimeout("tcp", addr.String(), t.DialTimeout)
	if err != nil {
		return nil, err
	}
	if tcp, ok := tcpConn.(*net.TCPConn); ok {
		_ = tcp.SetNoDelay(true)
		_ = tcp.SetKeepAlive(true)
		_ = tcp.SetKeepAlivePeriod(5 * time.Second)
	}
	return tcpConn, nil
}

func (t *tcpAliveConn) handleAlive() {
	if t.closed || t.handing {
		return
	}
	t.handing = true
	_ = t.Conn.Close() // 关掉旧的连接
	for !t.closed {
		if !t.hasAvailableNetFace() {
			time.Sleep(3 * time.Second)
			continue
		}
		conn, err := t.Dial(t.RemoteAddr())
		if err != nil {
			continue
		}
		t.mu.Lock()
		t.Conn = conn
		t.mu.Unlock()
		break
	}
	if t.closed { // 当连接被主动关闭时
		_ = t.Conn.Close() // 即使重连上也关闭
	}
	t.handing = false
}

func (t *tcpAliveConn) handleErr(err error) error {
	if err == nil {
		return nil
	}
	if t.closed {
		return err
	}
	// 延迟后返回. 通常上层代码在 for 循环中调用 Read/Write. 如果重连期间的调用响应过快, 则会导致上层日志写入频繁
	// 如果已主动调用 Close 则保持不变
	t.randSleep()
	msg := "tcpAliveConn handing: " + err.Error()
	return &Timeout{Msg: msg}
}

func (t *tcpAliveConn) randSleep() {
	minSleep := 900
	maxSleep := 3100
	randSleep := rand.IntN(maxSleep-minSleep) + minSleep
	time.Sleep(time.Duration(randSleep) * time.Millisecond)
}

func (t *tcpAliveConn) Read(b []byte) (n int, err error) {
	n, err = t.Conn.Read(b)
	if err != nil {
		go t.handleAlive()
	}
	return n, t.handleErr(err)
}

func (t *tcpAliveConn) Write(b []byte) (n int, err error) {
	n, err = t.Conn.Write(b)
	if err != nil {
		go t.handleAlive()
	}
	return n, t.handleErr(err)
}

func (t *tcpAliveConn) Close() error {
	if t.closed {
		return nil
	}
	t.closed = true
	return t.Conn.Close()
}

func DialTCP(address string) (net.Conn, error) {
	return DialTCPConfig(address, nil)
}

func DialTCPConfig(address string, config *Config) (net.Conn, error) {
	if _, err := net.ResolveTCPAddr("tcp", address); err != nil {
		return nil, err
	}
	if config == nil {
		config = (&Config{}).Client()
	}
	if config.DialTimeout <= 0 {
		config.DialTimeout = DialTimout
	}
	tcpConn, err := net.DialTimeout("tcp", address, config.DialTimeout)
	if err != nil {
		return nil, err
	}
	if tcp, ok := tcpConn.(*net.TCPConn); ok {
		_ = tcp.SetNoDelay(true)
		_ = tcp.SetKeepAlive(true)
		_ = tcp.SetKeepAlivePeriod(5 * time.Second)
	}
	conn := &TCPConn{
		Conn:   tcpConn,
		Config: config,
		mu:     sync.Mutex{},
	}
	if config.Reconnect {
		conn.Conn = &tcpAliveConn{
			Conn:        tcpConn,
			DialTimeout: config.DialTimeout,
			mu:          &conn.mu,
		}
	}
	return conn, nil
}

type listener struct {
	net.Listener
	config *Config
}

func (t *listener) Accept() (net.Conn, error) {
	tcpConn, err := t.Listener.Accept()
	if err != nil {
		return nil, err
	}
	conn := &TCPConn{
		Conn:   tcpConn,
		Config: t.config,
	}
	return conn, nil
}

func NewListener(ln net.Listener, config *Config) net.Listener {
	if config == nil {
		config = (&Config{}).Server()
	}
	return &listener{Listener: ln, config: config}
}

func ListenTCP(network, address string) (net.Listener, error) {
	tcpAddr, err := net.ResolveTCPAddr(network, address)
	if err != nil {
		return nil, err
	}
	ln, err := net.ListenTCP(network, tcpAddr)
	if err != nil {
		return nil, err
	}
	return NewListener(ln, nil), nil
}

func ListenTLS(network, address string, config *tls.Config) (net.Listener, error) {
	ln, err := ListenTCP(network, address)
	if err != nil {
		return nil, err
	}
	return tls.NewListener(ln, config), nil
}