fix: make port management better (#1926)

Author: sudharsan-selvaraj

src/CapabilityManager.ts

@@ -59,7 +59,7 @@ export async function androidCapabilities(
   caps.firstMatch[0]['appium:app'] = await findAppPath(caps);
   caps.firstMatch[0]['appium:udid'] = freeDevice.udid;
   caps.firstMatch[0]['appium:systemPort'] = await getFreePort(options.portRange);
-  caps.firstMatch[0]['appium:chromeDriverPort'] = await getPort();
+  caps.firstMatch[0]['appium:chromeDriverPort'] = await getFreePort(options.portRange);
   caps.firstMatch[0]['appium:adbRemoteHost'] = freeDevice.adbRemoteHost;
   caps.firstMatch[0]['appium:adbPort'] = freeDevice.adbPort;
   if (freeDevice.chromeDriverPath)

src/device-utils.ts

@@ -34,6 +34,7 @@ import DevicePlatform from './enums/Platform';
 import {
   cachePath,
   checkIfPathIsAbsolute,
+  getFreePort,
   isAppiumRunningAt,
   isDeviceFarmRunning,
   isMac,
@@ -241,7 +242,7 @@ export async function updateCapabilityForDevice(
 
   if (!device.hasOwnProperty('cloud')) {
     if (mergedCapabilites['appium:automationName']?.toLowerCase() === 'flutterintegration') {
-      capability.firstMatch[0]['appium:flutterSystemPort'] = await getPort();
+      capability.firstMatch[0]['appium:flutterSystemPort'] = await getFreePort(options.portRange);
     }
 
     if (device.platform.toLowerCase() == DevicePlatform.ANDROID) {

src/helpers.ts

@@ -15,6 +15,53 @@ import { FakeModuleLoader } from './fake-module-loader';
 import { IExternalModuleLoader } from './interfaces/IExternalModule';
 
 const APPIUM_VENDOR_PREFIX = 'appium:';
+
+// Port manager to track allocated ports
+class PortManager {
+  private allocatedPorts: Set<number> = new Set();
+
+  /**
+   * Check if a port is already allocated
+   */
+  isPortAllocated(port: number): boolean {
+    return this.allocatedPorts.has(port);
+  }
+
+  /**
+   * Allocate a port by adding it to the allocated set
+   */
+  allocatePort(port: number): void {
+    if (port && port > 0) {
+      this.allocatedPorts.add(port);
+      log.debug(`Port ${port} has been allocated`);
+    }
+  }
+
+  /**
+   * Release a port by removing it from the allocated set
+   */
+  releasePort(port: number): void {
+    if (port && port > 0 && this.allocatedPorts.has(port)) {
+      this.allocatedPorts.delete(port);
+      log.debug(`Port ${port} has been released`);
+    }
+  }
+
+  /**
+   * Release multiple ports
+   */
+  releasePorts(ports: (number | undefined | null)[]): void {
+    ports.forEach((port) => {
+      if (port && port > 0) {
+        this.releasePort(port);
+      }
+    });
+  }
+}
+
+// Singleton instance of port manager
+const portManager = new PortManager();
+
 export async function asyncForEach(
   array: string | any[],
   callback: {
@@ -69,16 +116,65 @@ export function checkIfPathIsAbsolute(configPath: string) {
   return path.isAbsolute(configPath);
 }
 
-export async function getFreePort(portRange?: string) {
+export async function getFreePort(portRange?: string): Promise<number> {
+  let port: number;
+
   if (portRange) {
     const range = portRange.split('-').map(Number);
     if (range.length !== 2 || isNaN(range[0]) || isNaN(range[1]) || range[0] > range[1]) {
       log.warn(`Invalid port range format: "${portRange}". Falling back to any free port.`);
+      port = await getPort();
+      // Check if allocated and try again if needed
+      while (portManager.isPortAllocated(port)) {
+        port = await getPort();
+      }
     } else {
-      return await getPort({ port: getPort.makeRange(range[0], range[1]) });
+      // Convert iterable to array and filter out allocated ports upfront
+      const portOptions = Array.from(getPort.makeRange(range[0], range[1]));
+      const availablePorts = portOptions.filter((p: number) => !portManager.isPortAllocated(p));
+
+      if (availablePorts.length === 0) {
+        log.warn(`All ports in range ${portRange} are allocated. Falling back to any free port.`);
+        port = await getPort();
+        // Check if allocated and try again if needed
+        while (portManager.isPortAllocated(port)) {
+          port = await getPort();
+        }
+      } else {
+        // Get a free port from the available (non-allocated) ports
+        // getPort will check system usage among these available ports
+        port = await getPort({ port: availablePorts });
+      }
+    }
+  } else {
+    port = await getPort();
+    // Check if the port is already allocated, keep trying until we find one that's not
+    while (portManager.isPortAllocated(port)) {
+      port = await getPort();
     }
   }
-  return await getPort();
+
+  // Allocate the port before returning it
+  portManager.allocatePort(port);
+  return port;
+}
+
+/**
+ * Release a port that was previously allocated
+ * @param port - The port number to release
+ */
+export function releasePort(port: number | undefined | null): void {
+  if (port && port > 0) {
+    portManager.releasePort(port);
+  }
+}
+
+/**
+ * Release multiple ports that were previously allocated
+ * @param ports - Array of port numbers to release
+ */
+export function releasePorts(ports: (number | undefined | null)[]): void {
+  portManager.releasePorts(ports);
 }
 
 export function nodeUrl(device: IDevice, basePath = ''): string {

src/plugin.ts

@@ -60,7 +60,7 @@ import _ from 'lodash';
 import { DeviceFarmApiClient } from './api-client';
 import { getDeviceFarmCapabilities } from './CapabilityManager';
 import { config, config as pluginConfig } from './config';
-import { getFreePort } from './helpers';
+import { getFreePort, releasePorts } from './helpers';
 import { ATDRepository } from './data-service/db';
 import { NodeService } from './data-service/node-service';
 import { addCLIArgs } from './data-service/pluginArgs';
@@ -716,6 +716,58 @@ class DevicePlugin extends BasePlugin {
         log.warn(`Error while releasing connection for device ${device.udid}. Error: ${err}`);
       }
     }
+
+    // Collect all ports used by the device and release them
+    const portsToRelease: (number | undefined | null)[] = [];
+
+    // Ports from device object
+    if (device) {
+      // iOS ports
+      if (device.wdaLocalPort) portsToRelease.push(device.wdaLocalPort);
+      if (device.mjpegServerPort) portsToRelease.push(device.mjpegServerPort);
+      if (device.goIOSAgentPort) portsToRelease.push(device.goIOSAgentPort);
+
+      // Android ports (stored on device if available)
+      if (device.systemPort) portsToRelease.push(device.systemPort);
+
+      // Ports from sessionResponse capabilities (Android and iOS)
+      // sessionResponse contains the capabilities returned from Appium
+      if (device.sessionResponse) {
+        const sessionResponse = device.sessionResponse;
+        // Check for ports with appium: prefix
+        const capabilities = sessionResponse.capabilities || sessionResponse;
+
+        // Android ports
+        const systemPort = capabilities['appium:systemPort'] || capabilities.systemPort;
+        const chromeDriverPort =
+          capabilities['appium:chromeDriverPort'] || capabilities.chromeDriverPort;
+        const flutterSystemPort =
+          capabilities['appium:flutterSystemPort'] || capabilities.flutterSystemPort;
+
+        // iOS ports
+        const wdaLocalPort = capabilities['appium:wdaLocalPort'] || capabilities.wdaLocalPort;
+
+        // Common ports
+        const mjpegServerPort =
+          capabilities['appium:mjpegServerPort'] || capabilities.mjpegServerPort;
+
+        if (systemPort) portsToRelease.push(systemPort);
+        if (chromeDriverPort) portsToRelease.push(chromeDriverPort);
+        if (flutterSystemPort) portsToRelease.push(flutterSystemPort);
+        if (wdaLocalPort) portsToRelease.push(wdaLocalPort);
+        if (mjpegServerPort) portsToRelease.push(mjpegServerPort);
+      }
+    }
+
+    // Release all collected ports
+    const validPorts = portsToRelease.filter((p) => p !== null && p !== undefined);
+    if (validPorts.length > 0) {
+      releasePorts(validPorts);
+      log.info(
+        `📱 Released ${validPorts.length} port(s) for session ${sessionId}: ${validPorts.join(', ')}`,
+      );
+    }
+
     return res;
   }
 }