Propagate request context to upstream requests

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 {

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 {

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)