easyai-ai-gateway/apps/api/internal/store/admission_lock.go

package store

import (
	"context"
	"errors"
	"sort"
	"sync"
	"time"

	"github.com/jackc/pgx/v5"
)

const (
	admissionLockRetryMin = 10 * time.Millisecond
	admissionLockRetryMax = 250 * time.Millisecond
)

var (
	errAdmissionLockBusy  = errors.New("task admission lock is busy")
	processAdmissionLocks = newAdmissionLocalLockSet()
)

type admissionLocalLock struct {
	token chan struct{}
	refs  int
}

type admissionLocalLockSet struct {
	mu    sync.Mutex
	locks map[string]*admissionLocalLock
}

func newAdmissionLocalLockSet() *admissionLocalLockSet {
	return &admissionLocalLockSet{locks: make(map[string]*admissionLocalLock)}
}

// acquire serializes contenders inside one API or Worker process before they
// consume a PostgreSQL connection. Sorted acquisition keeps operations that
// span task, platform-model, and user-group keys deadlock-free.
func (s *admissionLocalLockSet) acquire(ctx context.Context, keys []string) (func(), error) {
	keys = normalizedAdmissionLockKeys(keys)
	acquired := make([]struct {
		key  string
		lock *admissionLocalLock
	}, 0, len(keys))
	for _, key := range keys {
		lock := s.reference(key)
		select {
		case lock.token <- struct{}{}:
			acquired = append(acquired, struct {
				key  string
				lock *admissionLocalLock
			}{key: key, lock: lock})
		case <-ctx.Done():
			s.unreference(key, lock)
			for index := len(acquired) - 1; index >= 0; index-- {
				<-acquired[index].lock.token
				s.unreference(acquired[index].key, acquired[index].lock)
			}
			return nil, ctx.Err()
		}
	}
	return func() {
		for index := len(acquired) - 1; index >= 0; index-- {
			<-acquired[index].lock.token
			s.unreference(acquired[index].key, acquired[index].lock)
		}
	}, nil
}

func (s *admissionLocalLockSet) reference(key string) *admissionLocalLock {
	s.mu.Lock()
	defer s.mu.Unlock()
	lock := s.locks[key]
	if lock == nil {
		lock = &admissionLocalLock{token: make(chan struct{}, 1)}
		s.locks[key] = lock
	}
	lock.refs++
	return lock
}

func (s *admissionLocalLockSet) unreference(key string, lock *admissionLocalLock) {
	s.mu.Lock()
	defer s.mu.Unlock()
	lock.refs--
	if lock.refs == 0 && s.locks[key] == lock {
		delete(s.locks, key)
	}
}

func normalizedAdmissionLockKeys(keys []string) []string {
	seen := make(map[string]struct{}, len(keys))
	normalized := make([]string, 0, len(keys))
	for _, key := range keys {
		if key == "" {
			continue
		}
		if _, exists := seen[key]; exists {
			continue
		}
		seen[key] = struct{}{}
		normalized = append(normalized, key)
	}
	sort.Strings(normalized)
	return normalized
}

func admissionOperationLockKeys(input TaskAdmissionInput) []string {
	keys := []string{"task-admission:" + input.TaskID}
	for _, scope := range normalizedAdmissionScopes(input.Scopes) {
		keys = append(keys, admissionLockKey(scope))
	}
	return normalizedAdmissionLockKeys(keys)
}

func tryAdmissionTransactionLock(ctx context.Context, tx pgx.Tx, key string) error {
	var locked bool
	if err := tx.QueryRow(
		ctx,
		`SELECT pg_try_advisory_xact_lock(hashtextextended($1, 0))`,
		key,
	).Scan(&locked); err != nil {
		return err
	}
	if !locked {
		return errAdmissionLockBusy
	}
	return nil
}

func retryAdmissionOperation[T any](
	ctx context.Context,
	keys []string,
	operation func() (T, error),
) (T, error) {
	var zero T
	for attempt := 0; ; attempt++ {
		release, err := processAdmissionLocks.acquire(ctx, keys)
		if err != nil {
			return zero, err
		}
		result, operationErr := operation()
		release()
		if !errors.Is(operationErr, errAdmissionLockBusy) {
			return result, operationErr
		}
		if err := waitAdmissionLockRetry(ctx, attempt); err != nil {
			return zero, err
		}
	}
}

func waitAdmissionLockRetry(ctx context.Context, attempt int) error {
	delay := admissionLockRetryMin
	for index := 0; index < attempt && delay < admissionLockRetryMax; index++ {
		delay *= 2
	}
	if delay > admissionLockRetryMax {
		delay = admissionLockRetryMax
	}
	timer := time.NewTimer(delay)
	defer timer.Stop()
	select {
	case <-timer.C:
		return nil
	case <-ctx.Done():
		return ctx.Err()
	}
}