gpulock: GPU arbitration proxy for Ollama + ComfyUI

Implements SPEC.md: two listeners, one writer-preferring two-mode lock,
Ollama unload before image jobs, ComfyUI history polling + VRAM free,
optional model warm-up, healthz/metrics endpoints, streaming-safe
reverse proxies, Dockerfile and Gitea Actions CI.

internal/lock/lock_test.go

@@ -0,0 +1,228 @@
+package lock
+
+import (
+	"context"
+	"sync"
+	"testing"
+	"time"
+)
+
+func waitFor(t *testing.T, ch <-chan struct{}, what string) {
+	t.Helper()
+	select {
+	case <-ch:
+	case <-time.After(2 * time.Second):
+		t.Fatalf("timed out waiting for %s", what)
+	}
+}
+
+func assertBlocked(t *testing.T, ch <-chan struct{}, what string) {
+	t.Helper()
+	select {
+	case <-ch:
+		t.Fatalf("%s should still be blocked", what)
+	case <-time.After(50 * time.Millisecond):
+	}
+}
+
+func TestLLMConcurrency(t *testing.T) {
+	lk := New(nil)
+	ctx := context.Background()
+	for i := 0; i < 3; i++ {
+		if err := lk.AcquireLLM(ctx); err != nil {
+			t.Fatal(err)
+		}
+	}
+	state, n, pending := lk.Snapshot()
+	if state != StateLLM || n != 3 || pending {
+		t.Fatalf("got state=%s n=%d pending=%v", state, n, pending)
+	}
+	lk.ReleaseLLM()
+	lk.ReleaseLLM()
+	if state, _, _ := lk.Snapshot(); state != StateLLM {
+		t.Fatalf("state=%s, want llm", state)
+	}
+	lk.ReleaseLLM()
+	if state, n, _ := lk.Snapshot(); state != StateIdle || n != 0 {
+		t.Fatalf("got state=%s n=%d, want idle", state, n)
+	}
+}
+
+func TestImageWaitsForLLMDrainAndBlocksNewLLM(t *testing.T) {
+	lk := New(nil)
+	ctx := context.Background()
+
+	if err := lk.AcquireLLM(ctx); err != nil {
+		t.Fatal(err)
+	}
+	if err := lk.AcquireLLM(ctx); err != nil {
+		t.Fatal(err)
+	}
+
+	imageAcquired := make(chan struct{})
+	go func() {
+		if err := lk.AcquireImage(ctx); err != nil {
+			t.Error(err)
+		}
+		close(imageAcquired)
+	}()
+	assertBlocked(t, imageAcquired, "image acquire while LLMs in flight")
+
+	// While an image job is pending, new LLM requests must block even
+	// though LLM concurrency is otherwise allowed.
+	llmAcquired := make(chan struct{})
+	go func() {
+		if err := lk.AcquireLLM(ctx); err != nil {
+			t.Error(err)
+		}
+		close(llmAcquired)
+	}()
+	assertBlocked(t, llmAcquired, "LLM acquire while image pending")
+
+	lk.ReleaseLLM()
+	assertBlocked(t, imageAcquired, "image acquire with one LLM still in flight")
+
+	lk.ReleaseLLM()
+	waitFor(t, imageAcquired, "image acquire after drain")
+	assertBlocked(t, llmAcquired, "LLM acquire while image active")
+
+	state, _, pending := lk.Snapshot()
+	if state != StateImage || !pending {
+		t.Fatalf("got state=%s pending=%v", state, pending)
+	}
+
+	lk.ReleaseImage()
+	waitFor(t, llmAcquired, "LLM acquire after image release")
+	lk.ReleaseLLM()
+}
+
+func TestImageFIFO(t *testing.T) {
+	lk := New(nil)
+	ctx := context.Background()
+
+	if err := lk.AcquireImage(ctx); err != nil {
+		t.Fatal(err)
+	}
+
+	order := make(chan int, 2)
+	for i, id := range []int{1, 2} {
+		go func() {
+			if err := lk.AcquireImage(ctx); err != nil {
+				t.Error(err)
+				return
+			}
+			order <- id
+		}()
+		// Ensure this waiter is queued before the next one starts.
+		for deadline := time.Now().Add(2 * time.Second); ; {
+			lk.mu.Lock()
+			qlen := len(lk.imageQ)
+			lk.mu.Unlock()
+			if qlen == i+1 {
+				break
+			}
+			if time.Now().After(deadline) {
+				t.Fatalf("image waiter %d never queued", id)
+			}
+			time.Sleep(time.Millisecond)
+		}
+	}
+
+	lk.ReleaseImage()
+	if got := <-order; got != 1 {
+		t.Fatalf("first image job = %d, want 1", got)
+	}
+	select {
+	case got := <-order:
+		t.Fatalf("second image job %d acquired while first active", got)
+	case <-time.After(50 * time.Millisecond):
+	}
+	lk.ReleaseImage()
+	if got := <-order; got != 2 {
+		t.Fatalf("second image job = %d, want 2", got)
+	}
+	lk.ReleaseImage()
+}
+
+func TestContextCancelRemovesLLMWaiter(t *testing.T) {
+	lk := New(nil)
+	if err := lk.AcquireImage(context.Background()); err != nil {
+		t.Fatal(err)
+	}
+
+	ctx, cancel := context.WithCancel(context.Background())
+	done := make(chan error, 1)
+	go func() { done <- lk.AcquireLLM(ctx) }()
+	time.Sleep(20 * time.Millisecond)
+	cancel()
+	if err := <-done; err == nil {
+		t.Fatal("expected context error")
+	}
+
+	lk.ReleaseImage()
+	if state, n, _ := lk.Snapshot(); state != StateIdle || n != 0 {
+		t.Fatalf("got state=%s n=%d, want idle", state, n)
+	}
+}
+
+func TestContextCancelRemovesImageWaiter(t *testing.T) {
+	lk := New(nil)
+	ctx := context.Background()
+	if err := lk.AcquireLLM(ctx); err != nil {
+		t.Fatal(err)
+	}
+
+	wctx, cancel := context.WithCancel(ctx)
+	done := make(chan error, 1)
+	go func() { done <- lk.AcquireImage(wctx) }()
+	time.Sleep(20 * time.Millisecond)
+	cancel()
+	if err := <-done; err == nil {
+		t.Fatal("expected context error")
+	}
+
+	lk.ReleaseLLM()
+	// The cancelled waiter must be gone: a new LLM request acquires
+	// immediately instead of blocking behind it.
+	if err := lk.AcquireLLM(ctx); err != nil {
+		t.Fatal(err)
+	}
+	if _, _, pending := lk.Snapshot(); pending {
+		t.Fatal("image_pending=true after waiter was cancelled")
+	}
+	lk.ReleaseLLM()
+}
+
+// TestRace hammers the lock from many goroutines; run with -race.
+func TestRace(t *testing.T) {
+	lk := New(nil)
+	var wg sync.WaitGroup
+	for i := 0; i < 8; i++ {
+		wg.Add(1)
+		go func(i int) {
+			defer wg.Done()
+			for j := 0; j < 25; j++ {
+				if i%4 == 3 {
+					if err := lk.AcquireImage(context.Background()); err != nil {
+						t.Error(err)
+						return
+					}
+					lk.Snapshot()
+					lk.ReleaseImage()
+				} else {
+					if err := lk.AcquireLLM(context.Background()); err != nil {
+						t.Error(err)
+						return
+					}
+					lk.Snapshot()
+					lk.ReleaseLLM()
+				}
+			}
+		}(i)
+	}
+	wg.Wait()
+	state, n, pending := lk.Snapshot()
+	if state != StateIdle || n != 0 || pending {
+		t.Fatalf("leaked lock state: state=%s n=%d pending=%v", state, n, pending)
+	}
+}