fix: handle new downloads and ls restarts correctly

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/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()
     }
 }