From 4c300374cb62cfe0e7d7cd0518ef06cf53612ebd Mon Sep 17 00:00:00 2001 From: black1552 Date: Fri, 27 Feb 2026 10:06:26 +0800 Subject: [PATCH] =?UTF-8?q?feat(pool):=20=E4=BF=AE=E6=94=B9NewBadgerPool?= =?UTF-8?q?=E5=87=BD=E6=95=B0=E6=8E=A5=E5=8F=97=E6=95=B0=E6=8D=AE=E5=BA=93?= =?UTF-8?q?=E8=B7=AF=E5=BE=84=E5=8F=82=E6=95=B0=E5=B9=B6=E7=A7=BB=E9=99=A4?= =?UTF-8?q?GetAllConnIDs=E6=96=B9=E6=B3=95?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 修改NewBadgerPool函数签名以接受badgerDir字符串参数 - 将默认数据库路径配置改为通过参数传入 - 移除GetAllConnIDs方法及其相关实现代码 - 更新TCPServer和WebSocket Manager构造函数以传递数据库路径 - 修复TCP服务器和WebSocket示例中的初始化调用 --- pool/badger.go | 50 ++---------------------------------------- server/ws/example.go | 2 +- server/ws/websocket.go | 4 ++-- tcp/example.go | 2 +- tcp/tcp.go | 4 ++-- 5 files changed, 8 insertions(+), 54 deletions(-) diff --git a/pool/badger.go b/pool/badger.go index fe5ebee..03e3f3e 100644 --- a/pool/badger.go +++ b/pool/badger.go @@ -43,10 +43,10 @@ type BadgerPool struct { } // NewBadgerPool 创建BadgerDB连接池 -func NewBadgerPool() (*BadgerPool, error) { +func NewBadgerPool(badgerDir string) (*BadgerPool, error) { ctx, cancel := context.WithCancel(context.Background()) - db, err := badger.Open(badger.DefaultOptions(filepath.Join(gfile.Pwd(), "badger"))) + db, err := badger.Open(badger.DefaultOptions(badgerDir)) if err != nil { cancel() return nil, fmt.Errorf("failed to open badger db: %w", err) @@ -265,52 +265,6 @@ func (p *BadgerPool) Count() (int, error) { return count, nil } -// GetAllConnIDs 获取所有在线连接的ID列表 -func (p *BadgerPool) 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() - - // 从BadgerDB获取所有在线连接的ID - var ids []string - err := p.db.View(func(txn *badger.Txn) error { - opts := badger.DefaultIteratorOptions - opts.PrefetchSize = 10 - it := txn.NewIterator(opts) - defer it.Close() - - for it.Rewind(); it.Valid(); it.Next() { - item := it.Item() - var connInfo ConnectionInfo - err := item.Value(func(val []byte) error { - return json.Unmarshal(val, &connInfo) - }) - if err != nil { - return err - } - if connInfo.IsActive { - ids = append(ids, string(item.Key())) - } - } - return nil - }) - if err != nil { - return nil, fmt.Errorf("failed to get all connection IDs: %w", err) - } - - return ids, nil -} - // CleanupInactive 清理不活跃的连接 func (p *BadgerPool) CleanupInactive(duration time.Duration) error { allConns, err := p.GetAll() diff --git a/server/ws/example.go b/server/ws/example.go index 1831f55..7c06f90 100644 --- a/server/ws/example.go +++ b/server/ws/example.go @@ -19,7 +19,7 @@ func NewWs() *Manager { } // 2. 创建管理器 - m, err := NewManager(customConfig) + m, err := NewManager(customConfig, "./badger/ws") if err != nil { log.Fatalf("Failed to create manager: %v", err) } diff --git a/server/ws/websocket.go b/server/ws/websocket.go index 9dcb220..ac9b385 100644 --- a/server/ws/websocket.go +++ b/server/ws/websocket.go @@ -150,7 +150,7 @@ func (c *Config) Merge(other *Config) *Config { } // NewManager 创建连接管理器 -func NewManager(config *Config) (*Manager, error) { +func NewManager(config *Config, dbPath string) (*Manager, error) { defaultConfig := DefaultConfig() finalConfig := defaultConfig.Merge(config) // 初始化升级器 @@ -173,7 +173,7 @@ func NewManager(config *Config) (*Manager, error) { } // 初始化BadgerDB连接池 - badgerPool, err := pool.NewBadgerPool() + badgerPool, err := pool.NewBadgerPool(dbPath) if err != nil { return nil, fmt.Errorf("failed to create badger pool: %w", err) } diff --git a/tcp/example.go b/tcp/example.go index 7101501..6ad89e4 100644 --- a/tcp/example.go +++ b/tcp/example.go @@ -18,7 +18,7 @@ func Example() { } // 创建TCP服务器 - server, err := NewTCPServer("0.0.0.0:8888", config) + server, err := NewTCPServer("0.0.0.0:8888", config, "./badger/tcp") if err != nil { fmt.Printf("Failed to create server: %v\n", err) return diff --git a/tcp/tcp.go b/tcp/tcp.go index c248909..75ddff5 100644 --- a/tcp/tcp.go +++ b/tcp/tcp.go @@ -40,12 +40,12 @@ type ConnectionPool struct { } // NewTCPServer 创建一个新的TCP服务器 -func NewTCPServer(address string, config *TcpPoolConfig) (*TCPServer, error) { +func NewTCPServer(address string, config *TcpPoolConfig, dbPath string) (*TCPServer, error) { logger := g.Log(address) ctx, cancel := context.WithCancel(context.Background()) // 初始化BadgerDB连接池 - badgerPool, err := pool.NewBadgerPool() + badgerPool, err := pool.NewBadgerPool(dbPath) if err != nil { return nil, fmt.Errorf("failed to create badger pool: %w", err) }