fix(runner): 保留同平台重试的并发租约

同平台 retry_same 之前会在首个 attempt 失败时释放任务级 admission 租约,下一次 attempt 继续续租旧租约并进入 upstream_submission_unknown。

现在明确区分 attempt 自有租约与任务级 admission 租约,只在 attempt 结束时释放前者;补充回归测试覆盖同平台重试继续续租 admission 租约。

验证:env -u AI_GATEWAY_TEST_DATABASE_URL go test ./... -count=1;go vet ./...;gofmt;git diff --check。
This commit is contained in:
2026-08-03 21:06:23 +08:00
parent 3af6655e86
commit d129bcccbd
2 changed files with 40 additions and 8 deletions
+14 -8
View File
@@ -1357,20 +1357,18 @@ func (s *Service) runCandidate(
clientErr := &clients.ClientError{Code: "rate_limit", Message: err.Error(), Retryable: retryable}
return clients.Response{}, &localRateLimitError{clientErr: clientErr, cause: err, retryAfter: localRateLimitRetryAfter(err)}
}
attemptOwnedLeases := append([]store.ConcurrencyLease(nil), limitResult.Leases...)
if admittedPlatformModelID == candidate.PlatformModelID {
limitResult.Leases = append(limitResult.Leases, admittedLeases...)
limitResult.Leases, attemptOwnedLeases = bindAttemptConcurrencyLeases(limitResult.Leases, admittedLeases)
}
rateReservationsFinalized := false
retainAdmittedLeases := false
defer func() {
if !rateReservationsFinalized {
_ = s.store.ReleaseRateLimitReservations(context.WithoutCancel(ctx), limitResult.Reservations, "attempt_failed")
}
}()
defer func() {
if !retainAdmittedLeases {
_ = s.store.ReleaseConcurrencyLeases(context.WithoutCancel(ctx), limitResult.Leases)
}
_ = s.store.ReleaseConcurrencyLeases(context.WithoutCancel(ctx), attemptOwnedLeases)
}()
attemptID, err := s.store.CreateTaskAttempt(ctx, store.CreateTaskAttemptInput{
@@ -1798,12 +1796,20 @@ func (s *Service) runCandidate(
}); err != nil {
s.logger.Warn("record cache affinity observation failed", "error", err, "clientId", candidate.ClientID)
}
if admittedPlatformModelID == candidate.PlatformModelID {
retainAdmittedLeases = true
}
return response, nil
}
// bindAttemptConcurrencyLeases combines task-scoped admission leases with
// leases acquired by the current attempt. Only the latter are returned for
// attempt cleanup: admission owns its leases across retry_same and releases
// them when the task changes binding or reaches a terminal state.
func bindAttemptConcurrencyLeases(attemptLeases []store.ConcurrencyLease, admissionLeases []store.ConcurrencyLease) ([]store.ConcurrencyLease, []store.ConcurrencyLease) {
attemptOwned := append([]store.ConcurrencyLease(nil), attemptLeases...)
active := append([]store.ConcurrencyLease(nil), attemptLeases...)
active = append(active, admissionLeases...)
return active, attemptOwned
}
func (s *Service) observeProviderQuotaWait(metric string) {
observer, ok := s.billingMetrics.(interface {
ObserveProviderQuotaWait(string)