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))
+}