| 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340 | package networkimport (	"fmt"	"net"	"net/netip"	"sync"	"time")// TCPClient 用于所有使用 TCP 协议的客户端, 可以通过 Dial 创建此连接, 但通常应该是用 Client 接口而不是只用 TCPClient 结构体指针type TCPClient struct {	// reconnect 自动重连, 可多次开启或关闭, 开启后 Read / Write 遇到错误时会自动关闭连接然后使用 reconnecting 重连, 重连期间调用	// Read / Write 时会返回 ErrReconnect 错误, 因此可以通过此错误翻盘	reconnect bool	// connected 值为 false 时表示此连接由于超时或者服务器异常而被动关闭. 断开后调用 Read / Write 时会返回原始 socket 错误.	// 若 reconnect 值为 true 时则断开后会通过 reconnecting 重连, 重连期间调用 Read / Write 时返回 ErrReconnect 错误	connected bool	// closeManually 值为 true 时:	// 表示主动调用 Close 关闭连接, 此连接不可再重用	// 会使 reconnecting 失效	// 调用 Read / Write 时会返回 ErrClosed 错误	closeManually bool	// rDeadline 用于 Read 等待超时时间, 优先级高于 deadline	rDeadline time.Duration	// wDeadline 用于 Write 等待超时时间, 优先级高于 deadline	wDeadline time.Duration	// deadline 超时时间, 适用于 Read 和 Write, 当 rDeadline 和 wDeadline 不存在时生效	deadline time.Duration	conn net.Conn	mu sync.Mutex}// SetReadDeadline 设置 Read 超时时间, 优先级高于 SetDeadlinefunc (c *TCPClient) SetReadDeadline(timout time.Duration) {	c.rDeadline = timout}// SetWriteDeadline 设置 Write 超时时间, 优先级高于 SetDeadlinefunc (c *TCPClient) SetWriteDeadline(timout time.Duration) {	c.wDeadline = timout}// SetDeadline 设置 Read / Write 超时时间func (c *TCPClient) SetDeadline(timout time.Duration) {	c.deadline = timout}// SetReconnect 开启或关闭自动重连功能func (c *TCPClient) SetReconnect(r bool) {	c.reconnect = r}// Read 读取数据到 p 中, 使用 setReadDeadline 超时规则// 读取错误时:////	reconnect == true:  主动关闭连接并返回 ErrReconnect 错误, 重连期间调用 Read 时继续返回 ErrReconnect 错误//	reconnect == false: 返回原始错误//// 连接关闭时(connected == false):////	主动关闭(closeManually == true):  返回 ErrClosed//	开启自动重连时(reconnect == true): 返回 ErrReconnect//// 调用示例:// p := defaultPool.Get().([]byte)// defaultPool.Put(p)// b, err := Read(p)////	if err == ErrReconnect {//	   continue//	}////	if err != nil {//	   return//	}func (c *TCPClient) Read(p []byte) (n int, err error) {	if !c.connected {		if c.closeManually {			return 0, ErrClosed		}		if c.reconnect {			return 0, ErrReconnect		}	}	c.mu.Lock()	if err = c.setReadDeadline(); err != nil {		c.mu.Unlock()		return	}	n, err = c.conn.Read(p)	if err != nil {		if c.reconnect {			err = ErrReconnect		}		c.passiveClose()	}	c.mu.Unlock()	return}// Write 写入 p 至 conn, 使用 setWriteDeadline 超时规则// 写入错误时:////	reconnect == true:  主动关闭连接并返回 ErrReconnect 错误, 重连期间调用 Write 时继续返回 ErrReconnect 错误//	reconnect == false: 返回原始错误//// 连接关闭时(connected == false):////	主动关闭(closeManually == true):  返回 ErrClosed//	开启自动重连时(reconnect == true): 返回 ErrReconnect//// 调用示例:// n, err := Write(p)////	if err == ErrReconnect {//	   continue//	}////	if err != nil || len(p) != n {//	   return//	}func (c *TCPClient) Write(p []byte) (n int, err error) {	if !c.connected {		if c.closeManually {			return 0, ErrClosed		}		if c.reconnect {			return 0, ErrReconnect		}	}	c.mu.Lock()	defer c.mu.Unlock()	if err = c.setWriteDeadline(); err != nil {		return	}	n, err = c.conn.Write(p)	if err != nil {		if c.reconnect {			err = ErrReconnect		}		c.passiveClose()		return	}	if len(p) != n {		err = ErrNotFullyWrite	}	return}// Close 主动关闭连接// 调用后会关闭 reconnecting 线程, 关闭与服务器的连接, 并设置// closeManually = true// connected = falsefunc (c *TCPClient) Close() error {	if !c.connected {		return nil	}	c.mu.Lock()	_ = c.conn.Close()	c.closeManually = true	c.connected = false	c.mu.Unlock()	return nil}// setReadDeadline 设置 Read 读取超时, 必须在 Read 前调用. 优先级高于 deadline,// 当 rDeadline <= 0 时使用 deadline, 当两者都 <= 0 时则使用 DefaultReadTimoutfunc (c *TCPClient) setReadDeadline() error {	var timout time.Duration	if c.rDeadline > 0 {		timout = c.rDeadline	} else if c.deadline > 0 {		timout = c.deadline	} else {		timout = DefaultReadTimout	}	return c.conn.SetReadDeadline(time.Now().Add(timout))}// setWriteDeadline 设置 Write 读取超时, 必须在 Write 前调用. 优先级高于 deadline// 当 wDeadline <= 0 时使用 deadline, 当两者都 <= 0 时则使用 DefaultWriteTimoutfunc (c *TCPClient) setWriteDeadline() error {	var timout time.Duration	if c.wDeadline > 0 {		timout = c.wDeadline	} else if c.deadline > 0 {		timout = c.deadline	} else {		timout = DefaultWriteTimout	}	return c.conn.SetWriteDeadline(time.Now().Add(timout))}// passiveClose 被动关闭连接, 在 Read 和 Write 返回错误时在 mu 中调用.func (c *TCPClient) passiveClose() {	if c.connected && c.reconnect {		_ = c.conn.Close()		c.connected = false	}}// getAddr 获取服务器的 IP 和 Port, 用于 reconnecting// 即使 conn 被 Close 也可以正常获取func (c *TCPClient) getAddr() netip.AddrPort {	return c.conn.RemoteAddr().(*net.TCPAddr).AddrPort()}// reconnecting 每 1 秒检查一次连接, 当 closeManually == false 且 connected 和 reconnect == true 时使用 DefaultReconnectTimout 进行重连.// 主动调用 Close 会使 closeManually == true// Read 或 Write 遇到错误时满足 connected 和 reconnect == true (重连的条件)// 无限次重试, 直至连接成功func (c *TCPClient) reconnecting() {	for {		time.Sleep(1 * time.Second)		if c.closeManually {			return		}		if c.connected || !c.reconnect {			continue		}		addr := c.getAddr()		conn, err := net.DialTimeout(NetTCP, addr.String(), DefaultReconnectTimout)		if err == nil {			c.mu.Lock()			c.conn = (net.Conn)(nil)			c.conn = conn			c.connected = true			c.mu.Unlock()		}	}}// ModbusClient 用于 Modbus/TCP 服务器交互的快捷接口, 内部由 TCPClient 实现type ModbusClient struct {	conn Client}// WriteRead 写入 p 并读取返回数据, Write 或 Read 返回错误时, 见 Clientfunc (mc *ModbusClient) WriteRead(p []byte) ([]byte, error) {	n, err := mc.conn.Write(p)	if err != nil {		return nil, err	}	b := defaultPool.Get().([]byte)	defaultPool.Put(b)	n, err = mc.conn.Read(b)	if err != nil {		return nil, err	}	return Remake(b[:n]), nil}func (mc *ModbusClient) Close() error {	return mc.conn.Close()}// modbusStatus 实现 ModbusStatus 接口, 用于客户端需要实时获取服务器状态的场景, 详情见 getStatustype modbusStatus struct {	connected bool	e         error	b         []byte	msw       ModbusStatusWriter	conn      Client}// Get 数据来自 Modbus 服务器返回的数据. 仅保留最后一次服务器返回的数据// 当遇到非 ErrReconnect 的错误时应调用 Close 关闭此连接, 否则 getStatus 可能会一直返回 socket 错误func (ms *modbusStatus) Get() ([]byte, error) {	if !ms.connected {		return nil, ErrClosed	}	if ms.e == nil && cap(ms.b) == 0 {		return ms.Get()	}	return ms.b, ms.e}// Close 断开与服务器的连接, 关闭 getStatus 线程func (ms *modbusStatus) Close() error {	if !ms.connected {		return nil	}	ms.connected = false	ms.b = make([]byte, 0)	return ms.conn.Close()}// getStatus 每 1 秒调用 ModbusStatusWriter 接口创建数据并发送至 conn, 然后将返回的数据保存至 b// 如果期间遇到任何错误将会继续重试, 除非主动调用 Close 关闭func (ms *modbusStatus) getStatus() {	var (		i   int		b   []byte		err error	)	defer func() {		_ = ms.Close()	}()	for {		if !ms.connected {			return		}		time.Sleep(1 * time.Second)		b, err = ms.msw.Create()		if err != nil {			ms.e = fmt.Errorf("called ModbusStatusWrite.Create: %s", err)			return		}		if _, ms.e = ms.conn.Write(b); ms.e != nil {			continue		}		b = defaultPool.Get().([]byte)		defaultPool.Put(b)		i, ms.e = ms.conn.Read(b)		if ms.e != nil {			continue		}		ms.b = Remake(b[:i])	}}
 |