@@ -76,6 +76,7 @@ def __str__(self) -> str:
76
76
class HttpClient :
77
77
_DEFAULT_MAX_RETRY : int = 5
78
78
_DEFAULT_MAX_TIME : int = 60 * 10
79
+ _ACTIONS_TO_RETRY_ON = {ResponseAction .RETRY , ResponseAction .RATE_LIMITED }
79
80
80
81
def __init__ (
81
82
self ,
@@ -359,6 +360,17 @@ def _get_response_body(self, response: requests.Response) -> Optional[JsonType]:
359
360
except Exception :
360
361
return "The Content of the Response couldn't be decoded."
361
362
363
+ def _evict_key (self , prepared_request : requests .PreparedRequest ) -> None :
364
+ """
365
+ Addresses high memory consumption when enabling concurrency in https://github.com/airbytehq/oncall/issues/6821.
366
+
367
+ The `_request_attempt_count` attribute keeps growing as multiple requests are made using the same `http_client`.
368
+ To mitigate this issue, we evict keys for completed requests once we confirm that no further retries are needed.
369
+ This helps manage memory usage more efficiently while maintaining the necessary logic for retry attempts.
370
+ """
371
+ if prepared_request in self ._request_attempt_count :
372
+ del self ._request_attempt_count [prepared_request ]
373
+
362
374
def _handle_error_resolution (
363
375
self ,
364
376
response : Optional [requests .Response ],
@@ -367,6 +379,9 @@ def _handle_error_resolution(
367
379
error_resolution : ErrorResolution ,
368
380
exit_on_rate_limit : Optional [bool ] = False ,
369
381
) -> None :
382
+ if error_resolution .response_action not in self ._ACTIONS_TO_RETRY_ON :
383
+ self ._evict_key (request )
384
+
370
385
# Emit stream status RUNNING with the reason RATE_LIMITED to log that the rate limit has been reached
371
386
if error_resolution .response_action == ResponseAction .RATE_LIMITED :
372
387
# TODO: Update to handle with message repository when concurrent message repository is ready
0 commit comments