fix: improve download and restart logic

src/main/kotlin/io/snyk/plugin/Utils.kt

@@ -51,12 +51,12 @@ import snyk.common.ProductType
 import snyk.common.SnykCachedResults
 import snyk.common.UIComponentFinder
 import snyk.common.isSnykTenant
+import snyk.common.lsp.LanguageServerRestartListener
 import snyk.common.lsp.ScanInProgressKey
 import snyk.common.lsp.ScanIssue
 import snyk.common.lsp.ScanState
 import snyk.container.ContainerService
 import snyk.container.KubernetesImageCache
-import snyk.errorHandler.SentryErrorReporter
 import snyk.iac.IacScanService
 import java.io.File
 import java.io.FileNotFoundException
@@ -74,6 +74,7 @@ fun getIacService(project: Project): IacScanService? = project.serviceIfNotDispo
 fun getKubernetesImageCache(project: Project): KubernetesImageCache? = project.serviceIfNotDisposed()
 
 fun getSnykTaskQueueService(project: Project): SnykTaskQueueService? = project.serviceIfNotDisposed()
+fun getLanguageServerRestartListener(): LanguageServerRestartListener = ApplicationManager.getApplication().service()
 
 fun getSnykToolWindowPanel(project: Project): SnykToolWindowPanel? = project.serviceIfNotDisposed()
 
@@ -95,15 +96,15 @@ fun getContainerService(project: Project): ContainerService? = project.serviceIf
 
 fun getSnykCliAuthenticationService(project: Project?): SnykCliAuthenticationService? = project?.serviceIfNotDisposed()
 
-fun getSnykCliDownloaderService(): SnykCliDownloaderService = getApplicationService()
+fun getSnykCliDownloaderService(): SnykCliDownloaderService = ApplicationManager.getApplication().service()
 
 fun getSnykProjectSettingsService(project: Project): SnykProjectSettingsStateService? = project.serviceIfNotDisposed()
 
 fun getCliFile() = File(pluginSettings().cliPath)
 
 fun isCliInstalled(): Boolean = ApplicationManager.getApplication().isUnitTestMode || getCliFile().exists()
 
-fun pluginSettings(): SnykApplicationSettingsStateService = getApplicationService()
+fun pluginSettings(): SnykApplicationSettingsStateService = ApplicationManager.getApplication().service()
 
 fun getPluginPath() = PathManager.getPluginsPath() + "/snyk-intellij-plugin"
 
@@ -119,20 +120,10 @@ private inline fun <reified T : Any> Project.serviceIfNotDisposed(): T? {
     return try {
         getService(T::class.java)
     } catch (t: Throwable) {
-        SentryErrorReporter.captureException(t)
         null
     }
 }
 
-/**
- * Copy of [com.intellij.openapi.components.service] to make code compilable with jvm 11 bytecode (Idea 2022.1)
- */
-private inline fun <reified T : Any> getApplicationService(): T {
-    val serviceClass = T::class.java
-    return ApplicationManager.getApplication()?.getService(serviceClass)
-        ?: throw RuntimeException("Cannot find service ${serviceClass.name} (classloader=${serviceClass.classLoader})")
-}
-
 fun <L : Any> getSyncPublisher(project: Project, topic: Topic<L>): L? {
     val messageBus = project.messageBus
     if (messageBus.isDisposed) return null

src/main/kotlin/io/snyk/plugin/services/SnykTaskQueueService.kt

@@ -157,6 +157,7 @@ class SnykTaskQueueService(val project: Project) {
         })
     }
 
+    // FIXME this is currently not project, but app specific
     fun stopScan() {
         val languageServerWrapper = LanguageServerWrapper.getInstance()
 

src/main/kotlin/io/snyk/plugin/services/download/CliDownloader.kt

@@ -1,9 +1,12 @@
 package io.snyk.plugin.services.download
 
+import com.intellij.ide.impl.ProjectUtil
 import com.intellij.openapi.progress.ProgressIndicator
 import io.snyk.plugin.cli.Platform
 import io.snyk.plugin.pluginSettings
 import io.snyk.plugin.services.download.HttpRequestHelper.createRequest
+import io.snyk.plugin.ui.SnykBalloonNotificationHelper
+import snyk.common.lsp.LanguageServerWrapper
 import java.io.File
 import java.io.IOException
 import java.nio.file.AtomicMoveNotSupportedException
@@ -62,7 +65,12 @@ class CliDownloader {
             if (cliFile.exists()) {
                 cliFile.delete()
             }
+            val languageServerWrapper = LanguageServerWrapper.getInstance()
+            // prevent spawning of language server until files are moved
+            languageServerWrapper.isInitializing.lock()
             try {
+                // shutdown, so the binary can be updated
+                languageServerWrapper.shutdown()
                 Files.move(
                     downloadFile.toPath(),
                     cliFile.toPath(),
@@ -71,7 +79,14 @@ class CliDownloader {
                 )
             } catch (ignored: AtomicMoveNotSupportedException) {
                 // fallback to renameTo because of e
-                downloadFile.renameTo(cliFile)
+                val success = downloadFile.renameTo(cliFile)
+                if (!success) {
+                    val message =
+                        "CLI could not be updated. Please check if another process is using the CLI binary at ${pluginSettings().cliPath}"
+                    SnykBalloonNotificationHelper.showWarn(message, ProjectUtil.getActiveProject())
+                }
+            } finally {
+                languageServerWrapper.isInitializing.unlock()
             }
             cliFile.setExecutable(true)
             return cliFile

src/main/kotlin/io/snyk/plugin/services/download/CliDownloaderService.kt

@@ -15,7 +15,6 @@ import java.io.IOException
 import java.time.LocalDate
 import java.time.temporal.ChronoUnit
 import java.util.Date
-import java.util.concurrent.TimeUnit
 
 @Service
 class SnykCliDownloaderService {
@@ -71,14 +70,6 @@ class SnykCliDownloaderService {
 
             val languageServerWrapper = LanguageServerWrapper.getInstance()
             try {
-                if (languageServerWrapper.isInitialized) {
-                    try {
-                        languageServerWrapper.shutdown()
-                    } catch (e: RuntimeException) {
-                        logger<SnykCliDownloaderService>()
-                            .warn("Language server shutdown for download took too long, couldn't shutdown", e)
-                    }
-                }
                 downloader.downloadFile(cliFile, latestRelease, indicator)
                 pluginSettings().cliVersion = latestRelease
                 pluginSettings().lastCheckDate = Date()

src/main/kotlin/snyk/common/lsp/LanguageServerRestartListener.kt

@@ -0,0 +1,36 @@
+package snyk.common.lsp
+
+import com.intellij.openapi.Disposable
+import com.intellij.openapi.application.ApplicationManager
+import com.intellij.openapi.components.Service
+import com.intellij.openapi.components.service
+import com.intellij.openapi.util.Disposer
+import io.snyk.plugin.events.SnykCliDownloadListener
+import io.snyk.plugin.ui.toolwindow.SnykPluginDisposable
+
+@Service(Service.Level.APP)
+class LanguageServerRestartListener : Disposable {
+    private var disposed = false
+
+    fun isDisposed() = disposed
+    override fun dispose() {
+        this.disposed = true
+    }
+
+    companion object {
+        @JvmStatic
+        fun getInstance(): LanguageServerRestartListener = service()
+    }
+
+    init {
+        Disposer.register(SnykPluginDisposable.getInstance(), this)
+        ApplicationManager.getApplication().messageBus.connect()
+            .subscribe(SnykCliDownloadListener.CLI_DOWNLOAD_TOPIC, object : SnykCliDownloadListener {
+                override fun cliDownloadFinished(succeed: Boolean) {
+                    if (succeed && !disposed) {
+                        LanguageServerWrapper.getInstance().restart()
+                    }
+                }
+            })
+    }
+}

src/main/kotlin/snyk/common/lsp/LanguageServerWrapper.kt

@@ -1,6 +1,7 @@
 package snyk.common.lsp
 
 import com.google.gson.Gson
+import com.intellij.ide.impl.ProjectUtil
 import com.intellij.openapi.Disposable
 import com.intellij.openapi.application.ApplicationManager
 import com.intellij.openapi.components.service
@@ -60,6 +61,7 @@ import snyk.common.lsp.commands.COMMAND_LOGOUT
 import snyk.common.lsp.commands.COMMAND_REPORT_ANALYTICS
 import snyk.common.lsp.commands.COMMAND_WORKSPACE_FOLDER_SCAN
 import snyk.common.lsp.commands.ScanDoneEvent
+import snyk.common.lsp.progress.ProgressManager
 import snyk.common.lsp.settings.LanguageServerSettings
 import snyk.common.lsp.settings.SeverityFilter
 import snyk.pluginInfo
@@ -163,6 +165,8 @@ class LanguageServerWrapper(
             if (!listenerFuture.isDone) {
                 sendInitializeMessage()
                 isInitialized = true
+                // listen for downloads / restarts
+                LanguageServerRestartListener.getInstance()
             } else {
                 logger.warn("Language Server initialization did not succeed")
             }
@@ -183,7 +187,17 @@ class LanguageServerWrapper(
         messageProducerLogger.level = Level.OFF
         try {
             val shouldShutdown = lsIsAlive()
-            executorService.submit { if (shouldShutdown) languageServer.shutdown().get(1, TimeUnit.SECONDS) }
+            executorService.submit {
+                if (shouldShutdown) {
+                    val project = ProjectUtil.getActiveProject()
+                    if (project != null) {
+                        getSnykTaskQueueService(project)?.stopScan()
+                    }
+                    // cancel all progresses, as we can have more progresses than just scans
+                    ProgressManager.getInstance().cancelProgresses()
+                    languageServer.shutdown().get(1, TimeUnit.SECONDS)
+                }
+            }
         } catch (ignored: Exception) {
             // we don't care
         } finally {
@@ -405,11 +419,12 @@ class LanguageServerWrapper(
         }
     }
 
-    private fun restart() {
+    fun restart() {
         runInBackground("Snyk: restarting language server...") {
             shutdown()
             Thread.sleep(1000)
             ensureLanguageServerInitialized()
+            ProjectManager.getInstance().openProjects.forEach { project -> addContentRoots(project) }
         }
     }
 

src/main/kotlin/snyk/common/lsp/SnykLanguageClient.kt

@@ -59,7 +59,7 @@ class SnykLanguageClient :
     Disposable {
     val logger = Logger.getInstance("Snyk Language Server")
     val gson = Gson()
-    val progressManager = ProgressManager()
+    val progressManager = ProgressManager.getInstance()
 
     private var disposed = false
         get() {

src/main/kotlin/snyk/common/lsp/progress/ProgressManager.kt

@@ -18,6 +18,8 @@ package snyk.common.lsp.progress
 
 import com.intellij.ide.impl.ProjectUtil
 import com.intellij.openapi.Disposable
+import com.intellij.openapi.components.Service
+import com.intellij.openapi.components.service
 import com.intellij.openapi.progress.ProcessCanceledException
 import com.intellij.openapi.progress.ProgressIndicator
 import com.intellij.openapi.progress.ProgressManager
@@ -35,7 +37,7 @@ import java.util.concurrent.CompletableFuture
 import java.util.concurrent.ConcurrentHashMap
 import java.util.function.Function
 
-
+@Service
 class ProgressManager : Disposable {
     private val progresses: MutableMap<String, Progress> = ConcurrentHashMap<String, Progress>()
     private var disposed = false
@@ -231,5 +233,8 @@ class ProgressManager : Disposable {
                 Function.identity()
             ) { obj: Int -> obj.toString() }
         }
+
+        @JvmStatic
+        fun getInstance(): snyk.common.lsp.progress.ProgressManager = service()
     }
 }

src/test/kotlin/snyk/common/lsp/SnykLanguageClientTest.kt

@@ -27,6 +27,7 @@ import org.junit.Assert.assertEquals
 import org.junit.Assert.assertNotEquals
 import org.junit.Before
 import org.junit.Test
+import snyk.common.lsp.progress.ProgressManager
 import snyk.pluginInfo
 import snyk.trust.WorkspaceTrustService
 import java.nio.file.Files
@@ -54,6 +55,7 @@ class SnykLanguageClientTest {
         every { applicationMock.getService(WorkspaceTrustService::class.java) } returns trustServiceMock
 
         every { applicationMock.getService(ProjectManager::class.java) } returns projectManagerMock
+        every { applicationMock.getService(ProgressManager::class.java) } returns mockk(relaxed = true)
         every { applicationMock.isDisposed } returns false
         every { applicationMock.messageBus } returns mockk(relaxed = true)
 
@@ -84,8 +86,7 @@ class SnykLanguageClientTest {
     }
 
     @Test
-    fun applyEdit() {
-    }
+    fun applyEdit() {}
 
     @Test
     fun `refreshCodeLenses does not run when disposed`() {