Review LLM concurrency refactor

internal/framework/llm/scheduler_test.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"errors"
 	"reflect"
+	"runtime"
 	"sync"
 	"sync/atomic"
 	"testing"
@@ -42,7 +43,7 @@ func TestSchedulerFIFOOrdering(t *testing.T) {
 				t.Errorf("Run[%d] error: %v", id, runErr)
 			}
 		}()
-		time.Sleep(10 * time.Millisecond)
+		waitForQueueDepth(t, s, i+1)
 	}
 
 	firstRelease()
@@ -69,6 +70,7 @@ func TestSchedulerEnforcesMaxConcurrency(t *testing.T) {
 
 	var inFlight int32
 	var maxInFlight int32
+	release := make(chan struct{})
 	var wg sync.WaitGroup
 
 	for i := 0; i < 12; i++ {
@@ -84,7 +86,7 @@ func TestSchedulerEnforcesMaxConcurrency(t *testing.T) {
 						break
 					}
 				}
-				time.Sleep(20 * time.Millisecond)
+				<-release
 				atomic.AddInt32(&inFlight, -1)
 				return nil
 			})
@@ -94,6 +96,9 @@ func TestSchedulerEnforcesMaxConcurrency(t *testing.T) {
 		}()
 	}
 
+	waitForMinInFlight(t, &maxInFlight, 2)
+	close(release)
+
 	wg.Wait()
 	if got := atomic.LoadInt32(&maxInFlight); got > 2 {
 		t.Fatalf("expected max in-flight <= 2, got %d", got)
@@ -186,3 +191,33 @@ func TestSchedulerNoPermitLeakAfterQueuedCancellation(t *testing.T) {
 		t.Fatalf("expected scheduler to accept new work after cancellation, got %v", err)
 	}
 }
+
+func waitForQueueDepth(t *testing.T, s *Scheduler, want int) {
+	t.Helper()
+	deadline := time.Now().Add(250 * time.Millisecond)
+	for time.Now().Before(deadline) {
+		s.mu.Lock()
+		depth := len(s.queue)
+		s.mu.Unlock()
+		if depth >= want {
+			return
+		}
+		runtime.Gosched()
+	}
+	s.mu.Lock()
+	depth := len(s.queue)
+	s.mu.Unlock()
+	t.Fatalf("timed out waiting for queue depth >= %d (got %d)", want, depth)
+}
+
+func waitForMinInFlight(t *testing.T, maxInFlight *int32, want int32) {
+	t.Helper()
+	deadline := time.Now().Add(250 * time.Millisecond)
+	for time.Now().Before(deadline) {
+		if atomic.LoadInt32(maxInFlight) >= want {
+			return
+		}
+		runtime.Gosched()
+	}
+	t.Fatalf("timed out waiting for max in-flight >= %d (got %d)", want, atomic.LoadInt32(maxInFlight))
+}

internal/framework/proposal_generation/generate_test.go

@@ -6,6 +6,7 @@ import (
 	"os"
 	"path/filepath"
 	"reflect"
+	"runtime"
 	"strings"
 	"sync"
 	"sync/atomic"
@@ -54,6 +55,8 @@ func (s *countingScheduler) Run(ctx context.Context, fn func(context.Context) er
 type sleepingStructuredClient struct {
 	inFlight    int32
 	maxInFlight int32
+	entered     chan struct{}
+	release     chan struct{}
 }
 
 func (c *sleepingStructuredClient) CompleteStructured(ctx context.Context, req contracts.StructuredCompletionRequest, out any) (contracts.StructuredCompletionResponse, error) {
@@ -65,8 +68,11 @@ func (c *sleepingStructuredClient) CompleteStructured(ctx context.Context, req c
 			break
 		}
 	}
+	if c.entered != nil {
+		c.entered <- struct{}{}
+	}
 	select {
-	case <-time.After(20 * time.Millisecond):
+	case <-c.release:
 	case <-ctx.Done():
 		atomic.AddInt32(&c.inFlight, -1)
 		return contracts.StructuredCompletionResponse{}, ctx.Err()
@@ -309,7 +315,10 @@ func TestGenerateCandidatesRespectsSchedulerConcurrency(t *testing.T) {
 	if err != nil {
 		t.Fatalf("NewScheduler: %v", err)
 	}
-	client := &sleepingStructuredClient{}
+	client := &sleepingStructuredClient{
+		entered: make(chan struct{}, 16),
+		release: make(chan struct{}),
+	}
 	baseReq := defaultRequest(t)
 	baseReq.LLMClient = client
 	baseReq.Scheduler = scheduler
@@ -326,6 +335,9 @@ func TestGenerateCandidatesRespectsSchedulerConcurrency(t *testing.T) {
 			}
 		}(i)
 	}
+
+	waitForEntries(t, client.entered, 2)
+	close(client.release)
 	wg.Wait()
 
 	if got := atomic.LoadInt32(&client.maxInFlight); got > 2 {
@@ -335,3 +347,20 @@ func TestGenerateCandidatesRespectsSchedulerConcurrency(t *testing.T) {
 		t.Fatalf("expected observed concurrency of at least 2, got %d", got)
 	}
 }
+
+func waitForEntries(t *testing.T, entered <-chan struct{}, want int) {
+	t.Helper()
+	deadline := time.Now().Add(300 * time.Millisecond)
+	got := 0
+	for got < want && time.Now().Before(deadline) {
+		select {
+		case <-entered:
+			got++
+		default:
+			runtime.Gosched()
+		}
+	}
+	if got < want {
+		t.Fatalf("timed out waiting for %d entered calls (got %d)", want, got)
+	}
+}

internal/framework/runner/runner_test.go

@@ -179,6 +179,8 @@ func TestRunnerProposalSectionConcurrencyBoundedByProposalLLMConcurrency(t *test
 
 	var inFlight int32
 	var maxInFlight int32
+	entered := make(chan struct{}, len(transcript.Segments))
+	release := make(chan struct{})
 	r := New(fakeFactory{modules: map[string]contracts.TranscriptModule{
 		"m": fakeModule{key: "m", policy: proposals.ReplacementPolicyRequireUnique, proposeF: func(req contracts.ProposalRequest) ([]proposals.CorrectionProposal, error) {
 			current := atomic.AddInt32(&inFlight, 1)
@@ -188,19 +190,29 @@ func TestRunnerProposalSectionConcurrencyBoundedByProposalLLMConcurrency(t *test
 					break
 				}
 			}
-			time.Sleep(25 * time.Millisecond)
+			entered <- struct{}{}
+			<-release
 			atomic.AddInt32(&inFlight, -1)
 			return nil, nil
 		}},
 	}})
 
-	_, err := r.Run(context.Background(), RunInput{
-		Config:     &cfg,
-		Transcript: transcript,
-		ModuleSpecs: []contracts.ModuleRunSpec{
-			{ModuleKey: "m", InstanceName: "m"},
-		},
-	})
+	resultCh := make(chan error, 1)
+	go func() {
+		_, err := r.Run(context.Background(), RunInput{
+			Config:     &cfg,
+			Transcript: transcript,
+			ModuleSpecs: []contracts.ModuleRunSpec{
+				{ModuleKey: "m", InstanceName: "m"},
+			},
+		})
+		resultCh <- err
+	}()
+
+	waitForRunnerEntries(t, entered, 2, "proposal workers to enter")
+	close(release)
+
+	err := <-resultCh
 	if err != nil {
 		t.Fatalf("Run error: %v", err)
 	}
@@ -225,27 +237,52 @@ func TestRunnerProposalIndexOrderingIsDeterministicAcrossParallelSections(t *tes
 		{ID: 3, Text: "charlie three"},
 	}}
 
+	started := make(chan int, len(transcript.Segments))
+	releaseByID := map[int]chan struct{}{
+		1: make(chan struct{}),
+		2: make(chan struct{}),
+		3: make(chan struct{}),
+	}
+
 	r := New(fakeFactory{modules: map[string]contracts.TranscriptModule{
 		"m": fakeModule{key: "m", policy: proposals.ReplacementPolicyRequireUnique, proposeF: func(req contracts.ProposalRequest) ([]proposals.CorrectionProposal, error) {
 			if len(req.WorkingTranscript.Segments) != 1 {
 				t.Fatalf("expected one segment per section, got %d", len(req.WorkingTranscript.Segments))
 			}
 			seg := req.WorkingTranscript.Segments[0]
-			// Sleep inversely by ID to force completion order to differ from section order.
-			time.Sleep(time.Duration(4-seg.ID) * 10 * time.Millisecond)
+			started <- seg.ID
+			<-releaseByID[seg.ID]
 			return []proposals.CorrectionProposal{
 				{TargetSegmentID: seg.ID, OriginalText: seg.Text, CorrectedText: strings.ToUpper(seg.Text), Confidence: 1},
 			}, nil
 		}},
 	}})
 
-	out, err := r.Run(context.Background(), RunInput{
-		Config:     &cfg,
-		Transcript: transcript,
-		ModuleSpecs: []contracts.ModuleRunSpec{
-			{ModuleKey: "m", InstanceName: "m"},
-		},
-	})
+	resultCh := make(chan struct {
+		out RunOutput
+		err error
+	}, 1)
+	go func() {
+		out, err := r.Run(context.Background(), RunInput{
+			Config:     &cfg,
+			Transcript: transcript,
+			ModuleSpecs: []contracts.ModuleRunSpec{
+				{ModuleKey: "m", InstanceName: "m"},
+			},
+		})
+		resultCh <- struct {
+			out RunOutput
+			err error
+		}{out: out, err: err}
+	}()
+
+	waitForRunnerSectionIDs(t, started, map[int]struct{}{1: {}, 2: {}, 3: {}})
+	close(releaseByID[3])
+	close(releaseByID[2])
+	close(releaseByID[1])
+
+	result := <-resultCh
+	out, err := result.out, result.err
 	if err != nil {
 		t.Fatalf("Run error: %v", err)
 	}
@@ -269,6 +306,42 @@ func TestRunnerProposalIndexOrderingIsDeterministicAcrossParallelSections(t *tes
 	}
 }
 
+func waitForRunnerEntries(t *testing.T, entered <-chan struct{}, want int, label string) {
+	t.Helper()
+	deadline := time.After(350 * time.Millisecond)
+	got := 0
+	for got < want {
+		select {
+		case <-entered:
+			got++
+		case <-deadline:
+			t.Fatalf("timed out waiting for %d entries (%s), got %d", want, label, got)
+		}
+	}
+}
+
+func waitForRunnerSectionIDs(t *testing.T, started <-chan int, want map[int]struct{}) {
+	t.Helper()
+	deadline := time.After(350 * time.Millisecond)
+	seen := map[int]struct{}{}
+	for len(seen) < len(want) {
+		select {
+		case id := <-started:
+			seen[id] = struct{}{}
+		case <-deadline:
+			t.Fatalf("timed out waiting for section IDs %v, got %v", mapKeys(want), mapKeys(seen))
+		}
+	}
+}
+
+func mapKeys(values map[int]struct{}) []int {
+	keys := make([]int, 0, len(values))
+	for key := range values {
+		keys = append(keys, key)
+	}
+	return keys
+}
+
 func TestRunnerSkippedRecorded(t *testing.T) {
 	transcript := &schema.Transcript{Segments: []schema.Segment{{ID: 1, Speaker: "A", Start: 0, End: 1, Text: "word word"}}}
 	r := New(fakeFactory{modules: map[string]contracts.TranscriptModule{

internal/framework/validators/llm_validators_test.go

@@ -3,6 +3,7 @@ package validators
 import (
 	"context"
 	"errors"
+	"runtime"
 	"strings"
 	"sync"
 	"sync/atomic"
@@ -24,6 +25,8 @@ type fakeStructuredLLMClient struct {
 type sleepingValidationClient struct {
 	inFlight    int32
 	maxInFlight int32
+	entered     chan struct{}
+	release     chan struct{}
 }
 
 type boundedScheduler struct {
@@ -53,8 +56,11 @@ func (c *sleepingValidationClient) CompleteStructured(ctx context.Context, req S
 			break
 		}
 	}
+	if c.entered != nil {
+		c.entered <- struct{}{}
+	}
 	select {
-	case <-time.After(20 * time.Millisecond):
+	case <-c.release:
 	case <-ctx.Done():
 		atomic.AddInt32(&c.inFlight, -1)
 		return StructuredCompletionResponse{}, ctx.Err()
@@ -245,7 +251,10 @@ func TestLLMBackedValidatorUnknownProposalIndexFails(t *testing.T) {
 
 func TestLLMBackedValidatorRespectsSchedulerConcurrency(t *testing.T) {
 	scheduler := newBoundedScheduler(2)
-	client := &sleepingValidationClient{}
+	client := &sleepingValidationClient{
+		entered: make(chan struct{}, 16),
+		release: make(chan struct{}),
+	}
 	v, err := NewLLMBackedValidator("spoken_form_plausibility_review", LLMValidatorTypeSpokenFormPlausibility, "test-model")
 	if err != nil {
 		t.Fatalf("new validator error: %v", err)
@@ -264,6 +273,8 @@ func TestLLMBackedValidatorRespectsSchedulerConcurrency(t *testing.T) {
 			}
 		}()
 	}
+	waitForValidationEntries(t, client.entered, 2)
+	close(client.release)
 	wg.Wait()
 
 	if got := atomic.LoadInt32(&client.maxInFlight); got > 2 {
@@ -273,3 +284,20 @@ func TestLLMBackedValidatorRespectsSchedulerConcurrency(t *testing.T) {
 		t.Fatalf("expected observed concurrency of at least 2, got %d", got)
 	}
 }
+
+func waitForValidationEntries(t *testing.T, entered <-chan struct{}, want int) {
+	t.Helper()
+	deadline := time.Now().Add(300 * time.Millisecond)
+	got := 0
+	for got < want && time.Now().Before(deadline) {
+		select {
+		case <-entered:
+			got++
+		default:
+			runtime.Gosched()
+		}
+	}
+	if got < want {
+		t.Fatalf("timed out waiting for %d validator calls to enter (got %d)", want, got)
+	}
+}