fix(admission): 避免高并发准入锁占满连接池

将同进程相同准入键的请求先在内存中串行化，跨节点使用 pg_try_advisory_xact_lock 非阻塞竞争并在事务外退避，避免等待 advisory lock 时长期占用 PostgreSQL 连接。\n\n新增 64 路竞争回归测试，验证被占用锁下连接池峰值仅保留持锁者和单个尝试者；256 路 Gemini Base64 端到端压力通过，advisory wait 峰值为 0。
2026-07-31 06:31:24 +08:00
parent 4fa44f4c41
commit 3d9ae74b87
3 changed files with 431 additions and 7 deletions
@@ -0,0 +1,170 @@
+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()
+	}
+}