Add JetBrainsPluginClientManager to scan ports in parallel

This speeds up the port scan considerably (was 1s per port without server on Win).

[skip ci]

src/serena/jetbrains/jetbrains_plugin_client.py

@@ -2,9 +2,12 @@
 Client for the Serena JetBrains Plugin
 """
 
+import concurrent
 import json
 import logging
 import re
+import threading
+from concurrent.futures.thread import ThreadPoolExecutor
 from pathlib import Path
 from typing import Any, Literal, Optional, Self, TypeVar, cast
 
@@ -15,6 +18,7 @@ from sensai.util.string import ToStringMixin
 import serena.jetbrains.jetbrains_types as jb
 from serena.jetbrains.jetbrains_types import PluginStatusDTO
 from serena.project import Project
+from serena.util.class_decorators import singleton
 from serena.util.text_utils import render_html
 from serena.util.version import Version
 
@@ -59,6 +63,60 @@ class ServerNotFoundError(Exception):
     """Raised when the plugin's service is not found."""
 
 
+@singleton
+class JetBrainsPluginClientManager:
+    """
+    Manager for JetBrainsPluginClient instances, responsible for scanning ports to find available plugin instances
+    """
+
+    NUM_PORTS_TO_SCAN = 20
+
+    def __init__(self) -> None:
+        self._clients: dict[int, "JetBrainsPluginClient"] = {}
+        self._lock = threading.Lock()
+
+    def _submit_scan(self) -> list[concurrent.futures.Future["JetBrainsPluginClient"]]:
+        """
+        Performs a port scan to find available plugin instances in parallel.
+
+        :return: futures that will resolve to plugin clients for every port
+        """
+
+        def scan_port(port: int) -> JetBrainsPluginClient:
+            client = JetBrainsPluginClient(port)
+            with self._lock:
+                self._clients[port] = client
+            return client
+
+        futures = []
+        with ThreadPoolExecutor(max_workers=self.NUM_PORTS_TO_SCAN) as executor:
+            for i in range(self.NUM_PORTS_TO_SCAN):
+                future = executor.submit(scan_port, JetBrainsPluginClient.BASE_PORT + i)
+                futures.append(future)
+        return futures
+
+    def find_client(self, project_root: Path) -> "JetBrainsPluginClient":
+        plugin_paths_found = []
+        for future in self._submit_scan():
+            client = future.result()
+            if client.matches(project_root):
+                return client
+            elif client.project_root is not None:
+                plugin_paths_found.append(client.project_root)
+
+        log.warning(
+            "Searched for Serena JetBrains plugin service for project at %s but found no matching service. "
+            "Found plugin instances for the following project paths: %s",
+            project_root,
+            plugin_paths_found,
+        )
+        raise ServerNotFoundError(
+            f"Found no Serena service in a JetBrains IDE instance for the project at {project_root}. "
+            "STOP. Do not attempt any other tools or workarounds. Ask the user to open this folder as a project in a JetBrains IDE "
+            "with the Serena plugin installed and running!"
+        )
+
+
 class JetBrainsPluginClient(ToStringMixin):
     """
     Python client for the Serena Backend Service.
@@ -87,11 +145,11 @@ class JetBrainsPluginClient(ToStringMixin):
         self._session.headers.update({"Content-Type": "application/json", "Accept": "application/json"})
 
         # connect and obtain status
-        self._project_root: str | None = None
+        self.project_root: str | None = None
         self._plugin_version: Version | None = None
         try:
             status_response: PluginStatusDTO = cast(jb.PluginStatusDTO, self._make_request("GET", "/status"))
-            self._project_root = status_response["project_root"]
+            self.project_root = status_response["project_root"]
             self._plugin_version = Version(status_response["plugin_version"])
         except ConnectionError:  # expected if no server is running at the port
             pass
@@ -107,7 +165,7 @@ class JetBrainsPluginClient(ToStringMixin):
         cls._server_address = address
 
     def _tostring_includes(self) -> list[str]:
-        return ["_port", "_project_root", "_plugin_version"]
+        return ["_port", "project_root", "_plugin_version"]
 
     @classmethod
     def from_project(cls, project: Project) -> Self:
@@ -118,27 +176,9 @@ class JetBrainsPluginClient(ToStringMixin):
             if client.matches(resolved_path):
                 return client
 
-        plugin_paths_found = []
-        for port in range(cls.BASE_PORT, cls.BASE_PORT + 20):
-            client = JetBrainsPluginClient(port)
-            if client.matches(resolved_path):
-                log.info("Found matching %s", client)
-                cls._last_port = port
-                return client
-            elif client._project_root is not None:
-                plugin_paths_found.append(client._project_root)
-
-        log.warning(
-            "Searched for Serena JetBrains plugin service for project at %s but found no matching service. "
-            "Found plugin instances for the following project paths: %s",
-            resolved_path,
-            plugin_paths_found,
-        )
-        raise ServerNotFoundError(
-            f"Found no Serena service in a JetBrains IDE instance for the project at {resolved_path}. "
-            "STOP. Do not attempt any other tools or workarounds. Ask the user to open this folder as a project in a JetBrains IDE "
-            "with the Serena plugin installed and running!"
-        )
+        client = JetBrainsPluginClientManager().find_client(resolved_path)
+        cls._last_port = client._port
+        return client
 
     @staticmethod
     def _paths_match(resolved_serena_path: str, plugin_path: str) -> bool:
@@ -184,9 +224,9 @@ class JetBrainsPluginClient(ToStringMixin):
         :param resolved_path: the resolved project root path from Serena's perspective
         :return: whether this client instance matches the given project path
         """
-        if self._project_root is None:
+        if self.project_root is None:
             return False
-        return self._paths_match(str(resolved_path), self._project_root)
+        return self._paths_match(str(resolved_path), self.project_root)
 
     def is_version_at_least(self, *version_parts: int) -> bool:
         if self._plugin_version is None: