wangjia/sing-box
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

ccm,ocm: add balancer session rebalancing with per-credential interrupt

Browse Source 
When a sticky session's credential utilization exceeds the least-used
credential by a weight-adjusted threshold, force reassign all sessions
on that credential and cancel in-flight requests scoped to the balancer.

Threshold formula: effective = rebalance_threshold / planWeight, so a
config value of 20 triggers at 2% delta for Max 20x (w=10), 4% for
Max 5x (w=5), and 20% for Pro (w=1).
This commit is contained in:
世界
2026-03-14 17:15:32 +08:00
parent b97b9d9cfd
commit f871113832
8 changed files with 208 additions and 58 deletions
Download Patch File Download Diff File Expand all files Collapse all files
option/ccm.go
+4 -3
View File
@@ -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 {
option/ocm.go
+4 -3
View File
@@ -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 {
service/ccm/credential_external.go
+3 -3
View File
@@ -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,
}
}
service/ccm/credential_state.go
+96 -24
View File
@@ -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")
}
}
}
service/ccm/service.go
+2
View File
@@ -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)
service/ocm/credential_external.go
+3 -3
View File
@@ -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,
}
}
service/ocm/credential_state.go
+94 -22
View File
@@ -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")
}
}
}
service/ocm/service.go
+2
View File
@@ -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)