LokiStack/backend/services/daily_volume.go

package services

import (
	"context"
	"log"
	"sort"
	"time"

	"lokistack/models"
	"lokistack/repository"
)

// DailyVolumePoint 返回给前端的单日数据点
type DailyVolumePoint struct {
	Date    string `json:"date"`
	Bytes   int64  `json:"bytes"`
	Entries int64  `json:"entries"`
	// Partial 当天还没过完（或抓取时数据不全），前端用虚线/浅色区分
	Partial bool `json:"partial"`
	// FromHistory true=来自平台留存历史；false=刚从 Loki 实时查的
	FromHistory bool `json:"fromHistory"`
}

const (
	// dailyVolumeBackfillDays 启动时回补天数。
	// 取 2 天：昨天（完整）+ 今天（进行中），足够覆盖"重启导致漏抓"的常见情况。
	// 不设更大是因为 Loki retention 通常只有 24h，更早的已经查不到了。
	dailyVolumeBackfillDays = 2
	// dailyVolumeRetentionDays 平台自己留多久的历史。一年足够看年度趋势。
	dailyVolumeRetentionDays = 400
)

// dateKey 返回本地时区的 YYYY-MM-DD。
func dateKey(t time.Time) string {
	return t.In(time.Local).Format("2006-01-02")
}

// dayRange 返回某天本地时区的 [00:00, 次日 00:00) 纳秒区间。
func dayRange(day time.Time) (int64, int64) {
	local := day.In(time.Local)
	start := time.Date(local.Year(), local.Month(), local.Day(), 0, 0, 0, 0, time.Local)
	return start.UnixNano(), start.AddDate(0, 0, 1).UnixNano()
}

// CaptureDailyVolume 抓取指定日期的日志量并落库（按天幂等：同一天重复抓会覆盖）。
//
// 数据来源是 Loki 的 index/stats（只读索引元数据），实测毫秒级，
// 所以就算是补几十天历史也不会拖慢任何东西。
func CaptureDailyVolume(source *models.LogSource, day time.Time, partial bool) error {
	if source == nil {
		return nil
	}
	start, end := dayRange(day)
	// 当天未结束时截到"现在"，否则会把未来时间也查进去
	now := time.Now().UnixNano()
	if end > now {
		end = now
	}
	if start >= end {
		return nil
	}
	// 回补历史时，如果那天的起点已经超出 Loki 保留期（默认 24h），
	// Loki 里根本没那么多数据 —— 必须标成"不完整"，
	// 否则图上会把"数据被删了"误读成"那天日志量真的很少"。
	if start < time.Now().Add(-24*time.Hour).UnixNano() {
		partial = true
	}

	// 只查总量（不带 namespace 条件），一次拿到该源当天的条数与体积
	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
	defer cancel()
	stats, err := LokiStatsFor(ctx, source, `{namespace=~".+"}`, start, end)
	if err != nil {
		return err
	}

	date := dateKey(day)
	db := repository.GetDB()
	if db == nil {
		return nil
	}
	db.Lock()
	defer db.Unlock()

	nextID := uint(0)
	// 用下标取地址：`for _, snap := range` 拿到的是**副本**，
	// 改 snap 的字段不会写回切片（第一版就是这里写错了，
	// 导致同一天永远保持首次抓取的值、partial 不会更新）。
	for i := range db.DailyVolumeSnapshots {
		snap := &db.DailyVolumeSnapshots[i]
		if snap.ID > nextID {
			nextID = snap.ID
		}
		// 同一天同源：覆盖（幂等）
		if snap.Date == date && snap.SourceID == source.ID && snap.Namespace == "" {
			snap.Bytes = stats.Bytes
			snap.Entries = stats.Entries
			snap.Partial = partial
			snap.CapturedAt = time.Now()
			db.RequestSave()
			return nil
		}
	}
	nextID++
	db.DailyVolumeSnapshots = append(db.DailyVolumeSnapshots, models.DailyVolumeSnapshot{
		ID:         nextID,
		Date:       date,
		SourceID:   source.ID,
		Namespace:  "",
		Bytes:      stats.Bytes,
		Entries:    stats.Entries,
		Partial:    partial,
		CapturedAt: time.Now(),
	})
	db.RequestSave()
	return nil
}

// CaptureRecentDailyVolumes 回补最近 N 天（含今天），供启动时与定时任务调用。
func CaptureRecentDailyVolumes() {
	db := repository.GetDB()
	if db == nil {
		return
	}
	db.RLock()
	sources := make([]models.LogSource, 0, len(db.LogSources))
	for _, src := range db.LogSources {
		sources = append(sources, src)
	}
	db.RUnlock()

	now := time.Now()
	for _, src := range sources {
		src := src
		for offset := dailyVolumeBackfillDays - 1; offset >= 0; offset-- {
			day := now.AddDate(0, 0, -offset)
			// 今天是不完整的一天
			partial := offset == 0
			if err := CaptureDailyVolume(&src, day, partial); err != nil {
				log.Printf("[daily-volume] 抓取 %s %s 失败: %v", src.Name, dateKey(day), err)
			}
		}
	}
	pruneDailyVolumeSnapshots()
}

// pruneDailyVolumeSnapshots 丢弃超过保留期的历史。
func pruneDailyVolumeSnapshots() {
	db := repository.GetDB()
	if db == nil {
		return
	}
	cutoff := dateKey(time.Now().AddDate(0, 0, -dailyVolumeRetentionDays))
	db.Lock()
	defer db.Unlock()
	kept := make([]models.DailyVolumeSnapshot, 0, len(db.DailyVolumeSnapshots))
	for _, snap := range db.DailyVolumeSnapshots {
		if snap.Date >= cutoff {
			kept = append(kept, snap)
		}
	}
	if len(kept) != len(db.DailyVolumeSnapshots) {
		log.Printf("[daily-volume] 清理过期快照 %d 条", len(db.DailyVolumeSnapshots)-len(kept))
		db.DailyVolumeSnapshots = kept
		db.RequestSave()
	}
}

// GetDailyVolumeHistory 从平台留存的历史里取最近 days 天，按日期升序返回。
//
// 今天如果还没抓过，会实时补一次（这样用户切到页面时当天数字是最新的）。
func GetDailyVolumeHistory(sourceID uint, days int) []DailyVolumePoint {
	if days <= 0 {
		days = 7
	}
	if days > dailyVolumeRetentionDays {
		days = dailyVolumeRetentionDays
	}

	// 今天实时补抓，保证"今天"这根柱子是活的
	db := repository.GetDB()
	if db != nil {
		db.RLock()
		var target *models.LogSource
		for i := range db.LogSources {
			if db.LogSources[i].ID == sourceID {
				src := db.LogSources[i]
				target = &src
				break
			}
		}
		db.RUnlock()
		if target != nil {
			if err := CaptureDailyVolume(target, time.Now(), true); err != nil {
				log.Printf("[daily-volume] 实时补抓今天失败: %v", err)
			}
		}
	}

	from := dateKey(time.Now().AddDate(0, 0, -(days - 1)))
	points := []DailyVolumePoint{}
	db = repository.GetDB()
	if db == nil {
		return points
	}
	db.RLock()
	defer db.RUnlock()
	for _, snap := range db.DailyVolumeSnapshots {
		if snap.SourceID != sourceID || snap.Namespace != "" {
			continue
		}
		if snap.Date < from {
			continue
		}
		points = append(points, DailyVolumePoint{
			Date:        snap.Date,
			Bytes:       snap.Bytes,
			Entries:     snap.Entries,
			Partial:     snap.Partial,
			FromHistory: true,
		})
	}
	sort.Slice(points, func(i, j int) bool { return points[i].Date < points[j].Date })
	return points
}

// StartDailyVolumeCapture 后台每日抓取。
//
// 每 30 分钟检查一次：只要发现"昨天"还没落库就立刻补，
// 而不是死等 00:10 —— 平台可能整晚都没运行，定时器会错过那个点。
func StartDailyVolumeCapture() {
	go func() {
		// 启动先补最近两天（覆盖重启场景）
		time.Sleep(20 * time.Second)
		CaptureRecentDailyVolumes()

		ticker := time.NewTicker(30 * time.Minute)
		defer ticker.Stop()
		for range ticker.C {
			CaptureRecentDailyVolumes()
		}
	}()
}