wormhole-server/pkg/memory/memory.go

package memory

import (
	"runtime"
	"sync"
	"time"

	"github.com/sirupsen/logrus"
)

// BufferPool 缓冲区池
type BufferPool struct {
	pools  map[int]*sync.Pool
	sizes  []int
	logger *logrus.Logger
	stats  BufferStats
}

// BufferStats 缓冲区统计
type BufferStats struct {
	mu          sync.RWMutex
	Gets        int64 `json:"gets"`
	Puts        int64 `json:"puts"`
	News        int64 `json:"news"`
	Reuses      int64 `json:"reuses"`
	TotalAlloc  int64 `json:"totalAlloc"`
	TotalReused int64 `json:"totalReused"`
}

// MemoryMonitor 内存监控器
type MemoryMonitor struct {
	logger     *logrus.Logger
	ticker     *time.Ticker
	stopCh     chan struct{}
	thresholds Thresholds
	lastStats  runtime.MemStats
	callbacks  []MemoryCallback
	mu         sync.RWMutex
}

// Thresholds 内存阈值配置
type Thresholds struct {
	HeapAllocMB     int64 `json:"heapAllocMB"`
	HeapSysMB       int64 `json:"heapSysMB"`
	GCPercent       int   `json:"gcPercent"`
	ForceGCThreshMB int64 `json:"forceGCThreshMB"`
}

// MemoryCallback 内存回调函数
type MemoryCallback func(stats runtime.MemStats)

// Config 内存管理配置
type Config struct {
	BufferSizes        []int         `json:"bufferSizes"`
	MonitorInterval    time.Duration `json:"monitorInterval"`
	Thresholds         Thresholds    `json:"thresholds"`
	EnableAutoGC       bool          `json:"enableAutoGC"`
	EnableOptimization bool          `json:"enableOptimization"`
}

// Manager 内存管理器
type Manager struct {
	bufferPool *BufferPool
	monitor    *MemoryMonitor
	config     Config
	logger     *logrus.Logger
}

// NewManager 创建内存管理器
func NewManager(config Config, logger *logrus.Logger) *Manager {
	// 设置默认值
	if len(config.BufferSizes) == 0 {
		config.BufferSizes = []int{512, 1024, 2048, 4096, 8192, 16384, 32768, 65536}
	}
	if config.MonitorInterval == 0 {
		config.MonitorInterval = 30 * time.Second
	}
	if config.Thresholds.HeapAllocMB == 0 {
		config.Thresholds.HeapAllocMB = 100
	}
	if config.Thresholds.HeapSysMB == 0 {
		config.Thresholds.HeapSysMB = 200
	}
	if config.Thresholds.GCPercent == 0 {
		config.Thresholds.GCPercent = 100
	}
	if config.Thresholds.ForceGCThreshMB == 0 {
		config.Thresholds.ForceGCThreshMB = 500
	}

	manager := &Manager{
		config: config,
		logger: logger,
	}

	// 创建缓冲区池
	if config.EnableOptimization {
		manager.bufferPool = NewBufferPool(config.BufferSizes, logger)
	}

	// 创建内存监控器
	manager.monitor = NewMemoryMonitor(config.MonitorInterval, config.Thresholds, logger)

	// 启用自动GC
	if config.EnableAutoGC {
		manager.monitor.AddCallback(manager.autoGCCallback)
	}

	return manager
}

// NewBufferPool 创建缓冲区池
func NewBufferPool(sizes []int, logger *logrus.Logger) *BufferPool {
	bp := &BufferPool{
		pools:  make(map[int]*sync.Pool),
		sizes:  make([]int, len(sizes)),
		logger: logger,
	}

	copy(bp.sizes, sizes)

	// 为每个大小创建池
	for _, size := range sizes {
		sz := size // 捕获循环变量
		bp.pools[sz] = &sync.Pool{
			New: func() interface{} {
				bp.stats.incNews()
				bp.stats.addTotalAlloc(int64(sz))
				return make([]byte, sz)
			},
		}
	}

	return bp
}

// Get 获取缓冲区
func (bp *BufferPool) Get(size int) []byte {
	bp.stats.incGets()

	// 找到最合适的大小
	for _, poolSize := range bp.sizes {
		if size <= poolSize {
			buf := bp.pools[poolSize].Get().([]byte)
			bp.stats.incReuses()
			bp.stats.addTotalReused(int64(poolSize))
			return buf[:size]
		}
	}

	// 没有合适的池，创建新缓冲区
	bp.stats.incNews()
	bp.stats.addTotalAlloc(int64(size))
	return make([]byte, size)
}

// Put 归还缓冲区
func (bp *BufferPool) Put(buf []byte) {
	if buf == nil {
		return
	}

	bp.stats.incPuts()
	capacity := cap(buf)

	// 找到对应的池
	for _, poolSize := range bp.sizes {
		if capacity == poolSize {
			bp.pools[poolSize].Put(buf[:poolSize])
			return
		}
	}
}

// GetStats 获取缓冲区统计
func (bp *BufferPool) GetStats() BufferStats {
	bp.stats.mu.RLock()
	defer bp.stats.mu.RUnlock()
	return bp.stats
}

// NewMemoryMonitor 创建内存监控器
func NewMemoryMonitor(interval time.Duration, thresholds Thresholds, logger *logrus.Logger) *MemoryMonitor {
	monitor := &MemoryMonitor{
		logger:     logger,
		ticker:     time.NewTicker(interval),
		stopCh:     make(chan struct{}),
		thresholds: thresholds,
		callbacks:  make([]MemoryCallback, 0),
	}

	// 启动监控
	go monitor.run()

	return monitor
}

// AddCallback 添加内存回调
func (mm *MemoryMonitor) AddCallback(callback MemoryCallback) {
	mm.mu.Lock()
	mm.callbacks = append(mm.callbacks, callback)
	mm.mu.Unlock()
}

// Stop 停止监控
func (mm *MemoryMonitor) Stop() {
	close(mm.stopCh)
	mm.ticker.Stop()
}

// run 运行监控
func (mm *MemoryMonitor) run() {
	for {
		select {
		case <-mm.stopCh:
			return
		case <-mm.ticker.C:
			mm.check()
		}
	}
}

// check 检查内存状态
func (mm *MemoryMonitor) check() {
	var stats runtime.MemStats
	runtime.ReadMemStats(&stats)

	// 记录内存统计
	heapAllocMB := stats.HeapAlloc / 1024 / 1024
	heapSysMB := stats.HeapSys / 1024 / 1024

	mm.logger.WithFields(logrus.Fields{
		"heap_alloc_mb": heapAllocMB,
		"heap_sys_mb":   heapSysMB,
		"gc_num":        stats.NumGC,
		"goroutines":    runtime.NumGoroutine(),
	}).Debug("Memory stats")

	// 检查阈值
	if int64(heapAllocMB) > mm.thresholds.HeapAllocMB {
		mm.logger.WithField("heap_alloc_mb", heapAllocMB).Warn("Heap allocation threshold exceeded")
	}

	if int64(heapSysMB) > mm.thresholds.HeapSysMB {
		mm.logger.WithField("heap_sys_mb", heapSysMB).Warn("Heap system threshold exceeded")
	}

	// 强制GC阈值
	if int64(heapAllocMB) > mm.thresholds.ForceGCThreshMB {
		mm.logger.WithField("heap_alloc_mb", heapAllocMB).Info("Force GC triggered")
		runtime.GC()
	}

	// 调用回调函数
	mm.mu.RLock()
	for _, callback := range mm.callbacks {
		go callback(stats)
	}
	mm.mu.RUnlock()

	mm.lastStats = stats
}

// GetStats 获取最新内存统计
func (mm *MemoryMonitor) GetStats() runtime.MemStats {
	return mm.lastStats
}

// Manager 方法

// GetBuffer 获取缓冲区
func (m *Manager) GetBuffer(size int) []byte {
	if m.bufferPool != nil {
		return m.bufferPool.Get(size)
	}
	return make([]byte, size)
}

// PutBuffer 归还缓冲区
func (m *Manager) PutBuffer(buf []byte) {
	if m.bufferPool != nil {
		m.bufferPool.Put(buf)
	}
}

// GetBufferStats 获取缓冲区统计
func (m *Manager) GetBufferStats() BufferStats {
	if m.bufferPool != nil {
		return m.bufferPool.GetStats()
	}
	return BufferStats{}
}

// GetMemoryStats 获取内存统计
func (m *Manager) GetMemoryStats() runtime.MemStats {
	return m.monitor.GetStats()
}

// Stop 停止内存管理器
func (m *Manager) Stop() {
	if m.monitor != nil {
		m.monitor.Stop()
	}
}

// ForceGC 强制垃圾回收
func (m *Manager) ForceGC() {
	m.logger.Info("Manual GC triggered")
	runtime.GC()
}

// autoGCCallback 自动GC回调
func (m *Manager) autoGCCallback(stats runtime.MemStats) {
	heapAllocMB := stats.HeapAlloc / 1024 / 1024

	// 当堆分配超过阈值时触发GC
	if int64(heapAllocMB) > m.config.Thresholds.ForceGCThreshMB {
		m.ForceGC()
	}
}

// GetOverallStats 获取总体统计信息
func (m *Manager) GetOverallStats() map[string]interface{} {
	memStats := m.GetMemoryStats()
	bufferStats := m.GetBufferStats()

	return map[string]interface{}{
		"memory": map[string]interface{}{
			"heap_alloc_mb": memStats.HeapAlloc / 1024 / 1024,
			"heap_sys_mb":   memStats.HeapSys / 1024 / 1024,
			"gc_num":        memStats.NumGC,
			"goroutines":    runtime.NumGoroutine(),
		},
		"buffers": bufferStats,
		"config": map[string]interface{}{
			"optimization_enabled": m.config.EnableOptimization,
			"auto_gc_enabled":      m.config.EnableAutoGC,
			"buffer_sizes":         m.config.BufferSizes,
		},
	}
}

// BufferStats 方法

func (bs *BufferStats) incGets() {
	bs.mu.Lock()
	bs.Gets++
	bs.mu.Unlock()
}

func (bs *BufferStats) incPuts() {
	bs.mu.Lock()
	bs.Puts++
	bs.mu.Unlock()
}

func (bs *BufferStats) incNews() {
	bs.mu.Lock()
	bs.News++
	bs.mu.Unlock()
}

func (bs *BufferStats) incReuses() {
	bs.mu.Lock()
	bs.Reuses++
	bs.mu.Unlock()
}

func (bs *BufferStats) addTotalAlloc(size int64) {
	bs.mu.Lock()
	bs.TotalAlloc += size
	bs.mu.Unlock()
}

func (bs *BufferStats) addTotalReused(size int64) {
	bs.mu.Lock()
	bs.TotalReused += size
	bs.mu.Unlock()
}