diff --git a/internal/agent/canvas/node_body.go b/internal/agent/canvas/node_body.go index 7146cdf80a..d7fd6b791f 100644 --- a/internal/agent/canvas/node_body.go +++ b/internal/agent/canvas/node_body.go @@ -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 { diff --git a/internal/agent/canvas/node_body_fraction_test.go b/internal/agent/canvas/node_body_fraction_test.go new file mode 100644 index 0000000000..b5c115d34e --- /dev/null +++ b/internal/agent/canvas/node_body_fraction_test.go @@ -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) + } +} diff --git a/internal/agent/runtime/helpers.go b/internal/agent/runtime/helpers.go index fd6df02856..f9eb9d11a4 100644 --- a/internal/agent/runtime/helpers.go +++ b/internal/agent/runtime/helpers.go @@ -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) + } +} diff --git a/internal/agent/runtime/helpers_test.go b/internal/agent/runtime/helpers_test.go index 462e9c74ce..4bd265d331 100644 --- a/internal/agent/runtime/helpers_test.go +++ b/internal/agent/runtime/helpers_test.go @@ -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") + } +} diff --git a/internal/deepdoc/parser/pdf/parser.go b/internal/deepdoc/parser/pdf/parser.go index 432a5f4c3c..215e548455 100644 --- a/internal/deepdoc/parser/pdf/parser.go +++ b/internal/deepdoc/parser/pdf/parser.go @@ -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) diff --git a/internal/deepdoc/parser/pdf/parser_concurrency.go b/internal/deepdoc/parser/pdf/parser_concurrency.go index a6a96ccc63..07cdc44f3d 100644 --- a/internal/deepdoc/parser/pdf/parser_concurrency.go +++ b/internal/deepdoc/parser/pdf/parser_concurrency.go @@ -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 }) diff --git a/internal/deepdoc/parser/pdf/parser_page_progress_test.go b/internal/deepdoc/parser/pdf/parser_page_progress_test.go new file mode 100644 index 0000000000..a3b3305d4e --- /dev/null +++ b/internal/deepdoc/parser/pdf/parser_page_progress_test.go @@ -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") + } +} diff --git a/internal/deepdoc/parser/type/types.go b/internal/deepdoc/parser/type/types.go index aa76f4a697..92ded7bb62 100644 --- a/internal/deepdoc/parser/type/types.go +++ b/internal/deepdoc/parser/type/types.go @@ -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. diff --git a/internal/ingestion/component/extractor.go b/internal/ingestion/component/extractor.go index 10c82fd8af..60a55ffc5e 100644 --- a/internal/ingestion/component/extractor.go +++ b/internal/ingestion/component/extractor.go @@ -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() diff --git a/internal/ingestion/component/extractor_tag.go b/internal/ingestion/component/extractor_tag.go index 09d725700f..7480177e8d 100644 --- a/internal/ingestion/component/extractor_tag.go +++ b/internal/ingestion/component/extractor_tag.go @@ -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 }() diff --git a/internal/ingestion/component/knowledge_compiler/tree/claims.go b/internal/ingestion/component/knowledge_compiler/tree/claims.go index 5e612d1d2c..7b3f1af061 100644 --- a/internal/ingestion/component/knowledge_compiler/tree/claims.go +++ b/internal/ingestion/component/knowledge_compiler/tree/claims.go @@ -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)) } diff --git a/internal/ingestion/component/parser_dispatch.go b/internal/ingestion/component/parser_dispatch.go index 5770ec4e06..9552a2cab3 100644 --- a/internal/ingestion/component/parser_dispatch.go +++ b/internal/ingestion/component/parser_dispatch.go @@ -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 diff --git a/internal/ingestion/component/parser_dispatch_progress_test.go b/internal/ingestion/component/parser_dispatch_progress_test.go new file mode 100644 index 0000000000..dc2e961d64 --- /dev/null +++ b/internal/ingestion/component/parser_dispatch_progress_test.go @@ -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) + } +} diff --git a/internal/ingestion/component/progress_fraction_test.go b/internal/ingestion/component/progress_fraction_test.go new file mode 100644 index 0000000000..fb2805af0c --- /dev/null +++ b/internal/ingestion/component/progress_fraction_test.go @@ -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) + } +} diff --git a/internal/ingestion/component/tokenizer.go b/internal/ingestion/component/tokenizer.go index 62ab750b34..f1b8293e71 100644 --- a/internal/ingestion/component/tokenizer.go +++ b/internal/ingestion/component/tokenizer.go @@ -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 diff --git a/internal/ingestion/pipeline/pipeline.go b/internal/ingestion/pipeline/pipeline.go index 43dce51ea3..1551d64c45 100644 --- a/internal/ingestion/pipeline/pipeline.go +++ b/internal/ingestion/pipeline/pipeline.go @@ -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 { diff --git a/internal/ingestion/pipeline/pipeline_test.go b/internal/ingestion/pipeline/pipeline_test.go index 57b7c50f42..a09cccdf4f 100644 --- a/internal/ingestion/pipeline/pipeline_test.go +++ b/internal/ingestion/pipeline/pipeline_test.go @@ -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) + } +} diff --git a/internal/ingestion/service/ingestion_service.go b/internal/ingestion/service/ingestion_service.go index 16f5350501..8e885750a8 100644 --- a/internal/ingestion/service/ingestion_service.go +++ b/internal/ingestion/service/ingestion_service.go @@ -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 } diff --git a/internal/ingestion/service/progress_sink.go b/internal/ingestion/service/progress_sink.go index 35c117d80f..bef6f6fba0 100644 --- a/internal/ingestion/service/progress_sink.go +++ b/internal/ingestion/service/progress_sink.go @@ -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 -} diff --git a/internal/ingestion/service/progress_sink_test.go b/internal/ingestion/service/progress_sink_test.go index 7680b53f65..17c7f4a2e9 100644 --- a/internal/ingestion/service/progress_sink_test.go +++ b/internal/ingestion/service/progress_sink_test.go @@ -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) } } diff --git a/internal/ingestion/service/progress_tracker.go b/internal/ingestion/service/progress_tracker.go new file mode 100644 index 0000000000..b37daace71 --- /dev/null +++ b/internal/ingestion/service/progress_tracker.go @@ -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 +} diff --git a/internal/ingestion/service/progress_tracker_test.go b/internal/ingestion/service/progress_tracker_test.go new file mode 100644 index 0000000000..c6149a231b --- /dev/null +++ b/internal/ingestion/service/progress_tracker_test.go @@ -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) + } +} diff --git a/internal/parser/parser/pdf_parser_cgo.go b/internal/parser/parser/pdf_parser_cgo.go index 349de5b63c..959b275484 100644 --- a/internal/parser/parser/pdf_parser_cgo.go +++ b/internal/parser/parser/pdf_parser_cgo.go @@ -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, diff --git a/internal/parser/parser/pdf_parser_common.go b/internal/parser/parser/pdf_parser_common.go index b5de904d3a..fcde314cda 100644 --- a/internal/parser/parser/pdf_parser_common.go +++ b/internal/parser/parser/pdf_parser_common.go @@ -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 diff --git a/internal/service/ingestion_task_service.go b/internal/service/ingestion_task_service.go index ad7cbbdbce..0cb07dae14 100644 --- a/internal/service/ingestion_task_service.go +++ b/internal/service/ingestion_task_service.go @@ -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) -} diff --git a/internal/service/ingestion_task_service_test.go b/internal/service/ingestion_task_service_test.go index c9750269c2..a0b9f0042b 100644 --- a/internal/service/ingestion_task_service_test.go +++ b/internal/service/ingestion_task_service_test.go @@ -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)