Makes gen command retry with exponential backoff (kyuridenamida#174)

yuuki3655 · kyuridenamida · commit d5fffaa03091 · 2019-12-30T20:38:56.000+09:00
Previously, gen command retries with fixed delay of 1.5 seconds indefinitely. It could hurt the service especially when they are experiencing overloaded traffic. This patch changes the retry strategy with exponential backoff starting from 1.5 seconds to 60 seconds at maximum. The gen commands aborts by EnvironmentInitializationError exception after 10 attempts.
diff --git a/atcodertools/tools/envgen.py b/atcodertools/tools/envgen.py
@@ -5,7 +5,7 @@
 import sys
 import traceback
 from multiprocessing import Pool, cpu_count
-from time import sleep
+import time
 from typing import Tuple
 
 from colorama import Fore
@@ -34,6 +34,10 @@ class BannedFileDetectedError(Exception):
     pass
 
 
+class EnvironmentInitializationError(Exception):
+    pass
+
+
 def output_splitter():
     # for readability
     print("=================================================", file=sys.stderr)
@@ -165,16 +169,23 @@ def func(argv: Tuple[AtCoderClient, Problem, Config]):
 
 def prepare_contest(atcoder_client: AtCoderClient,
                     contest_id: str,
-                    config: Config):
-    retry_duration = 1.5
+                    config: Config,
+                    retry_delay_secs: float = 1.5,
+                    retry_max_delay_secs: float = 60,
+                    retry_max_tries: int = 10):
+    attempt_count = 1
     while True:
         problem_list = atcoder_client.download_problem_list(
             Contest(contest_id=contest_id))
         if problem_list:
             break
-        sleep(retry_duration)
+        if 0 < retry_max_tries < attempt_count:
+            raise EnvironmentInitializationError
         logger.warning(
-            "Failed to fetch. Will retry in {} seconds".format(retry_duration))
+            "Failed to fetch. Will retry in {} seconds. (Attempt {})".format(retry_delay_secs, attempt_count))
+        time.sleep(retry_delay_secs)
+        retry_delay_secs = min(retry_delay_secs * 2, retry_max_delay_secs)
+        attempt_count += 1
 
     tasks = [(atcoder_client,
               problem,
diff --git a/tests/test_envgen.py b/tests/test_envgen.py
@@ -2,14 +2,15 @@
 import shutil
 import tempfile
 import unittest
+from unittest import mock
 from os.path import relpath
 from logging import getLogger
 
 from atcodertools.client.atcoder import AtCoderClient
 from atcodertools.codegen.code_style_config import CodeStyleConfig
 from atcodertools.config.config import Config
 from atcodertools.config.etc_config import EtcConfig
-from atcodertools.tools.envgen import prepare_contest, main
+from atcodertools.tools.envgen import prepare_contest, main, EnvironmentInitializationError
 
 logger = getLogger(__name__)
 
@@ -76,6 +77,38 @@ def test_backup(self):
             )
         self.assertDirectoriesEqual(answer_data_dir_path, self.temp_dir)
 
+    @mock.patch('time.sleep')
+    def test_prepare_contest_aborts_after_max_retry_attempts(self, mock_sleep):
+        mock_client = mock.Mock(spec=AtCoderClient)
+        mock_client.download_problem_list.return_value = []
+        self.assertRaises(
+            EnvironmentInitializationError,
+            prepare_contest,
+            mock_client,
+            "agc029",
+            Config(
+                code_style_config=CodeStyleConfig(
+                    workspace_dir=self.temp_dir,
+                    template_file=TEMPLATE_PATH,
+                    lang="cpp",
+                ),
+                etc_config=EtcConfig(
+                    in_example_format="input_{}.txt",
+                    out_example_format="output_{}.txt"
+                ))
+        )
+        self.assertEqual(mock_sleep.call_count, 10)
+        mock_sleep.assert_has_calls([mock.call(1.5),
+                                     mock.call(3.0),
+                                     mock.call(6.0),
+                                     mock.call(12.0),
+                                     mock.call(24.0),
+                                     mock.call(48.0),
+                                     mock.call(60.0),
+                                     mock.call(60.0),
+                                     mock.call(60.0),
+                                     mock.call(60.0)])
+
     def assertDirectoriesEqual(self, expected_dir_path, dir_path):
         files1 = get_all_rel_file_paths(expected_dir_path)
         files2 = get_all_rel_file_paths(dir_path)
