Cyrene_v2/internal/runtime/model/usage_store.go

package model

import (
	"context"
	"database/sql"
	"fmt"
	"log/slog"
	"sync"
	"time"

	"git.yeij.top/AskaEth/Cyrene_v2/internal/logger"
)

// 本文件是**用量落库的写入侧**（切片 U1，主人裁决「2.4 落」）：一个有界队列 +
// 单写者 + 批量事务的异步写入器。
//
// 为什么不是「在 logUsage 里直接 INSERT」：那条语句跑在**每一次模型调用的返回
// 路径**上，库慢一点、锁一下，主人的每次对话就跟着慢一点 —— 记成本的开销绝不
// 能加到被记成本的那次调用上。这里因此把两者彻底解耦：
//
//  1. `RecordUsage` 只往**有界** channel 里塞一个样本，塞不进就**丢弃并限流
//     WARN**（绝不阻塞、绝不 panic、绝不无界增长）；
//  2. 一个常驻单写者把队列攒成批次，在一个事务里落库；
//  3. 整批失败只记一条带来源的 WARN（**不静默**），调用路径毫无感知。
//
// 写入样本里没有凭据、没有正文（见 UsageSample），这条红线在类型层面就成立。

// usageSource 是本文件每条日志的来源标签（跨模块可追溯定位，项目规范 §10）。
const usageSource = "model/usage_store"

const (
	// usageQueueCapacity 是落库队列的容量。它是**有界**的：满了就丢弃，
	// 而不是把模型调用拖住。
	usageQueueCapacity = 1024
	// usageBatchMax 是单次批量事务最多写入的行数。
	usageBatchMax = 128
	// usageFlushInterval 是队列非空时的最长等待：低频调用不会因为没攒满一批
	// 就一直不落库（面板看得到的是近 24h，不该差这几十秒）。
	usageFlushInterval = 2 * time.Second
	// usageDropWarnInterval 是「队列满、已丢弃样本」告警的限流间隔：真实部署里
	// 库挂了会让每一次调用都丢弃，不限流就会把日志刷爆。
	usageDropWarnInterval = 30 * time.Second
	// usageFinalFlushTimeout 是停止时把队列剩余样本写完的期限。
	usageFinalFlushTimeout = 5 * time.Second
)

// insertUsageSQL 写入一条用量样本。**session_id 当前恒为 NULL**（模型适配器跑在
// loop 之下、没有会话概念；见迁移 0017 的说明），列留着是为了接线时只改这一处。
const insertUsageSQL = `INSERT INTO llm_usage
    (session_id, purpose, adapter, model, prompt_tokens, completion_tokens,
     cache_hit_tokens, cache_miss_tokens, reasoning_tokens, total_tokens, took_ms)
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`

// UsageWriter 把用量样本异步落进 `llm_usage`。它实现 UsageSink，因此在模型适配器
// 里就是一个「可选出口」：注入即有、不注入则只记日志。
//
// 构造后必须 Start 才会真的落库；未 Start 时队列照样收、塞满后照样丢弃并告警
// （不静默失败）。Stop 必须在关库之前调用。
type UsageWriter struct {
	db  *sql.DB
	log *slog.Logger

	queue chan UsageSample
	done  chan struct{}

	mu       sync.Mutex
	cancel   context.CancelFunc
	started  bool
	dropped  int64
	lastWarn time.Time
}

// NewUsageWriter 构造写入器。db 为 nil（未接库）时返回 nil —— 调用方据此把
// 「无 sink」原样传给适配器，行为等于这一版之前。
func NewUsageWriter(db *sql.DB, log *slog.Logger) *UsageWriter {
	if db == nil {
		return nil
	}
	if log == nil {
		log = discardLogger(usageSource)
	}
	return &UsageWriter{
		db:    db,
		log:   logger.WithSource(log, usageSource),
		queue: make(chan UsageSample, usageQueueCapacity),
		done:  make(chan struct{}),
	}
}

// RecordUsage 实现 UsageSink：把样本塞进有界队列，塞不进就丢弃并（限流）告警。
//
// 它**绝不阻塞**：整个函数是一次非阻塞 channel 发送加一次可能的加锁自增，
// 因此模型调用路径的额外开销与库的健康状况无关。
func (w *UsageWriter) RecordUsage(sample UsageSample) {
	if w == nil {
		return
	}
	select {
	case w.queue <- sample:
	default:
		w.noteDrop()
	}
}

// Start 拉起单写者循环。重复调用是空操作；ctx 取消（或 Stop）后循环退出。
func (w *UsageWriter) Start(ctx context.Context) {
	if w == nil {
		return
	}
	w.mu.Lock()
	defer w.mu.Unlock()
	if w.started {
		return
	}
	w.started = true
	runCtx, cancel := context.WithCancel(ctx)
	w.cancel = cancel
	go w.loop(runCtx)
}

// Stop 停掉循环并等它退出，退出前尽力把队列里剩下的样本写完。未启动时是空操作。
// 它必须在关库**之前**调用（收尾写入要用那条连接池）。
func (w *UsageWriter) Stop() {
	if w == nil {
		return
	}
	w.mu.Lock()
	cancel := w.cancel
	started := w.started
	w.mu.Unlock()
	if !started || cancel == nil {
		return
	}
	cancel()
	<-w.done
}

// loop 是单写者：攒批 → 批量事务落库，直到 ctx 取消。
func (w *UsageWriter) loop(ctx context.Context) {
	defer close(w.done)
	ticker := time.NewTicker(usageFlushInterval)
	defer ticker.Stop()
	batch := make([]UsageSample, 0, usageBatchMax)
	flush := func() {
		if len(batch) == 0 {
			return
		}
		w.writeBatch(ctx, batch)
		batch = batch[:0]
	}
	for {
		select {
		case <-ctx.Done():
			// 收尾用**脱离取消**的上下文：ctx 已经取消，沿用它会让最后一批
			// 必然写失败，等于丢掉已经收下的样本。
			w.flushFinal(drainQueue(w.queue, batch))
			return
		case sample := <-w.queue:
			batch = append(batch, sample)
			if len(batch) >= usageBatchMax {
				flush()
			}
		case <-ticker.C:
			flush()
		}
	}
}

// flushFinal 在停机路径上把剩余样本写完（有期限，绝不无限等库）。
func (w *UsageWriter) flushFinal(remaining []UsageSample) {
	if len(remaining) == 0 {
		return
	}
	ctx, cancel := context.WithTimeout(context.Background(), usageFinalFlushTimeout)
	defer cancel()
	w.writeBatch(ctx, remaining)
}

// writeBatch 在一个事务里落一批样本。失败只记 WARN（带批次大小与原因），
// 绝不上抛 —— 调用路径早已返回，没有任何人在等这个结果。
func (w *UsageWriter) writeBatch(ctx context.Context, batch []UsageSample) {
	if err := insertUsageBatch(ctx, w.db, batch); err != nil {
		w.log.Warn("llm usage batch insert failed; batch dropped", "count", len(batch), "err", err)
	}
}

// noteDrop 记一次「队列满、样本被丢弃」，并按 usageDropWarnInterval 限流告警。
// 丢弃计数是累计值，因此告警里能看到这一段到底丢了多少条。
func (w *UsageWriter) noteDrop() {
	w.mu.Lock()
	w.dropped++
	dropped := w.dropped
	now := time.Now()
	shouldWarn := w.lastWarn.IsZero() || now.Sub(w.lastWarn) >= usageDropWarnInterval
	if shouldWarn {
		w.lastWarn = now
	}
	w.mu.Unlock()
	if shouldWarn {
		w.log.Warn("llm usage queue full; samples dropped (never blocks the model call)",
			"dropped", dropped, "capacity", usageQueueCapacity)
	}
}

// drainQueue 把队列里当前排着的样本全部取出（不阻塞），追加到 dst 之后返回。
func drainQueue(queue chan UsageSample, dst []UsageSample) []UsageSample {
	for {
		select {
		case sample := <-queue:
			dst = append(dst, sample)
		default:
			return dst
		}
	}
}

// insertUsageBatch 在一个事务里写入一批样本：单条 prepared statement 循环执行，
// 一次提交。用 prepared statement 而不是多值 INSERT，是为了避开各家 SQLite 对
// 绑定参数个数的不同上限（11 列 × 128 行）。
func insertUsageBatch(ctx context.Context, db *sql.DB, batch []UsageSample) error {
	if db == nil || len(batch) == 0 {
		return nil
	}
	tx, err := db.BeginTx(ctx, nil)
	if err != nil {
		return fmt.Errorf("begin usage batch: %w", err)
	}
	defer tx.Rollback()

	stmt, err := tx.PrepareContext(ctx, insertUsageSQL)
	if err != nil {
		return fmt.Errorf("prepare usage insert: %w", err)
	}
	defer stmt.Close()

	for i, sample := range batch {
		// session_id 传 NULL：见 insertUsageSQL 的说明。
		if _, err := stmt.ExecContext(ctx,
			nil,
			sample.Purpose,
			sample.Adapter,
			sample.Model,
			sample.PromptTokens,
			sample.CompletionTokens,
			sample.CacheHitTokens,
			sample.CacheMissTokens,
			sample.ReasoningTokens,
			sample.TotalTokens,
			sample.TookMs,
		); err != nil {
			return fmt.Errorf("insert usage sample %d: %w", i, err)
		}
	}
	if err := tx.Commit(); err != nil {
		return fmt.Errorf("commit usage batch: %w", err)
	}
	return nil
}