BCAI/easyai-ai-gateway
3
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

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

Browse Source 
将同进程相同准入键的请求先在内存中串行化，跨节点使用 pg_try_advisory_xact_lock 非阻塞竞争并在事务外退避，避免等待 advisory lock 时长期占用 PostgreSQL 连接。\n\n新增 64 路竞争回归测试，验证被占用锁下连接池峰值仅保留持锁者和单个尝试者；256 路 Gemini Base64 端到端压力通过，advisory wait 峰值为 0。
This commit is contained in:
wangbo
2026-07-31 06:31:24 +08:00
parent 4fa44f4c41
commit 3d9ae74b87
3 changed files with 431 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/api/internal/store/admission_queue.go
+40 -7
View File
@@ -130,13 +130,23 @@ func (s *Store) QueueTaskAdmissionWithHook(
if input.Mode != "async" {
return TaskAdmission{}, errors.New("queued admission registration requires asynchronous mode")
}
return retryAdmissionOperation(ctx, admissionOperationLockKeys(input), func() (TaskAdmission, error) {
return s.queueTaskAdmissionOnce(ctx, input, onQueued)
})
}
func (s *Store) queueTaskAdmissionOnce(
ctx context.Context,
input TaskAdmissionInput,
onQueued func(pgx.Tx) error,
) (TaskAdmission, error) {
tx, err := s.pool.Begin(ctx)
if err != nil {
return TaskAdmission{}, err
}
defer rollbackTransaction(tx)
if _, err := tx.Exec(ctx, `SELECT pg_advisory_xact_lock(hashtextextended($1, 0))`, "task-admission:"+input.TaskID); err != nil {
if err := tryAdmissionTransactionLock(ctx, tx, "task-admission:"+input.TaskID); err != nil {
return TaskAdmission{}, err
}
var taskActive bool
@@ -163,7 +173,7 @@ WHERE id = $1::uuid`, input.TaskID).Scan(&taskActive); err != nil {
)
}
for _, scope := range normalizedAdmissionScopes(lockScopes) {
if _, err := tx.Exec(ctx, `SELECT pg_advisory_xact_lock(hashtextextended($1, 0))`, admissionLockKey(scope)); err != nil {
if err := tryAdmissionTransactionLock(ctx, tx, admissionLockKey(scope)); err != nil {
return TaskAdmission{}, err
}
}
@@ -228,13 +238,23 @@ func (s *Store) tryTaskAdmission(
if err := validateTaskAdmissionInput(input); err != nil {
return TaskAdmissionResult{}, err
}
return retryAdmissionOperation(ctx, admissionOperationLockKeys(input), func() (TaskAdmissionResult, error) {
return s.tryTaskAdmissionOnce(ctx, input, onAdmitted)
})
}
func (s *Store) tryTaskAdmissionOnce(
ctx context.Context,
input TaskAdmissionInput,
onAdmitted func(pgx.Tx) error,
) (TaskAdmissionResult, error) {
tx, err := s.pool.Begin(ctx)
if err != nil {
return TaskAdmissionResult{}, err
}
defer rollbackTransaction(tx)
if _, err := tx.Exec(ctx, `SELECT pg_advisory_xact_lock(hashtextextended($1, 0))`, "task-admission:"+input.TaskID); err != nil {
if err := tryAdmissionTransactionLock(ctx, tx, "task-admission:"+input.TaskID); err != nil {
return TaskAdmissionResult{}, err
}
var taskActive bool
@@ -261,7 +281,7 @@ WHERE id = $1::uuid`, input.TaskID).Scan(&taskActive); err != nil {
)
}
for _, scope := range normalizedAdmissionScopes(lockScopes) {
if _, err := tx.Exec(ctx, `SELECT pg_advisory_xact_lock(hashtextextended($1, 0))`, admissionLockKey(scope)); err != nil {
if err := tryAdmissionTransactionLock(ctx, tx, admissionLockKey(scope)); err != nil {
return TaskAdmissionResult{}, err
}
}
@@ -452,12 +472,18 @@ func (s *Store) RebindWaitingTaskAdmission(ctx context.Context, input TaskAdmiss
if err := validateTaskAdmissionInput(input); err != nil {
return TaskAdmission{}, err
}
return retryAdmissionOperation(ctx, admissionOperationLockKeys(input), func() (TaskAdmission, error) {
return s.rebindWaitingTaskAdmissionOnce(ctx, input)
})
}
func (s *Store) rebindWaitingTaskAdmissionOnce(ctx context.Context, input TaskAdmissionInput) (TaskAdmission, error) {
tx, err := s.pool.Begin(ctx)
if err != nil {
return TaskAdmission{}, err
}
defer rollbackTransaction(tx)
if _, err := tx.Exec(ctx, `SELECT pg_advisory_xact_lock(hashtextextended($1, 0))`, "task-admission:"+input.TaskID); err != nil {
if err := tryAdmissionTransactionLock(ctx, tx, "task-admission:"+input.TaskID); err != nil {
return TaskAdmission{}, err
}
current, found, err := loadTaskAdmissionTx(ctx, tx, input.TaskID)
@@ -477,7 +503,7 @@ func (s *Store) RebindWaitingTaskAdmission(ctx context.Context, input TaskAdmiss
AdmissionScope{ScopeType: "user_group", ScopeKey: current.UserGroupID},
)
for _, scope := range normalizedAdmissionScopes(lockScopes) {
if _, err := tx.Exec(ctx, `SELECT pg_advisory_xact_lock(hashtextextended($1, 0))`, admissionLockKey(scope)); err != nil {
if err := tryAdmissionTransactionLock(ctx, tx, admissionLockKey(scope)); err != nil {
return TaskAdmission{}, err
}
}
@@ -629,12 +655,19 @@ WHERE admission.task_id = owned.task_id
}
func (s *Store) DeleteTaskAdmission(ctx context.Context, taskID string) error {
_, err := retryAdmissionOperation(ctx, []string{"task-admission:" + taskID}, func() (struct{}, error) {
return struct{}{}, s.deleteTaskAdmissionOnce(ctx, taskID)
})
return err
}
func (s *Store) deleteTaskAdmissionOnce(ctx context.Context, taskID string) error {
tx, err := s.pool.Begin(ctx)
if err != nil {
return err
}
defer rollbackTransaction(tx)
if _, err := tx.Exec(ctx, `SELECT pg_advisory_xact_lock(hashtextextended($1, 0))`, "task-admission:"+taskID); err != nil {
if err := tryAdmissionTransactionLock(ctx, tx, "task-admission:"+taskID); err != nil {
return err
}
if _, err := tx.Exec(ctx, `