feat(pool): 添加SQLite连接池实现并集成到TCP和WebSocket服务

- 新增pool包，包含ConnType连接类型和ConnectionInfo连接信息结构体 - 实现SQLitePool连接池，支持添加、获取、删除、更新连接操作 - 为TCP服务器集成SQLite连接池，存储连接信息到数据库 - 为WebSocket管理器集成SQLite连接池，存储连接信息到数据库 - 在TCP和WebSocket的连接生命周期中同步更新SQLite连接状态 - 添加GetAllConnIDs方法获取所有在线连接ID列表 - 在示例代码中添加错误处理和测试功能
2026-02-27 17:54:33 +08:00
parent d4f78124d8
commit 2520655fbd
7 changed files with 731 additions and 91 deletions
@@ -0,0 +1,25 @@
+package pool
+
+import (
+	"time"
+)
+
+// ConnType 连接类型
+type ConnType string
+
+const (
+	ConnTypeWebSocket ConnType = "websocket"
+	ConnTypeTCP       ConnType = "tcp"
+)
+
+// ConnectionInfo 连接信息
+type ConnectionInfo struct {
+	ID        string    `json:"id" gorm:"primaryKey"`
+	Type      ConnType  `json:"type" gorm:"index"`
+	Address   string    `json:"address"`
+	IsActive  bool      `json:"isActive" gorm:"index"`
+	LastUsed  time.Time `json:"lastUsed"`
+	CreatedAt time.Time `json:"createdAt"`
+	// 额外的连接数据，根据不同类型存储不同的信息
+	Data map[string]interface{} `json:"data" gorm:"-"`
+}
@@ -0,0 +1,247 @@
+package pool
+
+import (
+	"context"
+	"fmt"
+	"sync"
+	"time"
+
+	"git.magicany.cc/black1552/gf-common/db"
+	"gorm.io/gorm"
+)
+
+// SQLitePool SQLite连接池
+type SQLitePool struct {
+	db     *gorm.DB
+	mutex  sync.RWMutex
+	ctx    context.Context
+	cancel context.CancelFunc
+	// 内存缓存，提高并发性能
+	cache map[string]*ConnectionInfo
+}
+
+// NewSQLitePool 创建SQLite连接池
+func NewSQLitePool() (*SQLitePool, error) {
+	ctx, cancel := context.WithCancel(context.Background())
+
+	// 检查数据库连接是否正常
+	if db.Db == nil {
+		return nil, fmt.Errorf("database connection is not initialized")
+	}
+
+	// 自动迁移ConnectionInfo模型
+	err := db.Db.AutoMigrate(&ConnectionInfo{})
+	if err != nil {
+		cancel()
+		return nil, fmt.Errorf("failed to migrate connection info model: %w", err)
+	}
+
+	return &SQLitePool{
+		db:     db.Db,
+		ctx:    ctx,
+		cancel: cancel,
+		cache:  make(map[string]*ConnectionInfo),
+	}, nil
+}
+
+// Close 关闭连接池
+func (p *SQLitePool) Close() error {
+	p.cancel()
+	// SQLite连接由db包管理，不需要在这里关闭
+	return nil
+}
+
+// Add 添加连接
+func (p *SQLitePool) Add(conn *ConnectionInfo) error {
+	p.mutex.Lock()
+	defer p.mutex.Unlock()
+
+	// 存储到SQLite
+	result := p.db.Create(conn)
+	if result.Error != nil {
+		return fmt.Errorf("failed to store connection: %w", result.Error)
+	}
+
+	// 更新内存缓存
+	p.cache[conn.ID] = conn
+	return nil
+}
+
+// Get 获取连接
+func (p *SQLitePool) Get(connID string) (*ConnectionInfo, error) {
+	p.mutex.RLock()
+	// 先从内存缓存获取
+	if conn, ok := p.cache[connID]; ok {
+		p.mutex.RUnlock()
+		return conn, nil
+	}
+	p.mutex.RUnlock()
+
+	// 从SQLite获取
+	var connInfo ConnectionInfo
+	result := p.db.First(&connInfo, "id = ?", connID)
+	if result.Error != nil {
+		if result.Error == gorm.ErrRecordNotFound {
+			return nil, nil
+		}
+		return nil, fmt.Errorf("failed to get connection: %w", result.Error)
+	}
+
+	// 更新内存缓存
+	p.mutex.Lock()
+	p.cache[connID] = &connInfo
+	p.mutex.Unlock()
+
+	return &connInfo, nil
+}
+
+// Remove 移除连接
+func (p *SQLitePool) Remove(connID string) error {
+	p.mutex.Lock()
+	defer p.mutex.Unlock()
+
+	// 从SQLite删除
+	result := p.db.Delete(&ConnectionInfo{}, "id = ?", connID)
+	if result.Error != nil {
+		return fmt.Errorf("failed to remove connection: %w", result.Error)
+	}
+
+	// 从内存缓存删除
+	delete(p.cache, connID)
+	return nil
+}
+
+// Update 更新连接信息
+func (p *SQLitePool) Update(conn *ConnectionInfo) error {
+	p.mutex.Lock()
+	defer p.mutex.Unlock()
+
+	// 存储到SQLite
+	result := p.db.Save(conn)
+	if result.Error != nil {
+		return fmt.Errorf("failed to update connection: %w", result.Error)
+	}
+
+	// 更新内存缓存
+	p.cache[conn.ID] = conn
+	return nil
+}
+
+// GetAll 获取所有连接
+func (p *SQLitePool) GetAll() ([]*ConnectionInfo, error) {
+	p.mutex.RLock()
+	// 如果内存缓存不为空，直接返回缓存
+	if len(p.cache) > 0 {
+		conns := make([]*ConnectionInfo, 0, len(p.cache))
+		for _, conn := range p.cache {
+			conns = append(conns, conn)
+		}
+		p.mutex.RUnlock()
+		return conns, nil
+	}
+	p.mutex.RUnlock()
+
+	// 从SQLite获取所有连接
+	var conns []*ConnectionInfo
+	result := p.db.Find(&conns)
+	if result.Error != nil {
+		return nil, fmt.Errorf("failed to get all connections: %w", result.Error)
+	}
+
+	// 更新内存缓存
+	p.mutex.Lock()
+	for _, conn := range conns {
+		p.cache[conn.ID] = conn
+	}
+	p.mutex.Unlock()
+
+	return conns, nil
+}
+
+// GetByType 根据类型获取连接
+func (p *SQLitePool) GetByType(connType ConnType) ([]*ConnectionInfo, error) {
+	allConns, err := p.GetAll()
+	if err != nil {
+		return nil, err
+	}
+
+	var filtered []*ConnectionInfo
+	for _, conn := range allConns {
+		if conn.Type == connType {
+			filtered = append(filtered, conn)
+		}
+	}
+
+	return filtered, nil
+}
+
+// Count 获取连接数量
+func (p *SQLitePool) Count() (int, error) {
+	p.mutex.RLock()
+	// 如果内存缓存不为空，直接返回缓存大小
+	if len(p.cache) > 0 {
+		count := len(p.cache)
+		p.mutex.RUnlock()
+		return count, nil
+	}
+	p.mutex.RUnlock()
+
+	// 从SQLite统计数量
+	var count int64
+	result := p.db.Model(&ConnectionInfo{}).Count(&count)
+	if result.Error != nil {
+		return 0, fmt.Errorf("failed to count connections: %w", result.Error)
+	}
+
+	return int(count), nil
+}
+
+// GetAllConnIDs 获取所有在线连接的ID列表
+func (p *SQLitePool) GetAllConnIDs() ([]string, error) {
+	p.mutex.RLock()
+	// 如果内存缓存不为空，从缓存中提取在线连接的ID
+	if len(p.cache) > 0 {
+		ids := make([]string, 0, len(p.cache))
+		for id, conn := range p.cache {
+			if conn.IsActive {
+				ids = append(ids, id)
+			}
+		}
+		p.mutex.RUnlock()
+		return ids, nil
+	}
+	p.mutex.RUnlock()
+
+	// 从SQLite获取所有在线连接的ID
+	var conns []*ConnectionInfo
+	result := p.db.Where("is_active = ?", true).Find(&conns)
+	if result.Error != nil {
+		return nil, fmt.Errorf("failed to get all connection IDs: %w", result.Error)
+	}
+
+	ids := make([]string, 0, len(conns))
+	for _, conn := range conns {
+		ids = append(ids, conn.ID)
+	}
+
+	return ids, nil
+}
+
+// CleanupInactive 清理不活跃的连接
+func (p *SQLitePool) CleanupInactive(duration time.Duration) error {
+	allConns, err := p.GetAll()
+	if err != nil {
+		return err
+	}
+
+	now := time.Now()
+	for _, conn := range allConns {
+		if !conn.IsActive || now.Sub(conn.LastUsed) > duration {
+			if err := p.Remove(conn.ID); err != nil {
+				return err
+			}
+		}
+	}
+
+	return nil
+}