Use errgroup for bounded transcription work

internal/stage/transcribe.go

@@ -10,6 +10,8 @@ import (
 	"strings"
 	"sync"
 
+	"golang.org/x/sync/errgroup"
+
 	"gitea.maximumdirect.net/eric/narratio/internal/adapters/whisperx"
 	"gitea.maximumdirect.net/eric/narratio/internal/artifacts"
 	"gitea.maximumdirect.net/eric/narratio/internal/manifest"
@@ -93,106 +95,61 @@ func (transcribeStage) Run(ctx context.Context, env *Env, m *manifest.Manifest)
 		})
 	}
 
-	stageCtx, cancel := context.WithCancel(ctx)
-	defer cancel()
-
 	var (
-		wg        sync.WaitGroup
-		jobCh     = make(chan job)
 		mu        sync.Mutex
-		firstErr  error
 		completed int
 		perFile   = map[string]map[string]any{}
-		outputRef = map[string]artifacts.Ref{}
 	)
-
-	worker := func() {
-		defer wg.Done()
-		for j := range jobCh {
-			if stageCtx.Err() != nil {
-				return
+	group, groupCtx := errgroup.WithContext(ctx)
+	group.SetLimit(concurrency)
+	for _, job := range jobs {
+		job := job
+		group.Go(func() error {
+			if err := groupCtx.Err(); err != nil {
+				return err
 			}
-
-			res, runErr := env.WhisperX.Transcribe(stageCtx, whisperx.TranscribeRequest{
-				SpeakerID:               j.speakerID,
-				AudioPath:               j.audioPath,
-				OutputRawTranscriptPath: j.outPath,
+			res, runErr := env.WhisperX.Transcribe(groupCtx, whisperx.TranscribeRequest{
+				SpeakerID:               job.speakerID,
+				AudioPath:               job.audioPath,
+				OutputRawTranscriptPath: job.outPath,
 			})
 			if runErr != nil {
-				mu.Lock()
-				if firstErr == nil {
-					firstErr = fmt.Errorf("speaker %q (%s): %w", j.speakerID, filepath.Base(j.audioPath), runErr)
-					cancel()
-				}
-				mu.Unlock()
-				continue
+				return fmt.Errorf("speaker %q (%s): %w", job.speakerID, filepath.Base(job.audioPath), runErr)
 			}
 
-			if _, err := authoritativeOutputPath(j.outPath, res.OutputRawTranscriptPath); err != nil {
-				mu.Lock()
-				if firstErr == nil {
-					firstErr = fmt.Errorf("speaker %q: %w", j.speakerID, err)
-					cancel()
-				}
-				mu.Unlock()
-				continue
+			if _, err := authoritativeOutputPath(job.outPath, res.OutputRawTranscriptPath); err != nil {
+				return fmt.Errorf("speaker %q: %w", job.speakerID, err)
 			}
-			if err := validateTranscriptJSONFile(j.outPath); err != nil {
-				mu.Lock()
-				if firstErr == nil {
-					firstErr = fmt.Errorf("speaker %q output %q invalid: %w", j.speakerID, j.outPath, err)
-					cancel()
-				}
-				mu.Unlock()
-				continue
+			if err := validateTranscriptJSONFile(job.outPath); err != nil {
+				return fmt.Errorf("speaker %q output %q invalid: %w", job.speakerID, job.outPath, err)
 			}
 
 			meta := map[string]any{
 				"attempts":    res.Attempts,
 				"http_status": res.HTTPStatus,
 				"duration_ms": res.Duration.Milliseconds(),
-				"output_path": j.outPath,
+				"output_path": job.outPath,
 			}
 
 			mu.Lock()
-			perFile[j.speakerID] = meta
-			outputRef[j.speakerID] = artifacts.Ref{
-				Kind:         "transcript_raw",
-				Category:     "transcripts",
-				SessionID:    sessionID,
-				AbsolutePath: j.outPath,
-			}
+			perFile[job.speakerID] = meta
 			completed++
 			mu.Unlock()
-		}
-	}
-
-	wg.Add(concurrency)
-	for i := 0; i < concurrency; i++ {
-		go worker()
-	}
-
-	dispatched := 0
-dispatch:
-	for _, j := range jobs {
-		select {
-		case <-stageCtx.Done():
-			break dispatch
-		case jobCh <- j:
-			dispatched++
-		}
-	}
-	close(jobCh)
-	wg.Wait()
-
-	if firstErr != nil {
-		return nil, firstErr
+			return nil
+		})
 	}
+	groupErr := group.Wait()
+	mu.Lock()
+	completedCount := completed
+	mu.Unlock()
 	if err := ctx.Err(); err != nil {
-		return nil, fmt.Errorf("transcribe: canceled after dispatching %d of %d audio jobs and completing %d: %w", dispatched, len(jobs), completed, err)
+		return nil, fmt.Errorf("transcribe: canceled after completing %d of %d audio jobs: %w", completedCount, len(jobs), err)
 	}
-	if completed != len(jobs) {
-		return nil, fmt.Errorf("transcribe: completed %d of %d planned audio jobs after dispatching %d", completed, len(jobs), dispatched)
+	if groupErr != nil {
+		return nil, groupErr
+	}
+	if completedCount != len(jobs) {
+		return nil, fmt.Errorf("transcribe: completed %d of %d planned audio jobs", completedCount, len(jobs))
 	}
 
 	speakers := make([]string, 0, len(perFile))
@@ -206,10 +163,17 @@ dispatch:
 	outputPaths := make([]string, 0, len(speakers))
 	orderedPerFile := make(map[string]any, len(speakers))
 	for _, speaker := range speakers {
-		ref := outputRef[speaker]
-		runOutputPaths = append(runOutputPaths, ref.AbsolutePath)
+		var runOut string
+		for _, job := range jobs {
+			if job.speakerID == speaker {
+				runOut = job.outPath
+				break
+			}
+		}
+		runOutputPaths = append(runOutputPaths, runOut)
 		canonicalOut := filepath.Join(paths.TranscriptsRawDir, speaker+".json")
-		materialized, err := materializeRunLocalOutput(env.ArtifactStore, ref.AbsolutePath, canonicalOut, ref)
+		ref := artifacts.Ref{Kind: "transcript_raw", Category: "transcripts", SessionID: sessionID, AbsolutePath: runOut}
+		materialized, err := materializeRunLocalOutput(env.ArtifactStore, runOut, canonicalOut, ref)
 		if err != nil {
 			return nil, fmt.Errorf("transcribe: materialize %q output: %w", speaker, err)
 		}