150 lines
4.0 KiB
Go
150 lines
4.0 KiB
Go
// 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))
|
||
}
|
||
}
|