第一次提交

2026-01-30 15:51:42 +08:00
commit 7218ba6508
26 changed files with 3461 additions and 0 deletions
@@ -0,0 +1,47 @@
+package tcp
+
+import (
+	"fmt"
+	"time"
+)
+
+// Example 展示如何使用TCP服务
+func Example() {
+	// 创建配置
+	config := &TcpPoolConfig{
+		BufferSize:     2048,
+		MaxConnections: 100000,
+		ConnectTimeout: time.Second * 5,
+		ReadTimeout:    time.Second * 30,
+		WriteTimeout:   time.Second * 10,
+		MaxIdleTime:    time.Minute * 5,
+	}
+
+	// 创建TCP服务器
+	server := NewTCPServer("0.0.0.0:8888", config)
+
+	// 设置消息处理函数
+	server.SetMessageHandler(func(conn *TcpConnection, msg *TcpMessage) error {
+		fmt.Printf("Received message from %s: %s\n", conn.Id, string(msg.Data))
+
+		// 回显消息
+		return server.SendTo(conn.Id, []byte(fmt.Sprintf("Echo: %s", msg.Data)))
+	})
+
+	// 启动服务器
+	if err := server.Start(); err != nil {
+		fmt.Printf("Failed to start server: %v\n", err)
+		return
+	}
+
+	// 运行10秒后停止
+	fmt.Println("TCP server started. Running for 10 seconds...")
+	time.Sleep(time.Second * 10)
+
+	// 停止服务器
+	if err := server.Stop(); err != nil {
+		fmt.Printf("Failed to stop server: %v\n", err)
+	}
+
+	fmt.Println("TCP server stopped.")
+}
@@ -0,0 +1,280 @@
+package tcp
+
+import (
+	"context"
+	"fmt"
+	"sync"
+	"time"
+
+	"github.com/gogf/gf/v2/frame/g"
+	"github.com/gogf/gf/v2/net/gtcp"
+	"github.com/gogf/gf/v2/os/glog"
+	"github.com/gogf/gf/v2/os/grpool"
+	"github.com/gogf/gf/v2/os/gtime"
+)
+
+// MessageHandler 消息处理函数类型
+type MessageHandler func(conn *TcpConnection, msg *TcpMessage) error
+
+// TCPServer TCP服务器结构
+type TCPServer struct {
+	Address        string
+	Config         *TcpPoolConfig
+	Listener       *gtcp.Server
+	Connection     *ConnectionPool
+	Logger         *glog.Logger
+	MessageHandler MessageHandler
+	ctx            context.Context
+	cancel         context.CancelFunc
+	wg             sync.WaitGroup
+}
+
+// ConnectionPool 连接池结构
+type ConnectionPool struct {
+	connections map[string]*TcpConnection
+	mutex       sync.RWMutex
+	config      *TcpPoolConfig
+	logger      *glog.Logger
+}
+
+// NewTCPServer 创建一个新的TCP服务器
+func NewTCPServer(address string, config *TcpPoolConfig) *TCPServer {
+	logger := g.Log(address)
+	ctx, cancel := context.WithCancel(context.Background())
+
+	pool := &ConnectionPool{
+		connections: make(map[string]*TcpConnection),
+		config:      config,
+		logger:      logger,
+	}
+
+	server := &TCPServer{
+		Address:    address,
+		Config:     config,
+		Connection: pool,
+		Logger:     logger,
+		ctx:        ctx,
+		cancel:     cancel,
+	}
+
+	server.Listener = gtcp.NewServer(address, server.handleConnection)
+	return server
+}
+
+// SetMessageHandler 设置消息处理函数
+func (s *TCPServer) SetMessageHandler(handler MessageHandler) {
+	s.MessageHandler = handler
+}
+
+// Start 启动TCP服务器
+func (s *TCPServer) Start() error {
+	s.Logger.Info(s.ctx, fmt.Sprintf("TCP server starting on %s", s.Address))
+	go func() {
+		s.wg.Add(1)
+		defer s.wg.Done()
+		if err := s.Listener.Run(); err != nil {
+			s.Logger.Error(s.ctx, fmt.Sprintf("TCP server stopped with error: %v", err))
+		}
+	}()
+	return nil
+}
+
+// Stop 停止TCP服务器
+func (s *TCPServer) Stop() error {
+	s.Logger.Info(s.ctx, "TCP server stopping...")
+	s.cancel()
+	s.Listener.Close()
+	s.wg.Wait()
+	s.Connection.Clear()
+	s.Logger.Info(s.ctx, "TCP server stopped")
+	return nil
+}
+
+// handleConnection 处理新连接
+func (s *TCPServer) handleConnection(conn *gtcp.Conn) {
+	// 生成连接ID
+	connID := fmt.Sprintf("%s_%d", conn.RemoteAddr().String(), gtime.TimestampNano())
+
+	// 创建连接对象
+	tcpConn := &TcpConnection{
+		Id:        connID,
+		Address:   conn.RemoteAddr().String(),
+		Server:    *conn,
+		IsActive:  true,
+		LastUsed:  time.Now(),
+		CreatedAt: time.Now(),
+	}
+
+	// 添加到连接池
+	s.Connection.Add(tcpConn)
+	s.Logger.Info(s.ctx, fmt.Sprintf("New connection established: %s", connID))
+
+	// 启动消息接收协程
+	go s.receiveMessages(tcpConn)
+}
+
+// receiveMessages 接收消息
+func (s *TCPServer) receiveMessages(conn *TcpConnection) {
+	defer func() {
+		if err := recover(); err != nil {
+			s.Logger.Error(s.ctx, fmt.Sprintf("Panic in receiveMessages: %v", err))
+		}
+		s.Connection.Remove(conn.Id)
+		conn.Server.Close()
+		s.Logger.Info(s.ctx, fmt.Sprintf("Connection closed: %s", conn.Id))
+	}()
+
+	buffer := make([]byte, s.Config.BufferSize)
+	for {
+		select {
+		case <-s.ctx.Done():
+			return
+		default:
+			// 设置读取超时
+			conn.Server.SetReadDeadline(time.Now().Add(s.Config.ReadTimeout))
+
+			// 读取数据
+			n, err := conn.Server.Read(buffer)
+			if err != nil {
+				s.Logger.Error(s.ctx, fmt.Sprintf("Read error from %s: %v", conn.Id, err))
+				return
+			}
+
+			if n > 0 {
+				// 更新最后使用时间
+				conn.Mutex.Lock()
+				conn.LastUsed = time.Now()
+				conn.Mutex.Unlock()
+
+				// 处理消息
+				data := make([]byte, n)
+				copy(data, buffer[:n])
+
+				msg := &TcpMessage{
+					Id:        fmt.Sprintf("msg_%d", gtime.TimestampNano()),
+					ConnId:    conn.Id,
+					Data:      data,
+					Timestamp: time.Now(),
+					IsSend:    false,
+				}
+
+				// 使用协程池处理消息，避免阻塞
+				grpool.AddWithRecover(s.ctx, func(ctx context.Context) {
+					if s.MessageHandler != nil {
+						if err := s.MessageHandler(conn, msg); err != nil {
+							s.Logger.Error(s.ctx, fmt.Sprintf("Message handling error: %v", err))
+						}
+					}
+				}, func(ctx context.Context, err error) {
+					s.Logger.Error(ctx, fmt.Sprintf("Message handling error: %v", err))
+				})
+			}
+		}
+	}
+}
+
+// SendTo 发送消息到指定连接
+func (s *TCPServer) SendTo(connID string, data []byte) error {
+	conn := s.Connection.Get(connID)
+	if conn == nil {
+		return fmt.Errorf("connection not found: %s", connID)
+	}
+	return s.sendMessage(conn, data)
+}
+
+// SendToAll 发送消息到所有连接
+func (s *TCPServer) SendToAll(data []byte) error {
+	conns := s.Connection.GetAll()
+	for _, conn := range conns {
+		if err := s.sendMessage(conn, data); err != nil {
+			s.Logger.Error(s.ctx, fmt.Sprintf("Send to %s failed: %v", conn.Id, err))
+			// 继续发送给其他连接
+		}
+	}
+	return nil
+}
+
+// sendMessage 发送消息
+func (s *TCPServer) sendMessage(conn *TcpConnection, data []byte) error {
+	conn.Mutex.Lock()
+	defer conn.Mutex.Unlock()
+
+	// 设置写入超时
+	conn.Server.SetWriteDeadline(time.Now().Add(s.Config.WriteTimeout))
+
+	// 发送数据
+	_, err := conn.Server.Write(data)
+	if err != nil {
+		return err
+	}
+
+	// 更新最后使用时间
+	conn.LastUsed = time.Now()
+	return nil
+}
+
+// Kick 强制退出客户端
+func (s *TCPServer) Kick(connID string) error {
+	conn := s.Connection.Get(connID)
+	if conn == nil {
+		return fmt.Errorf("connection not found: %s", connID)
+	}
+
+	// 关闭连接
+	conn.Server.Close()
+	// 从连接池移除
+	s.Connection.Remove(connID)
+
+	s.Logger.Info(s.ctx, fmt.Sprintf("Kicked connection: %s", connID))
+	return nil
+}
+
+// Add 添加连接到连接池
+func (p *ConnectionPool) Add(conn *TcpConnection) {
+	p.mutex.Lock()
+	defer p.mutex.Unlock()
+	p.connections[conn.Id] = conn
+}
+
+// Get 获取连接
+func (p *ConnectionPool) Get(connID string) *TcpConnection {
+	p.mutex.RLock()
+	defer p.mutex.RUnlock()
+	return p.connections[connID]
+}
+
+// GetAll 获取所有连接
+func (p *ConnectionPool) GetAll() []*TcpConnection {
+	p.mutex.RLock()
+	defer p.mutex.RUnlock()
+
+	conns := make([]*TcpConnection, 0, len(p.connections))
+	for _, conn := range p.connections {
+		conns = append(conns, conn)
+	}
+	return conns
+}
+
+// Remove 从连接池移除连接
+func (p *ConnectionPool) Remove(connID string) {
+	p.mutex.Lock()
+	defer p.mutex.Unlock()
+	delete(p.connections, connID)
+}
+
+// Clear 清空连接池
+func (p *ConnectionPool) Clear() {
+	p.mutex.Lock()
+	defer p.mutex.Unlock()
+	for connID, conn := range p.connections {
+		conn.Server.Close()
+		delete(p.connections, connID)
+	}
+}
+
+// Count 获取连接数量
+func (p *ConnectionPool) Count() int {
+	p.mutex.RLock()
+	defer p.mutex.RUnlock()
+	return len(p.connections)
+}
@@ -0,0 +1,38 @@
+package tcp
+
+import (
+	"sync"
+	"time"
+
+	"github.com/gogf/gf/v2/net/gtcp"
+)
+
+// TcpPoolConfig TCP连接池配置
+type TcpPoolConfig struct {
+	BufferSize     int           `json:"bufferSize"`     // 缓冲区大小
+	MaxConnections int           `json:"maxConnections"` // 最大连接数
+	ConnectTimeout time.Duration `json:"connectTimeout"` // 连接超时时间
+	ReadTimeout    time.Duration `json:"readTimeout"`    // 读取超时时间
+	WriteTimeout   time.Duration `json:"writeTimeout"`   // 写入超时时间
+	MaxIdleTime    time.Duration `json:"maxIdleTime"`    // 最大空闲时间
+}
+
+// TcpConnection TCP连接结构
+type TcpConnection struct {
+	Id        string       `json:"id"`        // 连接ID
+	Address   string       `json:"address"`   // 连接地址
+	Server    gtcp.Conn    `json:"server"`    // 实际连接
+	IsActive  bool         `json:"isActive"`  // 是否活跃
+	LastUsed  time.Time    `json:"lastUsed"`  // 最后使用时间
+	CreatedAt time.Time    `json:"createdAt"` // 创建时间
+	Mutex     sync.RWMutex `json:"-"`         // 读写锁
+}
+
+// TcpMessage TCP消息结构
+type TcpMessage struct {
+	Id        string    `json:"id"`        // 消息ID
+	ConnId    string    `json:"connId"`    // 连接ID
+	Data      []byte    `json:"data"`      // 消息数据
+	Timestamp time.Time `json:"timestamp"` // 时间戳
+	IsSend    bool      `json:"isSend"`    // 是否是发送的消息
+}