Finalize bounded typed pipeline implementation

internal/framework/pipeline/runner_concurrency_test.go

@@ -178,6 +178,59 @@ func TestRunnerStartsLaneContinuationWhileOtherLaneExtractsRemain(t *testing.T)
 	}
 }
 
+func TestRunnerBoundsConcurrentLaneContinuations(t *testing.T) {
+	prepared := preparedConcurrentPipeline(t, 1)
+	base := prepared.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.lanes = lanes
+	prepared.resolved.ArtifactLanes = resolvedLanes
+	prepared.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)

internal/framework/pipeline/runner_concurrent.go

@@ -98,6 +98,7 @@ func (r *Runner) runLanes(parent context.Context, input RunInput, checkpoints Ch
 	jobs := make(chan extractJob, workerCount)
 	results := make(chan extractJobResult, workerCount)
 	completions := make(chan laneCompletion, len(states))
+	continuations := make(chan *laneExtractState, workerCount)
 
 	var workers sync.WaitGroup
 	for i := 0; i < workerCount; i++ {
@@ -130,26 +131,46 @@ func (r *Runner) runLanes(parent context.Context, input RunInput, checkpoints Ch
 		}
 	}()
 	go func() { workers.Wait(); close(results) }()
+	var continuationWorkers sync.WaitGroup
+	for i := 0; i < workerCount; i++ {
+		continuationWorkers.Add(1)
+		go func() {
+			defer continuationWorkers.Done()
+			for state := range continuations {
+				if err := ctx.Err(); err != nil {
+					completions <- laneCompletion{index: state.index, err: err}
+					continue
+				}
+				laneOutput, err := r.continueLane(ctx, input, checkpoints, loader, doc, sourceInput, sessionID, chunks, state)
+				completions <- laneCompletion{index: state.index, output: laneOutput, err: err}
+			}
+		}()
+	}
 
 	completedOutputs := make([]RunOutput, len(states))
 	var runErrors []orderedRunError
+	var pendingContinuations []*laneExtractState
 	launched, completed := 0, 0
-	launch := func(state *laneExtractState) {
-		launched++
-		go func() {
-			laneOutput, err := r.continueLane(ctx, input, checkpoints, loader, doc, sourceInput, sessionID, chunks, state)
-			completions <- laneCompletion{index: state.index, output: laneOutput, err: err}
-		}()
-	}
 	for _, state := range states {
 		if state.decision.Reused {
-			launch(state)
+			pendingContinuations = append(pendingContinuations, state)
 		}
 	}
 
 	resultChannel := results
-	for resultChannel != nil || completed < launched {
+	for resultChannel != nil || len(pendingContinuations) > 0 || completed < launched {
+		var continuationChannel chan<- *laneExtractState
+		var nextContinuation *laneExtractState
+		if len(pendingContinuations) > 0 && ctx.Err() == nil {
+			continuationChannel = continuations
+			nextContinuation = pendingContinuations[0]
+		} else if ctx.Err() != nil {
+			pendingContinuations = nil
+		}
 		select {
+		case continuationChannel <- nextContinuation:
+			pendingContinuations = pendingContinuations[1:]
+			launched++
 		case result, ok := <-resultChannel:
 			if !ok {
 				resultChannel = nil
@@ -171,7 +192,7 @@ func (r *Runner) runLanes(parent context.Context, input RunInput, checkpoints Ch
 					runErrors = append(runErrors, orderedRunError{stage: 0, lane: state.index, chunk: len(chunks), err: err})
 					cancel()
 				} else {
-					launch(state)
+					pendingContinuations = append(pendingContinuations, state)
 				}
 			}
 		case completion := <-completions:
@@ -183,6 +204,8 @@ func (r *Runner) runLanes(parent context.Context, input RunInput, checkpoints Ch
 			}
 		}
 	}
+	close(continuations)
+	continuationWorkers.Wait()
 	for i := range completedOutputs {
 		mergeLaneOutput(&output, completedOutputs[i])
 	}