golib/network/client.go

package network

import (
	"fmt"
	"io"
	"net"
	"sync"
	"time"
)

// TCPClient 用于所有使用 TCP 协议的客户端, 可以通过 Dial 创建此连接, 但通常应该是用 Client 接口而不是只用 TCPClient 结构体指针
type TCPClient struct {
	// reconnect 自动重连, 默认为 true, 当 Read / Write 遇到错误时主动断开连接并会通过 reconnecting 重连. 重连期间调用 Read / Write
	// 时会返回 ErrReconnect 错误. 当调用 Close 时 reconnect 会被更改为 false
	reconnect bool

	// connected 已连接, 默认为 true.
	// 调用 Close 后 connected 会被更改为 false
	// 值为 false 时表示已与服务器断开连接, 之后调用 Read / Write 时会返回原始 socket 错误.
	// 若 reconnect 值为 true 时则断开后会通过 reconnecting 重连, 重连期间调用 Read / Write 时会返回 ErrReconnect 错误.
	connected bool

	// rDeadline 用于 Read 等待超时时间, 优先级高于 deadline
	rDeadline time.Time
	// wDeadline 用于 Write 等待超时时间, 优先级高于 deadline
	wDeadline time.Time
	// deadline 超时时间, 适用于 Read 和 Write, 当 rDeadline 和 wDeadline 不存在时生效
	deadline time.Time

	// conn 服务器连接
	conn *ConnSafe

	mu sync.Mutex
}

// SetReadDeadline 设置 Read 超时时间, 优先级高于 SetDeadline
func (c *TCPClient) SetReadDeadline(t time.Time) error {
	c.rDeadline = t
	return nil
}

// SetWriteDeadline 设置 Write 超时时间, 优先级高于 SetDeadline
func (c *TCPClient) SetWriteDeadline(t time.Time) error {
	c.wDeadline = t
	return nil
}

// SetDeadline 设置 Read / Write 超时时间
func (c *TCPClient) SetDeadline(t time.Time) error {
	c.deadline = t
	return nil
}

// Read 读取数据到 p 中, 使用 setReadDeadline 超时规则
func (c *TCPClient) Read(p []byte) (n int, err error) {
	c.mu.Lock()
	defer c.mu.Unlock()

	if !c.connected {
		if c.reconnect {
			return 0, ErrReconnect
		}
		return 0, ErrClosed
	}

	if err = setReadDeadline(c.conn, c.rDeadline, c.deadline); err != nil {
		err = c.handleErr(err)
		return
	}

	n, err = c.conn.Read(p)
	if err != nil {
		err = c.handleErr(err)
	}
	return
}

// Write 写入 p 至 conn, 使用 setWriteDeadline 超时规则
func (c *TCPClient) Write(p []byte) (n int, err error) {
	c.mu.Lock()
	defer c.mu.Unlock()

	if !c.connected {
		if c.reconnect {
			return 0, ErrReconnect
		}
		return 0, ErrClosed
	}

	if err = setWriteDeadline(c.conn, c.wDeadline, c.deadline); err != nil {
		err = c.handleErr(err)
		return
	}

	n, err = c.conn.Write(p)
	if err != nil {
		err = c.handleErr(err)
	}
	return
}

// Close 主动关闭连接
func (c *TCPClient) Close() error {
	c.mu.Lock()
	defer c.mu.Unlock()

	if !c.connected {
		return nil
	}

	_ = c.conn.Close()
	c.reconnect = false
	c.connected = false
	return nil
}

func (c *TCPClient) LocalAddr() net.Addr {
	return c.conn.LocalAddr()
}

func (c *TCPClient) RemoteAddr() net.Addr {
	return c.conn.RemoteAddr()
}

// handleErr 当 err != nil 时, 若 connected == true && reconnect == true 则关闭连接并将 connected 更改为 ErrReconnect
func (c *TCPClient) handleErr(err error) error {
	if err == nil {
		return nil
	}
	if c.connected && c.reconnect {
		_ = c.conn.Close()
		c.connected = false
		return ErrReconnect
	}
	return err
}

// reconnecting 每 2 秒检查一次连接, 当 reconnect == true 且 connected == false 时使用 DefaultDialTimout 进行重连.
// 主动调用 Close 会使 reconnect == false
// 无限次重试, 直至连接成功
func (c *TCPClient) reconnecting() {
	t := time.NewTicker(2 * time.Second)
	for range t.C {
		if !c.reconnect {
			break
		}
		if c.connected {
			continue
		}
		addr := c.RemoteAddr().(*net.TCPAddr).AddrPort()
		conn, err := net.DialTimeout(NetTCP, addr.String(), DefaultDialTimout)
		if err == nil {
			c.mu.Lock()
			c.conn.Set(conn)
			c.connected = true
			c.mu.Unlock()
		}
	}
	t.Stop()
}

func createTCPClient(conn net.Conn) net.Conn {
	tc := new(TCPClient)
	tc.reconnect = true
	tc.connected = true
	tc.conn = &ConnSafe{}
	tc.conn.Set(conn)
	go tc.reconnecting()
	return tc
}

// modbusClient 实现 ModbusClient 接口, 用于客户端需要异步获取服务器状态的场景, 详情见 async
type modbusClient struct {
	connected bool

	e error
	b []byte
	p chan []byte

	data ModbusCreator
	conn net.Conn
}

// Get 数据来自 conn 服务器返回的数据. 仅保留最后一次服务器返回的数据
// 当遇到非 ErrReconnect 的错误时应调用 Close 关闭此连接, 否则 async 可能会一直返回 socket 错误
func (ms *modbusClient) Read(b []byte) (n int, err error) {
	if !ms.connected {
		return 0, ErrClosed
	}
	t := time.Now().Add(DefaultWriteTimout + DefaultModbusWriteInterval)
	for cap(ms.b) == 0 {
		timout := time.Now().Add(100 * time.Millisecond)
		if t.Equal(timout) || t.Before(timout) {
			return 0, ErrTimout
		}
		time.Sleep(100 * time.Millisecond)
	}
	copy(b, ms.b)
	return len(ms.b), ms.e
}

func (ms *modbusClient) Write(p []byte) (n int, err error) {
	if !ms.connected {
		return 0, ErrClosed
	}
	ms.p <- p
	return len(p), nil
}

// Close 断开与服务器的连接, 关闭 async 线程
func (ms *modbusClient) Close() error {
	if !ms.connected {
		return nil
	}
	ms.connected = false
	ms.b = make([]byte, 0)
	return ms.conn.Close()
}

func (ms *modbusClient) writeRead(p []byte) ([]byte, error) {
	if _, err := ms.conn.Write(p); err != nil {
		return nil, err
	}
	b := defaultPool.Get().(Bytes)
	defaultPool.Put(b)

	n, err := ms.conn.Read(b)
	if err != nil {
		return nil, err
	}
	return b[:n].Remake(), nil
}

// async 每 1 秒调用 ModbusCreator 接口创建数据并发送至 conn, 然后将返回的数据保存至 b
// 如果期间遇到任何错误将会继续重试, 除非主动调用 Close 关闭
func (ms *modbusClient) async() {
	t := time.NewTicker(DefaultModbusWriteInterval)
	defer func() {
		t.Stop()
		_ = ms.Close()
	}()

	for ms.connected {
		select {
		case p, ok := <-ms.p:
			if ok {
				ms.b, ms.e = ms.writeRead(p)
			}
		case <-t.C:
			// 如果创建数据失败则关闭连接
			b, err := ms.data.Create()
			if err != nil {
				ms.e = fmt.Errorf("modbusClient.Create: %s", err)
				return
			}
			ms.b, ms.e = ms.writeRead(b)
		}
	}
}

func createModbusClient(conn net.Conn, data ModbusCreator) io.ReadWriteCloser {
	ms := new(modbusClient)
	ms.connected = true
	ms.b = make([]byte, 0)
	ms.p = make(chan []byte, 1)
	ms.data = data
	ms.conn = conn
	go ms.async()
	return ms
}