New

internal/service/db_log_writer_service.go

@@ -0,0 +1,149 @@
+// Filename: internal/service/db_log_writer_service.go (全新文件)
+
+package service
+
+import (
+	"encoding/json"
+	"gemini-balancer/internal/models"
+	"gemini-balancer/internal/settings"
+	"gemini-balancer/internal/store"
+	"sync"
+	"time"
+
+	"github.com/sirupsen/logrus"
+	"gorm.io/gorm"
+)
+
+type DBLogWriterService struct {
+	db              *gorm.DB
+	store           store.Store
+	logger          *logrus.Entry
+	logBuffer       chan *models.RequestLog
+	stopChan        chan struct{}
+	wg              sync.WaitGroup
+	SettingsManager *settings.SettingsManager
+}
+
+func NewDBLogWriterService(db *gorm.DB, s store.Store, settings *settings.SettingsManager, logger *logrus.Logger) *DBLogWriterService {
+	cfg := settings.GetSettings()
+	bufferCapacity := cfg.LogBufferCapacity
+	if bufferCapacity <= 0 {
+		bufferCapacity = 1000
+	}
+	return &DBLogWriterService{
+		db:              db,
+		store:           s,
+		SettingsManager: settings,
+		logger:          logger.WithField("component", "DBLogWriter📝"),
+		// 使用配置值来创建缓冲区
+		logBuffer: make(chan *models.RequestLog, bufferCapacity),
+		stopChan:  make(chan struct{}),
+	}
+}
+
+func (s *DBLogWriterService) Start() {
+	s.wg.Add(2) // 一个用于事件监听，一个用于数据库写入
+
+	// 启动事件监听器
+	go s.eventListenerLoop()
+	// 启动数据库写入器
+	go s.dbWriterLoop()
+
+	s.logger.Info("DBLogWriterService started.")
+}
+
+func (s *DBLogWriterService) Stop() {
+	s.logger.Info("DBLogWriterService stopping...")
+	close(s.stopChan) // 通知所有goroutine停止
+	s.wg.Wait()       // 等待所有goroutine完成
+	s.logger.Info("DBLogWriterService stopped.")
+}
+
+// eventListenerLoop 负责从store接收事件并放入内存缓冲区
+func (s *DBLogWriterService) eventListenerLoop() {
+	defer s.wg.Done()
+
+	sub, err := s.store.Subscribe(models.TopicRequestFinished)
+	if err != nil {
+		s.logger.Fatalf("Failed to subscribe to topic %s: %v", models.TopicRequestFinished, err)
+		return
+	}
+	defer sub.Close()
+
+	s.logger.Info("Subscribed to request events for database logging.")
+
+	for {
+		select {
+		case msg := <-sub.Channel():
+			var event models.RequestFinishedEvent
+			if err := json.Unmarshal(msg.Payload, &event); err != nil {
+				s.logger.Errorf("Failed to unmarshal event for logging: %v", err)
+				continue
+			}
+
+			// 将事件中的日志部分放入缓冲区
+			select {
+			case s.logBuffer <- &event.RequestLog:
+			default:
+				s.logger.Warn("Log buffer is full. A log message might be dropped.")
+			}
+
+		case <-s.stopChan:
+			s.logger.Info("Event listener loop stopping.")
+			// 关闭缓冲区，以通知dbWriterLoop处理完剩余日志后退出
+			close(s.logBuffer)
+			return
+		}
+	}
+}
+
+// dbWriterLoop 负责从内存缓冲区批量读取日志并写入数据库
+func (s *DBLogWriterService) dbWriterLoop() {
+	defer s.wg.Done()
+
+	// 在启动时获取一次配置
+	cfg := s.SettingsManager.GetSettings()
+	batchSize := cfg.LogFlushBatchSize
+	if batchSize <= 0 {
+		batchSize = 100
+	}
+
+	flushTimeout := time.Duration(cfg.LogFlushIntervalSeconds) * time.Second
+	if flushTimeout <= 0 {
+		flushTimeout = 5 * time.Second
+	}
+	batch := make([]*models.RequestLog, 0, batchSize)
+	ticker := time.NewTicker(flushTimeout)
+	defer ticker.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)
+			}
+		}
+	}
+}
+
+// flushBatch 将一个批次的日志写入数据库
+func (s *DBLogWriterService) flushBatch(batch []*models.RequestLog) {
+	if err := s.db.CreateInBatches(batch, len(batch)).Error; err != nil {
+		s.logger.WithField("batch_size", len(batch)).WithError(err).Error("Failed to flush log batch to database.")
+	} else {
+		s.logger.Infof("Successfully flushed %d logs to database.", len(batch))
+	}
+}