feat: 十轮网关优化 - 安全加固/可观测性/性能/可靠性

- SSE Keepalive Ping (15s心跳防止代理断连) - Timing HTTP 头 (X-Timing-Queue/Inference/Total-Ms) - Adapter Request-ID 传播到后端 - Session 清理日志回调 - Server 安全加固 (ReadHeaderTimeout/MaxHeaderBytes 防 slowloris) - Usage Tracker 数据保留清理 (retentionDays + 定期清理) - Config Reload 后 Adapter Registry 更新 (RegisterIfAbsent + RWMutex) - Rate Limiter 空闲 Bucket 清理 (30分钟过期) - Shutdown Drain 超时可配置 (ShutdownDrainSeconds) - Config 模型字段校验增强 (provider/endpoint/actual_model) - Auth 过期 Key 自动清理 (5分钟扫描) - Admin API Rate Limiting - Adapter Health Check 独立超时 (每个 adapter 3s) - TCP 连接阶段超时 (DialContext 5s + KeepAlive 30s) - 幂等键缓存、审计日志、Gzip 中间件、CORS Expose Headers - Backpressure 响应头、熔断器 Prometheus 指标 - 连接池优化、Trace-ID 全链路传播
2026-08-03 15:43:11 +08:00
parent e7e98271d4
commit da9c8334d8
27 changed files with 3002 additions and 309 deletions
@@ -0,0 +1,293 @@
+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
+	}
+}