Merge pull request #647 from FlorentinD/wait-with-exponential-retry

Use exponential backoff for waiting

Author: DarthMax

graphdatascience/session/aura_api.py

@@ -97,10 +97,15 @@ def list_sessions(self, dbid: str) -> List[SessionDetails]:
         return [SessionDetails.fromJson(s) for s in response.json()]
 
     def wait_for_session_running(
-        self, session_id: str, dbid: str, sleep_time: float = 0.2, max_sleep_time: float = 300
+        self,
+        session_id: str,
+        dbid: str,
+        sleep_time: float = 0.2,
+        max_sleep_time: float = 10,
+        max_wait_time: float = 300,
     ) -> WaitResult:
         waited_time = 0.0
-        while waited_time <= max_sleep_time:
+        while waited_time < max_wait_time:
             session = self.list_session(session_id, dbid)
             if session is None:
                 return WaitResult.from_error(f"Session `{session_id}` for database `{dbid}` not found -- please retry")
@@ -114,6 +119,7 @@ def wait_for_session_running(
                 )
             waited_time += sleep_time
             time.sleep(sleep_time)
+            sleep_time = min(sleep_time * 2, max_sleep_time, max_wait_time - waited_time)
 
         return WaitResult.from_error(
             f"Session `{session_id}` for database `{dbid}` is not running after {waited_time} seconds"
@@ -204,10 +210,10 @@ def list_instance(self, instance_id: str) -> Optional[InstanceSpecificDetails]:
         return InstanceSpecificDetails.fromJson(raw_data)
 
     def wait_for_instance_running(
-        self, instance_id: str, sleep_time: float = 0.2, max_sleep_time: float = 300
+        self, instance_id: str, sleep_time: float = 0.2, max_sleep_time: float = 10, max_wait_time: float = 300
     ) -> WaitResult:
         waited_time = 0.0
-        while waited_time <= max_sleep_time:
+        while waited_time < max_wait_time:
             instance = self.list_instance(instance_id)
             if instance is None:
                 return WaitResult.from_error("Instance is not found -- please retry")
@@ -223,6 +229,7 @@ def wait_for_instance_running(
                 )
             waited_time += sleep_time
             time.sleep(sleep_time)
+            sleep_time = min(sleep_time * 2, max_sleep_time, max_wait_time - waited_time)
 
         return WaitResult.from_error(f"Instance is not running after waiting for {waited_time} seconds")
 

graphdatascience/tests/unit/test_aura_api.py

@@ -204,8 +204,8 @@ def test_dont_wait_forever_for_session(requests_mock: Mocker, caplog: LogCapture
 
     with caplog.at_level(logging.DEBUG):
         assert (
-            "Session `id0` for database `dbid-1` is not running after 0.8 seconds"
-            in api.wait_for_session_running("id0", "dbid-1", max_sleep_time=0.7).error
+            "Session `id0` for database `dbid-1` is not running after 0.7 seconds"
+            in api.wait_for_session_running("id0", "dbid-1", max_wait_time=0.7).error
         )
 
     assert "Session `id0` is not yet running. Current status: Creating Host: foo.bar. Retrying in 0.2" in caplog.text
@@ -486,8 +486,8 @@ def test_dont_wait_forever(requests_mock: Mocker, caplog: LogCaptureFixture) ->
 
     with caplog.at_level(logging.DEBUG):
         assert (
-            "Instance is not running after waiting for 0.8"
-            in api.wait_for_instance_running("id0", max_sleep_time=0.7).error
+            "Instance is not running after waiting for 0.7"
+            in api.wait_for_instance_running("id0", max_wait_time=0.7).error
         )
 
     assert "Instance `id0` is not yet running. Current status: creating. Retrying in 0.2 seconds..." in caplog.text

graphdatascience/tests/unit/test_aurads_sessions.py

@@ -79,9 +79,13 @@ def list_instance(self, instance_id: str) -> Optional[InstanceSpecificDetails]:
             return None
 
     def wait_for_instance_running(
-        self, instance_id: str, sleep_time: float = 0.2, max_sleep_time: float = 300
+        self,
+        instance_id: str,
+        sleep_time: float = 0.2,
+        max_sleep_time: float = 5,
+        max_wait_time: float = 5,
     ) -> WaitResult:
-        return super().wait_for_instance_running(instance_id, sleep_time=0.0001, max_sleep_time=0.001)
+        return super().wait_for_instance_running(instance_id, sleep_time=0.0001, max_wait_time=0.001)
 
     def tenant_details(self) -> TenantDetails:
         return TenantDetails(id=self._tenant_id, ds_type="fake-ds", regions_per_provider={"aws": {"leipzig-1"}})

graphdatascience/tests/unit/test_dedicated_sessions.py

@@ -119,14 +119,27 @@ def list_instance(self, instance_id: str) -> Optional[InstanceSpecificDetails]:
             return None
 
     def wait_for_session_running(
-        self, session_id: str, dbid: str, sleep_time: float = 0.2, max_sleep_time: float = 5
+        self,
+        session_id: str,
+        dbid: str,
+        sleep_time: float = 0.2,
+        max_sleep_time: float = 5,
+        max_wait_time: float = 5,
     ) -> WaitResult:
-        return super().wait_for_session_running(session_id, dbid, sleep_time=0.0001, max_sleep_time=0.001)
+        return super().wait_for_session_running(
+            session_id, dbid, sleep_time=0.0001, max_wait_time=0.001, max_sleep_time=0.001
+        )
 
     def wait_for_instance_running(
-        self, instance_id: str, sleep_time: float = 0.2, max_sleep_time: float = 300
+        self,
+        instance_id: str,
+        sleep_time: float = 0.2,
+        max_sleep_time: float = 5,
+        max_wait_time: float = 5,
     ) -> WaitResult:
-        return super().wait_for_instance_running(instance_id, sleep_time=0.0001, max_sleep_time=0.001)
+        return super().wait_for_instance_running(
+            instance_id, sleep_time=0.0001, max_wait_time=0.001, max_sleep_time=0.001
+        )
 
     def tenant_details(self) -> TenantDetails:
         return TenantDetails(id=self._tenant_id, ds_type="fake-ds", regions_per_provider={"aws": {"leipzig-1"}})