|
|
package client
|
|
|
|
|
|
import (
|
|
|
"context"
|
|
|
"io"
|
|
|
"net"
|
|
|
"sync"
|
|
|
"time"
|
|
|
|
|
|
"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客户端实现
|
|
|
type tcpClient struct {
|
|
|
config *client.Config
|
|
|
conn net.Conn
|
|
|
mu sync.RWMutex
|
|
|
connected bool
|
|
|
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
|
|
|
reconnectCh chan struct{}
|
|
|
reconnectStopCh chan struct{}
|
|
|
|
|
|
// 异步消息推送
|
|
|
messageCh chan []byte
|
|
|
messageErrCh chan error
|
|
|
onMessage func([]byte)
|
|
|
onError func(error)
|
|
|
|
|
|
// 控制是否已启动消息接收协程（仅在需要时启动，避免与同步Receive竞争读取）
|
|
|
receiverStarted bool
|
|
|
}
|
|
|
|
|
|
// NewTCPClient 创建TCP客户端
|
|
|
func NewTCPClient(config *client.Config) client.Client {
|
|
|
if config == nil {
|
|
|
config = client.DefaultConfig()
|
|
|
}
|
|
|
|
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
|
protocol := newApplicationProtocol(config.ApplicationProtocol)
|
|
|
|
|
|
c := &tcpClient{
|
|
|
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),
|
|
|
}
|
|
|
|
|
|
// 如果启用自动重连，启动重连goroutine
|
|
|
if config.AutoReconnect {
|
|
|
go c.reconnectLoop()
|
|
|
}
|
|
|
|
|
|
// 默认不启动消息接收goroutine，避免与同步Receive竞争读取
|
|
|
// 仅当设置了消息回调或显式需要异步接收时再启动
|
|
|
|
|
|
return c
|
|
|
}
|
|
|
|
|
|
// Connect 连接服务器
|
|
|
func (c *tcpClient) Connect() error {
|
|
|
return c.connect()
|
|
|
}
|
|
|
|
|
|
// connect 内部连接方法
|
|
|
func (c *tcpClient) connect() error {
|
|
|
c.mu.Lock()
|
|
|
defer c.mu.Unlock()
|
|
|
|
|
|
if c.connected {
|
|
|
return nil
|
|
|
}
|
|
|
|
|
|
// 解析地址
|
|
|
addr := c.config.Addr
|
|
|
if len(addr) > 6 && addr[:6] == "tcp://" {
|
|
|
addr = addr[6:]
|
|
|
}
|
|
|
|
|
|
// 创建连接
|
|
|
connectTimeout := c.config.ConnectTimeout
|
|
|
if connectTimeout <= 0 {
|
|
|
connectTimeout = 3 * time.Second
|
|
|
}
|
|
|
|
|
|
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{}{}:
|
|
|
default:
|
|
|
}
|
|
|
}
|
|
|
return err
|
|
|
}
|
|
|
|
|
|
c.conn = conn
|
|
|
c.connected = true
|
|
|
c.reconnectAttempts = 0 // 重置重连次数
|
|
|
c.unpacker = newApplicationUnpacker(c.protocol)
|
|
|
c.pending = nil
|
|
|
|
|
|
return nil
|
|
|
}
|
|
|
|
|
|
// Disconnect 断开连接
|
|
|
func (c *tcpClient) Disconnect() error {
|
|
|
c.mu.Lock()
|
|
|
defer c.mu.Unlock()
|
|
|
|
|
|
if !c.connected {
|
|
|
return nil
|
|
|
}
|
|
|
|
|
|
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:
|
|
|
}
|
|
|
}
|
|
|
|
|
|
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 {
|
|
|
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 writeTimeout > 0 {
|
|
|
conn.SetWriteDeadline(time.Now().Add(writeTimeout))
|
|
|
}
|
|
|
|
|
|
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) {
|
|
|
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")
|
|
|
}
|
|
|
conn := c.conn
|
|
|
readTimeout := c.config.ReadTimeout
|
|
|
if timeout > 0 {
|
|
|
readTimeout = timeout
|
|
|
}
|
|
|
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))
|
|
|
}
|
|
|
|
|
|
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
|
|
|
}
|
|
|
|
|
|
// 接收响应
|
|
|
return c.receiveWithTimeout(timeout)
|
|
|
}
|
|
|
|
|
|
// IsConnected 检查是否已连接
|
|
|
func (c *tcpClient) IsConnected() bool {
|
|
|
c.mu.RLock()
|
|
|
defer c.mu.RUnlock()
|
|
|
return c.connected
|
|
|
}
|
|
|
|
|
|
// Close 关闭客户端
|
|
|
func (c *tcpClient) Close() error {
|
|
|
c.closeOnce.Do(func() {
|
|
|
c.cancel()
|
|
|
close(c.reconnectStopCh)
|
|
|
})
|
|
|
c.mu.Lock()
|
|
|
defer c.mu.Unlock()
|
|
|
c.config.AutoReconnect = false // 禁用自动重连
|
|
|
if c.conn != nil {
|
|
|
c.conn.Close()
|
|
|
c.conn = nil
|
|
|
}
|
|
|
c.connected = false
|
|
|
c.pending = nil
|
|
|
return nil
|
|
|
}
|
|
|
|
|
|
// reconnectLoop 重连循环
|
|
|
func (c *tcpClient) reconnectLoop() {
|
|
|
for {
|
|
|
select {
|
|
|
case <-c.ctx.Done():
|
|
|
return
|
|
|
case <-c.reconnectStopCh:
|
|
|
return
|
|
|
case <-c.reconnectCh:
|
|
|
c.mu.RLock()
|
|
|
autoReconnect := c.config.AutoReconnect
|
|
|
connected := c.connected
|
|
|
attempts := c.reconnectAttempts
|
|
|
c.mu.RUnlock()
|
|
|
|
|
|
if !autoReconnect {
|
|
|
return
|
|
|
}
|
|
|
|
|
|
// 如果已经连接，跳过
|
|
|
if connected {
|
|
|
continue
|
|
|
}
|
|
|
|
|
|
// 检查最大重连次数
|
|
|
if c.config.MaxReconnectAttempts > 0 && attempts >= c.config.MaxReconnectAttempts {
|
|
|
if c.onError != nil {
|
|
|
c.onError(client.NewError("max reconnect attempts reached"))
|
|
|
}
|
|
|
return
|
|
|
}
|
|
|
|
|
|
// 等待重连间隔
|
|
|
time.Sleep(c.config.ReconnectInterval)
|
|
|
|
|
|
// 尝试重连
|
|
|
c.mu.Lock()
|
|
|
c.reconnectAttempts++
|
|
|
c.mu.Unlock()
|
|
|
|
|
|
if err := c.connect(); err != nil {
|
|
|
if c.onError != nil {
|
|
|
c.onError(err)
|
|
|
}
|
|
|
// 继续重连（延迟触发，避免立即重试）
|
|
|
go func() {
|
|
|
time.Sleep(1 * time.Second)
|
|
|
select {
|
|
|
case c.reconnectCh <- struct{}{}:
|
|
|
default:
|
|
|
}
|
|
|
}()
|
|
|
}
|
|
|
}
|
|
|
}
|
|
|
}
|
|
|
|
|
|
// messageReceiver 消息接收器（异步）
|
|
|
func (c *tcpClient) messageReceiver() {
|
|
|
for {
|
|
|
select {
|
|
|
case <-c.ctx.Done():
|
|
|
return
|
|
|
default:
|
|
|
c.mu.RLock()
|
|
|
conn := c.conn
|
|
|
connected := c.connected
|
|
|
c.mu.RUnlock()
|
|
|
|
|
|
if !connected || conn == nil {
|
|
|
time.Sleep(100 * time.Millisecond)
|
|
|
continue
|
|
|
}
|
|
|
|
|
|
// 设置读取超时
|
|
|
if c.config.ReadTimeout > 0 {
|
|
|
conn.SetReadDeadline(time.Now().Add(c.config.ReadTimeout))
|
|
|
}
|
|
|
|
|
|
msg, err := c.receiveWithTimeout(c.config.ReadTimeout)
|
|
|
if err != nil {
|
|
|
// 连接错误，触发重连
|
|
|
if c.config.AutoReconnect {
|
|
|
c.mu.Lock()
|
|
|
c.connected = false
|
|
|
if c.conn != nil {
|
|
|
c.conn.Close()
|
|
|
c.conn = nil
|
|
|
}
|
|
|
c.mu.Unlock()
|
|
|
select {
|
|
|
case c.reconnectCh <- struct{}{}:
|
|
|
default:
|
|
|
}
|
|
|
}
|
|
|
if c.onError != nil {
|
|
|
c.onError(err)
|
|
|
}
|
|
|
continue
|
|
|
}
|
|
|
|
|
|
// 发送消息到channel
|
|
|
select {
|
|
|
case c.messageCh <- msg:
|
|
|
default:
|
|
|
// channel已满，丢弃消息
|
|
|
}
|
|
|
|
|
|
// 调用回调
|
|
|
if c.onMessage != nil {
|
|
|
c.onMessage(msg)
|
|
|
}
|
|
|
}
|
|
|
}
|
|
|
}
|
|
|
|
|
|
// SetOnMessage 设置消息回调
|
|
|
func (c *tcpClient) SetOnMessage(fn func([]byte)) {
|
|
|
c.mu.Lock()
|
|
|
defer c.mu.Unlock()
|
|
|
c.onMessage = fn
|
|
|
// 首次设置回调时，启动消息接收协程
|
|
|
if !c.receiverStarted {
|
|
|
go c.messageReceiver()
|
|
|
c.receiverStarted = true
|
|
|
}
|
|
|
}
|
|
|
|
|
|
// SetOnError 设置错误回调
|
|
|
func (c *tcpClient) SetOnError(fn func(error)) {
|
|
|
c.mu.Lock()
|
|
|
defer c.mu.Unlock()
|
|
|
c.onError = fn
|
|
|
}
|
|
|
|
|
|
// MessageChannel 获取消息channel
|
|
|
func (c *tcpClient) MessageChannel() <-chan []byte {
|
|
|
return c.messageCh
|
|
|
}
|