Clean up leases on timeout and termination

When lease is not available we try for 300s to acquire it. Users can also interrupt this with Ctrl+C while waiting. Let's remove such lease requests since after termination the shell is gone and it leaves behind a mostly useless lease.

Author: michalskrivanek

packages/jumpstarter-cli/jumpstarter_cli/shell.py

@@ -1,9 +1,12 @@
 import sys
 from datetime import timedelta
 
+import anyio
 import click
+from anyio import create_task_group, get_cancelled_exc_class
 from jumpstarter_cli_common.config import opt_config
 from jumpstarter_cli_common.exceptions import handle_exceptions_with_reauthentication
+from jumpstarter_cli_common.signal import signal_handler
 
 from .common import opt_duration_partial, opt_selector
 from .login import relogin_client
@@ -12,6 +15,55 @@
 from jumpstarter.config.exporter import ExporterConfigV1Alpha1
 
 
+def _run_shell_with_lease(lease, exporter_logs, config, command):
+    """Run shell with lease context managers."""
+    def launch_remote_shell(path: str) -> int:
+        return launch_shell(
+            path, "remote", config.drivers.allow, config.drivers.unsafe,
+            config.shell.use_profiles, command=command
+        )
+
+    with lease.serve_unix() as path:
+        with lease.monitor():
+            if exporter_logs:
+                with lease.connect() as client:
+                    with client.log_stream():
+                        return launch_remote_shell(path)
+            else:
+                return launch_remote_shell(path)
+
+
+async def _shell_with_signal_handling(config, selector, lease_name, duration, exporter_logs, command):
+    """Handle lease acquisition and shell execution with signal handling."""
+    exit_code = 0
+    cancelled_exc_class = get_cancelled_exc_class()
+
+    try:
+        async with create_task_group() as tg:
+            tg.start_soon(signal_handler, tg.cancel_scope)
+            try:
+                try:
+                    async with anyio.from_thread.BlockingPortal() as portal:
+                        async with config.lease_async(selector, lease_name, duration, portal) as lease:
+                            exit_code = await anyio.to_thread.run_sync(
+                                _run_shell_with_lease, lease, exporter_logs, config, command
+                            )
+                except BaseExceptionGroup as eg:
+                    for exc in eg.exceptions:
+                        if isinstance(exc, TimeoutError):
+                            raise exc from None
+                    raise
+                except cancelled_exc_class:
+                    exit_code = 2
+            finally:
+                if not tg.cancel_scope.cancel_called:
+                    tg.cancel_scope.cancel()
+    except* TimeoutError:
+        exit_code = 1
+
+    return exit_code
+
+
 @click.command("shell")
 @opt_config()
 @click.argument("command", nargs=-1)
@@ -38,27 +90,9 @@ def shell(config, command: tuple[str, ...], lease_name, selector, duration, expo
 
     match config:
         case ClientConfigV1Alpha1():
-            exit_code = 0
-            def _launch_remote_shell(path: str) -> int:
-                return launch_shell(
-                    path,
-                    "remote",
-                    config.drivers.allow,
-                    config.drivers.unsafe,
-                    config.shell.use_profiles,
-                    command=command,
-                )
-
-            with config.lease(selector=selector, lease_name=lease_name, duration=duration) as lease:
-                with lease.serve_unix() as path:
-                    with lease.monitor():
-                        if exporter_logs:
-                            with lease.connect() as client:
-                                with client.log_stream():
-                                    exit_code = _launch_remote_shell(path)
-                        else:
-                            exit_code = _launch_remote_shell(path)
-            # we exit here to make sure that all the with clauses unwind
+            exit_code = anyio.run(
+                _shell_with_signal_handling, config, selector, lease_name, duration, exporter_logs, command
+            )
             sys.exit(exit_code)
 
         case ExporterConfigV1Alpha1():

packages/jumpstarter/jumpstarter/client/lease.py

@@ -9,7 +9,14 @@
 from datetime import datetime, timedelta
 from typing import Any, Self
 
-from anyio import AsyncContextManagerMixin, ContextManagerMixin, create_task_group, fail_after, sleep
+from anyio import (
+    AsyncContextManagerMixin,
+    CancelScope,
+    ContextManagerMixin,
+    create_task_group,
+    fail_after,
+    sleep,
+)
 from anyio.from_thread import BlockingPortal
 from grpc.aio import Channel
 from jumpstarter_protocol import jumpstarter_pb2, jumpstarter_pb2_grpc
@@ -99,6 +106,7 @@ async def request_async(self):
                 await self._create()
         else:
             await self._create()
+
         return await self._acquire()
 
     async def _acquire(self):
@@ -138,15 +146,21 @@ async def _acquire(self):
 
     @asynccontextmanager
     async def __asynccontextmanager__(self) -> AsyncGenerator[Self]:
-        value = await self.request_async()
         try:
+            value = await self.request_async()
             yield value
         finally:
             if self.release:
                 logger.info("Releasing Lease %s", self.name)
-                await self.svc.DeleteLease(
-                    name=self.name,
-                )
+                # Shield cleanup from cancellation to ensure it completes
+                with CancelScope(shield=True):
+                    try:
+                        with fail_after(30):
+                            await self.svc.DeleteLease(
+                                name=self.name,
+                            )
+                    except TimeoutError:
+                        logger.warning("Timeout while deleting lease %s during cleanup", self.name)
 
     @contextmanager
     def __contextmanager__(self) -> Generator[Self]: