From d049fee0845ad9fdb28d10592592301b47433089 Mon Sep 17 00:00:00 2001 From: ShukeBta <272197458+ShukeBta@users.noreply.github.com> Date: Sat, 27 Jun 2026 08:42:33 +0800 Subject: [PATCH] split token pending store helpers --- internal/service/token_svc.go | 135 ----------------------- internal/service/token_svc_pending.go | 147 ++++++++++++++++++++++++++ 2 files changed, 147 insertions(+), 135 deletions(-) create mode 100644 internal/service/token_svc_pending.go diff --git a/internal/service/token_svc.go b/internal/service/token_svc.go index 016e22e..5a6be94 100644 --- a/internal/service/token_svc.go +++ b/internal/service/token_svc.go @@ -26,8 +26,6 @@ const ( RefreshTokenLength = 32 ) -const loginRefreshTokenStoreTimeout = 750 * time.Millisecond - // Claims 是 JWT 载荷（复制自 middleware 以避免循环导入）。 type Claims struct { UserID string `json:"uid"` @@ -52,11 +50,6 @@ type TokenService struct { delayedStores map[string]pendingRefreshToken } -type pendingRefreshToken struct { - UserID string - ExpiresAt time.Time -} - // NewTokenService 创建令牌服务实例。 func NewTokenService(cfg *config.Config, log *zap.Logger, repo *repository.Container) *TokenService { return &TokenService{cfg: cfg, log: log, repo: repo, delayedStores: make(map[string]pendingRefreshToken)} @@ -130,56 +123,6 @@ func (s *TokenService) issuePair(ctx context.Context, userID, role, tier string, }, nil } -func (s *TokenService) storeRefreshTokenBestEffort(userID, tokenHash string, expiresAt time.Time) { - if !s.trackDelayedStore(userID, tokenHash, expiresAt) { - return - } - done := make(chan error, 1) - go func() { - ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) - defer cancel() - done <- s.storeRefreshToken(ctx, &model.RefreshToken{ - UserID: userID, - TokenHash: tokenHash, - ExpiresAt: expiresAt, - }) - }() - select { - case err := <-done: - s.finishBestEffortRefreshTokenStore(userID, tokenHash, expiresAt, err) - case <-time.After(loginRefreshTokenStoreTimeout): - if s.log != nil { - s.log.Warn("refresh token store delayed; login will continue", - zap.String("user_id", userID), - zap.Error(context.DeadlineExceeded)) - } - go func() { - err := <-done - s.finishBestEffortRefreshTokenStore(userID, tokenHash, expiresAt, err) - }() - } -} - -func (s *TokenService) finishBestEffortRefreshTokenStore(userID, tokenHash string, expiresAt time.Time, err error) { - if err == nil { - s.untrackDelayedStore(userID, tokenHash) - return - } - if repository.IsSQLiteBusyError(err) || errors.Is(err, context.DeadlineExceeded) || errors.Is(err, context.Canceled) { - if s.log != nil { - s.log.Warn("refresh token store delayed; login will continue", - zap.String("user_id", userID), - zap.Error(err)) - } - s.storeRefreshTokenEventually(userID, tokenHash, expiresAt) - return - } - s.untrackDelayedStore(userID, tokenHash) - if s.log != nil { - s.log.Warn("refresh token delayed store failed permanently", zap.String("user_id", userID), zap.Error(err)) - } -} - func (s *TokenService) storeRefreshToken(ctx context.Context, rt *model.RefreshToken) error { if err := s.repo.RefreshToken.Create(ctx, rt); err != nil { return err @@ -190,84 +133,6 @@ func (s *TokenService) storeRefreshToken(ctx context.Context, rt *model.RefreshT return nil } -func (s *TokenService) storeRefreshTokenEventually(userID, tokenHash string, expiresAt time.Time) { - defer s.untrackDelayedStore(userID, tokenHash) - delay := time.Second - for attempt := 1; attempt <= 8; attempt++ { - timer := time.NewTimer(delay) - <-timer.C - // 令牌可能已在等待期间被轮换/登出（从 pending 表移除）， - // 此时绝不能再写库，否则会复活一个已被替换的旧令牌。 - if _, stillPending := s.pendingDelayedStore(tokenHash); !stillPending { - return - } - ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) - err := s.storeRefreshToken(ctx, &model.RefreshToken{ - UserID: userID, - TokenHash: tokenHash, - ExpiresAt: expiresAt, - }) - cancel() - if err == nil { - return - } - if !repository.IsSQLiteBusyError(err) && !errors.Is(err, context.DeadlineExceeded) && !errors.Is(err, context.Canceled) { - if s.log != nil { - s.log.Warn("refresh token delayed store failed permanently", zap.String("user_id", userID), zap.Error(err)) - } - return - } - if s.log != nil && (attempt == 1 || attempt == 4 || attempt == 8) { - s.log.Warn("refresh token delayed store still waiting", - zap.String("user_id", userID), - zap.Int("attempt", attempt), - zap.Error(err)) - } - if delay < 60*time.Second { - delay *= 2 - } - } - if s.log != nil { - s.log.Warn("refresh token delayed store gave up", zap.String("user_id", userID)) - } -} - -func (s *TokenService) trackDelayedStore(userID, tokenHash string, expiresAt time.Time) bool { - if s == nil { - return false - } - s.delayedStoreMu.Lock() - defer s.delayedStoreMu.Unlock() - if s.delayedStores == nil { - s.delayedStores = make(map[string]pendingRefreshToken) - } - if _, ok := s.delayedStores[tokenHash]; ok { - return false - } - s.delayedStores[tokenHash] = pendingRefreshToken{UserID: userID, ExpiresAt: expiresAt} - return true -} - -func (s *TokenService) untrackDelayedStore(userID, tokenHash string) { - if s == nil { - return - } - s.delayedStoreMu.Lock() - delete(s.delayedStores, tokenHash) - s.delayedStoreMu.Unlock() -} - -// pendingDelayedStore 返回尚未落库的 refresh token 信息（如果存在）。 -func (s *TokenService) pendingDelayedStore(tokenHash string) (pendingRefreshToken, bool) { - if s == nil { - return pendingRefreshToken{}, false - } - s.delayedStoreMu.Lock() - defer s.delayedStoreMu.Unlock() - pending, ok := s.delayedStores[tokenHash] - return pending, ok -} - func (s *TokenService) maxActiveRefreshTokens(ctx context.Context) int { cfg := loadBotConfig(ctx, s.repo) if cfg.MaxLoggedClients < 1 { diff --git a/internal/service/token_svc_pending.go b/internal/service/token_svc_pending.go new file mode 100644 index 0000000..9ab296a --- /dev/null +++ b/internal/service/token_svc_pending.go @@ -0,0 +1,147 @@ +package service + +import ( + "context" + "errors" + "time" + + "go.uber.org/zap" + + "github.com/ShukeBta/MediaStationGo/internal/model" + "github.com/ShukeBta/MediaStationGo/internal/repository" +) + +const loginRefreshTokenStoreTimeout = 750 * time.Millisecond + +type pendingRefreshToken struct { + UserID string + ExpiresAt time.Time +} + +func (s *TokenService) storeRefreshTokenBestEffort(userID, tokenHash string, expiresAt time.Time) { + if !s.trackDelayedStore(userID, tokenHash, expiresAt) { + return + } + done := make(chan error, 1) + go func() { + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + done <- s.storeRefreshToken(ctx, &model.RefreshToken{ + UserID: userID, + TokenHash: tokenHash, + ExpiresAt: expiresAt, + }) + }() + select { + case err := <-done: + s.finishBestEffortRefreshTokenStore(userID, tokenHash, expiresAt, err) + case <-time.After(loginRefreshTokenStoreTimeout): + if s.log != nil { + s.log.Warn("refresh token store delayed; login will continue", + zap.String("user_id", userID), + zap.Error(context.DeadlineExceeded)) + } + go func() { + err := <-done + s.finishBestEffortRefreshTokenStore(userID, tokenHash, expiresAt, err) + }() + } +} + +func (s *TokenService) finishBestEffortRefreshTokenStore(userID, tokenHash string, expiresAt time.Time, err error) { + if err == nil { + s.untrackDelayedStore(userID, tokenHash) + return + } + if repository.IsSQLiteBusyError(err) || errors.Is(err, context.DeadlineExceeded) || errors.Is(err, context.Canceled) { + if s.log != nil { + s.log.Warn("refresh token store delayed; login will continue", + zap.String("user_id", userID), + zap.Error(err)) + } + s.storeRefreshTokenEventually(userID, tokenHash, expiresAt) + return + } + s.untrackDelayedStore(userID, tokenHash) + if s.log != nil { + s.log.Warn("refresh token delayed store failed permanently", zap.String("user_id", userID), zap.Error(err)) + } +} + +func (s *TokenService) storeRefreshTokenEventually(userID, tokenHash string, expiresAt time.Time) { + defer s.untrackDelayedStore(userID, tokenHash) + delay := time.Second + for attempt := 1; attempt <= 8; attempt++ { + timer := time.NewTimer(delay) + <-timer.C + // 令牌可能已在等待期间被轮换/登出（从 pending 表移除）， + // 此时绝不能再写库，否则会复活一个已被替换的旧令牌。 + if _, stillPending := s.pendingDelayedStore(tokenHash); !stillPending { + return + } + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + err := s.storeRefreshToken(ctx, &model.RefreshToken{ + UserID: userID, + TokenHash: tokenHash, + ExpiresAt: expiresAt, + }) + cancel() + if err == nil { + return + } + if !repository.IsSQLiteBusyError(err) && !errors.Is(err, context.DeadlineExceeded) && !errors.Is(err, context.Canceled) { + if s.log != nil { + s.log.Warn("refresh token delayed store failed permanently", zap.String("user_id", userID), zap.Error(err)) + } + return + } + if s.log != nil && (attempt == 1 || attempt == 4 || attempt == 8) { + s.log.Warn("refresh token delayed store still waiting", + zap.String("user_id", userID), + zap.Int("attempt", attempt), + zap.Error(err)) + } + if delay < 60*time.Second { + delay *= 2 + } + } + if s.log != nil { + s.log.Warn("refresh token delayed store gave up", zap.String("user_id", userID)) + } +} + +func (s *TokenService) trackDelayedStore(userID, tokenHash string, expiresAt time.Time) bool { + if s == nil { + return false + } + s.delayedStoreMu.Lock() + defer s.delayedStoreMu.Unlock() + if s.delayedStores == nil { + s.delayedStores = make(map[string]pendingRefreshToken) + } + if _, ok := s.delayedStores[tokenHash]; ok { + return false + } + s.delayedStores[tokenHash] = pendingRefreshToken{UserID: userID, ExpiresAt: expiresAt} + return true +} + +func (s *TokenService) untrackDelayedStore(userID, tokenHash string) { + if s == nil { + return + } + s.delayedStoreMu.Lock() + delete(s.delayedStores, tokenHash) + s.delayedStoreMu.Unlock() +} + +// pendingDelayedStore 返回尚未落库的 refresh token 信息（如果存在）。 +func (s *TokenService) pendingDelayedStore(tokenHash string) (pendingRefreshToken, bool) { + if s == nil { + return pendingRefreshToken{}, false + } + s.delayedStoreMu.Lock() + defer s.delayedStoreMu.Unlock() + pending, ok := s.delayedStores[tokenHash] + return pending, ok +}