From ddebe3c4a2bebfce8f24b0499b4b4642fb0e836c Mon Sep 17 00:00:00 2001 From: chuan Date: Sat, 15 Aug 2026 18:09:54 +0800 Subject: [PATCH] =?UTF-8?q?fix:=20=E4=BF=AE=E5=A4=8D=E5=A4=A7=E8=A7=84?= =?UTF-8?q?=E6=A8=A1=E6=95=B0=E6=8D=AE=E6=80=A7=E8=83=BD=E4=B8=8E=E6=B5=81?= =?UTF-8?q?=E5=BC=8F=E7=8A=B6=E6=80=81=E6=B8=85=E7=90=86?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- README.md | 2 +- docs/modules/billing-and-pricing.md | 2 + docs/modules/usage-and-statistics.md | 2 + go.mod | 2 + internal/plugin/app.go | 9 +- internal/plugin/app_test.go | 95 ++++++++++++ internal/plugin/pricing_management.go | 36 ++++- internal/repository/sqlite_access.go | 16 +- internal/repository/sqlite_billing.go | 2 +- .../repository/sqlite_performance_test.go | 145 ++++++++++++++++++ internal/repository/sqlite_pricing.go | 129 ++++++++++------ internal/repository/sqlite_query.go | 4 +- internal/repository/sqlite_summary.go | 56 ++++--- internal/repository/sqlite_usage.go | 28 +++- scripts/check-env.sh | 8 +- 15 files changed, 439 insertions(+), 97 deletions(-) create mode 100644 internal/repository/sqlite_performance_test.go diff --git a/README.md b/README.md index 3f334f2..3ee288e 100644 --- a/README.md +++ b/README.md @@ -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 diff --git a/docs/modules/billing-and-pricing.md b/docs/modules/billing-and-pricing.md index c8d2da2..a45a4db 100644 --- a/docs/modules/billing-and-pricing.md +++ b/docs/modules/billing-and-pricing.md @@ -124,6 +124,8 @@ 价格不存在时返回 HTTP 503 和 `billing_price_unavailable`，不会把未知成本的请求发送到上游。历史迁移数据仍可能因为当时没有价格而显示为成本不可用。 +管理员首次补充某个模型价格时，系统按固定小批次补算该模型尚未定价的历史 Usage。补算使用缺失成本专用索引，不会一次性把全部历史载入内存；每批独立提交，在线认证、额度检查和新用量写入可以继续执行。 + ## 不可变账目 | 账目类型 | 含义 | 金额方向 | diff --git a/docs/modules/usage-and-statistics.md b/docs/modules/usage-and-statistics.md index 87d4530..9d9b57a 100644 --- a/docs/modules/usage-and-statistics.md +++ b/docs/modules/usage-and-statistics.md @@ -111,6 +111,8 @@ CLIProxyAPI 会通过两个独立回调提供请求信息： 查询投影只保存明细检索所需的关联和排序字段，Token、成本、生命周期和账目仍以原始事实表为准。当前分页与组合筛选已按百万级记录场景设计，不要求将全部记录读入内存。 +SQLite 使用单 writer 和独立读连接池：计费与事实写入保持顺序一致，管理台分页和汇总不会占用 writer。今日用户汇总只扫描当天范围，最近使用时间通过 Key 时间索引定位，不随全部历史记录线性分组扫描。 + ## 数据一致性 - Usage 插入使用内容事件标识，完全相同的重复回调不会重复记录或扣费。 diff --git a/go.mod b/go.mod index 935e9fd..3715769 100644 --- a/go.mod +++ b/go.mod @@ -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 diff --git a/internal/plugin/app.go b/internal/plugin/app.go index fc791fa..d000581 100644 --- a/internal/plugin/app.go +++ b/internal/plugin/app.go @@ -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, diff --git a/internal/plugin/app_test.go b/internal/plugin/app_test.go index 13b658a..624ede2 100644 --- a/internal/plugin/app_test.go +++ b/internal/plugin/app_test.go @@ -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 { diff --git a/internal/plugin/pricing_management.go b/internal/plugin/pricing_management.go index 7a03b8b..686deb8 100644 --- a/internal/plugin/pricing_management.go +++ b/internal/plugin/pricing_management.go @@ -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}) } diff --git a/internal/repository/sqlite_access.go b/internal/repository/sqlite_access.go index b1c102b..4eee825 100644 --- a/internal/repository/sqlite_access.go +++ b/internal/repository/sqlite_access.go @@ -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 <> '' diff --git a/internal/repository/sqlite_billing.go b/internal/repository/sqlite_billing.go index b6b586a..66659da 100644 --- a/internal/repository/sqlite_billing.go +++ b/internal/repository/sqlite_billing.go @@ -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 } diff --git a/internal/repository/sqlite_performance_test.go b/internal/repository/sqlite_performance_test.go new file mode 100644 index 0000000..25c7b99 --- /dev/null +++ b/internal/repository/sqlite_performance_test.go @@ -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() +} diff --git a/internal/repository/sqlite_pricing.go b/internal/repository/sqlite_pricing.go index 4588da9..e73086d 100644 --- a/internal/repository/sqlite_pricing.go +++ b/internal/repository/sqlite_pricing.go @@ -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 } diff --git a/internal/repository/sqlite_query.go b/internal/repository/sqlite_query.go index 695197c..9eaa41e 100644 --- a/internal/repository/sqlite_query.go +++ b/internal/repository/sqlite_query.go @@ -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) } diff --git a/internal/repository/sqlite_summary.go b/internal/repository/sqlite_summary.go index 8a8d070..7047437 100644 --- a/internal/repository/sqlite_summary.go +++ b/internal/repository/sqlite_summary.go @@ -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)) diff --git a/internal/repository/sqlite_usage.go b/internal/repository/sqlite_usage.go index 99a61eb..9e6d3b4 100644 --- a/internal/repository/sqlite_usage.go +++ b/internal/repository/sqlite_usage.go @@ -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()) } diff --git a/scripts/check-env.sh b/scripts/check-env.sh index 2f593b5..f0abccf 100644 --- a/scripts/check-env.sh +++ b/scripts/check-env.sh @@ -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-essential（CGO 构建动态库需要）。' >&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' -