diff --git a/docs/api-usage.rst b/docs/api-usage.rst index 19b959317..bc45aa403 100644 --- a/docs/api-usage.rst +++ b/docs/api-usage.rst @@ -345,6 +345,25 @@ throttled, you can set this parameter to -1. This parameter is ignored if gl = gitlab.gitlab(url, token, api_version=4) gl.projects.list(all=True, max_retries=12) +You can provide custom retry time strategy by providing ``get_wait_time`` argument: + +.. code-block:: python + + import gitlab + + def get_custom_wait_time(response, retries): + retry_after = response.headers("Retry-After") + if retry_after: + return int(retry_after) + + return 2 ** cur_retries * 0.1 + + gl = gitlab.gitlab(url, token, api_version=4, get_wait_time=get_custom_wait_time) + gl.projects.list(all=True, max_retries=12) + +Note that the above ``get_custom_wait_time`` is identical to the default behaviour. + + .. warning:: You will get an Exception, if you then go over the rate limit of your GitLab instance. diff --git a/gitlab/__init__.py b/gitlab/__init__.py index 252074bfd..9780b3cf0 100644 --- a/gitlab/__init__.py +++ b/gitlab/__init__.py @@ -53,6 +53,64 @@ def _sanitize(value): return value +class DefaultWaitTimeStrategy(object): + def __call__(self, response, retries): + """Return wait time before next retry + + Args: + response (requests.Response): HTTP response from previous attempt + retries (int): Number of retries done + + Returns: + Number of seconds to wait before next retry + """ + + wait_time = self._get_from_response(response) + if wait_time is None: + wait_time = self._calculate_wait_time(response, retries) + + return wait_time + + def _get_from_response(self, response): + value = response.headers.get("Retry-After") + if value is None: + return None + + return int(value) + + def _calculate_wait_time(self, response, retries): + return 2 ** retries * 0.1 + + +class RequestThrottler(object): + def __init__(self, max_requests_per_period=600, period=60): + self.max_requests_per_period = max_requests_per_period + self.period = period + + self.calls_within_period = [] + + def __call__(self): + current_time = time.monotonic() + + cutoff_time = current_time - self.period + self._remove_calls_before(cutoff_time) + + if len(self.calls_within_period) >= self.max_requests_per_period: + new_cutoff_time = self.calls_within_period.pop(0) + 1 + self._remove_calls_before(current_time - self.period) + + self._wait(new_cutoff_time - cutoff_time) + + self.calls_within_period.append(time.monotonic()) + + def _remove_calls_before(self, cutoff_time): + while self.calls_within_period and self.calls_within_period[0] < cutoff_time: + self.calls_within_period.pop(0) + + def _wait(self, wait_time): + time.sleep(wait_time) + + class Gitlab(object): """Represents a GitLab server connection. @@ -70,6 +128,8 @@ class Gitlab(object): http_username (str): Username for HTTP authentication http_password (str): Password for HTTP authentication api_version (str): Gitlab API version to use (support for 4 only) + get_wait_time (callable): Callable returning number of seconds to wait + until next retry """ def __init__( @@ -85,6 +145,8 @@ def __init__( api_version="4", session=None, per_page=None, + get_wait_time=None, + throttle_requests=None, ): self._api_version = str(api_version) @@ -111,6 +173,9 @@ def __init__( self.per_page = per_page + self._get_wait_time = get_wait_time or DefaultWaitTimeStrategy() + self._throttle_requests = throttle_requests or RequestThrottler() + objects = importlib.import_module("gitlab.v%s.objects" % self._api_version) self._objects = objects @@ -523,6 +588,7 @@ def http_request( cur_retries = 0 while True: + self._throttle_requests() result = self.session.send(prepped, timeout=timeout, **settings) self._check_redirects(result) @@ -532,9 +598,7 @@ def http_request( if 429 == result.status_code and obey_rate_limit: if max_retries == -1 or cur_retries < max_retries: - wait_time = 2 ** cur_retries * 0.1 - if "Retry-After" in result.headers: - wait_time = int(result.headers["Retry-After"]) + wait_time = self._get_wait_time(result, cur_retries) cur_retries += 1 time.sleep(wait_time) continue diff --git a/gitlab/tests/test_gitlab.py b/gitlab/tests/test_gitlab.py index f9d4cc82e..a05a27e9d 100644 --- a/gitlab/tests/test_gitlab.py +++ b/gitlab/tests/test_gitlab.py @@ -21,6 +21,7 @@ import tempfile import json import unittest +import unittest.mock from httmock import HTTMock # noqa from httmock import response # noqa @@ -789,3 +790,125 @@ class MyGitlab(gitlab.Gitlab): gl = MyGitlab.from_config("one", [config_path]) self.assertIsInstance(gl, MyGitlab) os.unlink(config_path) + + +class TestRetryWaitTime(unittest.TestCase): + def setUp(self): + self.session_mock = unittest.mock.Mock(name="Session mock") + self.session_mock.prepare_request.return_value.url = "http://localhost" + self.session_mock.merge_environment_settings.return_value = {} + + @unittest.mock.patch("gitlab.time.sleep", name="sleep mock") + def test_default_retry_wait_time(self, sleep_mock): + self.gl = Gitlab( + "http://localhost", + private_token="private_token", + ssl_verify=True, + api_version=4, + session=self.session_mock, + ) + + self.session_mock.send.side_effect = [ + response(429, headers={"Retry-After": "60"}), + response(429, headers={"Retry-After": "180"}), + response(429), + response(429), + response(200), + ] + + http_r = self.gl.http_request("get", "/projects", max_retries=4) + + self.assertEqual(http_r.status_code, 200) + + self.assertEqual( + [ + unittest.mock.call(60), + unittest.mock.call(180), + unittest.mock.call(unittest.mock.ANY), + unittest.mock.call(unittest.mock.ANY), + ], + sleep_mock.call_args_list, + ) + + self.assertAlmostEqual( + sleep_mock.call_args_list[2][0][0], 2 ** 2 * 0.1, + ) + + self.assertAlmostEqual( + sleep_mock.call_args_list[3][0][0], 2 ** 3 * 0.1, + ) + + @unittest.mock.patch("gitlab.time.sleep", name="sleep mock") + def test_custom_retry_wait_time(self, sleep_mock): + self.gl = Gitlab( + "http://localhost", + private_token="private_token", + ssl_verify=True, + api_version=4, + session=self.session_mock, + get_wait_time=unittest.mock.Mock(side_effect=[100, 200]), + ) + + self.session_mock.send.side_effect = [ + response(429, headers={"Retry-After": "60"}), + response(429), + response(200), + ] + + http_r = self.gl.http_request("get", "/projects", max_retries=2) + + self.assertEqual(http_r.status_code, 200) + + self.assertEqual( + [unittest.mock.call(100), unittest.mock.call(200),], + sleep_mock.call_args_list, + ) + + +class TestRequestThrottler(unittest.TestCase): + def setUp(self): + self.throttler = gitlab.RequestThrottler(5, 10) + + monotonic_patcher = unittest.mock.patch( + "gitlab.time.monotonic", name="monotonic mock", return_value=0, + ) + self.monotonic_mock = monotonic_patcher.start() + self.addCleanup(monotonic_patcher.stop) + + sleep_patcher = unittest.mock.patch( + "gitlab.time.sleep", name="sleep mock", side_effect=self._sleep, + ) + self.sleep_mock = sleep_patcher.start() + self.addCleanup(sleep_patcher.stop) + + def _sleep(self, sleep_time): + self.monotonic_mock.return_value += sleep_time + + def test_throttling(self): + for _ in range(0, 5): + self.monotonic_mock.return_value += 1 + self.throttler() + + self.assertFalse(self.sleep_mock.called) + self.assertEqual([1, 2, 3, 4, 5], self.throttler.calls_within_period) + + self.monotonic_mock.return_value += 1 + self.assertEqual(self.monotonic_mock.return_value, 6) + self.throttler() + + self.sleep_mock.assert_called_once_with(6) + self.assertEqual(self.monotonic_mock.return_value, 12) + + self.assertEqual([2, 3, 4, 5, 12], self.throttler.calls_within_period) + + self.sleep_mock.reset_mock() + + self.monotonic_mock.return_value += 1 + self.throttler() + self.assertFalse(self.sleep_mock.called) + self.assertEqual([3, 4, 5, 12, 13], self.throttler.calls_within_period) + + self.monotonic_mock.return_value = 100 + self.throttler() + self.assertFalse(self.sleep_mock.called) + self.assertEqual([100], self.throttler.calls_within_period)