perf(worker): 以有界微批次提升准入吞吐

原因：线上 P24 同构验收中，单任务同步复制提交与全局容量锁串行化，使两个 Worker 的 48 个执行槽只能维持约 10–14 个运行任务，最老等待超过 15 分钟。影响：新增可配置的 1–32 条准入微批次，默认 8；租约、任务准入与唯一 River job 在一个有界事务内原子提交，并按确定顺序预锁任务和容量范围，避免整窗 48 条大事务和多 Dispatcher 死锁。并容忍 rebind 已被其他 Dispatcher 完成的幂等竞态。验证：Go 全量测试、go vet、真实 PostgreSQL 跨 Store 集成测试、ShellCheck、迁移安全检查、OpenAPI、前端 lint/test/build、Compose/Kubernetes 渲染及人工发布脚本均通过。
2026-08-01 20:26:47 +08:00
parent 92e328a575
commit c89c56ca65
13 changed files with 282 additions and 31 deletions
@@ -502,6 +502,99 @@ func (s *Store) TryTaskAdmissionBatchWithAdmittedHook(
 	return outcomes, nil
 }

+// TryTaskAdmissionAtomicBatchWithAdmittedHook admits a deliberately small
+// group of tasks in one transaction. The caller bounds the group size so the
+// global capacity lock spans one synchronous commit without returning to the
+// old failure mode where a full capacity window locked dozens of task rows.
+// Admission leases and unique River jobs either commit together for the whole
+// microbatch or roll back together.
+func (s *Store) TryTaskAdmissionAtomicBatchWithAdmittedHook(
+	ctx context.Context,
+	inputs []TaskAdmissionInput,
+	onAdmitted func(pgx.Tx, TaskAdmissionInput) error,
+) ([]TaskAdmissionBatchOutcome, error) {
+	if len(inputs) == 0 {
+		return nil, nil
+	}
+	for _, input := range inputs {
+		if err := validateTaskAdmissionInput(input); err != nil {
+			return nil, err
+		}
+	}
+	lockKeys := make([]string, 0, len(inputs)*3)
+	for _, input := range inputs {
+		lockKeys = append(lockKeys, admissionOperationLockKeys(input)...)
+	}
+	return retryAdmissionOperation(ctx, lockKeys, func() ([]TaskAdmissionBatchOutcome, error) {
+		return s.tryTaskAdmissionAtomicBatchOnce(ctx, inputs, onAdmitted, lockKeys)
+	})
+}
+
+func (s *Store) tryTaskAdmissionAtomicBatchOnce(
+	ctx context.Context,
+	inputs []TaskAdmissionInput,
+	onAdmitted func(pgx.Tx, TaskAdmissionInput) error,
+	lockKeys []string,
+) ([]TaskAdmissionBatchOutcome, error) {
+	tx, err := s.pool.Begin(ctx)
+	if err != nil {
+		return nil, err
+	}
+	defer rollbackTransaction(tx)
+	// Lock every task and scope in deterministic order. Without this pre-lock,
+	// two dispatchers could each hold a different task lock while waiting for
+	// the shared worker-capacity lock.
+	for _, key := range normalizedAdmissionLockKeys(lockKeys) {
+		if err := tryAdmissionTransactionLock(ctx, tx, key); err != nil {
+			return nil, err
+		}
+	}
+
+	outcomes := make([]TaskAdmissionBatchOutcome, 0, len(inputs))
+	notify := false
+	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 {
+			outcomes = append(outcomes, TaskAdmissionBatchOutcome{
+				TaskID: input.TaskID,
+				Err:    admissionErr,
+			})
+			return outcomes, admissionErr
+		}
+		outcomes = append(outcomes, TaskAdmissionBatchOutcome{
+			TaskID: input.TaskID,
+			Result: outcome.Result,
+			Err:    outcome.PostCommitErr,
+		})
+		if outcome.NotifyTaskID != "" {
+			notify = true
+		}
+		if !outcome.Result.Admitted {
+			break
+		}
+	}
+	if err := tx.Commit(ctx); err != nil {
+		return outcomes, err
+	}
+	if notify {
+		s.notifyTaskAdmissionBestEffort(ctx, "*")
+	}
+	return outcomes, nil
+}
+
 func validateNewAdmissionCapacity(states []admissionScopeState) error {
 	mustWait := false
 	for _, state := range states {
@@ -649,6 +649,99 @@ SELECT
 		}
 	}

+	atomicScope := batchScope
+	atomicScope.ScopeKey = "atomic-microbatch-" + suffix
+	atomicScope.ConcurrentLimit = 4
+	atomicTasks := []GatewayTask{createTask(true), createTask(true), createTask(true), createTask(true)}
+	atomicInputs := make([]TaskAdmissionInput, 0, len(atomicTasks))
+	for index, task := range atomicTasks {
+		input := inputFor(task, 300+index, "")
+		input.Scopes = []AdmissionScope{atomicScope}
+		if _, err := first.QueueTaskAdmissionWithHook(ctx, input, nil); err != nil {
+			t.Fatalf("queue atomic microbatch task %d: %v", index, err)
+		}
+		atomicInputs = append(atomicInputs, input)
+	}
+	atomicOutcomes, err := first.TryTaskAdmissionAtomicBatchWithAdmittedHook(
+		ctx,
+		atomicInputs,
+		func(tx pgx.Tx, input TaskAdmissionInput) error {
+			_, hookErr := tx.Exec(ctx, `
+UPDATE gateway_tasks
+SET river_job_id = 987654600
+WHERE id = $1::uuid`, input.TaskID)
+			return hookErr
+		},
+	)
+	if err != nil || len(atomicOutcomes) != len(atomicTasks) {
+		t.Fatalf("atomic microbatch outcomes=%+v err=%v", atomicOutcomes, err)
+	}
+	var microbatchAdmissions, microbatchLeases, microbatchRiverJobs int
+	atomicTaskIDs := []string{atomicTasks[0].ID, atomicTasks[1].ID, atomicTasks[2].ID, atomicTasks[3].ID}
+	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),
+  (SELECT count(*) FROM gateway_tasks WHERE id = ANY($1::uuid[]) AND river_job_id IS NOT NULL)`,
+		atomicTaskIDs,
+	).Scan(&microbatchAdmissions, &microbatchLeases, &microbatchRiverJobs); err != nil {
+		t.Fatalf("read atomic microbatch: %v", err)
+	}
+	if microbatchAdmissions != 4 || microbatchLeases != 4 || microbatchRiverJobs != 4 {
+		t.Fatalf("atomic microbatch admissions=%d leases=%d jobs=%d, want 4/4/4",
+			microbatchAdmissions, microbatchLeases, microbatchRiverJobs)
+	}
+	for _, task := range atomicTasks {
+		if err := first.DeleteTaskAdmission(ctx, task.ID); err != nil {
+			t.Fatalf("release atomic microbatch task %s: %v", task.ID, err)
+		}
+	}
+
+	atomicRollbackScope := batchScope
+	atomicRollbackScope.ScopeKey = "atomic-rollback-" + suffix
+	atomicRollbackScope.ConcurrentLimit = 2
+	atomicRollbackTasks := []GatewayTask{createTask(true), createTask(true)}
+	atomicRollbackInputs := make([]TaskAdmissionInput, 0, len(atomicRollbackTasks))
+	for index, task := range atomicRollbackTasks {
+		input := inputFor(task, 400+index, "")
+		input.Scopes = []AdmissionScope{atomicRollbackScope}
+		if _, err := first.QueueTaskAdmissionWithHook(ctx, input, nil); err != nil {
+			t.Fatalf("queue atomic rollback task %d: %v", index, err)
+		}
+		atomicRollbackInputs = append(atomicRollbackInputs, input)
+	}
+	_, err = first.TryTaskAdmissionAtomicBatchWithAdmittedHook(
+		ctx,
+		atomicRollbackInputs,
+		func(_ pgx.Tx, input TaskAdmissionInput) error {
+			if input.TaskID == atomicRollbackTasks[1].ID {
+				return batchHookFailure
+			}
+			return nil
+		},
+	)
+	if !errors.Is(err, batchHookFailure) {
+		t.Fatalf("atomic rollback error=%v, want synthetic failure", err)
+	}
+	var atomicRollbackWaiting, atomicRollbackLeases 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)`,
+		[]string{atomicRollbackTasks[0].ID, atomicRollbackTasks[1].ID},
+	).Scan(&atomicRollbackWaiting, &atomicRollbackLeases); err != nil {
+		t.Fatalf("read atomic rollback: %v", err)
+	}
+	if atomicRollbackWaiting != 2 || atomicRollbackLeases != 0 {
+		t.Fatalf("atomic rollback waiting=%d leases=%d, want 2/0",
+			atomicRollbackWaiting, atomicRollbackLeases)
+	}
+	for _, task := range atomicRollbackTasks {
+		if err := first.DeleteTaskAdmission(ctx, task.ID); err != nil {
+			t.Fatalf("release atomic rollback 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)