golib/gnet/net.go

package gnet

import (
	"context"
	"errors"
	"math/rand/v2"
	"net"
	"sync"
	"time"
)

const (
	ClientReadTimeout  = 10 * time.Second
	ClientWriteTimeout = 5 * time.Second
)

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

const (
	IdleTime = 1 * time.Second
)

const (
	DialTimeout = 10 * time.Second
)

const (
	MaxBuffSize = 4096
)

var (
	// ErrConnNotFound 连接不存在
	ErrConnNotFound = errors.New("gnet: connection not found")
	// ErrWaitingResponse 等待远程主机响应
	ErrWaitingResponse = errors.New("gnet: waiting for response from remote host")
	// ErrUnconnected 未能连接到远程主机; 用于开启 Config.Reconnect 后且需要告知上层逻辑真实的连接情况时使用
	ErrUnconnected = errors.New("gnet: connection not connected")
)

type Timeout struct {
	Msg string
}

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

func (t *Timeout) Error() string {
	if t.Msg == "" {
		return "network: timeout"
	}
	return 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 // Reconnect 自动重连. 仅用于客户端
	IgnoreError bool // IgnoreError 忽略首次连接时失败的错误, 用于 Reconnect 启用时. 仅用于客户端
	MuxBuff     int  // ReadMultiplexer.ReadMux Only
}

func (c *Config) Client() *Config {
	c.ReadTimeout = ClientReadTimeout
	c.WriteTimeout = ClientWriteTimeout
	c.DialTimeout = DialTimeout
	return c
}

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

type ConnStat interface {
	IsConnected() bool
	IsClosed() bool
}

func optimizationConn(conn net.Conn) net.Conn {
	if tcp, ok := conn.(*net.TCPConn); ok {
		_ = tcp.SetNoDelay(true)
		_ = tcp.SetKeepAlive(true)
		_ = tcp.SetKeepAlivePeriod(5 * time.Second)
	}
	return conn
}

type tcpAliveConn struct {
	address string
	net.Conn

	Config *Config
	buf    []byte
	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
}

// 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(address string, timeout time.Duration) (net.Conn, error) {
	tcpConn, err := net.DialTimeout("tcp", address, timeout)
	if err != nil {
		return nil, err
	}
	return optimizationConn(tcpConn), nil
}

func (t *tcpAliveConn) handleAlive() {
	if t.closed || t.handing {
		return
	}
	if !t.Config.Reconnect {
		_ = t.Close() // 如果未开启重连, 出现任何错误时都会主动关闭连接
		return
	}
	t.handing = true
	if t.Conn != nil {
		_ = t.Conn.Close() // 关掉旧的连接
	}
	for !t.closed {
		if !t.hasAvailableNetFace() {
			time.Sleep(DialTimeout)
			continue
		}
		conn, err := t.Dial(t.address, t.Config.DialTimeout)
		if err != nil {
			time.Sleep(DialTimeout)
			continue
		}
		t.mu.Lock()
		t.Conn = conn
		t.mu.Unlock()
		break
	}
	if t.closed { // 当连接被主动关闭时
		if t.Conn != nil {
			_ = t.Conn.Close() // 即使重连上也关闭
		}
	}
	t.handing = false
}

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

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

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

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

func (t *tcpAliveConn) Read(b []byte) (n int, err error) {
	t.mu.Lock()
	defer t.mu.Unlock()
	if err = t.setReadTimeout(); err != nil {
		return
	}
	if t.Conn == nil {
		return 0, t.handleErr(ErrWaitingResponse)
	}
	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) {
	t.mu.Lock()
	defer t.mu.Unlock()
	if err = t.setWriteTimeout(); err != nil {
		return
	}
	if t.Conn == nil {
		return 0, t.handleErr(ErrWaitingResponse)
	}
	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
	var err error
	if t.Conn != nil {
		err = t.Conn.Close()
	}
	t.buf = nil
	t.Conn = nil
	return err
}

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

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 = DialTimeout
	}
	if config.Reconnect && config.IgnoreError {
		conn := &tcpAliveConn{
			address: address,
			Conn:    nil,
			Config:  config,
		}
		go conn.handleAlive()
		return conn, nil
	}
	tcpConn, err := net.DialTimeout("tcp", address, config.DialTimeout)
	if err != nil {
		return nil, err
	}
	conn := &tcpAliveConn{
		address: address,
		Conn:    optimizationConn(tcpConn),
		Config:  config,
	}
	return conn, nil
}

func ReadWithContext(ctx context.Context, conn net.Conn, b []byte) (n int, err error) {
	done := make(chan struct{})
	stop := context.AfterFunc(ctx, func() {
		_ = conn.SetReadDeadline(time.Now())
		close(done)
	})
	n, err = conn.Read(b)
	if !stop() {
		<-done
		_ = conn.SetReadDeadline(time.Time{})
		if err == nil {
			err = ctx.Err()
		}
		return n, err
	}
	return n, err
}

func WriteWithContext(ctx context.Context, conn net.Conn, b []byte) (n int, err error) {
	done := make(chan struct{})
	stop := context.AfterFunc(ctx, func() {
		_ = conn.SetWriteDeadline(time.Now())
		close(done)
	})
	n, err = conn.Write(b)
	if !stop() {
		<-done
		_ = conn.SetWriteDeadline(time.Time{})
		if err == nil {
			err = ctx.Err()
		}
		return n, err
	}
	return n, err
}

type connWithContext struct {
	ctx context.Context
	net.Conn
}

func (c *connWithContext) Read(b []byte) (n int, err error) {
	return ReadWithContext(c.ctx, c.Conn, b)
}

func (c *connWithContext) Write(b []byte) (n int, err error) {
	return WriteWithContext(c.ctx, c.Conn, b)
}

func NewConnWithContext(ctx context.Context, conn net.Conn) net.Conn {
	return &connWithContext{ctx: ctx, Conn: conn}
}