New

internal/service/analytics_service.go

@@ -0,0 +1,197 @@
+// Filename: internal/service/analytics_service.go
+
+package service
+
+import (
+	"encoding/json"
+	"fmt"
+	"gemini-balancer/internal/db/dialect"
+	"gemini-balancer/internal/models"
+	"gemini-balancer/internal/store"
+	"strconv"
+	"strings"
+	"sync"
+	"time"
+
+	"github.com/sirupsen/logrus"
+	"gorm.io/gorm"
+)
+
+const (
+	flushLoopInterval = 1 * time.Minute
+)
+
+type AnalyticsServiceLogger struct{ *logrus.Entry }
+
+type AnalyticsService struct {
+	db       *gorm.DB
+	store    store.Store
+	logger   *logrus.Entry
+	stopChan chan struct{}
+	wg       sync.WaitGroup
+	dialect  dialect.DialectAdapter
+}
+
+func NewAnalyticsService(db *gorm.DB, s store.Store, logger *logrus.Logger, d dialect.DialectAdapter) *AnalyticsService {
+	return &AnalyticsService{
+		db:       db,
+		store:    s,
+		logger:   logger.WithField("component", "Analytics📊"),
+		stopChan: make(chan struct{}),
+		dialect:  d,
+	}
+}
+
+func (s *AnalyticsService) Start() {
+	s.wg.Add(2) // 2 (flushLoop, eventListener)
+	go s.flushLoop()
+	go s.eventListener()
+	s.logger.Info("AnalyticsService (Command Side) started.")
+}
+
+func (s *AnalyticsService) Stop() {
+	close(s.stopChan)
+	s.wg.Wait()
+	s.logger.Info("AnalyticsService stopped. Performing final data flush...")
+	s.flushToDB() // 停止前刷盘
+	s.logger.Info("AnalyticsService final data flush completed.")
+}
+
+func (s *AnalyticsService) eventListener() {
+	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("AnalyticsService subscribed to request events.")
+
+	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 analytics event: %v", err)
+				continue
+			}
+			s.handleAnalyticsEvent(&event)
+		case <-s.stopChan:
+			s.logger.Info("AnalyticsService stopping event listener.")
+			return
+		}
+	}
+}
+
+func (s *AnalyticsService) handleAnalyticsEvent(event *models.RequestFinishedEvent) {
+	if event.GroupID == 0 {
+		return
+	}
+	key := fmt.Sprintf("analytics:hourly:%s", time.Now().UTC().Format("2006-01-02T15"))
+	fieldPrefix := fmt.Sprintf("%d:%s", event.GroupID, event.ModelName)
+
+	pipe := s.store.Pipeline()
+	pipe.HIncrBy(key, fieldPrefix+":requests", 1)
+	if event.IsSuccess {
+		pipe.HIncrBy(key, fieldPrefix+":success", 1)
+	}
+	if event.PromptTokens > 0 {
+		pipe.HIncrBy(key, fieldPrefix+":prompt", int64(event.PromptTokens))
+	}
+	if event.CompletionTokens > 0 {
+		pipe.HIncrBy(key, fieldPrefix+":completion", int64(event.CompletionTokens))
+	}
+
+	if err := pipe.Exec(); err != nil {
+		s.logger.Warnf("[%s] Failed to record analytics event to store for group %d: %v", event.CorrelationID, event.GroupID, err)
+	}
+}
+
+func (s *AnalyticsService) flushLoop() {
+	defer s.wg.Done()
+	ticker := time.NewTicker(flushLoopInterval)
+	defer ticker.Stop()
+	for {
+		select {
+		case <-ticker.C:
+			s.flushToDB()
+		case <-s.stopChan:
+			return
+		}
+	}
+}
+
+func (s *AnalyticsService) flushToDB() {
+	now := time.Now().UTC()
+	keysToFlush := []string{
+		fmt.Sprintf("analytics:hourly:%s", now.Add(-1*time.Hour).Format("2006-01-02T15")),
+		fmt.Sprintf("analytics:hourly:%s", now.Format("2006-01-02T15")),
+	}
+
+	for _, key := range keysToFlush {
+		data, err := s.store.HGetAll(key)
+		if err != nil || len(data) == 0 {
+			continue
+		}
+
+		statsToFlush, parsedFields := s.parseStatsFromHash(now.Truncate(time.Hour), data)
+
+		if len(statsToFlush) > 0 {
+			upsertClause := s.dialect.OnConflictUpdateAll(
+				[]string{"time", "group_id", "model_name"},                                       // conflict columns
+				[]string{"request_count", "success_count", "prompt_tokens", "completion_tokens"}, // update columns
+			)
+			err := s.db.Clauses(upsertClause).Create(&statsToFlush).Error
+			if err != nil {
+				s.logger.Errorf("Failed to flush analytics data for key %s: %v", key, err)
+			} else {
+				s.logger.Infof("Successfully flushed %d records from key %s.", len(statsToFlush), key)
+				_ = s.store.HDel(key, parsedFields...)
+			}
+		}
+	}
+}
+
+func (s *AnalyticsService) parseStatsFromHash(t time.Time, data map[string]string) ([]models.StatsHourly, []string) {
+	tempAggregator := make(map[string]*models.StatsHourly)
+	var parsedFields []string
+	for field, valueStr := range data {
+		parts := strings.Split(field, ":")
+		if len(parts) != 3 {
+			continue
+		}
+		groupIDStr, modelName, counterType := parts[0], parts[1], parts[2]
+
+		aggKey := groupIDStr + ":" + modelName
+		if _, ok := tempAggregator[aggKey]; !ok {
+			gid, err := strconv.Atoi(groupIDStr)
+			if err != nil {
+				continue
+			}
+			tempAggregator[aggKey] = &models.StatsHourly{
+				Time:      t,
+				GroupID:   uint(gid),
+				ModelName: modelName,
+			}
+		}
+		val, _ := strconv.ParseInt(valueStr, 10, 64)
+		switch counterType {
+		case "requests":
+			tempAggregator[aggKey].RequestCount = val
+		case "success":
+			tempAggregator[aggKey].SuccessCount = val
+		case "prompt":
+			tempAggregator[aggKey].PromptTokens = val
+		case "completion":
+			tempAggregator[aggKey].CompletionTokens = val
+		}
+		parsedFields = append(parsedFields, field)
+	}
+	var result []models.StatsHourly
+	for _, stats := range tempAggregator {
+		if stats.RequestCount > 0 {
+			result = append(result, *stats)
+		}
+	}
+	return result, parsedFields
+}