package fqapi import ( "fanqie/internal/common" "fanqie/internal/config" "fmt" "sync" "time" ) var ( poolLogger = common.NewLogger("device_key_pool", "INFO", "") globalPool *DeviceKeyPool poolOnce sync.Once ) // DeviceKeyItem 设备密钥项 type DeviceKeyItem struct { ID int `json:"id"` DeviceHeader string `json:"device_header"` Key string `json:"key"` KeyVer int `json:"key_ver"` KeyTs int64 `json:"key_ts"` // 密钥注册时间戳(key_register_ts) IsValid bool `json:"is_valid"` CreatedAt time.Time `json:"created_at"` LastUsedAt time.Time `json:"last_used_at"` // 仅内存字段(不落库) KeyFetchedAt time.Time `json:"-"` // 密钥获取时间,超 24h 异步刷新 KeyRefreshAt time.Time `json:"-"` // 上次刷新尝试时间,防刷新风暴 } // DeviceKeyPool 设备密钥池 type DeviceKeyPool struct { mu sync.RWMutex currentIndex int items []*DeviceKeyItem enabled bool } // GetDeviceKeyPool 获取全局设备密钥池实例 func GetDeviceKeyPool() *DeviceKeyPool { poolOnce.Do(func() { globalPool = &DeviceKeyPool{ items: make([]*DeviceKeyItem, 0), enabled: false, } }) return globalPool } // Initialize 初始化设备密钥池 func (p *DeviceKeyPool) Initialize() error { cfg := config.GetConfig() if cfg == nil || !cfg.DeviceKeyPool.Enabled { poolLogger.Info("设备密钥池未启用") p.enabled = false return nil } p.enabled = true poolLogger.Info("开始初始化设备密钥池...") // 从数据库加载已有的设备 if err := p.loadFromDB(); err != nil { poolLogger.Warning("从数据库加载设备失败: " + err.Error()) } // 检查并补充设备到初始化数量 validCount := p.getValidCount() initialSize := cfg.DeviceKeyPool.InitialSize if validCount < initialSize { needCount := initialSize - validCount poolLogger.Info(fmt.Sprintf("当前有效设备数: %d, 需要补充: %d", validCount, needCount)) for i := 0; i < needCount; i++ { if err := p.addNewDevice(); err != nil { poolLogger.Error(fmt.Sprintf("添加设备失败: %v", err)) } } } poolLogger.Info(fmt.Sprintf("设备密钥池初始化完成,当前有效设备数: %d", p.getValidCount())) // last_used_at 批量落库(GetNext 不再同步写库) p.StartDevicePoolFlusher() return nil } // IsEnabled 检查设备密钥池是否启用 func (p *DeviceKeyPool) IsEnabled() bool { return p.enabled } // GetNext 获取下一个可用的设备密钥(轮询方式) func (p *DeviceKeyPool) GetNext() (*DeviceKeyItem, error) { if !p.enabled { return nil, fmt.Errorf("设备密钥池未启用") } p.mu.Lock() defer p.mu.Unlock() // 过滤有效设备 validItems := make([]*DeviceKeyItem, 0) for _, item := range p.items { if item.IsValid { validItems = append(validItems, item) } } if len(validItems) == 0 { return nil, fmt.Errorf("没有可用的设备") } // 轮询获取下一个设备 item := validItems[p.currentIndex%len(validItems)] p.currentIndex++ // 只更新内存中的最后使用时间,由后台 flusher 定期批量落库 // (此前每次取设备同步 UPDATE,SQLite 单写锁下是吞吐瓶颈) item.LastUsedAt = time.Now() // 密钥超 24h 未刷新:异步换新(本次请求仍用旧密钥,下次生效) if time.Since(item.KeyFetchedAt) > 20*time.Hour && time.Since(item.KeyRefreshAt) > 10*time.Minute { item.KeyRefreshAt = time.Now() header := item.DeviceHeader go func() { if ki, err := GetKey(header); err == nil { UpdateKeyInPool(header, ki) } else { poolLogger.Warning(fmt.Sprintf("刷新设备密钥失败: %v", err)) } }() } // 异步检查是否需要补充 go p.checkAndRefill() return item, nil } // MarkInvalid 标记设备为无效 func (p *DeviceKeyPool) MarkInvalid(deviceHeader string) error { if !p.enabled { return nil } p.mu.Lock() defer p.mu.Unlock() for _, item := range p.items { if item.DeviceHeader == deviceHeader { item.IsValid = false // poolLogger.Warning(fmt.Sprintf("设备 ID=%d 已标记为无效", item.ID)) // 更新数据库 if err := p.updateDeviceInDB(item); err != nil { return err } // 异步补充新设备 go p.checkAndRefill() return nil } } return fmt.Errorf("未找到设备") } // RemoveInvalid 移除所有无效设备并自动补充 func (p *DeviceKeyPool) RemoveInvalid() error { if !p.enabled { return nil } p.mu.Lock() // 从数据库中删除无效设备 if err := p.removeInvalidFromDB(); err != nil { p.mu.Unlock() return err } // 从内存中移除无效设备 validItems := make([]*DeviceKeyItem, 0) for _, item := range p.items { if item.IsValid { validItems = append(validItems, item) } } removedCount := len(p.items) - len(validItems) p.items = validItems if removedCount > 0 { poolLogger.Info(fmt.Sprintf("已移除 %d 个无效设备", removedCount)) // 立即检查并补充设备(使用已持有锁的版本) p.checkAndRefillLocked() } p.mu.Unlock() return nil } // GetStats 获取池统计信息 func (p *DeviceKeyPool) GetStats() map[string]interface{} { p.mu.RLock() defer p.mu.RUnlock() validCount := 0 invalidCount := 0 for _, item := range p.items { if item.IsValid { validCount++ } else { invalidCount++ } } return map[string]interface{}{ "enabled": p.enabled, "total": len(p.items), "valid": validCount, "invalid": invalidCount, "current_index": p.currentIndex, } } // checkAndRefill 检查并补充设备(私有方法,会自动加锁) func (p *DeviceKeyPool) checkAndRefill() { if !p.enabled { return } cfg := config.GetConfig() if cfg == nil { return } p.mu.Lock() validCount := p.getValidCount() p.mu.Unlock() minSize := cfg.DeviceKeyPool.MinSize maxSize := cfg.DeviceKeyPool.MaxSize if validCount < minSize { needCount := minSize - validCount // 限制不超过最大值 if validCount+needCount > maxSize { needCount = maxSize - validCount } if needCount > 0 { poolLogger.Info(fmt.Sprintf("有效设备数 %d 低于最小值 %d,开始补充 %d 个设备", validCount, minSize, needCount)) for i := 0; i < needCount; i++ { if err := p.addNewDevice(); err != nil { poolLogger.Error(fmt.Sprintf("补充设备失败: %v", err)) } } } } } // checkAndRefillLocked 检查并补充设备(已持有锁的内部版本) func (p *DeviceKeyPool) checkAndRefillLocked() { if !p.enabled { return } cfg := config.GetConfig() if cfg == nil { return } validCount := p.getValidCount() minSize := cfg.DeviceKeyPool.MinSize maxSize := cfg.DeviceKeyPool.MaxSize if validCount < minSize { needCount := minSize - validCount // 限制不超过最大值 if validCount+needCount > maxSize { needCount = maxSize - validCount } if needCount > 0 { poolLogger.Info(fmt.Sprintf("有效设备数 %d 低于最小值 %d,开始补充 %d 个设备", validCount, minSize, needCount)) // 临时释放锁,因为addNewDevice需要加锁 p.mu.Unlock() for i := 0; i < needCount; i++ { if err := p.addNewDevice(); err != nil { poolLogger.Error(fmt.Sprintf("补充设备失败: %v", err)) } } p.mu.Lock() } } } // getValidCount 获取有效设备数量(调用前需已加锁或在RLock内) func (p *DeviceKeyPool) getValidCount() int { count := 0 for _, item := range p.items { if item.IsValid { count++ } } return count } // addNewDevice 添加新设备(会自动加锁) func (p *DeviceKeyPool) addNewDevice() error { // 注册新设备 deviceHeader, err := Register() if err != nil { return fmt.Errorf("注册设备失败: %w", err) } // 获取密钥 keyInfo, err := GetKey(deviceHeader) if err != nil { return fmt.Errorf("获取密钥失败: %w", err) } // 创建设备项 item := &DeviceKeyItem{ DeviceHeader: deviceHeader, Key: keyInfo.Key, KeyVer: keyInfo.Ver, KeyTs: keyInfo.Ts, IsValid: true, CreatedAt: time.Now(), LastUsedAt: time.Now(), KeyFetchedAt: time.Now(), } // 保存到数据库 if err := p.saveDeviceToDB(item); err != nil { return fmt.Errorf("保存设备到数据库失败: %w", err) } // 添加到内存 p.mu.Lock() p.items = append(p.items, item) p.mu.Unlock() // poolLogger.Info(fmt.Sprintf("已添加新设备 ID=%d", item.ID)) return nil } // ===== 数据库操作方法 ===== // loadFromDB 从数据库加载设备 func (p *DeviceKeyPool) loadFromDB() error { if cacheDB == nil { return fmt.Errorf("数据库未初始化") } rows, err := cacheDB.Query(` SELECT id, device_header, key, key_ver, COALESCE(key_ts, 0), is_valid, created_at, last_used_at FROM device_pool ORDER BY id `) if err != nil { return err } defer rows.Close() items := make([]*DeviceKeyItem, 0) for rows.Next() { var item DeviceKeyItem var createdAt, lastUsedAt int64 err := rows.Scan( &item.ID, &item.DeviceHeader, &item.Key, &item.KeyVer, &item.KeyTs, &item.IsValid, &createdAt, &lastUsedAt, ) if err != nil { poolLogger.Error("扫描行失败: " + err.Error()) continue } item.CreatedAt = time.Unix(createdAt, 0) item.LastUsedAt = time.Unix(lastUsedAt, 0) item.KeyFetchedAt = time.Now() // 旧设备的密钥年龄未知,从加载起算 items = append(items, &item) } p.mu.Lock() p.items = items p.mu.Unlock() poolLogger.Info(fmt.Sprintf("从数据库加载了 %d 个设备", len(items))) return nil } // boolToInt 布尔转整数存储(is_valid 列为 INTEGER,PostgreSQL 严格要求类型匹配) func boolToInt(b bool) int { if b { return 1 } return 0 } // saveDeviceToDB 保存设备到数据库 func (p *DeviceKeyPool) saveDeviceToDB(item *DeviceKeyItem) error { if cacheDB == nil { return fmt.Errorf("数据库未初始化") } // RETURNING id 同时兼容 SQLite(3.35+) 与 PostgreSQL,替代 LastInsertId() err := cacheDB.QueryRow(q(` INSERT INTO device_pool (device_header, key, key_ver, key_ts, is_valid, created_at, last_used_at, updated_at) VALUES (?, ?, ?, ?, ?, ?, ?, ?) RETURNING id `), item.DeviceHeader, item.Key, item.KeyVer, item.KeyTs, boolToInt(item.IsValid), item.CreatedAt.Unix(), item.LastUsedAt.Unix(), item.CreatedAt.Unix()).Scan(&item.ID) return err } // updateDeviceInDB 更新设备信息 func (p *DeviceKeyPool) updateDeviceInDB(item *DeviceKeyItem) error { if cacheDB == nil { return fmt.Errorf("数据库未初始化") } _, err := cacheDB.Exec(q(` UPDATE device_pool SET device_header = ?, key = ?, key_ver = ?, key_ts = ?, is_valid = ?, last_used_at = ?, updated_at = ? WHERE id = ? `), item.DeviceHeader, item.Key, item.KeyVer, item.KeyTs, boolToInt(item.IsValid), item.LastUsedAt.Unix(), item.LastUsedAt.Unix(), item.ID) return err } // StartDevicePoolFlusher 启动 last_used_at 批量落库任务(GetNext 已改为只写内存) func (p *DeviceKeyPool) StartDevicePoolFlusher() { go func() { for range time.Tick(time.Minute) { FlushDevicePool() } }() poolLogger.Info("设备池 last_used_at 批量落库已启动(每分钟)") } // FlushDevicePool 将内存中的设备使用时间批量刷入数据库 // ponytail: 进程退出时最多丢 1 分钟使用记录(仅统计用途),无需精确停表 func FlushDevicePool() { p := GetDeviceKeyPool() if !p.enabled || cacheDB == nil { return } p.mu.RLock() items := append([]*DeviceKeyItem(nil), p.items...) p.mu.RUnlock() for _, it := range items { p.updateDeviceInDB(it) } } // removeInvalidFromDB 从数据库移除无效设备 func (p *DeviceKeyPool) removeInvalidFromDB() error { if cacheDB == nil { return fmt.Errorf("数据库未初始化") } result, err := cacheDB.Exec(`DELETE FROM device_pool WHERE is_valid = 0`) if err != nil { return err } affected, _ := result.RowsAffected() if affected > 0 { poolLogger.Info(fmt.Sprintf("从数据库删除了 %d 个无效设备", affected)) } return nil } // GetDeviceHeaderFromPool 从池中获取设备header(兼容旧接口) func GetDeviceHeaderFromPool() string { pool := GetDeviceKeyPool() if !pool.IsEnabled() { return "" } item, err := pool.GetNext() if err != nil { poolLogger.Error("从池中获取设备失败: " + err.Error()) return "" } return item.DeviceHeader } // GetKeyInfoFromPool 从池中获取密钥信息(兼容旧接口) func GetKeyInfoFromPool(deviceHeader string) *KeyInfo { pool := GetDeviceKeyPool() if !pool.IsEnabled() { return nil } pool.mu.RLock() defer pool.mu.RUnlock() for _, item := range pool.items { if item.DeviceHeader == deviceHeader && item.IsValid { return &KeyInfo{ Key: item.Key, Ver: item.KeyVer, } } } return nil } // UpdateKeyInPool 更新池中设备的密钥(含注册时间戳) func UpdateKeyInPool(deviceHeader string, ki *KeyInfo) error { pool := GetDeviceKeyPool() if !pool.IsEnabled() || ki == nil { return nil } pool.mu.Lock() defer pool.mu.Unlock() for _, item := range pool.items { if item.DeviceHeader == deviceHeader { item.Key = ki.Key item.KeyVer = ki.Ver item.KeyTs = ki.Ts item.KeyFetchedAt = time.Now() return pool.updateDeviceInDB(item) } } return fmt.Errorf("未找到设备") } // GetAnyValidDeviceHeader 获取任意一个有效设备的header(用于辅助功能) func GetAnyValidDeviceHeader() string { pool := GetDeviceKeyPool() if !pool.IsEnabled() { return "" } pool.mu.RLock() defer pool.mu.RUnlock() for _, item := range pool.items { if item.IsValid { return item.DeviceHeader } } return "" } // GetValidDevice 获取一个有效的设备 func GetValidDevice() (*DeviceKeyItem, error) { pool := GetDeviceKeyPool() if !pool.IsEnabled() { return nil, fmt.Errorf("设备池未启用") } return pool.GetNext() } // GetCurrentDevice 获取当前使用的设备(和GetValidDevice一样,返回下一个可用设备) func GetCurrentDevice() *DeviceKeyItem { pool := GetDeviceKeyPool() if !pool.IsEnabled() { return nil } item, _ := pool.GetNext() return item }