|
|
|
|
@ -2,11 +2,14 @@ package client
|
|
|
|
|
|
|
|
|
|
import (
|
|
|
|
|
"context"
|
|
|
|
|
"io"
|
|
|
|
|
"net"
|
|
|
|
|
"sync"
|
|
|
|
|
"time"
|
|
|
|
|
|
|
|
|
|
"github.com/noahlann/nnet/pkg/client"
|
|
|
|
|
"git.noahlan.cn/noahlan/nnet/v2/pkg/client"
|
|
|
|
|
protocolpkg "git.noahlan.cn/noahlan/nnet/v2/pkg/protocol"
|
|
|
|
|
unpackerpkg "git.noahlan.cn/noahlan/nnet/v2/pkg/unpacker"
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
// tcpClient TCP客户端实现
|
|
|
|
|
@ -18,6 +21,13 @@ type tcpClient struct {
|
|
|
|
|
ctx context.Context
|
|
|
|
|
cancel context.CancelFunc
|
|
|
|
|
readBuffer []byte
|
|
|
|
|
requestMu sync.Mutex
|
|
|
|
|
readMu sync.Mutex
|
|
|
|
|
writeMu sync.Mutex
|
|
|
|
|
closeOnce sync.Once
|
|
|
|
|
protocol protocolpkg.Protocol
|
|
|
|
|
unpacker unpackerpkg.Unpacker
|
|
|
|
|
pending [][]byte
|
|
|
|
|
|
|
|
|
|
// 自动重连相关
|
|
|
|
|
reconnectAttempts int
|
|
|
|
|
@ -41,16 +51,19 @@ func NewTCPClient(config *client.Config) client.Client {
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
|
|
|
protocol := newApplicationProtocol(config.ApplicationProtocol)
|
|
|
|
|
|
|
|
|
|
c := &tcpClient{
|
|
|
|
|
config: config,
|
|
|
|
|
ctx: ctx,
|
|
|
|
|
cancel: cancel,
|
|
|
|
|
readBuffer: make([]byte, 4096),
|
|
|
|
|
reconnectCh: make(chan struct{}, 1),
|
|
|
|
|
config: config,
|
|
|
|
|
ctx: ctx,
|
|
|
|
|
cancel: cancel,
|
|
|
|
|
readBuffer: make([]byte, 4096),
|
|
|
|
|
protocol: protocol,
|
|
|
|
|
unpacker: newApplicationUnpacker(protocol),
|
|
|
|
|
reconnectCh: make(chan struct{}, 1),
|
|
|
|
|
reconnectStopCh: make(chan struct{}),
|
|
|
|
|
messageCh: make(chan []byte, 100),
|
|
|
|
|
messageErrCh: make(chan error, 10),
|
|
|
|
|
messageCh: make(chan []byte, 100),
|
|
|
|
|
messageErrCh: make(chan error, 10),
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
// 如果启用自动重连，启动重连goroutine
|
|
|
|
|
@ -85,13 +98,36 @@ func (c *tcpClient) connect() error {
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
// 创建连接
|
|
|
|
|
dialer := &net.Dialer{
|
|
|
|
|
Timeout: c.config.ConnectTimeout,
|
|
|
|
|
connectTimeout := c.config.ConnectTimeout
|
|
|
|
|
if connectTimeout <= 0 {
|
|
|
|
|
connectTimeout = 3 * time.Second
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
conn, err := dialer.DialContext(c.ctx, "tcp", addr)
|
|
|
|
|
deadline := time.Now().Add(connectTimeout)
|
|
|
|
|
var conn net.Conn
|
|
|
|
|
var err error
|
|
|
|
|
for {
|
|
|
|
|
timeout := time.Until(deadline)
|
|
|
|
|
if timeout <= 0 {
|
|
|
|
|
break
|
|
|
|
|
}
|
|
|
|
|
if timeout > 200*time.Millisecond {
|
|
|
|
|
timeout = 200 * time.Millisecond
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
dialer := &net.Dialer{Timeout: timeout}
|
|
|
|
|
conn, err = dialer.DialContext(c.ctx, "tcp", addr)
|
|
|
|
|
if err == nil {
|
|
|
|
|
break
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
select {
|
|
|
|
|
case <-c.ctx.Done():
|
|
|
|
|
return c.ctx.Err()
|
|
|
|
|
case <-time.After(20 * time.Millisecond):
|
|
|
|
|
}
|
|
|
|
|
}
|
|
|
|
|
if err != nil {
|
|
|
|
|
// 如果启用自动重连，触发重连
|
|
|
|
|
if c.config.AutoReconnect {
|
|
|
|
|
select {
|
|
|
|
|
case c.reconnectCh <- struct{}{}:
|
|
|
|
|
@ -104,6 +140,8 @@ func (c *tcpClient) connect() error {
|
|
|
|
|
c.conn = conn
|
|
|
|
|
c.connected = true
|
|
|
|
|
c.reconnectAttempts = 0 // 重置重连次数
|
|
|
|
|
c.unpacker = newApplicationUnpacker(c.protocol)
|
|
|
|
|
c.pending = nil
|
|
|
|
|
|
|
|
|
|
return nil
|
|
|
|
|
}
|
|
|
|
|
@ -123,6 +161,7 @@ func (c *tcpClient) Disconnect() error {
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
c.connected = false
|
|
|
|
|
c.pending = nil
|
|
|
|
|
|
|
|
|
|
// 如果启用自动重连，触发重连
|
|
|
|
|
if c.config.AutoReconnect {
|
|
|
|
|
@ -135,62 +174,168 @@ func (c *tcpClient) Disconnect() error {
|
|
|
|
|
return nil
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
func (c *tcpClient) markDisconnected() {
|
|
|
|
|
c.mu.Lock()
|
|
|
|
|
defer c.mu.Unlock()
|
|
|
|
|
|
|
|
|
|
if c.conn != nil {
|
|
|
|
|
_ = c.conn.Close()
|
|
|
|
|
c.conn = nil
|
|
|
|
|
}
|
|
|
|
|
c.connected = false
|
|
|
|
|
c.pending = nil
|
|
|
|
|
|
|
|
|
|
if c.config.AutoReconnect {
|
|
|
|
|
select {
|
|
|
|
|
case c.reconnectCh <- struct{}{}:
|
|
|
|
|
default:
|
|
|
|
|
}
|
|
|
|
|
}
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
func shouldMarkDisconnected(err error) bool {
|
|
|
|
|
if err == nil {
|
|
|
|
|
return false
|
|
|
|
|
}
|
|
|
|
|
if netErr, ok := err.(net.Error); ok && netErr.Timeout() {
|
|
|
|
|
return false
|
|
|
|
|
}
|
|
|
|
|
return true
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
// Send 发送数据
|
|
|
|
|
func (c *tcpClient) Send(data []byte) error {
|
|
|
|
|
c.mu.RLock()
|
|
|
|
|
defer c.mu.RUnlock()
|
|
|
|
|
frame, err := encodeApplicationMessage(c.protocol, data)
|
|
|
|
|
if err != nil {
|
|
|
|
|
return err
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
c.mu.RLock()
|
|
|
|
|
if !c.connected || c.conn == nil {
|
|
|
|
|
c.mu.RUnlock()
|
|
|
|
|
return client.NewError("not connected")
|
|
|
|
|
}
|
|
|
|
|
conn := c.conn
|
|
|
|
|
writeTimeout := c.config.WriteTimeout
|
|
|
|
|
c.mu.RUnlock()
|
|
|
|
|
|
|
|
|
|
c.writeMu.Lock()
|
|
|
|
|
defer c.writeMu.Unlock()
|
|
|
|
|
|
|
|
|
|
if c.config.WriteTimeout > 0 {
|
|
|
|
|
c.conn.SetWriteDeadline(time.Now().Add(c.config.WriteTimeout))
|
|
|
|
|
if writeTimeout > 0 {
|
|
|
|
|
conn.SetWriteDeadline(time.Now().Add(writeTimeout))
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
_, err := c.conn.Write(data)
|
|
|
|
|
return err
|
|
|
|
|
for len(frame) > 0 {
|
|
|
|
|
n, err := conn.Write(frame)
|
|
|
|
|
if shouldMarkDisconnected(err) {
|
|
|
|
|
c.markDisconnected()
|
|
|
|
|
}
|
|
|
|
|
if err != nil {
|
|
|
|
|
return err
|
|
|
|
|
}
|
|
|
|
|
if n == 0 {
|
|
|
|
|
c.markDisconnected()
|
|
|
|
|
return io.ErrShortWrite
|
|
|
|
|
}
|
|
|
|
|
frame = frame[n:]
|
|
|
|
|
}
|
|
|
|
|
return nil
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
// Receive 接收数据
|
|
|
|
|
func (c *tcpClient) Receive() ([]byte, error) {
|
|
|
|
|
c.mu.RLock()
|
|
|
|
|
defer c.mu.RUnlock()
|
|
|
|
|
return c.receiveWithTimeout(0)
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
func (c *tcpClient) receiveWithTimeout(timeout time.Duration) ([]byte, error) {
|
|
|
|
|
c.readMu.Lock()
|
|
|
|
|
defer c.readMu.Unlock()
|
|
|
|
|
|
|
|
|
|
if len(c.pending) > 0 {
|
|
|
|
|
msg := c.pending[0]
|
|
|
|
|
c.pending[0] = nil
|
|
|
|
|
c.pending = c.pending[1:]
|
|
|
|
|
return msg, nil
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
c.mu.RLock()
|
|
|
|
|
if !c.connected || c.conn == nil {
|
|
|
|
|
c.mu.RUnlock()
|
|
|
|
|
return nil, client.NewError("not connected")
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
if c.config.ReadTimeout > 0 {
|
|
|
|
|
c.conn.SetReadDeadline(time.Now().Add(c.config.ReadTimeout))
|
|
|
|
|
conn := c.conn
|
|
|
|
|
readTimeout := c.config.ReadTimeout
|
|
|
|
|
if timeout > 0 {
|
|
|
|
|
readTimeout = timeout
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
n, err := c.conn.Read(c.readBuffer)
|
|
|
|
|
if err != nil {
|
|
|
|
|
return nil, err
|
|
|
|
|
bufferSize := len(c.readBuffer)
|
|
|
|
|
if bufferSize <= 0 {
|
|
|
|
|
bufferSize = 4096
|
|
|
|
|
}
|
|
|
|
|
protocol := c.protocol
|
|
|
|
|
c.mu.RUnlock()
|
|
|
|
|
|
|
|
|
|
for {
|
|
|
|
|
buffer := make([]byte, bufferSize)
|
|
|
|
|
if readTimeout > 0 {
|
|
|
|
|
conn.SetReadDeadline(time.Now().Add(readTimeout))
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
result := make([]byte, n)
|
|
|
|
|
copy(result, c.readBuffer[:n])
|
|
|
|
|
return result, nil
|
|
|
|
|
n, err := conn.Read(buffer)
|
|
|
|
|
if err != nil {
|
|
|
|
|
if shouldMarkDisconnected(err) {
|
|
|
|
|
c.markDisconnected()
|
|
|
|
|
}
|
|
|
|
|
return nil, err
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
if protocol == nil {
|
|
|
|
|
result := make([]byte, n)
|
|
|
|
|
copy(result, buffer[:n])
|
|
|
|
|
return result, nil
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
if c.unpacker == nil {
|
|
|
|
|
return decodeApplicationMessage(protocol, buffer[:n])
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
frames, _, _, err := c.unpacker.Unpack(buffer[:n])
|
|
|
|
|
if err != nil {
|
|
|
|
|
if shouldMarkDisconnected(err) {
|
|
|
|
|
c.markDisconnected()
|
|
|
|
|
}
|
|
|
|
|
return nil, err
|
|
|
|
|
}
|
|
|
|
|
if len(frames) == 0 {
|
|
|
|
|
continue
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
decoded := make([][]byte, 0, len(frames))
|
|
|
|
|
for _, frame := range frames {
|
|
|
|
|
msg, err := decodeApplicationMessage(protocol, frame)
|
|
|
|
|
if err != nil {
|
|
|
|
|
return nil, err
|
|
|
|
|
}
|
|
|
|
|
decoded = append(decoded, msg)
|
|
|
|
|
}
|
|
|
|
|
c.pending = append(c.pending, decoded[1:]...)
|
|
|
|
|
return decoded[0], nil
|
|
|
|
|
}
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
// Request 请求-响应（带超时）
|
|
|
|
|
func (c *tcpClient) Request(data []byte, timeout time.Duration) ([]byte, error) {
|
|
|
|
|
c.requestMu.Lock()
|
|
|
|
|
defer c.requestMu.Unlock()
|
|
|
|
|
|
|
|
|
|
// 发送请求
|
|
|
|
|
if err := c.Send(data); err != nil {
|
|
|
|
|
return nil, err
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
// 设置读取超时
|
|
|
|
|
oldTimeout := c.config.ReadTimeout
|
|
|
|
|
c.config.ReadTimeout = timeout
|
|
|
|
|
defer func() {
|
|
|
|
|
c.config.ReadTimeout = oldTimeout
|
|
|
|
|
}()
|
|
|
|
|
|
|
|
|
|
// 接收响应
|
|
|
|
|
return c.Receive()
|
|
|
|
|
return c.receiveWithTimeout(timeout)
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
// IsConnected 检查是否已连接
|
|
|
|
|
@ -202,8 +347,10 @@ func (c *tcpClient) IsConnected() bool {
|
|
|
|
|
|
|
|
|
|
// Close 关闭客户端
|
|
|
|
|
func (c *tcpClient) Close() error {
|
|
|
|
|
c.cancel()
|
|
|
|
|
close(c.reconnectStopCh)
|
|
|
|
|
c.closeOnce.Do(func() {
|
|
|
|
|
c.cancel()
|
|
|
|
|
close(c.reconnectStopCh)
|
|
|
|
|
})
|
|
|
|
|
c.mu.Lock()
|
|
|
|
|
defer c.mu.Unlock()
|
|
|
|
|
c.config.AutoReconnect = false // 禁用自动重连
|
|
|
|
|
@ -212,6 +359,7 @@ func (c *tcpClient) Close() error {
|
|
|
|
|
c.conn = nil
|
|
|
|
|
}
|
|
|
|
|
c.connected = false
|
|
|
|
|
c.pending = nil
|
|
|
|
|
return nil
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
@ -294,8 +442,7 @@ func (c *tcpClient) messageReceiver() {
|
|
|
|
|
conn.SetReadDeadline(time.Now().Add(c.config.ReadTimeout))
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
buffer := make([]byte, 4096)
|
|
|
|
|
n, err := conn.Read(buffer)
|
|
|
|
|
msg, err := c.receiveWithTimeout(c.config.ReadTimeout)
|
|
|
|
|
if err != nil {
|
|
|
|
|
// 连接错误，触发重连
|
|
|
|
|
if c.config.AutoReconnect {
|
|
|
|
|
@ -318,17 +465,15 @@ func (c *tcpClient) messageReceiver() {
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
// 发送消息到channel
|
|
|
|
|
message := make([]byte, n)
|
|
|
|
|
copy(message, buffer[:n])
|
|
|
|
|
select {
|
|
|
|
|
case c.messageCh <- message:
|
|
|
|
|
case c.messageCh <- msg:
|
|
|
|
|
default:
|
|
|
|
|
// channel已满，丢弃消息
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
// 调用回调
|
|
|
|
|
if c.onMessage != nil {
|
|
|
|
|
c.onMessage(message)
|
|
|
|
|
c.onMessage(msg)
|
|
|
|
|
}
|
|
|
|
|
}
|
|
|
|
|
}
|
|
|
|
|
@ -357,4 +502,3 @@ func (c *tcpClient) SetOnError(fn func(error)) {
|
|
|
|
|
func (c *tcpClient) MessageChannel() <-chan []byte {
|
|
|
|
|
return c.messageCh
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
