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

fix(ccm,ocm): remove external context from pollUsage/pollIfStale

Browse Source 
pollUsage(ctx) accepted caller context, and service_status.go passed
r.Context() which gets canceled on client disconnect or service shutdown.
This caused incrementPollFailures → interruptConnections on transient
cancellations. Each implementation now uses its own persistent context:
defaultCredential uses serviceContext, externalCredential uses
getReverseContext().
This commit is contained in:
世界
2026-03-18 00:54:01 +08:00
parent b119d08764
commit 3bcfdd5455
12 changed files with 35 additions and 28 deletions
Download Patch File Download Diff File Expand all files Collapse all files
service/ccm/credential.go
+1 -1
View File
@@ -120,7 +120,7 @@ type Credential interface {
setStatusSubscriber(*observable.Subscriber[struct{}])
start() error
pollUsage(ctx context.Context)
pollUsage()
lastUpdatedTime() time.Time
pollBackoff(base time.Duration) time.Duration
usageTrackerOrNil() *AggregatedUsage
service/ccm/credential_default.go
+6 -3
View File
@@ -139,6 +139,7 @@ func (c *defaultCredential) start() error {
c.logger.Warn("load usage statistics for ", c.tag, ": ", err)
}
}
go c.pollUsage()
return nil
}
@@ -516,7 +517,7 @@ func (c *defaultCredential) earliestReset() time.Time {
return earliest
}
func (c *defaultCredential) pollUsage(ctx context.Context) {
func (c *defaultCredential) pollUsage() {
if !c.pollAccess.TryLock() {
return
}
@@ -537,6 +538,7 @@ func (c *defaultCredential) pollUsage(ctx context.Context) {
return
}
ctx := c.serviceContext
httpClient := &http.Client{
Transport: c.forwardHTTPClient.Transport,
Timeout: 5 * time.Second,
@@ -633,11 +635,12 @@ func (c *defaultCredential) pollUsage(ctx context.Context) {
c.emitStatusUpdate()
if needsProfileFetch {
c.fetchProfile(ctx, httpClient, accessToken)
c.fetchProfile(httpClient, accessToken)
}
}
func (c *defaultCredential) fetchProfile(ctx context.Context, httpClient *http.Client, accessToken string) {
func (c *defaultCredential) fetchProfile(httpClient *http.Client, accessToken string) {
ctx := c.serviceContext
response, err := doHTTPWithRetry(ctx, httpClient, func() (*http.Request, error) {
request, err := http.NewRequestWithContext(ctx, http.MethodGet, claudeAPIBaseURL+"/api/oauth/profile", nil)
if err != nil {
service/ccm/credential_external.go
+3 -2
View File
@@ -574,13 +574,14 @@ func (c *externalCredential) doPollUsageRequest(ctx context.Context) (*http.Resp
return nil, E.New("no transport available")
}
func (c *externalCredential) pollUsage(ctx context.Context) {
func (c *externalCredential) pollUsage() {
if !c.pollAccess.TryLock() {
return
}
defer c.pollAccess.Unlock()
defer c.markUsagePollAttempted()
ctx := c.getReverseContext()
response, err := c.doPollUsageRequest(ctx)
if err != nil {
c.logger.Debug("poll usage for ", c.tag, ": ", err)
@@ -919,7 +920,7 @@ func (c *externalCredential) setReverseSession(session *yamux.Session) bool {
go c.statusStreamLoop()
}
if triggerUsageRefresh {
go c.pollUsage(c.getReverseContext())
go c.pollUsage()
}
if emitStatus {
c.emitStatusUpdate()
service/ccm/credential_provider.go
+5 -5
View File
@@ -16,7 +16,7 @@ type credentialProvider interface {
selectCredential(sessionID string, selection credentialSelection) (Credential, bool, error)
onRateLimited(sessionID string, credential Credential, resetAt time.Time, selection credentialSelection) Credential
linkProviderInterrupt(credential Credential, selection credentialSelection, onInterrupt func()) func() bool
pollIfStale(ctx context.Context)
pollIfStale()
allCredentials() []Credential
close()
}
@@ -58,7 +58,7 @@ func (p *singleCredentialProvider) onRateLimited(_ string, credential Credential
return nil
}
func (p *singleCredentialProvider) pollIfStale(ctx context.Context) {
func (p *singleCredentialProvider) pollIfStale() {
now := time.Now()
p.sessionAccess.Lock()
for id, createdAt := range p.sessions {
@@ -69,7 +69,7 @@ func (p *singleCredentialProvider) pollIfStale(ctx context.Context) {
p.sessionAccess.Unlock()
if time.Since(p.credential.lastUpdatedTime()) > p.credential.pollBackoff(defaultPollInterval) {
p.credential.pollUsage(ctx)
p.credential.pollUsage()
}
}
@@ -357,7 +357,7 @@ func (p *balancerProvider) pickRandom(filter func(Credential) bool) Credential {
return usable[rand.IntN(len(usable))]
}
func (p *balancerProvider) pollIfStale(ctx context.Context) {
func (p *balancerProvider) pollIfStale() {
now := time.Now()
p.sessionAccess.Lock()
for id, entry := range p.sessions {
@@ -377,7 +377,7 @@ func (p *balancerProvider) pollIfStale(ctx context.Context) {
for _, credential := range p.credentials {
if time.Since(credential.lastUpdatedTime()) > credential.pollBackoff(p.pollInterval) {
credential.pollUsage(ctx)
credential.pollUsage()
}
}
}
service/ccm/service_handler.go
+2 -2
View File
@@ -182,7 +182,7 @@ func (s *Service) ServeHTTP(w http.ResponseWriter, r *http.Request) {
return
}
provider.pollIfStale(s.ctx)
provider.pollIfStale()
anthropicBetaHeader := r.Header.Get("anthropic-beta")
if isFastModeRequest(anthropicBetaHeader) {
@@ -305,7 +305,7 @@ func (s *Service) ServeHTTP(w http.ResponseWriter, r *http.Request) {
if response.StatusCode != http.StatusOK && response.StatusCode != http.StatusTooManyRequests {
body, _ := io.ReadAll(response.Body)
s.logger.ErrorContext(ctx, "upstream error from ", selectedCredential.tagName(), ": status ", response.StatusCode, " ", string(body))
go selectedCredential.pollUsage(s.ctx)
go selectedCredential.pollUsage()
writeJSONError(w, r, http.StatusInternalServerError, "api_error",
"proxy request (status "+strconv.Itoa(response.StatusCode)+"): "+string(body))
return
service/ccm/service_status.go
+2 -2
View File
@@ -103,7 +103,7 @@ func (s *Service) handleStatusEndpoint(w http.ResponseWriter, r *http.Request) {
return
}
provider.pollIfStale(r.Context())
provider.pollIfStale()
status := s.computeAggregatedUtilization(provider, userConfig)
w.Header().Set("Content-Type", "application/json")
@@ -125,7 +125,7 @@ func (s *Service) handleStatusStream(w http.ResponseWriter, r *http.Request, pro
}
defer s.statusObserver.UnSubscribe(subscription)
provider.pollIfStale(r.Context())
provider.pollIfStale()
w.Header().Set("Content-Type", "application/json")
w.WriteHeader(http.StatusOK)
service/ocm/credential.go
+1 -1
View File
@@ -122,7 +122,7 @@ type Credential interface {
setOnBecameUnusable(fn func())
setStatusSubscriber(*observable.Subscriber[struct{}])
start() error
pollUsage(ctx context.Context)
pollUsage()
lastUpdatedTime() time.Time
pollBackoff(base time.Duration) time.Duration
usageTrackerOrNil() *AggregatedUsage
service/ocm/credential_default.go
+3 -1
View File
@@ -143,6 +143,7 @@ func (c *defaultCredential) start() error {
c.logger.Warn("load usage statistics for ", c.tag, ": ", err)
}
}
go c.pollUsage()
return nil
}
@@ -597,7 +598,7 @@ func (c *defaultCredential) ocmGetBaseURL() string {
return c.getBaseURL()
}
func (c *defaultCredential) pollUsage(ctx context.Context) {
func (c *defaultCredential) pollUsage() {
if !c.pollAccess.TryLock() {
return
}
@@ -621,6 +622,7 @@ func (c *defaultCredential) pollUsage(ctx context.Context) {
return
}
ctx := c.serviceContext
usageURL := strings.TrimSuffix(chatGPTBackendURL, "/codex") + "/wham/usage"
accountID := c.getAccountID()
service/ocm/credential_external.go
+3 -2
View File
@@ -615,13 +615,14 @@ func (c *externalCredential) doPollUsageRequest(ctx context.Context) (*http.Resp
return nil, E.New("no transport available")
}
func (c *externalCredential) pollUsage(ctx context.Context) {
func (c *externalCredential) pollUsage() {
if !c.pollAccess.TryLock() {
return
}
defer c.pollAccess.Unlock()
defer c.markUsagePollAttempted()
ctx := c.getReverseContext()
response, err := c.doPollUsageRequest(ctx)
if err != nil {
c.logger.Debug("poll usage for ", c.tag, ": ", err)
@@ -988,7 +989,7 @@ func (c *externalCredential) setReverseSession(session *yamux.Session) bool {
go c.statusStreamLoop()
}
if triggerUsageRefresh {
go c.pollUsage(c.getReverseContext())
go c.pollUsage()
}
if emitStatus {
c.emitStatusUpdate()
service/ocm/credential_provider.go
+5 -5
View File
@@ -16,7 +16,7 @@ type credentialProvider interface {
selectCredential(sessionID string, selection credentialSelection) (Credential, bool, error)
onRateLimited(sessionID string, credential Credential, resetAt time.Time, selection credentialSelection) Credential
linkProviderInterrupt(credential Credential, selection credentialSelection, onInterrupt func()) func() bool
pollIfStale(ctx context.Context)
pollIfStale()
allCredentials() []Credential
close()
}
@@ -58,7 +58,7 @@ func (p *singleCredentialProvider) onRateLimited(_ string, credential Credential
return nil
}
func (p *singleCredentialProvider) pollIfStale(ctx context.Context) {
func (p *singleCredentialProvider) pollIfStale() {
now := time.Now()
p.sessionAccess.Lock()
for id, createdAt := range p.sessions {
@@ -69,7 +69,7 @@ func (p *singleCredentialProvider) pollIfStale(ctx context.Context) {
p.sessionAccess.Unlock()
if time.Since(p.credential.lastUpdatedTime()) > p.credential.pollBackoff(defaultPollInterval) {
p.credential.pollUsage(ctx)
p.credential.pollUsage()
}
}
@@ -384,7 +384,7 @@ func (p *balancerProvider) pickRandom(filter func(Credential) bool) Credential {
return usable[rand.IntN(len(usable))]
}
func (p *balancerProvider) pollIfStale(ctx context.Context) {
func (p *balancerProvider) pollIfStale() {
now := time.Now()
p.sessionAccess.Lock()
for id, entry := range p.sessions {
@@ -404,7 +404,7 @@ func (p *balancerProvider) pollIfStale(ctx context.Context) {
for _, credential := range p.credentials {
if time.Since(credential.lastUpdatedTime()) > credential.pollBackoff(p.pollInterval) {
credential.pollUsage(ctx)
credential.pollUsage()
}
}
}
service/ocm/service_handler.go
+2 -2
View File
@@ -131,7 +131,7 @@ func (s *Service) ServeHTTP(w http.ResponseWriter, r *http.Request) {
return
}
provider.pollIfStale(s.ctx)
provider.pollIfStale()
selection := credentialSelectionForUser(userConfig)
@@ -285,7 +285,7 @@ func (s *Service) ServeHTTP(w http.ResponseWriter, r *http.Request) {
if response.StatusCode != http.StatusOK && response.StatusCode != http.StatusTooManyRequests {
body, _ := io.ReadAll(response.Body)
s.logger.ErrorContext(ctx, "upstream error from ", selectedCredential.tagName(), ": status ", response.StatusCode, " ", string(body))
go selectedCredential.pollUsage(s.ctx)
go selectedCredential.pollUsage()
writeJSONError(w, r, http.StatusInternalServerError, "api_error",
"proxy request (status "+strconv.Itoa(response.StatusCode)+"): "+string(body))
return
service/ocm/service_status.go
+2 -2
View File
@@ -103,7 +103,7 @@ func (s *Service) handleStatusEndpoint(w http.ResponseWriter, r *http.Request) {
return
}
provider.pollIfStale(r.Context())
provider.pollIfStale()
status := s.computeAggregatedUtilization(provider, userConfig)
w.Header().Set("Content-Type", "application/json")
@@ -125,7 +125,7 @@ func (s *Service) handleStatusStream(w http.ResponseWriter, r *http.Request, pro
}
defer s.statusObserver.UnSubscribe(subscription)
provider.pollIfStale(r.Context())
provider.pollIfStale()
w.Header().Set("Content-Type", "application/json")
w.WriteHeader(http.StatusOK)