feat: support Retry-After, API-Usage-Limit

insales/connection.py

@@ -2,6 +2,7 @@
 
 import datetime
 import time
+import threading
 import socket
 
 from base64 import b64encode
@@ -18,6 +19,13 @@
     from urllib import urlencode
 
 
+insales_lock = threading.Lock()
+
+
+def throttle_fn(curr: int, limit: int):
+    return 0.1 + (1.0 - (limit - curr)/limit)**16 * 8.5
+
+
 class ApiError(Exception):
     def __init__(self, msg, code=None):
         super(ApiError, self).__init__(msg)
@@ -28,15 +36,65 @@ class Connection(object):
     def __init__(self, account, api_key, password,
                  secure=False,
                  retry_on_503=False, retry_on_socket_error=False,
-                 retry_timeout=1, response_timeout=10):
+                 retry_timeout=1, response_timeout=10,
+                 throttle=False):
         self.account = account
         self.api_key = api_key
         self.password = password
         self.secure = secure
         self.retry_on_503 = retry_on_503
         self.retry_on_socket_error = retry_on_socket_error
-        self.retry_timeout = retry_timeout
+        self.retry_timeout = datetime.timedelta(seconds=retry_timeout)
         self.response_timeout = response_timeout
+        self.last_req_time = datetime.datetime.now()
+        self.retry_after = datetime.datetime.now()
+        self.max_wait_time = datetime.timedelta(seconds=60)
+        self.throttle = throttle != False
+        self.throttle_fn = throttle if callable(throttle) else throttle_fn
+
+    def get_retry_after(self):
+        return self.retry_after
+
+    def _set_retry_after(self, delta):
+        with insales_lock:
+            self.retry_after = self.last_req_time + min(
+                delta, self.max_wait_time
+            )
+
+    def _increase_retry_after(self, delta):
+        with insales_lock:
+            self.retry_after = max(self.last_req_time, self.retry_after) + min(
+                delta, self.max_wait_time
+            )
+
+    def _apply_retry_timeout(self):
+        self._set_retry_after(self.retry_timeout)
+
+    def _handle_retry_after_header(self, header):
+        try:
+            self._set_retry_after(datetime.timedelta(seconds=int(header)))
+        except ValueError:
+            self._apply_retry_timeout()
+
+    def _apply_usage_limit(self, header):
+        if header is None:
+            return
+
+        try:
+            curr_str, limit_str = header.split("/")
+            delta = datetime.timedelta(seconds=self.throttle_fn(
+                int(curr_str),
+                int(limit_str),
+            ))
+            self._increase_retry_after(delta)
+        except ValueError:
+            pass
+
+    def _wait_until_retry_after(self):
+        delta = self.retry_after - datetime.datetime.now()
+        if delta.total_seconds() > 0:
+            time.sleep(delta.total_seconds())
+
 
     def request(self, method, endpoint, qargs={}, data=None):
         path = self.format_path(endpoint, qargs)
@@ -48,24 +106,34 @@ def request(self, method, endpoint, qargs={}, data=None):
 
         done = False
         while not done:
+            self._wait_until_retry_after()
             try:
                 host = '%s.myinsales.ru' % self.account
                 if self.secure:
                     conn = HTTPSConnection(host, timeout=self.response_timeout)
                 else:
                     conn = HTTPConnection(host, timeout=self.response_timeout)
+                with insales_lock:
+                    self.last_req_time = datetime.datetime.now()
                 conn.request(method, path, headers=headers, body=data)
                 resp = conn.getresponse()
                 body = resp.read()
             except (socket.gaierror, socket.timeout, HTTPException):
                 if self.retry_on_socket_error:
-                    time.sleep(self.retry_timeout)
+                    self._apply_retry_timeout()
                     continue
                 else:
                     raise
 
+            if self.throttle:
+                self._apply_usage_limit(resp.getheader('API-Usage-Limit'))
+
             if resp.status == 503 and self.retry_on_503:
-                time.sleep(self.retry_timeout)
+                retry_after_header = resp.getheader('Retry-After')
+                if retry_after_header:
+                    self._handle_retry_after_header(retry_after_header)
+                else:
+                    self._apply_retry_timeout()
             else:
                 done = True