fix: 修复大规模数据性能与流式状态清理

This commit is contained in:
chuan
2026-08-15 18:09:54 +08:00
parent 072ed8f168
commit ddebe3c4a2
15 changed files with 439 additions and 97 deletions
+1 -1
View File
@@ -76,7 +76,7 @@ CLIProxyAPI 负责 HTTP 接入、协议转换、上游凭证和实际请求执
## 环境与构建(WSL/Linux
需要 Go 1.24+ 和 GCC。Ubuntu/Debian 可先准备 C 工具链:
需要支持 toolchain 自动切换的 Go 1.24+ 和 GCC;项目固定使用已包含安全修复的 Go 1.26.6 构建。Ubuntu/Debian 可先准备 C 工具链:
```bash
sudo apt-get update
+2
View File
@@ -124,6 +124,8 @@
价格不存在时返回 HTTP 503 和 `billing_price_unavailable`,不会把未知成本的请求发送到上游。历史迁移数据仍可能因为当时没有价格而显示为成本不可用。
管理员首次补充某个模型价格时,系统按固定小批次补算该模型尚未定价的历史 Usage。补算使用缺失成本专用索引,不会一次性把全部历史载入内存;每批独立提交,在线认证、额度检查和新用量写入可以继续执行。
## 不可变账目
| 账目类型 | 含义 | 金额方向 |
+2
View File
@@ -111,6 +111,8 @@ CLIProxyAPI 会通过两个独立回调提供请求信息:
查询投影只保存明细检索所需的关联和排序字段,Token、成本、生命周期和账目仍以原始事实表为准。当前分页与组合筛选已按百万级记录场景设计,不要求将全部记录读入内存。
SQLite 使用单 writer 和独立读连接池:计费与事实写入保持顺序一致,管理台分页和汇总不会占用 writer。今日用户汇总只扫描当天范围,最近使用时间通过 Key 时间索引定位,不随全部历史记录线性分组扫描。
## 数据一致性
- Usage 插入使用内容事件标识,完全相同的重复回调不会重复记录或扣费。
+2
View File
@@ -2,6 +2,8 @@ module cpa-ext
go 1.24
toolchain go1.26.6
require (
github.com/mattn/go-sqlite3 v1.14.48
gopkg.in/yaml.v3 v3.0.1
+7 -2
View File
@@ -17,6 +17,7 @@ import (
type App struct {
mu sync.RWMutex
priceMu sync.Mutex
config Config
usage *collection.Service
store *repository.SQLiteUsageRepository
@@ -230,6 +231,10 @@ func (a *App) handleRequestComplete(raw []byte) ([]byte, error) {
if a.usage == nil {
return nil, fmt.Errorf("用量数据库尚未初始化")
}
pendingManagedKeyID := ""
if pending, ok := a.pending.LoadAndDelete(completion.RequestID); ok {
pendingManagedKeyID, _ = pending.(string)
}
if err := a.store.CompleteBillingAdmission(context.Background(), completion.RequestID, completion.CompletedAt); err != nil {
return nil, fmt.Errorf("释放并发占用: %w", err)
}
@@ -238,8 +243,8 @@ func (a *App) handleRequestComplete(raw []byte) ([]byte, error) {
if key, keyErr := a.store.ManagedKeyByReference(context.Background(), managedKeyID); keyErr == nil {
managedKeyID = key.ID
}
if pending, ok := a.pending.LoadAndDelete(completion.RequestID); managedKeyID == "" && ok {
managedKeyID, _ = pending.(string)
if managedKeyID == "" {
managedKeyID = pendingManagedKeyID
}
if err := a.usage.ObserveRequest(context.Background(), collection.RequestRecord{
ManagedKeyID: managedKeyID, RequestID: completion.RequestID, TraceID: completion.TraceID,
+95
View File
@@ -186,6 +186,101 @@ func TestConcurrentUsage(t *testing.T) {
}
}
func TestSSELifecycleCleanupDoesNotAccumulatePendingState(t *testing.T) {
app := NewApp()
if _, err := app.HandleMethod(MethodPluginRegister, lifecycleRequest(t, SchemaVersion, testConfig(t, "enabled: true\ncodex_only: false\n"))); err != nil {
t.Fatal(err)
}
defer app.Shutdown()
if _, err := app.store.UpdateBilling(context.Background(), "key_default", managedaccess.BillingSettings{
QuotaMicros: 1_000_000_000, ResetPeriod: managedaccess.ResetNone, MaxConcurrency: 64,
}, time.Now().UTC()); err != nil {
t.Fatal(err)
}
scope := managedaccess.CallerScope("key_default")
outcomes := [...]RequestCompletionOutcome{
RequestCompletionSucceeded, RequestCompletionFailed, RequestCompletionCanceled, RequestCompletionRejected,
}
for index := range 128 {
requestID := fmt.Sprintf("sse-%d", index)
beforeRaw, _ := json.Marshal(RequestInterceptRequest{
RequestID: requestID, Model: "deepseek-v4-flash", RequestedModel: "deepseek-v4-flash",
Stream: true, Metadata: map[string]any{callerScopeMetadata: scope},
})
beforeResponse, err := app.HandleMethod(MethodRequestBefore, beforeRaw)
if err != nil {
t.Fatal(err)
}
var beforeEnvelope Envelope
_ = json.Unmarshal(beforeResponse, &beforeEnvelope)
var beforeResult RequestInterceptResponse
_ = json.Unmarshal(beforeEnvelope.Result, &beforeResult)
if beforeResult.Terminate {
t.Fatalf("SSE request %s was denied: %+v", requestID, beforeResult)
}
completionRaw, _ := json.Marshal(RequestCompletion{
RequestID: requestID, Model: "deepseek-v4-flash", Stream: true,
StartedAt: time.Now().Add(-time.Second), CompletedAt: time.Now(),
Outcome: outcomes[index%len(outcomes)], Metadata: map[string]any{callerScopeMetadata: scope},
})
if _, err := app.HandleMethod(MethodRequestComplete, completionRaw); err != nil {
t.Fatal(err)
}
}
pending := 0
app.pending.Range(func(_, _ any) bool { pending++; return true })
if pending != 0 {
t.Fatalf("pending SSE requests = %d", pending)
}
state, err := app.store.BillingState(context.Background(), "key_default", time.Now().UTC())
if err != nil || state.ActiveRequests != 0 {
t.Fatalf("SSE billing state=%+v err=%v", state, err)
}
}
func TestFailedSSECompletionStillClearsPendingState(t *testing.T) {
app := NewApp()
if _, err := app.HandleMethod(MethodPluginRegister, lifecycleRequest(t, SchemaVersion, testConfig(t, "enabled: true\ncodex_only: false\n"))); err != nil {
t.Fatal(err)
}
if _, err := app.store.UpdateBilling(context.Background(), "key_default", managedaccess.BillingSettings{
QuotaMicros: 1_000_000, ResetPeriod: managedaccess.ResetNone, MaxConcurrency: 4,
}, time.Now().UTC()); err != nil {
t.Fatal(err)
}
scope := managedaccess.CallerScope("key_default")
beforeRaw, _ := json.Marshal(RequestInterceptRequest{
RequestID: "sse-database-failure", Model: "deepseek-v4-flash", Stream: true,
Metadata: map[string]any{callerScopeMetadata: scope},
})
beforeResponse, err := app.HandleMethod(MethodRequestBefore, beforeRaw)
if err != nil {
t.Fatal(err)
}
var beforeEnvelope Envelope
_ = json.Unmarshal(beforeResponse, &beforeEnvelope)
var beforeResult RequestInterceptResponse
_ = json.Unmarshal(beforeEnvelope.Result, &beforeResult)
if beforeResult.Terminate {
t.Fatalf("SSE request was denied: %+v", beforeResult)
}
if err := app.store.Close(); err != nil {
t.Fatal(err)
}
completionRaw, _ := json.Marshal(RequestCompletion{
RequestID: "sse-database-failure", Stream: true, Outcome: RequestCompletionCanceled,
CompletedAt: time.Now(), Metadata: map[string]any{callerScopeMetadata: scope},
})
if _, err := app.HandleMethod(MethodRequestComplete, completionRaw); err == nil {
t.Fatal("completion unexpectedly succeeded against a closed database")
}
pending := 0
app.pending.Range(func(_, _ any) bool { pending++; return true })
if pending != 0 {
t.Fatalf("pending requests after failed completion = %d", pending)
}
}
func TestUnknownMethodReturnsErrorEnvelope(t *testing.T) {
raw, err := NewApp().HandleMethod("missing", nil)
if err != nil {
+29 -7
View File
@@ -63,16 +63,33 @@ func (a *App) putPrice(body []byte) ManagementResponse {
if err != nil {
return jsonManagementResponse(http.StatusBadRequest, map[string]any{"error": map[string]string{"code": "invalid_price", "message": err.Error()}})
}
a.priceMu.Lock()
defer a.priceMu.Unlock()
a.mu.Lock()
defer a.mu.Unlock()
if a.store == nil {
store := a.store
if store == nil {
a.mu.Unlock()
return jsonManagementResponse(http.StatusServiceUnavailable, map[string]any{"error": map[string]string{"code": "database_unavailable", "message": "价格数据库尚未初始化"}})
}
if err := a.store.UpsertPrice(context.Background(), policy); err != nil {
if err := store.UpsertPrice(context.Background(), policy); err != nil {
a.mu.Unlock()
return jsonManagementResponse(http.StatusInternalServerError, map[string]any{"error": map[string]string{"code": "database_error", "message": err.Error()}})
}
a.prices[normalizeModelName(policy.Model)] = policy
if _, err := a.store.BackfillMissingCosts(context.Background(), policy); err != nil {
a.mu.Unlock()
a.mu.RLock()
currentStore := a.store == store
a.mu.RUnlock()
if !currentStore {
return jsonManagementResponse(http.StatusServiceUnavailable, map[string]any{"error": map[string]string{"code": "database_reconfigured", "message": "价格已经保存,请在重新配置后重试历史补算"}})
}
if _, err := store.BackfillMissingCosts(context.Background(), policy); err != nil {
a.mu.RLock()
reconfigured := a.store != store
a.mu.RUnlock()
if reconfigured {
return jsonManagementResponse(http.StatusServiceUnavailable, map[string]any{"error": map[string]string{"code": "database_reconfigured", "message": "价格已经保存,重新配置中止了历史补算"}})
}
return jsonManagementResponse(http.StatusInternalServerError, map[string]any{"error": map[string]string{"code": "database_error", "message": err.Error()}})
}
return jsonManagementResponse(http.StatusOK, policyToDTO(policy))
@@ -87,14 +104,19 @@ func (a *App) deletePrice(body []byte) ManagementResponse {
if request.Model == "" {
return jsonManagementResponse(http.StatusBadRequest, map[string]any{"error": map[string]string{"code": "invalid_price", "message": "model 不能为空"}})
}
a.priceMu.Lock()
defer a.priceMu.Unlock()
a.mu.Lock()
defer a.mu.Unlock()
if a.store == nil {
store := a.store
if store == nil {
a.mu.Unlock()
return jsonManagementResponse(http.StatusServiceUnavailable, map[string]any{"error": map[string]string{"code": "database_unavailable", "message": "价格数据库尚未初始化"}})
}
if err := a.store.DeletePrice(context.Background(), request.Model); err != nil {
if err := store.DeletePrice(context.Background(), request.Model); err != nil {
a.mu.Unlock()
return jsonManagementResponse(http.StatusInternalServerError, map[string]any{"error": map[string]string{"code": "database_error", "message": err.Error()}})
}
defer a.mu.Unlock()
delete(a.prices, normalizeModelName(request.Model))
return jsonManagementResponse(http.StatusOK, map[string]any{"deleted": request.Model})
}
+8 -8
View File
@@ -16,7 +16,7 @@ var ErrManagedKeyNotFound = errors.New("managed key not found")
func (r *SQLiteUsageRepository) BootstrapManagedKey(ctx context.Context, name, secret string) (managedaccess.ManagedKey, error) {
var count int
if err := r.db.QueryRowContext(ctx, `SELECT COUNT(*) FROM managed_keys`).Scan(&count); err != nil {
if err := r.readDB.QueryRowContext(ctx, `SELECT COUNT(*) FROM managed_keys`).Scan(&count); err != nil {
return managedaccess.ManagedKey{}, fmt.Errorf("查询 Key 数量: %w", err)
}
if count > 0 {
@@ -176,7 +176,7 @@ func (r *SQLiteUsageRepository) ManagedKeyByReference(ctx context.Context, refer
func (r *SQLiteUsageRepository) scanManagedKey(ctx context.Context, where string, args ...any) (managedaccess.ManagedKey, error) {
var key managedaccess.ManagedKey
var createdAt, updatedAt string
err := r.db.QueryRowContext(ctx, `
err := r.readDB.QueryRowContext(ctx, `
SELECT id, name, secret, status, route_mode, upstream_account_id, all_models, created_at, updated_at
FROM managed_keys `+where, args...).Scan(&key.ID, &key.Name, &key.Secret, &key.Status, &key.RouteMode,
&key.UpstreamAccountID, &key.AllModels, &createdAt, &updatedAt)
@@ -197,7 +197,7 @@ func (r *SQLiteUsageRepository) ListManagedKeys(ctx context.Context, includeArch
if includeArchived {
where = ""
}
rows, err := r.db.QueryContext(ctx, `
rows, err := r.readDB.QueryContext(ctx, `
SELECT id, name, secret, status, route_mode, upstream_account_id, all_models, created_at, updated_at
FROM managed_keys `+where+` ORDER BY name COLLATE NOCASE`)
if err != nil {
@@ -232,7 +232,7 @@ FROM managed_keys `+where+` ORDER BY name COLLATE NOCASE`)
}
func (r *SQLiteUsageRepository) managedKeyModels(ctx context.Context, keyID string) ([]string, error) {
rows, err := r.db.QueryContext(ctx, `SELECT model FROM managed_key_models WHERE managed_key_id=? ORDER BY model`, keyID)
rows, err := r.readDB.QueryContext(ctx, `SELECT model FROM managed_key_models WHERE managed_key_id=? ORDER BY model`, keyID)
if err != nil {
return nil, fmt.Errorf("查询 Key 模型: %w", err)
}
@@ -288,7 +288,7 @@ ON CONFLICT(id) DO UPDATE SET cpa_auth_id=excluded.cpa_auth_id, provider=exclude
}
func (r *SQLiteUsageRepository) ListUpstreamAccounts(ctx context.Context) ([]managedaccess.UpstreamAccount, error) {
rows, err := r.db.QueryContext(ctx, `
rows, err := r.readDB.QueryContext(ctx, `
SELECT id, cpa_auth_id, cpa_auth_index, provider, display_name, status, status_message,
disabled, unavailable, priority, last_seen_at
FROM upstream_accounts ORDER BY provider, display_name COLLATE NOCASE`)
@@ -314,7 +314,7 @@ FROM upstream_accounts ORDER BY provider, display_name COLLATE NOCASE`)
func (r *SQLiteUsageRepository) UpstreamAccountByID(ctx context.Context, id string) (managedaccess.UpstreamAccount, error) {
var account managedaccess.UpstreamAccount
var lastSeen string
err := r.db.QueryRowContext(ctx, `
err := r.readDB.QueryRowContext(ctx, `
SELECT id, cpa_auth_id, cpa_auth_index, provider, display_name, status, status_message,
disabled, unavailable, priority, last_seen_at FROM upstream_accounts WHERE id=?`, id).Scan(
&account.ID, &account.CPAAuthID, &account.CPAAuthIndex, &account.Provider, &account.DisplayName,
@@ -332,7 +332,7 @@ SELECT id, cpa_auth_id, cpa_auth_index, provider, display_name, status, status_m
func (r *SQLiteUsageRepository) KeyStats(ctx context.Context, keyID string, today time.Time) (managedaccess.KeyStats, error) {
stats := managedaccess.KeyStats{KeyID: keyID}
var last sql.NullString
err := r.db.QueryRowContext(ctx, `
err := r.readDB.QueryRowContext(ctx, `
SELECT COUNT(*),
COALESCE(SUM(u.input_tokens),0), COALESCE(SUM(u.output_tokens),0),
COALESCE(SUM(u.total_tokens),0), COALESCE(SUM(u.cost_micros),0),
@@ -359,7 +359,7 @@ WHERE d.managed_key_id=?`, formatTime(today), formatTime(today), formatTime(toda
}
func (r *SQLiteUsageRepository) ModelSuggestions(ctx context.Context) ([]string, error) {
rows, err := r.db.QueryContext(ctx, `
rows, err := r.readDB.QueryContext(ctx, `
SELECT model FROM model_prices
UNION SELECT model FROM usage_records WHERE model <> ''
UNION SELECT requested_model FROM request_records WHERE requested_model <> ''
+1 -1
View File
@@ -237,7 +237,7 @@ FROM billing_ledger WHERE managed_key_id=?`
}
query += ` ORDER BY id DESC LIMIT ?`
args = append(args, limit)
rows, err := r.db.QueryContext(ctx, query, args...)
rows, err := r.readDB.QueryContext(ctx, query, args...)
if err != nil {
return nil, err
}
@@ -0,0 +1,145 @@
package repository
import (
"context"
"path/filepath"
"strings"
"testing"
"time"
"cpa-ext/internal/collection"
"cpa-ext/internal/pricing"
)
func TestSQLiteSeparatesReaderPoolFromWriter(t *testing.T) {
store, err := OpenSQLiteUsage(filepath.Join(t.TempDir(), "usage.db"))
if err != nil {
t.Fatal(err)
}
defer store.Close()
if store.readDB == nil || store.readDB == store.db {
t.Fatal("reader pool was not initialized separately")
}
if got := store.db.Stats().MaxOpenConnections; got != 1 {
t.Fatalf("writer connections = %d, want 1", got)
}
if got := store.readDB.Stats().MaxOpenConnections; got != 4 {
t.Fatalf("reader connections = %d, want 4", got)
}
}
func TestSQLiteWriterProgressesDuringReadTransaction(t *testing.T) {
store, err := OpenSQLiteUsage(filepath.Join(t.TempDir(), "usage.db"))
if err != nil {
t.Fatal(err)
}
defer store.Close()
readTx, err := store.readDB.BeginTx(context.Background(), nil)
if err != nil {
t.Fatal(err)
}
defer readTx.Rollback()
var count int
if err := readTx.QueryRow(`SELECT COUNT(*) FROM request_detail_index`).Scan(&count); err != nil {
t.Fatal(err)
}
done := make(chan error, 1)
go func() {
done <- store.Insert(context.Background(), collection.Record{
RequestedAt: time.Now().UTC(), Model: "model", TotalTokens: 1,
})
}()
select {
case err := <-done:
if err != nil {
t.Fatal(err)
}
case <-time.After(2 * time.Second):
t.Fatal("writer was blocked by the independent read transaction")
}
}
func TestBackfillMissingCostsUsesBoundedIndexedBatches(t *testing.T) {
store, err := OpenSQLiteUsage(filepath.Join(t.TempDir(), "usage.db"))
if err != nil {
t.Fatal(err)
}
defer store.Close()
_, err = store.db.Exec(`
WITH RECURSIVE sequence(value) AS (
SELECT 1 UNION ALL SELECT value+1 FROM sequence WHERE value<2505
)
INSERT INTO usage_records (requested_at, model, failed, input_tokens, output_tokens, total_tokens)
SELECT printf('2026-08-15T00:00:%02d.000000000Z', value % 60), 'MODEL-BATCH', 0, 1000, 100, 1100
FROM sequence`)
if err != nil {
t.Fatal(err)
}
policy := pricing.Policy{
Model: "model-batch",
Base: pricing.Rates{InputMicrosPer1M: 1_000_000, OutputMicrosPer1M: 2_000_000},
FastMultiplier: pricing.Ratio{Numerator: 5, Denominator: 2},
}
updated, err := store.BackfillMissingCosts(context.Background(), policy)
if err != nil {
t.Fatal(err)
}
if updated != 2505 {
t.Fatalf("updated = %d, want 2505", updated)
}
var missing int
if err := store.readDB.QueryRow(`SELECT COUNT(*) FROM usage_records WHERE cost_micros IS NULL`).Scan(&missing); err != nil {
t.Fatal(err)
}
if missing != 0 {
t.Fatalf("missing costs = %d, want 0", missing)
}
plan := explainPlan(t, store, `
SELECT id FROM usage_records
WHERE model = ? COLLATE NOCASE AND cost_micros IS NULL AND id > ? AND id <= ?
ORDER BY id LIMIT ?`, policy.Model, 0, 3000, costBackfillBatchSize)
if !strings.Contains(plan, "idx_usage_records_missing_cost_model") {
t.Fatalf("backfill query did not use missing-cost index: %s", plan)
}
}
func TestUsageDashboardUsesBoundedHistoryIndexes(t *testing.T) {
store, err := OpenSQLiteUsage(filepath.Join(t.TempDir(), "usage.db"))
if err != nil {
t.Fatal(err)
}
defer store.Close()
plan := explainPlan(t, store, usageDashboardUsersSQL, "2026-08-15T00:00:00Z")
if !strings.Contains(plan, "idx_request_detail_time") {
t.Fatalf("dashboard today query did not use time index: %s", plan)
}
if !strings.Contains(plan, "idx_request_detail_key") {
t.Fatalf("dashboard latest-use query did not use key index: %s", plan)
}
}
func explainPlan(t *testing.T, store *SQLiteUsageRepository, statement string, args ...any) string {
t.Helper()
rows, err := store.readDB.Query("EXPLAIN QUERY PLAN "+statement, args...)
if err != nil {
t.Fatal(err)
}
defer rows.Close()
var plan strings.Builder
for rows.Next() {
var id, parent, unused int
var detail string
if err := rows.Scan(&id, &parent, &unused, &detail); err != nil {
t.Fatal(err)
}
plan.WriteString(detail)
plan.WriteByte('\n')
}
if err := rows.Err(); err != nil {
t.Fatal(err)
}
return plan.String()
}
+79 -50
View File
@@ -10,7 +10,7 @@ import (
// ListPrices returns every configured exact-model policy.
func (r *SQLiteUsageRepository) ListPrices(ctx context.Context) ([]pricing.Policy, error) {
rows, err := r.db.QueryContext(ctx, `
rows, err := r.readDB.QueryContext(ctx, `
SELECT model, input_rate_micros, cache_read_rate_micros, cache_write_rate_micros, output_rate_micros,
long_context_enabled, long_context_threshold, long_context_comparison,
long_input_rate_micros, long_cache_read_rate_micros, long_cache_write_rate_micros, long_output_rate_micros,
@@ -96,18 +96,14 @@ ON CONFLICT(model) DO UPDATE SET
return nil
}
const costBackfillBatchSize = 1000
// BackfillMissingCosts 只补算尚未定价的历史记录,已经保存的账单金额不会随价格修改而变化。
// 每批记录独立提交,避免百万级历史记录一次性占满内存或长期占用 writer。
func (r *SQLiteUsageRepository) BackfillMissingCosts(ctx context.Context, policy pricing.Policy) (int64, error) {
if err := policy.Validate(); err != nil {
return 0, err
}
rows, err := r.db.QueryContext(ctx, `
SELECT id, input_tokens, cache_read_tokens, cache_write_tokens, output_tokens, service_tier, speed
FROM usage_records
WHERE model = ? AND cost_micros IS NULL`, policy.Model)
if err != nil {
return 0, fmt.Errorf("查询待补算用量: %w", err)
}
type pendingCost struct {
id int64
inputTokens int64
@@ -116,56 +112,89 @@ WHERE model = ? AND cost_micros IS NULL`, policy.Model)
outputTokens int64
serviceTier, speed string
}
var pending []pendingCost
for rows.Next() {
var item pendingCost
if err := rows.Scan(&item.id, &item.inputTokens, &item.cacheReadTokens, &item.cacheWriteTokens, &item.outputTokens, &item.serviceTier, &item.speed); err != nil {
_ = rows.Close()
return 0, fmt.Errorf("读取待补算用量: %w", err)
}
pending = append(pending, item)
var maxID int64
if err := r.readDB.QueryRowContext(ctx, `
SELECT COALESCE(MAX(id), 0) FROM usage_records
WHERE model = ? COLLATE NOCASE AND cost_micros IS NULL`, policy.Model).Scan(&maxID); err != nil {
return 0, fmt.Errorf("定位待补算用量范围: %w", err)
}
if err := rows.Err(); err != nil {
_ = rows.Close()
return 0, fmt.Errorf("遍历待补算用量: %w", err)
}
if err := rows.Close(); err != nil {
return 0, fmt.Errorf("关闭待补算查询: %w", err)
}
tx, err := r.db.BeginTx(ctx, nil)
if err != nil {
return 0, fmt.Errorf("开始补算事务: %w", err)
}
defer func() { _ = tx.Rollback() }()
var updated int64
for _, item := range pending {
result, calculateErr := pricing.Calculate(policy, pricing.Usage{
InputTokens: item.inputTokens, CacheReadTokens: item.cacheReadTokens,
CacheWriteTokens: item.cacheWriteTokens, OutputTokens: item.outputTokens,
ServiceTier: item.serviceTier, Speed: item.speed,
})
if calculateErr != nil {
continue
var lastID int64
for lastID < maxID {
rows, err := r.readDB.QueryContext(ctx, `
SELECT id, input_tokens, cache_read_tokens, cache_write_tokens, output_tokens, service_tier, speed
FROM usage_records
WHERE model = ? COLLATE NOCASE AND cost_micros IS NULL AND id > ? AND id <= ?
ORDER BY id LIMIT ?`, policy.Model, lastID, maxID, costBackfillBatchSize)
if err != nil {
return updated, fmt.Errorf("查询待补算用量: %w", err)
}
change, updateErr := tx.ExecContext(ctx, `
pending := make([]pendingCost, 0, costBackfillBatchSize)
for rows.Next() {
var item pendingCost
if err := rows.Scan(&item.id, &item.inputTokens, &item.cacheReadTokens, &item.cacheWriteTokens, &item.outputTokens, &item.serviceTier, &item.speed); err != nil {
_ = rows.Close()
return updated, fmt.Errorf("读取待补算用量: %w", err)
}
pending = append(pending, item)
}
if err := rows.Err(); err != nil {
_ = rows.Close()
return updated, fmt.Errorf("遍历待补算用量: %w", err)
}
if err := rows.Close(); err != nil {
return updated, fmt.Errorf("关闭待补算查询: %w", err)
}
if len(pending) == 0 {
break
}
lastID = pending[len(pending)-1].id
tx, err := r.db.BeginTx(ctx, nil)
if err != nil {
return updated, fmt.Errorf("开始补算事务: %w", err)
}
statement, err := tx.PrepareContext(ctx, `
UPDATE usage_records
SET cost_micros = ?, price_tier = ?, fast_requested = ?, fast_pricing_applied = ?,
price_multiplier_numerator = ?, price_multiplier_denominator = ?
WHERE id = ? AND cost_micros IS NULL`,
result.CostMicros, result.PriceTier, result.FastRequested, result.FastApplied,
result.MultiplierNumerator, result.MultiplierDenominator, item.id)
if updateErr != nil {
return 0, fmt.Errorf("补算用量价格: %w", updateErr)
WHERE id = ? AND cost_micros IS NULL`)
if err != nil {
_ = tx.Rollback()
return updated, fmt.Errorf("准备补算用量价格: %w", err)
}
count, countErr := change.RowsAffected()
if countErr != nil {
return 0, fmt.Errorf("读取补算数量: %w", countErr)
for _, item := range pending {
result, calculateErr := pricing.Calculate(policy, pricing.Usage{
InputTokens: item.inputTokens, CacheReadTokens: item.cacheReadTokens,
CacheWriteTokens: item.cacheWriteTokens, OutputTokens: item.outputTokens,
ServiceTier: item.serviceTier, Speed: item.speed,
})
if calculateErr != nil {
continue
}
change, updateErr := statement.ExecContext(ctx,
result.CostMicros, result.PriceTier, result.FastRequested, result.FastApplied,
result.MultiplierNumerator, result.MultiplierDenominator, item.id)
if updateErr != nil {
_ = statement.Close()
_ = tx.Rollback()
return updated, fmt.Errorf("补算用量价格: %w", updateErr)
}
count, countErr := change.RowsAffected()
if countErr != nil {
_ = statement.Close()
_ = tx.Rollback()
return updated, fmt.Errorf("读取补算数量: %w", countErr)
}
updated += count
}
if err := statement.Close(); err != nil {
_ = tx.Rollback()
return updated, fmt.Errorf("关闭补算语句: %w", err)
}
if err := tx.Commit(); err != nil {
return updated, fmt.Errorf("提交补算事务: %w", err)
}
updated += count
}
if err := tx.Commit(); err != nil {
return 0, fmt.Errorf("提交补算事务: %w", err)
}
return updated, nil
}
+2 -2
View File
@@ -143,7 +143,7 @@ func (r *SQLiteUsageRepository) QueryUsage(ctx context.Context, input collection
}
where, args := usageWhere(query)
var total int64
if err := r.db.QueryRowContext(ctx, `SELECT COUNT(*) FROM request_detail_index d WHERE `+where, args...).Scan(&total); err != nil {
if err := r.readDB.QueryRowContext(ctx, `SELECT COUNT(*) FROM request_detail_index d WHERE `+where, args...).Scan(&total); err != nil {
return collection.UsagePage{}, fmt.Errorf("统计请求明细: %w", err)
}
totalPages := 0
@@ -179,7 +179,7 @@ func (r *SQLiteUsageRepository) QueryUsage(ctx context.Context, input collection
} else {
statement += " LIMIT ? OFFSET ?"
}
rows, err := r.db.QueryContext(ctx, statement, limitArgs...)
rows, err := r.readDB.QueryContext(ctx, statement, limitArgs...)
if err != nil {
return collection.UsagePage{}, fmt.Errorf("查询请求明细: %w", err)
}
+38 -18
View File
@@ -2,57 +2,77 @@ package repository
import (
"context"
"database/sql"
"fmt"
"time"
"cpa-ext/internal/collection"
)
const usageDashboardUsersSQL = `
WITH today AS (
SELECT d.managed_key_id, COUNT(*) AS requests,
COALESCE(SUM(u.input_tokens),0) AS input_tokens,
COALESCE(SUM(u.output_tokens),0) AS output_tokens,
COALESCE(SUM(u.total_tokens),0) AS total_tokens,
COALESCE(SUM(u.cost_micros),0) AS cost_micros
FROM request_detail_index d INDEXED BY idx_request_detail_time
LEFT JOIN usage_records u ON u.id=d.usage_id
WHERE d.requested_at>=?
GROUP BY d.managed_key_id
), dashboard_keys AS (
SELECT id, name FROM managed_keys
UNION ALL
SELECT t.managed_key_id, '未识别'
FROM today t LEFT JOIN managed_keys k ON k.id=t.managed_key_id
WHERE k.id IS NULL
)
SELECT k.id, k.name,
COALESCE(t.requests,0), COALESCE(t.input_tokens,0), COALESCE(t.output_tokens,0),
COALESCE(t.total_tokens,0), COALESCE(t.cost_micros,0),
(SELECT recent.requested_at FROM request_detail_index recent
WHERE recent.managed_key_id=k.id
ORDER BY recent.requested_at DESC, recent.id DESC LIMIT 1)
FROM dashboard_keys k
LEFT JOIN today t ON t.managed_key_id=k.id
ORDER BY COALESCE(t.total_tokens,0) DESC, k.name COLLATE NOCASE`
func (r *SQLiteUsageRepository) UsageDashboard(ctx context.Context, today time.Time, days int) (collection.UsageDashboard, error) {
if days < 1 {
days = 7
}
start := time.Date(today.Year(), today.Month(), today.Day(), 0, 0, 0, 0, today.Location())
dashboard := collection.UsageDashboard{}
if err := r.db.QueryRowContext(ctx, `
if err := r.readDB.QueryRowContext(ctx, `
SELECT COUNT(*), COALESCE(SUM(u.input_tokens),0), COALESCE(SUM(u.output_tokens),0),
COALESCE(SUM(u.total_tokens),0), COALESCE(SUM(u.cost_micros),0)
FROM request_detail_index d LEFT JOIN usage_records u ON u.id=d.usage_id
WHERE d.requested_at>=?`, formatTime(start)).Scan(&dashboard.Today.Requests, &dashboard.Today.InputTokens, &dashboard.Today.OutputTokens, &dashboard.Today.TotalTokens, &dashboard.Today.CostMicros); err != nil {
return dashboard, fmt.Errorf("汇总今日用量: %w", err)
}
rows, err := r.db.QueryContext(ctx, `
SELECT d.managed_key_id, COALESCE(k.name, '未识别'),
SUM(CASE WHEN d.requested_at>=? THEN 1 ELSE 0 END),
COALESCE(SUM(CASE WHEN d.requested_at>=? THEN u.input_tokens ELSE 0 END),0),
COALESCE(SUM(CASE WHEN d.requested_at>=? THEN u.output_tokens ELSE 0 END),0),
COALESCE(SUM(CASE WHEN d.requested_at>=? THEN u.total_tokens ELSE 0 END),0),
COALESCE(SUM(CASE WHEN d.requested_at>=? THEN u.cost_micros ELSE 0 END),0),
MAX(d.requested_at)
FROM request_detail_index d
LEFT JOIN usage_records u ON u.id=d.usage_id
LEFT JOIN managed_keys k ON k.id=d.managed_key_id
GROUP BY d.managed_key_id, COALESCE(k.name, '未识别')
ORDER BY 6 DESC, 2`, formatTime(start), formatTime(start), formatTime(start), formatTime(start), formatTime(start))
rows, err := r.readDB.QueryContext(ctx, usageDashboardUsersSQL, formatTime(start))
if err != nil {
return dashboard, fmt.Errorf("汇总用户用量: %w", err)
}
for rows.Next() {
var item collection.UserUsageSummary
var last string
var last sql.NullString
if err := rows.Scan(&item.KeyID, &item.KeyAlias, &item.Today.Requests, &item.Today.InputTokens, &item.Today.OutputTokens, &item.Today.TotalTokens, &item.Today.CostMicros, &last); err != nil {
_ = rows.Close()
return dashboard, fmt.Errorf("读取用户用量汇总: %w", err)
}
if parsed, parseErr := time.Parse(time.RFC3339Nano, last); parseErr == nil {
item.LastUsedAt = &parsed
if last.Valid {
parsed, parseErr := time.Parse(time.RFC3339Nano, last.String)
if parseErr == nil {
item.LastUsedAt = &parsed
}
}
dashboard.Users = append(dashboard.Users, item)
}
_ = rows.Close()
dailyStart := start.AddDate(0, 0, -(days - 1))
rows, err = r.db.QueryContext(ctx, `
rows, err = r.readDB.QueryContext(ctx, `
SELECT DATE(d.requested_at, '+8 hours'), COALESCE(SUM(u.total_tokens),0)
FROM request_detail_index d LEFT JOIN usage_records u ON u.id=d.usage_id
WHERE d.requested_at>=? GROUP BY DATE(d.requested_at, '+8 hours') ORDER BY 1`, formatTime(dailyStart))
+24 -4
View File
@@ -59,6 +59,8 @@ CREATE TABLE IF NOT EXISTS usage_records (
CREATE INDEX IF NOT EXISTS idx_usage_records_requested_at ON usage_records(requested_at DESC, id DESC);
CREATE UNIQUE INDEX IF NOT EXISTS idx_usage_records_execution_id ON usage_records(execution_id) WHERE execution_id <> '';
CREATE INDEX IF NOT EXISTS idx_usage_records_model ON usage_records(model COLLATE NOCASE) WHERE model <> '';
CREATE INDEX IF NOT EXISTS idx_usage_records_missing_cost_model
ON usage_records(model COLLATE NOCASE, id) WHERE cost_micros IS NULL;
CREATE TABLE IF NOT EXISTS request_records (
request_id TEXT PRIMARY KEY,
@@ -234,9 +236,10 @@ const (
orphanLifecycleTieWindow = time.Millisecond
)
// SQLiteUsageRepository 使用单写连接保存全部用量记录
// SQLiteUsageRepository 使用单写连接保存事实,并用独立连接池承载只读查询
type SQLiteUsageRepository struct {
db *sql.DB
db *sql.DB
readDB *sql.DB
}
func OpenSQLiteUsage(databasePath string) (*SQLiteUsageRepository, error) {
@@ -301,6 +304,19 @@ SELECT k.id, 1, k.created_at, 0, 0 FROM managed_keys k`); err != nil {
_ = db.Close()
return nil, err
}
readDB, err := sql.Open("sqlite3", dsn)
if err != nil {
_ = db.Close()
return nil, fmt.Errorf("打开 SQLite 只读连接池: %w", err)
}
readDB.SetMaxOpenConns(4)
readDB.SetMaxIdleConns(4)
if err := readDB.Ping(); err != nil {
_ = readDB.Close()
_ = db.Close()
return nil, fmt.Errorf("连接 SQLite 只读连接池: %w", err)
}
repository.readDB = readDB
return repository, nil
}
@@ -516,7 +532,7 @@ func (r *SQLiteUsageRepository) ListRecent(ctx context.Context, limit int) ([]co
if scanLimit < limit {
scanLimit = limit
}
rows, err := r.db.QueryContext(ctx, `
rows, err := r.readDB.QueryContext(ctx, `
WITH combined AS (
SELECT
u.managed_key_id, u.request_id, u.execution_id, COALESCE(NULLIF(u.trace_id, ''), r.trace_id, '') AS trace_id,
@@ -738,5 +754,9 @@ func mergeLifecycleIntoUsage(usage, lifecycle collection.Record) collection.Reco
}
func (r *SQLiteUsageRepository) Close() error {
return r.db.Close()
var readErr error
if r.readDB != nil {
readErr = r.readDB.Close()
}
return errors.Join(readErr, r.db.Close())
}
+4 -4
View File
@@ -2,15 +2,15 @@
set -euo pipefail
command -v go >/dev/null 2>&1 || {
echo '缺少 Go,请先在 WSL 中安装 Go 1.24+ 并加入 PATH。' >&2
exit 1
echo '缺少 Go,请先在 WSL 中安装支持 toolchain 自动切换的 Go 1.24+ 并加入 PATH。' >&2
exit 1
}
command -v gcc >/dev/null 2>&1 || {
echo '缺少 gcc,请安装 build-essentialCGO 构建动态库需要)。' >&2
exit 1
}
echo "Go: $(go version)"
echo "Go bootstrap: $(go version)"
echo "Go effective: $(go env GOVERSION)"
echo "GCC: $(gcc --version | head -n 1)"
echo 'CLIProxyAPI target: ABI 1 / RPC schema 3'