From 5f919d88e1c7e8571819c12d08c559fa2a31a736 Mon Sep 17 00:00:00 2001
From: Ilya Kirillov <ilya.kirillov@jetbrains.com>
Date: Tue, 10 Jun 2025 11:21:59 +0200
Subject: [PATCH] [lsp] refactoring, move `tcpClient`/`tcpConnection` near
 `tcpServer` so the similar code is closer

LSP-166

GitOrigin-RevId: 10692c5c5e651c9ce4a8d5cadca3af1e5c228606
---
 .../com/jetbrains/lsp/implementation/tcp.kt   | 39 +++++++++++++++----
 1 file changed, 31 insertions(+), 8 deletions(-)

diff --git a/fleet/lsp.protocol/src/main/com/jetbrains/lsp/implementation/tcp.kt b/fleet/lsp.protocol/src/main/com/jetbrains/lsp/implementation/tcp.kt
index 2c212293ac3b..096a02a9abbf 100644
--- a/fleet/lsp.protocol/src/main/com/jetbrains/lsp/implementation/tcp.kt
+++ b/fleet/lsp.protocol/src/main/com/jetbrains/lsp/implementation/tcp.kt
@@ -1,17 +1,11 @@
 package com.jetbrains.lsp.implementation
 
 import fleet.util.logging.logger
-import kotlinx.coroutines.CoroutineScope
-import kotlinx.coroutines.CoroutineStart
-import kotlinx.coroutines.Dispatchers
-import kotlinx.coroutines.coroutineScope
-import kotlinx.coroutines.launch
-import kotlinx.coroutines.runInterruptible
-import kotlinx.coroutines.supervisorScope
+import kotlinx.coroutines.*
 import java.io.InputStream
 import java.io.OutputStream
 import java.net.ServerSocket
-import kotlin.io.use
+import java.net.Socket
 
 suspend fun tcpServer(port: Int = 0, server: suspend CoroutineScope.(InputStream, OutputStream) -> Unit) {
   ServerSocket(port).use { serverSocket ->
@@ -37,4 +31,33 @@ suspend fun tcpServer(port: Int = 0, server: suspend CoroutineScope.(InputStream
   }
 }
 
+
+/**
+ * VSC opens a **server** socket for LSP to connect to it.
+ */
+suspend fun tcpClient(port: Int, body: suspend CoroutineScope.(InputStream, OutputStream) -> Unit) {
+  val socket = runInterruptible(Dispatchers.IO) {
+    Socket("localhost", port)
+  }
+  socket.use {
+    coroutineScope {
+      body(socket.getInputStream(), socket.getOutputStream())
+    }
+  }
+}
+
+
+suspend fun tcpConnection(clientMode: Boolean, port: Int, body: suspend CoroutineScope.(InputStream, OutputStream) -> Unit) {
+  when {
+    clientMode -> {
+      tcpClient(port, body)
+    }
+
+    else -> {
+      tcpServer(port, body)
+    }
+  }
+}
+
+
 private val LOG = logger<LspClient>()