Skip to content

Commit efbabe8

Browse files
xOSxOS
committed
Update.
1 parent c85e4d3 commit efbabe8

File tree

5 files changed

+102
-105
lines changed

5 files changed

+102
-105
lines changed

‎cmd/dashboard/main.go‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@ import (
1414
flag "github.com/spf13/pflag"
1515
"github.com/xos/serverstatus/cmd/dashboard/controller"
1616
"github.com/xos/serverstatus/cmd/dashboard/rpc"
17+
"github.com/xos/serverstatus/db"
1718
"github.com/xos/serverstatus/model"
1819
"github.com/xos/serverstatus/proto"
1920
"github.com/xos/serverstatus/service/singleton"
@@ -111,7 +112,7 @@ func main() {
111112
}
112113

113114
log.Printf("BadgerDB将使用目录: %s", badgerPath)
114-
singleton.InitBadgerDBFromPath(badgerPath)
115+
db.InitBadgerDBFromPath(badgerPath)
115116
} else {
116117
// 默认使用SQLite
117118
log.Println("使用SQLite数据库...")

service/singleton/badger_adapter.go renamed to db/badger_init.go

Lines changed: 75 additions & 96 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
package singleton
1+
package db
22

33
import (
44
"encoding/json"
@@ -11,7 +11,6 @@ import (
1111
"strings"
1212
"time"
1313

14-
"github.com/xos/serverstatus/db"
1514
"github.com/xos/serverstatus/model"
1615
"github.com/xos/serverstatus/pkg/utils"
1716
)
@@ -25,7 +24,7 @@ func InitBadgerDBFromPath(path string) error {
2524
}
2625

2726
// 打开BadgerDB
28-
_, err := db.OpenDB(path)
27+
_, err := OpenDB(path)
2928
if err != nil {
3029
return fmt.Errorf("打开BadgerDB失败: %w", err)
3130
}
@@ -35,7 +34,7 @@ func InitBadgerDBFromPath(path string) error {
3534

3635
// 初始化BadgerDB（不进行自动迁移）
3736
log.Println("初始化BadgerDB...")
38-
if err := db.InitializeEmptyBadgerDB(); err != nil {
37+
if err := InitializeEmptyBadgerDB(); err != nil {
3938
return fmt.Errorf("初始化BadgerDB失败: %w", err)
4039
}
4140

@@ -52,27 +51,65 @@ func InitBadgerDBFromPath(path string) error {
5251
return nil
5352
}
5453

54+
// GetBadgerDBPath 获取BadgerDB数据库目录路径
55+
func GetBadgerDBPath(conf interface{}) string {
56+
// 默认BadgerDB路径
57+
defaultPath := "data/badger"
58+
59+
// 由于这个函数现在在 db 包中，我们需要通过参数传入配置
60+
// 或者从环境变量中获取配置
61+
if confMap, ok := conf.(map[string]interface{}); ok {
62+
if dbLocation, exists := confMap["DatabaseLocation"]; exists {
63+
if location, ok := dbLocation.(string); ok && location != "" {
64+
// 检查配置的路径是否看起来像SQLite文件路径
65+
if strings.HasSuffix(location, ".db") {
66+
// 如果是SQLite文件路径，转换为BadgerDB目录路径
67+
// 例如：data/sqlite.db -> data/badger
68+
dir := filepath.Dir(location)
69+
return filepath.Join(dir, "badger")
70+
} else {
71+
// 如果不是.db文件，假设它是目录路径，直接使用
72+
return location
73+
}
74+
}
75+
}
76+
}
77+
78+
return defaultPath
79+
}
80+
5581
// isSQLiteFileExists 检查SQLite数据库文件是否存在
56-
func isSQLiteFileExists() bool {
82+
func isSQLiteFileExists(conf interface{}) bool {
5783
// 获取SQLite数据库路径
58-
sqlitePath := getSQLitePath()
84+
sqlitePath := getSQLitePath(conf)
5985

6086
_, err := os.Stat(sqlitePath)
6187
return err == nil
6288
}
6389

6490
// getSQLitePath 获取SQLite数据库文件路径
65-
func getSQLitePath() string {
91+
func getSQLitePath(conf interface{}) string {
6692
// 默认路径
6793
defaultPath := "data/sqlite.db"
6894

69-
if Conf == nil {
95+
if conf == nil {
7096
return defaultPath
7197
}
7298

99+
// 从配置中获取数据库信息
100+
var dbType, dbLocation string
101+
if confMap, ok := conf.(map[string]interface{}); ok {
102+
if t, exists := confMap["DatabaseType"]; exists {
103+
dbType, _ = t.(string)
104+
}
105+
if l, exists := confMap["DatabaseLocation"]; exists {
106+
dbLocation, _ = l.(string)
107+
}
108+
}
109+
73110
// 如果当前数据库类型不是badger，直接使用配置的路径
74-
if Conf.DatabaseType != "badger" && Conf.DatabaseLocation != "" {
75-
return Conf.DatabaseLocation
111+
if dbType != "badger" && dbLocation != "" {
112+
return dbLocation
76113
}
77114

78115
// 如果当前是badger类型，需要推断SQLite的路径
@@ -85,13 +122,13 @@ func getSQLitePath() string {
85122

86123
// 如果配置了DatabaseLocation，也尝试将其作为SQLite路径
87124
// （可能用户之前配置的是SQLite路径，后来改为badger）
88-
if Conf.DatabaseLocation != "" {
125+
if dbLocation != "" {
89126
// 如果配置的路径看起来像SQLite文件，也加入候选
90-
if strings.HasSuffix(Conf.DatabaseLocation, ".db") {
91-
possiblePaths = append([]string{Conf.DatabaseLocation}, possiblePaths...)
127+
if strings.HasSuffix(dbLocation, ".db") {
128+
possiblePaths = append([]string{dbLocation}, possiblePaths...)
92129
} else {
93130
// 如果是目录，尝试在该目录下查找sqlite.db
94-
sqliteInConfigDir := filepath.Join(Conf.DatabaseLocation, "sqlite.db")
131+
sqliteInConfigDir := filepath.Join(dbLocation, "sqlite.db")
95132
possiblePaths = append([]string{sqliteInConfigDir}, possiblePaths...)
96133
}
97134
}
@@ -108,36 +145,10 @@ func getSQLitePath() string {
108145
return defaultPath
109146
}
110147

111-
// GetBadgerDBPath 获取BadgerDB数据库目录路径
112-
func GetBadgerDBPath() string {
113-
// 默认BadgerDB路径
114-
defaultPath := "data/badger"
115-
116-
if Conf == nil {
117-
return defaultPath
118-
}
119-
120-
// 如果配置了DatabaseLocation
121-
if Conf.DatabaseLocation != "" {
122-
// 检查配置的路径是否看起来像SQLite文件路径
123-
if strings.HasSuffix(Conf.DatabaseLocation, ".db") {
124-
// 如果是SQLite文件路径，转换为BadgerDB目录路径
125-
// 例如：data/sqlite.db -> data/badger
126-
dir := filepath.Dir(Conf.DatabaseLocation)
127-
return filepath.Join(dir, "badger")
128-
} else {
129-
// 如果不是.db文件，假设它是目录路径，直接使用
130-
return Conf.DatabaseLocation
131-
}
132-
}
133-
134-
return defaultPath
135-
}
136-
137148
// renameSQLiteFileToBackup 将SQLite数据库文件重命名为备份
138-
func renameSQLiteFileToBackup() error {
149+
func renameSQLiteFileToBackup(conf interface{}) error {
139150
// 获取SQLite数据库路径
140-
sqlitePath := getSQLitePath()
151+
sqlitePath := getSQLitePath(conf)
141152

142153
// 检查文件是否存在
143154
if _, err := os.Stat(sqlitePath); os.IsNotExist(err) {
@@ -151,14 +162,14 @@ func renameSQLiteFileToBackup() error {
151162
}
152163

153164
// migrateFromSQLiteToBadgerDB 从SQLite迁移数据到BadgerDB
154-
func migrateFromSQLiteToBadgerDB(badgerPath string) error {
165+
func migrateFromSQLiteToBadgerDB(badgerPath string, conf interface{}) error {
155166
// 获取SQLite数据库文件路径
156-
sqlitePath := getSQLitePath()
167+
sqlitePath := getSQLitePath(conf)
157168

158169
log.Printf("从SQLite路径迁移: %s 到 BadgerDB路径: %s", sqlitePath, badgerPath)
159170

160171
// 创建迁移实例
161-
migration, err := db.NewMigration(db.DB, sqlitePath)
172+
migration, err := NewMigration(DB, sqlitePath)
162173
if err != nil {
163174
return err
164175
}
@@ -175,31 +186,25 @@ func initializeBadgerDBIndexes() error {
175186
return nil
176187
}
177188

178-
// loadServersFromBadgerDB 从BadgerDB加载服务器列表到内存
179-
func loadServersFromBadgerDB() error {
189+
// LoadServersFromBadgerDB 从BadgerDB加载服务器列表到内存
190+
func LoadServersFromBadgerDB() ([]*model.Server, error) {
180191
log.Println("从BadgerDB加载服务器列表...")
181-
ServerLock.Lock()
182-
defer ServerLock.Unlock()
183-
184-
// 清空现有服务器列表
185-
ServerList = make(map[uint64]*model.Server)
186-
SortedServerList = make([]*model.Server, 0)
187-
SortedServerListForGuest = make([]*model.Server, 0)
188192

189193
// 使用ServerOps获取所有服务器
190-
serverOps := db.NewServerOps(db.DB)
191-
if serverOps == nil || db.DB == nil {
194+
serverOps := NewServerOps(DB)
195+
if serverOps == nil || DB == nil {
192196
log.Println("警告: BadgerDB或ServerOps未初始化")
193-
return errors.New("BadgerDB未初始化")
197+
return nil, errors.New("BadgerDB未初始化")
194198
}
195199

196200
servers, err := serverOps.GetAllServers()
197201
if err != nil {
198202
log.Printf("从BadgerDB获取服务器列表失败: %v", err)
199-
return err
203+
return nil, err
200204
}
201205

202-
// 初始化服务器列表
206+
// 处理服务器列表
207+
var processedServers []*model.Server
203208
for _, server := range servers {
204209
if server == nil {
205210
log.Println("警告: 跳过空的服务器记录")
@@ -277,7 +282,6 @@ func loadServersFromBadgerDB() error {
277282
log.Printf("为服务器 %s (ID: %d) 生成新Secret: %s", server.Name, server.ID, newSecret)
278283

279284
// 保存到数据库
280-
serverOps := db.NewServerOps(db.DB)
281285
if err := serverOps.SaveServer(server); err != nil {
282286
log.Printf("保存服务器 %s (ID: %d) 的新Secret到BadgerDB失败: %v", server.Name, server.ID, err)
283287
} else {
@@ -286,41 +290,18 @@ func loadServersFromBadgerDB() error {
286290
}
287291
}
288292

289-
// 添加到服务器映射
290-
ServerList[server.ID] = server
291-
}
292-
293-
// 构建排序后的服务器列表
294-
log.Println("构建排序的服务器列表...")
295-
// 刷新内存中的有序服务器列表
296-
SortedServerLock.Lock()
297-
defer SortedServerLock.Unlock()
298-
SortedServerList = []*model.Server{}
299-
SortedServerListForGuest = []*model.Server{}
300-
301-
for _, s := range ServerList {
302-
SortedServerList = append(SortedServerList, s)
303-
if !s.HideForGuest {
304-
SortedServerListForGuest = append(SortedServerListForGuest, s)
305-
}
293+
processedServers = append(processedServers, server)
306294
}
307295

308296
// 按照服务器排序值排序
309-
sort.SliceStable(SortedServerList, func(i, j int) bool {
310-
if SortedServerList[i].DisplayIndex == SortedServerList[j].DisplayIndex {
311-
return SortedServerList[i].ID < SortedServerList[j].ID
297+
sort.SliceStable(processedServers, func(i, j int) bool {
298+
if processedServers[i].DisplayIndex == processedServers[j].DisplayIndex {
299+
return processedServers[i].ID < processedServers[j].ID
312300
}
313-
return SortedServerList[i].DisplayIndex < SortedServerList[j].DisplayIndex
301+
return processedServers[i].DisplayIndex < processedServers[j].DisplayIndex
314302
})
315303

316-
sort.SliceStable(SortedServerListForGuest, func(i, j int) bool {
317-
if SortedServerListForGuest[i].DisplayIndex == SortedServerListForGuest[j].DisplayIndex {
318-
return SortedServerListForGuest[i].ID < SortedServerListForGuest[j].ID
319-
}
320-
return SortedServerListForGuest[i].DisplayIndex < SortedServerListForGuest[j].DisplayIndex
321-
})
322-
323-
return nil
304+
return processedServers, nil
324305
}
325306

326307
// startBadgerDBMaintenanceTasks 启动BadgerDB的后台维护任务
@@ -337,7 +318,7 @@ func startBadgerDBMaintenanceTasks() {
337318
select {
338319
case <-ticker.C:
339320
// 清理30天前的监控历史记录
340-
monitorHistoryOps := db.NewMonitorHistoryOps(db.DB)
321+
monitorHistoryOps := NewMonitorHistoryOps(DB)
341322
count, err := monitorHistoryOps.CleanupOldMonitorHistories(30 * 24 * time.Hour)
342323
if err != nil {
343324
log.Printf("清理过期监控历史记录失败: %v", err)
@@ -349,11 +330,9 @@ func startBadgerDBMaintenanceTasks() {
349330
}()
350331
}
351332

352-
// 以下是替代原有数据库操作的函数
353-
354333
// SaveServerToBadgerDB 保存服务器信息到BadgerDB
355334
func SaveServerToBadgerDB(server *model.Server) error {
356-
serverOps := db.NewServerOps(db.DB)
335+
serverOps := NewServerOps(DB)
357336
return serverOps.SaveServer(server)
358337
}
359338

@@ -371,13 +350,13 @@ func GetBadgerDBStats() map[string]interface{} {
371350

372351
// CleanBadgerDBMonitorHistory 清理BadgerDB中的监控历史记录
373352
func CleanBadgerDBMonitorHistory(days int) (int, error) {
374-
monitorHistoryOps := db.NewMonitorHistoryOps(db.DB)
353+
monitorHistoryOps := NewMonitorHistoryOps(DB)
375354
return monitorHistoryOps.CleanupOldMonitorHistories(time.Duration(days) * 24 * time.Hour)
376355
}
377356

378357
// CloseBadgerDB 关闭BadgerDB
379358
func CloseBadgerDB() {
380-
if db.DB != nil {
381-
db.DB.Close()
359+
if DB != nil {
360+
DB.Close()
382361
}
383362
}

‎service/singleton/alertsentinel.go‎

Lines changed: 19 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -394,9 +394,6 @@ func UpdateTrafficStats(serverID uint64, inTransfer, outTransfer uint64) {
394394
return
395395
}
396396

397-
// 流量总计
398-
totalTransfer := inTransfer + outTransfer
399-
400397
// 遍历所有报警规则，只更新包含此服务器的规则
401398
for _, alert := range Alerts {
402399
if !alert.Enabled() {
@@ -425,9 +422,25 @@ func UpdateTrafficStats(serverID uint64, inTransfer, outTransfer uint64) {
425422
}
426423
}
427424

428-
// 服务器在规则监控范围内，更新流量数据
425+
// 服务器在规则监控范围内，根据规则类型更新相应的流量数据
426+
var transferValue uint64
427+
switch alert.Rules[j].Type {
428+
case "transfer_in_cycle":
429+
// 只计算入站流量
430+
transferValue = inTransfer
431+
case "transfer_out_cycle":
432+
// 只计算出站流量
433+
transferValue = outTransfer
434+
case "transfer_all_cycle":
435+
// 计算总流量
436+
transferValue = inTransfer + outTransfer
437+
default:
438+
// 默认使用总流量（向后兼容）
439+
transferValue = inTransfer + outTransfer
440+
}
441+
429442
// 不论大小如何，总是更新最新值
430-
stats.Transfer[serverID] = totalTransfer
443+
stats.Transfer[serverID] = transferValue
431444

432445
// 更新服务器名称
433446
if serverName != "" {
@@ -452,7 +465,7 @@ func UpdateTrafficStats(serverID uint64, inTransfer, outTransfer uint64) {
452465
serverCopy.Host = server.Host
453466
}
454467
ServerLock.RUnlock()
455-
checkTrafficThresholds(alert, serverCopy, &alert.Rules[j], totalTransfer)
468+
checkTrafficThresholds(alert, serverCopy, &alert.Rules[j], transferValue)
456469
} else {
457470
ServerLock.RUnlock()
458471
}

‎service/singleton/server.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -192,7 +192,7 @@ func loadServers() {
192192

193193
// 保存到数据库
194194
if Conf.DatabaseType == "badger" {
195-
if err := SaveServerToBadgerDB(innerS); err != nil {
195+
if err := db.SaveServerToBadgerDB(innerS); err != nil {
196196
log.Printf("保存服务器 %s (ID: %d) 的新Secret到BadgerDB失败: %v", innerS.Name, innerS.ID, err)
197197
}
198198
} else {

‎service/singleton/singleton.go‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -82,7 +82,11 @@ func LoadSingleton() {
8282
var err error
8383

8484
// 获取正确的BadgerDB目录路径
85-
badgerPath := GetBadgerDBPath()
85+
confMap := map[string]interface{}{
86+
"DatabaseType": Conf.DatabaseType,
87+
"DatabaseLocation": Conf.DatabaseLocation,
88+
}
89+
badgerPath := db.GetBadgerDBPath(confMap)
8690
log.Printf("使用BadgerDB路径: %s", badgerPath)
8791

8892
db.DB, err = db.OpenDB(badgerPath)

0 commit comments

Comments
 (0)