Respect GitHub ratelimit to support more than 8 stacked PRs (#290)

Co-authored-by: Brian Amerige <[email protected]>

Author: ezyang

src/ghstack/github_real.py

@@ -3,12 +3,16 @@
 import json
 import logging
 import re
+import time
 from typing import Any, Dict, Optional, Sequence, Tuple, Union
 
 import requests
 
 import ghstack.github
 
+MAX_RETRIES = 5
+INITIAL_BACKOFF_SECONDS = 60
+
 
 class RealGitHubEndpoint(ghstack.github.GitHubEndpoint):
     """
@@ -155,32 +159,63 @@ def rest(self, method: str, path: str, **kwargs: Any) -> Any:
         }
 
         url = self.rest_endpoint.format(github_url=self.github_url) + "/" + path
-        logging.debug("# {} {}".format(method, url))
-        logging.debug("Request body:\n{}".format(json.dumps(kwargs, indent=1)))
 
-        resp: requests.Response = getattr(requests, method)(
-            url,
-            json=kwargs,
-            headers=headers,
-            proxies=self._proxies(),
-            verify=self.verify,
-            cert=self.cert,
-        )
+        backoff_seconds = INITIAL_BACKOFF_SECONDS
+        for attempt in range(0, MAX_RETRIES):
+            logging.debug("# {} {}".format(method, url))
+            logging.debug("Request body:\n{}".format(json.dumps(kwargs, indent=1)))
 
-        logging.debug("Response status: {}".format(resp.status_code))
+            resp: requests.Response = getattr(requests, method)(
+                url,
+                json=kwargs,
+                headers=headers,
+                proxies=self._proxies(),
+                verify=self.verify,
+                cert=self.cert,
+            )
 
-        try:
-            r = resp.json()
-        except ValueError:
-            logging.debug("Response body:\n{}".format(r.text))
-            raise
-        else:
-            pretty_json = json.dumps(r, indent=1)
-            logging.debug("Response JSON:\n{}".format(pretty_json))
+            logging.debug("Response status: {}".format(resp.status_code))
 
-        if resp.status_code == 404:
-            raise ghstack.github.NotFoundError(
-                """\
+            try:
+                r = resp.json()
+            except ValueError:
+                logging.debug("Response body:\n{}".format(r.text))
+                raise
+            else:
+                pretty_json = json.dumps(r, indent=1)
+                logging.debug("Response JSON:\n{}".format(pretty_json))
+
+            # Per Github rate limiting: https://docs.github.com/en/rest/using-the-rest-api/rate-limits-for-the-rest-api?apiVersion=2022-11-28#exceeding-the-rate-limit
+            if resp.status_code in (403, 429):
+                remaining_count = resp.headers.get("x-ratelimit-remaining")
+                reset_time = resp.headers.get("x-ratelimit-reset")
+
+                if remaining_count == "0" and reset_time:
+                    sleep_time = int(reset_time) - int(time.time())
+                    logging.warning(
+                        f"Rate limit exceeded. Sleeping until reset in {sleep_time} seconds."
+                    )
+                    time.sleep(sleep_time)
+                    continue
+                else:
+                    retry_after_seconds = resp.headers.get("retry-after")
+                    if retry_after_seconds:
+                        sleep_time = int(retry_after_seconds)
+                        logging.warning(
+                            f"Secondary rate limit hit. Sleeping for {sleep_time} seconds."
+                        )
+                    else:
+                        sleep_time = backoff_seconds
+                        logging.warning(
+                            f"Secondary rate limit hit. Sleeping for {sleep_time} seconds (exponential backoff)."
+                        )
+                        backoff_seconds *= 2
+                    time.sleep(sleep_time)
+                    continue
+
+            if resp.status_code == 404:
+                raise ghstack.github.NotFoundError(
+                    """\
 GitHub raised a 404 error on the request for
 {url}.
 Usually, this doesn't actually mean the page doesn't exist; instead, it
@@ -190,13 +225,15 @@ def rest(self, method: str, path: str, **kwargs: Any) -> Any:
 "public_repo" for permissions, and update ~/.ghstackrc with your new
 value.
 """.format(
-                    url=url, github_url=self.github_url
+                        url=url, github_url=self.github_url
+                    )
                 )
-            )
 
-        try:
-            resp.raise_for_status()
-        except requests.HTTPError:
-            raise RuntimeError(pretty_json)
+            try:
+                resp.raise_for_status()
+            except requests.HTTPError:
+                raise RuntimeError(pretty_json)
 
-        return r
+            return r
+
+        raise RuntimeError("Exceeded maximum retries due to GitHub rate limiting")

src/ghstack/submit.py

@@ -468,12 +468,6 @@ def run(self) -> List[DiffMeta]:
             raise RuntimeError(
                 "There appears to be no commits to process, based on the revs you passed me."
             )
-        elif commit_count > 8 and not self.force:
-            raise RuntimeError(
-                "Cowardly refusing to handle a stack with more than eight PRs.  "
-                "You are likely to get rate limited by GitHub if you try to create or "
-                "manipulate this many PRs.  You can bypass this throttle using --force"
-            )
 
         # This is not really accurate if you're doing a fancy pattern;
         # if this is a problem file us a bug.

test/submit/throttle.py.test

@@ -5,10 +5,6 @@ init_test()
 for i in range(10):
     commit(f"A{i}")
 
-assert_expected_raises_inline(
-    RuntimeError,
-    lambda: gh_submit("Initial"),
-    """Cowardly refusing to handle a stack with more than eight PRs.  You are likely to get rate limited by GitHub if you try to create or manipulate this many PRs.  You can bypass this throttle using --force""",
-)
+gh_submit("Initial")
 
 ok()