fixup use Semaphore

Tool: gitpod/catfood.gitpod.cloud

Author: mustard-mh

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

@@ -26,14 +26,18 @@ import java.util.*
 import java.util.concurrent.CompletableFuture
 import java.util.concurrent.ConcurrentHashMap
 import java.util.concurrent.atomic.AtomicReference
-import kotlinx.coroutines.sync.Mutex
-import kotlinx.coroutines.sync.withLock
+import kotlinx.coroutines.sync.Semaphore
+import kotlinx.coroutines.sync.withPermit
 
 @Suppress("UnstableApiUsage")
 abstract class AbstractGitpodPortForwardingService : GitpodPortForwardingService {
     companion object {
         const val FORWARDED_PORT_LABEL = "ForwardedByGitpod"
         const val EXPOSED_PORT_LABEL = "ExposedByGitpod"
+        private const val MAX_CONCURRENT_OPERATIONS = 10
+        private const val BATCH_SIZE = 10
+        private const val BATCH_DELAY = 100L
+        private const val DEBOUNCE_DELAY = 500L
     }
 
     private val perClientPortForwardingManager = service<PerClientPortForwardingManager>()
@@ -52,12 +56,8 @@ abstract class AbstractGitpodPortForwardingService : GitpodPortForwardingService
     // Debounce job for port updates
     private var debounceJob: Job? = null
 
-    // Batch size for port operations
-    private val batchSize = 10
-    private val batchDelay = 100L
-
-    // Mutex to limit concurrent operations
-    private val operationMutex = Mutex()
+    // Semaphore to limit concurrent operations
+    private val operationSemaphore = Semaphore(MAX_CONCURRENT_OPERATIONS)
 
     init { start() }
 
@@ -82,7 +82,6 @@ abstract class AbstractGitpodPortForwardingService : GitpodPortForwardingService
                     is InterruptedException, is CancellationException -> {
                         cancel("gitpod: Stopped observing ports list due to an expected interruption.")
                     }
-
                     else -> {
                         thisLogger().warn(
                             "gitpod: Got an error while trying to get ports list from Supervisor. " +
@@ -98,9 +97,7 @@ abstract class AbstractGitpodPortForwardingService : GitpodPortForwardingService
 
     private fun observePortsList(): CompletableFuture<Void> {
         val completableFuture = CompletableFuture<Void>()
-
         val statusServiceStub = StatusServiceGrpc.newStub(GitpodManager.supervisorChannel)
-
         val portsStatusRequest = Status.PortsStatusRequest.newBuilder().setObserve(true).build()
 
         val portsStatusResponseObserver = object :
@@ -110,12 +107,9 @@ abstract class AbstractGitpodPortForwardingService : GitpodPortForwardingService
             }
 
             override fun onNext(response: Status.PortsStatusResponse) {
-                // Cancel previous debounce job if exists
                 debounceJob?.cancel()
-
-                // Create new debounce job
                 debounceJob = runJob(lifetime) {
-                    delay(500)
+                    delay(DEBOUNCE_DELAY)
                     try {
                         if (syncInProgress.compareAndSet(false, true)) {
                             try {
@@ -142,7 +136,6 @@ abstract class AbstractGitpodPortForwardingService : GitpodPortForwardingService
         }
 
         statusServiceStub.portsStatus(portsStatusRequest, portsStatusResponseObserver)
-
         return completableFuture
     }
 
@@ -165,7 +158,6 @@ abstract class AbstractGitpodPortForwardingService : GitpodPortForwardingService
         val exposedPorts = servedPorts.filter { it.exposed?.url?.isNotBlank() ?: false }
         val portsNumbersFromNonServedPorts = portsList.filter { !it.served }.map { it.localPort }
 
-        // Clean up unused port lifetimes
         val allPortsToKeep = mutableSetOf<Int>()
 
         val servedPortsToStartForwarding = servedPorts.filter {
@@ -184,52 +176,35 @@ abstract class AbstractGitpodPortForwardingService : GitpodPortForwardingService
             .map { it.hostPortNumber }
             .filter { portsNumbersFromNonServedPorts.contains(it) || !portsNumbersFromPortsList.contains(it) }
 
-        // Process port changes in background with batching
         runJob(lifetime) {
             try {
-                forwardedPortsToStopForwarding.chunked(batchSize).forEach { batch ->
-                    batch.forEach { port ->
-                        operationMutex.withLock { stopForwarding(port) }
-                    }
-                    delay(batchDelay)
+                processPortsInBatches(forwardedPortsToStopForwarding) { port ->
+                    operationSemaphore.withPermit { stopForwarding(port) }
                 }
 
-                exposedPortsToStopExposingOnClient.chunked(batchSize).forEach { batch ->
-                    batch.forEach { port ->
-                        operationMutex.withLock { stopExposingOnClient(port) }
-                    }
-                    delay(batchDelay)
+                processPortsInBatches(exposedPortsToStopExposingOnClient) { port ->
+                    operationSemaphore.withPermit { stopExposingOnClient(port) }
                 }
 
-                servedPortsToStartForwarding.chunked(batchSize).forEach { batch ->
-                    batch.forEach { port ->
-                        operationMutex.withLock {
-                            startForwarding(port)
-                            allPortsToKeep.add(port.localPort)
-                        }
+                processPortsInBatches(servedPortsToStartForwarding) { port ->
+                    operationSemaphore.withPermit {
+                        startForwarding(port)
+                        allPortsToKeep.add(port.localPort)
                     }
-                    delay(batchDelay)
                 }
 
-                exposedPortsToStartExposingOnClient.chunked(batchSize).forEach { batch ->
-                    batch.forEach { port ->
-                        operationMutex.withLock {
-                            startExposingOnClient(port)
-                            allPortsToKeep.add(port.localPort)
-                        }
+                processPortsInBatches(exposedPortsToStartExposingOnClient) { port ->
+                    operationSemaphore.withPermit {
+                        startExposingOnClient(port)
+                        allPortsToKeep.add(port.localPort)
                     }
-                    delay(batchDelay)
                 }
 
-                // Update presentation for all ports in batches
-                portsList.chunked(batchSize).forEach { batch ->
+                processPortsInBatches(portsList) { port ->
                     application.invokeLater {
-                        batch.forEach {
-                            updatePortsPresentation(it)
-                            allPortsToKeep.add(it.localPort)
-                        }
+                        updatePortsPresentation(port)
+                        allPortsToKeep.add(port.localPort)
                     }
-                    delay(batchDelay)
                 }
 
                 cleanupUnusedLifetimes(allPortsToKeep)
@@ -239,9 +214,17 @@ abstract class AbstractGitpodPortForwardingService : GitpodPortForwardingService
         }
     }
 
+    private suspend fun <T> processPortsInBatches(ports: List<T>, action: suspend (T) -> Unit) {
+        ports.chunked(BATCH_SIZE).forEach { batch ->
+            batch.forEach { port ->
+                action(port)
+            }
+            delay(BATCH_DELAY)
+        }
+    }
+
     private fun cleanupUnusedLifetimes(portsToKeep: Set<Int>) {
-        val portsToRemove = portLifetimes.keys.filter { !portsToKeep.contains(it) }
-        portsToRemove.forEach { port ->
+        portLifetimes.keys.filter { !portsToKeep.contains(it) }.forEach { port ->
             portLifetimes[port]?.let { lifetime ->
                 thisLogger().debug("gitpod: Terminating lifetime for port $port")
                 lifetime.terminate()
@@ -251,9 +234,7 @@ abstract class AbstractGitpodPortForwardingService : GitpodPortForwardingService
     }
 
     private fun startForwarding(portStatus: PortsStatus) {
-        if (isLocalPortForwardingDisabled()) {
-            return
-        }
+        if (isLocalPortForwardingDisabled()) return
 
         val portLifetime = getOrCreatePortLifetime(portStatus.localPort)
 
@@ -332,12 +313,11 @@ abstract class AbstractGitpodPortForwardingService : GitpodPortForwardingService
         }
     }
 
-    private fun getOrCreatePortLifetime(port: Int): Lifetime {
-        return portLifetimes.computeIfAbsent(port) {
+    private fun getOrCreatePortLifetime(port: Int): Lifetime =
+        portLifetimes.computeIfAbsent(port) {
             thisLogger().debug("gitpod: Creating new lifetime for port $port")
             lifetime.createNested()
         }
-    }
 
     private fun terminatePortLifetime(port: Int) {
         portLifetimes[port]?.let { portLifetime ->
@@ -349,39 +329,42 @@ abstract class AbstractGitpodPortForwardingService : GitpodPortForwardingService
 
     private fun updatePortsPresentation(portStatus: PortsStatus) {
         perClientPortForwardingManager.getPorts(portStatus.localPort).forEach {
-            if (it.configuration.isForwardedPort()) {
-                it.presentation.name = portStatus.name
-                it.presentation.description = portStatus.description
-                it.presentation.tooltip = "Forwarded"
-                it.presentation.icon = RowIcon(AllIcons.Actions.Commit)
-            } else if (it.configuration.isExposedPort()) {
-                val isPubliclyExposed = (portStatus.exposed.visibility == Status.PortVisibility.public_visibility)
-
-                it.presentation.name = portStatus.name
-                it.presentation.description = portStatus.description
-                it.presentation.tooltip = "Exposed (${if (isPubliclyExposed) "Public" else "Private"})"
-                it.presentation.icon = if (isPubliclyExposed) {
-                    RowIcon(AllIcons.Actions.Commit)
-                } else {
-                    RowIcon(AllIcons.Actions.Commit, AllIcons.Diff.Lock)
+            when {
+                it.configuration.isForwardedPort() -> {
+                    it.presentation.name = portStatus.name
+                    it.presentation.description = portStatus.description
+                    it.presentation.tooltip = "Forwarded"
+                    it.presentation.icon = RowIcon(AllIcons.Actions.Commit)
+                }
+                it.configuration.isExposedPort() -> {
+                    val isPubliclyExposed = (portStatus.exposed.visibility == Status.PortVisibility.public_visibility)
+                    it.presentation.name = portStatus.name
+                    it.presentation.description = portStatus.description
+                    it.presentation.tooltip = "Exposed (${if (isPubliclyExposed) "Public" else "Private"})"
+                    it.presentation.icon = if (isPubliclyExposed) {
+                        RowIcon(AllIcons.Actions.Commit)
+                    } else {
+                        RowIcon(AllIcons.Actions.Commit, AllIcons.Diff.Lock)
+                    }
                 }
             }
         }
     }
 
-    override fun getLocalHostUriFromHostPort(hostPort: Int): Optional<URI> {
-        val forwardedPort = perClientPortForwardingManager.getPorts(hostPort).firstOrNull {
-            it.configuration.isForwardedPort()
-        } ?: return Optional.empty()
-
-        (forwardedPort.configuration as PortConfiguration.PerClientTcpForwarding).clientPortState.let {
-            return if (it is ClientPortState.Assigned) {
-                Optional.of(URIBuilder().setScheme("http").setHost(it.clientInterface).setPort(it.clientPort).build())
-            } else {
-                Optional.empty()
-            }
-        }
-    }
+    override fun getLocalHostUriFromHostPort(hostPort: Int): Optional<URI> =
+        perClientPortForwardingManager.getPorts(hostPort)
+            .firstOrNull { it.configuration.isForwardedPort() }
+            ?.let { forwardedPort ->
+                (forwardedPort.configuration as PortConfiguration.PerClientTcpForwarding)
+                    .clientPortState
+                    .let {
+                        if (it is ClientPortState.Assigned) {
+                            Optional.of(URIBuilder().setScheme("http").setHost(it.clientInterface).setPort(it.clientPort).build())
+                        } else {
+                            Optional.empty()
+                        }
+                    }
+            } ?: Optional.empty()
 
     override fun dispose() {
         portLifetimes.values.forEach { it.terminate() }