880 lines
34 KiB
Go
880 lines
34 KiB
Go
|
//go:build integration
|
||
|
|
|
||
|
|
package repository
|
||
|
|
|
||
|
|
import (
|
||
|
|
"context"
|
||
|
|
"errors"
|
||
|
|
"fmt"
|
||
|
|
"strconv"
|
||
|
|
"testing"
|
||
|
|
"time"
|
||
|
|
|
||
|
|
"github.com/Wei-Shaw/sub2api/internal/service"
|
||
|
|
"github.com/redis/go-redis/v9"
|
||
|
|
"github.com/stretchr/testify/require"
|
||
|
|
"github.com/stretchr/testify/suite"
|
||
|
|
)
|
||
|
|
|
||
|
|
// 测试用 TTL 配置（15 分钟，与默认值一致）
|
||
|
|
const testSlotTTLMinutes = 15
|
||
|
|
|
||
|
|
// 测试用 TTL Duration，用于 TTL 断言
|
||
|
|
var testSlotTTL = time.Duration(testSlotTTLMinutes) * time.Minute
|
||
|
|
|
||
|
|
type ConcurrencyCacheSuite struct {
|
||
|
|
IntegrationRedisSuite
|
||
|
|
cache service.ConcurrencyCache
|
||
|
|
rawCache *concurrencyCache
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestConcurrencyCacheSuite(t *testing.T) {
|
||
|
|
suite.Run(t, new(ConcurrencyCacheSuite))
|
||
|
|
}
|
||
|
|
|
||
|
|
func (s *ConcurrencyCacheSuite) SetupTest() {
|
||
|
|
s.IntegrationRedisSuite.SetupTest()
|
||
|
|
s.rawCache = NewConcurrencyCache(s.rdb, testSlotTTLMinutes, int(testSlotTTL.Seconds())).(*concurrencyCache)
|
||
|
|
s.cache = s.rawCache
|
||
|
|
}
|
||
|
|
|
||
|
|
type apiKeyConcurrencyCacheForTest interface {
|
||
|
|
TrackAPIKeySlot(ctx context.Context, apiKeyID int64, requestID string) error
|
||
|
|
ReleaseAPIKeySlot(ctx context.Context, apiKeyID int64, requestID string) error
|
||
|
|
GetAPIKeyConcurrencyBatch(ctx context.Context, apiKeyIDs []int64) (map[int64]int, error)
|
||
|
|
}
|
||
|
|
|
||
|
|
func (s *ConcurrencyCacheSuite) apiKeyConcurrencyCache() apiKeyConcurrencyCacheForTest {
|
||
|
|
cache, ok := s.cache.(apiKeyConcurrencyCacheForTest)
|
||
|
|
require.True(s.T(), ok)
|
||
|
|
return cache
|
||
|
|
}
|
||
|
|
|
||
|
|
func (s *ConcurrencyCacheSuite) TestOpenAIWSIngressAPIKeySlot_HardLimitRefreshAndRelease() {
|
||
|
|
apiKeyID := int64(9011)
|
||
|
|
firstLeaseID := "ingress-first"
|
||
|
|
secondLeaseID := "ingress-second"
|
||
|
|
|
||
|
|
ok, err := s.rawCache.AcquireOpenAIWSIngressLease(s.ctx, apiKeyID, 1, firstLeaseID)
|
||
|
|
require.NoError(s.T(), err)
|
||
|
|
require.True(s.T(), ok)
|
||
|
|
|
||
|
|
ok, err = s.rawCache.AcquireOpenAIWSIngressLease(s.ctx, apiKeyID, 1, secondLeaseID)
|
||
|
|
require.NoError(s.T(), err)
|
||
|
|
require.False(s.T(), ok, "a second live session must not exceed the API key limit")
|
||
|
|
|
||
|
|
ok, err = s.rawCache.RefreshOpenAIWSIngressLease(s.ctx, apiKeyID, firstLeaseID)
|
||
|
|
require.NoError(s.T(), err)
|
||
|
|
require.True(s.T(), ok, "the current owner must be able to refresh its lease")
|
||
|
|
|
||
|
|
require.NoError(s.T(), s.rawCache.ReleaseOpenAIWSIngressLease(s.ctx, apiKeyID, firstLeaseID))
|
||
|
|
ok, err = s.rawCache.AcquireOpenAIWSIngressLease(s.ctx, apiKeyID, 1, secondLeaseID)
|
||
|
|
require.NoError(s.T(), err)
|
||
|
|
require.True(s.T(), ok, "released capacity must become available immediately")
|
||
|
|
}
|
||
|
|
|
||
|
|
func (s *ConcurrencyCacheSuite) TestOpenAIWSIngressAPIKeySlot_ReapsCrashedLeaseWithoutDeletingLiveOtherInstance() {
|
||
|
|
apiKeyID := int64(9012)
|
||
|
|
key := openAIWSIngressLeaseKey(apiKeyID)
|
||
|
|
now, err := s.rawCache.redisUnixSeconds(s.ctx)
|
||
|
|
require.NoError(s.T(), err)
|
||
|
|
require.NoError(s.T(), s.rdb.ZAdd(s.ctx, key,
|
||
|
|
redis.Z{Score: float64(now - openAIWSIngressLeaseTTLSeconds - 1), Member: "crashed-instance"},
|
||
|
|
redis.Z{Score: float64(now), Member: "live-other-instance"},
|
||
|
|
).Err())
|
||
|
|
require.NoError(s.T(), s.rdb.Expire(s.ctx, key, time.Duration(openAIWSIngressLeaseTTLSeconds)*time.Second).Err())
|
||
|
|
|
||
|
|
ok, err := s.rawCache.AcquireOpenAIWSIngressLease(s.ctx, apiKeyID, 2, "new-instance")
|
||
|
|
require.NoError(s.T(), err)
|
||
|
|
require.True(s.T(), ok, "the crashed member should be reaped before enforcing the limit")
|
||
|
|
|
||
|
|
_, err = s.rdb.ZScore(s.ctx, key, "crashed-instance").Result()
|
||
|
|
require.ErrorIs(s.T(), err, redis.Nil)
|
||
|
|
_, err = s.rdb.ZScore(s.ctx, key, "live-other-instance").Result()
|
||
|
|
require.NoError(s.T(), err, "a live lease owned by another instance must be preserved")
|
||
|
|
count, err := s.rdb.ZCard(s.ctx, key).Result()
|
||
|
|
require.NoError(s.T(), err)
|
||
|
|
require.Equal(s.T(), int64(2), count)
|
||
|
|
}
|
||
|
|
|
||
|
|
func (s *ConcurrencyCacheSuite) TestLiveLease_CountsTowardRegularAccountAndUserLimits() {
|
||
|
|
liveCache, ok := s.cache.(service.LiveConcurrencyCache)
|
||
|
|
require.True(s.T(), ok)
|
||
|
|
|
||
|
|
accountID := int64(9101)
|
||
|
|
userID := int64(9102)
|
||
|
|
apiKeyID := int64(9103)
|
||
|
|
acquired, err := liveCache.AcquireLiveLease(
|
||
|
|
s.ctx,
|
||
|
|
accountID,
|
||
|
|
1,
|
||
|
|
userID,
|
||
|
|
1,
|
||
|
|
apiKeyID,
|
||
|
|
"live-integration",
|
||
|
|
false,
|
||
|
|
)
|
||
|
|
require.NoError(s.T(), err)
|
||
|
|
require.True(s.T(), acquired)
|
||
|
|
|
||
|
|
regularAccount, err := s.cache.AcquireAccountSlot(s.ctx, accountID, 1, "regular-account")
|
||
|
|
require.NoError(s.T(), err)
|
||
|
|
require.False(s.T(), regularAccount)
|
||
|
|
regularUser, err := s.cache.AcquireUserSlot(s.ctx, userID, 1, "regular-user")
|
||
|
|
require.NoError(s.T(), err)
|
||
|
|
require.False(s.T(), regularUser)
|
||
|
|
|
||
|
|
refreshed, err := liveCache.RefreshLiveLease(s.ctx, accountID, userID, apiKeyID, "live-integration")
|
||
|
|
require.NoError(s.T(), err)
|
||
|
|
require.True(s.T(), refreshed)
|
||
|
|
require.NoError(s.T(), liveCache.ReleaseLiveLease(s.ctx, accountID, userID, apiKeyID, "live-integration"))
|
||
|
|
regularAccount, err = s.cache.AcquireAccountSlot(s.ctx, accountID, 1, "regular-account")
|
||
|
|
require.NoError(s.T(), err)
|
||
|
|
require.True(s.T(), regularAccount)
|
||
|
|
}
|
||
|
|
|
||
|
|
func (s *ConcurrencyCacheSuite) TestAccountSlot_AcquireAndRelease() {
|
||
|
|
accountID := int64(10)
|
||
|
|
reqID1, reqID2, reqID3 := "req1", "req2", "req3"
|
||
|
|
|
||
|
|
ok, err := s.cache.AcquireAccountSlot(s.ctx, accountID, 2, reqID1)
|
||
|
|
require.NoError(s.T(), err, "AcquireAccountSlot 1")
|
||
|
|
require.True(s.T(), ok)
|
||
|
|
|
||
|
|
ok, err = s.cache.AcquireAccountSlot(s.ctx, accountID, 2, reqID2)
|
||
|
|
require.NoError(s.T(), err, "AcquireAccountSlot 2")
|
||
|
|
require.True(s.T(), ok)
|
||
|
|
|
||
|
|
ok, err = s.cache.AcquireAccountSlot(s.ctx, accountID, 2, reqID3)
|
||
|
|
require.NoError(s.T(), err, "AcquireAccountSlot 3")
|
||
|
|
require.False(s.T(), ok, "expected third acquire to fail")
|
||
|
|
|
||
|
|
cur, err := s.cache.GetAccountConcurrency(s.ctx, accountID)
|
||
|
|
require.NoError(s.T(), err, "GetAccountConcurrency")
|
||
|
|
require.Equal(s.T(), 2, cur, "concurrency mismatch")
|
||
|
|
|
||
|
|
require.NoError(s.T(), s.cache.ReleaseAccountSlot(s.ctx, accountID, reqID1), "ReleaseAccountSlot")
|
||
|
|
|
||
|
|
cur, err = s.cache.GetAccountConcurrency(s.ctx, accountID)
|
||
|
|
require.NoError(s.T(), err, "GetAccountConcurrency after release")
|
||
|
|
require.Equal(s.T(), 1, cur, "expected 1 after release")
|
||
|
|
}
|
||
|
|
|
||
|
|
func (s *ConcurrencyCacheSuite) TestAccountActiveIndex_AcquireAndRelease() {
|
||
|
|
accountID := int64(610)
|
||
|
|
member := strconv.FormatInt(accountID, 10)
|
||
|
|
reqID := "active-index-req"
|
||
|
|
|
||
|
|
now, err := s.rawCache.redisUnixSeconds(s.ctx)
|
||
|
|
require.NoError(s.T(), err)
|
||
|
|
|
||
|
|
ok, err := s.cache.AcquireAccountSlot(s.ctx, accountID, 2, reqID)
|
||
|
|
require.NoError(s.T(), err)
|
||
|
|
require.True(s.T(), ok)
|
||
|
|
|
||
|
|
score, err := s.rdb.ZScore(s.ctx, accountActiveIndexKey, member).Result()
|
||
|
|
require.NoError(s.T(), err)
|
||
|
|
require.Greater(s.T(), int64(score), now, "index score should be a future expiry")
|
||
|
|
|
||
|
|
require.NoError(s.T(), s.cache.ReleaseAccountSlot(s.ctx, accountID, reqID))
|
||
|
|
|
||
|
|
_, err = s.rdb.ZScore(s.ctx, accountActiveIndexKey, member).Result()
|
||
|
|
require.ErrorIs(s.T(), err, redis.Nil, "index member should be removed after load drops to zero")
|
||
|
|
}
|
||
|
|
|
||
|
|
func (s *ConcurrencyCacheSuite) TestAccountActiveIndex_WaitLifecycle() {
|
||
|
|
accountID := int64(611)
|
||
|
|
member := strconv.FormatInt(accountID, 10)
|
||
|
|
|
||
|
|
ok, err := s.cache.IncrementAccountWaitCount(s.ctx, accountID, 2)
|
||
|
|
require.NoError(s.T(), err)
|
||
|
|
require.True(s.T(), ok)
|
||
|
|
|
||
|
|
_, err = s.rdb.ZScore(s.ctx, accountActiveIndexKey, member).Result()
|
||
|
|
require.NoError(s.T(), err, "wait increment should register index member")
|
||
|
|
|
||
|
|
require.NoError(s.T(), s.cache.DecrementAccountWaitCount(s.ctx, accountID))
|
||
|
|
|
||
|
|
_, err = s.rdb.ZScore(s.ctx, accountActiveIndexKey, member).Result()
|
||
|
|
require.ErrorIs(s.T(), err, redis.Nil, "index member should be removed after wait drops to zero")
|
||
|
|
}
|
||
|
|
|
||
|
|
func (s *ConcurrencyCacheSuite) TestUserActiveIndex_AcquireAndRelease() {
|
||
|
|
userID := int64(612)
|
||
|
|
member := strconv.FormatInt(userID, 10)
|
||
|
|
reqID := "user-active-index-req"
|
||
|
|
|
||
|
|
ok, err := s.cache.AcquireUserSlot(s.ctx, userID, 2, reqID)
|
||
|
|
require.NoError(s.T(), err)
|
||
|
|
require.True(s.T(), ok)
|
||
|
|
|
||
|
|
_, err = s.rdb.ZScore(s.ctx, userActiveIndexKey, member).Result()
|
||
|
|
require.NoError(s.T(), err, "acquire should register user index member")
|
||
|
|
|
||
|
|
require.NoError(s.T(), s.cache.ReleaseUserSlot(s.ctx, userID, reqID))
|
||
|
|
|
||
|
|
_, err = s.rdb.ZScore(s.ctx, userActiveIndexKey, member).Result()
|
||
|
|
require.ErrorIs(s.T(), err, redis.Nil, "user index member should be removed after release")
|
||
|
|
}
|
||
|
|
|
||
|
|
func (s *ConcurrencyCacheSuite) TestAccountSlot_TTL() {
|
||
|
|
accountID := int64(11)
|
||
|
|
reqID := "req_ttl_test"
|
||
|
|
slotKey := fmt.Sprintf("%s%d", accountSlotKeyPrefix, accountID)
|
||
|
|
|
||
|
|
ok, err := s.cache.AcquireAccountSlot(s.ctx, accountID, 5, reqID)
|
||
|
|
require.NoError(s.T(), err, "AcquireAccountSlot")
|
||
|
|
require.True(s.T(), ok)
|
||
|
|
|
||
|
|
ttl, err := s.rdb.TTL(s.ctx, slotKey).Result()
|
||
|
|
require.NoError(s.T(), err, "TTL")
|
||
|
|
s.AssertTTLWithin(ttl, 1*time.Second, testSlotTTL)
|
||
|
|
}
|
||
|
|
|
||
|
|
func (s *ConcurrencyCacheSuite) TestAccountSlot_DuplicateReqID() {
|
||
|
|
accountID := int64(12)
|
||
|
|
reqID := "dup-req"
|
||
|
|
|
||
|
|
ok, err := s.cache.AcquireAccountSlot(s.ctx, accountID, 2, reqID)
|
||
|
|
require.NoError(s.T(), err)
|
||
|
|
require.True(s.T(), ok)
|
||
|
|
|
||
|
|
// Acquiring with same reqID should be idempotent
|
||
|
|
ok, err = s.cache.AcquireAccountSlot(s.ctx, accountID, 2, reqID)
|
||
|
|
require.NoError(s.T(), err)
|
||
|
|
require.True(s.T(), ok)
|
||
|
|
|
||
|
|
cur, err := s.cache.GetAccountConcurrency(s.ctx, accountID)
|
||
|
|
require.NoError(s.T(), err)
|
||
|
|
require.Equal(s.T(), 1, cur, "expected concurrency=1 (idempotent)")
|
||
|
|
}
|
||
|
|
|
||
|
|
func (s *ConcurrencyCacheSuite) TestAccountSlot_ReleaseIdempotent() {
|
||
|
|
accountID := int64(13)
|
||
|
|
reqID := "release-test"
|
||
|
|
|
||
|
|
ok, err := s.cache.AcquireAccountSlot(s.ctx, accountID, 1, reqID)
|
||
|
|
require.NoError(s.T(), err)
|
||
|
|
require.True(s.T(), ok)
|
||
|
|
|
||
|
|
require.NoError(s.T(), s.cache.ReleaseAccountSlot(s.ctx, accountID, reqID), "ReleaseAccountSlot")
|
||
|
|
// Releasing again should not error
|
||
|
|
require.NoError(s.T(), s.cache.ReleaseAccountSlot(s.ctx, accountID, reqID), "ReleaseAccountSlot again")
|
||
|
|
// Releasing non-existent should not error
|
||
|
|
require.NoError(s.T(), s.cache.ReleaseAccountSlot(s.ctx, accountID, "non-existent"), "ReleaseAccountSlot non-existent")
|
||
|
|
|
||
|
|
cur, err := s.cache.GetAccountConcurrency(s.ctx, accountID)
|
||
|
|
require.NoError(s.T(), err)
|
||
|
|
require.Equal(s.T(), 0, cur)
|
||
|
|
}
|
||
|
|
|
||
|
|
func (s *ConcurrencyCacheSuite) TestAccountSlot_MaxZero() {
|
||
|
|
accountID := int64(14)
|
||
|
|
reqID := "max-zero-test"
|
||
|
|
|
||
|
|
ok, err := s.cache.AcquireAccountSlot(s.ctx, accountID, 0, reqID)
|
||
|
|
require.NoError(s.T(), err)
|
||
|
|
require.False(s.T(), ok, "expected acquire to fail with max=0")
|
||
|
|
}
|
||
|
|
|
||
|
|
func (s *ConcurrencyCacheSuite) TestUserSlot_AcquireAndRelease() {
|
||
|
|
userID := int64(42)
|
||
|
|
reqID1, reqID2 := "req1", "req2"
|
||
|
|
|
||
|
|
ok, err := s.cache.AcquireUserSlot(s.ctx, userID, 1, reqID1)
|
||
|
|
require.NoError(s.T(), err, "AcquireUserSlot")
|
||
|
|
require.True(s.T(), ok)
|
||
|
|
|
||
|
|
ok, err = s.cache.AcquireUserSlot(s.ctx, userID, 1, reqID2)
|
||
|
|
require.NoError(s.T(), err, "AcquireUserSlot 2")
|
||
|
|
require.False(s.T(), ok, "expected second acquire to fail at max=1")
|
||
|
|
|
||
|
|
cur, err := s.cache.GetUserConcurrency(s.ctx, userID)
|
||
|
|
require.NoError(s.T(), err, "GetUserConcurrency")
|
||
|
|
require.Equal(s.T(), 1, cur, "expected concurrency=1")
|
||
|
|
|
||
|
|
require.NoError(s.T(), s.cache.ReleaseUserSlot(s.ctx, userID, reqID1), "ReleaseUserSlot")
|
||
|
|
// Releasing a non-existent slot should not error
|
||
|
|
require.NoError(s.T(), s.cache.ReleaseUserSlot(s.ctx, userID, "non-existent"), "ReleaseUserSlot non-existent")
|
||
|
|
|
||
|
|
cur, err = s.cache.GetUserConcurrency(s.ctx, userID)
|
||
|
|
require.NoError(s.T(), err, "GetUserConcurrency after release")
|
||
|
|
require.Equal(s.T(), 0, cur, "expected concurrency=0 after release")
|
||
|
|
}
|
||
|
|
|
||
|
|
func (s *ConcurrencyCacheSuite) TestUserSlot_TTL() {
|
||
|
|
userID := int64(200)
|
||
|
|
reqID := "req_ttl_test"
|
||
|
|
slotKey := fmt.Sprintf("%s%d", userSlotKeyPrefix, userID)
|
||
|
|
|
||
|
|
ok, err := s.cache.AcquireUserSlot(s.ctx, userID, 5, reqID)
|
||
|
|
require.NoError(s.T(), err, "AcquireUserSlot")
|
||
|
|
require.True(s.T(), ok)
|
||
|
|
|
||
|
|
ttl, err := s.rdb.TTL(s.ctx, slotKey).Result()
|
||
|
|
require.NoError(s.T(), err, "TTL")
|
||
|
|
s.AssertTTLWithin(ttl, 1*time.Second, testSlotTTL)
|
||
|
|
}
|
||
|
|
|
||
|
|
func (s *ConcurrencyCacheSuite) TestAPIKeySlot_TrackReleaseAndBatchCount() {
|
||
|
|
cache := s.apiKeyConcurrencyCache()
|
||
|
|
apiKeyID := int64(300)
|
||
|
|
emptyAPIKeyID := int64(301)
|
||
|
|
slotKey := fmt.Sprintf("%s%d", apiKeySlotKeyPrefix, apiKeyID)
|
||
|
|
|
||
|
|
require.NoError(s.T(), cache.TrackAPIKeySlot(s.ctx, apiKeyID, "req1"))
|
||
|
|
require.NoError(s.T(), cache.TrackAPIKeySlot(s.ctx, apiKeyID, "req2"))
|
||
|
|
|
||
|
|
counts, err := cache.GetAPIKeyConcurrencyBatch(s.ctx, []int64{apiKeyID, emptyAPIKeyID})
|
||
|
|
require.NoError(s.T(), err)
|
||
|
|
require.Equal(s.T(), map[int64]int{apiKeyID: 2, emptyAPIKeyID: 0}, counts)
|
||
|
|
|
||
|
|
ttl, err := s.rdb.TTL(s.ctx, slotKey).Result()
|
||
|
|
require.NoError(s.T(), err, "TTL")
|
||
|
|
s.AssertTTLWithin(ttl, 1*time.Second, testSlotTTL)
|
||
|
|
|
||
|
|
require.NoError(s.T(), cache.ReleaseAPIKeySlot(s.ctx, apiKeyID, "req1"))
|
||
|
|
counts, err = cache.GetAPIKeyConcurrencyBatch(s.ctx, []int64{apiKeyID})
|
||
|
|
require.NoError(s.T(), err)
|
||
|
|
require.Equal(s.T(), 1, counts[apiKeyID])
|
||
|
|
|
||
|
|
require.NoError(s.T(), cache.ReleaseAPIKeySlot(s.ctx, apiKeyID, "req2"))
|
||
|
|
counts, err = cache.GetAPIKeyConcurrencyBatch(s.ctx, []int64{apiKeyID})
|
||
|
|
require.NoError(s.T(), err)
|
||
|
|
require.Equal(s.T(), 0, counts[apiKeyID])
|
||
|
|
}
|
||
|
|
|
||
|
|
func (s *ConcurrencyCacheSuite) TestWaitQueue_IncrementAndDecrement() {
|
||
|
|
userID := int64(20)
|
||
|
|
waitKey := fmt.Sprintf("%s%d", waitQueueKeyPrefix, userID)
|
||
|
|
|
||
|
|
ok, err := s.cache.IncrementWaitCount(s.ctx, userID, 2)
|
||
|
|
require.NoError(s.T(), err, "IncrementWaitCount 1")
|
||
|
|
require.True(s.T(), ok)
|
||
|
|
|
||
|
|
ok, err = s.cache.IncrementWaitCount(s.ctx, userID, 2)
|
||
|
|
require.NoError(s.T(), err, "IncrementWaitCount 2")
|
||
|
|
require.True(s.T(), ok)
|
||
|
|
|
||
|
|
ok, err = s.cache.IncrementWaitCount(s.ctx, userID, 2)
|
||
|
|
require.NoError(s.T(), err, "IncrementWaitCount 3")
|
||
|
|
require.False(s.T(), ok, "expected wait increment over max to fail")
|
||
|
|
|
||
|
|
ttl, err := s.rdb.TTL(s.ctx, waitKey).Result()
|
||
|
|
require.NoError(s.T(), err, "TTL waitKey")
|
||
|
|
s.AssertTTLWithin(ttl, 1*time.Second, testSlotTTL)
|
||
|
|
|
||
|
|
require.NoError(s.T(), s.cache.DecrementWaitCount(s.ctx, userID), "DecrementWaitCount")
|
||
|
|
|
||
|
|
val, err := s.rdb.Get(s.ctx, waitKey).Int()
|
||
|
|
if !errors.Is(err, redis.Nil) {
|
||
|
|
require.NoError(s.T(), err, "Get waitKey")
|
||
|
|
}
|
||
|
|
require.Equal(s.T(), 1, val, "expected wait count 1")
|
||
|
|
}
|
||
|
|
|
||
|
|
func (s *ConcurrencyCacheSuite) TestWaitQueue_DecrementNoNegative() {
|
||
|
|
userID := int64(300)
|
||
|
|
waitKey := fmt.Sprintf("%s%d", waitQueueKeyPrefix, userID)
|
||
|
|
|
||
|
|
// Test decrement on non-existent key - should not error and should not create negative value
|
||
|
|
require.NoError(s.T(), s.cache.DecrementWaitCount(s.ctx, userID), "DecrementWaitCount on non-existent key")
|
||
|
|
|
||
|
|
// Verify no key was created or it's not negative
|
||
|
|
val, err := s.rdb.Get(s.ctx, waitKey).Int()
|
||
|
|
if !errors.Is(err, redis.Nil) {
|
||
|
|
require.NoError(s.T(), err, "Get waitKey")
|
||
|
|
}
|
||
|
|
require.GreaterOrEqual(s.T(), val, 0, "expected non-negative wait count after decrement on empty")
|
||
|
|
|
||
|
|
// Set count to 1, then decrement twice
|
||
|
|
ok, err := s.cache.IncrementWaitCount(s.ctx, userID, 5)
|
||
|
|
require.NoError(s.T(), err, "IncrementWaitCount")
|
||
|
|
require.True(s.T(), ok)
|
||
|
|
|
||
|
|
// Decrement once (1 -> 0)
|
||
|
|
require.NoError(s.T(), s.cache.DecrementWaitCount(s.ctx, userID), "DecrementWaitCount")
|
||
|
|
|
||
|
|
// Decrement again on 0 - should not go negative
|
||
|
|
require.NoError(s.T(), s.cache.DecrementWaitCount(s.ctx, userID), "DecrementWaitCount on zero")
|
||
|
|
|
||
|
|
// Verify count is 0, not negative
|
||
|
|
val, err = s.rdb.Get(s.ctx, waitKey).Int()
|
||
|
|
if !errors.Is(err, redis.Nil) {
|
||
|
|
require.NoError(s.T(), err, "Get waitKey after double decrement")
|
||
|
|
}
|
||
|
|
require.GreaterOrEqual(s.T(), val, 0, "expected non-negative wait count")
|
||
|
|
}
|
||
|
|
|
||
|
|
func (s *ConcurrencyCacheSuite) TestAccountWaitQueue_IncrementAndDecrement() {
|
||
|
|
accountID := int64(30)
|
||
|
|
waitKey := fmt.Sprintf("%s%d", accountWaitKeyPrefix, accountID)
|
||
|
|
|
||
|
|
ok, err := s.cache.IncrementAccountWaitCount(s.ctx, accountID, 2)
|
||
|
|
require.NoError(s.T(), err, "IncrementAccountWaitCount 1")
|
||
|
|
require.True(s.T(), ok)
|
||
|
|
|
||
|
|
ok, err = s.cache.IncrementAccountWaitCount(s.ctx, accountID, 2)
|
||
|
|
require.NoError(s.T(), err, "IncrementAccountWaitCount 2")
|
||
|
|
require.True(s.T(), ok)
|
||
|
|
|
||
|
|
ok, err = s.cache.IncrementAccountWaitCount(s.ctx, accountID, 2)
|
||
|
|
require.NoError(s.T(), err, "IncrementAccountWaitCount 3")
|
||
|
|
require.False(s.T(), ok, "expected account wait increment over max to fail")
|
||
|
|
|
||
|
|
ttl, err := s.rdb.TTL(s.ctx, waitKey).Result()
|
||
|
|
require.NoError(s.T(), err, "TTL account waitKey")
|
||
|
|
s.AssertTTLWithin(ttl, 1*time.Second, testSlotTTL)
|
||
|
|
|
||
|
|
require.NoError(s.T(), s.cache.DecrementAccountWaitCount(s.ctx, accountID), "DecrementAccountWaitCount")
|
||
|
|
|
||
|
|
val, err := s.rdb.Get(s.ctx, waitKey).Int()
|
||
|
|
if !errors.Is(err, redis.Nil) {
|
||
|
|
require.NoError(s.T(), err, "Get waitKey")
|
||
|
|
}
|
||
|
|
require.Equal(s.T(), 1, val, "expected account wait count 1")
|
||
|
|
}
|
||
|
|
|
||
|
|
func (s *ConcurrencyCacheSuite) TestCleanupStaleProcessSlots() {
|
||
|
|
// 预置迁移 marker，隔离一次性清扫，只验证索引驱动的清理路径。
|
||
|
|
require.NoError(s.T(), s.rdb.Set(s.ctx, legacyWaitSweepMarkerKey, "1", 0).Err())
|
||
|
|
accountID := int64(901)
|
||
|
|
userID := int64(902)
|
||
|
|
apiKeyID := int64(903)
|
||
|
|
unindexedAccountID := int64(1901)
|
||
|
|
accountKey := fmt.Sprintf("%s%d", accountSlotKeyPrefix, accountID)
|
||
|
|
userKey := fmt.Sprintf("%s%d", userSlotKeyPrefix, userID)
|
||
|
|
apiKeyKey := fmt.Sprintf("%s%d", apiKeySlotKeyPrefix, apiKeyID)
|
||
|
|
unindexedAccountKey := fmt.Sprintf("%s%d", accountSlotKeyPrefix, unindexedAccountID)
|
||
|
|
userWaitKey := fmt.Sprintf("%s%d", waitQueueKeyPrefix, userID)
|
||
|
|
accountWaitKey := fmt.Sprintf("%s%d", accountWaitKeyPrefix, accountID)
|
||
|
|
unindexedAccountWaitKey := fmt.Sprintf("%s%d", accountWaitKeyPrefix, unindexedAccountID)
|
||
|
|
|
||
|
|
now, err := s.rawCache.redisUnixSeconds(s.ctx)
|
||
|
|
require.NoError(s.T(), err)
|
||
|
|
require.NoError(s.T(), s.rdb.ZAdd(s.ctx, accountKey,
|
||
|
|
redis.Z{Score: float64(now), Member: "oldproc-1"},
|
||
|
|
redis.Z{Score: float64(now), Member: "keep-1"},
|
||
|
|
).Err())
|
||
|
|
require.NoError(s.T(), s.rdb.ZAdd(s.ctx, userKey,
|
||
|
|
redis.Z{Score: float64(now), Member: "oldproc-2"},
|
||
|
|
redis.Z{Score: float64(now), Member: "keep-2"},
|
||
|
|
).Err())
|
||
|
|
require.NoError(s.T(), s.rdb.ZAdd(s.ctx, unindexedAccountKey,
|
||
|
|
redis.Z{Score: float64(now), Member: "oldproc-unindexed"},
|
||
|
|
).Err())
|
||
|
|
require.NoError(s.T(), s.rdb.ZAdd(s.ctx, apiKeyKey,
|
||
|
|
redis.Z{Score: float64(now), Member: "oldproc-3"},
|
||
|
|
redis.Z{Score: float64(now), Member: "keep-3"},
|
||
|
|
).Err())
|
||
|
|
require.NoError(s.T(), s.rdb.Set(s.ctx, userWaitKey, 3, time.Minute).Err())
|
||
|
|
require.NoError(s.T(), s.rdb.Set(s.ctx, accountWaitKey, 2, time.Minute).Err())
|
||
|
|
require.NoError(s.T(), s.rdb.Set(s.ctx, unindexedAccountWaitKey, 2, time.Minute).Err())
|
||
|
|
require.NoError(s.T(), s.rdb.ZAdd(s.ctx, accountActiveIndexKey, redis.Z{
|
||
|
|
Score: float64(now + 60),
|
||
|
|
Member: strconv.FormatInt(accountID, 10),
|
||
|
|
}).Err())
|
||
|
|
require.NoError(s.T(), s.rdb.ZAdd(s.ctx, userActiveIndexKey, redis.Z{
|
||
|
|
Score: float64(now + 60),
|
||
|
|
Member: strconv.FormatInt(userID, 10),
|
||
|
|
}).Err())
|
||
|
|
|
||
|
|
require.NoError(s.T(), s.cache.CleanupStaleProcessSlots(s.ctx, "keep-"))
|
||
|
|
|
||
|
|
accountMembers, err := s.rdb.ZRange(s.ctx, accountKey, 0, -1).Result()
|
||
|
|
require.NoError(s.T(), err)
|
||
|
|
require.Equal(s.T(), []string{"keep-1"}, accountMembers)
|
||
|
|
|
||
|
|
userMembers, err := s.rdb.ZRange(s.ctx, userKey, 0, -1).Result()
|
||
|
|
require.NoError(s.T(), err)
|
||
|
|
require.Equal(s.T(), []string{"keep-2"}, userMembers)
|
||
|
|
|
||
|
|
// API Key 槽位（stats-only）不在启动清理范围内，靠分数裁剪与 key TTL 自愈。
|
||
|
|
apiKeyMembers, err := s.rdb.ZRange(s.ctx, apiKeyKey, 0, -1).Result()
|
||
|
|
require.NoError(s.T(), err)
|
||
|
|
require.ElementsMatch(s.T(), []string{"keep-3", "oldproc-3"}, apiKeyMembers)
|
||
|
|
|
||
|
|
_, err = s.rdb.Get(s.ctx, userWaitKey).Result()
|
||
|
|
require.True(s.T(), errors.Is(err, redis.Nil))
|
||
|
|
|
||
|
|
_, err = s.rdb.Get(s.ctx, accountWaitKey).Result()
|
||
|
|
require.True(s.T(), errors.Is(err, redis.Nil))
|
||
|
|
|
||
|
|
unindexedMembers, err := s.rdb.ZRange(s.ctx, unindexedAccountKey, 0, -1).Result()
|
||
|
|
require.NoError(s.T(), err)
|
||
|
|
require.Equal(s.T(), []string{"oldproc-unindexed"}, unindexedMembers)
|
||
|
|
_, err = s.rdb.Get(s.ctx, unindexedAccountWaitKey).Result()
|
||
|
|
require.NoError(s.T(), err)
|
||
|
|
}
|
||
|
|
|
||
|
|
func (s *ConcurrencyCacheSuite) TestGetAccountConcurrency_Missing() {
|
||
|
|
// When no slots exist, GetAccountConcurrency should return 0
|
||
|
|
cur, err := s.cache.GetAccountConcurrency(s.ctx, 999)
|
||
|
|
require.NoError(s.T(), err)
|
||
|
|
require.Equal(s.T(), 0, cur)
|
||
|
|
}
|
||
|
|
|
||
|
|
func (s *ConcurrencyCacheSuite) TestGetUserConcurrency_Missing() {
|
||
|
|
// When no slots exist, GetUserConcurrency should return 0
|
||
|
|
cur, err := s.cache.GetUserConcurrency(s.ctx, 999)
|
||
|
|
require.NoError(s.T(), err)
|
||
|
|
require.Equal(s.T(), 0, cur)
|
||
|
|
}
|
||
|
|
|
||
|
|
func (s *ConcurrencyCacheSuite) TestGetAccountsLoadBatch() {
|
||
|
|
s.T().Skip("TODO: Fix this test - CurrentConcurrency returns 0 instead of expected value in CI")
|
||
|
|
// Setup: Create accounts with different load states
|
||
|
|
account1 := int64(100)
|
||
|
|
account2 := int64(101)
|
||
|
|
account3 := int64(102)
|
||
|
|
|
||
|
|
// Account 1: 2/3 slots used, 1 waiting
|
||
|
|
ok, err := s.cache.AcquireAccountSlot(s.ctx, account1, 3, "req1")
|
||
|
|
require.NoError(s.T(), err)
|
||
|
|
require.True(s.T(), ok)
|
||
|
|
ok, err = s.cache.AcquireAccountSlot(s.ctx, account1, 3, "req2")
|
||
|
|
require.NoError(s.T(), err)
|
||
|
|
require.True(s.T(), ok)
|
||
|
|
ok, err = s.cache.IncrementAccountWaitCount(s.ctx, account1, 5)
|
||
|
|
require.NoError(s.T(), err)
|
||
|
|
require.True(s.T(), ok)
|
||
|
|
|
||
|
|
// Account 2: 1/2 slots used, 0 waiting
|
||
|
|
ok, err = s.cache.AcquireAccountSlot(s.ctx, account2, 2, "req3")
|
||
|
|
require.NoError(s.T(), err)
|
||
|
|
require.True(s.T(), ok)
|
||
|
|
|
||
|
|
// Account 3: 0/1 slots used, 0 waiting (idle)
|
||
|
|
|
||
|
|
// Query batch load
|
||
|
|
accounts := []service.AccountWithConcurrency{
|
||
|
|
{ID: account1, MaxConcurrency: 3},
|
||
|
|
{ID: account2, MaxConcurrency: 2},
|
||
|
|
{ID: account3, MaxConcurrency: 1},
|
||
|
|
}
|
||
|
|
|
||
|
|
loadMap, err := s.cache.GetAccountsLoadBatch(s.ctx, accounts)
|
||
|
|
require.NoError(s.T(), err)
|
||
|
|
require.Len(s.T(), loadMap, 3)
|
||
|
|
|
||
|
|
// Verify account1: (2 + 1) / 3 = 100%
|
||
|
|
load1 := loadMap[account1]
|
||
|
|
require.NotNil(s.T(), load1)
|
||
|
|
require.Equal(s.T(), account1, load1.AccountID)
|
||
|
|
require.Equal(s.T(), 2, load1.CurrentConcurrency)
|
||
|
|
require.Equal(s.T(), 1, load1.WaitingCount)
|
||
|
|
require.Equal(s.T(), 100, load1.LoadRate)
|
||
|
|
|
||
|
|
// Verify account2: (1 + 0) / 2 = 50%
|
||
|
|
load2 := loadMap[account2]
|
||
|
|
require.NotNil(s.T(), load2)
|
||
|
|
require.Equal(s.T(), account2, load2.AccountID)
|
||
|
|
require.Equal(s.T(), 1, load2.CurrentConcurrency)
|
||
|
|
require.Equal(s.T(), 0, load2.WaitingCount)
|
||
|
|
require.Equal(s.T(), 50, load2.LoadRate)
|
||
|
|
|
||
|
|
// Verify account3: (0 + 0) / 1 = 0%
|
||
|
|
load3 := loadMap[account3]
|
||
|
|
require.NotNil(s.T(), load3)
|
||
|
|
require.Equal(s.T(), account3, load3.AccountID)
|
||
|
|
require.Equal(s.T(), 0, load3.CurrentConcurrency)
|
||
|
|
require.Equal(s.T(), 0, load3.WaitingCount)
|
||
|
|
require.Equal(s.T(), 0, load3.LoadRate)
|
||
|
|
}
|
||
|
|
|
||
|
|
func (s *ConcurrencyCacheSuite) TestGetAccountsLoadBatch_Empty() {
|
||
|
|
// Test with empty account list
|
||
|
|
loadMap, err := s.cache.GetAccountsLoadBatch(s.ctx, []service.AccountWithConcurrency{})
|
||
|
|
require.NoError(s.T(), err)
|
||
|
|
require.Empty(s.T(), loadMap)
|
||
|
|
}
|
||
|
|
|
||
|
|
func (s *ConcurrencyCacheSuite) TestCleanupExpiredAccountSlots() {
|
||
|
|
accountID := int64(200)
|
||
|
|
slotKey := fmt.Sprintf("%s%d", accountSlotKeyPrefix, accountID)
|
||
|
|
|
||
|
|
// Acquire 3 slots
|
||
|
|
ok, err := s.cache.AcquireAccountSlot(s.ctx, accountID, 5, "req1")
|
||
|
|
require.NoError(s.T(), err)
|
||
|
|
require.True(s.T(), ok)
|
||
|
|
ok, err = s.cache.AcquireAccountSlot(s.ctx, accountID, 5, "req2")
|
||
|
|
require.NoError(s.T(), err)
|
||
|
|
require.True(s.T(), ok)
|
||
|
|
ok, err = s.cache.AcquireAccountSlot(s.ctx, accountID, 5, "req3")
|
||
|
|
require.NoError(s.T(), err)
|
||
|
|
require.True(s.T(), ok)
|
||
|
|
|
||
|
|
// Verify 3 slots exist
|
||
|
|
cur, err := s.cache.GetAccountConcurrency(s.ctx, accountID)
|
||
|
|
require.NoError(s.T(), err)
|
||
|
|
require.Equal(s.T(), 3, cur)
|
||
|
|
|
||
|
|
// Manually set old timestamps for req1 and req2 (simulate expired slots)
|
||
|
|
now := time.Now().Unix()
|
||
|
|
expiredTime := now - int64(testSlotTTL.Seconds()) - 10 // 10 seconds past TTL
|
||
|
|
err = s.rdb.ZAdd(s.ctx, slotKey, redis.Z{Score: float64(expiredTime), Member: "req1"}).Err()
|
||
|
|
require.NoError(s.T(), err)
|
||
|
|
err = s.rdb.ZAdd(s.ctx, slotKey, redis.Z{Score: float64(expiredTime), Member: "req2"}).Err()
|
||
|
|
require.NoError(s.T(), err)
|
||
|
|
|
||
|
|
// Run cleanup
|
||
|
|
err = s.cache.CleanupExpiredAccountSlots(s.ctx, accountID)
|
||
|
|
require.NoError(s.T(), err)
|
||
|
|
|
||
|
|
// Verify only 1 slot remains (req3)
|
||
|
|
cur, err = s.cache.GetAccountConcurrency(s.ctx, accountID)
|
||
|
|
require.NoError(s.T(), err)
|
||
|
|
require.Equal(s.T(), 1, cur)
|
||
|
|
|
||
|
|
// Verify req3 still exists
|
||
|
|
members, err := s.rdb.ZRange(s.ctx, slotKey, 0, -1).Result()
|
||
|
|
require.NoError(s.T(), err)
|
||
|
|
require.Len(s.T(), members, 1)
|
||
|
|
require.Equal(s.T(), "req3", members[0])
|
||
|
|
}
|
||
|
|
|
||
|
|
func (s *ConcurrencyCacheSuite) TestCleanupExpiredAccountSlots_NoExpired() {
|
||
|
|
accountID := int64(201)
|
||
|
|
|
||
|
|
// Acquire 2 fresh slots
|
||
|
|
ok, err := s.cache.AcquireAccountSlot(s.ctx, accountID, 5, "req1")
|
||
|
|
require.NoError(s.T(), err)
|
||
|
|
require.True(s.T(), ok)
|
||
|
|
ok, err = s.cache.AcquireAccountSlot(s.ctx, accountID, 5, "req2")
|
||
|
|
require.NoError(s.T(), err)
|
||
|
|
require.True(s.T(), ok)
|
||
|
|
|
||
|
|
// Run cleanup (should not remove anything)
|
||
|
|
err = s.cache.CleanupExpiredAccountSlots(s.ctx, accountID)
|
||
|
|
require.NoError(s.T(), err)
|
||
|
|
|
||
|
|
// Verify both slots still exist
|
||
|
|
cur, err := s.cache.GetAccountConcurrency(s.ctx, accountID)
|
||
|
|
require.NoError(s.T(), err)
|
||
|
|
require.Equal(s.T(), 2, cur)
|
||
|
|
}
|
||
|
|
|
||
|
|
func (s *ConcurrencyCacheSuite) TestCleanupExpiredAccountSlotKeys() {
|
||
|
|
now, err := s.rawCache.redisUnixSeconds(s.ctx)
|
||
|
|
require.NoError(s.T(), err)
|
||
|
|
expiredTime := now - int64(testSlotTTL.Seconds()) - 10
|
||
|
|
accountKeyWithFresh := fmt.Sprintf("%s%d", accountSlotKeyPrefix, 301)
|
||
|
|
accountKeyExpiredOnly := fmt.Sprintf("%s%d", accountSlotKeyPrefix, 302)
|
||
|
|
userKey := fmt.Sprintf("%s%d", userSlotKeyPrefix, 303)
|
||
|
|
unindexedAccountKey := fmt.Sprintf("%s%d", accountSlotKeyPrefix, 304)
|
||
|
|
|
||
|
|
require.NoError(s.T(), s.rdb.ZAdd(s.ctx, accountKeyWithFresh,
|
||
|
|
redis.Z{Score: float64(expiredTime), Member: "expired"},
|
||
|
|
redis.Z{Score: float64(now), Member: "fresh"},
|
||
|
|
).Err())
|
||
|
|
require.NoError(s.T(), s.rdb.ZAdd(s.ctx, accountKeyExpiredOnly,
|
||
|
|
redis.Z{Score: float64(expiredTime), Member: "expired-only"},
|
||
|
|
).Err())
|
||
|
|
require.NoError(s.T(), s.rdb.ZAdd(s.ctx, userKey,
|
||
|
|
redis.Z{Score: float64(expiredTime), Member: "user-expired"},
|
||
|
|
).Err())
|
||
|
|
require.NoError(s.T(), s.rdb.ZAdd(s.ctx, unindexedAccountKey,
|
||
|
|
redis.Z{Score: float64(expiredTime), Member: "unindexed-expired"},
|
||
|
|
).Err())
|
||
|
|
require.NoError(s.T(), s.rdb.ZAdd(s.ctx, accountActiveIndexKey,
|
||
|
|
redis.Z{Score: float64(now), Member: "301"},
|
||
|
|
redis.Z{Score: float64(now), Member: "302"},
|
||
|
|
).Err())
|
||
|
|
|
||
|
|
require.NoError(s.T(), s.cache.CleanupExpiredAccountSlotKeys(s.ctx))
|
||
|
|
|
||
|
|
accountMembers, err := s.rdb.ZRange(s.ctx, accountKeyWithFresh, 0, -1).Result()
|
||
|
|
require.NoError(s.T(), err)
|
||
|
|
require.Equal(s.T(), []string{"fresh"}, accountMembers)
|
||
|
|
|
||
|
|
exists, err := s.rdb.Exists(s.ctx, accountKeyExpiredOnly).Result()
|
||
|
|
require.NoError(s.T(), err)
|
||
|
|
require.EqualValues(s.T(), 0, exists)
|
||
|
|
|
||
|
|
userMembers, err := s.rdb.ZRange(s.ctx, userKey, 0, -1).Result()
|
||
|
|
require.NoError(s.T(), err)
|
||
|
|
require.Equal(s.T(), []string{"user-expired"}, userMembers)
|
||
|
|
|
||
|
|
unindexedMembers, err := s.rdb.ZRange(s.ctx, unindexedAccountKey, 0, -1).Result()
|
||
|
|
require.NoError(s.T(), err)
|
||
|
|
require.Equal(s.T(), []string{"unindexed-expired"}, unindexedMembers)
|
||
|
|
|
||
|
|
score, err := s.rdb.ZScore(s.ctx, accountActiveIndexKey, "301").Result()
|
||
|
|
require.NoError(s.T(), err)
|
||
|
|
require.Greater(s.T(), int64(score), now)
|
||
|
|
_, err = s.rdb.ZScore(s.ctx, accountActiveIndexKey, "302").Result()
|
||
|
|
require.ErrorIs(s.T(), err, redis.Nil)
|
||
|
|
}
|
||
|
|
|
||
|
|
func (s *ConcurrencyCacheSuite) TestCleanupExpiredAccountSlotKeys_ReapsUserIndex() {
|
||
|
|
now, err := s.rawCache.redisUnixSeconds(s.ctx)
|
||
|
|
require.NoError(s.T(), err)
|
||
|
|
expiredScore := float64(now - 10)
|
||
|
|
userKeyWithFresh := fmt.Sprintf("%s%d", userSlotKeyPrefix, 401)
|
||
|
|
|
||
|
|
// 401 有真实负载但索引 score 已过期：应刷新而不是删除。
|
||
|
|
require.NoError(s.T(), s.rdb.ZAdd(s.ctx, userKeyWithFresh,
|
||
|
|
redis.Z{Score: float64(now), Member: "fresh"},
|
||
|
|
).Err())
|
||
|
|
// 402 无任何负载：过期索引 member 应被回收。
|
||
|
|
// 非法 member 也应随过期候选一并清除。
|
||
|
|
require.NoError(s.T(), s.rdb.ZAdd(s.ctx, userActiveIndexKey,
|
||
|
|
redis.Z{Score: expiredScore, Member: "401"},
|
||
|
|
redis.Z{Score: expiredScore, Member: "402"},
|
||
|
|
redis.Z{Score: expiredScore, Member: "not-a-user-id"},
|
||
|
|
).Err())
|
||
|
|
|
||
|
|
require.NoError(s.T(), s.cache.CleanupExpiredAccountSlotKeys(s.ctx))
|
||
|
|
|
||
|
|
score, err := s.rdb.ZScore(s.ctx, userActiveIndexKey, "401").Result()
|
||
|
|
require.NoError(s.T(), err)
|
||
|
|
require.Greater(s.T(), int64(score), now, "loaded user should be re-scheduled, not dropped")
|
||
|
|
|
||
|
|
_, err = s.rdb.ZScore(s.ctx, userActiveIndexKey, "402").Result()
|
||
|
|
require.ErrorIs(s.T(), err, redis.Nil, "idle expired user member should be reaped")
|
||
|
|
|
||
|
|
_, err = s.rdb.ZScore(s.ctx, userActiveIndexKey, "not-a-user-id").Result()
|
||
|
|
require.ErrorIs(s.T(), err, redis.Nil, "invalid member should be reaped")
|
||
|
|
}
|
||
|
|
|
||
|
|
func (s *ConcurrencyCacheSuite) TestCleanupStaleProcessSlots_LegacyWaitSweepRunsOnce() {
|
||
|
|
unindexedAccountWaitKey := fmt.Sprintf("%s%d", accountWaitKeyPrefix, 2901)
|
||
|
|
unindexedUserWaitKey := fmt.Sprintf("%s%d", waitQueueKeyPrefix, 2902)
|
||
|
|
require.NoError(s.T(), s.rdb.Set(s.ctx, unindexedAccountWaitKey, 5, time.Minute).Err())
|
||
|
|
require.NoError(s.T(), s.rdb.Set(s.ctx, unindexedUserWaitKey, 3, time.Minute).Err())
|
||
|
|
|
||
|
|
// 首次运行：marker 不存在，一次性清扫删除所有遗留等待计数（含未入索引的）。
|
||
|
|
require.NoError(s.T(), s.cache.CleanupStaleProcessSlots(s.ctx, "keep-"))
|
||
|
|
|
||
|
|
_, err := s.rdb.Get(s.ctx, unindexedAccountWaitKey).Result()
|
||
|
|
require.ErrorIs(s.T(), err, redis.Nil, "legacy account wait key should be swept on first startup")
|
||
|
|
_, err = s.rdb.Get(s.ctx, unindexedUserWaitKey).Result()
|
||
|
|
require.ErrorIs(s.T(), err, redis.Nil, "legacy user wait key should be swept on first startup")
|
||
|
|
|
||
|
|
exists, err := s.rdb.Exists(s.ctx, legacyWaitSweepMarkerKey).Result()
|
||
|
|
require.NoError(s.T(), err)
|
||
|
|
require.EqualValues(s.T(), 1, exists, "sweep marker should be set after first run")
|
||
|
|
|
||
|
|
// 再次运行：marker 已存在，未入索引的等待计数不再被触碰。
|
||
|
|
require.NoError(s.T(), s.rdb.Set(s.ctx, unindexedAccountWaitKey, 5, time.Minute).Err())
|
||
|
|
require.NoError(s.T(), s.cache.CleanupStaleProcessSlots(s.ctx, "keep-"))
|
||
|
|
val, err := s.rdb.Get(s.ctx, unindexedAccountWaitKey).Int()
|
||
|
|
require.NoError(s.T(), err, "sweep must not run twice")
|
||
|
|
require.Equal(s.T(), 5, val)
|
||
|
|
}
|
||
|
|
|
||
|
|
func (s *ConcurrencyCacheSuite) TestCleanupStaleProcessSlots_ProcessesExpiredIndexMembers() {
|
||
|
|
// score 已过期的索引成员往往正是崩溃进程留下的残留，启动清理必须覆盖它们。
|
||
|
|
require.NoError(s.T(), s.rdb.Set(s.ctx, legacyWaitSweepMarkerKey, "1", 0).Err())
|
||
|
|
accountID := int64(3901)
|
||
|
|
userID := int64(3902)
|
||
|
|
accountKey := fmt.Sprintf("%s%d", accountSlotKeyPrefix, accountID)
|
||
|
|
userKey := fmt.Sprintf("%s%d", userSlotKeyPrefix, userID)
|
||
|
|
accountWaitKey := fmt.Sprintf("%s%d", accountWaitKeyPrefix, accountID)
|
||
|
|
|
||
|
|
now, err := s.rawCache.redisUnixSeconds(s.ctx)
|
||
|
|
require.NoError(s.T(), err)
|
||
|
|
require.NoError(s.T(), s.rdb.ZAdd(s.ctx, accountKey,
|
||
|
|
redis.Z{Score: float64(now), Member: "oldproc-1"},
|
||
|
|
).Err())
|
||
|
|
require.NoError(s.T(), s.rdb.ZAdd(s.ctx, userKey,
|
||
|
|
redis.Z{Score: float64(now), Member: "oldproc-2"},
|
||
|
|
).Err())
|
||
|
|
require.NoError(s.T(), s.rdb.Set(s.ctx, accountWaitKey, 4, time.Minute).Err())
|
||
|
|
// 索引 score 设为过去时刻，模拟长时间停机后索引已“过期”。
|
||
|
|
require.NoError(s.T(), s.rdb.ZAdd(s.ctx, accountActiveIndexKey, redis.Z{
|
||
|
|
Score: float64(now - 100),
|
||
|
|
Member: strconv.FormatInt(accountID, 10),
|
||
|
|
}).Err())
|
||
|
|
require.NoError(s.T(), s.rdb.ZAdd(s.ctx, userActiveIndexKey, redis.Z{
|
||
|
|
Score: float64(now - 100),
|
||
|
|
Member: strconv.FormatInt(userID, 10),
|
||
|
|
}).Err())
|
||
|
|
|
||
|
|
require.NoError(s.T(), s.cache.CleanupStaleProcessSlots(s.ctx, "keep-"))
|
||
|
|
|
||
|
|
exists, err := s.rdb.Exists(s.ctx, accountKey).Result()
|
||
|
|
require.NoError(s.T(), err)
|
||
|
|
require.EqualValues(s.T(), 0, exists, "stale slot key of expired index member should be purged")
|
||
|
|
|
||
|
|
exists, err = s.rdb.Exists(s.ctx, userKey).Result()
|
||
|
|
require.NoError(s.T(), err)
|
||
|
|
require.EqualValues(s.T(), 0, exists)
|
||
|
|
|
||
|
|
_, err = s.rdb.Get(s.ctx, accountWaitKey).Result()
|
||
|
|
require.ErrorIs(s.T(), err, redis.Nil, "wait counter of expired index member should be deleted")
|
||
|
|
|
||
|
|
_, err = s.rdb.ZScore(s.ctx, accountActiveIndexKey, strconv.FormatInt(accountID, 10)).Result()
|
||
|
|
require.ErrorIs(s.T(), err, redis.Nil, "emptied member should be removed from index")
|
||
|
|
_, err = s.rdb.ZScore(s.ctx, userActiveIndexKey, strconv.FormatInt(userID, 10)).Result()
|
||
|
|
require.ErrorIs(s.T(), err, redis.Nil)
|
||
|
|
}
|
||
|
|
|
||
|
|
func (s *ConcurrencyCacheSuite) TestCleanupStaleProcessSlots_RemovesOldPrefixesAndWaitCounters() {
|
||
|
|
// 预置迁移 marker，确保等待计数删除来自索引驱动路径而非一次性清扫。
|
||
|
|
require.NoError(s.T(), s.rdb.Set(s.ctx, legacyWaitSweepMarkerKey, "1", 0).Err())
|
||
|
|
accountID := int64(901)
|
||
|
|
userID := int64(902)
|
||
|
|
accountSlotKey := fmt.Sprintf("%s%d", accountSlotKeyPrefix, accountID)
|
||
|
|
userSlotKey := fmt.Sprintf("%s%d", userSlotKeyPrefix, userID)
|
||
|
|
userWaitKey := fmt.Sprintf("%s%d", waitQueueKeyPrefix, userID)
|
||
|
|
accountWaitKey := fmt.Sprintf("%s%d", accountWaitKeyPrefix, accountID)
|
||
|
|
|
||
|
|
now, err := s.rawCache.redisUnixSeconds(s.ctx)
|
||
|
|
require.NoError(s.T(), err)
|
||
|
|
require.NoError(s.T(), s.rdb.ZAdd(s.ctx, accountSlotKey,
|
||
|
|
redis.Z{Score: float64(now), Member: "oldproc-1"},
|
||
|
|
redis.Z{Score: float64(now), Member: "activeproc-1"},
|
||
|
|
).Err())
|
||
|
|
require.NoError(s.T(), s.rdb.Expire(s.ctx, accountSlotKey, testSlotTTL).Err())
|
||
|
|
require.NoError(s.T(), s.rdb.ZAdd(s.ctx, userSlotKey,
|
||
|
|
redis.Z{Score: float64(now), Member: "oldproc-2"},
|
||
|
|
redis.Z{Score: float64(now), Member: "activeproc-2"},
|
||
|
|
).Err())
|
||
|
|
require.NoError(s.T(), s.rdb.Expire(s.ctx, userSlotKey, testSlotTTL).Err())
|
||
|
|
require.NoError(s.T(), s.rdb.Set(s.ctx, userWaitKey, 3, testSlotTTL).Err())
|
||
|
|
require.NoError(s.T(), s.rdb.Set(s.ctx, accountWaitKey, 2, testSlotTTL).Err())
|
||
|
|
require.NoError(s.T(), s.rdb.ZAdd(s.ctx, accountActiveIndexKey, redis.Z{
|
||
|
|
Score: float64(now + 60),
|
||
|
|
Member: strconv.FormatInt(accountID, 10),
|
||
|
|
}).Err())
|
||
|
|
require.NoError(s.T(), s.rdb.ZAdd(s.ctx, userActiveIndexKey, redis.Z{
|
||
|
|
Score: float64(now + 60),
|
||
|
|
Member: strconv.FormatInt(userID, 10),
|
||
|
|
}).Err())
|
||
|
|
|
||
|
|
require.NoError(s.T(), s.cache.CleanupStaleProcessSlots(s.ctx, "activeproc-"))
|
||
|
|
|
||
|
|
accountMembers, err := s.rdb.ZRange(s.ctx, accountSlotKey, 0, -1).Result()
|
||
|
|
require.NoError(s.T(), err)
|
||
|
|
require.Equal(s.T(), []string{"activeproc-1"}, accountMembers)
|
||
|
|
|
||
|
|
userMembers, err := s.rdb.ZRange(s.ctx, userSlotKey, 0, -1).Result()
|
||
|
|
require.NoError(s.T(), err)
|
||
|
|
require.Equal(s.T(), []string{"activeproc-2"}, userMembers)
|
||
|
|
|
||
|
|
_, err = s.rdb.Get(s.ctx, userWaitKey).Result()
|
||
|
|
require.ErrorIs(s.T(), err, redis.Nil)
|
||
|
|
_, err = s.rdb.Get(s.ctx, accountWaitKey).Result()
|
||
|
|
require.ErrorIs(s.T(), err, redis.Nil)
|
||
|
|
}
|
||
|
|
|
||
|
|
func (s *ConcurrencyCacheSuite) TestCleanupStaleProcessSlots_DeletesEmptySlotKeys() {
|
||
|
|
accountID := int64(903)
|
||
|
|
accountSlotKey := fmt.Sprintf("%s%d", accountSlotKeyPrefix, accountID)
|
||
|
|
now, err := s.rawCache.redisUnixSeconds(s.ctx)
|
||
|
|
require.NoError(s.T(), err)
|
||
|
|
require.NoError(s.T(), s.rdb.ZAdd(s.ctx, accountSlotKey, redis.Z{Score: float64(now), Member: "oldproc-1"}).Err())
|
||
|
|
require.NoError(s.T(), s.rdb.Expire(s.ctx, accountSlotKey, testSlotTTL).Err())
|
||
|
|
require.NoError(s.T(), s.rdb.ZAdd(s.ctx, accountActiveIndexKey, redis.Z{
|
||
|
|
Score: float64(now + 60),
|
||
|
|
Member: strconv.FormatInt(accountID, 10),
|
||
|
|
}).Err())
|
||
|
|
|
||
|
|
require.NoError(s.T(), s.cache.CleanupStaleProcessSlots(s.ctx, "activeproc-"))
|
||
|
|
|
||
|
|
exists, err := s.rdb.Exists(s.ctx, accountSlotKey).Result()
|
||
|
|
require.NoError(s.T(), err)
|
||
|
|
require.EqualValues(s.T(), 0, exists)
|
||
|
|
}
|