gpu-turnstile/internal/lock/lock.go

// Package lock implements the two-mode GPU arbitration lock: any number of
// concurrent LLM requests ("readers") or exactly one image job ("writer"),
// with image jobs taking priority over newly arriving LLM requests.
package lock

import (
	"context"
	"log/slog"
	"sync"
)

// State is the current GPU occupancy state.
type State string

const (
	StateIdle  State = "idle"
	StateLLM   State = "llm"
	StateImage State = "image"
)

type imageWaiter struct{ id uint64 }

// Lock is a writer-preferring two-mode lock. The zero value is not usable;
// construct with New.
type Lock struct {
	mu     sync.Mutex
	change chan struct{} // closed and replaced on every state change

	n           int  // LLM requests in flight
	imageActive bool // an image job holds the GPU
	imageQ      []imageWaiter
	nextID      uint64

	log *slog.Logger
}

// New returns a ready-to-use Lock. log may be nil; if set, every state
// transition is logged at debug level.
func New(log *slog.Logger) *Lock {
	return &Lock{change: make(chan struct{}), log: log}
}

// broadcast wakes all waiters. Call with mu held.
func (l *Lock) broadcast() {
	close(l.change)
	l.change = make(chan struct{})
}

func (l *Lock) logTransition(msg string, args ...any) {
	if l.log != nil {
		l.log.Debug(msg, args...)
	}
}

// AcquireLLM blocks until no image job is active or pending, then registers
// one in-flight LLM request. Returns ctx.Err() if the context is cancelled
// while waiting; no state is changed in that case.
func (l *Lock) AcquireLLM(ctx context.Context) error {
	l.mu.Lock()
	for l.imageActive || len(l.imageQ) > 0 {
		ch := l.change
		l.mu.Unlock()
		select {
		case <-ctx.Done():
			return ctx.Err()
		case <-ch:
		}
		l.mu.Lock()
	}
	l.n++
	n := l.n
	l.mu.Unlock()
	l.logTransition("lock transition", "state", StateLLM, "llm_inflight", n)
	return nil
}

// TryAcquireLLM acquires one in-flight LLM slot without waiting and
// reports whether it succeeded. It fails when an image job is active or
// pending.
func (l *Lock) TryAcquireLLM() bool {
	l.mu.Lock()
	if l.imageActive || len(l.imageQ) > 0 {
		l.mu.Unlock()
		return false
	}
	l.n++
	n := l.n
	l.mu.Unlock()
	l.logTransition("lock transition", "state", StateLLM, "llm_inflight", n)
	return true
}

// ReleaseLLM marks one LLM request as finished.
func (l *Lock) ReleaseLLM() {
	l.mu.Lock()
	l.n--
	n := l.n
	if l.n == 0 {
		l.broadcast()
	}
	l.mu.Unlock()
	if n == 0 {
		l.logTransition("lock transition", "state", StateIdle, "llm_inflight", 0)
	}
}

// AcquireImage queues the caller FIFO behind other image jobs, blocks new
// LLM requests immediately, and waits for in-flight LLM requests to drain
// before granting exclusive GPU access. Returns ctx.Err() if the context is
// cancelled while waiting; the caller is removed from the queue.
func (l *Lock) AcquireImage(ctx context.Context) error {
	l.mu.Lock()
	l.nextID++
	w := imageWaiter{id: l.nextID}
	l.imageQ = append(l.imageQ, w)
	pending := len(l.imageQ)
	l.broadcast() // new LLM requests must block from now on
	l.mu.Unlock()
	l.logTransition("image job pending", "image_queue", pending)

	for {
		l.mu.Lock()
		if l.imageQ[0].id == w.id && l.n == 0 && !l.imageActive {
			l.imageQ = l.imageQ[1:]
			l.imageActive = true
			l.mu.Unlock()
			l.logTransition("lock transition", "state", StateImage)
			return nil
		}
		ch := l.change
		l.mu.Unlock()
		select {
		case <-ctx.Done():
			l.mu.Lock()
			for i, q := range l.imageQ {
				if q.id == w.id {
					l.imageQ = append(l.imageQ[:i], l.imageQ[i+1:]...)
					break
				}
			}
			l.broadcast()
			l.mu.Unlock()
			l.logTransition("image job wait cancelled", "image_queue", len(l.imageQ))
			return ctx.Err()
		case <-ch:
		}
	}
}

// ReleaseImage frees the GPU after an image job.
func (l *Lock) ReleaseImage() {
	l.mu.Lock()
	l.imageActive = false
	l.broadcast()
	l.mu.Unlock()
	l.logTransition("lock transition", "state", StateIdle)
}

// Snapshot reports the current state, the number of in-flight LLM requests,
// and whether an image job is active or waiting.
func (l *Lock) Snapshot() (state State, llmInflight int, imagePending bool) {
	l.mu.Lock()
	defer l.mu.Unlock()
	switch {
	case l.imageActive:
		state = StateImage
	case l.n > 0:
		state = StateLLM
	default:
		state = StateIdle
	}
	return state, l.n, l.imageActive || len(l.imageQ) > 0
}