Add configurable backoff

src/ghstack/cli.py

@@ -61,6 +61,8 @@ def cli_context(
             oauth_token=config.github_oauth,
             proxy=config.proxy,
             github_url=config.github_url,
+            max_retries=config.max_retries,
+            initial_backoff_seconds=config.initial_backoff_seconds,
         )
         yield shell, config, github
 

src/ghstack/config.py

@@ -111,6 +111,11 @@ def get_gh_cli_credentials(
         ("reviewer", Optional[str]),
         # Default labels to add to new pull requests (comma-separated labels)
         ("label", Optional[str]),
+        # Retry/backoff tuning
+        ("max_retries", int),
+        # The initial backoff time to use, in seconds.  We will double this
+        # time for each retry.
+        ("initial_backoff_seconds", int),
         # Command to generate a per-PR update description from diff contents
         ("automsg", Optional[str]),
     ],
@@ -334,6 +339,16 @@ def read_config(
     else:
         automsg = None
 
+    if config.has_option("ghstack", "max_retries"):
+        max_retries = config.getint("ghstack", "max_retries")
+    else:
+        max_retries = 5
+
+    if config.has_option("ghstack", "initial_backoff_seconds"):
+        initial_backoff_seconds = config.getint("ghstack", "initial_backoff_seconds")
+    else:
+        initial_backoff_seconds = 60
+
     if write_back:
         with open(config_path, "w") as f:
             config.write(f)
@@ -351,6 +366,8 @@ def read_config(
         remote_name=remote_name,
         reviewer=reviewer,
         label=label,
+        max_retries=max_retries,
+        initial_backoff_seconds=initial_backoff_seconds,
         automsg=automsg,
     )
     logging.debug(f"conf = {conf}")

src/ghstack/github_real.py

@@ -13,9 +13,6 @@
 
 import ghstack.github
 
-MAX_RETRIES = 5
-INITIAL_BACKOFF_SECONDS = 60
-
 
 class RealGitHubEndpoint(ghstack.github.GitHubEndpoint):
     """
@@ -61,19 +58,26 @@ def rest_endpoint(self) -> str:
     cert: Optional[Union[str, Tuple[str, str]]]
     _session: Optional[aiohttp.ClientSession]
 
+    max_retries: int
+    initial_backoff_seconds: int
+
     def __init__(
         self,
         oauth_token: Optional[str],
         github_url: str,
         proxy: Optional[str] = None,
         verify: Optional[Union[str, bool]] = None,
         cert: Optional[Union[str, Tuple[str, str]]] = None,
+        max_retries: int = 5,
+        initial_backoff_seconds: int = 60,
     ):
         self.oauth_token = oauth_token
         self.proxy = proxy
         self.github_url = github_url
         self.verify = verify
         self.cert = cert
+        self.max_retries = max_retries
+        self.initial_backoff_seconds = initial_backoff_seconds
         self._rest_request_ids = itertools.count(1)
         self._session = None
 
@@ -212,11 +216,11 @@ async def arest(self, method: str, path: str, **kwargs: Any) -> Any:
         if aiohttp_ssl is not None:
             request_kwargs["ssl"] = aiohttp_ssl
 
-        backoff_seconds = INITIAL_BACKOFF_SECONDS
+        backoff_seconds = self.initial_backoff_seconds
         request_id = next(self._rest_request_ids)
         log_prefix = f"rest[{request_id}]"
         session = self._get_session()
-        for attempt in range(0, MAX_RETRIES):
+        for attempt in range(0, self.max_retries):
             logging.debug("# %s %s %s", log_prefix, method, url)
             logging.debug(
                 "%s request body:\n%s", log_prefix, json.dumps(kwargs, indent=1)
@@ -239,32 +243,32 @@ async def arest(self, method: str, path: str, **kwargs: Any) -> Any:
                 if resp.status in (403, 429):
                     remaining_count = resp.headers.get("x-ratelimit-remaining")
                     reset_time = resp.headers.get("x-ratelimit-reset")
+                    more_attempts = attempt < self.max_retries - 1
 
                     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."
-                        )
-                        await asyncio.sleep(sleep_time)
-                        continue
+                        sleep_time = max(0, int(reset_time) - int(time.time()))
+                        if more_attempts:
+                            logging.warning(
+                                f"Rate limit exceeded. Sleeping until reset in {sleep_time} seconds."
+                            )
+                            await asyncio.sleep(sleep_time)
+                            continue
                     # GitHub doesn't document the content of these messages, but this
                     # seems to be an accurate way to find secondary rate limits.  Any
                     # other reason for 403 or 429 will fall through to the error below.
                     elif "rate limit" in resp_text.lower():
                         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
+                            backoff_seconds *= 2
+                        if more_attempts:
                             logging.warning(
-                                f"Secondary rate limit hit. Sleeping for {sleep_time} seconds (exponential backoff)."
+                                f"Secondary rate limit hit. Sleeping for {sleep_time} seconds."
                             )
-                            backoff_seconds *= 2
-                        await asyncio.sleep(sleep_time)
-                        continue
+                            await asyncio.sleep(max(0, sleep_time))
+                            continue
 
                 if resp.status == 404:
                     raise ghstack.github.NotFoundError(