github/sub2api
1
0
Fork 0
mirror of https://github.com/Wei-Shaw/sub2api.git synced 2026-10-07 17:08:33 +08:00
Code Issues Packages Projects Releases Wiki Activity

优化快照发布临时分配

Browse Source
This commit is contained in:
jjaw
2026-07-21 03:07:01 +08:00
parent e625ce3b3b
commit af3b6bbd68
3 changed files with 68 additions and 59 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/internal/repository/scheduler_cache.go
+12 -53
View File
@@ -389,13 +389,15 @@ func (c *schedulerCache) SetSnapshot(ctx context.Context, bucket service.Schedul
if err != nil {
return err
}
if err := c.writeSnapshotVersion(ctx, bucket, version, accounts); err != nil {
// 快照成员最终只依赖可编码账号的有序 ID；直接复用 ID 路径，避免为
// 随后立即丢弃的完整 Account 再分配一份临时切片。
if _, err := c.writeSnapshotVersionAndReturnAccountIDs(ctx, bucket, version, accounts); err != nil {
return err
}
return c.activateSnapshotVersion(ctx, bucket, token, version)
}
// SetSnapshotAndReturnAccountIDs 完整发布快照，并返回 writeAccounts 实际接受的有序账号 ID。
// SetSnapshotAndReturnAccountIDs 完整发布快照，并返回实际成功编码并写入的有序账号 ID。
// 该可选能力只供同一重建批次复用，返回前仍会完成版本激活与 fencing 校验。
func (c *schedulerCache) SetSnapshotAndReturnAccountIDs(ctx context.Context, bucket service.SchedulerBucket, token service.SchedulerBucketWriteToken, accounts []service.Account) ([]int64, error) {
if !token.ValidFor(bucket) {
@@ -447,14 +449,6 @@ func (c *schedulerCache) allocateSnapshotVersion(ctx context.Context, bucket ser
return strconv.FormatInt(result, 10), nil
}
func (c *schedulerCache) writeSnapshotVersion(ctx context.Context, bucket service.SchedulerBucket, version string, accounts []service.Account) error {
cacheableAccounts, err := c.writeAccounts(ctx, accounts)
if err != nil {
return err
}
return c.writeSnapshotAccounts(ctx, bucket, version, cacheableAccounts)
}
func (c *schedulerCache) writeSnapshotVersionAndReturnAccountIDs(ctx context.Context, bucket service.SchedulerBucket, version string, accounts []service.Account) ([]int64, error) {
accountIDs, err := c.writeAccountIDs(ctx, accounts)
if err != nil {
@@ -466,20 +460,6 @@ func (c *schedulerCache) writeSnapshotVersionAndReturnAccountIDs(ctx context.Con
return accountIDs, nil
}
func (c *schedulerCache) writeSnapshotAccounts(ctx context.Context, bucket service.SchedulerBucket, version string, accounts []service.Account) error {
if len(accounts) == 0 {
return nil
}
members := make([]redis.Z, 0, len(accounts))
for idx, account := range accounts {
members = append(members, redis.Z{
Score: float64(idx),
Member: strconv.FormatInt(account.ID, 10),
})
}
return c.writeSnapshotMembers(ctx, bucket, version, members)
}
func (c *schedulerCache) writeSnapshotAccountIDs(ctx context.Context, bucket service.SchedulerBucket, version string, accountIDs []int64) error {
members := schedulerSnapshotMembers(accountIDs)
return c.writeSnapshotMembers(ctx, bucket, version, members)
@@ -572,11 +552,11 @@ func (c *schedulerCache) SetAccount(ctx context.Context, account *service.Accoun
if account == nil || account.ID <= 0 {
return nil
}
cacheableAccounts, err := c.writeAccounts(ctx, []service.Account{*account})
accountIDs, err := c.writeAccountIDs(ctx, []service.Account{*account})
if err != nil {
return err
}
if len(cacheableAccounts) == 0 {
if len(accountIDs) == 0 {
return c.DeleteAccount(ctx, account.ID)
}
return nil
@@ -719,29 +699,13 @@ func decodeCachedAccount(val any) (*service.Account, error) {
return &account, nil
}
func (c *schedulerCache) writeAccounts(ctx context.Context, accounts []service.Account) ([]service.Account, error) {
cacheableAccounts, _, err := c.writeAccountPayloads(ctx, accounts, false)
return cacheableAccounts, err
}
func (c *schedulerCache) writeAccountIDs(ctx context.Context, accounts []service.Account) ([]int64, error) {
_, accountIDs, err := c.writeAccountPayloads(ctx, accounts, true)
return accountIDs, err
}
func (c *schedulerCache) writeAccountPayloads(ctx context.Context, accounts []service.Account, collectIDs bool) ([]service.Account, []int64, error) {
if len(accounts) == 0 {
return nil, nil, nil
return nil, nil
}
pipe := c.rdb.Pipeline()
var cacheableAccounts []service.Account
var accountIDs []int64
if collectIDs {
accountIDs = make([]int64, 0, len(accounts))
} else {
cacheableAccounts = make([]service.Account, 0, len(accounts))
}
accountIDs := make([]int64, 0, len(accounts))
pending := 0
flush := func() error {
if pending == 0 {
@@ -768,24 +732,19 @@ func (c *schedulerCache) writeAccountPayloads(ctx context.Context, accounts []se
id := strconv.FormatInt(account.ID, 10)
pipe.Set(ctx, schedulerAccountKey(id), fullPayload, 0)
pipe.Set(ctx, schedulerAccountMetaKey(id), metaPayload, 0)
// 复用路径只保留有序 ID，避免先物化完整账号切片再做第二次扫描。
if collectIDs {
accountIDs = append(accountIDs, account.ID)
} else {
cacheableAccounts = append(cacheableAccounts, account)
}
accountIDs = append(accountIDs, account.ID)
pending++
if pending >= c.writeChunkSize {
if err := flush(); err != nil {
return nil, nil, err
return nil, err
}
}
}
if err := flush(); err != nil {
return nil, nil, err
return nil, err
}
return cacheableAccounts, accountIDs, nil
return accountIDs, nil
}
func marshalSchedulerCacheAccount(account service.Account) ([]byte, []byte, error) {
backend/internal/repository/scheduler_cache_unit_test.go
+55 -5
View File
@@ -33,18 +33,17 @@ func newSchedulerCacheUnitWithRedis(t *testing.T) (*schedulerCache, *miniredis.M
return cache, mr
}
func TestSchedulerCacheWriteAccountsSkipsUnencodableTimes(t *testing.T) {
func TestSchedulerCacheWriteAccountIDsSkipsUnencodableTimes(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{
accountIDs, err := cache.writeAccountIDs(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)
require.Equal(t, []int64{111}, accountIDs)
cached, err := cache.GetAccount(ctx, 111)
require.NoError(t, err)
@@ -144,6 +143,56 @@ func TestSchedulerCacheSnapshotAccountIDReusePreservesPayloadAndMembers(t *testi
require.Nil(t, missing)
}
func TestSchedulerCacheSetSnapshotMatchesIDPublishing(t *testing.T) {
ctx := context.Background()
cache, _ := newSchedulerCacheUnitWithRedis(t)
invalidTime := time.Date(10000, time.January, 1, 0, 0, 0, 0, time.UTC)
validOne := service.Account{
ID: 721,
Name: "first",
Platform: service.PlatformOpenAI,
Type: service.AccountTypeOAuth,
Credentials: map[string]any{"model_mapping": map[string]any{"source": "target"}},
Extra: map[string]any{"mixed_scheduling": true},
GroupIDs: []int64{21},
}
validTwo := service.Account{ID: 722, Name: "second", Platform: service.PlatformOpenAI, Type: service.AccountTypeAPIKey}
invalid := service.Account{ID: 799, Platform: service.PlatformOpenAI, Type: service.AccountTypeAPIKey, ExpiresAt: &invalidTime}
accounts := []service.Account{validOne, invalid, validTwo, validOne}
normal := service.SchedulerBucket{GroupID: 21, Platform: service.PlatformOpenAI, Mode: service.SchedulerModeSingle}
normalToken, err := cache.CaptureBucketWriteToken(ctx, normal)
require.NoError(t, err)
require.NoError(t, cache.SetSnapshot(ctx, normal, normalToken, accounts))
fullBefore, err := cache.rdb.Get(ctx, schedulerAccountKey("721")).Bytes()
require.NoError(t, err)
metaBefore, err := cache.rdb.Get(ctx, schedulerAccountMetaKey("721")).Bytes()
require.NoError(t, err)
idOnly := service.SchedulerBucket{GroupID: 21, Platform: service.PlatformOpenAI, Mode: service.SchedulerModeForced}
idOnlyToken, err := cache.CaptureBucketWriteToken(ctx, idOnly)
require.NoError(t, err)
accountIDs, err := cache.SetSnapshotAndReturnAccountIDs(ctx, idOnly, idOnlyToken, accounts)
require.NoError(t, err)
require.Equal(t, []int64{721, 722, 721}, accountIDs)
fullAfter, err := cache.rdb.Get(ctx, schedulerAccountKey("721")).Bytes()
require.NoError(t, err)
metaAfter, err := cache.rdb.Get(ctx, schedulerAccountMetaKey("721")).Bytes()
require.NoError(t, err)
require.Equal(t, fullBefore, fullAfter, "普通快照和 ID 发布必须写入相同完整账号 payload")
require.Equal(t, metaBefore, metaAfter, "普通快照和 ID 发布必须写入相同元数据 payload")
for _, bucket := range []service.SchedulerBucket{normal, idOnly} {
version, err := cache.rdb.Get(ctx, schedulerBucketKey(schedulerActivePrefix, bucket)).Result()
require.NoError(t, err)
members, err := cache.rdb.ZRange(ctx, schedulerSnapshotKey(bucket, version), 0, -1).Result()
require.NoError(t, err)
require.Equal(t, []string{"722", "721"}, members, bucket.String())
}
}
func TestSchedulerCacheSnapshotAccountIDReuseKeepsEmptySnapshotSemantics(t *testing.T) {
ctx := context.Background()
cache := newSchedulerCacheUnit(t)
@@ -558,7 +607,8 @@ func TestSchedulerCacheActivationIsFencedAfterRetire(t *testing.T) {
require.NoError(t, err)
version, err := cache.allocateSnapshotVersion(ctx, bucket, token)
require.NoError(t, err)
require.NoError(t, cache.writeSnapshotVersion(ctx, bucket, version, []service.Account{account}))
_, err = cache.writeSnapshotVersionAndReturnAccountIDs(ctx, bucket, version, []service.Account{account})
require.NoError(t, err)
// Deterministic race C: retirement and authoritative reopen both happen after
// INCR/write but before the old writer activates.
backend/internal/service/scheduler_snapshot_service.go
+1 -1
View File
@@ -981,7 +981,7 @@ func (s *SchedulerSnapshotService) setRebuildSnapshot(
return err
}
if queries.remaining[key] > 1 {
// 必须保存 writeAccounts 实际接受的有序 ID，不能从原账号切片重新推导；
// 必须保存实际成功编码并写入的有序 ID，不能从原账号切片重新推导；
// 否则不可编码账号会只出现在后续桶中，破坏两个快照的成员一致性。
// 返回切片由当前批次独占，直接接管可避免 10k 账号场景再次复制。
queries.snapshotAccountIDs[key] = accountIDs