Skip to content

Commit fea1143

Browse files
author
Janne Rönkkö
committed
feat: Add throttling API requests
1 parent b858088 commit fea1143

2 files changed

Lines changed: 87 additions & 11 deletions

File tree

‎gitlab/__init__.py‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -82,6 +82,35 @@ def _calculate_wait_time(self, response, retries):
8282
return 2 ** retries * 0.1
8383

8484

85+
class RequestThrottler(object):
86+
def __init__(self, max_requests_per_period=60, period=600):
87+
self.max_requests_per_period = max_requests_per_period
88+
self.period = period
89+
90+
self.calls_within_period = []
91+
92+
def __call__(self):
93+
current_time = time.monotonic()
94+
95+
cutoff_time = current_time - self.period
96+
self._remove_calls_before(cutoff_time)
97+
98+
if len(self.calls_within_period) >= self.max_requests_per_period:
99+
new_cutoff_time = self.calls_within_period.pop(0) + 1
100+
self._remove_calls_before(current_time - self.period)
101+
102+
self._wait(new_cutoff_time - cutoff_time)
103+
104+
self.calls_within_period.append(time.monotonic())
105+
106+
def _remove_calls_before(self, cutoff_time):
107+
while self.calls_within_period and self.calls_within_period[0] < cutoff_time:
108+
self.calls_within_period.pop(0)
109+
110+
def _wait(self, wait_time):
111+
time.sleep(wait_time)
112+
113+
85114
class Gitlab(object):
86115
"""Represents a GitLab server connection.
87116
@@ -117,6 +146,7 @@ def __init__(
117146
session=None,
118147
per_page=None,
119148
get_wait_time=None,
149+
throttle_requests=None,
120150
):
121151

122152
self._api_version = str(api_version)
@@ -144,6 +174,7 @@ def __init__(
144174
self.per_page = per_page
145175

146176
self._get_wait_time = get_wait_time or DefaultWaitTimeStrategy()
177+
self._throttle_requests = throttle_requests or RequestThrottler()
147178

148179
objects = importlib.import_module("gitlab.v%s.objects" % self._api_version)
149180
self._objects = objects
@@ -557,6 +588,7 @@ def http_request(
557588
cur_retries = 0
558589

559590
while True:
591+
self._throttle_requests()
560592
result = self.session.send(prepped, timeout=timeout, **settings)
561593

562594
self._check_redirects(result)

‎gitlab/tests/test_gitlab.py‎

Lines changed: 55 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -794,11 +794,11 @@ class MyGitlab(gitlab.Gitlab):
794794

795795
class TestRetryWaitTime(unittest.TestCase):
796796
def setUp(self):
797-
self.session_mock = unittest.mock.Mock(name='Session mock')
797+
self.session_mock = unittest.mock.Mock(name="Session mock")
798798
self.session_mock.prepare_request.return_value.url = "http://localhost"
799799
self.session_mock.merge_environment_settings.return_value = {}
800800

801-
@unittest.mock.patch('gitlab.time.sleep', name='sleep mock')
801+
@unittest.mock.patch("gitlab.time.sleep", name="sleep mock")
802802
def test_default_retry_wait_time(self, sleep_mock):
803803
self.gl = Gitlab(
804804
"http://localhost",
@@ -831,16 +831,14 @@ def test_default_retry_wait_time(self, sleep_mock):
831831
)
832832

833833
self.assertAlmostEqual(
834-
sleep_mock.call_args_list[2][0][0],
835-
2 ** 2 * 0.1,
834+
sleep_mock.call_args_list[2][0][0], 2 ** 2 * 0.1,
836835
)
837836

838837
self.assertAlmostEqual(
839-
sleep_mock.call_args_list[3][0][0],
840-
2 ** 3 * 0.1,
838+
sleep_mock.call_args_list[3][0][0], 2 ** 3 * 0.1,
841839
)
842840

843-
@unittest.mock.patch('gitlab.time.sleep', name='sleep mock')
841+
@unittest.mock.patch("gitlab.time.sleep", name="sleep mock")
844842
def test_custom_retry_wait_time(self, sleep_mock):
845843
self.gl = Gitlab(
846844
"http://localhost",
@@ -862,9 +860,55 @@ def test_custom_retry_wait_time(self, sleep_mock):
862860
self.assertEqual(http_r.status_code, 200)
863861

864862
self.assertEqual(
865-
[
866-
unittest.mock.call(100),
867-
unittest.mock.call(200),
868-
],
863+
[unittest.mock.call(100), unittest.mock.call(200),],
869864
sleep_mock.call_args_list,
870865
)
866+
867+
868+
class TestRequestThrottler(unittest.TestCase):
869+
def setUp(self):
870+
self.throttler = gitlab.RequestThrottler(5, 10)
871+
872+
monotonic_patcher = unittest.mock.patch(
873+
"gitlab.time.monotonic", name="monotonic mock", return_value=0,
874+
)
875+
self.monotonic_mock = monotonic_patcher.start()
876+
self.addCleanup(monotonic_patcher.stop)
877+
878+
sleep_patcher = unittest.mock.patch(
879+
"gitlab.time.sleep", name="sleep mock", side_effect=self._sleep,
880+
)
881+
self.sleep_mock = sleep_patcher.start()
882+
self.addCleanup(sleep_patcher.stop)
883+
884+
def _sleep(self, sleep_time):
885+
self.monotonic_mock.return_value += sleep_time
886+
887+
def test_throttling(self):
888+
for _ in range(0, 5):
889+
self.monotonic_mock.return_value += 1
890+
self.throttler()
891+
892+
self.assertFalse(self.sleep_mock.called)
893+
self.assertEqual([1, 2, 3, 4, 5], self.throttler.calls_within_period)
894+
895+
self.monotonic_mock.return_value += 1
896+
self.assertEqual(self.monotonic_mock.return_value, 6)
897+
self.throttler()
898+
899+
self.sleep_mock.assert_called_once_with(6)
900+
self.assertEqual(self.monotonic_mock.return_value, 12)
901+
902+
self.assertEqual([2, 3, 4, 5, 12], self.throttler.calls_within_period)
903+
904+
self.sleep_mock.reset_mock()
905+
906+
self.monotonic_mock.return_value += 1
907+
self.throttler()
908+
self.assertFalse(self.sleep_mock.called)
909+
self.assertEqual([3, 4, 5, 12, 13], self.throttler.calls_within_period)
910+
911+
self.monotonic_mock.return_value = 100
912+
self.throttler()
913+
self.assertFalse(self.sleep_mock.called)
914+
self.assertEqual([100], self.throttler.calls_within_period)

0 commit comments

Comments
 (0)