fix 一括生成で止まらないようにする

Author: usuyuki

backend/cmd/subscriber/main.go

@@ -23,6 +23,7 @@ import (
 	"github.com/prometheus/client_golang/prometheus/promhttp"
 	"github.com/redis/rueidis"
 	"github.com/sirupsen/logrus"
+	"golang.org/x/time/rate"
 )
 
 // getTaskTimeout 環境変数からタスクタイムアウトを取得(デフォルト600秒)
@@ -167,6 +168,10 @@ func runSubscriber(app *container.SubscriberApp, cleanup *container.Cleanup, log
 		}
 	}()
 
+	// Gemini API レートリミッター（3000 RPM = 50 RPS、バースト50）
+	// チャンク分割・embedding生成の両方に適用してレートリミット超過を防ぐ
+	geminiRateLimiter := rate.NewLimiter(rate.Every(time.Minute/3000), 50)
+
 	logger.WithField("max_concurrent_jobs", app.SubscriberConfig.MaxConcurrentJobs).Info("Subscriber is listening for messages...")
 
 	// Create context for subscription that can be cancelled
@@ -239,7 +244,7 @@ func runSubscriber(app *container.SubscriberApp, cleanup *container.Cleanup, log
 						}()
 
 						start := time.Now()
-						err := processMessage(subCtx, app.DB, app.Redis, app.LLMFactory, app.LockService, msg.Message, logger)
+						err := processMessage(subCtx, app.DB, app.Redis, app.LLMFactory, app.LockService, geminiRateLimiter, msg.Message, logger)
 						duration := time.Since(start)
 
 						// メトリクス更新は processMessage 内で行う
@@ -336,7 +341,7 @@ func runSubscriber(app *container.SubscriberApp, cleanup *container.Cleanup, log
 	return nil
 }
 
-func processMessage(ctx context.Context, db *sql.DB, redisClient rueidis.Client, llmFactory container.LLMClientFactory, lockService container.LockService, payload string, logger *logrus.Entry) error {
+func processMessage(ctx context.Context, db *sql.DB, redisClient rueidis.Client, llmFactory container.LLMClientFactory, lockService container.LockService, geminiRateLimiter *rate.Limiter, payload string, logger *logrus.Entry) error {
 	start := time.Now()
 
 	// まずメッセージタイプを確認
@@ -413,7 +418,7 @@ func processMessage(ctx context.Context, db *sql.DB, redisClient rueidis.Client,
 			messagesProcessedCounter.WithLabelValues("diary_embedding", "error").Inc()
 			return fmt.Errorf("failed to unmarshal diary embedding message: %w", unmarshalErr)
 		}
-		err = generateDiaryEmbedding(ctx, db, llmFactory, message.UserID, message.DiaryID, logger)
+		err = generateDiaryEmbedding(ctx, db, llmFactory, geminiRateLimiter, message.UserID, message.DiaryID, logger)
 		if err != nil {
 			messagesProcessedCounter.WithLabelValues("diary_embedding", "error").Inc()
 		} else {
@@ -1091,7 +1096,7 @@ func generateDiaryHighlightWithLLM(ctx context.Context, db *sql.DB, llmFactory c
 	return highlights, nil
 }
 
-func generateDiaryEmbedding(ctx context.Context, db *sql.DB, llmFactory container.LLMClientFactory, userID, diaryID string, logger *logrus.Entry) error {
+func generateDiaryEmbedding(ctx context.Context, db *sql.DB, llmFactory container.LLMClientFactory, geminiRateLimiter *rate.Limiter, userID, diaryID string, logger *logrus.Entry) error {
 	logger.WithFields(logrus.Fields{
 		"user_id":  userID,
 		"diary_id": diaryID,
@@ -1140,6 +1145,9 @@ func generateDiaryEmbedding(ctx context.Context, db *sql.DB, llmFactory containe
 	}()
 
 	// 4. 日記を話題ごとのチャンクに分割する（失敗時は日記全体を1チャンクとしてフォールバック）
+	if err := geminiRateLimiter.Wait(ctx); err != nil {
+		return fmt.Errorf("rate limiter cancelled before SplitDiaryIntoChunks: %w", err)
+	}
 	chunkDataList, err := geminiClient.SplitDiaryIntoChunks(ctx, diaryContent)
 	if err != nil {
 		logger.WithFields(logrus.Fields{
@@ -1172,6 +1180,10 @@ func generateDiaryEmbedding(ctx context.Context, db *sql.DB, llmFactory containe
 			defer embWg.Done()
 			// 時間的クエリの精度向上のため日付情報を先頭に付与する
 			enrichedChunk := fmt.Sprintf("%d年%d月%d日の日記:\n%s", diaryDate.Year(), int(diaryDate.Month()), diaryDate.Day(), cd.Content)
+			if waitErr := geminiRateLimiter.Wait(ctx); waitErr != nil {
+				embResults[idx] = embeddingResult{err: fmt.Errorf("rate limiter cancelled before GenerateEmbedding for chunk %d: %w", idx, waitErr)}
+				return
+			}
 			embedding, err := geminiClient.GenerateEmbedding(ctx, enrichedChunk, true)
 			if err != nil {
 				embResults[idx] = embeddingResult{err: fmt.Errorf("failed to generate embedding for chunk %d: %w", idx, err)}

backend/go.mod

@@ -267,6 +267,7 @@ require (
 	golang.org/x/sync v0.20.0 // indirect
 	golang.org/x/sys v0.42.0 // indirect
 	golang.org/x/text v0.35.0 // indirect
+	golang.org/x/time v0.15.0 // indirect
 	golang.org/x/tools v0.42.0 // indirect
 	golang.org/x/tools/go/expect v0.1.1-deprecated // indirect
 	golang.org/x/tools/go/packages/packagestest v0.1.1-deprecated // indirect

backend/go.sum

@@ -921,6 +921,8 @@ golang.org/x/text v0.9.0/go.mod h1:e1OnstbJyHTd6l/uOt8jFFHp6TRDWZR/bV3emEE/zU8=
 golang.org/x/text v0.13.0/go.mod h1:TvPlkZtksWOMsz7fbANvkp4WM8x/WCo/om8BMLbz+aE=
 golang.org/x/text v0.35.0 h1:JOVx6vVDFokkpaq1AEptVzLTpDe9KGpj5tR4/X+ybL8=
 golang.org/x/text v0.35.0/go.mod h1:khi/HExzZJ2pGnjenulevKNX1W67CUy0AsXcNubPGCA=
+golang.org/x/time v0.15.0 h1:bbrp8t3bGUeFOx08pvsMYRTCVSMk89u4tKbNOZbp88U=
+golang.org/x/time v0.15.0/go.mod h1:Y4YMaQmXwGQZoFaVFk4YpCt4FLQMYKZe9oeV/f4MSno=
 golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
 golang.org/x/tools v0.0.0-20190311212946-11955173bddd/go.mod h1:LCzVGOaR6xXOjkQ3onu1FJEFr0SW1gC7cKk1uF8kGRs=
 golang.org/x/tools v0.0.0-20190425163242-31fd60d6bfdc/go.mod h1:RgjU9mgBXZiqYHBnxXauZ1Gv1EHHAz9KjViQ78xBX0Q=

backend/service/diary/service.go

@@ -1223,8 +1223,23 @@ func (s *DiaryEntry) RegenerateAllEmbeddings(
 		return nil, status.Error(codes.Internal, "Redis not configured")
 	}
 
-	// 同一ユーザーによる同時実行を防ぐ分散ロック（TTL: 10分）
-	regenLock := lock.NewDistributedLock(s.Redis, lock.EmbeddingRegenLockKey(userIDStr), 10*time.Minute)
+	// embedding未生成の日記ID一覧を取得（ロック取得前に件数を確認してTTLを決定する）
+	diaryIDs, err := database.DiaryIDsWithoutEmbeddings(ctx, s.DB, userID)
+	if err != nil {
+		return nil, status.Errorf(codes.Internal, "Failed to query diaries: %v", err)
+	}
+
+	// 件数に応じてロックTTLを動的に計算する（1件あたり10秒、最小10分・最大24時間）
+	lockTTL := time.Duration(len(diaryIDs)) * 10 * time.Second
+	if lockTTL < 10*time.Minute {
+		lockTTL = 10 * time.Minute
+	}
+	if lockTTL > 24*time.Hour {
+		lockTTL = 24 * time.Hour
+	}
+
+	// 同一ユーザーによる同時実行を防ぐ分散ロック
+	regenLock := lock.NewDistributedLock(s.Redis, lock.EmbeddingRegenLockKey(userIDStr), lockTTL)
 	acquired, err := regenLock.TryLock(ctx)
 	if err != nil {
 		return nil, status.Errorf(codes.Internal, "Failed to acquire lock: %v", err)
@@ -1238,12 +1253,6 @@ func (s *DiaryEntry) RegenerateAllEmbeddings(
 		}
 	}()
 
-	// embedding未生成の日記ID一覧を取得
-	diaryIDs, err := database.DiaryIDsWithoutEmbeddings(ctx, s.DB, userID)
-	if err != nil {
-		return nil, status.Errorf(codes.Internal, "Failed to query diaries: %v", err)
-	}
-
 	// 各日記のembedding生成メッセージをキューに追加
 	// 手動再生成のため当日の日記も即時処理する（on-saveとは異なる）
 	var count int32