Skip to content

Commit 683175b

Browse files
xOSxOS
committed
Update.
1 parent 5f40dd3 commit 683175b

File tree

3 files changed

+167
-97
lines changed

3 files changed

+167
-97
lines changed

‎service/singleton/notification.go‎

Lines changed: 42 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ import (
66
"sync"
77
"time"
88

9+
"github.com/xos/serverstatus/db"
910
"github.com/xos/serverstatus/model"
1011
)
1112

@@ -30,24 +31,58 @@ func loadNotifications() {
3031
notificationsLock.Lock()
3132
defer notificationsLock.Unlock()
3233

33-
var notifications []model.Notification
34-
if err := DB.Find(&notifications).Error; err != nil {
35-
panic(err)
34+
var notifications []*model.Notification
35+
36+
// 根据数据库类型选择不同的加载方式
37+
if Conf.DatabaseType == "badger" {
38+
// 使用 BadgerDB 加载通知
39+
if db.DB != nil {
40+
notificationOps := db.NewNotificationOps(db.DB)
41+
var err error
42+
notifications, err = notificationOps.GetAllNotifications()
43+
if err != nil {
44+
log.Printf("从 BadgerDB 加载通知失败: %v", err)
45+
return
46+
}
47+
} else {
48+
log.Println("警告: BadgerDB 未初始化")
49+
return
50+
}
51+
} else {
52+
// 使用 GORM (SQLite) 加载通知
53+
if err := DB.Find(&notifications).Error; err != nil {
54+
log.Printf("从 SQLite 加载通知失败: %v", err)
55+
return
56+
}
3657
}
58+
3759
for i := 0; i < len(notifications); i++ {
3860
// 旧版本的Tag可能不存在 自动设置为默认值
3961
if notifications[i].Tag == "" {
40-
SetDefaultNotificationTagInDB(&notifications[i])
62+
SetDefaultNotificationTagInDB(notifications[i])
4163
}
42-
AddNotificationToList(&notifications[i])
64+
AddNotificationToList(notifications[i])
4365
}
4466
}
4567

4668
// SetDefaultNotificationTagInDB 设置默认通知方式的 Tag
4769
func SetDefaultNotificationTagInDB(n *model.Notification) {
4870
n.Tag = "default"
49-
if err := DB.Save(n).Error; err != nil {
50-
log.Println("NG>> SetDefaultNotificationTagInDB 错误: ", err)
71+
72+
// 根据数据库类型选择不同的保存方式
73+
if Conf.DatabaseType == "badger" {
74+
// 使用 BadgerDB 保存
75+
if db.DB != nil {
76+
notificationOps := db.NewNotificationOps(db.DB)
77+
if err := notificationOps.SaveNotification(n); err != nil {
78+
log.Println("NG>> SetDefaultNotificationTagInDB 错误 (BadgerDB): ", err)
79+
}
80+
}
81+
} else {
82+
// 使用 GORM 保存
83+
if err := DB.Save(n).Error; err != nil {
84+
log.Println("NG>> SetDefaultNotificationTagInDB 错误 (SQLite): ", err)
85+
}
5186
}
5287
}
5388

‎service/singleton/server.go‎

Lines changed: 75 additions & 45 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ import (
88
"time"
99

1010
"github.com/jinzhu/copier"
11+
"github.com/xos/serverstatus/db"
1112
"github.com/xos/serverstatus/model"
1213
"github.com/xos/serverstatus/pkg/utils"
1314
)
@@ -33,10 +34,39 @@ func InitServer() {
3334
// loadServers 加载服务器列表并根据ID排序
3435
func loadServers() {
3536
InitServer()
36-
var servers []model.Server
37-
DB.Find(&servers)
37+
38+
var servers []*model.Server
39+
40+
// 根据数据库类型选择不同的加载方式
41+
if Conf.DatabaseType == "badger" {
42+
// 使用 BadgerDB 加载服务器
43+
if db.DB != nil {
44+
serverOps := db.NewServerOps(db.DB)
45+
var err error
46+
servers, err = serverOps.GetAllServers()
47+
if err != nil {
48+
log.Printf("从 BadgerDB 加载服务器列表失败: %v", err)
49+
return
50+
}
51+
} else {
52+
log.Println("警告: BadgerDB 未初始化")
53+
return
54+
}
55+
} else {
56+
// 使用 GORM (SQLite) 加载服务器
57+
var sqliteServers []model.Server
58+
DB.Find(&sqliteServers)
59+
for _, s := range sqliteServers {
60+
servers = append(servers, &s)
61+
}
62+
}
63+
3864
for _, s := range servers {
3965
innerS := s
66+
// 如果是指针，不需要再取地址
67+
if innerS == nil {
68+
continue
69+
}
4070

4171
// 初始化基本对象
4272
innerS.State = &model.HostState{}
@@ -51,23 +81,23 @@ func loadServers() {
5181
innerS.LastActive = time.Now().Add(-24 * time.Hour)
5282
}
5383

54-
// 从数据库加载Host信息
55-
var hostJSONStr string
56-
if err := DB.Raw("SELECT host_json FROM servers WHERE id = ?", innerS.ID).Scan(&hostJSONStr).Error; err == nil && len(hostJSONStr) > 0 {
57-
// 创建Host对象并解析数据
84+
// 初始化主机信息
85+
if innerS.Host == nil {
86+
innerS.Host = &model.Host{}
87+
innerS.Host.Initialize()
88+
}
89+
90+
// 解析主机信息
91+
if innerS.HostJSON != "" && len(innerS.HostJSON) > 0 {
5892
host := &model.Host{}
59-
if err := utils.Json.Unmarshal([]byte(hostJSONStr), host); err != nil {
93+
if err := utils.Json.Unmarshal([]byte(innerS.HostJSON), host); err != nil {
6094
log.Printf("解析服务器 %s 的Host数据失败: %v", innerS.Name, err)
6195
// 创建空的Host对象作为后备
6296
host = &model.Host{}
6397
}
6498
// 确保Host对象正确初始化
6599
host.Initialize()
66100
innerS.Host = host
67-
} else {
68-
// 如果数据库中没有Host数据，创建空的Host对象
69-
innerS.Host = &model.Host{}
70-
innerS.Host.Initialize()
71101
}
72102

73103
// 加载离线前的最后状态
@@ -109,7 +139,7 @@ func loadServers() {
109139
}
110140

111141
innerS.TaskCloseLock = new(sync.Mutex)
112-
ServerList[innerS.ID] = &innerS
142+
ServerList[innerS.ID] = innerS
113143
SecretToID[innerS.Secret] = innerS.ID
114144
ServerTagToIDList[innerS.Tag] = append(ServerTagToIDList[innerS.Tag], innerS.ID)
115145
}
@@ -191,11 +221,11 @@ func CleanupServerState() {
191221
ServerLock.RLock()
192222
var serversToCleanup []*model.Server
193223
now := time.Now()
194-
224+
195225
for _, server := range ServerList {
196226
// 需要清理的条件：长时间离线或长时间未活动
197227
if (!server.IsOnline && now.Sub(server.LastActive) > 24*time.Hour) ||
198-
(server.IsOnline && now.Sub(server.LastActive) > 5*time.Minute) {
228+
(server.IsOnline && now.Sub(server.LastActive) > 5*time.Minute) {
199229
serversToCleanup = append(serversToCleanup, server)
200230
}
201231
}
@@ -217,7 +247,7 @@ func CleanupServerState() {
217247
server.State = nil
218248
server.LastStateBeforeOffline = nil
219249
server.TaskStream = nil
220-
250+
221251
// 安全关闭通道，使用异步处理防止死锁
222252
if server.TaskClose != nil && server.TaskCloseLock != nil {
223253
go func(s *model.Server) {
@@ -227,7 +257,7 @@ func CleanupServerState() {
227257
s.TaskCloseLock.Lock()
228258
done <- true
229259
}()
230-
260+
231261
select {
232262
case <-done:
233263
defer s.TaskCloseLock.Unlock()
@@ -260,7 +290,7 @@ func CleanupServerState() {
260290
s.TaskCloseLock.Lock()
261291
done <- true
262292
}()
263-
293+
264294
select {
265295
case <-done:
266296
defer s.TaskCloseLock.Unlock()
@@ -289,63 +319,63 @@ func CleanupServerState() {
289319
// 这个函数专门设计为防死锁，使用超时机制和非阻塞操作
290320
func SafeCleanupServerState() {
291321
log.Printf("开始安全服务器状态清理...")
292-
322+
293323
// 使用带超时的Server锁获取
294324
type lockResult struct {
295325
acquired bool
296326
servers []*model.Server
297327
}
298-
328+
299329
lockCh := make(chan lockResult, 1)
300-
330+
301331
// 尝试获取锁的goroutine
302332
go func() {
303333
result := lockResult{acquired: false}
304-
334+
305335
// 尝试在500ms内获取锁
306336
done := make(chan bool, 1)
307337
go func() {
308338
ServerLock.RLock()
309339
done <- true
310340
}()
311-
341+
312342
select {
313343
case <-done:
314344
// 成功获取锁，快速复制服务器列表
315345
defer ServerLock.RUnlock()
316-
346+
317347
now := time.Now()
318348
for _, server := range ServerList {
319-
if server != nil &&
320-
((!server.IsOnline && now.Sub(server.LastActive) > 12*time.Hour) ||
321-
(server.IsOnline && now.Sub(server.LastActive) > 3*time.Minute)) {
349+
if server != nil &&
350+
((!server.IsOnline && now.Sub(server.LastActive) > 12*time.Hour) ||
351+
(server.IsOnline && now.Sub(server.LastActive) > 3*time.Minute)) {
322352
result.servers = append(result.servers, server)
323353
}
324354
}
325355
result.acquired = true
326-
356+
327357
case <-time.After(500 * time.Millisecond):
328358
// 超时，跳过这次清理
329359
log.Printf("SafeCleanupServerState: 获取ServerLock超时，跳过清理")
330360
}
331-
361+
332362
lockCh <- result
333363
}()
334-
364+
335365
// 等待锁获取结果
336366
select {
337367
case result := <-lockCh:
338368
if !result.acquired {
339369
log.Printf("SafeCleanupServerState: 无法获取锁，清理跳过")
340370
return
341371
}
342-
372+
343373
log.Printf("SafeCleanupServerState: 发现 %d 个需要清理的服务器", len(result.servers))
344-
374+
345375
// 同步清理每个服务器，使用超时保护
346376
cleanedCount := 0
347377
skippedCount := 0
348-
378+
349379
for _, server := range result.servers {
350380
cleaned := cleanupSingleServerState(server)
351381
if cleaned {
@@ -354,9 +384,9 @@ func SafeCleanupServerState() {
354384
skippedCount++
355385
}
356386
}
357-
387+
358388
log.Printf("SafeCleanupServerState: 清理完成，成功清理 %d 个，跳过 %d 个", cleanedCount, skippedCount)
359-
389+
360390
case <-time.After(1 * time.Second):
361391
log.Printf("SafeCleanupServerState: 整体操作超时")
362392
}
@@ -367,10 +397,10 @@ func cleanupSingleServerState(server *model.Server) bool {
367397
if server == nil {
368398
return false
369399
}
370-
400+
371401
now := time.Now()
372402
cleaned := false
373-
403+
374404
// 清理长时间离线的服务器状态
375405
if !server.IsOnline && now.Sub(server.LastActive) > 12*time.Hour {
376406
// 保存最后状态到数据库（非阻塞）
@@ -387,27 +417,27 @@ func cleanupSingleServerState(server *model.Server) bool {
387417
}()
388418
}
389419
}
390-
420+
391421
// 清理内存状态
392422
server.State = nil
393423
server.LastStateBeforeOffline = nil
394424
cleaned = true
395425
}
396-
426+
397427
// 清理任务连接（使用超时锁）
398428
if server.TaskCloseLock != nil {
399429
lockAcquired := make(chan bool, 1)
400-
430+
401431
// 尝试获取TaskCloseLock
402432
go func() {
403433
server.TaskCloseLock.Lock()
404434
lockAcquired <- true
405435
}()
406-
436+
407437
select {
408438
case <-lockAcquired:
409439
defer server.TaskCloseLock.Unlock()
410-
440+
411441
if server.TaskClose != nil {
412442
// 非阻塞发送关闭信号
413443
select {
@@ -416,20 +446,20 @@ func cleanupSingleServerState(server *model.Server) bool {
416446
default:
417447
log.Printf("服务器 %s 的任务通道已满或已关闭，直接关闭", server.Name)
418448
}
419-
449+
420450
// 安全关闭通道
421451
close(server.TaskClose)
422452
server.TaskClose = nil
423453
}
424-
454+
425455
server.TaskStream = nil
426456
cleaned = true
427-
457+
428458
case <-time.After(100 * time.Millisecond):
429459
log.Printf("获取服务器 %s 的TaskCloseLock超时，跳过任务清理", server.Name)
430460
}
431461
}
432-
462+
433463
return cleaned
434464
}
435465

0 commit comments

Comments
 (0)