proxy-go/internal/metrics/persistence.go

package metrics

import (
	"encoding/json"
	"fmt"
	"log"
	"os"
	"path/filepath"
	"proxy-go/internal/utils"
	"runtime"
	"strconv"
	"sync"
	"sync/atomic"
	"time"
)

// MetricsStorage 指标存储结构
type MetricsStorage struct {
	collector      *Collector
	saveInterval   time.Duration
	dataDir        string
	stopChan       chan struct{}
	wg             sync.WaitGroup
	lastSaveTime   time.Time
	mutex          sync.RWMutex
	statusCodeFile string
}

// NewMetricsStorage 创建新的指标存储
func NewMetricsStorage(collector *Collector, dataDir string, saveInterval time.Duration) *MetricsStorage {
	if saveInterval < time.Minute {
		saveInterval = time.Minute
	}

	return &MetricsStorage{
		collector:      collector,
		saveInterval:   saveInterval,
		dataDir:        dataDir,
		stopChan:       make(chan struct{}),
		statusCodeFile: filepath.Join(dataDir, "status_codes.json"),
	}
}

// Start 启动定时保存任务
func (ms *MetricsStorage) Start() error {
	// 确保数据目录存在
	if err := os.MkdirAll(ms.dataDir, 0755); err != nil {
		return fmt.Errorf("创建数据目录失败: %v", err)
	}

	// 尝试加载现有数据
	if err := ms.LoadMetrics(); err != nil {
		log.Printf("[MetricsStorage] 加载指标数据失败: %v", err)
		// 加载失败不影响启动
	}

	ms.wg.Add(1)
	go ms.runSaveTask()
	log.Printf("[MetricsStorage] 指标存储服务已启动，保存间隔: %v", ms.saveInterval)
	return nil
}

// Stop 停止定时保存任务
func (ms *MetricsStorage) Stop() {
	close(ms.stopChan)
	ms.wg.Wait()

	// 在停止前保存一次数据
	if err := ms.SaveMetrics(); err != nil {
		log.Printf("[MetricsStorage] 停止时保存指标数据失败: %v", err)
	}

	log.Printf("[MetricsStorage] 指标存储服务已停止")
}

// runSaveTask 运行定时保存任务
func (ms *MetricsStorage) runSaveTask() {
	defer ms.wg.Done()

	ticker := time.NewTicker(ms.saveInterval)
	defer ticker.Stop()

	for {
		select {
		case <-ticker.C:
			if err := ms.SaveMetrics(); err != nil {
				log.Printf("[MetricsStorage] 保存指标数据失败: %v", err)
			}
		case <-ms.stopChan:
			return
		}
	}
}

// SaveMetrics 保存指标数据
func (ms *MetricsStorage) SaveMetrics() error {
	start := time.Now()
	log.Printf("[MetricsStorage] 开始保存指标数据...")

	// 获取当前指标数据
	stats := ms.collector.GetStats()

	// 保存状态码统计
	if err := saveJSONToFile(ms.statusCodeFile, stats["status_code_stats"]); err != nil {
		return fmt.Errorf("保存状态码统计失败: %v", err)
	}

	// 不再保存引用来源统计，因为它现在只保存在内存中

	// 单独保存延迟分布
	if latencyStats, ok := stats["latency_stats"].(map[string]interface{}); ok {
		if distribution, ok := latencyStats["distribution"]; ok {
			if err := saveJSONToFile(filepath.Join(ms.dataDir, "latency_distribution.json"), distribution); err != nil {
				log.Printf("[MetricsStorage] 保存延迟分布失败: %v", err)
			}
		}
	}

	// 强制进行一次GC
	runtime.GC()

	// 打印内存使用情况
	var mem runtime.MemStats
	runtime.ReadMemStats(&mem)

	log.Printf("[MetricsStorage] 指标数据保存完成，耗时: %v, 内存使用: %s",
		time.Since(start), utils.FormatBytes(int64(mem.Alloc)))
	return nil
}

// LoadMetrics 加载指标数据
func (ms *MetricsStorage) LoadMetrics() error {
	start := time.Now()
	log.Printf("[MetricsStorage] 开始加载指标数据...")

	// 不再加载 basicMetrics（metrics.json）

	// 1. 加载状态码统计（如果文件存在）
	if fileExists(ms.statusCodeFile) {
		var statusCodeStats map[string]interface{}
		if err := loadJSONFromFile(ms.statusCodeFile, &statusCodeStats); err != nil {
			log.Printf("[MetricsStorage] 加载状态码统计失败: %v", err)
		} else {
			// 由于新的 StatusCodeStats 结构，我们需要手动设置值
			loadedCount := 0
			for codeStr, countValue := range statusCodeStats {
				// 解析状态码
				if code, err := strconv.Atoi(codeStr); err == nil {
					// 解析计数值
					var count int64
					switch v := countValue.(type) {
					case float64:
						count = int64(v)
					case int64:
						count = v
					case int:
						count = int64(v)
					default:
						continue
					}

					// 手动设置到新的 StatusCodeStats 结构中
					ms.collector.statusCodeStats.mu.Lock()
					if _, exists := ms.collector.statusCodeStats.stats[code]; !exists {
						ms.collector.statusCodeStats.stats[code] = new(int64)
					}
					atomic.StoreInt64(ms.collector.statusCodeStats.stats[code], count)
					ms.collector.statusCodeStats.mu.Unlock()
					loadedCount++
				}
			}
			log.Printf("[MetricsStorage] 成功加载了 %d 条状态码统计", loadedCount)
		}
	}

	// 不再加载引用来源统计，因为它现在只保存在内存中

	// 3. 加载延迟分布（如果文件存在）
	latencyDistributionFile := filepath.Join(ms.dataDir, "latency_distribution.json")
	if fileExists(latencyDistributionFile) {
		var distribution map[string]interface{}
		if err := loadJSONFromFile(latencyDistributionFile, &distribution); err != nil {
			log.Printf("[MetricsStorage] 加载延迟分布失败: %v", err)
		} else {
			// 由于新的 LatencyBuckets 结构，我们需要手动设置值
			for bucket, count := range distribution {
				countValue, ok := count.(float64)
				if !ok {
					continue
				}

				// 根据桶名称设置对应的值
				switch bucket {
				case "lt10ms":
					atomic.StoreInt64(&ms.collector.latencyBuckets.lt10ms, int64(countValue))
				case "10-50ms":
					atomic.StoreInt64(&ms.collector.latencyBuckets.ms10_50, int64(countValue))
				case "50-200ms":
					atomic.StoreInt64(&ms.collector.latencyBuckets.ms50_200, int64(countValue))
				case "200-1000ms":
					atomic.StoreInt64(&ms.collector.latencyBuckets.ms200_1000, int64(countValue))
				case "gt1s":
					atomic.StoreInt64(&ms.collector.latencyBuckets.gt1s, int64(countValue))
				}
			}
			log.Printf("[MetricsStorage] 加载了延迟分布数据")
		}
	}
	// 强制进行一次GC
	runtime.GC()

	// 打印内存使用情况
	var mem runtime.MemStats
	runtime.ReadMemStats(&mem)

	log.Printf("[MetricsStorage] 指标数据加载完成，耗时: %v, 内存使用: %s",
		time.Since(start), utils.FormatBytes(int64(mem.Alloc)))
	return nil
}

// GetLastSaveTime 获取最后保存时间
func (ms *MetricsStorage) GetLastSaveTime() time.Time {
	ms.mutex.RLock()
	defer ms.mutex.RUnlock()
	return ms.lastSaveTime
}

// 辅助函数：保存JSON到文件
func saveJSONToFile(filename string, data interface{}) error {
	// 创建临时文件
	tempFile := filename + ".tmp"

	// 将数据编码为JSON
	jsonData, err := json.MarshalIndent(data, "", "  ")
	if err != nil {
		return err
	}

	// 写入临时文件
	if err := os.WriteFile(tempFile, jsonData, 0644); err != nil {
		return err
	}

	// 重命名临时文件为目标文件（原子操作）
	return os.Rename(tempFile, filename)
}

// 辅助函数：从文件加载JSON
func loadJSONFromFile(filename string, data interface{}) error {
	// 读取文件内容
	jsonData, err := os.ReadFile(filename)
	if err != nil {
		return err
	}

	// 解码JSON数据
	return json.Unmarshal(jsonData, data)
}

// 辅助函数：检查文件是否存在
func fileExists(filename string) bool {
	_, err := os.Stat(filename)
	return err == nil
}