diff --git a/backend/internal/repository/scheduler_cache.go b/backend/internal/repository/scheduler_cache.go index c8e1fe14e0..c68bcf96e6 100644 --- a/backend/internal/repository/scheduler_cache.go +++ b/backend/internal/repository/scheduler_cache.go @@ -4,6 +4,7 @@ import ( "context" "encoding/json" "fmt" + "log/slog" "strconv" "time" @@ -163,14 +164,15 @@ func (c *schedulerCache) SetSnapshot(ctx context.Context, bucket service.Schedul versionStr := strconv.FormatInt(version, 10) snapshotKey := schedulerSnapshotKey(bucket, versionStr) - if err := c.writeAccounts(ctx, accounts); err != nil { + cacheableAccounts, err := c.writeAccounts(ctx, accounts) + if err != nil { return err } - if len(accounts) > 0 { + if len(cacheableAccounts) > 0 { // 使用序号作为 score,保持数据库返回的排序语义。 - members := make([]redis.Z, 0, len(accounts)) - for idx, account := range accounts { + members := make([]redis.Z, 0, len(cacheableAccounts)) + for idx, account := range cacheableAccounts { members = append(members, redis.Z{ Score: float64(idx), Member: strconv.FormatInt(account.ID, 10), @@ -224,7 +226,14 @@ func (c *schedulerCache) SetAccount(ctx context.Context, account *service.Accoun if account == nil || account.ID <= 0 { return nil } - return c.writeAccounts(ctx, []service.Account{*account}) + cacheableAccounts, err := c.writeAccounts(ctx, []service.Account{*account}) + if err != nil { + return err + } + if len(cacheableAccounts) == 0 { + return c.DeleteAccount(ctx, account.ID) + } + return nil } func (c *schedulerCache) DeleteAccount(ctx context.Context, accountID int64) error { @@ -262,13 +271,14 @@ func (c *schedulerCache) UpdateLastUsed(ctx context.Context, updates map[int64]t return err } account.LastUsedAt = ptrTime(updates[ids[i]]) - updated, err := json.Marshal(account) + updated, metaPayload, err := marshalSchedulerCacheAccount(*account) if err != nil { - return err - } - metaPayload, err := json.Marshal(buildSchedulerMetadataAccount(*account)) - if err != nil { - return err + slog.Warn("scheduler cache removes account with unencodable payload", + "account_id", ids[i], + "error", err, + ) + pipe.Del(ctx, keys[i], schedulerAccountMetaKey(strconv.FormatInt(ids[i], 10))) + continue } pipe.Set(ctx, keys[i], updated, 0) pipe.Set(ctx, schedulerAccountMetaKey(strconv.FormatInt(ids[i], 10)), metaPayload, 0) @@ -359,12 +369,13 @@ func decodeCachedAccount(val any) (*service.Account, error) { return &account, nil } -func (c *schedulerCache) writeAccounts(ctx context.Context, accounts []service.Account) error { +func (c *schedulerCache) writeAccounts(ctx context.Context, accounts []service.Account) ([]service.Account, error) { if len(accounts) == 0 { - return nil + return nil, nil } pipe := c.rdb.Pipeline() + cacheableAccounts := make([]service.Account, 0, len(accounts)) pending := 0 flush := func() error { if pending == 0 { @@ -379,27 +390,43 @@ func (c *schedulerCache) writeAccounts(ctx context.Context, accounts []service.A } for _, account := range accounts { - fullPayload, err := json.Marshal(account) + fullPayload, metaPayload, err := marshalSchedulerCacheAccount(account) if err != nil { - return err - } - metaPayload, err := json.Marshal(buildSchedulerMetadataAccount(account)) - if err != nil { - return err + slog.Warn("scheduler cache skips account with unencodable payload", + "account_id", account.ID, + "error", err, + ) + continue } id := strconv.FormatInt(account.ID, 10) pipe.Set(ctx, schedulerAccountKey(id), fullPayload, 0) pipe.Set(ctx, schedulerAccountMetaKey(id), metaPayload, 0) + cacheableAccounts = append(cacheableAccounts, account) pending++ if pending >= c.writeChunkSize { if err := flush(); err != nil { - return err + return nil, err } } } - return flush() + if err := flush(); err != nil { + return nil, err + } + return cacheableAccounts, nil +} + +func marshalSchedulerCacheAccount(account service.Account) ([]byte, []byte, error) { + fullPayload, err := json.Marshal(account) + if err != nil { + return nil, nil, fmt.Errorf("marshal account: %w", err) + } + metaPayload, err := json.Marshal(buildSchedulerMetadataAccount(account)) + if err != nil { + return nil, nil, fmt.Errorf("marshal account metadata: %w", err) + } + return fullPayload, metaPayload, nil } func (c *schedulerCache) mgetChunked(ctx context.Context, keys []string) ([]any, error) { diff --git a/backend/internal/repository/scheduler_cache_unit_test.go b/backend/internal/repository/scheduler_cache_unit_test.go index 19c4cc4f36..ecca7f3892 100644 --- a/backend/internal/repository/scheduler_cache_unit_test.go +++ b/backend/internal/repository/scheduler_cache_unit_test.go @@ -3,12 +3,78 @@ package repository import ( + "context" "testing" + "time" "github.com/Wei-Shaw/sub2api/internal/service" + "github.com/alicebob/miniredis/v2" + "github.com/redis/go-redis/v9" "github.com/stretchr/testify/require" ) +func newSchedulerCacheUnit(t *testing.T) *schedulerCache { + t.Helper() + mr := miniredis.RunT(t) + rdb := redis.NewClient(&redis.Options{Addr: mr.Addr()}) + t.Cleanup(func() { _ = rdb.Close() }) + cache, ok := newSchedulerCacheWithChunkSizes(rdb, defaultSchedulerSnapshotMGetChunkSize, defaultSchedulerSnapshotWriteChunkSize).(*schedulerCache) + require.True(t, ok) + return cache +} + +func TestSchedulerCacheWriteAccountsSkipsUnencodableTimes(t *testing.T) { + ctx := context.Background() + cache := newSchedulerCacheUnit(t) + invalidTime := time.Date(10000, time.January, 1, 0, 0, 0, 0, time.UTC) + + cacheable, err := cache.writeAccounts(ctx, []service.Account{ + {ID: 111, Platform: service.PlatformOpenAI, Type: service.AccountTypeAPIKey}, + {ID: 112, Platform: service.PlatformOpenAI, Type: service.AccountTypeAPIKey, ExpiresAt: &invalidTime}, + }) + require.NoError(t, err) + require.Len(t, cacheable, 1) + require.Equal(t, int64(111), cacheable[0].ID) + + cached, err := cache.GetAccount(ctx, 111) + require.NoError(t, err) + require.NotNil(t, cached) + + invalid, err := cache.GetAccount(ctx, 112) + require.NoError(t, err) + require.Nil(t, invalid) +} + +func TestSchedulerCacheSetAccountClearsUnencodablePayload(t *testing.T) { + ctx := context.Background() + cache := newSchedulerCacheUnit(t) + + account := service.Account{ID: 113, Platform: service.PlatformOpenAI, Type: service.AccountTypeAPIKey} + require.NoError(t, cache.SetAccount(ctx, &account)) + + invalidTime := time.Date(10000, time.January, 1, 0, 0, 0, 0, time.UTC) + account.ExpiresAt = &invalidTime + require.NoError(t, cache.SetAccount(ctx, &account)) + + cached, err := cache.GetAccount(ctx, account.ID) + require.NoError(t, err) + require.Nil(t, cached) +} + +func TestSchedulerCacheUpdateLastUsedClearsUnencodablePayload(t *testing.T) { + ctx := context.Background() + cache := newSchedulerCacheUnit(t) + account := service.Account{ID: 114, Platform: service.PlatformOpenAI, Type: service.AccountTypeAPIKey} + require.NoError(t, cache.SetAccount(ctx, &account)) + + invalidTime := time.Date(10000, time.January, 1, 0, 0, 0, 0, time.UTC) + require.NoError(t, cache.UpdateLastUsed(ctx, map[int64]time.Time{account.ID: invalidTime})) + + cached, err := cache.GetAccount(ctx, account.ID) + require.NoError(t, err) + require.Nil(t, cached) +} + func TestBuildSchedulerMetadataAccount_KeepsOpenAIWSFlags(t *testing.T) { account := service.Account{ ID: 42,