perf(worker): 批量填充异步执行槽

跨地域同步复制下逐任务准入会为每个 River Job 支付一次事务提交 RTT，导致 P24 实际只能维持少量运行任务。调度器现在按全局容量窗口准备 FIFO 批次，在同一 PostgreSQL 事务内逐项校验队首、创建租约并插入唯一 River Job，一次提交即可填满可用槽；任一 Hook 失败时整批原子回滚。\n\n验收压力采样同时排除正在终止的 Ready Pod，避免滚动切换期间误连已移除容器。\n\n验证：真实 PostgreSQL 批量提交及整批回滚集成测试、Go 全量测试、go vet、runner/store race、gofmt、bash -n、ShellCheck、迁移安全检查通过。
2026-07-31 07:59:57 +08:00
parent 14d13ad3a7
commit 5dd7765ae3
4 changed files with 364 additions and 82 deletions
@@ -85,6 +85,12 @@ type TaskAdmissionResult struct {
 	Leases        []ConcurrencyLease
 }

+type TaskAdmissionBatchOutcome struct {
+	TaskID string
+	Result TaskAdmissionResult
+	Err    error
+}
+
 type TaskAdmissionMetricsSnapshot struct {
 	QueueDepth             int
 	WaitingSync            int
@@ -254,23 +260,44 @@ func (s *Store) tryTaskAdmissionOnce(
 	}
 	defer rollbackTransaction(tx)

-	if err := tryAdmissionTransactionLock(ctx, tx, "task-admission:"+input.TaskID); err != nil {
+	outcome, err := s.tryTaskAdmissionTx(ctx, tx, input, onAdmitted)
+	if err != nil {
 		return TaskAdmissionResult{}, err
 	}
+	if err := tx.Commit(ctx); err != nil {
+		return TaskAdmissionResult{}, err
+	}
+	return outcome.Result, outcome.PostCommitErr
+}
+
+type taskAdmissionTxOutcome struct {
+	Result        TaskAdmissionResult
+	PostCommitErr error
+}
+
+func (s *Store) tryTaskAdmissionTx(
+	ctx context.Context,
+	tx pgx.Tx,
+	input TaskAdmissionInput,
+	onAdmitted func(pgx.Tx) error,
+) (taskAdmissionTxOutcome, error) {
+	if err := tryAdmissionTransactionLock(ctx, tx, "task-admission:"+input.TaskID); err != nil {
+		return taskAdmissionTxOutcome{}, err
+	}
 	var taskActive bool
 	if err := tx.QueryRow(ctx, `
 SELECT status IN ('queued', 'running')
 FROM gateway_tasks
 WHERE id = $1::uuid`, input.TaskID).Scan(&taskActive); err != nil {
-		return TaskAdmissionResult{}, err
+		return taskAdmissionTxOutcome{}, err
 	}
 	if !taskActive {
-		return TaskAdmissionResult{}, ErrTaskExecutionFinished
+		return taskAdmissionTxOutcome{}, ErrTaskExecutionFinished
 	}

 	admission, found, err := loadTaskAdmissionTx(ctx, tx, input.TaskID)
 	if err != nil {
-		return TaskAdmissionResult{}, err
+		return taskAdmissionTxOutcome{}, err
 	}
 	scopes := normalizedAdmissionScopes(input.Scopes)
 	lockScopes := append([]AdmissionScope{}, scopes...)
@@ -282,13 +309,13 @@ WHERE id = $1::uuid`, input.TaskID).Scan(&taskActive); err != nil {
 	}
 	for _, scope := range normalizedAdmissionScopes(lockScopes) {
 		if err := tryAdmissionTransactionLock(ctx, tx, admissionLockKey(scope)); err != nil {
-			return TaskAdmissionResult{}, err
+			return taskAdmissionTxOutcome{}, err
 		}
 	}
 	if found && admission.Status == "admitted" {
 		activeLeases, leaseErr := activeTaskAdmissionLeaseCountTx(ctx, tx, input.TaskID)
 		if leaseErr != nil {
-			return TaskAdmissionResult{}, leaseErr
+			return taskAdmissionTxOutcome{}, leaseErr
 		}
 		bindingChanged := admission.PlatformID != input.PlatformID ||
 			admission.PlatformModelID != input.PlatformModelID ||
@@ -297,78 +324,72 @@ WHERE id = $1::uuid`, input.TaskID).Scan(&taskActive); err != nil {
 			result := TaskAdmissionResult{Admission: admission, Admitted: true}
 			if onAdmitted != nil {
 				if err := onAdmitted(tx); err != nil {
-					return TaskAdmissionResult{}, err
+					return taskAdmissionTxOutcome{}, err
 				}
 			}
-			if err := tx.Commit(ctx); err != nil {
-				return TaskAdmissionResult{}, err
-			}
-			return result, nil
+			return taskAdmissionTxOutcome{Result: result}, nil
 		}
 		if bindingChanged {
 			targetStates, stateErr := admissionScopeStatesTx(ctx, tx, scopes)
 			if stateErr != nil {
-				return TaskAdmissionResult{}, stateErr
+				return taskAdmissionTxOutcome{}, stateErr
 			}
 			if admissionErr := validateNewAdmissionCapacity(targetStates); admissionErr != nil {
-				return TaskAdmissionResult{}, admissionErr
+				return taskAdmissionTxOutcome{}, admissionErr
 			}
 		}
 		admission, err = resetTaskAdmissionToWaitingTx(ctx, tx, input)
 		if err != nil {
-			return TaskAdmissionResult{}, err
+			return taskAdmissionTxOutcome{}, err
 		}
 	}
 	if found && !admission.WaitDeadlineAt.After(time.Now()) {
 		if err := expireTaskAdmissionTx(ctx, tx, input.TaskID); err != nil {
-			return TaskAdmissionResult{}, err
+			return taskAdmissionTxOutcome{}, err
 		}
 		if err := notifyTaskAdmissionTx(ctx, tx, "*"); err != nil {
-			return TaskAdmissionResult{}, err
+			return taskAdmissionTxOutcome{}, err
 		}
-		if err := tx.Commit(ctx); err != nil {
-			return TaskAdmissionResult{}, err
-		}
-		return TaskAdmissionResult{}, &QueueTimeoutError{TaskID: input.TaskID}
+		return taskAdmissionTxOutcome{
+			PostCommitErr: &QueueTimeoutError{TaskID: input.TaskID},
+		}, nil
 	}

 	scopeStates, err := admissionScopeStatesTx(ctx, tx, scopes)
 	if err != nil {
-		return TaskAdmissionResult{}, err
+		return taskAdmissionTxOutcome{}, err
 	}
 	if !found {
 		if admissionErr := validateNewAdmissionCapacity(scopeStates); admissionErr != nil {
-			return TaskAdmissionResult{}, admissionErr
+			return taskAdmissionTxOutcome{}, admissionErr
 		}
 		deadline := admissionDeadline(scopes)
 		admission, err = insertTaskAdmissionTx(ctx, tx, input, deadline)
 		if err != nil {
-			return TaskAdmissionResult{}, err
+			return taskAdmissionTxOutcome{}, err
 		}
 		found = true
 	} else if input.Mode == "sync" && strings.TrimSpace(input.WaiterID) != "" {
 		admission, err = renewTaskAdmissionWaiterTx(ctx, tx, input.TaskID, input.WaiterID)
 		if err != nil {
-			return TaskAdmissionResult{}, err
+			return taskAdmissionTxOutcome{}, err
 		}
 	}

 	head, err := isTaskAdmissionHeadTx(ctx, tx, admission, scopes)
 	if err != nil {
-		return TaskAdmissionResult{}, err
+		return taskAdmissionTxOutcome{}, err
 	}
 	if !head {
-		if err := tx.Commit(ctx); err != nil {
-			return TaskAdmissionResult{}, err
-		}
-		return TaskAdmissionResult{Admission: admission}, nil
+		return taskAdmissionTxOutcome{
+			Result: TaskAdmissionResult{Admission: admission},
+		}, nil
 	}
 	for _, state := range scopeStates {
 		if state.Saturated {
-			if err := tx.Commit(ctx); err != nil {
-				return TaskAdmissionResult{}, err
-			}
-			return TaskAdmissionResult{Admission: admission}, nil
+			return taskAdmissionTxOutcome{
+				Result: TaskAdmissionResult{Admission: admission},
+			}, nil
 		}
 	}

@@ -390,34 +411,98 @@ WHERE id = $1::uuid`, input.TaskID).Scan(&taskActive); err != nil {
 		})
 		if err != nil {
 			if errors.Is(err, ErrRateLimited) {
-				if err := tx.Commit(ctx); err != nil {
-					return TaskAdmissionResult{}, err
-				}
-				return TaskAdmissionResult{Admission: admission}, nil
+				return taskAdmissionTxOutcome{
+					Result: TaskAdmissionResult{Admission: admission},
+				}, nil
 			}
-			return TaskAdmissionResult{}, err
+			return taskAdmissionTxOutcome{}, err
 		}
 		leases = append(leases, lease)
 	}
 	admission, err = markTaskAdmissionAdmittedTx(ctx, tx, input.TaskID)
 	if err != nil {
-		return TaskAdmissionResult{}, err
+		return taskAdmissionTxOutcome{}, err
 	}
 	// Continue the FIFO chain after one task is admitted. Every API process
 	// only probes the global head it owns, so free capacity is filled without
 	// broadcasting an advisory-lock attempt to every waiter.
 	if err := notifyTaskAdmissionTx(ctx, tx, "*"); err != nil {
-		return TaskAdmissionResult{}, err
+		return taskAdmissionTxOutcome{}, err
 	}
 	if onAdmitted != nil {
 		if err := onAdmitted(tx); err != nil {
-			return TaskAdmissionResult{}, err
+			return taskAdmissionTxOutcome{}, err
 		}
 	}
-	if err := tx.Commit(ctx); err != nil {
-		return TaskAdmissionResult{}, err
+	return taskAdmissionTxOutcome{
+		Result: TaskAdmissionResult{
+			Admission:     admission,
+			Admitted:      true,
+			NewlyAdmitted: true,
+			Leases:        leases,
+		},
+	}, nil
+}
+
+// TryTaskAdmissionBatchWithAdmittedHook fills a FIFO capacity window in one
+// transaction. The hook inserts every unique River job in that same
+// transaction, preserving the admission/job crash boundary while avoiding one
+// synchronous replication round trip per task.
+func (s *Store) TryTaskAdmissionBatchWithAdmittedHook(
+	ctx context.Context,
+	inputs []TaskAdmissionInput,
+	onAdmitted func(pgx.Tx, TaskAdmissionInput) error,
+) ([]TaskAdmissionBatchOutcome, error) {
+	if len(inputs) == 0 {
+		return nil, nil
 	}
-	return TaskAdmissionResult{Admission: admission, Admitted: true, NewlyAdmitted: true, Leases: leases}, nil
+	lockKeys := make([]string, 0, len(inputs)*3)
+	for _, input := range inputs {
+		if err := validateTaskAdmissionInput(input); err != nil {
+			return nil, err
+		}
+		lockKeys = append(lockKeys, admissionOperationLockKeys(input)...)
+	}
+	return retryAdmissionOperation(ctx, lockKeys, func() ([]TaskAdmissionBatchOutcome, error) {
+		tx, err := s.pool.Begin(ctx)
+		if err != nil {
+			return nil, err
+		}
+		defer rollbackTransaction(tx)
+
+		outcomes := make([]TaskAdmissionBatchOutcome, 0, len(inputs))
+		for _, input := range inputs {
+			hook := func(tx pgx.Tx) error {
+				if onAdmitted == nil {
+					return nil
+				}
+				return onAdmitted(tx, input)
+			}
+			outcome, admissionErr := s.tryTaskAdmissionTx(ctx, tx, input, hook)
+			if errors.Is(admissionErr, ErrTaskExecutionFinished) {
+				outcomes = append(outcomes, TaskAdmissionBatchOutcome{
+					TaskID: input.TaskID,
+					Err:    admissionErr,
+				})
+				continue
+			}
+			if admissionErr != nil {
+				return nil, admissionErr
+			}
+			outcomes = append(outcomes, TaskAdmissionBatchOutcome{
+				TaskID: input.TaskID,
+				Result: outcome.Result,
+				Err:    outcome.PostCommitErr,
+			})
+			if outcome.PostCommitErr == nil && !outcome.Result.Admitted {
+				break
+			}
+		}
+		if err := tx.Commit(ctx); err != nil {
+			return nil, err
+		}
+		return outcomes, nil
+	})
 }

 func validateNewAdmissionCapacity(states []admissionScopeState) error {
@@ -501,6 +501,140 @@ WHERE id = $1::uuid`, recoveryGraceTask.ID); err != nil {
 	if !recoveryGraceFound {
 		t.Fatal("generic River recovery did not claim a task beyond the recovery grace period")
 	}
+
+	batchScope := AdmissionScope{
+		ScopeType:       "worker_capacity",
+		ScopeKey:        "batch-" + suffix,
+		ScopeName:       "batch capacity",
+		ConcurrentLimit: 3,
+		Amount:          1,
+		LeaseTTLSeconds: 120,
+		QueueLimit:      100,
+		MaxWaitSeconds:  600,
+	}
+	batchTasks := []GatewayTask{createTask(true), createTask(true), createTask(true)}
+	batchInputs := make([]TaskAdmissionInput, 0, len(batchTasks))
+	batchJobIDs := make(map[string]int64, len(batchTasks))
+	for index, task := range batchTasks {
+		input := inputFor(task, 100+index, "")
+		input.Scopes = []AdmissionScope{batchScope}
+		if _, err := first.QueueTaskAdmissionWithHook(ctx, input, nil); err != nil {
+			t.Fatalf("queue batch task %d: %v", index, err)
+		}
+		batchInputs = append(batchInputs, input)
+		batchJobIDs[task.ID] = int64(987654400 + index)
+	}
+	batchOutcomes, err := first.TryTaskAdmissionBatchWithAdmittedHook(
+		ctx,
+		batchInputs,
+		func(tx pgx.Tx, input TaskAdmissionInput) error {
+			_, hookErr := tx.Exec(ctx, `
+UPDATE gateway_tasks
+SET river_job_id = $2
+WHERE id = $1::uuid`, input.TaskID, batchJobIDs[input.TaskID])
+			return hookErr
+		},
+	)
+	if err != nil || len(batchOutcomes) != len(batchTasks) {
+		t.Fatalf("batch admission outcomes=%+v err=%v", batchOutcomes, err)
+	}
+	for index, outcome := range batchOutcomes {
+		if outcome.Err != nil || !outcome.Result.Admitted || !outcome.Result.NewlyAdmitted ||
+			len(outcome.Result.Leases) != 1 {
+			t.Fatalf("batch outcome %d=%+v", index, outcome)
+		}
+	}
+	var batchAdmissions, batchLeases, batchRiverJobs int
+	if err := first.pool.QueryRow(ctx, `
+SELECT
+  (SELECT count(*)
+   FROM gateway_task_admissions
+   WHERE task_id = ANY($1::uuid[]) AND status = 'admitted'),
+  (SELECT count(*)
+   FROM gateway_concurrency_leases
+   WHERE task_id = ANY($1::uuid[]) AND released_at IS NULL AND expires_at > now()),
+  (SELECT count(*)
+   FROM gateway_tasks
+   WHERE id = ANY($1::uuid[]) AND river_job_id IS NOT NULL)`,
+		[]string{batchTasks[0].ID, batchTasks[1].ID, batchTasks[2].ID},
+	).Scan(&batchAdmissions, &batchLeases, &batchRiverJobs); err != nil {
+		t.Fatalf("read committed batch admission: %v", err)
+	}
+	if batchAdmissions != 3 || batchLeases != 3 || batchRiverJobs != 3 {
+		t.Fatalf(
+			"batch admissions=%d leases=%d River jobs=%d, want 3/3/3",
+			batchAdmissions,
+			batchLeases,
+			batchRiverJobs,
+		)
+	}
+	for _, task := range batchTasks {
+		if err := first.DeleteTaskAdmission(ctx, task.ID); err != nil {
+			t.Fatalf("release batch task %s: %v", task.ID, err)
+		}
+	}
+
+	rollbackScope := batchScope
+	rollbackScope.ScopeKey = "batch-rollback-" + suffix
+	rollbackScope.ConcurrentLimit = 2
+	rollbackTasks := []GatewayTask{createTask(true), createTask(true)}
+	rollbackInputs := make([]TaskAdmissionInput, 0, len(rollbackTasks))
+	for index, task := range rollbackTasks {
+		input := inputFor(task, 200+index, "")
+		input.Scopes = []AdmissionScope{rollbackScope}
+		if _, err := first.QueueTaskAdmissionWithHook(ctx, input, nil); err != nil {
+			t.Fatalf("queue rollback batch task %d: %v", index, err)
+		}
+		rollbackInputs = append(rollbackInputs, input)
+	}
+	batchHookFailure := errors.New("synthetic batch hook failure")
+	_, err = first.TryTaskAdmissionBatchWithAdmittedHook(
+		ctx,
+		rollbackInputs,
+		func(tx pgx.Tx, input TaskAdmissionInput) error {
+			if input.TaskID == rollbackTasks[1].ID {
+				return batchHookFailure
+			}
+			_, hookErr := tx.Exec(ctx, `
+UPDATE gateway_tasks
+SET river_job_id = 987654499
+WHERE id = $1::uuid`, input.TaskID)
+			return hookErr
+		},
+	)
+	if !errors.Is(err, batchHookFailure) {
+		t.Fatalf("batch hook failure=%v, want synthetic failure", err)
+	}
+	var rollbackWaiting, rollbackLeases, rollbackRiverJobs int
+	if err := first.pool.QueryRow(ctx, `
+SELECT
+  (SELECT count(*)
+   FROM gateway_task_admissions
+   WHERE task_id = ANY($1::uuid[]) AND status = 'waiting'),
+  (SELECT count(*)
+   FROM gateway_concurrency_leases
+   WHERE task_id = ANY($1::uuid[]) AND released_at IS NULL),
+  (SELECT count(*)
+   FROM gateway_tasks
+   WHERE id = ANY($1::uuid[]) AND river_job_id IS NOT NULL)`,
+		[]string{rollbackTasks[0].ID, rollbackTasks[1].ID},
+	).Scan(&rollbackWaiting, &rollbackLeases, &rollbackRiverJobs); err != nil {
+		t.Fatalf("read rolled back batch admission: %v", err)
+	}
+	if rollbackWaiting != 2 || rollbackLeases != 0 || rollbackRiverJobs != 0 {
+		t.Fatalf(
+			"rolled back batch waiting=%d leases=%d River jobs=%d, want 2/0/0",
+			rollbackWaiting,
+			rollbackLeases,
+			rollbackRiverJobs,
+		)
+	}
+	for _, task := range rollbackTasks {
+		if err := first.DeleteTaskAdmission(ctx, task.ID); err != nil {
+			t.Fatalf("release rollback batch task %s: %v", task.ID, err)
+		}
+	}
+
 	result, err = second.TryTaskAdmission(ctx, inputFor(queuedAtomicTask, 100, ""))
 	if err != nil || !result.Admitted || len(result.Leases) != 1 {
 		t.Fatalf("worker-time queued admission result=%+v err=%v", result, err)