Merge branch 'main' into close_ended_lease

Author: evakhoni

packages/jumpstarter/jumpstarter/client/lease.py

@@ -23,6 +23,7 @@
 from grpc.aio import AioRpcError, Channel
 from jumpstarter_protocol import jumpstarter_pb2, jumpstarter_pb2_grpc
 from rich.console import Console
+from tenacity import retry, retry_if_exception_type, wait_exponential_jitter
 
 from .exceptions import LeaseError
 from jumpstarter.client import client_from_path
@@ -80,6 +81,24 @@ async def get(self):
             svc = ClientService(channel=self.channel, namespace=self.namespace)
             return await svc.GetLease(name=self.name)
 
+    @retry(
+        wait=wait_exponential_jitter(initial=1, max=120, jitter=1),
+        retry=retry_if_exception_type(ConnectionError),
+        reraise=True,
+    )
+    async def _get_with_retry(self):
+        """Get lease with exponential backoff retry on ConnectionError.
+
+        Retries with exponential backoff and jitter indefinitely when ConnectionError occurs.
+        The wait time between retries is capped at 2 minutes (120 seconds).
+        Jitter helps prevent thundering herd problems when multiple clients retry simultaneously.
+        """
+        try:
+            return await self.get()
+        except ConnectionError as e:
+            logger.error("Error while getting lease %s: %s", self.name, e)
+            raise
+
     def request(self):
         """Request a lease, or verifies a lease which was already created.
 
@@ -139,12 +158,11 @@ async def _acquire(self):
                 with LeaseAcquisitionSpinner(self.name) as spinner:
                     while True:
                         logger.debug("Polling Lease %s", self.name)
-                        result = await self.get()
-
+                        result = await self._get_with_retry()
                         # lease ready
                         if condition_true(result.conditions, "Ready"):
                             logger.debug("Lease %s acquired", self.name)
-                            spinner.update_status(f"Lease {self.name} acquired successfully!")
+                            spinner.update_status(f"Lease {self.name} acquired successfully!", force=True)
                             self.exporter_name = result.exporter
                             break
 
@@ -322,6 +340,8 @@ def __init__(self, lease_name: str | None = None):
         self.start_time = None
         self._should_show_spinner = self._is_terminal_available() and not self._is_non_interactive()
         self._current_message = None
+        self._last_log_time = None
+        self._log_throttle_interval = timedelta(minutes=5)
 
     def _is_non_interactive(self) -> bool:
         """Check if the user desires a NONINTERACTIVE environment."""
@@ -349,18 +369,32 @@ def __exit__(self, exc_type, exc_val, exc_tb):
         if self.spinner:
             self.spinner.stop()
 
-    def update_status(self, message: str):
-        """Update the spinner status message."""
+    def update_status(self, message: str, force: bool = False):
+        """Update the spinner status message.
+
+        :param message: The status message to display
+        :param force: If True, always log the message even when throttling (default: False)
+        """
         if self.spinner and self._should_show_spinner:
             self._current_message = f"[blue]{message}[/blue]"
             elapsed = datetime.now() - self.start_time
             elapsed_str = str(elapsed).split(".")[0]  # Remove microseconds
             self.spinner.update(f"{self._current_message} [dim]({elapsed_str})[/dim]")
         else:
             # Log info message when no console is available
-            elapsed = datetime.now() - self.start_time
-            elapsed_str = str(elapsed).split(".")[0]  # Remove microseconds
-            logger.info(f"{message} ({elapsed_str})")
+            # Throttle updates to at most every 5 minutes unless forced
+            now = datetime.now()
+            should_log = (
+                force
+                or self._last_log_time is None
+                or (now - self._last_log_time) >= self._log_throttle_interval
+            )
+
+            if should_log:
+                elapsed = now - self.start_time
+                elapsed_str = str(elapsed).split(".")[0]  # Remove microseconds
+                logger.info(f"{message} ({elapsed_str})")
+                self._last_log_time = now
 
     def tick(self):
         """Update the spinner with current elapsed time without changing the message."""

packages/jumpstarter/jumpstarter/client/lease_test.py

@@ -228,3 +228,103 @@ def test_start_time_initialization_in_context(self):
         with spinner:
             assert spinner.start_time is not None
             assert isinstance(spinner.start_time, datetime)
+
+    def test_throttling_first_update_logged(self, caplog):
+        """Test that the first update is always logged when console is not available."""
+        with patch.object(LeaseAcquisitionSpinner, '_is_terminal_available', return_value=False):
+            spinner = LeaseAcquisitionSpinner("test-lease")
+            spinner.start_time = datetime.now()
+
+            with caplog.at_level(logging.INFO):
+                spinner.update_status("First message")
+
+            assert "First message" in caplog.text
+            assert spinner._last_log_time is not None
+
+    def test_throttling_second_update_within_interval_not_logged(self, caplog):
+        """Test that updates within 5 minutes are not logged."""
+        with patch.object(LeaseAcquisitionSpinner, '_is_terminal_available', return_value=False):
+            spinner = LeaseAcquisitionSpinner("test-lease")
+            spinner.start_time = datetime.now()
+            spinner._last_log_time = datetime.now() - timedelta(minutes=2)  # 2 minutes ago
+
+            with caplog.at_level(logging.INFO):
+                spinner.update_status("Second message")
+
+            # Should not log because only 2 minutes have passed
+            assert "Second message" not in caplog.text
+
+    def test_throttling_update_after_interval_logged(self, caplog):
+        """Test that updates after 5 minutes are logged."""
+        with patch.object(LeaseAcquisitionSpinner, '_is_terminal_available', return_value=False):
+            spinner = LeaseAcquisitionSpinner("test-lease")
+            spinner.start_time = datetime.now()
+            spinner._last_log_time = datetime.now() - timedelta(minutes=6)  # 6 minutes ago
+
+            with caplog.at_level(logging.INFO):
+                spinner.update_status("After interval message")
+
+            assert "After interval message" in caplog.text
+            assert spinner._last_log_time is not None
+
+    def test_throttling_forced_update_always_logged(self, caplog):
+        """Test that forced updates are always logged regardless of throttle interval."""
+        with patch.object(LeaseAcquisitionSpinner, '_is_terminal_available', return_value=False):
+            spinner = LeaseAcquisitionSpinner("test-lease")
+            spinner.start_time = datetime.now()
+            spinner._last_log_time = datetime.now() - timedelta(minutes=1)  # 1 minute ago
+
+            with caplog.at_level(logging.INFO):
+                spinner.update_status("Forced message", force=True)
+
+            assert "Forced message" in caplog.text
+            assert spinner._last_log_time is not None
+
+    def test_throttling_multiple_updates_only_logs_when_needed(self, caplog):
+        """Test that multiple rapid updates only log at appropriate intervals."""
+        with patch.object(LeaseAcquisitionSpinner, '_is_terminal_available', return_value=False):
+            spinner = LeaseAcquisitionSpinner("test-lease")
+            spinner.start_time = datetime.now()
+
+            with caplog.at_level(logging.INFO):
+                # First update should be logged
+                spinner.update_status("Message 1")
+                assert "Message 1" in caplog.text
+
+                # Set last log time to recent
+                spinner._last_log_time = datetime.now() - timedelta(minutes=1)
+
+                # Second update should not be logged (within interval)
+                spinner.update_status("Message 2")
+                assert "Message 2" not in caplog.text
+
+                # Third update should not be logged (within interval)
+                spinner.update_status("Message 3")
+                assert "Message 3" not in caplog.text
+
+                # Set last log time to past the interval
+                spinner._last_log_time = datetime.now() - timedelta(minutes=6)
+
+                # Fourth update should be logged (past interval)
+                spinner.update_status("Message 4")
+                assert "Message 4" in caplog.text
+
+    def test_throttling_not_applied_when_console_available(self):
+        """Test that throttling is not applied when console is available."""
+        with patch.object(LeaseAcquisitionSpinner, '_is_terminal_available', return_value=True):
+            spinner = LeaseAcquisitionSpinner("test-lease")
+            spinner.start_time = datetime.now()
+
+            mock_spinner = Mock()
+            spinner.spinner = mock_spinner
+
+            # Multiple updates should all call update() regardless of throttle
+            spinner.update_status("Message 1")
+            spinner.update_status("Message 2")
+            spinner.update_status("Message 3")
+
+            # All should be called even if we set a recent last_log_time
+            spinner._last_log_time = datetime.now() - timedelta(minutes=1)
+            spinner.update_status("Message 4")
+
+            assert mock_spinner.update.call_count == 4

packages/jumpstarter/pyproject.toml

@@ -19,6 +19,7 @@ dependencies = [
     "xdg-base-dirs>=6.0.2",
     "pydantic-settings>=2.9.1",
     "rich>=14.0.0",
+    "tenacity>=8.2.0",
 ]
 
 [dependency-groups]