gemini-banlancer/internal/service/db_log_writer_service.go

// Filename: internal/service/db_log_writer_service.go
package service

import (
	"context"
	"encoding/json"
	"sync"
	"sync/atomic"
	"time"

	"gemini-balancer/internal/models"
	"gemini-balancer/internal/settings"
	"gemini-balancer/internal/store"

	"github.com/sirupsen/logrus"
	"gorm.io/gorm"
)

type DBLogWriterService struct {
	db              *gorm.DB
	store           store.Store
	logger          *logrus.Entry
	settingsManager *settings.SettingsManager

	logBuffer chan *models.RequestLog
	stopChan  chan struct{}
	wg        sync.WaitGroup
	ctx       context.Context
	cancel    context.CancelFunc

	// 统计指标
	totalReceived  atomic.Uint64
	totalFlushed   atomic.Uint64
	totalDropped   atomic.Uint64
	flushCount     atomic.Uint64
	lastFlushTime  time.Time
	lastFlushMutex sync.RWMutex
}

func NewDBLogWriterService(
	db *gorm.DB,
	s store.Store,
	settingsManager *settings.SettingsManager,
	logger *logrus.Logger,
) *DBLogWriterService {
	cfg := settingsManager.GetSettings()
	bufferCapacity := cfg.LogBufferCapacity
	if bufferCapacity <= 0 {
		bufferCapacity = 1000
	}

	ctx, cancel := context.WithCancel(context.Background())

	return &DBLogWriterService{
		db:              db,
		store:           s,
		settingsManager: settingsManager,
		logger:          logger.WithField("component", "DBLogWriter📝"),
		logBuffer:       make(chan *models.RequestLog, bufferCapacity),
		stopChan:        make(chan struct{}),
		ctx:             ctx,
		cancel:          cancel,
		lastFlushTime:   time.Now(),
	}
}

func (s *DBLogWriterService) Start() {
	s.wg.Add(2)
	go s.eventListenerLoop()
	go s.dbWriterLoop()

	// 定期输出统计信息
	s.wg.Add(1)
	go s.metricsReporter()

	s.logger.WithFields(logrus.Fields{
		"buffer_capacity": cap(s.logBuffer),
	}).Info("DBLogWriterService started")
}

func (s *DBLogWriterService) Stop() {
	s.logger.Info("DBLogWriterService stopping...")
	close(s.stopChan)
	s.cancel() // 取消上下文
	s.wg.Wait()

	// 输出最终统计
	s.logger.WithFields(logrus.Fields{
		"total_received": s.totalReceived.Load(),
		"total_flushed":  s.totalFlushed.Load(),
		"total_dropped":  s.totalDropped.Load(),
		"flush_count":    s.flushCount.Load(),
	}).Info("DBLogWriterService stopped")
}

// 事件监听循环
func (s *DBLogWriterService) eventListenerLoop() {
	defer s.wg.Done()

	sub, err := s.store.Subscribe(s.ctx, models.TopicRequestFinished)
	if err != nil {
		s.logger.WithError(err).Error("Failed to subscribe to request events, log writing disabled")
		return
	}
	defer func() {
		if err := sub.Close(); err != nil {
			s.logger.WithError(err).Warn("Failed to close subscription")
		}
	}()

	s.logger.Info("Subscribed to request events for database logging")

	for {
		select {
		case msg := <-sub.Channel():
			s.handleMessage(msg)

		case <-s.stopChan:
			s.logger.Info("Event listener loop stopping")
			close(s.logBuffer)
			return

		case <-s.ctx.Done():
			s.logger.Info("Event listener context cancelled")
			close(s.logBuffer)
			return
		}
	}
}

// 处理单条消息
func (s *DBLogWriterService) handleMessage(msg *store.Message) {
	var event models.RequestFinishedEvent
	if err := json.Unmarshal(msg.Payload, &event); err != nil {
		s.logger.WithError(err).Error("Failed to unmarshal request event")
		return
	}

	s.totalReceived.Add(1)

	select {
	case s.logBuffer <- &event.RequestLog:
		// 成功入队
	default:
		// 缓冲区满，丢弃日志
		dropped := s.totalDropped.Add(1)
		if dropped%100 == 1 { // 每100条丢失输出一次警告
			s.logger.WithFields(logrus.Fields{
				"total_dropped":   dropped,
				"buffer_capacity": cap(s.logBuffer),
				"buffer_len":      len(s.logBuffer),
			}).Warn("Log buffer full, messages being dropped")
		}
	}
}

// 数据库写入循环
func (s *DBLogWriterService) dbWriterLoop() {
	defer s.wg.Done()

	cfg := s.settingsManager.GetSettings()
	batchSize := cfg.LogFlushBatchSize
	if batchSize <= 0 {
		batchSize = 100
	}
	flushInterval := time.Duration(cfg.LogFlushIntervalSeconds) * time.Second
	if flushInterval <= 0 {
		flushInterval = 5 * time.Second
	}

	s.logger.WithFields(logrus.Fields{
		"batch_size":     batchSize,
		"flush_interval": flushInterval,
	}).Info("DB writer loop started")

	batch := make([]*models.RequestLog, 0, batchSize)
	ticker := time.NewTicker(flushInterval)
	defer ticker.Stop()

	// 配置热更新检查（每分钟）
	configTicker := time.NewTicker(1 * time.Minute)
	defer configTicker.Stop()

	for {
		select {
		case logEntry, ok := <-s.logBuffer:
			if !ok {
				// 通道关闭，刷新剩余日志
				if len(batch) > 0 {
					s.flushBatch(batch)
				}
				s.logger.Info("DB writer loop finished")
				return
			}

			batch = append(batch, logEntry)
			if len(batch) >= batchSize {
				s.flushBatch(batch)
				batch = make([]*models.RequestLog, 0, batchSize)
			}

		case <-ticker.C:
			if len(batch) > 0 {
				s.flushBatch(batch)
				batch = make([]*models.RequestLog, 0, batchSize)
			}

		case <-configTicker.C:
			// 热更新配置
			cfg := s.settingsManager.GetSettings()
			newBatchSize := cfg.LogFlushBatchSize
			if newBatchSize <= 0 {
				newBatchSize = 100
			}
			newFlushInterval := time.Duration(cfg.LogFlushIntervalSeconds) * time.Second
			if newFlushInterval <= 0 {
				newFlushInterval = 5 * time.Second
			}

			if newBatchSize != batchSize {
				s.logger.WithFields(logrus.Fields{
					"old": batchSize,
					"new": newBatchSize,
				}).Info("Batch size updated")
				batchSize = newBatchSize
				if len(batch) >= batchSize {
					s.flushBatch(batch)
					batch = make([]*models.RequestLog, 0, batchSize)
				}
			}

			if newFlushInterval != flushInterval {
				s.logger.WithFields(logrus.Fields{
					"old": flushInterval,
					"new": newFlushInterval,
				}).Info("Flush interval updated")
				flushInterval = newFlushInterval
				ticker.Reset(flushInterval)
			}
		}
	}
}

// 批量刷写到数据库
func (s *DBLogWriterService) flushBatch(batch []*models.RequestLog) {
	if len(batch) == 0 {
		return
	}

	start := time.Now()
	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
	defer cancel()

	err := s.db.WithContext(ctx).CreateInBatches(batch, len(batch)).Error
	duration := time.Since(start)

	s.lastFlushMutex.Lock()
	s.lastFlushTime = time.Now()
	s.lastFlushMutex.Unlock()

	if err != nil {
		s.logger.WithFields(logrus.Fields{
			"batch_size": len(batch),
			"duration":   duration,
		}).WithError(err).Error("Failed to flush log batch to database")
	} else {
		flushed := s.totalFlushed.Add(uint64(len(batch)))
		flushCount := s.flushCount.Add(1)

		// 只在慢写入或大批量时输出日志
		if duration > 1*time.Second || len(batch) > 500 {
			s.logger.WithFields(logrus.Fields{
				"batch_size":    len(batch),
				"duration":      duration,
				"total_flushed": flushed,
				"flush_count":   flushCount,
			}).Info("Log batch flushed to database")
		} else {
			s.logger.WithFields(logrus.Fields{
				"batch_size": len(batch),
				"duration":   duration,
			}).Debug("Log batch flushed to database")
		}
	}
}

// 定期输出统计信息
func (s *DBLogWriterService) metricsReporter() {
	defer s.wg.Done()

	ticker := time.NewTicker(5 * time.Minute)
	defer ticker.Stop()

	for {
		select {
		case <-ticker.C:
			s.reportMetrics()
		case <-s.stopChan:
			return
		case <-s.ctx.Done():
			return
		}
	}
}

func (s *DBLogWriterService) reportMetrics() {
	s.lastFlushMutex.RLock()
	lastFlush := s.lastFlushTime
	s.lastFlushMutex.RUnlock()

	received := s.totalReceived.Load()
	flushed := s.totalFlushed.Load()
	dropped := s.totalDropped.Load()
	pending := uint64(len(s.logBuffer))

	s.logger.WithFields(logrus.Fields{
		"received":     received,
		"flushed":      flushed,
		"dropped":      dropped,
		"pending":      pending,
		"flush_count":  s.flushCount.Load(),
		"last_flush":   time.Since(lastFlush).Round(time.Second),
		"buffer_usage": float64(pending) / float64(cap(s.logBuffer)) * 100,
		"success_rate": float64(flushed) / float64(received) * 100,
	}).Info("DBLogWriter metrics")
}

// GetMetrics 返回当前统计指标（供监控使用）
func (s *DBLogWriterService) GetMetrics() map[string]interface{} {
	s.lastFlushMutex.RLock()
	lastFlush := s.lastFlushTime
	s.lastFlushMutex.RUnlock()

	return map[string]interface{}{
		"total_received":  s.totalReceived.Load(),
		"total_flushed":   s.totalFlushed.Load(),
		"total_dropped":   s.totalDropped.Load(),
		"flush_count":     s.flushCount.Load(),
		"buffer_pending":  len(s.logBuffer),
		"buffer_capacity": cap(s.logBuffer),
		"last_flush_ago":  time.Since(lastFlush).Seconds(),
	}
}