diff --git a/service/ccm/service.go b/service/ccm/service.go index 2760c348a..58bdd3787 100644 --- a/service/ccm/service.go +++ b/service/ccm/service.go @@ -458,7 +458,7 @@ func (s *Service) ServeHTTP(w http.ResponseWriter, r *http.Request) { return } - requestContext := selectedCredential.wrapRequestContext(r.Context()) + requestContext := selectedCredential.wrapRequestContext(ctx) provider.wrapProviderInterrupt(selectedCredential, requestContext) defer func() { requestContext.cancelRequest() @@ -497,7 +497,7 @@ func (s *Service) ServeHTTP(w http.ResponseWriter, r *http.Request) { response.Body.Close() s.logger.InfoContext(ctx, "retrying with credential ", nextCredential.tagName(), " after 429 from ", selectedCredential.tagName()) requestContext.cancelRequest() - requestContext = nextCredential.wrapRequestContext(r.Context()) + requestContext = nextCredential.wrapRequestContext(ctx) provider.wrapProviderInterrupt(nextCredential, requestContext) retryRequest, buildErr := nextCredential.buildProxyRequest(requestContext, r, bodyBytes, s.httpHeaders) if buildErr != nil { diff --git a/service/ocm/service.go b/service/ocm/service.go index 376659c4b..cd7909dd4 100644 --- a/service/ocm/service.go +++ b/service/ocm/service.go @@ -499,7 +499,7 @@ func (s *Service) ServeHTTP(w http.ResponseWriter, r *http.Request) { s.logger.DebugContext(ctx, logParts...) } - requestContext := selectedCredential.wrapRequestContext(r.Context()) + requestContext := selectedCredential.wrapRequestContext(ctx) provider.wrapProviderInterrupt(selectedCredential, requestContext) defer func() { requestContext.cancelRequest() @@ -539,7 +539,7 @@ func (s *Service) ServeHTTP(w http.ResponseWriter, r *http.Request) { response.Body.Close() s.logger.InfoContext(ctx, "retrying with credential ", nextCredential.tagName(), " after 429 from ", selectedCredential.tagName()) requestContext.cancelRequest() - requestContext = nextCredential.wrapRequestContext(r.Context()) + requestContext = nextCredential.wrapRequestContext(ctx) provider.wrapProviderInterrupt(nextCredential, requestContext) retryRequest, buildErr := nextCredential.buildProxyRequest(requestContext, r, bodyBytes, s.httpHeaders) if buildErr != nil { diff --git a/service/ocm/service_websocket.go b/service/ocm/service_websocket.go index 17178e8c2..fcffaae96 100644 --- a/service/ocm/service_websocket.go +++ b/service/ocm/service_websocket.go @@ -96,12 +96,18 @@ func (s *Service) handleWebSocket( ) { var ( err error + requestContext *credentialRequestContext upstreamConn net.Conn upstreamBufferedReader *bufio.Reader upstreamResponseHeaders http.Header statusCode int statusResponseBody string ) + defer func() { + if requestContext != nil { + requestContext.cancelRequest() + } + }() for { accessToken, accessErr := selectedCredential.getAccessToken() @@ -179,10 +185,15 @@ func (s *Service) handleWebSocket( }, } - upstreamConn, upstreamBufferedReader, _, err = upstreamDialer.Dial(s.ctx, upstreamURL) + requestContext = selectedCredential.wrapRequestContext(ctx) + provider.wrapProviderInterrupt(selectedCredential, requestContext) + upstreamConn, upstreamBufferedReader, _, err = upstreamDialer.Dial(requestContext, upstreamURL) if err == nil { + requestContext.releaseCredentialInterrupt() break } + requestContext.cancelRequest() + requestContext = nil if statusCode == http.StatusTooManyRequests { resetAt := parseOCMRateLimitResetFromHeaders(upstreamResponseHeaders) nextCredential := provider.onRateLimited(sessionID, selectedCredential, resetAt, credentialFilter)