diff --git a/option/ccm.go b/option/ccm.go index dd55a4ba4..6d9d14423 100644 --- a/option/ccm.go +++ b/option/ccm.go @@ -91,9 +91,10 @@ type CCMDefaultCredentialOptions struct { } type CCMBalancerCredentialOptions struct { - Strategy string `json:"strategy,omitempty"` - Credentials badoption.Listable[string] `json:"credentials"` - PollInterval badoption.Duration `json:"poll_interval,omitempty"` + Strategy string `json:"strategy,omitempty"` + Credentials badoption.Listable[string] `json:"credentials"` + PollInterval badoption.Duration `json:"poll_interval,omitempty"` + RebalanceThreshold float64 `json:"rebalance_threshold,omitempty"` } type CCMExternalCredentialOptions struct { diff --git a/option/ocm.go b/option/ocm.go index e508abae7..6015b506c 100644 --- a/option/ocm.go +++ b/option/ocm.go @@ -91,9 +91,10 @@ type OCMDefaultCredentialOptions struct { } type OCMBalancerCredentialOptions struct { - Strategy string `json:"strategy,omitempty"` - Credentials badoption.Listable[string] `json:"credentials"` - PollInterval badoption.Duration `json:"poll_interval,omitempty"` + Strategy string `json:"strategy,omitempty"` + Credentials badoption.Listable[string] `json:"credentials"` + PollInterval badoption.Duration `json:"poll_interval,omitempty"` + RebalanceThreshold float64 `json:"rebalance_threshold,omitempty"` } type OCMExternalCredentialOptions struct { diff --git a/service/ccm/credential_external.go b/service/ccm/credential_external.go index a5781d6f6..ade2d8361 100644 --- a/service/ccm/credential_external.go +++ b/service/ccm/credential_external.go @@ -465,9 +465,9 @@ func (c *externalCredential) wrapRequestContext(parent context.Context) *credent cancel() }) return &credentialRequestContext{ - Context: derived, - releaseFunc: stop, - cancelFunc: cancel, + Context: derived, + releaseFuncs: []func() bool{stop}, + cancelFunc: cancel, } } diff --git a/service/ccm/credential_state.go b/service/ccm/credential_state.go index 83db5a197..490c6148f 100644 --- a/service/ccm/credential_state.go +++ b/service/ccm/credential_state.go @@ -110,15 +110,21 @@ type defaultCredential struct { type credentialRequestContext struct { context.Context - releaseOnce sync.Once - cancelOnce sync.Once - releaseFunc func() bool - cancelFunc context.CancelFunc + releaseOnce sync.Once + cancelOnce sync.Once + releaseFuncs []func() bool + cancelFunc context.CancelFunc +} + +func (c *credentialRequestContext) addInterruptLink(stop func() bool) { + c.releaseFuncs = append(c.releaseFuncs, stop) } func (c *credentialRequestContext) releaseCredentialInterrupt() { c.releaseOnce.Do(func() { - c.releaseFunc() + for _, f := range c.releaseFuncs { + f() + } }) } @@ -504,9 +510,9 @@ func (c *defaultCredential) wrapRequestContext(parent context.Context) *credenti cancel() }) return &credentialRequestContext{ - Context: derived, - releaseFunc: stop, - cancelFunc: cancel, + Context: derived, + releaseFuncs: []func() bool{stop}, + cancelFunc: cancel, } } @@ -851,6 +857,7 @@ func (c *defaultCredential) buildProxyRequest(ctx context.Context, original *htt type credentialProvider interface { selectCredential(sessionID string, filter func(credential) bool) (credential, bool, error) onRateLimited(sessionID string, cred credential, resetAt time.Time, filter func(credential) bool) credential + wrapProviderInterrupt(cred credential, requestContext *credentialRequestContext) pollIfStale(ctx context.Context) allCredentials() []credential close() @@ -913,6 +920,8 @@ func (p *singleCredentialProvider) allCredentials() []credential { return []credential{p.cred} } +func (p *singleCredentialProvider) wrapProviderInterrupt(_ credential, _ *credentialRequestContext) {} + func (p *singleCredentialProvider) close() {} const sessionExpiry = 24 * time.Hour @@ -922,27 +931,37 @@ type sessionEntry struct { createdAt time.Time } -// balancerProvider assigns sessions to credentials based on a configurable strategy. -type balancerProvider struct { - credentials []credential - strategy string - roundRobinIndex atomic.Uint64 - pollInterval time.Duration - sessionMutex sync.RWMutex - sessions map[string]sessionEntry - logger log.ContextLogger +type credentialInterruptEntry struct { + context context.Context + cancel context.CancelFunc } -func newBalancerProvider(credentials []credential, strategy string, pollInterval time.Duration, logger log.ContextLogger) *balancerProvider { +// balancerProvider assigns sessions to credentials based on a configurable strategy. +type balancerProvider struct { + credentials []credential + strategy string + roundRobinIndex atomic.Uint64 + pollInterval time.Duration + rebalanceThreshold float64 + sessionMutex sync.RWMutex + sessions map[string]sessionEntry + interruptAccess sync.Mutex + credentialInterrupts map[string]credentialInterruptEntry + logger log.ContextLogger +} + +func newBalancerProvider(credentials []credential, strategy string, pollInterval time.Duration, rebalanceThreshold float64, logger log.ContextLogger) *balancerProvider { if pollInterval <= 0 { pollInterval = defaultPollInterval } return &balancerProvider{ - credentials: credentials, - strategy: strategy, - pollInterval: pollInterval, - sessions: make(map[string]sessionEntry), - logger: logger, + credentials: credentials, + strategy: strategy, + pollInterval: pollInterval, + rebalanceThreshold: rebalanceThreshold, + sessions: make(map[string]sessionEntry), + credentialInterrupts: make(map[string]credentialInterruptEntry), + logger: logger, } } @@ -954,6 +973,20 @@ func (p *balancerProvider) selectCredential(sessionID string, filter func(creden if exists { for _, cred := range p.credentials { if cred.tagName() == entry.tag && (filter == nil || filter(cred)) && cred.isUsable() { + if p.rebalanceThreshold > 0 && (p.strategy == "" || p.strategy == "least_used") { + better := p.pickLeastUsed(filter) + if better != nil && better.tagName() != cred.tagName() { + effectiveThreshold := p.rebalanceThreshold / cred.planWeight() + delta := cred.weeklyUtilization() - better.weeklyUtilization() + if delta > effectiveThreshold { + p.logger.Info("rebalancing away from ", cred.tagName(), + ": utilization delta ", delta, "% exceeds effective threshold ", + effectiveThreshold, "% (weight ", cred.planWeight(), ")") + p.rebalanceCredential(cred.tagName()) + break + } + } + } return cred, false, nil } } @@ -977,6 +1010,40 @@ func (p *balancerProvider) selectCredential(sessionID string, filter func(creden return best, isNew, nil } +func (p *balancerProvider) rebalanceCredential(tag string) { + p.interruptAccess.Lock() + if entry, loaded := p.credentialInterrupts[tag]; loaded { + entry.cancel() + } + ctx, cancel := context.WithCancel(context.Background()) + p.credentialInterrupts[tag] = credentialInterruptEntry{context: ctx, cancel: cancel} + p.interruptAccess.Unlock() + + p.sessionMutex.Lock() + for id, entry := range p.sessions { + if entry.tag == tag { + delete(p.sessions, id) + } + } + p.sessionMutex.Unlock() +} + +func (p *balancerProvider) wrapProviderInterrupt(cred credential, requestContext *credentialRequestContext) { + tag := cred.tagName() + p.interruptAccess.Lock() + entry, loaded := p.credentialInterrupts[tag] + if !loaded { + ctx, cancel := context.WithCancel(context.Background()) + entry = credentialInterruptEntry{context: ctx, cancel: cancel} + p.credentialInterrupts[tag] = entry + } + p.interruptAccess.Unlock() + stop := context.AfterFunc(entry.context, func() { + requestContext.cancelOnce.Do(requestContext.cancelFunc) + }) + requestContext.addInterruptLink(stop) +} + func (p *balancerProvider) onRateLimited(sessionID string, cred credential, resetAt time.Time, filter func(credential) bool) credential { cred.markRateLimited(resetAt) if sessionID != "" { @@ -1166,6 +1233,8 @@ func (p *fallbackProvider) allCredentials() []credential { return p.credentials } +func (p *fallbackProvider) wrapProviderInterrupt(_ credential, _ *credentialRequestContext) {} + func (p *fallbackProvider) close() {} func allCredentialsUnavailableError(credentials []credential) error { @@ -1247,7 +1316,7 @@ func buildCredentialProviders( if err != nil { return nil, nil, err } - providers[credOpt.Tag] = newBalancerProvider(subCredentials, credOpt.BalancerOptions.Strategy, time.Duration(credOpt.BalancerOptions.PollInterval), logger) + providers[credOpt.Tag] = newBalancerProvider(subCredentials, credOpt.BalancerOptions.Strategy, time.Duration(credOpt.BalancerOptions.PollInterval), credOpt.BalancerOptions.RebalanceThreshold, logger) case "fallback": subCredentials, err := resolveCredentialTags(credOpt.FallbackOptions.Credentials, allCredentialMap, credOpt.Tag) if err != nil { @@ -1346,6 +1415,9 @@ func validateCCMOptions(options option.CCMServiceOptions) error { default: return E.New("credential ", cred.Tag, ": unknown balancer strategy: ", cred.BalancerOptions.Strategy) } + if cred.BalancerOptions.RebalanceThreshold < 0 { + return E.New("credential ", cred.Tag, ": rebalance_threshold must not be negative") + } } } diff --git a/service/ccm/service.go b/service/ccm/service.go index 81e3b38a5..2760c348a 100644 --- a/service/ccm/service.go +++ b/service/ccm/service.go @@ -459,6 +459,7 @@ func (s *Service) ServeHTTP(w http.ResponseWriter, r *http.Request) { } requestContext := selectedCredential.wrapRequestContext(r.Context()) + provider.wrapProviderInterrupt(selectedCredential, requestContext) defer func() { requestContext.cancelRequest() }() @@ -497,6 +498,7 @@ func (s *Service) ServeHTTP(w http.ResponseWriter, r *http.Request) { s.logger.InfoContext(ctx, "retrying with credential ", nextCredential.tagName(), " after 429 from ", selectedCredential.tagName()) requestContext.cancelRequest() requestContext = nextCredential.wrapRequestContext(r.Context()) + provider.wrapProviderInterrupt(nextCredential, requestContext) retryRequest, buildErr := nextCredential.buildProxyRequest(requestContext, r, bodyBytes, s.httpHeaders) if buildErr != nil { s.logger.ErrorContext(ctx, "retry request: ", buildErr) diff --git a/service/ocm/credential_external.go b/service/ocm/credential_external.go index 2a09c84ad..3d46f8fa8 100644 --- a/service/ocm/credential_external.go +++ b/service/ocm/credential_external.go @@ -502,9 +502,9 @@ func (c *externalCredential) wrapRequestContext(parent context.Context) *credent cancel() }) return &credentialRequestContext{ - Context: derived, - releaseFunc: stop, - cancelFunc: cancel, + Context: derived, + releaseFuncs: []func() bool{stop}, + cancelFunc: cancel, } } diff --git a/service/ocm/credential_state.go b/service/ocm/credential_state.go index ea7c621bf..7a6c8ef5e 100644 --- a/service/ocm/credential_state.go +++ b/service/ocm/credential_state.go @@ -110,15 +110,21 @@ type defaultCredential struct { type credentialRequestContext struct { context.Context - releaseOnce sync.Once - cancelOnce sync.Once - releaseFunc func() bool - cancelFunc context.CancelFunc + releaseOnce sync.Once + cancelOnce sync.Once + releaseFuncs []func() bool + cancelFunc context.CancelFunc +} + +func (c *credentialRequestContext) addInterruptLink(stop func() bool) { + c.releaseFuncs = append(c.releaseFuncs, stop) } func (c *credentialRequestContext) releaseCredentialInterrupt() { c.releaseOnce.Do(func() { - c.releaseFunc() + for _, f := range c.releaseFuncs { + f() + } }) } @@ -518,9 +524,9 @@ func (c *defaultCredential) wrapRequestContext(parent context.Context) *credenti cancel() }) return &credentialRequestContext{ - Context: derived, - releaseFunc: stop, - cancelFunc: cancel, + Context: derived, + releaseFuncs: []func() bool{stop}, + cancelFunc: cancel, } } @@ -848,6 +854,7 @@ func (c *defaultCredential) buildProxyRequest(ctx context.Context, original *htt type credentialProvider interface { selectCredential(sessionID string, filter func(credential) bool) (credential, bool, error) onRateLimited(sessionID string, cred credential, resetAt time.Time, filter func(credential) bool) credential + wrapProviderInterrupt(cred credential, requestContext *credentialRequestContext) pollIfStale(ctx context.Context) allCredentials() []credential close() @@ -909,6 +916,8 @@ func (p *singleCredentialProvider) allCredentials() []credential { return []credential{p.cred} } +func (p *singleCredentialProvider) wrapProviderInterrupt(_ credential, _ *credentialRequestContext) {} + func (p *singleCredentialProvider) close() {} const sessionExpiry = 24 * time.Hour @@ -918,30 +927,40 @@ type sessionEntry struct { createdAt time.Time } +type credentialInterruptEntry struct { + context context.Context + cancel context.CancelFunc +} + type balancerProvider struct { - credentials []credential - strategy string - roundRobinIndex atomic.Uint64 - pollInterval time.Duration - sessionMutex sync.RWMutex - sessions map[string]sessionEntry - logger log.ContextLogger + credentials []credential + strategy string + roundRobinIndex atomic.Uint64 + pollInterval time.Duration + rebalanceThreshold float64 + sessionMutex sync.RWMutex + sessions map[string]sessionEntry + interruptAccess sync.Mutex + credentialInterrupts map[string]credentialInterruptEntry + logger log.ContextLogger } func compositeCredentialSelectable(cred credential) bool { return !cred.ocmIsAPIKeyMode() } -func newBalancerProvider(credentials []credential, strategy string, pollInterval time.Duration, logger log.ContextLogger) *balancerProvider { +func newBalancerProvider(credentials []credential, strategy string, pollInterval time.Duration, rebalanceThreshold float64, logger log.ContextLogger) *balancerProvider { if pollInterval <= 0 { pollInterval = defaultPollInterval } return &balancerProvider{ - credentials: credentials, - strategy: strategy, - pollInterval: pollInterval, - sessions: make(map[string]sessionEntry), - logger: logger, + credentials: credentials, + strategy: strategy, + pollInterval: pollInterval, + rebalanceThreshold: rebalanceThreshold, + sessions: make(map[string]sessionEntry), + credentialInterrupts: make(map[string]credentialInterruptEntry), + logger: logger, } } @@ -953,6 +972,20 @@ func (p *balancerProvider) selectCredential(sessionID string, filter func(creden if exists { for _, cred := range p.credentials { if cred.tagName() == entry.tag && compositeCredentialSelectable(cred) && (filter == nil || filter(cred)) && cred.isUsable() { + if p.rebalanceThreshold > 0 && (p.strategy == "" || p.strategy == "least_used") { + better := p.pickLeastUsed(filter) + if better != nil && better.tagName() != cred.tagName() { + effectiveThreshold := p.rebalanceThreshold / cred.planWeight() + delta := cred.weeklyUtilization() - better.weeklyUtilization() + if delta > effectiveThreshold { + p.logger.Info("rebalancing away from ", cred.tagName(), + ": utilization delta ", delta, "% exceeds effective threshold ", + effectiveThreshold, "% (weight ", cred.planWeight(), ")") + p.rebalanceCredential(cred.tagName()) + break + } + } + } return cred, false, nil } } @@ -976,6 +1009,40 @@ func (p *balancerProvider) selectCredential(sessionID string, filter func(creden return best, isNew, nil } +func (p *balancerProvider) rebalanceCredential(tag string) { + p.interruptAccess.Lock() + if entry, loaded := p.credentialInterrupts[tag]; loaded { + entry.cancel() + } + ctx, cancel := context.WithCancel(context.Background()) + p.credentialInterrupts[tag] = credentialInterruptEntry{context: ctx, cancel: cancel} + p.interruptAccess.Unlock() + + p.sessionMutex.Lock() + for id, entry := range p.sessions { + if entry.tag == tag { + delete(p.sessions, id) + } + } + p.sessionMutex.Unlock() +} + +func (p *balancerProvider) wrapProviderInterrupt(cred credential, requestContext *credentialRequestContext) { + tag := cred.tagName() + p.interruptAccess.Lock() + entry, loaded := p.credentialInterrupts[tag] + if !loaded { + ctx, cancel := context.WithCancel(context.Background()) + entry = credentialInterruptEntry{context: ctx, cancel: cancel} + p.credentialInterrupts[tag] = entry + } + p.interruptAccess.Unlock() + stop := context.AfterFunc(entry.context, func() { + requestContext.cancelOnce.Do(requestContext.cancelFunc) + }) + requestContext.addInterruptLink(stop) +} + func (p *balancerProvider) onRateLimited(sessionID string, cred credential, resetAt time.Time, filter func(credential) bool) credential { cred.markRateLimited(resetAt) if sessionID != "" { @@ -1169,6 +1236,8 @@ func (p *fallbackProvider) allCredentials() []credential { return p.credentials } +func (p *fallbackProvider) wrapProviderInterrupt(_ credential, _ *credentialRequestContext) {} + func (p *fallbackProvider) close() {} func allRateLimitedError(credentials []credential) error { @@ -1232,7 +1301,7 @@ func buildOCMCredentialProviders( if err != nil { return nil, nil, err } - providers[credOpt.Tag] = newBalancerProvider(subCredentials, credOpt.BalancerOptions.Strategy, time.Duration(credOpt.BalancerOptions.PollInterval), logger) + providers[credOpt.Tag] = newBalancerProvider(subCredentials, credOpt.BalancerOptions.Strategy, time.Duration(credOpt.BalancerOptions.PollInterval), credOpt.BalancerOptions.RebalanceThreshold, logger) case "fallback": subCredentials, err := resolveCredentialTags(credOpt.FallbackOptions.Credentials, allCredentialMap, credOpt.Tag) if err != nil { @@ -1339,6 +1408,9 @@ func validateOCMOptions(options option.OCMServiceOptions) error { default: return E.New("credential ", cred.Tag, ": unknown balancer strategy: ", cred.BalancerOptions.Strategy) } + if cred.BalancerOptions.RebalanceThreshold < 0 { + return E.New("credential ", cred.Tag, ": rebalance_threshold must not be negative") + } } } diff --git a/service/ocm/service.go b/service/ocm/service.go index 94a98c665..376659c4b 100644 --- a/service/ocm/service.go +++ b/service/ocm/service.go @@ -500,6 +500,7 @@ func (s *Service) ServeHTTP(w http.ResponseWriter, r *http.Request) { } requestContext := selectedCredential.wrapRequestContext(r.Context()) + provider.wrapProviderInterrupt(selectedCredential, requestContext) defer func() { requestContext.cancelRequest() }() @@ -539,6 +540,7 @@ func (s *Service) ServeHTTP(w http.ResponseWriter, r *http.Request) { s.logger.InfoContext(ctx, "retrying with credential ", nextCredential.tagName(), " after 429 from ", selectedCredential.tagName()) requestContext.cancelRequest() requestContext = nextCredential.wrapRequestContext(r.Context()) + provider.wrapProviderInterrupt(nextCredential, requestContext) retryRequest, buildErr := nextCredential.buildProxyRequest(requestContext, r, bodyBytes, s.httpHeaders) if buildErr != nil { s.logger.ErrorContext(ctx, "retry request: ", buildErr)