Skip to content

Commit f95652b

Browse files
fix: add request count with eviction for not retriable requests to fix memory problem (#92)
Co-authored-by: Aldo Gonzalez <[email protected]>
1 parent 65ed26f commit f95652b

File tree

4 files changed

+600
-470
lines changed

4 files changed

+600
-470
lines changed

airbyte_cdk/sources/streams/http/http_client.py

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -76,6 +76,7 @@ def __str__(self) -> str:
7676
class HttpClient:
7777
_DEFAULT_MAX_RETRY: int = 5
7878
_DEFAULT_MAX_TIME: int = 60 * 10
79+
_ACTIONS_TO_RETRY_ON = {ResponseAction.RETRY, ResponseAction.RATE_LIMITED}
7980

8081
def __init__(
8182
self,
@@ -359,6 +360,17 @@ def _get_response_body(self, response: requests.Response) -> Optional[JsonType]:
359360
except Exception:
360361
return "The Content of the Response couldn't be decoded."
361362

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+
362374
def _handle_error_resolution(
363375
self,
364376
response: Optional[requests.Response],
@@ -367,6 +379,9 @@ def _handle_error_resolution(
367379
error_resolution: ErrorResolution,
368380
exit_on_rate_limit: Optional[bool] = False,
369381
) -> None:
382+
if error_resolution.response_action not in self._ACTIONS_TO_RETRY_ON:
383+
self._evict_key(request)
384+
370385
# Emit stream status RUNNING with the reason RATE_LIMITED to log that the rate limit has been reached
371386
if error_resolution.response_action == ResponseAction.RATE_LIMITED:
372387
# TODO: Update to handle with message repository when concurrent message repository is ready

0 commit comments

Comments
 (0)