From c5a6b682ecd6438228404d0d097cf119928ccc23 Mon Sep 17 00:00:00 2001 From: dushixiang Date: Sat, 10 Jan 2026 18:28:19 +0800 Subject: [PATCH] =?UTF-8?q?=E4=BF=AE=E5=A4=8D=E6=89=8B=E5=8A=A8=E8=B0=83?= =?UTF-8?q?=E6=95=B4=E5=B7=B2=E4=BD=BF=E7=94=A8=E6=B5=81=E9=87=8F=E8=A2=AB?= =?UTF-8?q?=E8=A6=86=E7=9B=96=E7=9A=84=E9=97=AE=E9=A2=98?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- internal/models/agent.go | 1 + internal/service/traffic_service.go | 118 +++++++++++++++++++++------- 2 files changed, 91 insertions(+), 28 deletions(-) diff --git a/internal/models/agent.go b/internal/models/agent.go index f8001ca..2ba571a 100644 --- a/internal/models/agent.go +++ b/internal/models/agent.go @@ -37,6 +37,7 @@ type TrafficStatsData struct { Type string `json:"type"` // 统计类型: "recv"进站, "send"出站, "both"全部 Limit uint64 `json:"limit"` // 流量限额(字节), 0表示不限制 Used uint64 `json:"used"` // 当前周期已使用流量(字节) + UsedOffset int64 `json:"usedOffset"` // 手动调整流量偏移(字节), 允许为负数 ResetDay int `json:"resetDay"` // 流量重置日期(1-31), 0表示不自动重置 PeriodStart int64 `json:"periodStart"` // 当前周期开始时间(时间戳毫秒) BaselineRecv uint64 `json:"baselineRecv"` // 当前周期流量基线(BytesRecvTotal) diff --git a/internal/service/traffic_service.go b/internal/service/traffic_service.go index f97e22d..0433e69 100644 --- a/internal/service/traffic_service.go +++ b/internal/service/traffic_service.go @@ -7,6 +7,7 @@ import ( "github.com/dushixiang/pika/internal/models" "github.com/dushixiang/pika/internal/repo" + "go.uber.org/zap" "gorm.io/datatypes" "gorm.io/gorm" @@ -26,6 +27,13 @@ func NewTrafficService(logger *zap.Logger, db *gorm.DB) *TrafficService { } } +func (s *TrafficService) updateById(ctx context.Context, agentID string, stats *models.TrafficStatsData) error { + return s.agentRepo.UpdateById(ctx, &models.Agent{ + ID: agentID, + TrafficStats: datatypes.NewJSONType(*stats), + }) +} + // UpdateAgentTraffic 更新探针流量统计(每次上报网络指标时调用) func (s *TrafficService) UpdateAgentTraffic(ctx context.Context, agentID string, currentRecvTotal, currentSentTotal uint64) error { agent, err := s.agentRepo.FindById(ctx, agentID) @@ -46,16 +54,17 @@ func (s *TrafficService) UpdateAgentTraffic(ctx context.Context, agentID string, stats.Type = "recv" // 默认只统计进站流量，向后兼容 } - // 初始化基线(首次统计) + now := time.Now().UnixMilli() + baselineChanged := false + + // 初始化基线(首次统计或基线丢失) if stats.BaselineRecv == 0 && stats.BaselineSend == 0 { stats.BaselineRecv = currentRecvTotal stats.BaselineSend = currentSentTotal - stats.Used = 0 if stats.PeriodStart == 0 { - stats.PeriodStart = time.Now().UnixMilli() + stats.PeriodStart = now } - agent.TrafficStats = datatypes.NewJSONType(stats) - return s.agentRepo.UpdateById(ctx, &agent) + baselineChanged = true } // 检测计数器重置(探针重启) @@ -76,23 +85,15 @@ func (s *TrafficService) UpdateAgentTraffic(ctx context.Context, agentID string, if sendReset { stats.BaselineSend = currentSentTotal } - // 保持 Used 不变,避免丢失已统计的流量 - } else { - // 根据配置的类型计算使用量 - switch stats.Type { - case "recv": // 只统计进站流量 - stats.Used = currentRecvTotal - stats.BaselineRecv - case "send": // 只统计出站流量 - stats.Used = currentSentTotal - stats.BaselineSend - case "both": // 统计全部流量 - recvUsed := currentRecvTotal - stats.BaselineRecv - sendUsed := currentSentTotal - stats.BaselineSend - stats.Used = recvUsed + sendUsed - default: // 默认只统计进站流量 - stats.Used = currentRecvTotal - stats.BaselineRecv - } + baselineChanged = true } + baseDelta := calculateBaseDelta(stats.Type, currentRecvTotal, currentSentTotal, stats.BaselineRecv, stats.BaselineSend) + if baselineChanged { + stats.UsedOffset = calculateOffset(stats.Used, baseDelta) + } + stats.Used = applyOffset(baseDelta, stats.UsedOffset) + // 检查告警(如果配置了限额) if stats.Limit > 0 { s.checkTrafficAlerts(ctx, &agent, &stats) @@ -102,7 +103,7 @@ func (s *TrafficService) UpdateAgentTraffic(ctx context.Context, agentID string, agent.TrafficStats = datatypes.NewJSONType(stats) // 更新数据库 - return s.agentRepo.UpdateById(ctx, &agent) + return s.updateById(ctx, agentID, &stats) } // checkTrafficAlerts 检查并发送流量告警 @@ -220,6 +221,7 @@ func (s *TrafficService) UpdateTrafficConfig(ctx context.Context, agentID string // 如果是首次启用、禁用后重新启用、修改重置日期或修改流量类型，重置流量统计 if (!oldEnabled && stats.Enabled) || (stats.Enabled && resetDay != oldResetDay) || (stats.Enabled && stats.PeriodStart == 0) || (stats.Enabled && oldType != stats.Type && oldType != "") { stats.Used = 0 + stats.UsedOffset = 0 stats.PeriodStart = now stats.BaselineRecv = 0 // 下次上报时会设置正确的基线 stats.BaselineSend = 0 @@ -230,12 +232,14 @@ func (s *TrafficService) UpdateTrafficConfig(ctx context.Context, agentID string // 如果提供了 used 参数且大于 0，则使用该值 if used > 0 { + stats.UsedOffset = addOffset(stats.UsedOffset, calculateOffset(used, stats.Used)) stats.Used = used } // 如果禁用流量统计，清空相关数据 if !stats.Enabled { stats.Used = 0 + stats.UsedOffset = 0 stats.PeriodStart = 0 stats.BaselineRecv = 0 stats.BaselineSend = 0 @@ -244,9 +248,7 @@ func (s *TrafficService) UpdateTrafficConfig(ctx context.Context, agentID string stats.AlertSent100 = false } - agent.TrafficStats = datatypes.NewJSONType(stats) - agent.UpdatedAt = now - return s.agentRepo.UpdateById(ctx, &agent) + return s.updateById(ctx, agentID, &stats) } // GetTrafficStats 获取流量统计信息 @@ -313,6 +315,7 @@ func (s *TrafficService) ResetAgentTraffic(ctx context.Context, agentID string) stats := agent.TrafficStats.Data() stats.Used = 0 + stats.UsedOffset = 0 stats.BaselineRecv = 0 // 下次上报时会设置正确的基线 stats.BaselineSend = 0 stats.PeriodStart = now @@ -320,14 +323,11 @@ func (s *TrafficService) ResetAgentTraffic(ctx context.Context, agentID string) stats.AlertSent90 = false stats.AlertSent100 = false - agent.TrafficStats = datatypes.NewJSONType(stats) - agent.UpdatedAt = now - s.logger.Info("探针流量已重置", zap.String("agentId", agentID), zap.String("agentName", agent.Name)) - return s.agentRepo.UpdateById(ctx, &agent) + return s.updateById(ctx, agentID, &stats) } // CheckAndResetTraffic 检查并重置所有到期的探针流量(定时任务调用) @@ -409,6 +409,68 @@ func calculateNextResetDate(periodStart time.Time, resetDay int) time.Time { return time.Date(nextYear, nextMonth, nextMonthResetDay, 0, 0, 0, 0, location) } +func calculateBaseDelta(trafficType string, currentRecvTotal, currentSentTotal, baselineRecv, baselineSend uint64) uint64 { + switch trafficType { + case "recv": + return currentRecvTotal - baselineRecv + case "send": + return currentSentTotal - baselineSend + case "both": + recvUsed := currentRecvTotal - baselineRecv + sendUsed := currentSentTotal - baselineSend + return recvUsed + sendUsed + default: + return currentRecvTotal - baselineRecv + } +} + +func calculateOffset(targetUsed, baseDelta uint64) int64 { + const maxInt64 = int64(^uint64(0) >> 1) + const minInt64 = -maxInt64 - 1 + + if targetUsed >= baseDelta { + diff := targetUsed - baseDelta + if diff > uint64(maxInt64) { + return maxInt64 + } + return int64(diff) + } + + diff := baseDelta - targetUsed + if diff > uint64(maxInt64) { + return minInt64 + } + return -int64(diff) +} + +func applyOffset(baseDelta uint64, offset int64) uint64 { + if offset >= 0 { + return baseDelta + uint64(offset) + } + const minInt64 = -1 - (int64(^uint64(0) >> 1)) + if offset == minInt64 { + return 0 + } + neg := uint64(-offset) + if neg > baseDelta { + return 0 + } + return baseDelta - neg +} + +func addOffset(base, delta int64) int64 { + const maxInt64 = int64(^uint64(0) >> 1) + const minInt64 = -maxInt64 - 1 + + if delta > 0 && base > maxInt64-delta { + return maxInt64 + } + if delta < 0 && base < minInt64-delta { + return minInt64 + } + return base + delta +} + // TrafficStats 流量统计信息 type TrafficStats struct { Enabled bool `json:"enabled"`