282 lines
8.2 KiB
Go
282 lines
8.2 KiB
Go
package storage
|
||
|
||
import (
|
||
"log"
|
||
"os"
|
||
"sort"
|
||
"strconv"
|
||
"sync"
|
||
"sync/atomic"
|
||
"time"
|
||
)
|
||
|
||
// ─── 召回元数据「累积延迟更新」缓冲（2026-09-12 优化B）─────────────────────────
|
||
//
|
||
// 问题（根因）：recall 读路径原先对**每条结果**同步调用 lancedb.Update
|
||
// （recall_count / last_recalled_at / freshness）。LanceDB 是 MVCC 存储——
|
||
// **每次 update 提交产生一个版本**。实测生产：15 版本/分钟、_versions 目录
|
||
// 膨胀到 17.72G（真实数据仅 22M，75,327 个 .manifest，放大 800 倍）。
|
||
//
|
||
// 方案：
|
||
//
|
||
// Record() 读路径只把增量写进内存缓冲；**同一记忆在一个窗口内多次命中合并成 1 条 delta**
|
||
// Flush() 后台按窗口（默认 300s）+ 阈值（默认 256 条不同记忆）批量提交一次；
|
||
// Rust 侧 lancedb_update_batch 把整批**按 delta 分组**，每组一次
|
||
// update(`recall_count + delta` + id IN (...)) 提交（lance 不支持 CASE WHEN）
|
||
// → 每批版本数 = 不同 delta 的个数（delta 绝大多数为 1，故通常 1 个版本；
|
||
// 旧实现：每行 1 个版本）
|
||
//
|
||
// 语义取舍（明确记录，便于日后审计）：
|
||
// - last_recalled_at / freshness 最多延迟一个窗口（分钟级）。遗忘/衰减判定以「天」为单位，无影响。
|
||
// - 进程崩溃会丢最后一个窗口的增量（召回统计，不是记忆数据本身），可接受。
|
||
// - recall_count 由 Rust 侧「读库现值 + delta」计算（不是读 Go 缓存），顺带修掉旧实现里
|
||
// 「本地缓存 +1 后再 +1」导致的计数漂移。
|
||
type RecallWriteItem struct {
|
||
ID string `json:"id"`
|
||
Delta int `json:"delta"`
|
||
}
|
||
|
||
// BatchRecallUpdater 单事务批量更新接口。
|
||
// 由 Rust IPC 后端（RustLanceDBClient）实现；其他后端不支持时自动退化为逐条 Update。
|
||
type BatchRecallUpdater interface {
|
||
UpdateRecallBatch(table string, items []RecallWriteItem, lastRecalledAt string) (int64, error)
|
||
}
|
||
|
||
const (
|
||
defaultRecallFlushInterval = 300 * time.Second
|
||
defaultRecallFlushMaxItems = 256
|
||
)
|
||
|
||
type pendingRecall struct {
|
||
delta int
|
||
lastAt time.Time
|
||
}
|
||
|
||
// RecallWriteBuffer 累积召回元数据写，延迟批量落盘。
|
||
type RecallWriteBuffer struct {
|
||
mu sync.Mutex
|
||
pending map[string]*pendingRecall
|
||
ldb LanceDB
|
||
batch BatchRecallUpdater
|
||
interval time.Duration
|
||
maxItems int
|
||
|
||
stopCh chan struct{}
|
||
doneCh chan struct{}
|
||
stopped bool
|
||
startOnce sync.Once
|
||
|
||
flushCount int64
|
||
flushItems int64
|
||
flushRows int64
|
||
flushErrors int64
|
||
fallbackRows int64
|
||
}
|
||
|
||
// RecallWriteBufferInstance 进程级单例（与 SearchCacheInstance / CoOccurTrackerInstance 同模式）
|
||
var RecallWriteBufferInstance *RecallWriteBuffer
|
||
|
||
// NewRecallWriteBuffer 创建缓冲并启动后台 flush 协程；interval<=0 表示只按阈值触发。
|
||
func NewRecallWriteBuffer(ldb LanceDB, interval time.Duration, maxItems int) *RecallWriteBuffer {
|
||
if maxItems <= 0 {
|
||
maxItems = defaultRecallFlushMaxItems
|
||
}
|
||
b := &RecallWriteBuffer{
|
||
pending: make(map[string]*pendingRecall),
|
||
ldb: ldb,
|
||
interval: interval,
|
||
maxItems: maxItems,
|
||
stopCh: make(chan struct{}),
|
||
doneCh: make(chan struct{}),
|
||
}
|
||
// Rust IPC 后端支持单事务批量更新（CASE WHEN 一次提交）
|
||
if bu, ok := ldb.(BatchRecallUpdater); ok {
|
||
b.batch = bu
|
||
}
|
||
go b.loop()
|
||
return b
|
||
}
|
||
|
||
// InitRecallWriteBuffer 初始化进程级单例（server 启动时调用一次）。
|
||
// 窗口可用环境变量 RECALL_WRITE_FLUSH_SECONDS 覆盖（运维/验证用）。
|
||
func InitRecallWriteBuffer(ldb LanceDB) *RecallWriteBuffer {
|
||
if RecallWriteBufferInstance != nil {
|
||
return RecallWriteBufferInstance
|
||
}
|
||
interval := defaultRecallFlushInterval
|
||
if v := os.Getenv("RECALL_WRITE_FLUSH_SECONDS"); v != "" {
|
||
if n, err := strconv.Atoi(v); err == nil && n >= 0 {
|
||
interval = time.Duration(n) * time.Second
|
||
}
|
||
}
|
||
b := NewRecallWriteBuffer(ldb, interval, defaultRecallFlushMaxItems)
|
||
RecallWriteBufferInstance = b
|
||
log.Printf("[recall-buffer] 已启用累积延迟更新: flush 窗口=%s, 阈值=%d 条, 单事务批量=%v",
|
||
interval, b.maxItems, b.batch != nil)
|
||
return b
|
||
}
|
||
|
||
// StopRecallWriteBuffer 停止单例并落盘最后一个窗口（进程优雅退出时调用）。
|
||
func StopRecallWriteBuffer() {
|
||
if b := RecallWriteBufferInstance; b != nil {
|
||
b.Stop()
|
||
}
|
||
}
|
||
|
||
func (b *RecallWriteBuffer) loop() {
|
||
defer close(b.doneCh)
|
||
if b.interval <= 0 {
|
||
<-b.stopCh
|
||
return
|
||
}
|
||
t := time.NewTicker(b.interval)
|
||
defer t.Stop()
|
||
for {
|
||
select {
|
||
case <-b.stopCh:
|
||
return
|
||
case <-t.C:
|
||
b.Flush()
|
||
}
|
||
}
|
||
}
|
||
|
||
// Record 累积一次召回命中的记忆 ID（同窗口内同 ID 合并 delta）。
|
||
// 返回当前待写条数；达到阈值时异步触发一次 flush（不阻塞读路径）。
|
||
func (b *RecallWriteBuffer) Record(ids []string) int {
|
||
now := time.Now()
|
||
b.mu.Lock()
|
||
for _, id := range ids {
|
||
if id == "" {
|
||
continue
|
||
}
|
||
p, ok := b.pending[id]
|
||
if !ok {
|
||
p = &pendingRecall{}
|
||
b.pending[id] = p
|
||
}
|
||
p.delta++
|
||
p.lastAt = now
|
||
}
|
||
n := len(b.pending)
|
||
b.mu.Unlock()
|
||
|
||
if n >= b.maxItems {
|
||
go b.Flush()
|
||
}
|
||
return n
|
||
}
|
||
|
||
// Flush 取出当前窗口的全部增量并批量提交（每批最多 1 个 LanceDB 版本）。
|
||
func (b *RecallWriteBuffer) Flush() (int, error) {
|
||
b.mu.Lock()
|
||
if len(b.pending) == 0 {
|
||
b.mu.Unlock()
|
||
return 0, nil
|
||
}
|
||
batch := make([]RecallWriteItem, 0, len(b.pending))
|
||
var lastAt time.Time
|
||
for id, p := range b.pending {
|
||
if p.delta <= 0 {
|
||
continue
|
||
}
|
||
batch = append(batch, RecallWriteItem{ID: id, Delta: p.delta})
|
||
if p.lastAt.After(lastAt) {
|
||
lastAt = p.lastAt
|
||
}
|
||
}
|
||
b.pending = make(map[string]*pendingRecall)
|
||
b.mu.Unlock()
|
||
|
||
if len(batch) == 0 {
|
||
return 0, nil
|
||
}
|
||
sort.Slice(batch, func(i, j int) bool { return batch[i].ID < batch[j].ID })
|
||
ts := lastAt.Format(time.RFC3339)
|
||
|
||
var rows int64
|
||
var err error
|
||
usedFallback := false
|
||
if b.batch != nil {
|
||
rows, err = b.batch.UpdateRecallBatch("memories", batch, ts)
|
||
}
|
||
if b.batch == nil || err != nil {
|
||
usedFallback = true
|
||
if err != nil {
|
||
log.Printf("[recall-buffer] 单事务批量提交失败(%v) → 退化逐条 Update（数据不丢，版本数不优化）", err)
|
||
atomic.AddInt64(&b.flushErrors, 1)
|
||
}
|
||
rows = 0
|
||
for _, it := range batch {
|
||
if uerr := b.ldb.Update("memories", it.ID, map[string]any{
|
||
"recall_count": map[string]string{"$inc": "1"},
|
||
"last_recalled_at": ts,
|
||
"freshness": "verified",
|
||
}); uerr != nil {
|
||
log.Printf("[recall-buffer] 逐条 Update 失败 id=%s: %v", it.ID, uerr)
|
||
continue
|
||
}
|
||
rows++
|
||
}
|
||
atomic.AddInt64(&b.fallbackRows, rows)
|
||
}
|
||
|
||
// 落盘成功后同步进程内缓存（缓存只是热数据，权威值在 LanceDB）
|
||
if rows > 0 {
|
||
_local.mu.Lock()
|
||
for _, it := range batch {
|
||
if m, ok := _local.memories[it.ID]; ok {
|
||
m.RecallCount += it.Delta
|
||
if !lastAt.IsZero() {
|
||
m.LastRecalledAt = lastAt
|
||
}
|
||
}
|
||
}
|
||
_local.mu.Unlock()
|
||
}
|
||
|
||
atomic.AddInt64(&b.flushCount, 1)
|
||
atomic.AddInt64(&b.flushItems, int64(len(batch)))
|
||
atomic.AddInt64(&b.flushRows, rows)
|
||
log.Printf("[recall-buffer] flush: 合并 %d 条记忆 → 落盘 %d 行, 单事务=%v, fallback=%v",
|
||
len(batch), rows, !usedFallback, usedFallback)
|
||
return int(rows), nil
|
||
}
|
||
|
||
// Stop 停止后台 flush 协程，并把最后一个窗口落盘（进程优雅退出时调用）。
|
||
func (b *RecallWriteBuffer) Stop() {
|
||
b.mu.Lock()
|
||
if b.stopped {
|
||
b.mu.Unlock()
|
||
return
|
||
}
|
||
b.stopped = true
|
||
b.mu.Unlock()
|
||
|
||
close(b.stopCh)
|
||
select {
|
||
case <-b.doneCh:
|
||
case <-time.After(10 * time.Second):
|
||
log.Printf("[recall-buffer] 停止超时，直接落盘剩余窗口")
|
||
}
|
||
b.Flush()
|
||
}
|
||
|
||
// Stats 观测用（flush 频率 / 合并率 / 版本写入行数）。
|
||
func (b *RecallWriteBuffer) Stats() map[string]interface{} {
|
||
b.mu.Lock()
|
||
pendingN := len(b.pending)
|
||
b.mu.Unlock()
|
||
return map[string]interface{}{
|
||
"pending": pendingN,
|
||
"flush_count": atomic.LoadInt64(&b.flushCount),
|
||
"flush_items": atomic.LoadInt64(&b.flushItems),
|
||
"flush_rows": atomic.LoadInt64(&b.flushRows),
|
||
"flush_errors": atomic.LoadInt64(&b.flushErrors),
|
||
"fallback_rows": atomic.LoadInt64(&b.fallbackRows),
|
||
"window": b.interval.String(),
|
||
"max_items": b.maxItems,
|
||
"single_tx": b.batch != nil,
|
||
}
|
||
}
|