Merge pull request #65 from komari-monitor/dev

fix: #57 api.Records引入线程安全的泛型映射

api/Common.go

@@ -15,7 +15,7 @@ import (
 )
 
 var (
-	Records map[string][]common.Report = make(map[string][]common.Report)
+	Records = utils.NewSafeMap[string, []common.Report]()
 )
 
 type TerminalSession struct {
@@ -31,8 +31,10 @@ var TerminalSessions = make(map[string]*TerminalSession)
 func SaveClientReportToDB() error {
 	lastMinute := time.Now().Add(-time.Minute * 1).Unix()
 	record := []models.Record{}
+	var err error
+
 	// 遍历所有的客户端记录
-	for uuid, reports := range Records {
+	Records.Range(func(uuid string, reports []common.Report) bool {
 		// 删除一分钟前的记录
 		filtered := reports[:0]
 		for _, r := range reports {
@@ -40,24 +42,24 @@ func SaveClientReportToDB() error {
 				filtered = append(filtered, r)
 			}
 		}
-		Records[uuid] = filtered
+		Records.Store(uuid, filtered)
 
 		if uuid == "" {
-			continue
+			return true
 		}
 
 		r := utils.AverageReport(uuid, time.Now(), filtered)
-
 		record = append(record, r)
 
 		db := dbcore.GetDBInstance()
-		err := db.Model(&models.Record{}).Where("client = ?", uuid).Create(record).Error
+		err = db.Model(&models.Record{}).Create(record).Error
 		if err != nil {
-			return err
+			return false
 		}
-	}
+		return true
+	})
 
-	return nil
+	return err
 }
 
 type Response struct {

api/admin/ping.go

@@ -62,9 +62,9 @@ func EditPingTask(c *gin.Context) {
 	if err := tasks.EditPingTask(req.Tasks); err != nil {
 		api.RespondError(c, http.StatusInternalServerError, err.Error())
 	} else {
-		for _, task := range req.Tasks {
-			tasks.DeletePingRecords([]uint{task.Id})
-		}
+		// for _, task := range req.Tasks {
+		// 	tasks.DeletePingRecords([]uint{task.Id})
+		// }
 		api.RespondSuccess(c, nil)
 	}
 }

api/admin/test/test.go

@@ -10,11 +10,11 @@ import (
 	"github.com/komari-monitor/komari/utils/messageSender"
 )
 
-func TestTelegram(c *gin.Context) {
+func TestSendMessage(c *gin.Context) {
 	message := "Komari test message."
 	err := messageSender.SendTextMessage(message)
 	if err != nil {
-		api.RespondError(c, 500, "Failed to send Telegram message: "+err.Error())
+		api.RespondError(c, 500, "Failed to send message: "+err.Error())
 		return
 	}
 	api.RespondSuccess(c, nil)

api/client/report.go

@@ -178,13 +178,15 @@ func WebSocketReport(c *gin.Context) {
 }
 
 func SaveClientReport(uuid string, report common.Report) error {
-	if api.Records[uuid] == nil {
-		api.Records[uuid] = []common.Report{}
+	reports, _ := api.Records.Get(uuid)
+	if reports == nil {
+		reports = []common.Report{}
 	}
 	if report.CPU.Usage < 0.01 {
 		report.CPU.Usage = 0.01
 	}
-	api.Records[uuid] = append(api.Records[uuid], report)
+	reports = append(reports, report)
+	api.Records.Set(uuid, reports)
 
 	return nil
 }

api/latestRecord.go

@@ -10,5 +10,6 @@ func GetClientRecentRecords(c *gin.Context) {
 		RespondError(c, 400, "UUID is required")
 		return
 	}
-	RespondSuccess(c, Records[uuid])
+	records, _ := Records.Get(uuid)
+	RespondSuccess(c, records)
 }

api/record/record.go

@@ -69,17 +69,38 @@ func GetPingRecords(c *gin.Context) {
 	}
 
 	if len(records) > 0 {
+		// 获取当前属于该 uuid 的 pingTasks
 		pingTasks, err := tasks.GetAllPingTasks()
 		if err != nil {
 			api.RespondError(c, 500, "Failed to fetch ping tasks: "+err.Error())
 			return
 		}
+
+		taskIdSet := make(map[uint]struct{})
+		for _, r := range records {
+			taskIdSet[r.TaskId] = struct{}{}
+		}
+
 		tasksList := make([]gin.H, 0, len(pingTasks))
 		for _, t := range pingTasks {
-			tasksList = append(tasksList, gin.H{
-				"id":   t.Id,
-				"name": t.Name,
-			})
+			// 只保留分配给该 uuid 的任务
+			found := false
+			for _, client := range t.Clients {
+				if client == uuid {
+					found = true
+					break
+				}
+			}
+			if !found {
+				continue
+			}
+			// 只返回有记录的任务
+			if _, ok := taskIdSet[t.Id]; ok {
+				tasksList = append(tasksList, gin.H{
+					"id":   t.Id,
+					"name": t.Name,
+				})
+			}
 		}
 		response["tasks"] = tasksList
 	}

cmd/server.go

@@ -120,8 +120,8 @@ var ServerCmd = &cobra.Command{
 			// test
 			testGroup := adminAuthrized.Group("/test")
 			{
-				testGroup.GET("/telegram", test.TestTelegram)
 				testGroup.GET("/geoip", test.TestGeoIp)
+				testGroup.POST("/sendMessage", test.TestSendMessage)
 			}
 			// update
 			updateGroup := adminAuthrized.Group("/update")
@@ -290,7 +290,7 @@ func DoScheduledWork() {
 			records.DeleteRecordBefore(time.Now().Add(-time.Hour * time.Duration(cfg.RecordPreserveTime)))
 			records.CompactRecord()
 			tasks.ClearTaskResultsByTimeBefore(time.Now().Add(-time.Hour * time.Duration(cfg.RecordPreserveTime)))
-			tasks.DeletePingRecordsBefore(time.Now().Add(-time.Hour * time.Duration(cfg.RecordPreserveTime)))
+			tasks.DeletePingRecordsBefore(time.Now().Add(-time.Hour * time.Duration(cfg.PingRecordPreserveTime)))
 			logOperation.RemoveOldLogs()
 		case <-minute.C:
 			api.SaveClientReportToDB()

database/records/records.go

@@ -43,11 +43,10 @@ func GetRecordsByClientAndTime(uuid string, start, end time.Time) ([]models.Reco
 	var records []models.Record
 
 	threeHoursAgo := time.Now().Add(-3 * time.Hour)
-	interval := end.Sub(start)
 
+	var recentRecords []models.Record
+	recentStart := start
 	if end.After(threeHoursAgo) {
-		var recentRecords []models.Record
-		recentStart := start
 		if recentStart.Before(threeHoursAgo) {
 			recentStart = threeHoursAgo
 		}
@@ -56,35 +55,37 @@ func GetRecordsByClientAndTime(uuid string, start, end time.Time) ([]models.Reco
 			log.Printf("Error fetching recent records for client %s between %s and %s: %v", uuid, recentStart, end, err)
 			return nil, err
 		}
-		if interval > 4*time.Hour {
-			// 按15分钟分组，每组只保留一条（取最新一条）
-			grouped := make(map[string]models.Record)
-			for _, rec := range recentRecords {
-				key := rec.Time.Truncate(15 * time.Minute).Format(time.RFC3339)
-				if old, ok := grouped[key]; !ok || rec.Time.After(old.Time) {
-					grouped[key] = rec
-				}
-			}
-			var groupedList []models.Record
-			for _, rec := range grouped {
-				groupedList = append(groupedList, rec)
-			}
-			sort.Slice(groupedList, func(i, j int) bool {
-				return groupedList[i].Time.Before(groupedList[j].Time)
-			})
-			records = append(records, groupedList...)
-		} else {
-			// 查询区间不超过，直接返回全部recentRecords
-			records = append(records, recentRecords...)
-		}
 	}
 
 	var long_term []models.Record
 	err := db.Table("records_long_term").Where("client = ? AND time >= ? AND time <= ?", uuid, start, end).Order("time ASC").Find(&long_term).Error
 	if err != nil {
 		log.Printf("Error fetching long-term records for client %s between %s and %s: %v", uuid, start, end, err)
+		return recentRecords, nil
+	}
+
+	if len(long_term) == 0 {
+		// 没有查到long_term，返回全部recentRecords
+		records = append(records, recentRecords...)
 		return records, nil
 	}
+
+	// 查到了long_term，recentRecords按15分钟分组，每组只保留一条（取最新一条）
+	grouped := make(map[string]models.Record)
+	for _, rec := range recentRecords {
+		key := rec.Time.Truncate(15 * time.Minute).Format(time.RFC3339)
+		if old, ok := grouped[key]; !ok || rec.Time.After(old.Time) {
+			grouped[key] = rec
+		}
+	}
+	var groupedList []models.Record
+	for _, rec := range grouped {
+		groupedList = append(groupedList, rec)
+	}
+	sort.Slice(groupedList, func(i, j int) bool {
+		return groupedList[i].Time.Before(groupedList[j].Time)
+	})
+	records = append(records, groupedList...)
 	records = append(records, long_term...)
 	return records, nil
 }

utils/safaMap.go

@@ -0,0 +1,88 @@
+package utils
+
+import (
+	"sync"
+)
+
+// SafeMap 是一个线程安全的泛型映射类型
+// 使用 sync.RWMutex 来确保并发访问的安全性
+type SafeMap[K comparable, V any] struct {
+	mu    sync.RWMutex // 读写锁，用于保护 items 的并发访问
+	items map[K]V      // 存储键值对的底层映射
+}
+
+// NewSafeMap 创建一个新的 SafeMap 实例
+// 返回一个指向 SafeMap 的指针
+func NewSafeMap[K comparable, V any]() *SafeMap[K, V] {
+	return &SafeMap[K, V]{
+		items: make(map[K]V), // 初始化底层映射
+	}
+}
+
+// Set 添加或更新映射中的键值对
+// 使用写锁保护操作
+func (sm *SafeMap[K, V]) Set(key K, value V) {
+	sm.mu.Lock()
+	defer sm.mu.Unlock()
+	sm.items[key] = value
+}
+
+// Store 添加或更新映射中的键值对
+// 是 Set 方法的别名
+func (sm *SafeMap[K, V]) Store(key K, value V) {
+	sm.Set(key, value)
+}
+
+// Get 根据键从映射中检索值
+// 返回值和一个布尔值，指示键是否存在
+func (sm *SafeMap[K, V]) Get(key K) (V, bool) {
+	sm.mu.RLock()
+	defer sm.mu.RUnlock()
+
+	value, exists := sm.items[key]
+	return value, exists
+}
+
+// Delete 从映射中移除指定键的键值对
+// 使用写锁保护操作
+func (sm *SafeMap[K, V]) Delete(key K) {
+	sm.mu.Lock()
+	defer sm.mu.Unlock()
+
+	delete(sm.items, key)
+}
+
+// Keys 返回映射中所有键的切片
+// 使用读锁保护操作
+func (sm *SafeMap[K, V]) Keys() []K {
+	sm.mu.RLock()
+	defer sm.mu.RUnlock()
+
+	keys := make([]K, 0, len(sm.items))
+	for key := range sm.items {
+		keys = append(keys, key)
+	}
+	return keys
+}
+
+// Len 返回映射中键值对的数量
+// 使用读锁保护操作
+func (sm *SafeMap[K, V]) Len() int {
+	sm.mu.RLock()
+	defer sm.mu.RUnlock()
+
+	return len(sm.items)
+}
+
+// Range 遍历映射中的所有键值对
+// 接受一个函数作为参数，如果函数返回 false，则停止遍历
+func (sm *SafeMap[K, V]) Range(f func(key K, value V) bool) {
+	sm.mu.RLock()
+	defer sm.mu.RUnlock()
+
+	for key, value := range sm.items {
+		if !f(key, value) {
+			break
+		}
+	}
+}