Update Context for store

internal/service/analytics_service.go

@@ -1,8 +1,8 @@
 // Filename: internal/service/analytics_service.go
-
 package service
 
 import (
+	"context"
 	"encoding/json"
 	"fmt"
 	"gemini-balancer/internal/db/dialect"
@@ -43,7 +43,7 @@ func NewAnalyticsService(db *gorm.DB, s store.Store, logger *logrus.Logger, d di
 }
 
 func (s *AnalyticsService) Start() {
-	s.wg.Add(2) // 2 (flushLoop, eventListener)
+	s.wg.Add(2)
 	go s.flushLoop()
 	go s.eventListener()
 	s.logger.Info("AnalyticsService (Command Side) started.")
@@ -53,13 +53,13 @@ func (s *AnalyticsService) Stop() {
 	close(s.stopChan)
 	s.wg.Wait()
 	s.logger.Info("AnalyticsService stopped. Performing final data flush...")
-	s.flushToDB() // 停止前刷盘
+	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)
+	sub, err := s.store.Subscribe(context.Background(), models.TopicRequestFinished)
 	if err != nil {
 		s.logger.Fatalf("Failed to subscribe to topic %s: %v", models.TopicRequestFinished, err)
 		return
@@ -87,9 +87,10 @@ func (s *AnalyticsService) handleAnalyticsEvent(event *models.RequestFinishedEve
 	if event.RequestLog.GroupID == nil {
 		return
 	}
+	ctx := context.Background()
 	key := fmt.Sprintf("analytics:hourly:%s", time.Now().UTC().Format("2006-01-02T15"))
 	fieldPrefix := fmt.Sprintf("%d:%s", *event.RequestLog.GroupID, event.RequestLog.ModelName)
-	pipe := s.store.Pipeline()
+	pipe := s.store.Pipeline(ctx)
 	pipe.HIncrBy(key, fieldPrefix+":requests", 1)
 	if event.RequestLog.IsSuccess {
 		pipe.HIncrBy(key, fieldPrefix+":success", 1)
@@ -120,6 +121,7 @@ func (s *AnalyticsService) flushLoop() {
 }
 
 func (s *AnalyticsService) flushToDB() {
+	ctx := context.Background()
 	now := time.Now().UTC()
 	keysToFlush := []string{
 		fmt.Sprintf("analytics:hourly:%s", now.Add(-1*time.Hour).Format("2006-01-02T15")),
@@ -127,7 +129,7 @@ func (s *AnalyticsService) flushToDB() {
 	}
 
 	for _, key := range keysToFlush {
-		data, err := s.store.HGetAll(key)
+		data, err := s.store.HGetAll(ctx, key)
 		if err != nil || len(data) == 0 {
 			continue
 		}
@@ -136,15 +138,15 @@ func (s *AnalyticsService) flushToDB() {
 
 		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
+				[]string{"time", "group_id", "model_name"},
+				[]string{"request_count", "success_count", "prompt_tokens", "completion_tokens"},
 			)
-			err := s.db.Clauses(upsertClause).Create(&statsToFlush).Error
+			err := s.db.WithContext(ctx).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...)
+				_ = s.store.HDel(ctx, key, parsedFields...)
 			}
 		}
 	}