296 lines
7.1 KiB
Go
296 lines
7.1 KiB
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)
|
|
}
|
|
}
|
|
|
|
func TestExternalHoldBlocksBoth(t *testing.T) {
|
|
lk := New(nil)
|
|
ctx := context.Background()
|
|
|
|
lk.SetExternal("game.exe (pid 42)")
|
|
if lk.TryAcquireLLM() {
|
|
t.Fatal("TryAcquireLLM succeeded during external hold")
|
|
}
|
|
if got := lk.External(); got != "game.exe (pid 42)" {
|
|
t.Fatalf("External() = %q", got)
|
|
}
|
|
if state, _, _ := lk.Snapshot(); state != StateExternal {
|
|
t.Fatalf("state=%s, want external", state)
|
|
}
|
|
|
|
llmAcquired := make(chan struct{})
|
|
go func() {
|
|
if err := lk.AcquireLLM(ctx); err != nil {
|
|
t.Error(err)
|
|
}
|
|
close(llmAcquired)
|
|
}()
|
|
assertBlocked(t, llmAcquired, "LLM acquire during external hold")
|
|
|
|
imageAcquired := make(chan struct{})
|
|
go func() {
|
|
if err := lk.AcquireImage(ctx); err != nil {
|
|
t.Error(err)
|
|
}
|
|
close(imageAcquired)
|
|
}()
|
|
assertBlocked(t, imageAcquired, "image acquire during external hold")
|
|
|
|
lk.ClearExternal()
|
|
// The queued image job wins over the LLM waiter (image priority).
|
|
waitFor(t, imageAcquired, "image acquire after ClearExternal")
|
|
assertBlocked(t, llmAcquired, "LLM acquire while image active")
|
|
lk.ReleaseImage()
|
|
waitFor(t, llmAcquired, "LLM acquire after image release")
|
|
lk.ReleaseLLM()
|
|
}
|
|
|
|
func TestExternalHoldDoesNotPreempt(t *testing.T) {
|
|
lk := New(nil)
|
|
ctx := context.Background()
|
|
if err := lk.AcquireLLM(ctx); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
|
|
lk.SetExternal("game.exe")
|
|
// In-flight LLM work keeps the llm state; the hold blocks new grants.
|
|
if state, _, _ := lk.Snapshot(); state != StateLLM {
|
|
t.Fatalf("state=%s, want llm while work in flight", state)
|
|
}
|
|
if lk.TryAcquireLLM() {
|
|
t.Fatal("TryAcquireLLM succeeded during external hold")
|
|
}
|
|
lk.ReleaseLLM()
|
|
if state, _, _ := lk.Snapshot(); state != StateExternal {
|
|
t.Fatalf("state=%s, want external after drain", state)
|
|
}
|
|
lk.ClearExternal()
|
|
if state, _, _ := lk.Snapshot(); state != StateIdle {
|
|
t.Fatalf("state=%s, want idle", state)
|
|
}
|
|
}
|