[lsp] fix lsp server freeze when a client disconnects during channel read

- Using ktor here seems to help here
- Also, introduce a multiclient mode flag to stope the server after a client disconnects (according to the LSP spec)

LSP-166

GitOrigin-RevId: a739d097e02156828d68973865cacc484e38ff81

Author: darthorimar

kotlin-lsp/BUILD.bazel

@@ -51,6 +51,9 @@ jvm_library(
     "@community//fleet/rhizomedb",
     "@community//fleet/kernel",
     "@community//platform/service-container",
+    "@community//libraries/io",
+    "@lib//:io-ktor-io",
+    "@lib//:ktor-network-tls",
   ]
 )
 ### auto-generated section `build language-server.kotlin-lsp` end

kotlin-lsp/language-server.kotlin-lsp.iml

@@ -58,5 +58,7 @@
     <orderEntry type="module" module-name="fleet.rhizomedb" />
     <orderEntry type="module" module-name="fleet.kernel" />
     <orderEntry type="module" module-name="intellij.platform.serviceContainer" />
+    <orderEntry type="library" name="ktor-network-tls" level="project" />
+    <orderEntry type="library" name="kotlinx-io-core" level="project" />
   </component>
 </module>

kotlin-lsp/src/com/jetbrains/ls/kotlinLsp/KotlinLspServer.kt

@@ -22,14 +22,14 @@ import com.jetbrains.ls.kotlinLsp.util.addKotlinStdlib
 import com.jetbrains.ls.kotlinLsp.util.logSystemInfo
 import com.jetbrains.ls.snapshot.api.impl.core.createServerStarterAnalyzerImpl
 import com.jetbrains.lsp.implementation.*
-import kotlinx.coroutines.*
+import kotlinx.coroutines.CompletableDeferred
+import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.awaitCancellation
+import kotlinx.coroutines.runBlocking
 import org.jetbrains.kotlin.idea.base.plugin.artifacts.KotlinArtifacts
 import org.jetbrains.kotlin.idea.compiler.configuration.KotlinPluginLayoutMode
 import org.jetbrains.kotlin.idea.compiler.configuration.KotlinPluginLayoutModeProvider
 import org.jetbrains.kotlin.idea.compiler.configuration.isRunningFromSources
-import java.io.InputStream
-import java.io.OutputStream
-import java.net.Socket
 import kotlin.io.path.absolutePathString
 import kotlin.io.path.createTempDirectory
 
@@ -45,11 +45,19 @@ private class RunKotlinLspCommand : CliktCommand(name = "kotlin-lsp") {
         .help("Whether the Kotlin LSP server is used in client mode. If not set, server mode will be used with a port specified by `${::socket.name}`")
         .validate { if (it && stdio) fail("Can't use stdio mode with client mode") }
 
+    val multiclient: Boolean by option().flag()
+        .help("Whether the Kotlin LSP server is used in multiclient mode. If not set, server will be shut down after the first client disconnects.`")
+        .validate {
+            if (it && stdio) fail("Stdio mode doesn't support multiclient mode")
+            if (it && client) fail("Client mode doesn't support multiclient mode")
+        }
+
+
     private fun createRunConfig(): KotlinLspServerRunConfig {
         val mode = when {
             stdio -> KotlinLspServerMode.Stdio
-            client -> KotlinLspServerMode.Socket.Client(socket)
-            else -> KotlinLspServerMode.Socket.Server(socket)
+            client -> KotlinLspServerMode.Socket(TcpConnectionConfig.Client(port = socket))
+            else -> KotlinLspServerMode.Socket(TcpConnectionConfig.Server(port = socket, isMulticlient = multiclient))
         }
         return KotlinLspServerRunConfig(mode)
     }
@@ -76,16 +84,17 @@ private fun run(runConfig: KotlinLspServerRunConfig) {
                 KotlinLspServerMode.Stdio -> {
                     val stdout = System.out
                     System.setOut(System.err)
-                    handleRequests(System.`in`, stdout, config, mode)
+                    stdioConnection(System.`in`, stdout) { connection ->
+                        handleRequests(connection, config, mode)
+                    }
                 }
 
                 is KotlinLspServerMode.Socket -> {
                     logSystemInfo()
                     tcpConnection(
-                        clientMode = mode is KotlinLspServerMode.Socket.Client,
-                        port = mode.port,
+                        mode.config,
                     ) { connection ->
-                        handleRequests(connection.inputStream, connection.outputStream, config, mode)
+                        handleRequests(connection, config, mode)
                     }
                 }
             }
@@ -94,11 +103,19 @@ private fun run(runConfig: KotlinLspServerRunConfig) {
 }
 
 context(LSServerContext)
-private suspend fun handleRequests(input: InputStream, output: OutputStream, config: LSConfiguration, mode: KotlinLspServerMode) {
-    withBaseProtocolFraming(input, output) { incoming, outgoing ->
+private suspend fun handleRequests(connection: LspConnection, config: LSConfiguration, mode: KotlinLspServerMode) {
+    val shutdownOnExitSignal = when (mode) {
+        is KotlinLspServerMode.Socket -> when (val tcpConfig = mode.config) {
+            is TcpConnectionConfig.Client -> true
+            is TcpConnectionConfig.Server -> !tcpConfig.isMulticlient
+        }
+        KotlinLspServerMode.Stdio -> true
+    }
+    val exitSignal = if (shutdownOnExitSignal) CompletableDeferred<Unit>() else null
+
+    withBaseProtocolFraming(connection, exitSignal) { incoming, outgoing ->
         withServer {
-            val exitSignal = CompletableDeferred<Unit>()
-            val handler = createLspHandlers(config, exitSignal, clientMode = mode is KotlinLspServerMode.Socket.Client)
+            val handler = createLspHandlers(config, exitSignal)
 
             withLsp(
                 incoming,
@@ -108,7 +125,11 @@ private suspend fun handleRequests(input: InputStream, output: OutputStream, con
                     Client.contextElement(lspClient)
                 },
             ) { lsp ->
-                exitSignal.await()
+                if (exitSignal != null) {
+                    exitSignal.await()
+                } else {
+                    awaitCancellation()
+                }
             }
         }
     }
@@ -164,12 +185,12 @@ fun createConfiguration(
 }
 
 context(LSServer)
-fun createLspHandlers(config: LSConfiguration, exitSignal: CompletableDeferred<Unit>, clientMode: Boolean = false): LspHandlers {
+fun createLspHandlers(config: LSConfiguration, exitSignal: CompletableDeferred<Unit>?): LspHandlers {
     with(config) {
         return lspHandlers {
             initializeRequest()
             setTraceNotification()
-            shutdownRequest(clientMode, exitSignal)
+            shutdownRequest(exitSignal)
             fileUpdateRequests()
             features()
         }

kotlin-lsp/src/com/jetbrains/ls/kotlinLsp/KotlinLspServerRunConfig.kt

@@ -1,16 +1,13 @@
 package com.jetbrains.ls.kotlinLsp
 
+import com.jetbrains.lsp.implementation.TcpConnectionConfig
+
 class KotlinLspServerRunConfig(
     val mode: KotlinLspServerMode
 )
 
 sealed interface KotlinLspServerMode {
     object Stdio : KotlinLspServerMode
 
-    sealed interface Socket : KotlinLspServerMode {
-        val port: Int
-
-        class Server(override val port: Int) : Socket
-        class Client(override val port: Int) : Socket
-    }
-}
+    data class Socket(val config: TcpConnectionConfig) : KotlinLspServerMode
+}

kotlin-lsp/src/com/jetbrains/ls/kotlinLsp/ktorHack/SinkByteWriteChannel.kt

@@ -0,0 +1,93 @@
+/*
+ * Copyright 2014-2025 JetBrains s.r.o and contributors. Use of this source code is governed by the Apache 2.0 license.
+ */
+
+@file:Suppress("INVISIBLE_REFERENCE")
+/*
+Copied from https://github.com/ktorio/ktor/pull/4594 until ktor is not updated in monorepo
+ */
+
+package com.jetbrains.ls.kotlinLsp.ktorHack
+
+import io.ktor.utils.io.ByteWriteChannel
+import io.ktor.utils.io.CLOSED
+import io.ktor.utils.io.CloseToken
+import io.ktor.utils.io.InternalAPI
+import kotlinx.io.IOException
+import kotlinx.io.RawSink
+import kotlinx.io.Sink
+import kotlinx.io.asSink
+import kotlinx.io.buffered
+import java.io.OutputStream
+
+/**
+ * Creates a [io.ktor.utils.io.ByteWriteChannel] that writes to this [Sink].
+ *
+ * Example usage:
+ * ```kotlin
+ * suspend fun writeMessage(raw: RawSink) {
+ *     val channel = raw.asByteWriteChannel()
+ *     channel.writeByte(42)
+ *     channel.flushAndClose()
+ * }
+ *
+ * val buffer = Buffer()
+ * writeMessage(buffer)
+ * buffer.readByte() // 42
+ * ```
+ *
+ * Please note that the channel will be buffered even if the sink is not.
+ */
+public fun RawSink.asByteWriteChannel(): ByteWriteChannel = SinkByteWriteChannel(this)
+
+@Suppress("INVISIBLE_REFERENCE")
+internal class SinkByteWriteChannel(origin: RawSink) : ByteWriteChannel {
+    @Volatile
+    var closed: CloseToken? = null
+    private val buffer = origin.buffered()
+
+    override val isClosedForWrite: Boolean
+        get() = closed != null
+
+    override val closedCause: Throwable?
+        get() = closed?.cause
+
+    @InternalAPI
+    override val writeBuffer: Sink
+        get() {
+            if (isClosedForWrite) throw closedCause ?: IOException("Channel is closed for write")
+            return buffer
+        }
+
+    @OptIn(InternalAPI::class)
+    override suspend fun flush() {
+        writeBuffer.flush()
+    }
+
+    @OptIn(InternalAPI::class)
+    override suspend fun flushAndClose() {
+        writeBuffer.flush()
+        closed = CLOSED
+    }
+
+    @OptIn(InternalAPI::class)
+    override fun cancel(cause: Throwable?) {
+        val token = if (cause == null) CLOSED else CloseToken(cause)
+        closed = token
+    }
+}
+
+/**
+ * Converts this [OutputStream] into a [ByteWriteChannel], enabling asynchronous writing of byte sequences.
+ *
+ * ```kotlin
+ * val outputStream: OutputStream = FileOutputStream("file.txt")
+ * val channel: ByteWriteChannel = outputStream.asByteWriteChannel()
+ * channel.writeFully("Hello, World!".toByteArray())
+ * channel.flushAndClose() // Ensure the data is written to the OutputStream
+ * ```
+ *
+ * All operations on the [ByteWriteChannel] are buffered: the underlying [OutputStream] will be receiving bytes
+ * when the [ByteWriteChannel.flush] happens.
+ */
+public fun OutputStream.asByteWriteChannel(): ByteWriteChannel = asSink().asByteWriteChannel()

kotlin-lsp/src/com/jetbrains/ls/kotlinLsp/requests/core/shutdown.kt

@@ -9,16 +9,12 @@ import com.jetbrains.lsp.protocol.Shutdown
 import kotlinx.coroutines.CompletableDeferred
 
 context(LSServer, LSConfiguration)
-internal fun LspHandlersBuilder.shutdownRequest(clientMode: Boolean, exitSignal: CompletableDeferred<Unit>) {
+internal fun LspHandlersBuilder.shutdownRequest(exitSignal: CompletableDeferred<Unit>?) {
     request(Shutdown) {
         // TODO All the requests and notifications after this one should return InvalidRequest
         // https://microsoft.github.io/language-server-protocol/specifications/lsp/3.17/specification/#shutdown
     }
     notification(ExitNotificationType) {
-        if (clientMode) {
-            exitSignal.complete(Unit)
-        } else {
-            // do nothing, keep the server running
-        }
+        exitSignal?.complete(Unit)
     }
 }

kotlin-lsp/src/com/jetbrains/ls/kotlinLsp/stdio.kt

@@ -0,0 +1,39 @@
+// Copyright 2000-2025 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
+package com.jetbrains.ls.kotlinLsp
+
+import com.jetbrains.lsp.implementation.LspConnection
+import io.ktor.utils.io.ByteWriteChannel
+import com.jetbrains.ls.kotlinLsp.ktorHack.asByteWriteChannel
+import io.ktor.utils.io.cancel
+import io.ktor.utils.io.jvm.javaio.toByteReadChannel
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.coroutineScope
+import java.io.InputStream
+import java.io.OutputStream
+
+suspend fun stdioConnection(
+    inputStream: InputStream,
+    outputStream: OutputStream,
+    body: suspend CoroutineScope.(LspConnection) -> Unit
+) {
+    coroutineScope {
+        body(StdioConnection(inputStream, outputStream))
+    }
+}
+
+private class StdioConnection(
+    inputStream: InputStream,
+    outputStream: OutputStream,
+) : LspConnection {
+    override val input = inputStream.toByteReadChannel()
+    override val output: ByteWriteChannel = outputStream.asByteWriteChannel()
+
+    override fun close() {
+        input.cancel()
+        output.cancel(kotlinx.io.IOException("cancelled"))
+    }
+
+    override fun isAlive(): Boolean {
+        return !input.isClosedForRead || !output.isClosedForWrite
+    }
+}