@@ -14,14 +14,21 @@ import com.intellij.execution.process.ProcessOutputType
14
14
import com.intellij.openapi.Disposable
15
15
import com.intellij.openapi.components.Service
16
16
import com.intellij.openapi.components.service
17
+ import com.intellij.openapi.components.serviceIfCreated
17
18
import com.intellij.openapi.project.Project
18
19
import com.intellij.openapi.util.Disposer
19
20
import com.intellij.openapi.util.Key
20
21
import com.intellij.util.io.await
21
22
import kotlinx.coroutines.CoroutineScope
23
+ import kotlinx.coroutines.Deferred
24
+ import kotlinx.coroutines.Job
22
25
import kotlinx.coroutines.TimeoutCancellationException
26
+ import kotlinx.coroutines.async
23
27
import kotlinx.coroutines.launch
24
- import kotlinx.coroutines.time.withTimeout
28
+ import kotlinx.coroutines.runBlocking
29
+ import kotlinx.coroutines.sync.Mutex
30
+ import kotlinx.coroutines.sync.withLock
31
+ import kotlinx.coroutines.withTimeout
25
32
import org.eclipse.lsp4j.ClientCapabilities
26
33
import org.eclipse.lsp4j.ClientInfo
27
34
import org.eclipse.lsp4j.FileOperationsWorkspaceCapabilities
@@ -35,6 +42,7 @@ import org.eclipse.lsp4j.jsonrpc.Launcher
35
42
import org.eclipse.lsp4j.launch.LSPLauncher
36
43
import org.slf4j.event.Level
37
44
import software.aws.toolkits.core.utils.getLogger
45
+ import software.aws.toolkits.core.utils.info
38
46
import software.aws.toolkits.core.utils.warn
39
47
import software.aws.toolkits.jetbrains.isDeveloperMode
40
48
import software.aws.toolkits.jetbrains.services.amazonq.lsp.encryption.JwtEncryptionManager
@@ -48,8 +56,8 @@ import java.io.PrintWriter
48
56
import java.io.StringWriter
49
57
import java.net.URI
50
58
import java.nio.charset.StandardCharsets
51
- import java.time.Duration
52
59
import java.util.concurrent.Future
60
+ import kotlin.time.Duration.Companion.seconds
53
61
54
62
// https://github.com/redhat-developer/lsp4ij/blob/main/src/main/java/com/redhat/devtools/lsp4ij/server/LSPProcessListener.java
55
63
// JB impl and redhat both use a wrapper to handle input buffering issue
@@ -86,28 +94,89 @@ internal class LSPProcessListener : ProcessListener {
86
94
87
95
@Service(Service .Level .PROJECT )
88
96
class AmazonQLspService (private val project : Project , private val cs : CoroutineScope ) : Disposable {
89
- private var instance: AmazonQServerInstance ? = null
97
+ private var instance: Deferred < AmazonQServerInstance >
90
98
91
- init {
92
- cs.launch {
93
- // manage lifecycle RAII-like so we can restart at arbitrary time
94
- // and suppress IDE error if server fails to start
99
+ // dont allow lsp commands if server is restarting
100
+ private val mutex = Mutex (false )
101
+
102
+ private fun start () = cs.async {
103
+ // manage lifecycle RAII-like so we can restart at arbitrary time
104
+ // and suppress IDE error if server fails to start
105
+ var attempts = 0
106
+ while (attempts < 3 ) {
95
107
try {
96
- instance = AmazonQServerInstance (project, cs).also {
97
- Disposer .register(this @AmazonQLspService, it)
108
+ return @async withTimeout(30 .seconds) {
109
+ val instance = AmazonQServerInstance (project, cs).also {
110
+ Disposer .register(this @AmazonQLspService, it)
111
+ }
112
+ // wait for handshake to complete
113
+ instance.initializer.join()
114
+
115
+ instance
98
116
}
99
117
} catch (e: Exception ) {
100
118
LOG .warn(e) { " Failed to start LSP server" }
101
119
}
120
+ attempts++
102
121
}
122
+
123
+ error(" Failed to start LSP server in 3 attempts" )
124
+ }
125
+
126
+ init {
127
+ instance = start()
103
128
}
104
129
105
130
override fun dispose () {
106
131
}
107
132
133
+ suspend fun restart () = mutex.withLock {
134
+ // stop if running
135
+ instance.let {
136
+ if (it.isActive) {
137
+ // not even running yet
138
+ return
139
+ }
140
+
141
+ try {
142
+ val i = it.await()
143
+ if (i.initializer.isActive) {
144
+ // not initialized
145
+ return
146
+ }
147
+
148
+ Disposer .dispose(i)
149
+ } catch (e: Exception ) {
150
+ LOG .info(e) { " Exception while disposing LSP server" }
151
+ }
152
+ }
153
+
154
+ instance = start()
155
+ }
156
+
157
+ suspend fun execute (runnable : suspend (AmazonQLanguageServer ) -> Unit ) {
158
+ val lsp = withTimeout(10 .seconds) {
159
+ val holder = mutex.withLock { instance }.await()
160
+ holder.initializer.join()
161
+
162
+ holder.languageServer
163
+ }
164
+
165
+ runnable(lsp)
166
+ }
167
+
168
+ fun executeSync (runnable : suspend (AmazonQLanguageServer ) -> Unit ) {
169
+ runBlocking(cs.coroutineContext) {
170
+ execute(runnable)
171
+ }
172
+ }
173
+
108
174
companion object {
109
175
private val LOG = getLogger<AmazonQLspService >()
110
176
fun getInstance (project : Project ) = project.service<AmazonQLspService >()
177
+
178
+ fun executeIfRunning (project : Project , runnable : (AmazonQLanguageServer ) -> Unit ) =
179
+ project.serviceIfCreated<AmazonQLspService >()?.executeSync(runnable)
111
180
}
112
181
}
113
182
@@ -116,12 +185,13 @@ private class AmazonQServerInstance(private val project: Project, private val cs
116
185
117
186
private val launcher: Launcher <AmazonQLanguageServer >
118
187
119
- private val languageServer: AmazonQLanguageServer
188
+ val languageServer: AmazonQLanguageServer
120
189
get() = launcher.remoteProxy
121
190
122
191
@Suppress(" ForbiddenVoid" )
123
192
private val launcherFuture: Future <Void >
124
193
private val launcherHandler: KillableProcessHandler
194
+ val initializer: Job
125
195
126
196
private fun createClientCapabilities (): ClientCapabilities =
127
197
ClientCapabilities ().apply {
@@ -213,12 +283,12 @@ private class AmazonQServerInstance(private val project: Project, private val cs
213
283
214
284
launcherFuture = launcher.startListening()
215
285
216
- cs.launch {
286
+ initializer = cs.launch {
217
287
// encryption info must be sent within 5s or Flare process will exit
218
288
encryptionManager.writeInitializationPayload(launcherHandler.process.outputStream)
219
289
220
290
val initializeResult = try {
221
- withTimeout(Duration .ofSeconds( 10 ) ) {
291
+ withTimeout(5 .seconds ) {
222
292
languageServer.initialize(createInitializeParams()).await()
223
293
}
224
294
} catch (_: TimeoutCancellationException ) {
0 commit comments