676 lines
25 KiB
Go
676 lines
25 KiB
Go
package pipeline
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"fmt"
|
|
"reflect"
|
|
"strings"
|
|
"sync"
|
|
"sync/atomic"
|
|
"testing"
|
|
"time"
|
|
|
|
"gitea.maximumdirect.net/eric/notarius/internal/core/source"
|
|
"gitea.maximumdirect.net/eric/notarius/internal/framework/contracts"
|
|
frameworkllm "gitea.maximumdirect.net/eric/notarius/internal/framework/llm"
|
|
)
|
|
|
|
type countingProvider struct {
|
|
active atomic.Int32
|
|
maximum atomic.Int32
|
|
calls atomic.Int32
|
|
}
|
|
|
|
func (p *countingProvider) CompleteStructured(ctx context.Context, _ contracts.StructuredCompletionRequest, _ any) (contracts.StructuredCompletionResponse, error) {
|
|
p.calls.Add(1)
|
|
current := p.active.Add(1)
|
|
defer p.active.Add(-1)
|
|
for {
|
|
seen := p.maximum.Load()
|
|
if current <= seen || p.maximum.CompareAndSwap(seen, current) {
|
|
break
|
|
}
|
|
}
|
|
select {
|
|
case <-time.After(5 * time.Millisecond):
|
|
return contracts.StructuredCompletionResponse{}, nil
|
|
case <-ctx.Done():
|
|
return contracts.StructuredCompletionResponse{}, ctx.Err()
|
|
}
|
|
}
|
|
|
|
func preparedConcurrentPipeline(t *testing.T, chunkCount int) *PreparedPipeline {
|
|
t.Helper()
|
|
catalog := typedResolutionCatalog(t, completeTypedCatalogOptions())
|
|
resolved, err := ResolvePipeline(typedResolutionProfile(), ResolveOptions{}, catalog)
|
|
if err != nil {
|
|
t.Fatalf("ResolvePipeline() error = %v", err)
|
|
}
|
|
prepared, err := Prepare(resolved, registriesFromModuleCatalog(catalog), ModuleDependencies{})
|
|
if err != nil {
|
|
t.Fatalf("Prepare() error = %v", err)
|
|
}
|
|
doc := typedTestDocumentWithUnits(chunkCount)
|
|
if adapter, ok := prepared.input.(*typedTestInput); ok {
|
|
adapter.doc = doc
|
|
} else {
|
|
t.Fatalf("prepared input = %T, want *typedTestInput", prepared.input)
|
|
}
|
|
prepared.chunker = &typedTestChunker{key: "typed/chunk", plan: typedTestPlan(doc)}
|
|
return prepared
|
|
}
|
|
|
|
type orderedLaneSpec struct {
|
|
id string
|
|
profile string
|
|
}
|
|
|
|
func preparedOrderedPipeline(t *testing.T, chunkCount int, specs ...orderedLaneSpec) *PreparedPipeline {
|
|
t.Helper()
|
|
catalog := typedResolutionCatalog(t, completeTypedCatalogOptions())
|
|
profile := typedResolutionProfile()
|
|
profile.Artifacts = nil
|
|
for index, spec := range specs {
|
|
lane, ok := typedResolutionProfile().Artifacts[spec.profile]
|
|
if !ok {
|
|
t.Fatalf("typed lane profile %q is not defined", spec.profile)
|
|
}
|
|
profile.Steps = append(profile.Steps, PipelineStepProfile{
|
|
ID: fmt.Sprintf("step-%d", index+1),
|
|
Artifacts: map[string]ArtifactLaneProfile{spec.id: lane},
|
|
})
|
|
}
|
|
resolved, err := ResolvePipeline(profile, ResolveOptions{}, catalog)
|
|
if err != nil {
|
|
t.Fatalf("ResolvePipeline() error = %v", err)
|
|
}
|
|
prepared, err := Prepare(resolved, registriesFromModuleCatalog(catalog), ModuleDependencies{})
|
|
if err != nil {
|
|
t.Fatalf("Prepare() error = %v", err)
|
|
}
|
|
doc := typedTestDocumentWithUnits(chunkCount)
|
|
if adapter, ok := prepared.input.(*typedTestInput); ok {
|
|
adapter.doc = doc
|
|
} else {
|
|
t.Fatalf("prepared input = %T, want *typedTestInput", prepared.input)
|
|
}
|
|
prepared.chunker = &typedTestChunker{key: "typed/chunk", plan: typedTestPlan(doc)}
|
|
return prepared
|
|
}
|
|
|
|
type countingOrderedOutput struct {
|
|
calls atomic.Int32
|
|
}
|
|
|
|
func (o *countingOrderedOutput) Key() string { return "typed/output" }
|
|
func (o *countingOrderedOutput) Encode(context.Context, contracts.OutputRequest) (contracts.OutputResult, error) {
|
|
o.calls.Add(1)
|
|
return contracts.OutputResult{}, nil
|
|
}
|
|
|
|
type countingInputAdapter struct {
|
|
contracts.InputAdapter
|
|
calls atomic.Int32
|
|
}
|
|
|
|
func (a *countingInputAdapter) Parse(ctx context.Context, request contracts.ParseRequest) (*source.SourceDocument, error) {
|
|
a.calls.Add(1)
|
|
return a.InputAdapter.Parse(ctx, request)
|
|
}
|
|
|
|
type cancelingDebugRecorder struct {
|
|
cancel context.CancelFunc
|
|
onName string
|
|
once sync.Once
|
|
}
|
|
|
|
func (r *cancelingDebugRecorder) Enabled() bool { return true }
|
|
|
|
func (r *cancelingDebugRecorder) WriteJSON(name string, _ any) error {
|
|
if name == r.onName {
|
|
r.once.Do(r.cancel)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (*cancelingDebugRecorder) WriteBytes(string, []byte) error { return nil }
|
|
|
|
type cancelingOutputEncoder struct {
|
|
contracts.OutputEncoder
|
|
cancel context.CancelFunc
|
|
calls atomic.Int32
|
|
}
|
|
|
|
func (e *cancelingOutputEncoder) Encode(context.Context, contracts.OutputRequest) (contracts.OutputResult, error) {
|
|
e.calls.Add(1)
|
|
e.cancel()
|
|
return contracts.OutputResult{
|
|
Files: []contracts.OutputFile{{Name: "result.txt", ContentType: "text/plain", Bytes: []byte("result")}},
|
|
Warnings: []contracts.Warning{{ReasonCode: "returned-after-cancel"}},
|
|
}, nil
|
|
}
|
|
|
|
func TestRunnerSkipsInputAdapterWhenContextIsAlreadyCanceled(t *testing.T) {
|
|
prepared := preparedConcurrentPipeline(t, 1)
|
|
adapter := &countingInputAdapter{InputAdapter: prepared.input}
|
|
prepared.input = adapter
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
cancel()
|
|
|
|
_, err := New().Run(ctx, RunInput{Prepared: prepared, RawInput: []byte("input")})
|
|
if !errors.Is(err, context.Canceled) {
|
|
t.Fatalf("Run() error = %v, want context.Canceled", err)
|
|
}
|
|
if got := adapter.calls.Load(); got != 0 {
|
|
t.Fatalf("input Parse calls = %d, want none", got)
|
|
}
|
|
}
|
|
|
|
func TestRunnerSkipsOutputEncodingAfterDebugCancellation(t *testing.T) {
|
|
prepared := preparedConcurrentPipeline(t, 1)
|
|
encoder := &countingOrderedOutput{}
|
|
prepared.output = encoder
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
|
|
_, err := New().Run(ctx, RunInput{
|
|
Prepared: prepared,
|
|
RawInput: []byte("input"),
|
|
Debug: &cancelingDebugRecorder{cancel: cancel, onName: "output/input.json"},
|
|
})
|
|
if !errors.Is(err, context.Canceled) {
|
|
t.Fatalf("Run() error = %v, want context.Canceled", err)
|
|
}
|
|
if got := encoder.calls.Load(); got != 0 {
|
|
t.Fatalf("output Encode calls = %d, want none", got)
|
|
}
|
|
}
|
|
|
|
func TestRunnerDiscardsOutputReturnedAfterCancellation(t *testing.T) {
|
|
prepared := preparedConcurrentPipeline(t, 1)
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
encoder := &cancelingOutputEncoder{OutputEncoder: prepared.output, cancel: cancel}
|
|
prepared.output = encoder
|
|
|
|
output, err := New().Run(ctx, RunInput{Prepared: prepared, RawInput: []byte("input")})
|
|
if !errors.Is(err, context.Canceled) {
|
|
t.Fatalf("Run() error = %v, want context.Canceled", err)
|
|
}
|
|
if got := encoder.calls.Load(); got != 1 {
|
|
t.Fatalf("output Encode calls = %d, want one", got)
|
|
}
|
|
if len(output.OutputFiles) != 0 {
|
|
t.Fatalf("output files = %#v, want none", output.OutputFiles)
|
|
}
|
|
for _, warning := range output.Warnings {
|
|
if warning.ReasonCode == "returned-after-cancel" {
|
|
t.Fatalf("output warnings include encoder warning after cancellation")
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestRunnerFailsGeneratedHandoffBeforeConsumerExtraction(t *testing.T) {
|
|
prepared := preparedConcurrentPipeline(t, 1)
|
|
input := &countingInputAdapter{InputAdapter: prepared.input}
|
|
prepared.input = input
|
|
preparedLane := &prepared.Steps[0].lanes[0]
|
|
lane := &preparedLane.resolved
|
|
binding := ReferenceBinding{
|
|
Stage: StageExtract,
|
|
LaneID: lane.ID,
|
|
SlotName: "generated",
|
|
Artifact: &ArtifactReference{Step: "producer", Lane: "source"},
|
|
}
|
|
lane.ExtractReferences.Bindings = []ReferenceBinding{binding}
|
|
lane.ExtractReferences.ReferenceSet = contracts.ReferenceSet{Slots: map[string]contracts.ResolvedReferenceSlot{
|
|
"generated": {Slot: contracts.ReferenceSlot{Name: "generated", AcceptedArtifactKinds: []contracts.ArtifactKind{"test/notes"}}},
|
|
}}
|
|
prepared.resolved.Steps[0].ArtifactLanes[0].ExtractReferences = CloneReferenceTarget(lane.ExtractReferences)
|
|
var extractCalls atomic.Int32
|
|
preparedLane.typed.extract = func(context.Context, any, contracts.TypedExtractionRequest) (erasedTypedResult, error) {
|
|
extractCalls.Add(1)
|
|
return erasedTypedResult{Value: codecNotes{}}, nil
|
|
}
|
|
|
|
_, err := New().Run(context.Background(), RunInput{Prepared: prepared, RawInput: []byte("input")})
|
|
if err == nil || !strings.Contains(err.Error(), "generated dependency") {
|
|
t.Fatalf("Run() error = %v, want generated dependency failure", err)
|
|
}
|
|
if got := input.calls.Load(); got != 1 {
|
|
t.Fatalf("input Parse calls = %d, want one before handoff", got)
|
|
}
|
|
if got := extractCalls.Load(); got != 0 {
|
|
t.Fatalf("consumer extract calls = %d, want zero", got)
|
|
}
|
|
}
|
|
|
|
func TestRunnerExecutesOrderedStepsWithHardBarriers(t *testing.T) {
|
|
prepared := preparedOrderedPipeline(t, 1,
|
|
orderedLaneSpec{id: "notes", profile: "notes"},
|
|
orderedLaneSpec{id: "score", profile: "score"},
|
|
)
|
|
var mu sync.Mutex
|
|
var events []string
|
|
record := func(event string) {
|
|
mu.Lock()
|
|
events = append(events, event)
|
|
mu.Unlock()
|
|
}
|
|
first := &prepared.Steps[0].lanes[0]
|
|
second := &prepared.Steps[1].lanes[0]
|
|
first.typed.extract = func(context.Context, any, contracts.TypedExtractionRequest) (erasedTypedResult, error) {
|
|
return erasedTypedResult{Value: codecNotes{Items: []string{"first"}}}, nil
|
|
}
|
|
originalFirstNormalize := first.typed.normalize
|
|
first.typed.normalize = func(ctx context.Context, implementation any, request contracts.TypedNormalizeRequest[any]) (erasedTypedResult, error) {
|
|
result, err := originalFirstNormalize(ctx, implementation, request)
|
|
if err == nil {
|
|
record("first-normalize-done")
|
|
}
|
|
return result, err
|
|
}
|
|
second.typed.extract = func(context.Context, any, contracts.TypedExtractionRequest) (erasedTypedResult, error) {
|
|
record("second-extract-started")
|
|
return erasedTypedResult{Value: codecScore{Value: 2}}, nil
|
|
}
|
|
|
|
output, err := New().Run(context.Background(), RunInput{Prepared: prepared, RawInput: []byte("input"), ExtractWorkers: 2})
|
|
if err != nil {
|
|
t.Fatalf("Run() error = %v", err)
|
|
}
|
|
mu.Lock()
|
|
gotEvents := append([]string(nil), events...)
|
|
mu.Unlock()
|
|
if want := []string{"first-normalize-done", "second-extract-started"}; !reflect.DeepEqual(gotEvents, want) {
|
|
t.Fatalf("ordered events = %#v, want %#v", gotEvents, want)
|
|
}
|
|
if got := []string{output.NormalizeOutputs[0].LaneID, output.NormalizeOutputs[1].LaneID}; !reflect.DeepEqual(got, []string{"notes", "score"}) {
|
|
t.Fatalf("normalized lane order = %#v, want notes then score", got)
|
|
}
|
|
}
|
|
|
|
func TestRunnerStopsLaterOrderedStepsAfterFailure(t *testing.T) {
|
|
prepared := preparedOrderedPipeline(t, 1,
|
|
orderedLaneSpec{id: "first", profile: "notes"},
|
|
orderedLaneSpec{id: "failure", profile: "score"},
|
|
orderedLaneSpec{id: "later", profile: "notes"},
|
|
)
|
|
first := &prepared.Steps[0].lanes[0]
|
|
second := &prepared.Steps[1].lanes[0]
|
|
third := &prepared.Steps[2].lanes[0]
|
|
first.typed.extract = func(context.Context, any, contracts.TypedExtractionRequest) (erasedTypedResult, error) {
|
|
return erasedTypedResult{Value: codecNotes{Items: []string{"first"}}}, nil
|
|
}
|
|
second.typed.extract = func(context.Context, any, contracts.TypedExtractionRequest) (erasedTypedResult, error) {
|
|
return erasedTypedResult{}, errors.New("ordered step extraction failed")
|
|
}
|
|
var thirdCalls atomic.Int32
|
|
third.typed.extract = func(context.Context, any, contracts.TypedExtractionRequest) (erasedTypedResult, error) {
|
|
thirdCalls.Add(1)
|
|
return erasedTypedResult{Value: codecNotes{Items: []string{"later"}}}, nil
|
|
}
|
|
encoder := &countingOrderedOutput{}
|
|
prepared.output = encoder
|
|
|
|
output, err := New().Run(context.Background(), RunInput{Prepared: prepared, RawInput: []byte("input"), ExtractWorkers: 2})
|
|
if err == nil || !strings.Contains(err.Error(), `execute pipeline step "step-2"`) {
|
|
t.Fatalf("Run() error = %v, want step-scoped failure", err)
|
|
}
|
|
if len(output.NormalizeOutputs) != 1 || output.NormalizeOutputs[0].LaneID != "first" {
|
|
t.Fatalf("completed outputs = %#v, want only the first step output", output.NormalizeOutputs)
|
|
}
|
|
if got := thirdCalls.Load(); got != 0 {
|
|
t.Fatalf("later step extract calls = %d, want zero", got)
|
|
}
|
|
if got := encoder.calls.Load(); got != 0 {
|
|
t.Fatalf("output encoder calls = %d, want zero after failure", got)
|
|
}
|
|
}
|
|
|
|
func installExtractOperation(prepared *PreparedPipeline, laneIndex int, operation func(context.Context, contracts.TypedExtractionRequest) (erasedTypedResult, error)) {
|
|
prepared.Steps[0].lanes[laneIndex].typed.extract = func(ctx context.Context, _ any, request contracts.TypedExtractionRequest) (erasedTypedResult, error) {
|
|
return operation(ctx, request)
|
|
}
|
|
}
|
|
|
|
func typedValueForLane(laneIndex int, chunkIndex int) any {
|
|
if laneIndex == 0 {
|
|
return codecNotes{Items: []string{fmt.Sprintf("chunk-%d", chunkIndex)}}
|
|
}
|
|
return codecScore{Value: chunkIndex}
|
|
}
|
|
|
|
func TestRunnerBoundsExtractJobsAndStabilizesReverseCompletion(t *testing.T) {
|
|
prepared := preparedConcurrentPipeline(t, 3)
|
|
started := make(chan int, 6)
|
|
release := make([]chan struct{}, 6)
|
|
for i := range release {
|
|
release[i] = make(chan struct{})
|
|
}
|
|
var active atomic.Int32
|
|
var maximum atomic.Int32
|
|
for laneIndex := range prepared.Steps[0].lanes {
|
|
lane := laneIndex
|
|
installExtractOperation(prepared, lane, func(ctx context.Context, request contracts.TypedExtractionRequest) (erasedTypedResult, error) {
|
|
jobIndex := request.Chunk.Index*2 + lane
|
|
current := active.Add(1)
|
|
defer active.Add(-1)
|
|
for {
|
|
seen := maximum.Load()
|
|
if current <= seen || maximum.CompareAndSwap(seen, current) {
|
|
break
|
|
}
|
|
}
|
|
started <- jobIndex
|
|
select {
|
|
case <-release[jobIndex]:
|
|
case <-ctx.Done():
|
|
return erasedTypedResult{}, ctx.Err()
|
|
}
|
|
return erasedTypedResult{Value: typedValueForLane(lane, request.Chunk.Index), Warnings: []contracts.Warning{{Scope: fmt.Sprintf("lane-%d/chunk-%d", lane, request.Chunk.Index), ReasonCode: "observed", Message: "ordered"}}}, nil
|
|
})
|
|
}
|
|
|
|
type outcome struct {
|
|
output RunOutput
|
|
err error
|
|
}
|
|
done := make(chan outcome, 1)
|
|
go func() {
|
|
output, err := New().Run(context.Background(), RunInput{Prepared: prepared, RawInput: []byte("input"), ExtractWorkers: 2})
|
|
done <- outcome{output: output, err: err}
|
|
}()
|
|
for pair := 0; pair < 3; pair++ {
|
|
first, second := <-started, <-started
|
|
want := map[int]bool{pair * 2: true, pair*2 + 1: true}
|
|
if !want[first] || !want[second] || first == second {
|
|
t.Fatalf("started jobs = %d, %d; want round-robin pair %d", first, second, pair)
|
|
}
|
|
close(release[second])
|
|
close(release[first])
|
|
}
|
|
result := <-done
|
|
if result.err != nil {
|
|
t.Fatalf("Run() error = %v", result.err)
|
|
}
|
|
if got := maximum.Load(); got != 2 {
|
|
t.Fatalf("maximum concurrent extract jobs = %d, want 2", got)
|
|
}
|
|
wantScopes := []string{"lane-0/chunk-0", "lane-0/chunk-1", "lane-0/chunk-2", "lane-1/chunk-0", "lane-1/chunk-1", "lane-1/chunk-2"}
|
|
gotScopes := make([]string, len(result.output.Warnings))
|
|
for i := range result.output.Warnings {
|
|
gotScopes[i] = result.output.Warnings[i].Scope
|
|
}
|
|
if !reflect.DeepEqual(gotScopes, wantScopes) {
|
|
t.Fatalf("warning order = %#v, want %#v", gotScopes, wantScopes)
|
|
}
|
|
}
|
|
|
|
func TestRunnerStartsLaneContinuationWhileOtherLaneExtractsRemain(t *testing.T) {
|
|
prepared := preparedConcurrentPipeline(t, 2)
|
|
releaseOtherLane := make(chan struct{})
|
|
mergeStarted := make(chan struct{}, 1)
|
|
installExtractOperation(prepared, 0, func(_ context.Context, request contracts.TypedExtractionRequest) (erasedTypedResult, error) {
|
|
return erasedTypedResult{Value: typedValueForLane(0, request.Chunk.Index)}, nil
|
|
})
|
|
installExtractOperation(prepared, 1, func(ctx context.Context, request contracts.TypedExtractionRequest) (erasedTypedResult, error) {
|
|
select {
|
|
case <-releaseOtherLane:
|
|
return erasedTypedResult{Value: typedValueForLane(1, request.Chunk.Index)}, nil
|
|
case <-ctx.Done():
|
|
return erasedTypedResult{}, ctx.Err()
|
|
}
|
|
})
|
|
originalMerge := prepared.Steps[0].lanes[0].typed.merge
|
|
prepared.Steps[0].lanes[0].typed.merge = func(ctx context.Context, implementation any, request contracts.TypedMergeRequest[any]) (erasedTypedResult, error) {
|
|
select {
|
|
case mergeStarted <- struct{}{}:
|
|
default:
|
|
}
|
|
return originalMerge(ctx, implementation, request)
|
|
}
|
|
done := make(chan error, 1)
|
|
go func() {
|
|
_, err := New().Run(context.Background(), RunInput{Prepared: prepared, RawInput: []byte("input"), ExtractWorkers: 2})
|
|
done <- err
|
|
}()
|
|
select {
|
|
case <-mergeStarted:
|
|
close(releaseOtherLane)
|
|
case <-time.After(2 * time.Second):
|
|
t.Fatal("first lane merge did not start while the second lane remained in extract")
|
|
}
|
|
if err := <-done; err != nil {
|
|
t.Fatalf("Run() error = %v", err)
|
|
}
|
|
}
|
|
|
|
func TestRunnerBoundsConcurrentLaneContinuations(t *testing.T) {
|
|
prepared := preparedConcurrentPipeline(t, 1)
|
|
base := prepared.Steps[0].lanes[0]
|
|
lanes := make([]preparedLaneExecutor, 8)
|
|
resolvedLanes := make([]ResolvedArtifactLane, len(lanes))
|
|
publicLanes := make([]PreparedArtifactLane, len(lanes))
|
|
var active atomic.Int32
|
|
var maximum atomic.Int32
|
|
for i := range lanes {
|
|
lane := base
|
|
lane.resolved.ID = fmt.Sprintf("lane-%02d", i)
|
|
lane.typed = &preparedTypedLane{
|
|
extractor: base.typed.extractor,
|
|
merger: base.typed.merger,
|
|
normalizer: base.typed.normalizer,
|
|
extract: func(_ context.Context, _ any, request contracts.TypedExtractionRequest) (erasedTypedResult, error) {
|
|
return erasedTypedResult{Value: codecNotes{Items: []string{request.Chunk.ID}}}, nil
|
|
},
|
|
merge: func(_ context.Context, _ any, _ contracts.TypedMergeRequest[any]) (erasedTypedResult, error) {
|
|
current := active.Add(1)
|
|
defer active.Add(-1)
|
|
for {
|
|
seen := maximum.Load()
|
|
if current <= seen || maximum.CompareAndSwap(seen, current) {
|
|
break
|
|
}
|
|
}
|
|
time.Sleep(5 * time.Millisecond)
|
|
return erasedTypedResult{Value: codecNotes{}}, nil
|
|
},
|
|
normalize: base.typed.normalize,
|
|
codec: base.typed.codec,
|
|
}
|
|
lanes[i] = lane
|
|
resolvedLanes[i] = lane.resolved
|
|
publicLanes[i] = PreparedArtifactLane{Resolved: lane.resolved}
|
|
}
|
|
prepared.Steps[0].lanes = lanes
|
|
prepared.resolved.Steps[0].ArtifactLanes = resolvedLanes
|
|
prepared.Steps[0].ArtifactLanes = publicLanes
|
|
|
|
output, err := New().Run(context.Background(), RunInput{Prepared: prepared, RawInput: []byte("input"), ExtractWorkers: 2})
|
|
if err != nil {
|
|
t.Fatalf("Run() error = %v", err)
|
|
}
|
|
if len(output.NormalizeOutputs) != len(lanes) {
|
|
t.Fatalf("normalize outputs = %d, want %d", len(output.NormalizeOutputs), len(lanes))
|
|
}
|
|
if got := maximum.Load(); got != 2 {
|
|
t.Fatalf("maximum concurrent lane continuations = %d, want 2", got)
|
|
}
|
|
}
|
|
|
|
func TestRunnerSelectsFrameworkErrorByStableLaneOrder(t *testing.T) {
|
|
prepared := preparedConcurrentPipeline(t, 1)
|
|
ready := make(chan struct{}, 2)
|
|
release := make(chan struct{})
|
|
for laneIndex := range prepared.Steps[0].lanes {
|
|
lane := laneIndex
|
|
installExtractOperation(prepared, lane, func(context.Context, contracts.TypedExtractionRequest) (erasedTypedResult, error) {
|
|
ready <- struct{}{}
|
|
<-release
|
|
return erasedTypedResult{}, fmt.Errorf("lane-%d failure", lane)
|
|
})
|
|
}
|
|
done := make(chan error, 1)
|
|
go func() {
|
|
_, err := New().Run(context.Background(), RunInput{Prepared: prepared, RawInput: []byte("input"), ExtractWorkers: 2})
|
|
done <- err
|
|
}()
|
|
<-ready
|
|
<-ready
|
|
close(release)
|
|
if err := <-done; err == nil || !strings.Contains(err.Error(), "lane-0 failure") {
|
|
t.Fatalf("Run() error = %v, want first resolved lane failure", err)
|
|
}
|
|
}
|
|
|
|
func TestRunnerSelectsMergeErrorBeforeEarlierLaneNormalizeError(t *testing.T) {
|
|
prepared := preparedConcurrentPipeline(t, 1)
|
|
for laneIndex := range prepared.Steps[0].lanes {
|
|
lane := laneIndex
|
|
installExtractOperation(prepared, lane, func(_ context.Context, request contracts.TypedExtractionRequest) (erasedTypedResult, error) {
|
|
return erasedTypedResult{Value: typedValueForLane(lane, request.Chunk.Index)}, nil
|
|
})
|
|
}
|
|
ready := make(chan struct{}, 2)
|
|
release := make(chan struct{})
|
|
prepared.Steps[0].lanes[0].typed.normalize = func(context.Context, any, contracts.TypedNormalizeRequest[any]) (erasedTypedResult, error) {
|
|
ready <- struct{}{}
|
|
<-release
|
|
return erasedTypedResult{}, errors.New("earlier lane normalize failure")
|
|
}
|
|
prepared.Steps[0].lanes[1].typed.merge = func(context.Context, any, contracts.TypedMergeRequest[any]) (erasedTypedResult, error) {
|
|
ready <- struct{}{}
|
|
<-release
|
|
return erasedTypedResult{}, errors.New("later lane merge failure")
|
|
}
|
|
done := make(chan error, 1)
|
|
go func() {
|
|
_, err := New().Run(context.Background(), RunInput{Prepared: prepared, RawInput: []byte("input"), ExtractWorkers: 2})
|
|
done <- err
|
|
}()
|
|
<-ready
|
|
<-ready
|
|
close(release)
|
|
if err := <-done; err == nil || !strings.Contains(err.Error(), "later lane merge failure") {
|
|
t.Fatalf("Run() error = %v, want merge failure before normalize failure", err)
|
|
}
|
|
}
|
|
|
|
func TestRunnerKeepsSuccessfulChunksWhenAnotherExtractIsRejected(t *testing.T) {
|
|
prepared := preparedConcurrentPipeline(t, 2)
|
|
for laneIndex := range prepared.Steps[0].lanes {
|
|
lane := laneIndex
|
|
installExtractOperation(prepared, lane, func(_ context.Context, request contracts.TypedExtractionRequest) (erasedTypedResult, error) {
|
|
return erasedTypedResult{Value: typedValueForLane(lane, request.Chunk.Index)}, nil
|
|
})
|
|
}
|
|
validator := &prepared.Steps[0].lanes[0].extractValidators.validators[0]
|
|
validator.typedValidate = func(_ context.Context, _ any, target typedValidationTarget) (contracts.ValidationResult, error) {
|
|
if target.chunk != nil && target.chunk.Index == 0 {
|
|
return contracts.ValidationResult{Approved: false, ReasonCode: "expected_rejection", Message: "rejected by test"}, nil
|
|
}
|
|
return contracts.ValidationResult{Approved: true}, nil
|
|
}
|
|
output, err := New().Run(context.Background(), RunInput{Prepared: prepared, RawInput: []byte("input"), ExtractWorkers: 2})
|
|
if err != nil {
|
|
t.Fatalf("Run() error = %v", err)
|
|
}
|
|
if len(output.Rejected) != 1 || output.Rejected[0].LaneID != prepared.Steps[0].lanes[0].resolved.ID || output.Rejected[0].ChunkIndex != 0 {
|
|
t.Fatalf("rejections = %#v, want the first lane's first chunk", output.Rejected)
|
|
}
|
|
if len(output.NormalizeOutputs) != 2 {
|
|
t.Fatalf("normalize outputs = %d, want both lanes to continue", len(output.NormalizeOutputs))
|
|
}
|
|
if output.Manifest.ValidationStatus != "rejected" {
|
|
t.Fatalf("validation status = %q, want rejected", output.Manifest.ValidationStatus)
|
|
}
|
|
}
|
|
|
|
func TestRunnerSeparatelyBoundsExtractJobsAndSharedProviderCalls(t *testing.T) {
|
|
prepared := preparedConcurrentPipeline(t, 3)
|
|
provider := &countingProvider{}
|
|
scheduler, err := frameworkllm.NewScheduler(2)
|
|
if err != nil {
|
|
t.Fatalf("NewScheduler() error = %v", err)
|
|
}
|
|
client := frameworkllm.NewScheduledClient(provider, scheduler)
|
|
var jobActive atomic.Int32
|
|
var jobMaximum atomic.Int32
|
|
var attempts sync.Map
|
|
for laneIndex := range prepared.Steps[0].lanes {
|
|
lane := laneIndex
|
|
prepared.Steps[0].lanes[lane].resolved.Extract.Retries = 1
|
|
installExtractOperation(prepared, lane, func(ctx context.Context, request contracts.TypedExtractionRequest) (erasedTypedResult, error) {
|
|
current := jobActive.Add(1)
|
|
defer jobActive.Add(-1)
|
|
for {
|
|
seen := jobMaximum.Load()
|
|
if current <= seen || jobMaximum.CompareAndSwap(seen, current) {
|
|
break
|
|
}
|
|
}
|
|
if _, callErr := client.CompleteStructured(ctx, contracts.StructuredCompletionRequest{StageName: "extract"}, nil); callErr != nil {
|
|
return erasedTypedResult{}, callErr
|
|
}
|
|
key := fmt.Sprintf("%d/%d", lane, request.Chunk.Index)
|
|
countValue, _ := attempts.LoadOrStore(key, &atomic.Int32{})
|
|
count := countValue.(*atomic.Int32).Add(1)
|
|
if lane == 0 && request.Chunk.Index == 0 && count == 1 {
|
|
return erasedTypedResult{}, errors.New("retry extract")
|
|
}
|
|
return erasedTypedResult{Value: typedValueForLane(lane, request.Chunk.Index)}, nil
|
|
})
|
|
validator := &prepared.Steps[0].lanes[lane].extractValidators.validators[0]
|
|
validator.typedValidate = func(ctx context.Context, _ any, _ typedValidationTarget) (contracts.ValidationResult, error) {
|
|
if _, callErr := client.CompleteStructured(ctx, contracts.StructuredCompletionRequest{StageName: "validate"}, nil); callErr != nil {
|
|
return contracts.ValidationResult{}, callErr
|
|
}
|
|
return contracts.ValidationResult{Approved: true}, nil
|
|
}
|
|
}
|
|
output, err := New().Run(context.Background(), RunInput{Prepared: prepared, RawInput: []byte("input"), ExtractWorkers: 3})
|
|
if err != nil {
|
|
t.Fatalf("Run() error = %v", err)
|
|
}
|
|
if len(output.NormalizeOutputs) != 2 {
|
|
t.Fatalf("normalize outputs = %d, want 2", len(output.NormalizeOutputs))
|
|
}
|
|
if got := jobMaximum.Load(); got != 3 {
|
|
t.Fatalf("maximum extract jobs = %d, want 3", got)
|
|
}
|
|
if got := provider.maximum.Load(); got != 2 {
|
|
t.Fatalf("maximum provider calls = %d, want 2", got)
|
|
}
|
|
if got := provider.calls.Load(); got != 13 {
|
|
t.Fatalf("provider calls = %d, want 13 including retry and validators", got)
|
|
}
|
|
}
|
|
|
|
func TestRunnerReturnsParentCancellationAndStopsQueuedExtracts(t *testing.T) {
|
|
prepared := preparedConcurrentPipeline(t, 4)
|
|
started := make(chan struct{}, 8)
|
|
for laneIndex := range prepared.Steps[0].lanes {
|
|
lane := laneIndex
|
|
installExtractOperation(prepared, lane, func(ctx context.Context, request contracts.TypedExtractionRequest) (erasedTypedResult, error) {
|
|
started <- struct{}{}
|
|
<-ctx.Done()
|
|
return erasedTypedResult{Value: typedValueForLane(lane, request.Chunk.Index)}, ctx.Err()
|
|
})
|
|
}
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
done := make(chan error, 1)
|
|
go func() {
|
|
_, err := New().Run(ctx, RunInput{Prepared: prepared, RawInput: []byte("input"), ExtractWorkers: 2})
|
|
done <- err
|
|
}()
|
|
<-started
|
|
<-started
|
|
cancel()
|
|
if err := <-done; !errors.Is(err, context.Canceled) {
|
|
t.Fatalf("Run() error = %v, want context.Canceled", err)
|
|
}
|
|
if got := len(started); got != 0 {
|
|
t.Fatalf("additional started jobs = %d, want none after cancellation", got)
|
|
}
|
|
}
|