Fast port allocation checker: move to its own file

Author: olethanh

src/aleph/vm/models.py

@@ -32,7 +32,7 @@
 )
 from aleph.vm.network.firewall import add_port_redirect_rule, remove_port_redirect_rule
 from aleph.vm.network.interfaces import TapInterface
-from aleph.vm.network.port_availability_checker import get_available_host_port
+from aleph.vm.network.port_availability_checker import fast_get_available_host_port
 from aleph.vm.orchestrator.metrics import (
     ExecutionRecord,
     delete_record,
@@ -65,9 +65,6 @@ def to_dict(self):
         return self.__dict__
 
 
-LAST_ASSIGNED_HOST_PORT = 24000
-
-
 class VmExecution:
     """
     Control the execution of a VM on a high level.
@@ -136,13 +133,9 @@ async def update_port_redirects(self, requested_ports: dict[int, dict[str, bool]
                     host_port = current["host"]
                     remove_port_redirect_rule(interface, host_port, vm_port, protocol)
             del self.mapped_ports[vm_port]
-
         for vm_port in redirect_to_add:
             target = requested_ports[vm_port]
-            host_port = get_available_host_port(start_port=LAST_ASSIGNED_HOST_PORT)
-            LAST_ASSIGNED_HOST_PORT = host_port
-            if LAST_ASSIGNED_HOST_PORT > 65535:
-                LAST_ASSIGNED_HOST_PORT = 24000
+            host_port = fast_get_available_host_port()
 
             for protocol in SUPPORTED_PROTOCOL_FOR_REDIRECT:
                 if target[protocol]:

src/aleph/vm/network/port_availability_checker.py

@@ -1,13 +1,12 @@
 import socket
-from typing import Optional
 
 from aleph.vm.network.firewall import check_nftables_redirections
 
 MIN_DYNAMIC_PORT = 24000
 MAX_PORT = 65535
 
 
-def get_available_host_port(start_port: Optional[int] = None) -> int:
+def get_available_host_port(start_port: int | None = None) -> int:
     """Find an available port on the host system.
 
     Args:
@@ -37,7 +36,31 @@ def get_available_host_port(start_port: Optional[int] = None) -> int:
 
             return port
 
-        except (socket.error, OSError):
+        except OSError:
             pass
 
     raise RuntimeError(f"No available ports found in range {MIN_DYNAMIC_PORT}-{MAX_PORT}")
+
+
+LAST_ASSIGNED_HOST_PORT = MIN_DYNAMIC_PORT
+
+
+def fast_get_available_host_port() -> int:
+    """Find an available port on the host system.
+    Use a global state to not start as each check may take several seconds and return a resulta faster
+
+    Args:
+        start_port: Optional starting port number. If not provided, starts from MIN_DYNAMIC_PORT
+
+    Returns:
+        An available port number
+
+    Raises:
+        RuntimeError: If no ports are available in the valid range
+    """
+    global LAST_ASSIGNED_HOST_PORT  # noqa: PLW0603
+    host_port = get_available_host_port(start_port=LAST_ASSIGNED_HOST_PORT)
+    LAST_ASSIGNED_HOST_PORT = host_port
+    if LAST_ASSIGNED_HOST_PORT > MAX_PORT:
+        LAST_ASSIGNED_HOST_PORT = MIN_DYNAMIC_PORT
+    return host_port