From 942eff81fb8b33207b9f4df919e6fdfdab723c29 Mon Sep 17 00:00:00 2001
From: black1552 <black1552@163.com>
Date: Fri, 27 Feb 2026 10:15:02 +0800
Subject: [PATCH] =?UTF-8?q?feat(pool):=20=E6=B7=BB=E5=8A=A0=E8=8E=B7?=
 =?UTF-8?q?=E5=8F=96=E6=89=80=E6=9C=89=E5=9C=A8=E7=BA=BF=E8=BF=9E=E6=8E=A5?=
 =?UTF-8?q?ID=E7=9A=84=E5=8A=9F=E8=83=BD?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

- 实现了 GetAllConnIDs 方法来获取在线连接的ID列表
- 添加了内存缓存检查以提高性能
- 实现了从 BadgerDB 迭代获取所有连接记录的逻辑
- 添加了连接活跃状态过滤功能
- 实现了 JSON 反序列化连接信息的处理
- 添加了错误处理和日志记录机制
---
 pool/badger.go | 45 +++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 45 insertions(+)

diff --git a/pool/badger.go b/pool/badger.go
index 9438a0b..16ac8c7 100644
--- a/pool/badger.go
+++ b/pool/badger.go
@@ -215,6 +215,51 @@ func (p *BadgerPool) GetAll() ([]*ConnectionInfo, error) {
 	return conns, 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 connections: %w", err)
+	}
+	return ids, nil
+}
+
 // GetByType 根据类型获取连接
 func (p *BadgerPool) GetByType(connType ConnType) ([]*ConnectionInfo, error) {
 	allConns, err := p.GetAll()