Skip to content

Commit 6c5f2f7

Browse files
S-furiS-furi
committed
feat: use WebFlux-based web server for lsp-based completions
This commit introduces a new dedicated module for completions. This module leverages a `WebFlux` web-server for its springboot application in order to exploit reactive patterns both in the LSP RESTful approach and WebSocket. This helps in preventing clashes with default compiler application which uses Servlet and forced us to use an old implementation of WebSockets (`TextWebSocketHandler`).
1 parent 4adcc02 commit 6c5f2f7

File tree

55 files changed

+782
-621
lines changed

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

55 files changed

+782
-621
lines changed

‎completions/build.gradle.kts‎

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,29 @@
1+
2+
plugins {
3+
id("base-kotlin-jvm-conventions")
4+
alias(libs.plugins.spring.dependency.management)
5+
alias(libs.plugins.spring.boot)
6+
alias(libs.plugins.kotlin.plugin.spring)
7+
}
8+
9+
dependencies {
10+
implementation("org.springframework.boot:spring-boot-starter-webflux")
11+
implementation(libs.org.eclipse.lsp4j)
12+
implementation(libs.kotlinx.coroutines.reactor)
13+
implementation(libs.kotlinx.serialization.core.jvm)
14+
implementation(libs.kotlinx.serialization.json.jvm)
15+
implementation(project(":executors", configuration = "default"))
16+
implementation(project(":common", configuration = "default"))
17+
implementation(project(":dependencies"))
18+
19+
testImplementation(libs.kotlin.test)
20+
testImplementation(libs.bundles.testcontainers)
21+
testImplementation(libs.rector.test)
22+
testImplementation("org.springframework.boot:spring-boot-starter-test") {
23+
exclude(group = "org.junit.vintage", module = "junit-vintage-engine")
24+
}
25+
}
26+
27+
tasks.test {
28+
useJUnitPlatform()
29+
}

‎completions/src/main/kotlin/completions/CompletionsApplication.kt‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,11 @@
1+
package completions
2+
3+
import org.springframework.boot.autoconfigure.SpringBootApplication
4+
import org.springframework.boot.runApplication
5+
6+
@SpringBootApplication
7+
class CompletionsApplication
8+
9+
fun main(args: Array<String>) {
10+
runApplication<CompletionsApplication>(*args)
11+
}

‎completions/src/main/kotlin/completions/configuration/WebSocketConfig.kt‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,18 @@
1+
package completions.configuration
2+
3+
import completions.controllers.ws.LspCompletionWebSocketHandler
4+
import org.springframework.context.annotation.Bean
5+
import org.springframework.context.annotation.Configuration
6+
import org.springframework.web.reactive.HandlerMapping
7+
import org.springframework.web.reactive.handler.SimpleUrlHandlerMapping
8+
import org.springframework.web.reactive.socket.server.support.WebSocketHandlerAdapter
9+
10+
@Configuration
11+
class WebSocketConfig {
12+
@Bean
13+
fun webSocketHandlerAdapter(): WebSocketHandlerAdapter = WebSocketHandlerAdapter()
14+
15+
@Bean
16+
fun webSocketMapping(handler: LspCompletionWebSocketHandler): HandlerMapping =
17+
SimpleUrlHandlerMapping(mapOf("/completions/lsp/complete" to handler), 1)
18+
}

‎completions/src/main/kotlin/completions/controllers/CompletionsRestController.kt‎

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,23 @@
1+
package completions.controllers
2+
3+
import completions.lsp.components.LspCompletionProvider
4+
import completions.model.Project
5+
import model.Completion
6+
import org.springframework.web.bind.annotation.PostMapping
7+
import org.springframework.web.bind.annotation.RequestBody
8+
import org.springframework.web.bind.annotation.RequestMapping
9+
import org.springframework.web.bind.annotation.RequestParam
10+
import org.springframework.web.bind.annotation.RestController
11+
12+
@RestController
13+
@RequestMapping(value = ["/api/complete"])
14+
class CompletionsRestController(
15+
private val lspCompletionProvider: LspCompletionProvider
16+
) {
17+
@PostMapping("/lsp")
18+
suspend fun complete(
19+
@RequestBody project: Project,
20+
@RequestParam line: Int,
21+
@RequestParam ch: Int,
22+
): List<Completion> = lspCompletionProvider.complete(project, line, ch)
23+
}

‎completions/src/main/kotlin/completions/controllers/ws/LspCompletionWebSocketHandler.kt‎

Lines changed: 155 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,155 @@
1+
@file:Suppress("ReactiveStreamsUnusedPublisher")
2+
3+
package completions.controllers.ws
4+
5+
import completions.model.Project
6+
import com.fasterxml.jackson.databind.ObjectMapper
7+
import com.fasterxml.jackson.module.kotlin.readValue
8+
import com.fasterxml.jackson.module.kotlin.registerKotlinModule
9+
import completions.lsp.KotlinLspProxy
10+
import completions.lsp.StatefulKotlinLspProxy.onClientConnected
11+
import completions.lsp.StatefulKotlinLspProxy.onClientDisconnected
12+
import completions.lsp.components.LspCompletionProvider
13+
import kotlinx.coroutines.CoroutineName
14+
import kotlinx.coroutines.CoroutineScope
15+
import kotlinx.coroutines.DelicateCoroutinesApi
16+
import kotlinx.coroutines.Dispatchers
17+
import kotlinx.coroutines.Job
18+
import kotlinx.coroutines.SupervisorJob
19+
import kotlinx.coroutines.cancelChildren
20+
import kotlinx.coroutines.channels.Channel
21+
import kotlinx.coroutines.launch
22+
import model.Completion
23+
import org.slf4j.LoggerFactory
24+
import org.springframework.stereotype.Component
25+
import org.springframework.web.reactive.socket.WebSocketHandler
26+
import org.springframework.web.reactive.socket.WebSocketMessage
27+
import org.springframework.web.reactive.socket.WebSocketSession
28+
import reactor.core.publisher.Flux
29+
import reactor.core.publisher.Mono
30+
import reactor.core.publisher.Sinks
31+
import kotlin.coroutines.cancellation.CancellationException
32+
33+
@Component
34+
class LspCompletionWebSocketHandler(
35+
private val lspProxy: KotlinLspProxy,
36+
private val lspCompletionProvider: LspCompletionProvider,
37+
) : WebSocketHandler {
38+
private val job = SupervisorJob()
39+
40+
private val logger = LoggerFactory.getLogger(LspCompletionWebSocketHandler::class.java)
41+
42+
override fun handle(session: WebSocketSession): Mono<Void?> {
43+
val sessionId = session.id
44+
45+
val sink = Sinks.many().unicast().onBackpressureBuffer<WebSocketMessage>()
46+
val outbound: Flux<WebSocketMessage> = sink.asFlux()
47+
48+
val sessionJob = SupervisorJob(job)
49+
val sessionScope = CoroutineScope(Dispatchers.IO + sessionJob + CoroutineName("LspWS-$sessionId"))
50+
51+
sessionScope.launch {
52+
runCatching { lspProxy.requireAvailable() }
53+
.onSuccess {
54+
lspProxy.onClientConnected(sessionId)
55+
sink.emit(session, WSResponse.Init(sessionId))
56+
}
57+
}
58+
59+
val requestChannel = Channel<WsCompletionRequest>(Channel.UNLIMITED)
60+
61+
val inbound = session.receive()
62+
.map { it.payloadAsText }
63+
.doOnNext { payload ->
64+
runCatching { objectMapper.readValue<WsCompletionRequest>(payload) }
65+
.onFailure { sink.emit(session, WSResponse.Error("Failed to parse request: $payload")) }
66+
.onSuccess { requestChannel.trySend(it) }
67+
}.then()
68+
69+
setupCompletionWorker(session, sessionScope, requestChannel, sink)
70+
71+
val sendMono = session.send(outbound).doFinally { cleanup(session, sessionJob, sink) }
72+
return Mono.`when`(inbound, sendMono)
73+
.doOnError { logger.warn("WS session error for client $sessionId: ${it.message}") }
74+
.doFinally { cleanup(session, sessionJob, sink) }
75+
}
76+
77+
@OptIn(DelicateCoroutinesApi::class)
78+
private fun setupCompletionWorker(
79+
session: WebSocketSession,
80+
sessionScope: CoroutineScope,
81+
requestChannel: Channel<WsCompletionRequest>,
82+
sink: Sinks.Many<WebSocketMessage>,
83+
) = sessionScope.launch {
84+
try {
85+
while (!requestChannel.isClosedForReceive) {
86+
val first = requestChannel.receiveCatching().getOrNull() ?: break
87+
var req = first
88+
89+
while (true) {
90+
val next = requestChannel.tryReceive().getOrNull() ?: break
91+
sink.emit(session, WSResponse.Discarded(req.requestId))
92+
req = next
93+
}
94+
95+
val available = runCatching { lspProxy.requireAvailable() }.isSuccess
96+
if (!available) {
97+
sink.emit(session, WSResponse.Error("Lsp client is not available"))
98+
continue
99+
}
100+
101+
try {
102+
val completions = lspCompletionProvider.complete(
103+
clientId = session.id,
104+
project = req.project,
105+
line = req.line,
106+
ch = req.ch,
107+
applyFuzzyRanking = true,
108+
)
109+
sink.emit(session, WSResponse.Completions(completions, req.requestId))
110+
} catch (e: Exception) {
111+
logger.warn("Completion processing failed for client ${session.id}:", e)
112+
sink.emit(session, WSResponse.Error(e.message ?: "Unknown error", req.requestId))
113+
}
114+
}
115+
} catch (t: Throwable) {
116+
if (t !is CancellationException) {
117+
logger.warn("Error collecting responses for client ${session.id}: ${t.message}")
118+
sink.emit(session, WSResponse.Error(t.message ?: "Error retrieving completions"))
119+
}
120+
} finally {
121+
requestChannel.close()
122+
sessionScope.coroutineContext.cancelChildren()
123+
}
124+
}
125+
126+
private fun cleanup(session: WebSocketSession, sessionJob: Job, sink: Sinks.Many<WebSocketMessage>) {
127+
lspProxy.onClientDisconnected(session.id)
128+
sessionJob.cancel()
129+
sink.tryEmitComplete()
130+
session.close()
131+
}
132+
}
133+
134+
private fun Sinks.Many<WebSocketMessage>.emit(session: WebSocketSession, response: WSResponse) =
135+
tryEmitNext(session.textMessage(response.toJson()))
136+
137+
sealed interface WSResponse {
138+
val requestId: String?
139+
140+
open class Error(val message: String, override val requestId: String? = null) : WSResponse
141+
data class Init(val sessionId: String, override val requestId: String? = null) : WSResponse
142+
data class Completions(val completions: List<Completion>, override val requestId: String? = null) : WSResponse
143+
data class Discarded(override val requestId: String) : Error("discarded", requestId)
144+
145+
fun toJson(): String = objectMapper.writeValueAsString(this)
146+
}
147+
148+
private data class WsCompletionRequest(
149+
val requestId: String,
150+
val project: Project,
151+
val line: Int,
152+
val ch: Int,
153+
)
154+
155+
private val objectMapper = ObjectMapper().apply { registerKotlinModule() }

‎src/main/kotlin/com/compiler/server/service/lsp/KotlinLspProxy.kt‎ renamed to ‎completions/src/main/kotlin/completions/lsp/KotlinLspProxy.kt‎

Lines changed: 19 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,10 @@
1-
package com.compiler.server.service.lsp
1+
package completions.lsp
22

3-
import com.compiler.server.service.lsp.components.LspProject
4-
import com.compiler.server.model.Project
5-
import com.compiler.server.model.ProjectFile
6-
import com.compiler.server.service.lsp.client.LspClient
7-
import com.compiler.server.service.lsp.client.ReconnectingLspClient
3+
import completions.lsp.client.LspClient
4+
import completions.lsp.client.ReconnectingLspClient
5+
import completions.lsp.components.LspProject
6+
import completions.model.Project
7+
import completions.model.ProjectFile
88
import kotlinx.coroutines.CompletableDeferred
99
import kotlinx.coroutines.CoroutineName
1010
import kotlinx.coroutines.CoroutineScope
@@ -94,7 +94,7 @@ class KotlinLspProxy {
9494
* Initialize the LSP client. This method must be called before any other method in this
9595
* class. It is recommended to call this method when the **spring** application context is initialized.
9696
*
97-
* [workspacePath] is the path ([[java.net.URI.path]]) to the root project directory,
97+
* [workspacePath] is the path ([[URI.path]]) to the root project directory,
9898
* where the project must be a project supported by [Kotlin-LSP](https://github.com/Kotlin/kotlin-lsp).
9999
* The workspace will not contain users' files, but it can be used to store common files,
100100
* to specify kotlin/java versions, project-wide imported libraries and so on.
@@ -103,7 +103,7 @@ class KotlinLspProxy {
103103
* @param clientName the name of the client, defaults to "lsp-proxy"
104104
*/
105105
suspend fun initializeClient(
106-
workspacePath: String = LSP_REMOTE_WORKSPACE_ROOT.path,
106+
workspacePath: String = lspRemoteWorkspaceRoot().path,
107107
clientName: String = "kotlin-compiler-server"
108108
) {
109109
if (!::lspClient.isInitialized) {
@@ -174,24 +174,24 @@ class KotlinLspProxy {
174174
}
175175

176176
/**
177-
* [LSP_REMOTE_WORKSPACE_ROOT] is the workspace that the LSP will point to, while
178-
* [LSP_LOCAL_WORKSPACE_ROOT] is the local workspace that the LSP client is running on.
177+
* [lspRemoteWorkspaceRoot] is the workspace that the LSP will point to, while
178+
* [lspLocalWorkspaceRoot] is the local workspace that the LSP client is running on.
179179
* They are usually the same if the LSP client is running on the same machine as the server,
180-
* otherwise [LSP_REMOTE_WORKSPACE_ROOT] will have to be set wrt to server's local workspace.
180+
* otherwise [lspRemoteWorkspaceRoot] will have to be set wrt to server's local workspace.
181181
*
182-
* Note that [LSP_REMOTE_WORKSPACE_ROOT] is the most important one, since it will be used
182+
* Note that [lspRemoteWorkspaceRoot] is the most important one, since it will be used
183183
* from the LSP analyzer to resolve the project's dependencies.
184184
*/
185185
companion object {
186186
private val logger = LoggerFactory.getLogger(KotlinLspProxy::class.java)
187187

188-
val LSP_REMOTE_WORKSPACE_ROOT: URI = Path.of(
188+
fun lspRemoteWorkspaceRoot(): URI =Path.of(
189189
System.getProperty("LSP_REMOTE_WORKSPACE_ROOT")
190190
?: System.getenv("LSP_REMOTE_WORKSPACE_ROOT")
191191
?: defaultWorkspacePath()
192192
).toUri()
193193

194-
val LSP_LOCAL_WORKSPACE_ROOT: URI = Path.of(
194+
fun lspLocalWorkspaceRoot(): URI = Path.of(
195195
System.getProperty("LSP_LOCAL_WORKSPACE_ROOT")
196196
?: System.getenv("LSP_LOCAL_WORKSPACE_ROOT")
197197
?: defaultWorkspacePath()
@@ -200,7 +200,11 @@ class KotlinLspProxy {
200200
private fun defaultWorkspacePath(): String =
201201
System.getProperty("LSP_USERS_PROJECTS_ROOT")
202202
?: System.getProperty("LSP_USERS_PROJECTS_ROOT")
203-
?: ("lsp-users-projects-root")
203+
?: run {
204+
logger.warn("LSP_USERS_PROJECTS_ROOT system property not set, using default")
205+
KotlinLspProxy::class.java.getResource("/lsp-users-projects-root")?.path
206+
?: error("Could not find default workspace path")
207+
}
204208
}
205209
}
206210

‎src/main/kotlin/com/compiler/server/service/lsp/LspAvailabilityAdvice.kt‎ renamed to ‎completions/src/main/kotlin/completions/lsp/LspAvailabilityAdvice.kt‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
package com.compiler.server.service.lsp
1+
package completions.lsp
22

33
import org.springframework.http.HttpStatus
44
import org.springframework.http.ResponseEntity

‎src/main/kotlin/com/compiler/server/service/lsp/LspCompletionParser.kt‎ renamed to ‎completions/src/main/kotlin/completions/lsp/LspCompletionParser.kt‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
package com.compiler.server.service.lsp
1+
package completions.lsp
22

33
import org.eclipse.lsp4j.CompletionItem
44
import model.Completion

‎src/main/kotlin/com/compiler/server/service/lsp/LspUnavailableException.kt‎ renamed to ‎completions/src/main/kotlin/completions/lsp/LspUnavailableException.kt‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
package com.compiler.server.service.lsp
1+
package completions.lsp
22

33
import java.lang.RuntimeException
44

‎src/main/kotlin/com/compiler/server/service/lsp/client/KotlinLanguageClient.kt‎ renamed to ‎completions/src/main/kotlin/completions/lsp/client/KotlinLanguageClient.kt‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
package com.compiler.server.service.lsp.client
1+
package completions.lsp.client
22

33
import org.eclipse.lsp4j.MessageActionItem
44
import org.eclipse.lsp4j.MessageParams

0 commit comments

Comments
 (0)