From 8f46c97d47608d258d44a334b6d02428b8659bc1 Mon Sep 17 00:00:00 2001
From: Sajid Alam <sajid_alam@mckinsey.com>
Date: Wed, 6 Nov 2024 10:24:53 +0000
Subject: [PATCH] automatically increment port

Signed-off-by: Sajid Alam <sajid_alam@mckinsey.com>
---
 package/kedro_viz/launchers/cli/run.py | 18 ++++--------------
 package/kedro_viz/launchers/utils.py   | 23 +++++++++++++++++++++++
 2 files changed, 27 insertions(+), 14 deletions(-)

diff --git a/package/kedro_viz/launchers/cli/run.py b/package/kedro_viz/launchers/cli/run.py
index 02d27e84f..f869185a5 100644
--- a/package/kedro_viz/launchers/cli/run.py
+++ b/package/kedro_viz/launchers/cli/run.py
@@ -115,8 +115,8 @@ def run(
     from kedro_viz.launchers.utils import (
         _PYPROJECT,
         _check_viz_up,
+        _find_available_port,
         _find_kedro_project,
-        _is_port_in_use,
         _start_browser,
         _wait_for,
         display_cli_message,
@@ -147,17 +147,8 @@ def run(
             "yellow",
         )
 
-    # Check if the port is already in use
-    if _is_port_in_use(host, port):
-        display_cli_message(
-            f"Error: Port {port} is already in use. Kedro Viz could not start.",
-            "red",
-        )
-        display_cli_message(
-            "Please specify a different port using the '--port' option.",
-            "red",
-        )
-        return
+    port = _find_available_port(host, port)
+
     try:
         if port in _VIZ_PROCESSES and _VIZ_PROCESSES[port].is_alive():
             _VIZ_PROCESSES[port].terminate()
@@ -198,8 +189,7 @@ def run(
                 target=run_server, daemon=False, kwargs={**run_server_kwargs}
             )
 
-        display_cli_message("Starting Kedro Viz ...", "green")
-
+        display_cli_message(f"Starting Kedro Viz on port {port}...", "green")
         viz_process.start()
 
         _VIZ_PROCESSES[port] = viz_process
diff --git a/package/kedro_viz/launchers/utils.py b/package/kedro_viz/launchers/utils.py
index b74af23af..50f8e6e84 100644
--- a/package/kedro_viz/launchers/utils.py
+++ b/package/kedro_viz/launchers/utils.py
@@ -3,6 +3,7 @@
 
 import logging
 import socket
+import sys
 import webbrowser
 from pathlib import Path
 from time import sleep, time
@@ -86,6 +87,28 @@ def _is_port_in_use(host: str, port: int):
         return s.connect_ex((host, port)) == 0
 
 
+def _find_available_port(host: str, start_port: int, max_attempts: int = 5) -> int:
+    max_port = start_port + max_attempts - 1
+    port = start_port
+    while port <= max_port:
+        if not _is_port_in_use(host, port):
+            return port
+        display_cli_message(
+            f"Port {port} is already in use. Trying the next port...",
+            "yellow",
+        )
+        port += 1
+    display_cli_message(
+        f"Error: All ports in the range {start_port}-{max_port} are in use.",
+        "red",
+    )
+    display_cli_message(
+        "Please specify a different port using the '--port' option.",
+        "red",
+    )
+    sys.exit(1)
+
+
 def _is_localhost(host: str) -> bool:
     """Check whether a host is a localhost"""
     return host in ("127.0.0.1", "localhost", "0.0.0.0")