[JetBrains] Improve port forwarding logic

Tool: gitpod/catfood.gitpod.cloud

Author: mustard-mh

components/ide/jetbrains/backend-plugin/src/main/kotlin/io/gitpod/jetbrains/remote/AbstractGitpodPortForwardingService.kt

@@ -23,6 +23,8 @@ import kotlinx.coroutines.future.asDeferred
 import org.apache.http.client.utils.URIBuilder
 import java.util.*
 import java.util.concurrent.CompletableFuture
+import java.util.concurrent.ConcurrentHashMap
+import java.util.concurrent.atomic.AtomicReference
 
 @Suppress("UnstableApiUsage")
 abstract class AbstractGitpodPortForwardingService : GitpodPortForwardingService {
@@ -35,6 +37,15 @@ abstract class AbstractGitpodPortForwardingService : GitpodPortForwardingService
     private val ignoredPortsForNotificationService = service<GitpodIgnoredPortsForNotificationService>()
     private val lifetime = Lifetime.Eternal.createNested()
 
+    // Store current observed ports and their lifetime references
+    private val portLifetimes = ConcurrentHashMap<Int, Lifetime>()
+
+    // Prevent multiple sync operations from executing simultaneously
+    private val syncInProgress = AtomicReference(false)
+
+    // Last observed ports list, used to prevent duplicate processing
+    private var lastPortsList = listOf<PortsStatus>()
+
     init { start() }
 
     private fun start() {
@@ -86,7 +97,17 @@ abstract class AbstractGitpodPortForwardingService : GitpodPortForwardingService
             }
 
             override fun onNext(response: Status.PortsStatusResponse) {
-                application.invokeLater { syncPortsListWithClient(response) }
+                application.invokeLater {
+                    if (syncInProgress.compareAndSet(false, true)) {
+                        try {
+                            syncPortsListWithClient(response)
+                        } finally {
+                            syncInProgress.set(false)
+                        }
+                    } else {
+                        thisLogger().debug("gitpod: Sync already in progress, skipping update")
+                    }
+                }
             }
 
             override fun onCompleted() {
@@ -110,44 +131,97 @@ abstract class AbstractGitpodPortForwardingService : GitpodPortForwardingService
     private fun syncPortsListWithClient(response: Status.PortsStatusResponse) {
         val ignoredPorts = ignoredPortsForNotificationService.getIgnoredPorts()
         val portsList = response.portsList.filter { !ignoredPorts.contains(it.localPort) }
+
+        if (portsList == lastPortsList) {
+            thisLogger().debug("gitpod: Port list unchanged, skipping update")
+            return
+        }
+        lastPortsList = portsList
+
         val portsNumbersFromPortsList = portsList.map { it.localPort }
         val servedPorts = portsList.filter { it.served }
         val exposedPorts = servedPorts.filter { it.exposed?.url?.isNotBlank() ?: false }
         val portsNumbersFromNonServedPorts = portsList.filter { !it.served }.map { it.localPort }
+
+        // Clean up unused port lifetimes
+        val allCurrentPorts = perClientPortForwardingManager.getPorts()
+        val allPortsToKeep = mutableSetOf<Int>()
+
+        // Find ports that need to start forwarding
         val servedPortsToStartForwarding = servedPorts.filter {
             perClientPortForwardingManager.getPorts(it.localPort).none { p -> p.labels.contains(FORWARDED_PORT_LABEL) }
         }
+
+        // Find ports that need to start exposing on client
         val exposedPortsToStartExposingOnClient = exposedPorts.filter {
             perClientPortForwardingManager.getPorts(it.localPort).none { p -> p.labels.contains(EXPOSED_PORT_LABEL) }
         }
+
+        // Find ports that need to stop forwarding
         val forwardedPortsToStopForwarding = perClientPortForwardingManager.getPorts(FORWARDED_PORT_LABEL)
             .map { it.hostPortNumber }
             .filter { portsNumbersFromNonServedPorts.contains(it) || !portsNumbersFromPortsList.contains(it) }
+
+        // Find ports that need to stop exposing on client
         val exposedPortsToStopExposingOnClient = perClientPortForwardingManager.getPorts(EXPOSED_PORT_LABEL)
             .map { it.hostPortNumber }
             .filter { portsNumbersFromNonServedPorts.contains(it) || !portsNumbersFromPortsList.contains(it) }
 
-        servedPortsToStartForwarding.forEach { startForwarding(it) }
+        forwardedPortsToStopForwarding.forEach { stopForwarding(it) }
+        exposedPortsToStopExposingOnClient.forEach { stopExposingOnClient(it) }
 
-        exposedPortsToStartExposingOnClient.forEach { startExposingOnClient(it) }
+        servedPortsToStartForwarding.forEach {
+            startForwarding(it)
+            allPortsToKeep.add(it.localPort)
+        }
 
-        forwardedPortsToStopForwarding.forEach { stopForwarding(it) }
+        exposedPortsToStartExposingOnClient.forEach {
+            startExposingOnClient(it)
+            allPortsToKeep.add(it.localPort)
+        }
 
-        exposedPortsToStopExposingOnClient.forEach { stopExposingOnClient(it) }
+        portsList.forEach {
+            updatePortsPresentation(it)
+            allPortsToKeep.add(it.localPort)
+        }
+
+        cleanupUnusedLifetimes(allPortsToKeep)
+    }
 
-        portsList.forEach { updatePortsPresentation(it) }
+    private fun cleanupUnusedLifetimes(portsToKeep: Set<Int>) {
+        val portsToRemove = portLifetimes.keys.filter { !portsToKeep.contains(it) }
+        portsToRemove.forEach { port ->
+            portLifetimes[port]?.let { lifetime ->
+                thisLogger().debug("gitpod: Terminating lifetime for port $port")
+                lifetime.terminate()
+                portLifetimes.remove(port)
+            }
+        }
     }
 
     private fun startForwarding(portStatus: PortsStatus) {
         if (isLocalPortForwardingDisabled()) {
             return
         }
+
+        val portLifetime = getOrCreatePortLifetime(portStatus.localPort)
+
         try {
-            perClientPortForwardingManager.forwardPort(
+            thisLogger().debug("gitpod: Starting forwarding for port ${portStatus.localPort}")
+            val port = perClientPortForwardingManager.forwardPort(
                 portStatus.localPort,
                 PortType.TCP,
                 setOf(FORWARDED_PORT_LABEL),
             )
+
+            portLifetime.onTerminationOrNow {
+                thisLogger().debug("gitpod: Cleaning up port ${portStatus.localPort} due to lifetime termination")
+                try {
+                    perClientPortForwardingManager.removePort(port)
+                } catch (e: Exception) {
+                    thisLogger().warn("gitpod: Failed to remove port on lifetime termination", e)
+                }
+            }
         } catch (throwable: Throwable) {
             if (throwable !is PortAlreadyForwardedException) {
                 thisLogger().warn("gitpod: Caught an exception while forwarding port: ${throwable.message}")
@@ -156,23 +230,70 @@ abstract class AbstractGitpodPortForwardingService : GitpodPortForwardingService
     }
 
     private fun stopForwarding(hostPort: Int) {
-        perClientPortForwardingManager.getPorts(hostPort)
+        thisLogger().debug("gitpod: Stopping forwarding for port $hostPort")
+        val portsToRemove = perClientPortForwardingManager.getPorts(hostPort)
             .filter { it.labels.contains(FORWARDED_PORT_LABEL) }
-            .forEach { perClientPortForwardingManager.removePort(it) }
+
+        terminatePortLifetime(hostPort)
+
+        portsToRemove.forEach {
+            try {
+                perClientPortForwardingManager.removePort(it)
+            } catch (e: Exception) {
+                thisLogger().warn("gitpod: Failed to remove forwarded port $hostPort", e)
+            }
+        }
     }
 
     private fun startExposingOnClient(portStatus: PortsStatus) {
-        perClientPortForwardingManager.exposePort(
+        val portLifetime = getOrCreatePortLifetime(portStatus.localPort)
+
+        thisLogger().debug("gitpod: Starting exposing for port ${portStatus.localPort}")
+        val port = perClientPortForwardingManager.exposePort(
             portStatus.localPort,
             portStatus.exposed.url,
             setOf(EXPOSED_PORT_LABEL),
         )
+
+        portLifetime.onTerminationOrNow {
+            thisLogger().debug("gitpod: Cleaning up exposed port ${portStatus.localPort} due to lifetime termination")
+            try {
+                perClientPortForwardingManager.removePort(port)
+            } catch (e: Exception) {
+                thisLogger().warn("gitpod: Failed to remove exposed port on lifetime termination", e)
+            }
+        }
     }
 
     private fun stopExposingOnClient(hostPort: Int) {
-        perClientPortForwardingManager.getPorts(hostPort)
+        thisLogger().debug("gitpod: Stopping exposing for port $hostPort")
+        val portsToRemove = perClientPortForwardingManager.getPorts(hostPort)
             .filter { it.labels.contains(EXPOSED_PORT_LABEL) }
-            .forEach { perClientPortForwardingManager.removePort(it) }
+
+        terminatePortLifetime(hostPort)
+
+        portsToRemove.forEach {
+            try {
+                perClientPortForwardingManager.removePort(it)
+            } catch (e: Exception) {
+                thisLogger().warn("gitpod: Failed to remove exposed port $hostPort", e)
+            }
+        }
+    }
+
+    private fun getOrCreatePortLifetime(port: Int): Lifetime {
+        return portLifetimes.computeIfAbsent(port) {
+            thisLogger().debug("gitpod: Creating new lifetime for port $port")
+            lifetime.createNested()
+        }
+    }
+
+    private fun terminatePortLifetime(port: Int) {
+        portLifetimes[port]?.let { portLifetime ->
+            thisLogger().debug("gitpod: Terminating lifetime for port $port")
+            portLifetime.terminate()
+            portLifetimes.remove(port)
+        }
     }
 
     private fun updatePortsPresentation(portStatus: PortsStatus) {
@@ -212,6 +333,8 @@ abstract class AbstractGitpodPortForwardingService : GitpodPortForwardingService
     }
 
     override fun dispose() {
+        portLifetimes.values.forEach { it.terminate() }
+        portLifetimes.clear()
         lifetime.terminate()
     }
 }