mirror of
https://github.com/tiennm99/goclaw.git
synced 2026-10-14 06:17:41 +00:00
fix(vault): aggregate enrichment progress across per-agent batches
After rescan agent_key fix, docs split into per-agent batch queues. Each batch was calling Start()/Finish() independently, causing the progress bar to flash 0/1 per agent instead of showing global progress. Now: handler Start(total) sets global count, worker batches call TrackBatch()/MarkBatchDone() for lifecycle, AddDone(n) auto-completes when done >= total. Progress bar shows smooth 0→N across all agents.
This commit is contained in:
1 parent
d58e0366fd
commit
5b52ce27a8
2 files changed
+49
-23
No files matched your search
@@ -9,13 +9,17 @@ import (
|
||||
)
|
||||
|
||||
// EnrichProgress tracks enrichment pipeline progress and broadcasts via WS events.
|
||||
// Lifecycle: handler calls Start(total) once with the global count,
|
||||
// worker batches call AddDone(n) as docs complete, and MarkBatchDone()
|
||||
// when a per-agent batch finishes. Auto-completes when done >= total.
|
||||
type EnrichProgress struct {
|
||||
mu sync.Mutex
|
||||
msgBus bus.EventPublisher
|
||||
tenantID uuid.UUID
|
||||
total int
|
||||
done int
|
||||
running bool
|
||||
mu sync.Mutex
|
||||
msgBus bus.EventPublisher
|
||||
tenantID uuid.UUID
|
||||
total int
|
||||
done int
|
||||
running bool
|
||||
activeBatches int // number of worker batches in flight
|
||||
}
|
||||
|
||||
// NewEnrichProgress creates a progress tracker that broadcasts to WS clients.
|
||||
@@ -49,33 +53,62 @@ func (p *EnrichProgress) broadcast(e EnrichEvent) {
|
||||
bus.BroadcastForTenant(p.msgBus, protocol.EventVaultEnrichProgress, p.tenantID, e)
|
||||
}
|
||||
|
||||
// Start signals enrichment with updated total and tenant scope.
|
||||
// Only resets done counter when starting a new batch (was idle).
|
||||
// Subsequent calls within the same batch just update the total.
|
||||
// Start signals enrichment with the global total. Called by the HTTP rescan
|
||||
// handler ONCE with the full count. Resets counters for a fresh run.
|
||||
func (p *EnrichProgress) Start(total int, tenantID uuid.UUID) {
|
||||
p.mu.Lock()
|
||||
defer p.mu.Unlock()
|
||||
if !p.running {
|
||||
p.done = 0
|
||||
}
|
||||
p.done = 0
|
||||
p.total = total
|
||||
p.tenantID = tenantID
|
||||
p.running = true
|
||||
p.broadcast(EnrichEvent{Phase: "enriching", Done: p.done, Total: p.total, Running: true})
|
||||
p.activeBatches = 0
|
||||
p.broadcast(EnrichEvent{Phase: "enriching", Done: 0, Total: total, Running: true})
|
||||
}
|
||||
|
||||
// TrackBatch increments the active batch counter. Called by worker when
|
||||
// a per-agent batch starts processing.
|
||||
func (p *EnrichProgress) TrackBatch() {
|
||||
p.mu.Lock()
|
||||
defer p.mu.Unlock()
|
||||
p.activeBatches++
|
||||
}
|
||||
|
||||
// AddDone increments completed count by n and broadcasts progress.
|
||||
// Auto-completes when done >= total.
|
||||
func (p *EnrichProgress) AddDone(n int) {
|
||||
p.mu.Lock()
|
||||
defer p.mu.Unlock()
|
||||
p.done += n
|
||||
if p.done >= p.total && p.total > 0 {
|
||||
p.broadcast(EnrichEvent{Phase: "complete", Done: p.done, Total: p.total, Running: false})
|
||||
p.running = false
|
||||
return
|
||||
}
|
||||
p.broadcast(EnrichEvent{Phase: "enriching", Done: p.done, Total: p.total, Running: true})
|
||||
}
|
||||
|
||||
// MarkBatchDone decrements the active batch counter. If all batches
|
||||
// finished and done < total (some docs failed), force-complete.
|
||||
func (p *EnrichProgress) MarkBatchDone() {
|
||||
p.mu.Lock()
|
||||
defer p.mu.Unlock()
|
||||
p.activeBatches--
|
||||
if p.activeBatches <= 0 && p.running {
|
||||
p.broadcast(EnrichEvent{Phase: "complete", Done: p.done, Total: p.total, Running: false})
|
||||
p.running = false
|
||||
p.activeBatches = 0
|
||||
}
|
||||
}
|
||||
|
||||
// Finish signals the enrichment pipeline has completed and resets counters.
|
||||
// Kept for backward compat but prefer AddDone auto-complete or MarkBatchDone.
|
||||
func (p *EnrichProgress) Finish() {
|
||||
p.mu.Lock()
|
||||
defer p.mu.Unlock()
|
||||
if !p.running {
|
||||
return // already completed via auto-complete
|
||||
}
|
||||
p.broadcast(EnrichEvent{Phase: "complete", Done: p.done, Total: p.total, Running: false})
|
||||
p.done = 0
|
||||
p.total = 0
|
||||
|
||||
@@ -12,7 +12,6 @@ import (
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"github.com/nextlevelbuilder/goclaw/internal/bus"
|
||||
"github.com/nextlevelbuilder/goclaw/internal/eventbus"
|
||||
"github.com/nextlevelbuilder/goclaw/internal/providers"
|
||||
@@ -133,26 +132,20 @@ type enriched struct {
|
||||
// Items are chunked into enrichBatchSize groups so bulk rescan doesn't
|
||||
// overwhelm the LLM provider with hundreds of concurrent requests.
|
||||
func (w *enrichWorker) processBatch(ctx context.Context, key string) {
|
||||
w.progress.TrackBatch()
|
||||
var totalQueued int
|
||||
|
||||
for {
|
||||
items := w.queue.Drain(key)
|
||||
if len(items) == 0 {
|
||||
if w.queue.TryFinish(key) {
|
||||
if totalQueued > 0 {
|
||||
w.progress.Finish()
|
||||
}
|
||||
w.progress.MarkBatchDone()
|
||||
return
|
||||
}
|
||||
continue
|
||||
}
|
||||
|
||||
totalQueued += len(items)
|
||||
tenantID := uuid.Nil
|
||||
if len(items) > 0 {
|
||||
tenantID, _ = uuid.Parse(items[0].TenantID)
|
||||
}
|
||||
w.progress.Start(totalQueued, tenantID)
|
||||
|
||||
// Process in chunks of enrichBatchSize, up to enrichMaxConcurrent in parallel.
|
||||
var wg sync.WaitGroup
|
||||
@@ -172,7 +165,7 @@ func (w *enrichWorker) processBatch(ctx context.Context, key string) {
|
||||
wg.Wait()
|
||||
|
||||
if w.queue.TryFinish(key) {
|
||||
w.progress.Finish()
|
||||
w.progress.MarkBatchDone()
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
Reference in new issue
Block a user