golib/v2/gnet/net.go

package gnet

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

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

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

const (
	WriteInterval = 1 * time.Second
)

const (
	DialTimout = 2 * 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)
}

type Config struct {
	ReadTimout  time.Duration
	WriteTimout time.Duration
	Timout      time.Duration // Read and Write
	DialTimout  time.Duration
}

func (c *Config) Client() *Config {
	c.ReadTimout = ClientReadTimout
	c.WriteTimout = ClientWriteTimout
	c.DialTimout = DialTimout
	return c
}

func (c *Config) Server() *Config {
	c.ReadTimout = ServerReadTimout
	c.WriteTimout = ServerWriteTimeout
	return c
}

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

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

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

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

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

type tcpAliveConn struct {
	net.Conn
	config *Config
	mu     sync.Mutex

	handing bool
	closed  bool
}

func (t *tcpAliveConn) Dial(network, address string, config *Config) (net.Conn, error) {
	conn, err := DialTCPConfig(network, address, config)
	if err != nil {
		return nil, err
	}
	return &tcpAliveConn{Conn: conn, config: config}, nil
}

func (t *tcpAliveConn) handleAlive(force bool) {
	if t.closed {
		return
	}
	if !force && t.handing {
		return
	}
	t.handing = true
	_ = t.Conn.Close() // 关掉旧的连接
	rAddr := t.RemoteAddr()
	conn, err := t.Dial(rAddr.Network(), rAddr.String(), t.config)
	if err != nil {
		t.handleAlive(true)
		return
	}
	t.mu.Lock()
	t.Conn = conn
	t.handing = false
	t.mu.Unlock()
}

func (t *tcpAliveConn) handleErr(err error) error {
	if t.closed {
		return err
	}
	if t.handing {
		msg := "tcpAliveConn handing: "
		if err == nil {
			msg = msg + "..."
		} else {
			msg = msg + err.Error()
		}
		return &Timeout{Msg: msg}
	}
	return err
}

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

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

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

func Client(conn net.Conn, config *Config) *TCPConn {
	if config == nil {
		config = (&Config{}).Client()
	}
	client := &TCPConn{
		Conn:   conn,
		Config: config,
	}
	return client
}

func DialTCP(network, address string) (net.Conn, error) {
	return DialTCPConfig(network, address, &Config{})
}

func DialTCPConfig(network, address string, config *Config) (*TCPConn, error) {
	tcpAddr, err := net.ResolveTCPAddr(network, address)
	if err != nil {
		return nil, err
	}
	if config.DialTimout <= 0 {
		config.DialTimout = DialTimout
	}
	tcpConn, err := net.DialTimeout(network, tcpAddr.String(), config.DialTimout)
	if err != nil {
		return nil, err
	}
	return Client(tcpConn, config), nil
}

func DialTCPAlive(network, address string, config *Config) (net.Conn, error) {
	var dialer tcpAliveConn
	return dialer.Dial(network, address, config)
}

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
}