feat: 重新支持apikey费用排序功能

2026-01-23 00:53:33 +00:00 · 2025-11-28 15:32:50 +08:00
parent d9476230c6
commit 28caa93d99
5 changed files with 1165 additions and 20 deletions
--- a/src/services/apiKeyService.js
+++ b/src/services/apiKeyService.js
@@ -158,6 +158,14 @@ class ApiKeyService {
    // 保存API Key数据并建立哈希映射
    await redis.setApiKey(keyId, keyData, hashedKey)

+    // 同步添加到费用排序索引
+    try {
+      const costRankService = require('./costRankService')
+      await costRankService.addKeyToIndexes(keyId)
+    } catch (err) {
+      logger.warn(`Failed to add key ${keyId} to cost rank indexes:`, err.message)
+    }
+
    logger.success(`🔑 Generated new API key: ${name} (${keyId})`)

    return {
@@ -756,6 +764,14 @@ class ApiKeyService {
        await redis.deleteApiKeyHash(keyData.apiKey)
      }

+      // 从费用排序索引中移除
+      try {
+        const costRankService = require('./costRankService')
+        await costRankService.removeKeyFromIndexes(keyId)
+      } catch (err) {
+        logger.warn(`Failed to remove key ${keyId} from cost rank indexes:`, err.message)
+      }
+
      logger.success(`🗑️ Soft deleted API key: ${keyId} by ${deletedBy} (${deletedByType})`)

      return { success: true }
@@ -807,6 +823,14 @@ class ApiKeyService {
        })
      }

+      // 重新添加到费用排序索引
+      try {
+        const costRankService = require('./costRankService')
+        await costRankService.addKeyToIndexes(keyId)
+      } catch (err) {
+        logger.warn(`Failed to add restored key ${keyId} to cost rank indexes:`, err.message)
+      }
+
      logger.success(`✅ Restored API key: ${keyId} by ${restoredBy} (${restoredByType})`)

      return { success: true, apiKey: updatedData }
--- a/src/services/costRankService.js
+++ b/src/services/costRankService.js
@@ -0,0 +1,591 @@
+/**
+ * 费用排序索引服务
+ *
+ * 为 API Keys 提供按费用排序的功能，使用 Redis Sorted Set 预计算排序索引
+ * 支持 today/7days/30days/all 四种固定时间范围的预计算索引
+ * 支持 custom 时间范围的实时计算
+ *
+ * 设计原则：
+ * - 只计算未删除的 API Key
+ * - 使用原子操作避免竞态条件
+ * - 提供增量更新接口供 API Key 创建/删除时调用
+ */
+
+const redis = require('../models/redis')
+const logger = require('../utils/logger')
+
+// ============================================================================
+// 常量配置
+// ============================================================================
+
+/** 时间范围更新间隔配置（省资源模式） */
+const UPDATE_INTERVALS = {
+  today: 10 * 60 * 1000, // 10分钟
+  '7days': 30 * 60 * 1000, // 30分钟
+  '30days': 60 * 60 * 1000, // 1小时
+  all: 2 * 60 * 60 * 1000 // 2小时
+}
+
+/** 支持的时间范围列表 */
+const VALID_TIME_RANGES = ['today', '7days', '30days', 'all']
+
+/** 分布式锁超时时间（秒） */
+const LOCK_TTL = 300
+
+/** 批处理大小 */
+const BATCH_SIZE = 100
+
+// ============================================================================
+// Redis Key 生成器（集中管理 key 格式）
+// ============================================================================
+
+const RedisKeys = {
+  /** 费用排序索引 Sorted Set */
+  rankKey: (timeRange) => `cost_rank:${timeRange}`,
+
+  /** 临时索引 key（用于原子替换） */
+  tempRankKey: (timeRange) => `cost_rank:${timeRange}:temp:${Date.now()}`,
+
+  /** 索引元数据 Hash */
+  metaKey: (timeRange) => `cost_rank_meta:${timeRange}`,
+
+  /** 更新锁 */
+  lockKey: (timeRange) => `cost_rank_lock:${timeRange}`,
+
+  /** 每日费用 */
+  dailyCost: (keyId, date) => `usage:cost:daily:${keyId}:${date}`,
+
+  /** 总费用 */
+  totalCost: (keyId) => `usage:cost:total:${keyId}`
+}
+
+// ============================================================================
+// CostRankService 类
+// ============================================================================
+
+class CostRankService {
+  constructor() {
+    this.timers = {}
+    this.isInitialized = false
+  }
+
+  // --------------------------------------------------------------------------
+  // 生命周期管理
+  // --------------------------------------------------------------------------
+
+  /**
+   * 初始化服务：启动定时任务
+   * 幂等设计：多次调用只会初始化一次
+   */
+  async initialize() {
+    // 先清理可能存在的旧定时器（支持热重载）
+    this._clearAllTimers()
+
+    if (this.isInitialized) {
+      logger.warn('CostRankService already initialized, re-initializing...')
+    }
+
+    logger.info('🔄 Initializing CostRankService...')
+
+    try {
+      // 启动时立即更新所有索引（异步，不阻塞启动）
+      this.updateAllRanks().catch((err) => {
+        logger.error('Failed to initialize cost ranks:', err)
+      })
+
+      // 设置定时更新
+      for (const [timeRange, interval] of Object.entries(UPDATE_INTERVALS)) {
+        this.timers[timeRange] = setInterval(() => {
+          this.updateRank(timeRange).catch((err) => {
+            logger.error(`Failed to update cost rank for ${timeRange}:`, err)
+          })
+        }, interval)
+      }
+
+      this.isInitialized = true
+      logger.success('✅ CostRankService initialized')
+    } catch (error) {
+      logger.error('❌ Failed to initialize CostRankService:', error)
+      throw error
+    }
+  }
+
+  /**
+   * 关闭服务：清理定时器
+   */
+  shutdown() {
+    this._clearAllTimers()
+    this.isInitialized = false
+    logger.info('CostRankService shutdown')
+  }
+
+  /**
+   * 清理所有定时器
+   * @private
+   */
+  _clearAllTimers() {
+    for (const timer of Object.values(this.timers)) {
+      clearInterval(timer)
+    }
+    this.timers = {}
+  }
+
+  // --------------------------------------------------------------------------
+  // 索引更新（全量）
+  // --------------------------------------------------------------------------
+
+  /**
+   * 更新所有时间范围的索引
+   */
+  async updateAllRanks() {
+    for (const timeRange of VALID_TIME_RANGES) {
+      try {
+        await this.updateRank(timeRange)
+      } catch (error) {
+        logger.error(`Failed to update rank for ${timeRange}:`, error)
+      }
+    }
+  }
+
+  /**
+   * 更新指定时间范围的排序索引
+   * @param {string} timeRange - 时间范围
+   */
+  async updateRank(timeRange) {
+    const client = redis.getClient()
+    if (!client) {
+      logger.warn('Redis client not available, skipping cost rank update')
+      return
+    }
+
+    const lockKey = RedisKeys.lockKey(timeRange)
+    const rankKey = RedisKeys.rankKey(timeRange)
+    const metaKey = RedisKeys.metaKey(timeRange)
+
+    // 获取分布式锁
+    const acquired = await client.set(lockKey, '1', 'NX', 'EX', LOCK_TTL)
+    if (!acquired) {
+      logger.debug(`Skipping ${timeRange} rank update - another update in progress`)
+      return
+    }
+
+    const startTime = Date.now()
+
+    try {
+      // 标记为更新中
+      await client.hset(metaKey, 'status', 'updating')
+
+      // 1. 获取所有未删除的 API Key IDs
+      const keyIds = await this._getActiveApiKeyIds()
+
+      if (keyIds.length === 0) {
+        // 无数据时清空索引
+        await client.del(rankKey)
+        await this._updateMeta(client, metaKey, startTime, 0)
+        return
+      }
+
+      // 2. 计算日期范围
+      const dateRange = this._getDateRange(timeRange)
+
+      // 3. 分批计算费用
+      const costs = await this._calculateCostsInBatches(keyIds, dateRange)
+
+      // 4. 原子更新索引（使用临时 key + RENAME 避免竞态条件）
+      await this._atomicUpdateIndex(client, rankKey, costs)
+
+      // 5. 更新元数据
+      await this._updateMeta(client, metaKey, startTime, keyIds.length)
+
+      logger.info(
+        `📊 Updated cost rank for ${timeRange}: ${keyIds.length} keys in ${Date.now() - startTime}ms`
+      )
+    } catch (error) {
+      await client.hset(metaKey, 'status', 'failed')
+      logger.error(`Failed to update cost rank for ${timeRange}:`, error)
+      throw error
+    } finally {
+      await client.del(lockKey)
+    }
+  }
+
+  /**
+   * 原子更新索引（避免竞态条件）
+   * @private
+   */
+  async _atomicUpdateIndex(client, rankKey, costs) {
+    if (costs.size === 0) {
+      await client.del(rankKey)
+      return
+    }
+
+    // 使用临时 key 构建新索引
+    const tempKey = `${rankKey}:temp:${Date.now()}`
+
+    try {
+      // 构建 ZADD 参数
+      const members = []
+      costs.forEach((cost, keyId) => {
+        members.push(cost, keyId)
+      })
+
+      // 写入临时 key
+      await client.zadd(tempKey, ...members)
+
+      // 原子替换（RENAME 是原子操作）
+      await client.rename(tempKey, rankKey)
+    } catch (error) {
+      // 清理临时 key
+      await client.del(tempKey).catch(() => {})
+      throw error
+    }
+  }
+
+  /**
+   * 更新元数据
+   * @private
+   */
+  async _updateMeta(client, metaKey, startTime, keyCount) {
+    await client.hmset(metaKey, {
+      lastUpdate: new Date().toISOString(),
+      keyCount: keyCount.toString(),
+      status: 'ready',
+      updateDuration: (Date.now() - startTime).toString()
+    })
+  }
+
+  // --------------------------------------------------------------------------
+  // 索引增量更新（供外部调用）
+  // --------------------------------------------------------------------------
+
+  /**
+   * 添加 API Key 到所有索引（创建 API Key 时调用）
+   * @param {string} keyId - API Key ID
+   */
+  async addKeyToIndexes(keyId) {
+    const client = redis.getClient()
+    if (!client) {
+      return
+    }
+
+    try {
+      const pipeline = client.pipeline()
+
+      // 将新 Key 添加到所有索引，初始分数为 0
+      for (const timeRange of VALID_TIME_RANGES) {
+        pipeline.zadd(RedisKeys.rankKey(timeRange), 0, keyId)
+      }
+
+      await pipeline.exec()
+      logger.debug(`Added key ${keyId} to cost rank indexes`)
+    } catch (error) {
+      logger.error(`Failed to add key ${keyId} to cost rank indexes:`, error)
+    }
+  }
+
+  /**
+   * 从所有索引中移除 API Key（删除 API Key 时调用）
+   * @param {string} keyId - API Key ID
+   */
+  async removeKeyFromIndexes(keyId) {
+    const client = redis.getClient()
+    if (!client) {
+      return
+    }
+
+    try {
+      const pipeline = client.pipeline()
+
+      // 从所有索引中移除
+      for (const timeRange of VALID_TIME_RANGES) {
+        pipeline.zrem(RedisKeys.rankKey(timeRange), keyId)
+      }
+
+      await pipeline.exec()
+      logger.debug(`Removed key ${keyId} from cost rank indexes`)
+    } catch (error) {
+      logger.error(`Failed to remove key ${keyId} from cost rank indexes:`, error)
+    }
+  }
+
+  // --------------------------------------------------------------------------
+  // 查询接口
+  // --------------------------------------------------------------------------
+
+  /**
+   * 获取排序后的 keyId 列表
+   * @param {string} timeRange - 时间范围
+   * @param {string} sortOrder - 排序方向 'asc' | 'desc'
+   * @param {number} offset - 偏移量
+   * @param {number} limit - 限制数量，-1 表示全部
+   * @returns {Promise<string[]>} keyId 列表
+   */
+  async getSortedKeyIds(timeRange, sortOrder = 'desc', offset = 0, limit = -1) {
+    const client = redis.getClient()
+    if (!client) {
+      throw new Error('Redis client not available')
+    }
+
+    const rankKey = RedisKeys.rankKey(timeRange)
+    const end = limit === -1 ? -1 : offset + limit - 1
+
+    if (sortOrder === 'desc') {
+      return await client.zrevrange(rankKey, offset, end)
+    } else {
+      return await client.zrange(rankKey, offset, end)
+    }
+  }
+
+  /**
+   * 获取 Key 的费用分数
+   * @param {string} timeRange - 时间范围
+   * @param {string} keyId - API Key ID
+   * @returns {Promise<number>} 费用
+   */
+  async getKeyCost(timeRange, keyId) {
+    const client = redis.getClient()
+    if (!client) {
+      return 0
+    }
+
+    const score = await client.zscore(RedisKeys.rankKey(timeRange), keyId)
+    return score ? parseFloat(score) : 0
+  }
+
+  /**
+   * 批量获取多个 Key 的费用分数
+   * @param {string} timeRange - 时间范围
+   * @param {string[]} keyIds - API Key ID 列表
+   * @returns {Promise<Map<string, number>>} keyId -> cost
+   */
+  async getBatchKeyCosts(timeRange, keyIds) {
+    const client = redis.getClient()
+    if (!client || keyIds.length === 0) {
+      return new Map()
+    }
+
+    const rankKey = RedisKeys.rankKey(timeRange)
+    const costs = new Map()
+
+    const pipeline = client.pipeline()
+    keyIds.forEach((keyId) => {
+      pipeline.zscore(rankKey, keyId)
+    })
+    const results = await pipeline.exec()
+
+    keyIds.forEach((keyId, index) => {
+      const [err, score] = results[index]
+      costs.set(keyId, err || !score ? 0 : parseFloat(score))
+    })
+
+    return costs
+  }
+
+  /**
+   * 获取所有排序索引的状态
+   * @returns {Promise<Object>} 各时间范围的状态
+   */
+  async getRankStatus() {
+    const client = redis.getClient()
+    if (!client) {
+      return {}
+    }
+
+    const status = {}
+
+    for (const timeRange of VALID_TIME_RANGES) {
+      const meta = await client.hgetall(RedisKeys.metaKey(timeRange))
+      status[timeRange] = {
+        lastUpdate: meta.lastUpdate || null,
+        keyCount: parseInt(meta.keyCount || 0),
+        status: meta.status || 'unknown',
+        updateDuration: parseInt(meta.updateDuration || 0)
+      }
+    }
+
+    return status
+  }
+
+  /**
+   * 强制刷新指定时间范围的索引
+   * @param {string} timeRange - 时间范围，不传则刷新全部
+   */
+  async forceRefresh(timeRange = null) {
+    if (timeRange) {
+      await this.updateRank(timeRange)
+    } else {
+      await this.updateAllRanks()
+    }
+  }
+
+  // --------------------------------------------------------------------------
+  // Custom 时间范围实时计算
+  // --------------------------------------------------------------------------
+
+  /**
+   * 计算 custom 时间范围的费用（实时计算，排除已删除的 Key）
+   * @param {string} startDate - 开始日期 YYYY-MM-DD
+   * @param {string} endDate - 结束日期 YYYY-MM-DD
+   * @returns {Promise<Map<string, number>>} keyId -> cost
+   */
+  async calculateCustomRangeCosts(startDate, endDate) {
+    const client = redis.getClient()
+    if (!client) {
+      throw new Error('Redis client not available')
+    }
+
+    logger.info(`📊 Calculating custom range costs: ${startDate} to ${endDate}`)
+    const startTime = Date.now()
+
+    // 1. 获取所有未删除的 API Key IDs
+    const keyIds = await this._getActiveApiKeyIds()
+
+    if (keyIds.length === 0) {
+      return new Map()
+    }
+
+    // 2. 分批计算费用
+    const costs = await this._calculateCostsInBatches(keyIds, { startDate, endDate })
+
+    const duration = Date.now() - startTime
+    logger.info(`📊 Custom range costs calculated: ${keyIds.length} keys in ${duration}ms`)
+
+    return costs
+  }
+
+  // --------------------------------------------------------------------------
+  // 私有辅助方法
+  // --------------------------------------------------------------------------
+
+  /**
+   * 获取所有未删除的 API Key IDs
+   * @private
+   * @returns {Promise<string[]>}
+   */
+  async _getActiveApiKeyIds() {
+    // 使用现有的 scanApiKeyIds 获取所有 ID
+    const allKeyIds = await redis.scanApiKeyIds()
+
+    if (allKeyIds.length === 0) {
+      return []
+    }
+
+    // 批量获取 API Key 数据，过滤已删除的
+    const allKeys = await redis.batchGetApiKeys(allKeyIds)
+
+    return allKeys.filter((k) => !k.isDeleted).map((k) => k.id)
+  }
+
+  /**
+   * 分批计算费用
+   * @private
+   */
+  async _calculateCostsInBatches(keyIds, dateRange) {
+    const costs = new Map()
+
+    for (let i = 0; i < keyIds.length; i += BATCH_SIZE) {
+      const batch = keyIds.slice(i, i + BATCH_SIZE)
+      const batchCosts = await this._calculateBatchCosts(batch, dateRange)
+      batchCosts.forEach((cost, keyId) => costs.set(keyId, cost))
+    }
+
+    return costs
+  }
+
+  /**
+   * 批量计算费用
+   * @private
+   */
+  async _calculateBatchCosts(keyIds, dateRange) {
+    const client = redis.getClient()
+    const costs = new Map()
+
+    if (dateRange.useTotal) {
+      // 'all' 时间范围：直接读取 total cost
+      const pipeline = client.pipeline()
+      keyIds.forEach((keyId) => {
+        pipeline.get(RedisKeys.totalCost(keyId))
+      })
+      const results = await pipeline.exec()
+
+      keyIds.forEach((keyId, index) => {
+        const [err, value] = results[index]
+        costs.set(keyId, err ? 0 : parseFloat(value || 0))
+      })
+    } else {
+      // 特定日期范围：汇总每日费用
+      const dates = this._getDatesBetween(dateRange.startDate, dateRange.endDate)
+
+      const pipeline = client.pipeline()
+      keyIds.forEach((keyId) => {
+        dates.forEach((date) => {
+          pipeline.get(RedisKeys.dailyCost(keyId, date))
+        })
+      })
+      const results = await pipeline.exec()
+
+      let resultIndex = 0
+      keyIds.forEach((keyId) => {
+        let totalCost = 0
+        dates.forEach(() => {
+          const [err, value] = results[resultIndex++]
+          if (!err && value) {
+            totalCost += parseFloat(value)
+          }
+        })
+        costs.set(keyId, totalCost)
+      })
+    }
+
+    return costs
+  }
+
+  /**
+   * 获取日期范围配置
+   * @private
+   */
+  _getDateRange(timeRange) {
+    const now = new Date()
+    const today = redis.getDateStringInTimezone(now)
+
+    switch (timeRange) {
+      case 'today':
+        return { startDate: today, endDate: today }
+      case '7days': {
+        const d7 = new Date(now)
+        d7.setDate(d7.getDate() - 6)
+        return { startDate: redis.getDateStringInTimezone(d7), endDate: today }
+      }
+      case '30days': {
+        const d30 = new Date(now)
+        d30.setDate(d30.getDate() - 29)
+        return { startDate: redis.getDateStringInTimezone(d30), endDate: today }
+      }
+      case 'all':
+        return { useTotal: true }
+      default:
+        throw new Error(`Invalid time range: ${timeRange}`)
+    }
+  }
+
+  /**
+   * 获取两个日期之间的所有日期
+   * @private
+   */
+  _getDatesBetween(startDate, endDate) {
+    const dates = []
+    const current = new Date(startDate)
+    const end = new Date(endDate)
+
+    while (current <= end) {
+      dates.push(
+        `${current.getFullYear()}-${String(current.getMonth() + 1).padStart(2, '0')}-${String(current.getDate()).padStart(2, '0')}`
+      )
+      current.setDate(current.getDate() + 1)
+    }
+
+    return dates
+  }
+}
+
+module.exports = new CostRankService()