diff --git a/.gitignore b/.gitignore index 8854d7e..655e33d 100644 --- a/.gitignore +++ b/.gitignore @@ -84,21 +84,10 @@ node_modules/ .cache .parcel-cache -# Next.js build output -.next - -# Nuxt.js build / generate output -.nuxt -dist - # Gatsby files .cache/ public -# Storybook build outputs -.out -.storybook-out - # Temporary folders tmp/ temp/ @@ -226,5 +215,3 @@ charts/*.tgz run-local.sh start-local.sh dev-setup.sh - -cmd/gpt-load/dist diff --git a/internal/keypool/provider.go b/internal/keypool/provider.go index ff2381c..85f605a 100644 --- a/internal/keypool/provider.go +++ b/internal/keypool/provider.go @@ -187,13 +187,26 @@ func (p *KeyProvider) handleFailure(keyID uint, keyHashKey, activeKeysListKey st // LoadKeysFromDB 从数据库加载所有分组和密钥,并填充到 Store 中。 func (p *KeyProvider) LoadKeysFromDB() error { + initFlagKey := "initialization:db_keys_loaded" + + exists, err := p.store.Exists(initFlagKey) + if err != nil { + return fmt.Errorf("failed to check initialization flag: %w", err) + } + + if exists { + logrus.Debug("Keys have already been loaded into the store. Skipping.") + return nil + } + + logrus.Debug("First time startup, loading keys from DB...") // 1. 分批从数据库加载并使用 Pipeline 写入 Redis allActiveKeyIDs := make(map[uint][]any) batchSize := 1000 var batchKeys []*models.APIKey - err := p.db.Model(&models.APIKey{}).FindInBatches(&batchKeys, batchSize, func(tx *gorm.DB, batch int) error { + err = p.db.Model(&models.APIKey{}).FindInBatches(&batchKeys, batchSize, func(tx *gorm.DB, batch int) error { logrus.Debugf("Processing batch %d with %d keys...", batch, len(batchKeys)) var pipeline store.Pipeliner @@ -242,6 +255,10 @@ func (p *KeyProvider) LoadKeysFromDB() error { } } + if err := p.store.Set(initFlagKey, []byte("1"), 0); err != nil { + logrus.WithField("flagKey", initFlagKey).Error("Failed to set initialization flag after loading keys") + } + return nil }