gpu-turnstile/internal/lock/lock_test.go

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