sift/internal/digest/background.go

package digest

import (
	"context"
	"sync"
	"time"

	"git.wntrmute.dev/kyle/sift/internal/dcache"
)

// maxBgQueue is the number of documents the background worker may hold at once.
const maxBgQueue = 4

// bgWorker is the single background goroutine that embeds the chunks of a long document after a
// bm25 answer queued it. It owns a bounded, deduplicated queue and yields the Dixie gate to
// interactive digests one batch at a time.
type bgWorker struct {
	engine *Engine

	mu      sync.Mutex
	cond    *sync.Cond
	queue   []embedJob
	seen    map[string]struct{}
	inProg  bool
	started bool
}

type embedJob struct {
	chunks      []string
	contentHash string
	callerClass string
}

// worker returns the engine's background worker, creating it once. It does not start the goroutine.
func (e *Engine) worker() *bgWorker {
	e.bgOnce.Do(func() {
		w := &bgWorker{engine: e}
		w.cond = sync.NewCond(&w.mu)
		w.seen = map[string]struct{}{}
		e.bg = w
	})
	return e.bg
}

// queueForEmbed asks the background worker to embed a document's chunks. The document is
// deduplicated by contentHash+"|"+callerClass; a full queue drops it (and logs) so a later
// question re-queues it.
func (e *Engine) queueForEmbed(chunks []string, contentHash, callerClass string) {
	if e == nil {
		return
	}
	w := e.worker()
	w.mu.Lock()
	defer w.mu.Unlock()
	if !w.started {
		w.started = true
		go w.run()
	}
	key := contentHash + "|" + callerClass
	if _, ok := w.seen[key]; ok {
		return
	}
	if len(w.queue) >= maxBgQueue {
		// Leave it un-seen: a later question re-queues it.
		if e.Log != nil {
			e.Log.Warn("background embed queue full")
		}
		return
	}
	w.seen[key] = struct{}{}
	w.queue = append(w.queue, embedJob{chunks: chunks, contentHash: contentHash, callerClass: callerClass})
	w.cond.Signal()
}

func (w *bgWorker) run() {
	for {
		w.mu.Lock()
		for len(w.queue) == 0 {
			w.cond.Wait()
		}
		job := w.queue[0]
		w.queue = w.queue[1:]
		w.inProg = true
		w.mu.Unlock()

		w.process(job)

		// Forget the document whether it finished or not: dedupe covers only queued and
		// in-progress work, so a later question re-queues a document whose embedding failed.
		w.mu.Lock()
		delete(w.seen, job.contentHash+"|"+job.callerClass)
		w.inProg = false
		w.mu.Unlock()
	}
}

// process embeds the first MaxDocChunks chunks that lack a vector, EmbedBatch at a time, yielding
// the gate between batches and storing each vector. On any embedding error it drops the rest of
// the document.
func (w *bgWorker) process(job embedJob) {
	e := w.engine
	key := func(i int) string {
		return dcache.EmbedKey(job.contentHash+"/sp1200", i, e.embedModel(), job.callerClass)
	}
	var idx []int
	for i := range job.chunks {
		if i >= MaxDocChunks {
			break
		}
		if _, ok := e.Cache.GetEmbed(key(i)); ok {
			continue
		}
		idx = append(idx, i)
		if len(idx) >= MaxDocChunks {
			break
		}
	}
	for b := 0; b < len(idx); b += EmbedBatch {
		j := b + EmbedBatch
		if j > len(idx) {
			j = len(idx)
		}
		batch := make([]string, j-b)
		for k := range batch {
			batch[k] = job.chunks[idx[b+k]]
		}
		w.acquireGate()
		ectx, cancel := context.WithTimeout(context.Background(), e.timeout())
		vecs, err := e.embed(ectx, batch)
		cancel()
		e.Gate.Release()
		if err != nil {
			if e.Log != nil {
				e.Log.Warn("background embed", "err", truncateErr(err, 512))
			}
			return
		}
		for k := range batch {
			if k < len(vecs) && len(vecs[k]) > 0 {
				e.Cache.PutEmbed(key(idx[b+k]), vecs[k])
			}
		}
	}
}

// acquireGate takes the Dixie gate, sleeping a second after each refusal (queue full or wait
// expired) so interactive digests get their turns. Background work has no deadline: it waits.
func (w *bgWorker) acquireGate() {
	for {
		if ok, _ := w.engine.Gate.Acquire(context.Background()); ok {
			return
		}
		time.Sleep(time.Second)
	}
}

// WaitBackground returns when the background worker is idle: the queue empty and nothing embedding.
func (e *Engine) WaitBackground(ctx context.Context) error {
	w := e.worker()
	for {
		w.mu.Lock()
		idle := len(w.queue) == 0 && !w.inProg
		w.mu.Unlock()
		if idle {
			return nil
		}
		if err := ctx.Err(); err != nil {
			return err
		}
		time.Sleep(20 * time.Millisecond)
	}
}