feat: add method to abort connections and improve new connections handler

Author: hlouzada

spyder/plugins/remoteclient/api/manager/base.py

@@ -83,7 +83,7 @@ def __init__(
 
         self.__installing_server = False
         self.__starting_server = False
-        self.__creating_connection = False
+        self.__connection_task: asyncio.Task | None = None
 
         # For logging
         self.logger = logging.getLogger(
@@ -100,9 +100,10 @@ def __init__(
     async def __create_events(self):
         self.__server_installed = asyncio.Event()
         self.__starting_event = asyncio.Event()
-        self.__connection_established = asyncio.Event()
+        self.__abort_requested = asyncio.Event()
+        self.__lock = asyncio.Lock()
 
-    def _emit_connection_status(self, status, message):
+    def _emit_connection_status(self, status: ConnectionStatus, message: str):
         if self._plugin is not None:
             self._plugin.sig_connection_status_changed.emit(
                 ConnectionInfo(
@@ -119,11 +120,22 @@ def server_started(self):
         return self.__starting_event.is_set() and not self.__starting_server
 
     @property
-    def connected(self):
-        """Check if SSH connection is open."""
+    def connected(self) -> bool:
+        """Check if a connection is currently established."""
         return (
-            self.__connection_established.is_set()
-            and not self.__creating_connection
+            self.__connection_task is not None
+            and self.__connection_task.done()
+            and not self.__connection_task.cancelled()
+            and not self.__connection_task.exception()
+            and self.__connection_task.result() is True
+        )
+
+    @property
+    def is_connecting(self) -> bool:
+        """Check if a connection attempt is ongoing."""
+        return (
+            self.__connection_task is not None
+            and not self.__connection_task.done()
         )
 
     @property
@@ -162,7 +174,7 @@ async def close(self):
         await self.close_connection()
 
     def _handle_connection_lost(self, exc: Exception | None = None):
-        self.__connection_established.clear()
+        self.__connection_task = None
         self.__starting_event.clear()
         self._port_forwarder = None
         if exc:
@@ -232,6 +244,10 @@ async def start_remote_server(self):
             return self.server_started
 
         self.__starting_server = True
+        self._emit_connection_status(
+            ConnectionStatus.Starting,
+            _("Starting Spyder remote services."),
+        )
         try:
             if await self._start_remote_server():
                 self.__starting_event.set()
@@ -278,20 +294,64 @@ async def _install_remote_server(self):
         raise NotImplementedError(msg)
 
     async def create_new_connection(self) -> bool:
-        if self.__creating_connection:
-            await self.__connection_established.wait()
-            return self.connected
+        """Create a new connection, supporting user abort."""
+        if self.connected:
+            self.logger.debug(
+                "Atempting to create a new connection with an existing for %s",
+                self.server_name,
+            )
+            await self.close_connection()
 
-        self.__creating_connection = True
-        try:
-            if await self._create_new_connection():
-                self.__connection_established.set()
-                return True
-        finally:
-            self.__creating_connection = False
+        async with self.__lock:
+            if self.__connection_task is not None:
+                return self.connected
 
-        self.__connection_established.clear()
-        return False
+            self._emit_connection_status(
+                ConnectionStatus.Connecting,
+                _("We're establishing the connection. Please be patient"),
+            )
+
+            self.__abort_requested.clear()
+            self.__connection_task = asyncio.create_task(self._create_new_connection())
+
+            await asyncio.wait(
+                [self.__connection_task, self.__abort_requested.wait()],
+                return_when=asyncio.FIRST_COMPLETED,
+            )
+
+            # User aborted
+            if self.__abort_requested.is_set():
+                self.__connection_task.cancel()
+                self.logger.warning("Connection attempt aborted by user")
+                self._emit_connection_status(
+                    ConnectionStatus.Inactive,
+                    _("The connection attempt was cancelled"),
+                )
+                return False
+
+            # Connection completed
+            try:
+                if await self.__connection_task:
+                    self._emit_connection_status(
+                        ConnectionStatus.Connected,
+                        _("The connection was successfully established"),
+                    )
+                    return True
+            except BaseException:
+                self.logger.exception(
+                    "Error creating a new connection for %s", self.server_name
+                )
+                self._emit_connection_status(
+                    ConnectionStatus.Error,
+                    _("There was an error establishing the connection"),
+                )
+
+            return False
+
+    async def abort_connection(self):
+        """Abort an ongoing connection attempt."""
+        if self.is_connecting:
+            self.__abort_requested.set()
 
     @abstractmethod
     async def _create_new_connection(self) -> bool:
@@ -337,7 +397,7 @@ async def close_connection(self):
 
     def _reset_connection_stablished(self):
         """Reset the connection status."""
-        self.__connection_established.clear()
+        self.__connection_task = None
 
     @staticmethod
     def get_free_port():

spyder/plugins/remoteclient/api/manager/jupyterhub.py

@@ -217,18 +217,7 @@ async def _install_remote_server(self):
         return False
 
     async def _create_new_connection(self) -> bool:
-        if self.connected:
-            self.logger.debug(
-                "Atempting to create a new connection with an existing for %s",
-                self.server_name,
-            )
-            await self.close_connection()
-
-        self._emit_connection_status(
-            ConnectionStatus.Connecting,
-            _("We're establishing the connection. Please be patient"),
-        )
-
+        """Create a new SSH connection."""
         self.logger.debug("Connecting to jupyterhub at %s", self.hub_url)
 
         self._session = aiohttp.ClientSession(
@@ -237,13 +226,9 @@ async def _create_new_connection(self) -> bool:
         )
 
         user_data = None
-        try:
-            async with self._session.get("hub/api/user") as response:
-                if response.ok:
-                    user_data = await response.json()
-        except aiohttp.ClientError:
-            self.logger.exception("Error connecting to JupyterHub")
-            return False
+        async with self._session.get("hub/api/user") as response:
+            if response.ok:
+                user_data = await response.json()
 
         if user_data is None:
             self.logger.error(

spyder/plugins/remoteclient/api/manager/ssh.py

@@ -329,7 +329,7 @@ async def _install_remote_server(self):
         return True
 
     async def _create_new_connection(self) -> bool:
-        """Creates a new SSH connection to the remote server machine.
+        """Create a new SSH connection to the remote server machine.
 
         Args
         ----
@@ -341,37 +341,18 @@ async def _create_new_connection(self) -> bool:
         bool
             True if the connection was successful, False otherwise.
         """
-        if self.connected:
-            self.logger.debug(
-                f"Atempting to create a new connection with an existing for "
-                f"{self.peer_host}"
-            )
-            await self.close_connection()
-
-        self._emit_connection_status(
-            ConnectionStatus.Connecting,
-            _("We're establishing the connection. Please be patient"),
-        )
-
         connect_kwargs = {
             k: v
             for k, v in self.options.items()
             if k not in self._extra_options
         }
         self.logger.debug("Opening SSH connection")
-        try:
-            self._ssh_connection = await asyncssh.connect(
-                **connect_kwargs, client_factory=self.client_factory
-            )
-        except (OSError, asyncssh.Error) as e:
-            self.logger.error(f"Failed to open ssh connection: {e}")
-            self._emit_connection_status(
-                ConnectionStatus.Error,
-                _("It was not possible to open a connection to this machine"),
-            )
-            return False
 
-        self.logger.info(f"SSH connection opened for {self.peer_host}")
+        self._ssh_connection = await asyncssh.connect(
+            **connect_kwargs, client_factory=self.client_factory
+        )
+
+        self.logger.info("SSH connection established for %s", self.peer_host)
 
         return True