mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-11-10 21:32:43 +08:00
feat(ingestion): report sub-component parsing progress instead of a frozen 20% (#19965)
This commit is contained in:
1 parent
3bc38fe97a
commit
12e6d888db
26 files changed
+1485
-161
No files matched your search
@@ -257,6 +257,10 @@ func realComponentBodyWithOptions(cpnID, componentClass string, comp runtime.Com
|
||||
// model deadlines remain enforced at the model driver.
|
||||
cctx, cancel := context.WithCancel(ctx)
|
||||
defer cancel()
|
||||
// Bind the node id into the fraction reporter so the component's
|
||||
// runtime.ReportComponentFraction calls are attributed to this node
|
||||
// without the component knowing its own cpnID.
|
||||
cctx = runtime.BindComponentFraction(cctx, cpnID)
|
||||
|
||||
var out map[string]any
|
||||
invokeErr := runtime.TrackProgress(cpnID, runtime.ProgressCallbackFromContext(ctx), func() error {
|
||||
|
||||
@@ -0,0 +1,86 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
package canvas
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
|
||||
"ragflow/internal/agent/runtime"
|
||||
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
// fractionComponent reports a fixed fraction from inside Invoke, the way a
|
||||
// real component reports pages parsed or chunks embedded.
|
||||
type fractionComponent struct {
|
||||
fraction float64
|
||||
}
|
||||
|
||||
func (f *fractionComponent) Name() string { return "fraction" }
|
||||
|
||||
func (f *fractionComponent) Invoke(ctx context.Context, _ *gorm.DB, in map[string]any) (map[string]any, error) {
|
||||
runtime.ReportComponentFraction(ctx, f.fraction)
|
||||
return in, nil
|
||||
}
|
||||
|
||||
func (f *fractionComponent) Stream(_ context.Context, _ map[string]any) (<-chan map[string]any, error) {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
func (f *fractionComponent) Inputs() map[string]string { return nil }
|
||||
func (f *fractionComponent) Outputs() map[string]string { return nil }
|
||||
|
||||
// TestRealComponentBody_BindsFractionToNodeID verifies realComponentBody is
|
||||
// the single injection point that attributes fraction reports to the running
|
||||
// node: the component reports a bare fraction and the run-level callback
|
||||
// receives it under the node's cpnID. Two nodes of the same class must not
|
||||
// cross-attribute.
|
||||
func TestRealComponentBody_BindsFractionToNodeID(t *testing.T) {
|
||||
var got [][2]interface{}
|
||||
ctx := runtime.WithProgressFractionCallback(t.Context(), func(component string, fraction float64) {
|
||||
got = append(got, [2]interface{}{component, fraction})
|
||||
})
|
||||
|
||||
bodyA := realComponentBody("Parser:aaa", "TestFraction", &fractionComponent{fraction: 0.25})
|
||||
bodyB := realComponentBody("Parser:bbb", "TestFraction", &fractionComponent{fraction: 0.75})
|
||||
if _, err := bodyA(ctx, map[string]any{}); err != nil {
|
||||
t.Fatalf("bodyA: %v", err)
|
||||
}
|
||||
if _, err := bodyB(ctx, map[string]any{}); err != nil {
|
||||
t.Fatalf("bodyB: %v", err)
|
||||
}
|
||||
|
||||
if len(got) != 2 {
|
||||
t.Fatalf("fraction reports = %d, want 2", len(got))
|
||||
}
|
||||
if got[0][0] != "Parser:aaa" || got[0][1] != 0.25 {
|
||||
t.Fatalf("report 0 = %v, want [Parser:aaa 0.25]", got[0])
|
||||
}
|
||||
if got[1][0] != "Parser:bbb" || got[1][1] != 0.75 {
|
||||
t.Fatalf("report 1 = %v, want [Parser:bbb 0.75]", got[1])
|
||||
}
|
||||
}
|
||||
|
||||
// TestRealComponentBody_NoFractionSinkIsNoop verifies a headless run (no
|
||||
// run-level callback) pays nothing and components still report freely.
|
||||
func TestRealComponentBody_NoFractionSinkIsNoop(t *testing.T) {
|
||||
body := realComponentBody("test-cpn", "TestFraction", &fractionComponent{fraction: 0.5})
|
||||
if _, err := body(t.Context(), map[string]any{}); err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
}
|
||||
@@ -251,3 +251,53 @@ func ProgressCallbackFromContext(ctx context.Context) ProgressCallback {
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// ProgressFractionCallback receives an in-flight component's completion
|
||||
// fraction (0..1) — pages parsed, chunks embedded. Unlike lifecycle events,
|
||||
// fractions are high-frequency and purely observational: they refine the
|
||||
// progress percentage between two lifecycle events and carry no persistence
|
||||
// contract. component is the node id (cpnID), the same identity
|
||||
// ProgressEvent.Component carries, so sinks key both channels identically.
|
||||
type ProgressFractionCallback func(component string, fraction float64)
|
||||
|
||||
// progressFractionCBKey carries the run-level ProgressFractionCallback.
|
||||
type progressFractionCBKey struct{}
|
||||
|
||||
// componentFractionReporterKey carries the per-node closure that
|
||||
// ReportComponentFraction invokes. It is pre-bound with the cpnID by
|
||||
// BindComponentFraction so components report a bare fraction and can never
|
||||
// misattribute progress to another node.
|
||||
type componentFractionReporterKey struct{}
|
||||
|
||||
// WithProgressFractionCallback attaches a run-level fraction sink to ctx.
|
||||
// A nil callback is valid and keeps components observer-independent.
|
||||
func WithProgressFractionCallback(ctx context.Context, cb ProgressFractionCallback) context.Context {
|
||||
return context.WithValue(ctx, progressFractionCBKey{}, cb)
|
||||
}
|
||||
|
||||
// BindComponentFraction derives a context whose ReportComponentFraction calls
|
||||
// carry the given component id. The canvas framework (realComponentBody) is
|
||||
// the single caller, mirroring how it owns TrackProgress; when no run-level
|
||||
// callback is attached it returns ctx unchanged so headless runs pay nothing.
|
||||
func BindComponentFraction(ctx context.Context, component string) context.Context {
|
||||
cb, _ := ctx.Value(progressFractionCBKey{}).(ProgressFractionCallback)
|
||||
if cb == nil {
|
||||
return ctx
|
||||
}
|
||||
return context.WithValue(ctx, componentFractionReporterKey{}, func(fraction float64) {
|
||||
cb(component, fraction)
|
||||
})
|
||||
}
|
||||
|
||||
// ReportComponentFraction forwards a component's in-flight completion
|
||||
// fraction to the run-level sink. Components call it from their progress
|
||||
// loops without knowing their own node id; the framework-bound closure
|
||||
// supplies the attribution. No-op when no sink is attached.
|
||||
func ReportComponentFraction(ctx context.Context, fraction float64) {
|
||||
if ctx == nil {
|
||||
return
|
||||
}
|
||||
if report, ok := ctx.Value(componentFractionReporterKey{}).(func(float64)); ok && report != nil {
|
||||
report(fraction)
|
||||
}
|
||||
}
|
||||
@@ -387,3 +387,66 @@ func TestTrackElapsed_NilMapFromFn(t *testing.T) {
|
||||
t.Error("missing _elapsed_time after nil-map input")
|
||||
}
|
||||
}
|
||||
|
||||
// TestReportComponentFraction_NoopWithoutBinding verifies a component can
|
||||
// report fractions unconditionally: with no run-level callback bound the
|
||||
// report is dropped, keeping headless runs observer-independent.
|
||||
func TestReportComponentFraction_NoopWithoutBinding(t *testing.T) {
|
||||
ReportComponentFraction(context.Background(), 0.5)
|
||||
ReportComponentFraction(nil, 0.5)
|
||||
|
||||
// A run callback that was never bound to a component must not fire either.
|
||||
ctx := WithProgressFractionCallback(context.Background(), func(string, float64) {
|
||||
t.Fatal("callback fired without BindComponentFraction")
|
||||
})
|
||||
ReportComponentFraction(ctx, 0.5)
|
||||
}
|
||||
|
||||
// TestBindComponentFraction_AttributesReports verifies the framework binding
|
||||
// supplies the component id: the run-level callback receives (cpnID, frac)
|
||||
// while the component itself only reports a bare fraction.
|
||||
func TestBindComponentFraction_AttributesReports(t *testing.T) {
|
||||
var mu sync.Mutex
|
||||
var got []struct {
|
||||
component string
|
||||
fraction float64
|
||||
}
|
||||
ctx := WithProgressFractionCallback(context.Background(), func(component string, fraction float64) {
|
||||
mu.Lock()
|
||||
defer mu.Unlock()
|
||||
got = append(got, struct {
|
||||
component string
|
||||
fraction float64
|
||||
}{component, fraction})
|
||||
})
|
||||
|
||||
ctxA := BindComponentFraction(ctx, "Parser:aaa")
|
||||
ctxB := BindComponentFraction(ctx, "Parser:bbb")
|
||||
ReportComponentFraction(ctxA, 0.25)
|
||||
ReportComponentFraction(ctxB, 0.75)
|
||||
ReportComponentFraction(ctxA, 0.5)
|
||||
|
||||
mu.Lock()
|
||||
defer mu.Unlock()
|
||||
if len(got) != 3 {
|
||||
t.Fatalf("callback calls = %d, want 3", len(got))
|
||||
}
|
||||
want := []struct {
|
||||
component string
|
||||
fraction float64
|
||||
}{{"Parser:aaa", 0.25}, {"Parser:bbb", 0.75}, {"Parser:aaa", 0.5}}
|
||||
for i, w := range want {
|
||||
if got[i] != w {
|
||||
t.Fatalf("call %d = %+v, want %+v", i, got[i], w)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestBindComponentFraction_NoCallbackReturnsSameContext verifies the binding
|
||||
// is free when no run-level sink is attached.
|
||||
func TestBindComponentFraction_NoCallbackReturnsSameContext(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
if bound := BindComponentFraction(ctx, "Parser:aaa"); bound != ctx {
|
||||
t.Fatal("BindComponentFraction allocated a context without a run-level callback")
|
||||
}
|
||||
}
|
||||
@@ -366,6 +366,12 @@ func (p *Parser) processPageBoxes(ctx context.Context, pageImg image.Image, char
|
||||
// visibility, while context cancellation still stops new dispatch and lets
|
||||
// in-flight work observe ctx.Err().
|
||||
//
|
||||
// OnPageDone is reported by the worker itself as its page finishes (see
|
||||
// pageProgress), not by the collection loop below: the loop only starts after
|
||||
// every page has been submitted, and SubmitTo blocks once the worker queue is
|
||||
// full, so collector-side reporting would stay silent until submission
|
||||
// finished, then deliver every completion in one burst.
|
||||
//
|
||||
// Pages are returned sorted by page number so callers can stream them
|
||||
// directly into downstream assembly without re-sorting.
|
||||
func (p *Parser) runPageWorkers(ctx context.Context, engine pdf.PDFEngine,
|
||||
@@ -389,6 +395,7 @@ func (p *Parser) runPageWorkers(ctx context.Context, engine pdf.PDFEngine,
|
||||
|
||||
type pageTaskResult = utility.WorkerPoolResult[pageTask, pageResult]
|
||||
resultCh := make(chan pageTaskResult, len(pages))
|
||||
progress := &pageProgress{total: len(pages), onDone: p.Config.OnPageDone}
|
||||
|
||||
submitted := 0
|
||||
for _, pg := range pages {
|
||||
@@ -398,6 +405,7 @@ func (p *Parser) runPageWorkers(ctx context.Context, engine pdf.PDFEngine,
|
||||
pageNumber: pg,
|
||||
docAnalyzer: docAnalyzer,
|
||||
tb: tb,
|
||||
progress: progress,
|
||||
}
|
||||
if err := parserPageWorkerPool().SubmitTo(ctx, task, resultCh); err != nil {
|
||||
recordErr(err)
|
||||
|
||||
@@ -57,6 +57,39 @@ type pageTask struct {
|
||||
pageNumber int
|
||||
docAnalyzer pdf.DocAnalyzer
|
||||
tb pdf.TableBuilder
|
||||
// progress reports this run's page completions; the worker fires the
|
||||
// caller's callback as soon as its page finishes.
|
||||
progress *pageProgress
|
||||
}
|
||||
|
||||
// pageProgress serialises page-completion reporting for one runPageWorkers
|
||||
// run. The counter and the callback share one lock so completions are reported
|
||||
// in order and never concurrently — the contract ParserConfig.OnPageDone
|
||||
// documents — even though the pages themselves run on parallel workers.
|
||||
//
|
||||
// Reporting from the worker is what ties the callback to the page that just
|
||||
// finished. The collection loop cannot run until every page has been
|
||||
// submitted, and SubmitTo blocks once the worker queue is full, so on a
|
||||
// document larger than the pool's worker+queue capacity a collector-side
|
||||
// callback stays silent while the pool drains the overflow and then reports
|
||||
// every completion behind it in one burst.
|
||||
type pageProgress struct {
|
||||
mu sync.Mutex
|
||||
done int
|
||||
total int
|
||||
onDone func(done, total int)
|
||||
}
|
||||
|
||||
// finish records one completed page and reports it. It is a no-op when no
|
||||
// callback is configured (the default), so an unwatched parse pays nothing.
|
||||
func (p *pageProgress) finish() {
|
||||
if p.onDone == nil {
|
||||
return
|
||||
}
|
||||
p.mu.Lock()
|
||||
defer p.mu.Unlock()
|
||||
p.done++
|
||||
p.onDone(p.done, p.total)
|
||||
}
|
||||
|
||||
const pageWorkerQueueFactor = 4
|
||||
@@ -81,6 +114,7 @@ func parserPageWorkerPool() *utility.WorkerPool[pageTask, pageResult] {
|
||||
}
|
||||
pagePool = utility.NewWorkerPool(workers, workers*pageWorkerQueueFactor,
|
||||
func(ctx context.Context, task pageTask) (pageResult, error) {
|
||||
defer task.progress.finish()
|
||||
return task.parser.processPage(ctx, task.engine, task.pageNumber,
|
||||
task.docAnalyzer, task.tb), nil
|
||||
})
|
||||
|
||||
@@ -0,0 +1,129 @@
|
||||
package pdf
|
||||
|
||||
import (
|
||||
"image"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
pdf "ragflow/internal/deepdoc/parser/pdf/type"
|
||||
)
|
||||
|
||||
// TestParseRaw_OnPageDone_ReportsEveryPage verifies the page callback fires
|
||||
// once per collected page with a monotonically increasing done count and the
|
||||
// total number of submitted pages.
|
||||
func TestParseRaw_OnPageDone_ReportsEveryPage(t *testing.T) {
|
||||
eng := makePageTaggedEngine(6)
|
||||
|
||||
type observation struct{ done, total int }
|
||||
var seen []observation
|
||||
cfg := pdf.DefaultParserConfig()
|
||||
cfg.OnPageDone = func(done, total int) {
|
||||
seen = append(seen, observation{done, total})
|
||||
}
|
||||
p := NewParser(cfg)
|
||||
|
||||
if _, err := p.ParseRaw(t.Context(), eng, &MockDocAnalyzer{Healthy: true}); err != nil {
|
||||
t.Fatalf("ParseRaw: %v", err)
|
||||
}
|
||||
|
||||
if len(seen) != 6 {
|
||||
t.Fatalf("expected 6 callbacks, got %d: %v", len(seen), seen)
|
||||
}
|
||||
for i, obs := range seen {
|
||||
if obs.done != i+1 {
|
||||
t.Errorf("callback %d: done = %d, want %d", i, obs.done, i+1)
|
||||
}
|
||||
if obs.total != 6 {
|
||||
t.Errorf("callback %d: total = %d, want 6", i, obs.total)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestParseRaw_OnPageDone_TotalIsSubmittedPages verifies the reported total
|
||||
// reflects the page-range restriction rather than the document page count, so
|
||||
// a caller computing done/total reaches 1.0 when the run finishes.
|
||||
func TestParseRaw_OnPageDone_TotalIsSubmittedPages(t *testing.T) {
|
||||
eng := makePageTaggedEngine(10)
|
||||
|
||||
var totals []int
|
||||
cfg := pdf.DefaultParserConfig()
|
||||
cfg.Pages = [][]int{{1, 3}}
|
||||
cfg.OnPageDone = func(_, total int) {
|
||||
totals = append(totals, total)
|
||||
}
|
||||
p := NewParser(cfg)
|
||||
|
||||
if _, err := p.ParseRaw(t.Context(), eng, &MockDocAnalyzer{Healthy: true}); err != nil {
|
||||
t.Fatalf("ParseRaw: %v", err)
|
||||
}
|
||||
|
||||
if len(totals) != 3 {
|
||||
t.Fatalf("expected 3 callbacks, got %d: %v", len(totals), totals)
|
||||
}
|
||||
for i, total := range totals {
|
||||
if total != 3 {
|
||||
t.Errorf("callback %d: total = %d, want 3", i, total)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestParseRaw_OnPageDone_NilIsNoop is the regression guard: an unset callback
|
||||
// must not change parsing.
|
||||
func TestParseRaw_OnPageDone_NilIsNoop(t *testing.T) {
|
||||
eng := makePageTaggedEngine(4)
|
||||
p := NewParser(pdf.DefaultParserConfig())
|
||||
|
||||
result, err := p.ParseRaw(t.Context(), eng, &MockDocAnalyzer{Healthy: true})
|
||||
if err != nil {
|
||||
t.Fatalf("ParseRaw: %v", err)
|
||||
}
|
||||
if len(result.PageHeight) != 4 {
|
||||
t.Errorf("expected 4 parsed pages, got %d", len(result.PageHeight))
|
||||
}
|
||||
}
|
||||
|
||||
// TestParseRaw_OnPageDone_FiresBeforeSubmissionDrains is the reporting-lag
|
||||
// regression guard. It parses more pages than the worker pool can ever hold
|
||||
// (max 12 workers + 4x queue = 60), so a callback driven by the collection
|
||||
// loop could not run until the whole document had been submitted — every
|
||||
// non-first page here blocks its worker until the first page has been
|
||||
// reported, which a collector-side callback only reaches after they give up.
|
||||
func TestParseRaw_OnPageDone_FiresBeforeSubmissionDrains(t *testing.T) {
|
||||
const numPages = 256
|
||||
|
||||
firstDone := make(chan struct{})
|
||||
giveUp := make(chan struct{})
|
||||
var firstOnce, timeoutOnce sync.Once
|
||||
var late atomic.Bool
|
||||
|
||||
eng := makePageTaggedEngine(numPages)
|
||||
eng.RenderPageImageFunc = func(pg int, _ float64) (image.Image, error) {
|
||||
if pg != 0 {
|
||||
select {
|
||||
case <-firstDone:
|
||||
case <-giveUp:
|
||||
case <-time.After(5 * time.Second):
|
||||
late.Store(true)
|
||||
timeoutOnce.Do(func() { close(giveUp) })
|
||||
}
|
||||
}
|
||||
return image.NewRGBA(image.Rect(0, 0, eng.RenderW, eng.RenderH)), nil
|
||||
}
|
||||
|
||||
cfg := pdf.DefaultParserConfig()
|
||||
cfg.OnPageDone = func(done, _ int) {
|
||||
if done == 1 {
|
||||
firstOnce.Do(func() { close(firstDone) })
|
||||
}
|
||||
}
|
||||
p := NewParser(cfg)
|
||||
|
||||
if _, err := p.ParseRaw(t.Context(), eng, &MockDocAnalyzer{Healthy: true}); err != nil {
|
||||
t.Fatalf("ParseRaw: %v", err)
|
||||
}
|
||||
if late.Load() {
|
||||
t.Error("first page completion was reported only after other pages gave up waiting: OnPageDone is not firing when the page finishes")
|
||||
}
|
||||
}
|
||||
@@ -254,6 +254,15 @@ type ParserConfig struct {
|
||||
// cross-page text repetition) before TextMerge can fold a header box into
|
||||
// a body section.
|
||||
RemoveHeaderFooter bool
|
||||
// OnPageDone, when set, is called as each page finishes, from the worker
|
||||
// that parsed it — not after every page has been submitted — so the first
|
||||
// report arrives with the first completed page on any document size (done
|
||||
// counts completed pages, in completion order; total is the number of
|
||||
// pages to process). Calls are ordered and serialized but run on page
|
||||
// workers, so the callback must be fast and non-blocking. It lets callers
|
||||
// surface parse progress without the parser knowing about any progress
|
||||
// sink; nil disables the callback at zero cost.
|
||||
OnPageDone func(done, total int)
|
||||
}
|
||||
|
||||
// DefaultParserConfig returns a ParserConfig with sensible defaults.
|
||||
|
||||
@@ -41,6 +41,7 @@ import (
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
eschema "github.com/cloudwego/eino/schema"
|
||||
@@ -293,7 +294,7 @@ func NewExtractorComponent(params map[string]any) (runtime.Component, error) {
|
||||
// self.chat_mdl; the Go port exposes it explicitly).
|
||||
func (c *ExtractorComponent) Inputs() map[string]string {
|
||||
return map[string]string{
|
||||
"chunks": "List of map[string]any from upstream Tokenizer. Each entry must carry a string 'text' field. Optional — when absent the LLM is called once with the resolved args.",
|
||||
"chunks": "List of map[string]any from the upstream Chunker. Each entry must carry a string 'text' field. Optional — when absent the LLM is called once with the resolved args.",
|
||||
"llm_id": "Optional per-call LLM id override. Falls back to Param.LLMID when absent.",
|
||||
}
|
||||
}
|
||||
@@ -622,17 +623,20 @@ func (c *ExtractorComponent) Invoke(ctx context.Context, db *gorm.DB, inputs map
|
||||
}, nil
|
||||
}
|
||||
|
||||
keywordsOn := c.Param.Keywords.TopN > 0
|
||||
tagsOn := c.Param.Tags.TopN > 0
|
||||
remainingOn := c.Param.Questions.TopN > 0 || c.Param.Summary.Enabled || c.Param.Metadata.Enabled
|
||||
phaseWindows := newFractionSplit(keywordsOn, tagsOn, remainingOn)
|
||||
|
||||
if err := runtime.WithTimeout(ctx, extractorTimeout, func(timeoutCtx context.Context) error {
|
||||
// Phase 1: Keywords extraction (if enabled), running across chunks via extractorPool.
|
||||
if c.Param.Keywords.TopN > 0 {
|
||||
if err := c.runAutoKeywordsPool(timeoutCtx, db, in); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := c.runAutoKeywordsPool(timeoutCtx, db, in, phaseWindows.take(keywordsOn)); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
// Phase 2: Tag phase (if enabled), benefiting from title and freshly extracted keywords.
|
||||
if c.Param.Tags.TopN > 0 {
|
||||
tagged, tagErr := c.runAutoTags(timeoutCtx, db, in)
|
||||
if tagsOn {
|
||||
tagged, tagErr := c.runAutoTags(timeoutCtx, db, in, phaseWindows.take(tagsOn))
|
||||
if tagErr != nil {
|
||||
return tagErr
|
||||
}
|
||||
@@ -640,7 +644,7 @@ func (c *ExtractorComponent) Invoke(ctx context.Context, db *gorm.DB, inputs map
|
||||
}
|
||||
|
||||
// Phase 3: Remaining extractions (Questions, Summary, Metadata) via extractorPool.
|
||||
return c.runRemainingExtractions(timeoutCtx, db, in)
|
||||
return c.runRemainingExtractions(timeoutCtx, db, in, phaseWindows.take(remainingOn))
|
||||
}); err != nil {
|
||||
return nil, fmt.Errorf("extractor: %w", err)
|
||||
}
|
||||
@@ -844,9 +848,53 @@ func (c *ExtractorComponent) runAutoSummary(ctx context.Context, db *gorm.DB, in
|
||||
return nil
|
||||
}
|
||||
|
||||
// fractionWindow is the slice of a component's 0..1 fraction budget owned by one
|
||||
// sequential phase. Splitting the budget keeps the reported fraction climbing
|
||||
// across phases: the run-level progress tracker keeps a high-water mark, so a
|
||||
// second phase restarting from zero would freeze the bar until the component
|
||||
// exits.
|
||||
type fractionWindow struct{ base, span float64 }
|
||||
|
||||
// report maps a phase completion ratio into the window and forwards it to the
|
||||
// run-level fraction channel. The zero window reports nothing, which keeps
|
||||
// disabled phases free of budget arithmetic.
|
||||
func (w fractionWindow) report(ctx context.Context, ratio float64) {
|
||||
if w.span <= 0 {
|
||||
return
|
||||
}
|
||||
ratio = min(max(ratio, 0), 1)
|
||||
runtime.ReportComponentFraction(ctx, w.base+w.span*ratio)
|
||||
}
|
||||
|
||||
// fractionSplit hands out one equal window per enabled phase, in phase order.
|
||||
type fractionSplit struct {
|
||||
phases int
|
||||
next int
|
||||
}
|
||||
|
||||
func newFractionSplit(enabled ...bool) *fractionSplit {
|
||||
s := &fractionSplit{}
|
||||
for _, on := range enabled {
|
||||
if on {
|
||||
s.phases++
|
||||
}
|
||||
}
|
||||
return s
|
||||
}
|
||||
|
||||
func (s *fractionSplit) take(enabled bool) fractionWindow {
|
||||
if !enabled {
|
||||
return fractionWindow{}
|
||||
}
|
||||
span := 1 / float64(s.phases)
|
||||
w := fractionWindow{base: float64(s.next) * span, span: span}
|
||||
s.next++
|
||||
return w
|
||||
}
|
||||
|
||||
// runAutoKeywordsPool dispatches keyword extraction across all chunks concurrently
|
||||
// using extractorPool before the tagging stage.
|
||||
func (c *ExtractorComponent) runAutoKeywordsPool(ctx context.Context, db *gorm.DB, in extractorInputs) error {
|
||||
func (c *ExtractorComponent) runAutoKeywordsPool(ctx context.Context, db *gorm.DB, in extractorInputs, w fractionWindow) error {
|
||||
if c.Param.Keywords.TopN <= 0 || len(in.chunks) == 0 {
|
||||
return nil
|
||||
}
|
||||
@@ -869,12 +917,12 @@ func (c *ExtractorComponent) runAutoKeywordsPool(ctx context.Context, db *gorm.D
|
||||
}
|
||||
futs = append(futs, f)
|
||||
}
|
||||
return awaitFutures(ctx, futs)
|
||||
return awaitFutures(ctx, futs, w)
|
||||
}
|
||||
|
||||
// runRemainingExtractions dispatches auto questions / summary / metadata
|
||||
// extractions across all chunks concurrently using extractorPool.
|
||||
func (c *ExtractorComponent) runRemainingExtractions(ctx context.Context, db *gorm.DB, in extractorInputs) error {
|
||||
func (c *ExtractorComponent) runRemainingExtractions(ctx context.Context, db *gorm.DB, in extractorInputs, w fractionWindow) error {
|
||||
if c.Param.Questions.TopN <= 0 && !c.Param.Summary.Enabled && !c.Param.Metadata.Enabled {
|
||||
return nil
|
||||
}
|
||||
@@ -892,7 +940,7 @@ func (c *ExtractorComponent) runRemainingExtractions(ctx context.Context, db *go
|
||||
}
|
||||
futs = append(futs, f)
|
||||
}
|
||||
return awaitFutures(ctx, futs)
|
||||
return awaitFutures(ctx, futs, w)
|
||||
}
|
||||
|
||||
func (c *ExtractorComponent) remainingExtractionJob(ctx context.Context, db *gorm.DB, in extractorInputs, idx int, ck map[string]any, chunkText string) extractorJob {
|
||||
@@ -916,14 +964,24 @@ func (c *ExtractorComponent) remainingExtractionJob(ctx context.Context, db *gor
|
||||
}
|
||||
}
|
||||
|
||||
func awaitFutures(ctx context.Context, futs []utility.WorkerPoolFuture[extractorJob, struct{}]) error {
|
||||
// awaitFutures waits for every submitted job and returns the first failure.
|
||||
// Each settled future advances the phase's fraction window, so a chunk-heavy
|
||||
// LLM extraction reports in-flight progress instead of going quiet until the
|
||||
// whole phase is done.
|
||||
func awaitFutures(ctx context.Context, futs []utility.WorkerPoolFuture[extractorJob, struct{}], w fractionWindow) error {
|
||||
if len(futs) == 0 {
|
||||
return nil
|
||||
}
|
||||
var firstErr error
|
||||
var emu sync.Mutex
|
||||
var wg sync.WaitGroup
|
||||
var done atomic.Int64
|
||||
total := int64(len(futs))
|
||||
for _, f := range futs {
|
||||
wg.Add(1)
|
||||
go func(f utility.WorkerPoolFuture[extractorJob, struct{}]) {
|
||||
defer wg.Done()
|
||||
defer func() { w.report(ctx, float64(done.Add(1))/float64(total)) }()
|
||||
res, werr := f.Wait(ctx)
|
||||
if werr != nil {
|
||||
emu.Lock()
|
||||
|
||||
@@ -15,6 +15,7 @@ import (
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
"github.com/cespare/xxhash/v2"
|
||||
@@ -556,7 +557,10 @@ func isHighConfidenceMatch(matched *schema.TaggedChunk) bool {
|
||||
return false
|
||||
}
|
||||
|
||||
func (c *ExtractorComponent) runAutoTags(ctx context.Context, db *gorm.DB, in extractorInputs) ([]map[string]any, error) {
|
||||
func (c *ExtractorComponent) runAutoTags(ctx context.Context, db *gorm.DB, in extractorInputs, w fractionWindow) ([]map[string]any, error) {
|
||||
// Hand the whole window to the next phase on every exit path, including the
|
||||
// skip paths below that do no measurable work.
|
||||
defer w.report(ctx, 1)
|
||||
lang := in.lang
|
||||
if lang == "" {
|
||||
lang = detectTextLanguage(in.chunks)
|
||||
@@ -600,11 +604,14 @@ func (c *ExtractorComponent) runAutoTags(ctx context.Context, db *gorm.DB, in ex
|
||||
inv := getExtractorChatInvoker()
|
||||
sem := make(chan struct{}, taggerLLMConcurrency)
|
||||
var wg sync.WaitGroup
|
||||
var tagged atomic.Int64
|
||||
total := int64(len(docsToTag))
|
||||
|
||||
for i := range docsToTag {
|
||||
wg.Add(1)
|
||||
go func(idx int) {
|
||||
defer wg.Done()
|
||||
defer func() { w.report(ctx, float64(tagged.Add(1))/float64(total)) }()
|
||||
select {
|
||||
case sem <- struct{}{}:
|
||||
defer func() { <-sem }()
|
||||
|
||||
@@ -254,6 +254,7 @@ func ExtractClaimsForChunks(ctx context.Context, deps common.Deps, llmID string,
|
||||
// which advances the counter only for finished batches; workers
|
||||
// finish out of order, so the figure jumps by batch size).
|
||||
completed += len(batch)
|
||||
runtime.ReportComponentFraction(ctx, float64(completed)/float64(total))
|
||||
if reporter.shouldReport(completed) {
|
||||
runtime.ReportProgressMessage(ctx, "Compiler", fmt.Sprintf("tree-template: extracting claims for chunk %d/%d", completed, total))
|
||||
}
|
||||
|
||||
@@ -28,6 +28,7 @@ import (
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
"ragflow/internal/agent/runtime"
|
||||
"ragflow/internal/common"
|
||||
"ragflow/internal/ingestion/component/schema"
|
||||
"ragflow/internal/parser/parser"
|
||||
@@ -96,6 +97,7 @@ func dispatchParse(ctx context.Context, fileType utility.FileType, filename stri
|
||||
return parser.ParseResult{Err: fmt.Errorf("parser: resolve %q: %w", fileType, err)}
|
||||
}
|
||||
configureParserFromSetups(p, fileType, setups)
|
||||
attachPDFPageProgress(ctx, p)
|
||||
|
||||
res := p.ParseWithResult(ctx, filename, data)
|
||||
if res.Err != nil {
|
||||
@@ -112,6 +114,24 @@ func dispatchParse(ctx context.Context, fileType utility.FileType, filename stri
|
||||
return res
|
||||
}
|
||||
|
||||
// attachPDFPageProgress wires the DeepDOC page callback so a long PDF
|
||||
// parse reports in-flight completion instead of staying silent until
|
||||
// the component exits. Only the local DeepDOC backend fires it; remote
|
||||
// engines (MinerU, PaddleOCR, Docling, …) never call the callback and
|
||||
// the fraction stays at zero.
|
||||
func attachPDFPageProgress(ctx context.Context, p parser.ParseResultProducer) {
|
||||
pdfParser, ok := p.(*parser.PDFParser)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
pdfParser.OnPageDone = func(done, total int) {
|
||||
if total <= 0 {
|
||||
return
|
||||
}
|
||||
runtime.ReportComponentFraction(ctx, float64(done)/float64(total))
|
||||
}
|
||||
}
|
||||
|
||||
// fileTypeFromInputs derives the parser-library extension form
|
||||
// (utility.FileType) from the upstream inputs. The result is the
|
||||
// value passed to parser.GetParser, whose switch arms are keyed
|
||||
|
||||
@@ -0,0 +1,76 @@
|
||||
// Page-progress wiring tests pin the DeepDOC → component fraction bridge:
|
||||
// a PDF parse reports done/total pages through the framework-bound fraction
|
||||
// channel, and every other parser family stays untouched.
|
||||
|
||||
package component
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
|
||||
"ragflow/internal/agent/runtime"
|
||||
"ragflow/internal/parser/parser"
|
||||
)
|
||||
|
||||
type fractionObservation struct {
|
||||
component string
|
||||
fraction float64
|
||||
}
|
||||
|
||||
// boundFractionContext mirrors the canvas framework: a run-level callback plus
|
||||
// the per-node binding that supplies attribution.
|
||||
func boundFractionContext(t *testing.T, component string, seen *[]fractionObservation) context.Context {
|
||||
t.Helper()
|
||||
ctx := runtime.WithProgressFractionCallback(t.Context(), func(c string, f float64) {
|
||||
*seen = append(*seen, fractionObservation{c, f})
|
||||
})
|
||||
return runtime.BindComponentFraction(ctx, component)
|
||||
}
|
||||
|
||||
func TestAttachPDFPageProgress_ReportsFraction(t *testing.T) {
|
||||
var seen []fractionObservation
|
||||
ctx := boundFractionContext(t, "Parser:pdf", &seen)
|
||||
|
||||
pdfParser := parser.NewPDFParser()
|
||||
attachPDFPageProgress(ctx, pdfParser)
|
||||
if pdfParser.OnPageDone == nil {
|
||||
t.Fatal("expected OnPageDone to be wired")
|
||||
}
|
||||
|
||||
pdfParser.OnPageDone(3, 12)
|
||||
pdfParser.OnPageDone(12, 12)
|
||||
|
||||
want := []fractionObservation{{"Parser:pdf", 0.25}, {"Parser:pdf", 1}}
|
||||
if len(seen) != len(want) {
|
||||
t.Fatalf("observations = %v, want %v", seen, want)
|
||||
}
|
||||
for i := range want {
|
||||
if seen[i] != want[i] {
|
||||
t.Errorf("observation %d = %v, want %v", i, seen[i], want[i])
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestAttachPDFPageProgress_IgnoresUnknownTotal(t *testing.T) {
|
||||
var seen []fractionObservation
|
||||
ctx := boundFractionContext(t, "Parser:pdf", &seen)
|
||||
|
||||
pdfParser := parser.NewPDFParser()
|
||||
attachPDFPageProgress(ctx, pdfParser)
|
||||
pdfParser.OnPageDone(1, 0)
|
||||
|
||||
if len(seen) != 0 {
|
||||
t.Errorf("expected no fraction for total=0, got %v", seen)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAttachPDFPageProgress_LeavesOtherParsersAlone(t *testing.T) {
|
||||
var seen []fractionObservation
|
||||
ctx := boundFractionContext(t, "Parser:txt", &seen)
|
||||
|
||||
attachPDFPageProgress(ctx, parser.NewTextParser())
|
||||
|
||||
if len(seen) != 0 {
|
||||
t.Errorf("expected no fraction from a non-PDF parser, got %v", seen)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,237 @@
|
||||
// Fraction instrumentation tests pin the component-side progress signal:
|
||||
// the extractor's phase windows stay monotonic across sequential phases,
|
||||
// awaitFutures reports every settled job, and the tokenizer's embedding
|
||||
// batches advance the fraction (cache hits counted as finished work).
|
||||
|
||||
package component
|
||||
|
||||
import (
|
||||
"context"
|
||||
"sync"
|
||||
"testing"
|
||||
|
||||
"ragflow/internal/agent/runtime"
|
||||
"ragflow/internal/ingestion/component/schema"
|
||||
"ragflow/internal/utility"
|
||||
)
|
||||
|
||||
// fractionRecorder is a concurrency-safe fraction sink, since the extractor
|
||||
// reports from several goroutines at once.
|
||||
type fractionRecorder struct {
|
||||
mu sync.Mutex
|
||||
seen []fractionObservation
|
||||
}
|
||||
|
||||
func (r *fractionRecorder) context(t *testing.T, component string) context.Context {
|
||||
t.Helper()
|
||||
ctx := runtime.WithProgressFractionCallback(t.Context(), func(c string, f float64) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
r.seen = append(r.seen, fractionObservation{c, f})
|
||||
})
|
||||
return runtime.BindComponentFraction(ctx, component)
|
||||
}
|
||||
|
||||
func (r *fractionRecorder) fractions() []float64 {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
out := make([]float64, 0, len(r.seen))
|
||||
for _, obs := range r.seen {
|
||||
out = append(out, obs.fraction)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func TestFractionSplit_HandsOutMonotonicWindows(t *testing.T) {
|
||||
split := newFractionSplit(true, false, true)
|
||||
|
||||
first := split.take(true)
|
||||
skipped := split.take(false)
|
||||
second := split.take(true)
|
||||
|
||||
if first.base != 0 || first.span != 0.5 {
|
||||
t.Errorf("first window = %+v, want base 0 span 0.5", first)
|
||||
}
|
||||
if skipped.base != 0 || skipped.span != 0 {
|
||||
t.Errorf("disabled window = %+v, want the zero window", skipped)
|
||||
}
|
||||
if second.base != 0.5 || second.span != 0.5 {
|
||||
t.Errorf("second window = %+v, want base 0.5 span 0.5", second)
|
||||
}
|
||||
|
||||
rec := &fractionRecorder{}
|
||||
ctx := rec.context(t, "Extractor:abc")
|
||||
first.report(ctx, 0.5)
|
||||
skipped.report(ctx, 1)
|
||||
second.report(ctx, 1)
|
||||
|
||||
got := rec.fractions()
|
||||
want := []float64{0.25, 1}
|
||||
if len(got) != len(want) {
|
||||
t.Fatalf("fractions = %v, want %v", got, want)
|
||||
}
|
||||
for i := range want {
|
||||
if got[i] != want[i] {
|
||||
t.Errorf("fraction %d = %v, want %v", i, got[i], want[i])
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestFractionWindow_ClampsRatio(t *testing.T) {
|
||||
rec := &fractionRecorder{}
|
||||
ctx := rec.context(t, "Extractor:abc")
|
||||
w := fractionWindow{base: 0.5, span: 0.5}
|
||||
|
||||
w.report(ctx, -1)
|
||||
w.report(ctx, 2)
|
||||
|
||||
got := rec.fractions()
|
||||
if len(got) != 2 || got[0] != 0.5 || got[1] != 1 {
|
||||
t.Errorf("fractions = %v, want [0.5 1]", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAwaitFutures_ReportsEverySettledJob(t *testing.T) {
|
||||
rec := &fractionRecorder{}
|
||||
ctx := rec.context(t, "Extractor:abc")
|
||||
|
||||
futs := make([]utility.WorkerPoolFuture[extractorJob, struct{}], 0, 4)
|
||||
for i := 0; i < 4; i++ {
|
||||
f, err := extractorPool.Submit(ctx, func() error { return nil })
|
||||
if err != nil {
|
||||
t.Fatalf("submit: %v", err)
|
||||
}
|
||||
futs = append(futs, f)
|
||||
}
|
||||
|
||||
w := fractionWindow{base: 0.5, span: 0.5}
|
||||
if err := awaitFutures(ctx, futs, w); err != nil {
|
||||
t.Fatalf("awaitFutures: %v", err)
|
||||
}
|
||||
|
||||
got := rec.fractions()
|
||||
if len(got) != 4 {
|
||||
t.Fatalf("expected 4 reports, got %d: %v", len(got), got)
|
||||
}
|
||||
highest := got[0]
|
||||
for i, f := range got {
|
||||
if f < w.base || f > w.base+w.span {
|
||||
t.Errorf("report %d = %v, outside window [%v %v]", i, f, w.base, w.base+w.span)
|
||||
}
|
||||
highest = max(highest, f)
|
||||
}
|
||||
// Goroutines report concurrently, so only the highest value proves the
|
||||
// phase reached the end of its window.
|
||||
if highest != w.base+w.span {
|
||||
t.Errorf("highest report = %v, want %v", highest, w.base+w.span)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAwaitFutures_NoJobsNoReports(t *testing.T) {
|
||||
rec := &fractionRecorder{}
|
||||
ctx := rec.context(t, "Extractor:abc")
|
||||
|
||||
if err := awaitFutures(ctx, nil, fractionWindow{span: 1}); err != nil {
|
||||
t.Fatalf("awaitFutures: %v", err)
|
||||
}
|
||||
if got := rec.fractions(); len(got) != 0 {
|
||||
t.Errorf("expected no reports for an empty phase, got %v", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestEmbedChunks_ReportsBatchFractions(t *testing.T) {
|
||||
t.Setenv("TOKENIZER_EMBEDDING_BATCH_SIZE", "2")
|
||||
|
||||
comp, stub := withStubEmbedder(t, 4)
|
||||
comp.param.Fields = []string{"text"}
|
||||
|
||||
rec := &fractionRecorder{}
|
||||
ctx := rec.context(t, "Tokenizer:abc")
|
||||
chunks := []schema.ChunkDoc{
|
||||
chunkWithID("c1", "alpha"),
|
||||
chunkWithID("c2", "beta"),
|
||||
chunkWithID("c3", "gamma"),
|
||||
chunkWithID("c4", "delta"),
|
||||
}
|
||||
|
||||
if _, _, err := comp.embedChunks(ctx, "tenant", "kb", "", chunks, newMemCacheStore()); err != nil {
|
||||
t.Fatalf("embedChunks: %v", err)
|
||||
}
|
||||
if stub.calls.Load() != 2 {
|
||||
t.Fatalf("expected 2 embedding batches, got %d", stub.calls.Load())
|
||||
}
|
||||
|
||||
got := rec.fractions()
|
||||
want := []float64{0.5, 1}
|
||||
if len(got) != len(want) {
|
||||
t.Fatalf("fractions = %v, want %v", got, want)
|
||||
}
|
||||
for i := range want {
|
||||
if got[i] != want[i] {
|
||||
t.Errorf("fraction %d = %v, want %v", i, got[i], want[i])
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestEmbedChunks_CountsCacheHitsAsDone(t *testing.T) {
|
||||
t.Setenv("TOKENIZER_EMBEDDING_BATCH_SIZE", "1")
|
||||
|
||||
comp, _ := withStubEmbedder(t, 4)
|
||||
comp.param.Fields = []string{"text"}
|
||||
|
||||
warm := []schema.ChunkDoc{chunkWithID("c1", "alpha"), chunkWithID("c2", "beta")}
|
||||
store := newMemCacheStore()
|
||||
if _, _, err := comp.embedChunks(t.Context(), "tenant", "kb", "", warm, store); err != nil {
|
||||
t.Fatalf("warm embedChunks: %v", err)
|
||||
}
|
||||
|
||||
rec := &fractionRecorder{}
|
||||
ctx := rec.context(t, "Tokenizer:abc")
|
||||
chunks := append(append([]schema.ChunkDoc(nil), warm...),
|
||||
chunkWithID("c3", "gamma"), chunkWithID("c4", "delta"))
|
||||
|
||||
if _, _, err := comp.embedChunks(ctx, "tenant", "kb", "", chunks, store); err != nil {
|
||||
t.Fatalf("embedChunks: %v", err)
|
||||
}
|
||||
|
||||
// Two of the four chunks come from the cache, so the first freshly embedded
|
||||
// chunk already completes 3/4 of the work.
|
||||
got := rec.fractions()
|
||||
want := []float64{0.75, 1}
|
||||
if len(got) != len(want) {
|
||||
t.Fatalf("fractions = %v, want %v", got, want)
|
||||
}
|
||||
for i := range want {
|
||||
if got[i] != want[i] {
|
||||
t.Errorf("fraction %d = %v, want %v", i, got[i], want[i])
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestEmbedChunks_AllCacheHitsReportsCompletion covers the run where every
|
||||
// content embedding is a cache hit: the batch loop never executes, so the
|
||||
// phase has to report its completion separately or it stays silent.
|
||||
func TestEmbedChunks_AllCacheHitsReportsCompletion(t *testing.T) {
|
||||
t.Setenv("TOKENIZER_EMBEDDING_BATCH_SIZE", "1")
|
||||
|
||||
comp, _ := withStubEmbedder(t, 4)
|
||||
comp.param.Fields = []string{"text"}
|
||||
|
||||
store := newMemCacheStore()
|
||||
chunks := []schema.ChunkDoc{chunkWithID("c1", "alpha"), chunkWithID("c2", "beta")}
|
||||
if _, _, err := comp.embedChunks(t.Context(), "tenant", "kb", "", chunks, store); err != nil {
|
||||
t.Fatalf("warm embedChunks: %v", err)
|
||||
}
|
||||
|
||||
rec := &fractionRecorder{}
|
||||
ctx := rec.context(t, "Tokenizer:abc")
|
||||
if _, _, err := comp.embedChunks(ctx, "tenant", "kb", "", chunks, store); err != nil {
|
||||
t.Fatalf("embedChunks: %v", err)
|
||||
}
|
||||
|
||||
got := rec.fractions()
|
||||
want := []float64{1}
|
||||
if len(got) != len(want) || got[0] != want[0] {
|
||||
t.Fatalf("fractions = %v, want %v (all-cached work must report completion)", got, want)
|
||||
}
|
||||
}
|
||||
@@ -465,6 +465,9 @@ func (c *TokenizerComponent) embedChunks(ctx context.Context, tenantID, kbID, na
|
||||
// truncs[i] is the embedded text for texts[i] / pairs[i]; carried so the
|
||||
// Set-side key matches the Get-side key for freshly embedded content.
|
||||
truncs := make([]string, 0, len(chunks))
|
||||
// cacheHits counts the chunks served from the per-chunk embedding cache;
|
||||
// they are already-resolved work, so the progress fraction starts from them.
|
||||
cacheHits := 0
|
||||
for i, ck := range chunks {
|
||||
raw := concatFields(ck, c.param.Fields)
|
||||
txt := htmlTableRE.ReplaceAllString(raw, " ")
|
||||
@@ -487,6 +490,7 @@ func (c *TokenizerComponent) embedChunks(ctx context.Context, tenantID, kbID, na
|
||||
var vec []float64
|
||||
if err := json.Unmarshal([]byte(cached), &vec); err == nil && len(vec) > 0 {
|
||||
resolved[i] = &resolvedVec{content: vec, isHit: true, trunc: trunc}
|
||||
cacheHits++
|
||||
continue
|
||||
}
|
||||
}
|
||||
@@ -539,6 +543,7 @@ func (c *TokenizerComponent) embedChunks(ctx context.Context, tenantID, kbID, na
|
||||
if batchSize <= 0 {
|
||||
return nil, 0, fmt.Errorf("tokenizer: embedder reported non-positive batch size %d", batchSize)
|
||||
}
|
||||
embedTotal := cacheHits + len(texts)
|
||||
for start := 0; start < len(texts); start += batchSize {
|
||||
end := start + batchSize
|
||||
if end > len(texts) {
|
||||
@@ -555,6 +560,13 @@ func (c *TokenizerComponent) embedChunks(ctx context.Context, tenantID, kbID, na
|
||||
tokenCount += result.TokenCount
|
||||
}
|
||||
contentResults = append(contentResults, batchResults...)
|
||||
runtime.ReportComponentFraction(ctx, float64(cacheHits+end)/float64(embedTotal))
|
||||
}
|
||||
if len(texts) == 0 && embedTotal > 0 {
|
||||
// With no cache misses the batch loop never runs, so the phase would
|
||||
// otherwise report nothing at all; every content embedding is already
|
||||
// resolved, which is the loop's completion value.
|
||||
runtime.ReportComponentFraction(ctx, 1)
|
||||
}
|
||||
|
||||
titleWeight := c.param.FilenameEmbdWeight
|
||||
|
||||
@@ -239,6 +239,7 @@ func (p *Pipeline) Run(ctx context.Context, inputs map[string]any, overrideParam
|
||||
// observability concern, not a data dependency.
|
||||
runCtx = runtime.WithProgressCallback(runCtx, p.componentProgressCallback(ctx))
|
||||
runCtx = runtime.WithProgressMessageCallback(runCtx, p.componentProgressMessageCallback(ctx))
|
||||
runCtx = runtime.WithProgressFractionCallback(runCtx, p.componentFractionCallback(ctx))
|
||||
|
||||
current := cloneMapOrEmpty(inputs)
|
||||
|
||||
@@ -392,6 +393,28 @@ type detailedProgressSink interface {
|
||||
OnComponentMessage(ctx context.Context, taskID, documentID, component, message string)
|
||||
}
|
||||
|
||||
// fractionProgressSink is the optional interface through which the pipeline
|
||||
// forwards in-flight component fractions (pages parsed, chunks embedded) to
|
||||
// the sink's progress tracker. Mirrors detailedProgressSink: a sink that does
|
||||
// not implement it simply receives no fraction channel.
|
||||
type fractionProgressSink interface {
|
||||
OnComponentFraction(ctx context.Context, component string, fraction float64)
|
||||
}
|
||||
|
||||
// componentFractionCallback forwards fraction reports to the sink. The sink
|
||||
// only mutates in-memory state here (its flusher owns persistence), so the
|
||||
// run context is passed through without the WithoutCancel fallback that the
|
||||
// I/O-bound callbacks need.
|
||||
func (p *Pipeline) componentFractionCallback(ctx context.Context) runtime.ProgressFractionCallback {
|
||||
sink, ok := p.sink.(fractionProgressSink)
|
||||
if !ok {
|
||||
return nil
|
||||
}
|
||||
return func(component string, fraction float64) {
|
||||
sink.OnComponentFraction(ctx, component, fraction)
|
||||
}
|
||||
}
|
||||
|
||||
func (p *Pipeline) componentProgressMessageCallback(ctx context.Context) runtime.ProgressMessageCallback {
|
||||
sink, ok := p.sink.(detailedProgressSink)
|
||||
if !ok {
|
||||
|
||||
@@ -446,3 +446,95 @@ func TestValidatePipeline_DisallowMultipleExtractors(t *testing.T) {
|
||||
t.Errorf("expected at most 1 Extractor error message, got: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// fractionStage reports a fixed in-flight fraction from Invoke, the way a
|
||||
// parser reports pages or a tokenizer reports embedded chunks.
|
||||
type fractionStage struct {
|
||||
fraction float64
|
||||
}
|
||||
|
||||
func (f *fractionStage) Invoke(ctx context.Context, _ *gorm.DB, inputs map[string]any) (map[string]any, error) {
|
||||
runtime.ReportComponentFraction(ctx, f.fraction)
|
||||
return cloneMapOrEmpty(inputs), nil
|
||||
}
|
||||
func (f *fractionStage) Inputs() map[string]string { return map[string]string{"name": "string"} }
|
||||
func (f *fractionStage) Outputs() map[string]string { return map[string]string{"output": "any"} }
|
||||
|
||||
// fractionRecordingSink records fraction reports alongside the base
|
||||
// ProgressSink methods, exercising the optional-interface assertion.
|
||||
type fractionRecordingSink struct {
|
||||
recordingSink
|
||||
muFractions sync.Mutex
|
||||
fractions map[string]float64
|
||||
}
|
||||
|
||||
func (r *fractionRecordingSink) OnComponentFraction(_ context.Context, component string, fraction float64) {
|
||||
r.muFractions.Lock()
|
||||
defer r.muFractions.Unlock()
|
||||
if r.fractions == nil {
|
||||
r.fractions = map[string]float64{}
|
||||
}
|
||||
r.fractions[component] = fraction
|
||||
}
|
||||
|
||||
// TestPipelineRunForwardsFractionsToSink verifies the pipeline assembles the
|
||||
// fraction channel end to end: a component's bare ReportComponentFraction
|
||||
// reaches the sink's optional OnComponentFraction under the node's cpnID.
|
||||
func TestPipelineRunForwardsFractionsToSink(t *testing.T) {
|
||||
stage := &fractionStage{fraction: 0.42}
|
||||
const name = "p.FractionStage"
|
||||
runtime.MustRegister(name, runtime.CategoryIngestion,
|
||||
func(_ string, _ map[string]any) (runtime.Component, error) { return stage, nil },
|
||||
runtime.Metadata{Version: "1.0.0"})
|
||||
|
||||
sink := &fractionRecordingSink{}
|
||||
pipe, err := NewPipelineFromDSL([]byte(`{
|
||||
"dsl": {
|
||||
"components": {
|
||||
"begin": {"obj": {"component_name": "Begin", "params": {}}, "downstream": ["a"]},
|
||||
"a": {"obj": {"component_name": "`+name+`", "params": {}}, "upstream": ["begin"]}
|
||||
},
|
||||
"path": ["begin", "a"],
|
||||
"graph": {"nodes": []}
|
||||
}
|
||||
}`), "task-frac", WithProgressSink(sink), WithDocumentID("doc-frac"))
|
||||
if err != nil {
|
||||
t.Fatalf("NewPipelineFromDSL: %v", err)
|
||||
}
|
||||
if _, err := pipe.Run(t.Context(), map[string]any{"name": "doc-frac"}, nil); err != nil {
|
||||
t.Fatalf("Run: %v", err)
|
||||
}
|
||||
|
||||
sink.muFractions.Lock()
|
||||
defer sink.muFractions.Unlock()
|
||||
if got := sink.fractions["a"]; got != 0.42 {
|
||||
t.Fatalf("fraction for node a = %v, want 0.42 (fractions=%v)", got, sink.fractions)
|
||||
}
|
||||
}
|
||||
|
||||
// TestPipelineRunWithoutFractionSink verifies a sink that does not implement
|
||||
// the optional fraction interface still runs: the channel is simply absent.
|
||||
func TestPipelineRunWithoutFractionSink(t *testing.T) {
|
||||
stage := &fractionStage{fraction: 0.9}
|
||||
const name = "p.FractionStageNoSink"
|
||||
runtime.MustRegister(name, runtime.CategoryIngestion,
|
||||
func(_ string, _ map[string]any) (runtime.Component, error) { return stage, nil },
|
||||
runtime.Metadata{Version: "1.0.0"})
|
||||
|
||||
pipe, err := NewPipelineFromDSL([]byte(`{
|
||||
"dsl": {
|
||||
"components": {
|
||||
"begin": {"obj": {"component_name": "Begin", "params": {}}, "downstream": ["a"]},
|
||||
"a": {"obj": {"component_name": "`+name+`", "params": {}}, "upstream": ["begin"]}
|
||||
},
|
||||
"path": ["begin", "a"],
|
||||
"graph": {"nodes": []}
|
||||
}
|
||||
}`), "task-nofrac", WithProgressSink(&recordingSink{}))
|
||||
if err != nil {
|
||||
t.Fatalf("NewPipelineFromDSL: %v", err)
|
||||
}
|
||||
if _, err := pipe.Run(t.Context(), map[string]any{"name": "x"}, nil); err != nil {
|
||||
t.Fatalf("Run: %v", err)
|
||||
}
|
||||
}
|
||||
@@ -1195,7 +1195,14 @@ func (e *Ingestor) defaultRunDocumentTask(ctx context.Context, ingestionTask *en
|
||||
if ingestionTask.PipelineLogID != nil {
|
||||
pipelineLogID = *ingestionTask.PipelineLogID
|
||||
}
|
||||
result, err := executor.WithProgressSink(newProgressSink(ctx, e.ingestionTaskSvc, pipelineLogID)).Execute(docTaskCtx.Ctx)
|
||||
sink := newProgressSink(ctx, e.ingestionTaskSvc, pipelineLogID)
|
||||
// Close is idempotent and intentionally called twice: the defer is the
|
||||
// panic backstop (an unwind out of Execute must still stop the flusher),
|
||||
// and the explicit call pins the final flush ahead of docState.apply,
|
||||
// whose terminal process_duration the flush would otherwise overwrite.
|
||||
defer sink.Close()
|
||||
result, err := executor.WithProgressSink(sink).Execute(docTaskCtx.Ctx)
|
||||
sink.Close()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
@@ -19,28 +19,49 @@ package service
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"sync/atomic"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
agentruntime "ragflow/internal/agent/runtime"
|
||||
"ragflow/internal/common"
|
||||
"ragflow/internal/dao"
|
||||
"ragflow/internal/ingestion/pipeline"
|
||||
servicepkg "ragflow/internal/service"
|
||||
documentpkg "ragflow/internal/service/document"
|
||||
)
|
||||
|
||||
// progressFlushInterval bounds how often the flusher persists the in-memory
|
||||
// runProgress to the document row. Lifecycle events and fraction reports only
|
||||
// mutate memory; the ticker coalesces them into at most one UPDATE per
|
||||
// interval regardless of event rate, so this is the worst-case lag between a
|
||||
// component boundary and its mirrored percent.
|
||||
const progressFlushInterval = time.Second
|
||||
|
||||
// progressSink implements pipeline.ProgressSink. It records a run's component
|
||||
// events and total through the service layer, and mirrors only the numeric
|
||||
// progress state to the owning document. All writes are best-effort so a
|
||||
// events and total through the service layer, accumulates completion percent
|
||||
// in memory (runProgress), and mirrors that percent into the owning document
|
||||
// row from a single flusher goroutine. All writes are best-effort so a
|
||||
// reporting failure never aborts the pipeline run.
|
||||
type progressSink struct {
|
||||
taskSvc *servicepkg.IngestionTaskService
|
||||
docSvc docProgressSvc
|
||||
pipelineLogID string
|
||||
// total is the component-count denominator cached from OnComponentTotal.
|
||||
// It is Store-d once in the Run goroutine and Load-ed by OnComponentProgress,
|
||||
// which eino fires from concurrent parallel-branch goroutines. Atomic because
|
||||
// the two access paths share no other synchronization.
|
||||
total atomic.Int64
|
||||
baseCtx context.Context
|
||||
|
||||
progress *runProgress
|
||||
|
||||
// mu guards docID, which eino's parallel-branch callbacks bind
|
||||
// concurrently from lifecycle events.
|
||||
mu sync.Mutex
|
||||
docID string
|
||||
|
||||
closed chan struct{}
|
||||
closeOnce sync.Once
|
||||
wg sync.WaitGroup
|
||||
|
||||
// lastFlushed dedupes ticker writes while percent is unchanged. Only the
|
||||
// flusher goroutine touches it, and Close re-flushes after wg.Wait, so no
|
||||
// additional synchronization is needed.
|
||||
lastFlushed float64
|
||||
}
|
||||
|
||||
// docProgressSvc is the subset of *service.DocumentService the sink needs to
|
||||
@@ -57,15 +78,79 @@ func newProgressSink(ctx context.Context, taskSvc *servicepkg.IngestionTaskServi
|
||||
// OnComponentProgress (and thus docSvc) can fire from multiple goroutines;
|
||||
// a lazy check-then-act here would be a data race. The sink owns no
|
||||
// server-config dependency, so this is safe in any environment.
|
||||
return &progressSink{
|
||||
s := &progressSink{
|
||||
taskSvc: taskSvc,
|
||||
docSvc: documentpkg.NewDocumentService(),
|
||||
pipelineLogID: pipelineLogID,
|
||||
baseCtx: ctx,
|
||||
progress: newRunProgress(),
|
||||
closed: make(chan struct{}),
|
||||
}
|
||||
s.wg.Add(1)
|
||||
go s.runFlusher(ctx)
|
||||
return s
|
||||
}
|
||||
|
||||
// Close stops the flusher and performs one final forced flush so the last
|
||||
// in-memory percent reaches the document row before the caller writes the
|
||||
// terminal progress state. It is idempotent and safe to call from any
|
||||
// goroutine. The final flush detaches from the (possibly cancelled) run
|
||||
// context; callers must invoke it before returning from the run.
|
||||
func (s *progressSink) Close() {
|
||||
s.closeOnce.Do(func() {
|
||||
close(s.closed)
|
||||
s.wg.Wait()
|
||||
ctx, cancel := context.WithTimeout(context.WithoutCancel(s.baseCtx), 5*time.Second)
|
||||
defer cancel()
|
||||
s.flush(ctx, true)
|
||||
})
|
||||
}
|
||||
|
||||
func (s *progressSink) runFlusher(ctx context.Context) {
|
||||
defer s.wg.Done()
|
||||
ticker := time.NewTicker(progressFlushInterval)
|
||||
defer ticker.Stop()
|
||||
for {
|
||||
select {
|
||||
case <-s.closed:
|
||||
return
|
||||
case <-ticker.C:
|
||||
s.flush(ctx, false)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// flush persists the current percent to the bound document. Ticker flushes
|
||||
// skip when percent is unchanged since the last successful write; the final
|
||||
// flush on Close is forced.
|
||||
func (s *progressSink) flush(ctx context.Context, force bool) {
|
||||
s.mu.Lock()
|
||||
docID := s.docID
|
||||
s.mu.Unlock()
|
||||
if docID == "" {
|
||||
return
|
||||
}
|
||||
p := s.progress.Percent()
|
||||
if !force && p == s.lastFlushed {
|
||||
return
|
||||
}
|
||||
if err := s.docSvc.UpdateRunState(ctx, docID, p); err != nil {
|
||||
common.Warn(fmt.Sprintf("progressSink: flush progress for document %s failed: %v", docID, err))
|
||||
return
|
||||
}
|
||||
s.lastFlushed = p
|
||||
}
|
||||
|
||||
func (s *progressSink) bindDocument(docID string) {
|
||||
s.mu.Lock()
|
||||
if s.docID == "" {
|
||||
s.docID = docID
|
||||
}
|
||||
s.mu.Unlock()
|
||||
}
|
||||
|
||||
func (s *progressSink) OnComponentTotal(ctx context.Context, taskID string, total int) {
|
||||
s.total.Store(int64(total))
|
||||
s.progress.SetTotal(total)
|
||||
if err := s.taskSvc.UpdateComponentTotal(ctx, taskID, total); err != nil {
|
||||
common.Error(fmt.Sprintf("progressSink: update component_total for task %s failed: %v", taskID, err), err)
|
||||
}
|
||||
@@ -75,24 +160,23 @@ func (s *progressSink) OnComponentProgress(ctx context.Context, ev pipeline.Prog
|
||||
if err := s.taskSvc.RecordLifecycle(ctx, s.pipelineLogID, ev.TaskID, ev.Component, ev.Phase, ev.Message); err != nil {
|
||||
common.Error(fmt.Sprintf("progressSink: record component progress for task %s failed: %v", ev.TaskID, err), err)
|
||||
}
|
||||
if ev.DocumentID == "" {
|
||||
return
|
||||
if ev.Phase == int(agentruntime.PhaseExit) {
|
||||
s.progress.MarkDone(ev.Component)
|
||||
}
|
||||
total := s.total.Load()
|
||||
agg, err := s.taskSvc.AggregateTaskProgressByPipelineLogID(ctx, s.pipelineLogID, int(total))
|
||||
if err != nil {
|
||||
common.Error(fmt.Sprintf("progressSink: aggregate task progress for task %s failed: %v", ev.TaskID, err), err)
|
||||
return
|
||||
}
|
||||
if agg == nil || total <= 0 {
|
||||
return
|
||||
}
|
||||
progress := deriveDocumentProgress(agg, int(total))
|
||||
if err = s.docSvc.UpdateRunState(ctx, ev.DocumentID, progress); err != nil {
|
||||
common.Error(fmt.Sprintf("progressSink: update progress state for document %s task %s failed: %v", ev.DocumentID, ev.TaskID, err), err)
|
||||
if ev.DocumentID != "" {
|
||||
s.bindDocument(ev.DocumentID)
|
||||
}
|
||||
}
|
||||
|
||||
// OnComponentFraction records an in-flight component's 0..1 completion
|
||||
// fraction (pages parsed, chunks embedded). It deliberately does not flush
|
||||
// directly: fraction reports are high-frequency and the ticker coalesces them.
|
||||
// The pipeline reaches this method through an optional-interface assertion,
|
||||
// mirroring detailedProgressSink.
|
||||
func (s *progressSink) OnComponentFraction(_ context.Context, component string, frac float64) {
|
||||
s.progress.SetFrac(component, frac)
|
||||
}
|
||||
|
||||
// OnComponentMessage records detailed compiler-stage information without
|
||||
// creating a lifecycle row, so completion percent remains based on actual
|
||||
// canvas components.
|
||||
@@ -104,13 +188,3 @@ func (s *progressSink) OnComponentMessage(ctx context.Context, taskID, _ string,
|
||||
common.Error(fmt.Sprintf("progressSink: record message for task %s failed", taskID), err)
|
||||
}
|
||||
}
|
||||
|
||||
// deriveDocumentProgress computes the document-level progress (0..1) from the
|
||||
// aggregated ingestion_task_log.
|
||||
func deriveDocumentProgress(agg *dao.TaskProgress, total int) float64 {
|
||||
progress := agg.Percent / 100
|
||||
if progress > 1 {
|
||||
progress = 1
|
||||
}
|
||||
return progress
|
||||
}
|
||||
@@ -21,6 +21,7 @@ import (
|
||||
"runtime"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"ragflow/internal/dao"
|
||||
"ragflow/internal/ingestion/pipeline"
|
||||
@@ -49,8 +50,8 @@ func TestProgressSink_CanConstructDocumentServiceWithoutServerConfig(t *testing.
|
||||
// TestProgressSink_EagerlyConstructsDocumentService ensures the sink builds its
|
||||
// DocumentService at construction time rather than lazily on the first progress
|
||||
// event. Lazy construction is a data race under eino's parallel-branch progress
|
||||
// callbacks (see TestProgressSink_OnComponentProgress_NoDataRace); eager
|
||||
// construction makes docSvc immutable after newProgressSink returns.
|
||||
// callbacks (see TestProgressSink_DocService_NoDataRace); eager construction
|
||||
// makes docSvc immutable after newProgressSink returns.
|
||||
func TestProgressSink_EagerlyConstructsDocumentService(t *testing.T) {
|
||||
db := testutil.SetupTestDB(t)
|
||||
cleanup := testutil.ReplaceDBForTest(t, db)
|
||||
@@ -58,6 +59,7 @@ func TestProgressSink_EagerlyConstructsDocumentService(t *testing.T) {
|
||||
|
||||
ctx := t.Context()
|
||||
sink := newProgressSink(ctx, servicepkg.NewIngestionTaskService(), "run-1")
|
||||
defer sink.Close()
|
||||
if sink.docSvc == nil {
|
||||
t.Fatal("expected sink to eagerly construct its DocumentService, got nil (lazy)")
|
||||
}
|
||||
@@ -82,6 +84,7 @@ func TestProgressSink_DocService_NoDataRace(t *testing.T) {
|
||||
// must already be constructed (not lazily built mid-call) when the
|
||||
// goroutines below race into docSvc.
|
||||
sink := newProgressSink(ctx, servicepkg.NewIngestionTaskService(), "run-1")
|
||||
defer sink.Close()
|
||||
|
||||
const n = 30
|
||||
var wg sync.WaitGroup
|
||||
@@ -98,12 +101,10 @@ func TestProgressSink_DocService_NoDataRace(t *testing.T) {
|
||||
wg.Wait()
|
||||
}
|
||||
|
||||
// TestProgressSink_Total_NoDataRace guards the total denominator against being
|
||||
// a non-atomic shared field. OnComponentTotal (writer, Run goroutine) and
|
||||
// OnComponentProgress (reader, concurrent eino branches) share total; a plain
|
||||
// int is a data race per the Go memory model. The read is hit directly on the
|
||||
// field rather than through OnComponentProgress because the latter serializes
|
||||
// on the single test-DB connection before reaching the read, masking the race.
|
||||
// TestProgressSink_Total_NoDataRace guards the shared run-progress state
|
||||
// against unsynchronized access. OnComponentTotal (writer, Run goroutine) and
|
||||
// the flusher/Percent readers (concurrent eino branches, flusher goroutine)
|
||||
// share the tracker; its mutex must make concurrent SetTotal/Percent safe.
|
||||
func TestProgressSink_Total_NoDataRace(t *testing.T) {
|
||||
db := testutil.SetupTestDB(t)
|
||||
cleanup := testutil.ReplaceDBForTest(t, db)
|
||||
@@ -112,6 +113,7 @@ func TestProgressSink_Total_NoDataRace(t *testing.T) {
|
||||
|
||||
ctx := t.Context()
|
||||
sink := newProgressSink(ctx, servicepkg.NewIngestionTaskService(), "run-1")
|
||||
defer sink.Close()
|
||||
|
||||
const n = 30
|
||||
var wg sync.WaitGroup
|
||||
@@ -121,7 +123,7 @@ func TestProgressSink_Total_NoDataRace(t *testing.T) {
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
<-start
|
||||
sink.OnComponentTotal(ctx, taskID, 5) // writes s.total
|
||||
sink.OnComponentTotal(ctx, taskID, 5) // writes tracker total
|
||||
}()
|
||||
}
|
||||
for i := 0; i < n; i++ {
|
||||
@@ -129,7 +131,7 @@ func TestProgressSink_Total_NoDataRace(t *testing.T) {
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
<-start
|
||||
v := sink.total.Load() // reads s.total atomically
|
||||
v := sink.progress.Percent() // reads tracker under mutex
|
||||
runtime.KeepAlive(v)
|
||||
}()
|
||||
}
|
||||
@@ -138,18 +140,42 @@ func TestProgressSink_Total_NoDataRace(t *testing.T) {
|
||||
}
|
||||
|
||||
type stubDocProgressSvc struct {
|
||||
mu sync.Mutex
|
||||
stateCalls int
|
||||
stateDocID string
|
||||
stateProgress float64
|
||||
}
|
||||
|
||||
func (s *stubDocProgressSvc) UpdateRunState(_ context.Context, docID string, progress float64) error {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
s.stateCalls++
|
||||
s.stateDocID = docID
|
||||
s.stateProgress = progress
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *stubDocProgressSvc) snapshot() (calls int, docID string, progress float64) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
return s.stateCalls, s.stateDocID, s.stateProgress
|
||||
}
|
||||
|
||||
// ctxAwareStubDocProgressSvc fails like a real DB call when the flush context
|
||||
// is cancelled, so tests can prove the Close final flush (detached context)
|
||||
// still succeeds after the run context is cancelled, while ticker flushes
|
||||
// bound to the run context would fail.
|
||||
type ctxAwareStubDocProgressSvc struct {
|
||||
stubDocProgressSvc
|
||||
}
|
||||
|
||||
func (s *ctxAwareStubDocProgressSvc) UpdateRunState(ctx context.Context, docID string, progress float64) error {
|
||||
if err := ctx.Err(); err != nil {
|
||||
return err
|
||||
}
|
||||
return s.stubDocProgressSvc.UpdateRunState(ctx, docID, progress)
|
||||
}
|
||||
|
||||
// TestProgressSinkPersistsViaService verifies the sink is the single writer of
|
||||
// ingestion_task.component_total, ingestion_task_log, and document run-progress
|
||||
// - all through the service layer, not the DAO.
|
||||
@@ -180,6 +206,9 @@ func TestProgressSinkPersistsViaService(t *testing.T) {
|
||||
Phase: 1,
|
||||
Message: "Parser Done",
|
||||
})
|
||||
// Close joins the flusher and performs the final forced flush, making the
|
||||
// mirrored state deterministic regardless of ticker timing.
|
||||
sink.Close()
|
||||
|
||||
logs, err := dao.NewIngestionTaskLogDAO().ListLogsByPipelineLogID(ctx, db, "run-1")
|
||||
if err != nil {
|
||||
@@ -192,10 +221,12 @@ func TestProgressSinkPersistsViaService(t *testing.T) {
|
||||
t.Fatalf("unexpected log row: %+v", logs[0])
|
||||
}
|
||||
|
||||
// 1 of 2 components done -> RUNNING (run "1"), progress 0.5. The
|
||||
// event stream owns text, so this must not write document.progress_msg.
|
||||
if stub.stateCalls != 1 || stub.stateDocID != docID || stub.stateProgress != 0.5 {
|
||||
t.Fatalf("UpdateRunState = calls:%d doc:%q progress:%v, want 1/%q/0.5", stub.stateCalls, stub.stateDocID, stub.stateProgress, docID)
|
||||
// 1 of 2 components done -> progress 0.5 mirrored from the in-memory
|
||||
// tracker. The event stream owns text, so this must not write
|
||||
// document.progress_msg (UpdateRunState only touches progress/duration).
|
||||
calls, gotDocID, gotProgress := stub.snapshot()
|
||||
if calls < 1 || gotDocID != docID || gotProgress != 0.5 {
|
||||
t.Fatalf("UpdateRunState = calls:%d doc:%q progress:%v, want >=1/%q/0.5", calls, gotDocID, gotProgress, docID)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -206,6 +237,7 @@ func TestProgressSinkWritesLifecycleEventForCapturedRun(t *testing.T) {
|
||||
_, _, _, taskID := testutil.SeedTestData(t, db, testutil.WithPipelineID("flow-1"))
|
||||
|
||||
sink := newProgressSink(t.Context(), servicepkg.NewIngestionTaskService(), "run-1")
|
||||
defer sink.Close()
|
||||
sink.OnComponentProgress(t.Context(), pipeline.ProgressEvent{
|
||||
TaskID: taskID,
|
||||
Component: "Parser",
|
||||
@@ -230,7 +262,8 @@ func TestProgressSinkWritesLifecycleEventForCapturedRun(t *testing.T) {
|
||||
}
|
||||
|
||||
// TestProgressSinkEmptyDocumentIDSkipsMirror verifies the log row is still
|
||||
// recorded when no owning document is bound, but the document mirror is skipped.
|
||||
// recorded when no owning document is bound, but the document mirror is
|
||||
// skipped - including the Close final flush.
|
||||
func TestProgressSinkEmptyDocumentIDSkipsMirror(t *testing.T) {
|
||||
db := testutil.SetupTestDB(t)
|
||||
cleanup := testutil.ReplaceDBForTest(t, db)
|
||||
@@ -248,6 +281,7 @@ func TestProgressSinkEmptyDocumentIDSkipsMirror(t *testing.T) {
|
||||
Phase: 1,
|
||||
Message: "Chunker Done",
|
||||
})
|
||||
sink.Close()
|
||||
|
||||
logs, err := dao.NewIngestionTaskLogDAO().ListLogsByPipelineLogID(ctx, db, "run-1")
|
||||
if err != nil {
|
||||
@@ -256,71 +290,121 @@ func TestProgressSinkEmptyDocumentIDSkipsMirror(t *testing.T) {
|
||||
if len(logs) != 1 {
|
||||
t.Fatalf("expected 1 component-progress row, got %d", len(logs))
|
||||
}
|
||||
if stub.stateCalls != 0 {
|
||||
t.Fatalf("UpdateRunState calls = %d, want 0 (no document bound)", stub.stateCalls)
|
||||
if calls, _, _ := stub.snapshot(); calls != 0 {
|
||||
t.Fatalf("UpdateRunState calls = %d, want 0 (no document bound)", calls)
|
||||
}
|
||||
}
|
||||
|
||||
// TestDeriveDocumentProgress exercises every branch of the run-label derivation
|
||||
// logic. The function is called from OnComponentProgress with a non-nil agg
|
||||
// (guarded by the caller), so the nil case is documented as a known panic.
|
||||
func TestDeriveDocumentProgress(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
agg *dao.TaskProgress
|
||||
total int
|
||||
wantProg float64
|
||||
}{
|
||||
{
|
||||
name: "failed component progress",
|
||||
agg: &dao.TaskProgress{Failed: 1, Done: 0, Running: 0, Percent: 0},
|
||||
total: 5,
|
||||
wantProg: 0.0,
|
||||
},
|
||||
{
|
||||
name: "all done",
|
||||
agg: &dao.TaskProgress{Failed: 0, Done: 5, Running: 0, Percent: 100},
|
||||
total: 5,
|
||||
wantProg: 1.0,
|
||||
},
|
||||
{
|
||||
name: "partial done",
|
||||
agg: &dao.TaskProgress{Failed: 0, Done: 3, Running: 0, Percent: 60},
|
||||
total: 5,
|
||||
wantProg: 0.6,
|
||||
},
|
||||
{
|
||||
name: "running only",
|
||||
agg: &dao.TaskProgress{Failed: 0, Done: 0, Running: 2, Percent: 0},
|
||||
total: 5,
|
||||
wantProg: 0.0,
|
||||
},
|
||||
{
|
||||
name: "nothing started",
|
||||
agg: &dao.TaskProgress{Failed: 0, Done: 0, Running: 0, Percent: 0},
|
||||
total: 5,
|
||||
wantProg: 0.0,
|
||||
},
|
||||
{
|
||||
name: "total zero, nothing done",
|
||||
agg: &dao.TaskProgress{Failed: 0, Done: 0, Running: 0, Percent: 0},
|
||||
total: 0,
|
||||
wantProg: 0.0,
|
||||
},
|
||||
{
|
||||
name: "failed with 100 percent",
|
||||
agg: &dao.TaskProgress{Failed: 1, Done: 5, Running: 0, Percent: 100},
|
||||
total: 5,
|
||||
wantProg: 1.0,
|
||||
},
|
||||
}
|
||||
// TestProgressSinkFractionsCoalesceIntoCloseFlush verifies high-frequency
|
||||
// fraction reports never write per event: they only mutate the tracker, and
|
||||
// the ticker plus the final flush on Close are the only writers.
|
||||
func TestProgressSinkFractionsCoalesceIntoCloseFlush(t *testing.T) {
|
||||
db := testutil.SetupTestDB(t)
|
||||
cleanup := testutil.ReplaceDBForTest(t, db)
|
||||
defer cleanup()
|
||||
_, _, docID, taskID := testutil.SeedTestData(t, db, testutil.WithPipelineID("flow-1"))
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
prog := deriveDocumentProgress(tt.agg, tt.total)
|
||||
if prog != tt.wantProg {
|
||||
t.Errorf("progress = %v, want %v", prog, tt.wantProg)
|
||||
ctx := t.Context()
|
||||
sink := newProgressSink(ctx, servicepkg.NewIngestionTaskService(), "run-1")
|
||||
stub := &stubDocProgressSvc{}
|
||||
sink.docSvc = stub
|
||||
|
||||
sink.OnComponentTotal(ctx, taskID, 2)
|
||||
// Bind the document without a lifecycle exit (enter phase, no MarkDone).
|
||||
sink.OnComponentProgress(ctx, pipeline.ProgressEvent{
|
||||
TaskID: taskID,
|
||||
DocumentID: docID,
|
||||
Component: "Parser",
|
||||
Phase: 0,
|
||||
Message: "Parser Started",
|
||||
})
|
||||
for i := 1; i <= 100; i++ {
|
||||
sink.OnComponentFraction(ctx, "Parser", float64(i)/100)
|
||||
}
|
||||
sink.Close()
|
||||
|
||||
// Ticker (at most one write, since a single value change is what the
|
||||
// ticker would see) + Close final (1) = at most 2 writes for 100 reports.
|
||||
calls, gotDocID, gotProgress := stub.snapshot()
|
||||
if calls > 2 {
|
||||
t.Fatalf("UpdateRunState calls = %d, want <= 2 (fractions must coalesce)", calls)
|
||||
}
|
||||
if gotDocID != docID {
|
||||
t.Fatalf("docID = %q, want %q", gotDocID, docID)
|
||||
}
|
||||
if gotProgress != 0.5 {
|
||||
t.Fatalf("progress = %v, want 0.5 (frac 1.0 of 2 components)", gotProgress)
|
||||
}
|
||||
}
|
||||
|
||||
// TestProgressSinkTickerFlushesLifecycleProgressWithoutClose verifies lifecycle
|
||||
// progress reaches the document row during the run instead of waiting for
|
||||
// Close: the flusher's ticker mirrors a changed percent within one interval.
|
||||
func TestProgressSinkTickerFlushesLifecycleProgressWithoutClose(t *testing.T) {
|
||||
db := testutil.SetupTestDB(t)
|
||||
cleanup := testutil.ReplaceDBForTest(t, db)
|
||||
defer cleanup()
|
||||
_, _, docID, taskID := testutil.SeedTestData(t, db, testutil.WithPipelineID("flow-1"))
|
||||
|
||||
ctx := t.Context()
|
||||
sink := newProgressSink(ctx, servicepkg.NewIngestionTaskService(), "run-1")
|
||||
defer sink.Close()
|
||||
stub := &stubDocProgressSvc{}
|
||||
sink.docSvc = stub
|
||||
|
||||
sink.OnComponentTotal(ctx, taskID, 4)
|
||||
sink.OnComponentProgress(ctx, pipeline.ProgressEvent{
|
||||
TaskID: taskID,
|
||||
DocumentID: docID,
|
||||
Component: "File",
|
||||
Phase: 1,
|
||||
Message: "File Done",
|
||||
})
|
||||
|
||||
deadline := time.Now().Add(3 * time.Second)
|
||||
for {
|
||||
if calls, gotDocID, gotProgress := stub.snapshot(); calls >= 1 {
|
||||
if gotDocID != docID || gotProgress != 0.25 {
|
||||
t.Fatalf("ticker flush = doc:%q progress:%v, want %q/0.25", gotDocID, gotProgress, docID)
|
||||
}
|
||||
})
|
||||
return
|
||||
}
|
||||
if time.Now().After(deadline) {
|
||||
t.Fatal("ticker did not flush within 3s")
|
||||
}
|
||||
time.Sleep(10 * time.Millisecond)
|
||||
}
|
||||
}
|
||||
|
||||
// TestProgressSinkCloseFlushSurvivesCancelledRunContext verifies the final
|
||||
// flush detaches from the run context: a stopped/cancelled run still mirrors
|
||||
// its last in-memory percent instead of losing it to ctx.Err().
|
||||
func TestProgressSinkCloseFlushSurvivesCancelledRunContext(t *testing.T) {
|
||||
db := testutil.SetupTestDB(t)
|
||||
cleanup := testutil.ReplaceDBForTest(t, db)
|
||||
defer cleanup()
|
||||
_, _, docID, taskID := testutil.SeedTestData(t, db, testutil.WithPipelineID("flow-1"))
|
||||
|
||||
baseCtx, cancel := context.WithCancel(t.Context())
|
||||
cancel()
|
||||
sink := newProgressSink(baseCtx, servicepkg.NewIngestionTaskService(), "run-1")
|
||||
stub := &ctxAwareStubDocProgressSvc{}
|
||||
sink.docSvc = stub
|
||||
|
||||
// Events run on a live context (the pipeline wraps sink calls with
|
||||
// progressSinkContext), but the flusher's base context is dead.
|
||||
sink.OnComponentTotal(t.Context(), taskID, 2)
|
||||
sink.OnComponentProgress(t.Context(), pipeline.ProgressEvent{
|
||||
TaskID: taskID,
|
||||
DocumentID: docID,
|
||||
Component: "Parser",
|
||||
Phase: 1,
|
||||
Message: "Parser Done",
|
||||
})
|
||||
sink.Close()
|
||||
|
||||
calls, gotDocID, gotProgress := stub.snapshot()
|
||||
if calls != 1 || gotDocID != docID || gotProgress != 0.5 {
|
||||
t.Fatalf("final flush = calls:%d doc:%q progress:%v, want 1/%q/0.5 (ticker flushes must fail on cancelled ctx, Close flush must succeed)", calls, gotDocID, gotProgress, docID)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,115 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
package service
|
||||
|
||||
import (
|
||||
"math"
|
||||
"sync"
|
||||
)
|
||||
|
||||
// runProgress accumulates one pipeline run's component progress in memory.
|
||||
// Percent is (completed components + sum of in-flight fractions) / total,
|
||||
// guarded to never decrease: the flusher persists it as document.progress,
|
||||
// and a regressing progress bar reads as a restart even though the run is
|
||||
// advancing. Components are keyed by canvas component id, the identity both
|
||||
// lifecycle events and fraction reports carry.
|
||||
//
|
||||
// The tracker is born empty on every run: the pipeline has no partial-resume
|
||||
// entry point, so an MQ redelivery replays the whole graph and there is no
|
||||
// prior state worth seeding from.
|
||||
type runProgress struct {
|
||||
mu sync.Mutex
|
||||
total int
|
||||
done map[string]struct{}
|
||||
frac map[string]float64
|
||||
highWater float64
|
||||
}
|
||||
|
||||
func newRunProgress() *runProgress {
|
||||
return &runProgress{
|
||||
done: make(map[string]struct{}),
|
||||
frac: make(map[string]float64),
|
||||
}
|
||||
}
|
||||
|
||||
// SetTotal records the component-count denominator from OnComponentTotal.
|
||||
func (r *runProgress) SetTotal(total int) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
r.total = total
|
||||
}
|
||||
|
||||
// MarkDone counts a component as complete on its lifecycle exit event and
|
||||
// drops its fraction so the two never double-count.
|
||||
func (r *runProgress) MarkDone(component string) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
r.done[component] = struct{}{}
|
||||
delete(r.frac, component)
|
||||
}
|
||||
|
||||
// SetFrac records an in-flight component's 0..1 completion fraction. Reports
|
||||
// for already-completed components are ignored: a late fraction (e.g. a
|
||||
// trailing page callback racing the exit event) must not pull percent down.
|
||||
// A fraction below the one already stored for the component is dropped for
|
||||
// the same reason: per-job reports settle out of order (extractor futures,
|
||||
// tag fan-out), so a trailing lower value must not overwrite the higher one.
|
||||
func (r *runProgress) SetFrac(component string, frac float64) {
|
||||
// NaN passes both clamps below and would then poison the sticky
|
||||
// high-water mark forever; drop it instead of storing it.
|
||||
if math.IsNaN(frac) {
|
||||
return
|
||||
}
|
||||
if frac < 0 {
|
||||
frac = 0
|
||||
}
|
||||
if frac > 1 {
|
||||
frac = 1
|
||||
}
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
if _, ok := r.done[component]; ok {
|
||||
return
|
||||
}
|
||||
if cur, ok := r.frac[component]; ok && frac < cur {
|
||||
return
|
||||
}
|
||||
r.frac[component] = frac
|
||||
}
|
||||
|
||||
// Percent returns the run's 0..1 progress, clamped and monotonic. It returns
|
||||
// 0 until OnComponentTotal supplies the denominator.
|
||||
func (r *runProgress) Percent() float64 {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
if r.total <= 0 {
|
||||
return 0
|
||||
}
|
||||
sum := float64(len(r.done))
|
||||
for _, f := range r.frac {
|
||||
sum += f
|
||||
}
|
||||
p := sum / float64(r.total)
|
||||
if p > 1 {
|
||||
p = 1
|
||||
}
|
||||
if p < r.highWater {
|
||||
return r.highWater
|
||||
}
|
||||
r.highWater = p
|
||||
return p
|
||||
}
|
||||
@@ -0,0 +1,163 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
package service
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"math"
|
||||
"sync"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func almostEqual(a, b float64) bool { return math.Abs(a-b) < 1e-9 }
|
||||
|
||||
func TestRunProgressZeroTotal(t *testing.T) {
|
||||
r := newRunProgress()
|
||||
r.MarkDone("File")
|
||||
r.SetFrac("Parser:abc", 0.5)
|
||||
if got := r.Percent(); got != 0 {
|
||||
t.Fatalf("Percent without total = %v, want 0", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRunProgressDoneAndFrac(t *testing.T) {
|
||||
r := newRunProgress()
|
||||
r.SetTotal(5)
|
||||
r.MarkDone("File")
|
||||
r.SetFrac("Parser:abc", 0.5)
|
||||
if got := r.Percent(); !almostEqual(got, 0.3) {
|
||||
t.Fatalf("Percent = %v, want 0.3 ((1 done + 0.5 frac)/5)", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRunProgressMarkDoneClearsFrac(t *testing.T) {
|
||||
r := newRunProgress()
|
||||
r.SetTotal(5)
|
||||
r.SetFrac("Parser:abc", 0.5)
|
||||
r.MarkDone("Parser:abc")
|
||||
if got := r.Percent(); !almostEqual(got, 0.2) {
|
||||
t.Fatalf("Percent = %v, want 0.2 (frac must not double-count after done)", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRunProgressSetFracAfterDoneIgnored(t *testing.T) {
|
||||
r := newRunProgress()
|
||||
r.SetTotal(5)
|
||||
r.MarkDone("Parser:abc")
|
||||
r.SetFrac("Parser:abc", 0.5)
|
||||
if got := r.Percent(); !almostEqual(got, 0.2) {
|
||||
t.Fatalf("Percent = %v, want 0.2 (late frac for done component ignored)", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRunProgressMonotonicOnFracRegression(t *testing.T) {
|
||||
r := newRunProgress()
|
||||
r.SetTotal(5)
|
||||
r.SetFrac("Parser:abc", 0.9)
|
||||
if got := r.Percent(); !almostEqual(got, 0.18) {
|
||||
t.Fatalf("Percent = %v, want 0.18", got)
|
||||
}
|
||||
r.SetFrac("Parser:abc", 0.1)
|
||||
if got := r.Percent(); !almostEqual(got, 0.18) {
|
||||
t.Fatalf("Percent after regression = %v, want high-water 0.18", got)
|
||||
}
|
||||
r.SetFrac("Parser:abc", 0.95)
|
||||
if got := r.Percent(); !almostEqual(got, 0.19) {
|
||||
t.Fatalf("Percent after advance = %v, want 0.19", got)
|
||||
}
|
||||
}
|
||||
|
||||
// TestRunProgressKeepsHighestFracPerComponent pins the per-component max:
|
||||
// per-job reports settle out of order, so a stale lower fraction must not
|
||||
// replace the higher one stored for the same component. The second
|
||||
// component's report lifts the sum above the high-water mark, which would
|
||||
// otherwise hide the overwrite.
|
||||
func TestRunProgressKeepsHighestFracPerComponent(t *testing.T) {
|
||||
r := newRunProgress()
|
||||
r.SetTotal(4)
|
||||
r.SetFrac("Extractor:abc", 0.9)
|
||||
if got := r.Percent(); !almostEqual(got, 0.225) {
|
||||
t.Fatalf("Percent = %v, want 0.225", got)
|
||||
}
|
||||
r.SetFrac("Extractor:abc", 0.2)
|
||||
r.SetFrac("Extractor:def", 0.1)
|
||||
if got := r.Percent(); !almostEqual(got, 0.25) {
|
||||
t.Fatalf("Percent = %v, want 0.25 (0.9 + 0.1; the stale 0.2 must not overwrite 0.9)", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRunProgressClampsFracAndPercent(t *testing.T) {
|
||||
r := newRunProgress()
|
||||
r.SetTotal(2)
|
||||
r.SetFrac("a", -0.5)
|
||||
r.SetFrac("b", 1.5)
|
||||
if got := r.Percent(); !almostEqual(got, 0.5) {
|
||||
t.Fatalf("Percent = %v, want 0.5 (frac clamped to [0,1])", got)
|
||||
}
|
||||
r.MarkDone("a")
|
||||
r.MarkDone("b")
|
||||
r.MarkDone("c")
|
||||
if got := r.Percent(); got != 1 {
|
||||
t.Fatalf("Percent = %v, want 1 (clamped when done exceeds total)", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRunProgressIgnoresNaN(t *testing.T) {
|
||||
r := newRunProgress()
|
||||
r.SetTotal(5)
|
||||
r.SetFrac("Parser:abc", 0.4)
|
||||
if got := r.Percent(); !almostEqual(got, 0.08) {
|
||||
t.Fatalf("Percent = %v, want 0.08", got)
|
||||
}
|
||||
|
||||
r.SetFrac("Parser:abc", math.NaN())
|
||||
if got := r.Percent(); !almostEqual(got, 0.08) {
|
||||
t.Fatalf("Percent after NaN report = %v, want 0.08 (NaN must not poison the high-water mark)", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRunProgressIdempotentRepeatedReports(t *testing.T) {
|
||||
r := newRunProgress()
|
||||
r.SetTotal(5)
|
||||
for i := 0; i < 3; i++ {
|
||||
r.MarkDone("File")
|
||||
r.SetFrac("Parser:abc", 0.5)
|
||||
}
|
||||
if got := r.Percent(); !almostEqual(got, 0.3) {
|
||||
t.Fatalf("Percent = %v, want 0.3 (repeated reports must be idempotent)", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRunProgressConcurrentReports(t *testing.T) {
|
||||
r := newRunProgress()
|
||||
r.SetTotal(100)
|
||||
var wg sync.WaitGroup
|
||||
for i := 0; i < 100; i++ {
|
||||
wg.Add(1)
|
||||
go func(i int) {
|
||||
defer wg.Done()
|
||||
name := fmt.Sprintf("comp-%d", i)
|
||||
r.SetFrac(name, 0.5)
|
||||
r.MarkDone(name)
|
||||
r.Percent()
|
||||
}(i)
|
||||
}
|
||||
wg.Wait()
|
||||
if got := r.Percent(); got != 1 {
|
||||
t.Fatalf("Percent = %v, want 1 after all components done", got)
|
||||
}
|
||||
}
|
||||
@@ -37,6 +37,7 @@ func (p *PDFParser) ParseWithResult(ctx context.Context, filename string, data [
|
||||
cfg.Pages = p.Pages
|
||||
cfg.RemoveTOC = p.RemoveTOC
|
||||
cfg.RemoveHeaderFooter = p.RemoveHeaderFooter
|
||||
cfg.OnPageDone = p.OnPageDone
|
||||
parser := deepdocpdf.NewParser(cfg)
|
||||
res := parsePDFWithDeepDocOptions(ctx, filename, data, pdfPostProcessOptions{
|
||||
outputFormat: p.OutputFormat,
|
||||
|
||||
@@ -68,7 +68,12 @@ type PDFParser struct {
|
||||
// Pages restricts parsing to these 1-indexed inclusive page ranges.
|
||||
// nil/empty means parse all pages. Populated by ConfigureFromSetup from
|
||||
// the filetype setup map and forwarded to the deepdoc ParserConfig.
|
||||
Pages [][]int
|
||||
Pages [][]int
|
||||
// OnPageDone, when set, is forwarded to the deepdoc ParserConfig so the
|
||||
// caller observes per-page parse progress (done/total). Only the deepdoc
|
||||
// backend invokes it; remote engines poll opaque HTTP jobs and never call
|
||||
// it, so the fraction simply stays where it was.
|
||||
OnPageDone func(done, total int)
|
||||
MinerUAPIServer string
|
||||
MinerUAPIKey string
|
||||
MinerUBackend string
|
||||
|
||||
@@ -1063,9 +1063,3 @@ func (s *IngestionTaskService) truncateIngestionEventMessage(message string) str
|
||||
}
|
||||
return marker
|
||||
}
|
||||
|
||||
// AggregateTaskProgressByPipelineLogID returns component progress for one
|
||||
// immutable ingestion run.
|
||||
func (s *IngestionTaskService) AggregateTaskProgressByPipelineLogID(ctx context.Context, pipelineLogID string, total int) (*dao.TaskProgress, error) {
|
||||
return s.ingestionTaskLogDAO.AggregateProgressByPipelineLogID(ctx, dao.DB, pipelineLogID, total)
|
||||
}
|
||||
@@ -1118,34 +1118,6 @@ func TestIngestionTaskServiceRecordLifecyclePersistsRunScopedLifecycleEvent(t *t
|
||||
}
|
||||
}
|
||||
|
||||
func TestIngestionTaskServiceAggregateTaskProgressByRunClassifiesByPhase(t *testing.T) {
|
||||
db := setupServiceTestDB(t)
|
||||
pushServiceDB(t, db)
|
||||
insertTestIngestionTask(t, "task-1", "user-1", "doc-1", "kb-1")
|
||||
ctx := t.Context()
|
||||
|
||||
svc := NewIngestionTaskService()
|
||||
if err := svc.RecordLifecycle(ctx, "run-1", "task-1", "Parser", 1, "Parser Done"); err != nil {
|
||||
t.Fatalf("record Parser: %v", err)
|
||||
}
|
||||
if err := svc.RecordLifecycle(ctx, "run-1", "task-1", "Chunker", 0, "Chunker Started"); err != nil {
|
||||
t.Fatalf("record Chunker: %v", err)
|
||||
}
|
||||
if err := svc.RecordLifecycle(ctx, "run-2", "task-1", "Parser", 0, "other run"); err != nil {
|
||||
t.Fatalf("record other run: %v", err)
|
||||
}
|
||||
agg, err := svc.AggregateTaskProgressByPipelineLogID(ctx, "run-1", 2)
|
||||
if err != nil {
|
||||
t.Fatalf("AggregateTaskProgressByPipelineLogID failed: %v", err)
|
||||
}
|
||||
if agg.Done != 1 || agg.Running != 1 || agg.Failed != 0 {
|
||||
t.Fatalf("aggregate = %+v, want Done=1 Running=1 Failed=0", agg)
|
||||
}
|
||||
if agg.Percent != 50 {
|
||||
t.Fatalf("percent = %v, want 50", agg.Percent)
|
||||
}
|
||||
}
|
||||
|
||||
func TestIngestionTaskServiceRecordMessageClearsLifecycleFields(t *testing.T) {
|
||||
db := setupServiceTestDB(t)
|
||||
pushServiceDB(t, db)
|
||||
|
||||
Reference in new issue
Block a user