package usage import ( "database/sql" "fmt" "sync" "time" _ "github.com/mattn/go-sqlite3" ) // AppUsage 记录单个应用的用量统计。 type AppUsage struct { AppID string `json:"app_id"` RequestCount int64 `json:"request_count"` InputTokens int64 `json:"input_tokens"` OutputTokens int64 `json:"output_tokens"` TotalTokens int64 `json:"total_tokens"` ErrorCount int64 `json:"error_count"` WindowStart string `json:"window_start"` WindowEnd string `json:"window_end"` } // Tracker 管理 per-app 用量统计(内存滑动窗口 + SQLite 持久化)。 type Tracker struct { mu sync.RWMutex apps map[string]*appCounters window time.Duration db *sql.DB logCh chan usageLogEntry // 异步写入 channel stopCh chan struct{} retentionDays int // 历史数据保留天数 } type appCounters struct { requestCount int64 inputTokens int64 outputTokens int64 errorCount int64 windowStart time.Time } type usageLogEntry struct { appID string inputTokens int outputTokens int isError bool recordedAt time.Time } // NewTracker 创建一个用量统计器。 // window 为统计窗口时长(如 1 小时),窗口到期后自动重置。 // dbPath 为 SQLite 持久化路径,空字符串则仅用内存。 // retentionDays 为历史数据保留天数,超过此天数的数据自动清理(0 表示不清理)。 func NewTracker(window time.Duration, dbPath string, retentionDays int) (*Tracker, error) { t := &Tracker{ apps: make(map[string]*appCounters), window: window, logCh: make(chan usageLogEntry, 1024), stopCh: make(chan struct{}), retentionDays: retentionDays, } if dbPath != "" { db, err := sql.Open("sqlite3", dbPath) if err != nil { return nil, fmt.Errorf("open usage db: %w", err) } if err := initUsageDB(db); err != nil { return nil, fmt.Errorf("init usage db: %w", err) } t.db = db go t.logWriter() // 启动异步写入 goroutine } return t, nil } func initUsageDB(db *sql.DB) error { schema := ` CREATE TABLE IF NOT EXISTS usage_logs ( id INTEGER PRIMARY KEY AUTOINCREMENT, app_id TEXT NOT NULL, input_tokens INTEGER NOT NULL DEFAULT 0, output_tokens INTEGER NOT NULL DEFAULT 0, is_error INTEGER NOT NULL DEFAULT 0, recorded_at TEXT NOT NULL ); CREATE INDEX IF NOT EXISTS idx_usage_app ON usage_logs(app_id); CREATE INDEX IF NOT EXISTS idx_usage_time ON usage_logs(recorded_at);` _, err := db.Exec(schema) return err } // Record 记录一次请求的用量。 func (t *Tracker) Record(appID string, inputTokens, outputTokens int, isError bool) { now := time.Now() // 更新内存计数器 t.mu.Lock() c, ok := t.apps[appID] if !ok { c = &appCounters{windowStart: now} t.apps[appID] = c } // 窗口过期则重置 if now.Sub(c.windowStart) > t.window { c.requestCount = 0 c.inputTokens = 0 c.outputTokens = 0 c.errorCount = 0 c.windowStart = now } c.requestCount++ c.inputTokens += int64(inputTokens) c.outputTokens += int64(outputTokens) if isError { c.errorCount++ } t.mu.Unlock() // 异步写入 SQLite(非阻塞,channel 满则丢弃) if t.db != nil { select { case t.logCh <- usageLogEntry{appID, inputTokens, outputTokens, isError, now}: default: // channel 满则丢弃,避免阻塞主流程 } } } // logWriter 后台 goroutine,从 channel 读取用量日志并批量写入 SQLite。 func (t *Tracker) logWriter() { batch := make([]usageLogEntry, 0, 64) ticker := time.NewTicker(5 * time.Second) defer ticker.Stop() flush := func() { if len(batch) == 0 { return } tx, err := t.db.Begin() if err != nil { batch = batch[:0] return } stmt, _ := tx.Prepare(`INSERT INTO usage_logs (app_id, input_tokens, output_tokens, is_error, recorded_at) VALUES (?, ?, ?, ?, ?)`) if stmt != nil { for _, e := range batch { errInt := 0 if e.isError { errInt = 1 } stmt.Exec(e.appID, e.inputTokens, e.outputTokens, errInt, e.recordedAt.Format(time.RFC3339)) } stmt.Close() } tx.Commit() batch = batch[:0] } for { select { case e := <-t.logCh: batch = append(batch, e) if len(batch) >= 64 { flush() } case <-ticker.C: flush() // 定期清理过期数据 if t.retentionDays > 0 { t.cleanupOldLogs() } case <-t.stopCh: // 排空 channel 后 flush for len(t.logCh) > 0 { batch = append(batch, <-t.logCh) } flush() return } } } // Get 返回指定应用的当前用量快照。 func (t *Tracker) Get(appID string) AppUsage { t.mu.RLock() defer t.mu.RUnlock() c, ok := t.apps[appID] if !ok { return AppUsage{AppID: appID} } return AppUsage{ AppID: appID, RequestCount: c.requestCount, InputTokens: c.inputTokens, OutputTokens: c.outputTokens, TotalTokens: c.inputTokens + c.outputTokens, ErrorCount: c.errorCount, WindowStart: c.windowStart.Format(time.RFC3339), WindowEnd: c.windowStart.Add(t.window).Format(time.RFC3339), } } // GetAll 返回所有应用的用量快照。 func (t *Tracker) GetAll() []AppUsage { t.mu.RLock() defer t.mu.RUnlock() result := make([]AppUsage, 0, len(t.apps)) for appID, c := range t.apps { result = append(result, AppUsage{ AppID: appID, RequestCount: c.requestCount, InputTokens: c.inputTokens, OutputTokens: c.outputTokens, TotalTokens: c.inputTokens + c.outputTokens, ErrorCount: c.errorCount, WindowStart: c.windowStart.Format(time.RFC3339), WindowEnd: c.windowStart.Add(t.window).Format(time.RFC3339), }) } return result } // GetHistory 返回指定应用的历史用量(从 SQLite 查询,按小时聚合)。 // hours 为查询的时间范围(最近 N 小时)。 func (t *Tracker) GetHistory(appID string, hours int) ([]map[string]any, error) { if t.db == nil { return nil, fmt.Errorf("persistence not enabled") } if hours <= 0 { hours = 24 } since := time.Now().Add(-time.Duration(hours) * time.Hour).Format(time.RFC3339) rows, err := t.db.Query( `SELECT substr(recorded_at, 1, 13) as hour_bucket, COUNT(*) as request_count, SUM(input_tokens) as input_tokens, SUM(output_tokens) as output_tokens, SUM(is_error) as error_count FROM usage_logs WHERE app_id = ? AND recorded_at >= ? GROUP BY hour_bucket ORDER BY hour_bucket DESC`, appID, since, ) if err != nil { return nil, fmt.Errorf("query usage history: %w", err) } defer rows.Close() var result []map[string]any for rows.Next() { var hourBucket string var reqCount, inputTokens, outputTokens, errCount int64 if err := rows.Scan(&hourBucket, &reqCount, &inputTokens, &outputTokens, &errCount); err != nil { return nil, err } result = append(result, map[string]any{ "hour": hourBucket, "request_count": reqCount, "input_tokens": inputTokens, "output_tokens": outputTokens, "total_tokens": inputTokens + outputTokens, "error_count": errCount, }) } return result, rows.Err() } // Close 关闭数据库连接。 func (t *Tracker) Close() error { close(t.stopCh) if t.db != nil { return t.db.Close() } return nil } // cleanupOldLogs 清理超过保留期限的用量日志。 func (t *Tracker) cleanupOldLogs() { cutoff := time.Now().Add(-time.Duration(t.retentionDays) * 24 * time.Hour).Format(time.RFC3339) if _, err := t.db.Exec(`DELETE FROM usage_logs WHERE recorded_at < ?`, cutoff); err != nil { // 清理失败不影响主流程,下次再试 return } }