bugfix: Try to make ending taks more reliable

I've seen an issue with build import not finishing ever, I am not sure how that could happen, but I am trying to make it more reliable by adding an additional field to know when task was finished and ending the task on progress in that case.
scalameta · Oct 10, 2024 · 82e4d40 · 82e4d40
1 parent 74052cc
commit 82e4d40
Show file tree

Hide file tree

Showing 3 changed files with 42 additions and 20 deletions.
diff --git a/metals/src/main/scala/scala/meta/internal/metals/ForwardingMetalsBuildClient.scala b/metals/src/main/scala/scala/meta/internal/metals/ForwardingMetalsBuildClient.scala
@@ -177,7 +177,7 @@ final class ForwardingMetalsBuildClient(
           compilations.remove(target).foreach(_.end())
 
           val name = info.getDisplayName
-          val token =
+          val (_, token) =
             workDoneProgress.startProgress(
               s"Compiling $name",
               withProgress = true,

diff --git a/metals/src/main/scala/scala/meta/internal/metals/WorkDoneProgress.scala b/metals/src/main/scala/scala/meta/internal/metals/WorkDoneProgress.scala
@@ -5,9 +5,11 @@ import java.util.concurrent.ConcurrentHashMap
 import java.util.concurrent.ScheduledExecutorService
 import java.util.concurrent.ScheduledFuture
 import java.util.concurrent.TimeUnit
+import java.util.concurrent.atomic.AtomicBoolean
 
 import scala.concurrent.ExecutionContext
 import scala.concurrent.Future
+import scala.util.control.NonFatal
 
 import scala.meta.internal.metals.MetalsEnrichments._
 import scala.meta.internal.metals.WorkDoneProgress.Token
@@ -32,6 +34,7 @@ class WorkDoneProgress(
       maybeProgress: Option[TaskProgress],
   ) {
     val timer = new Timer(time)
+    val wasFinished = new AtomicBoolean(false)
     def additionalMessage: Option[String] =
       if (showTimer) {
         val seconds = timer.elapsedSeconds
@@ -79,15 +82,15 @@ class WorkDoneProgress(
       withProgress: Boolean = false,
       showTimer: Boolean = true,
       onCancel: Option[() => Unit] = None,
-  ): Future[Token] = {
+  ): (Task, Future[Token]) = {
     val uuid = UUID.randomUUID().toString()
     val token = messages.Either.forLeft[String, Integer](uuid)
 
     val optProgress = Option.when(withProgress)(TaskProgress.empty)
     val task = Task(onCancel, showTimer, optProgress)
     taskMap.put(token, task)
 
-    client
+    val tokenFuture = client
       .createProgress(new WorkDoneProgressCreateParams(token))
       .asScala
       .map { _ =>
@@ -107,6 +110,7 @@ class WorkDoneProgress(
         client.notifyProgress(new ProgressParams(token, notification))
         token
       }
+    (task, tokenFuture)
   }
 
   def notifyProgress(
@@ -130,16 +134,22 @@ class WorkDoneProgress(
   }
 
   private def notifyProgress(token: Token, task: Task): Unit = {
-    val report = new WorkDoneProgressReport()
-    task.maybeProgress.foreach { progress =>
-      report.setPercentage(progress.percentage)
+    // make sure we don't update if a task was finished
+    if (task.wasFinished.get()) {
+      endProgress(Future.successful(token))
+      taskMap.remove(token)
+    } else {
+      val report = new WorkDoneProgressReport()
+      task.maybeProgress.foreach { progress =>
+        report.setPercentage(progress.percentage)
+      }
+      task.additionalMessage.foreach(report.setMessage)
+      val notification =
+        messages.Either.forLeft[WorkDoneProgressNotification, Object](
+          report
+        )
+      client.notifyProgress(new ProgressParams(token, notification))
     }
-    task.additionalMessage.foreach(report.setMessage)
-    val notification =
-      messages.Either.forLeft[WorkDoneProgressNotification, Object](
-        report
-      )
-    client.notifyProgress(new ProgressParams(token, notification))
   }
 
   def endProgress(token: Future[Token]): Future[Unit] =
@@ -151,8 +161,11 @@ class WorkDoneProgress(
           messages.Either.forLeft[WorkDoneProgressNotification, Object](end)
         client.notifyProgress(new ProgressParams(token, params))
       }
-      .recover { case _: NullPointerException =>
-      // no such value in the task map, task already ended or cancelled
+      .recover {
+        case _: NullPointerException =>
+        // no such value in the task map, task already ended or cancelled
+        case NonFatal(e) =>
+          scribe.error("Could not end a progress task", e)
       }
 
   def trackFuture[T](
@@ -161,16 +174,22 @@ class WorkDoneProgress(
       onCancel: Option[() => Unit] = None,
       showTimer: Boolean = true,
   )(implicit ec: ExecutionContext): Future[T] = {
-    val token =
+    val (task, token) =
       startProgress(message, onCancel = onCancel, showTimer = showTimer)
-    value.onComplete(_ => endProgress(token))
+    value.onComplete { _ =>
+      task.wasFinished.set(true)
+      endProgress(token)
+    }
     value
   }
 
   def trackBlocking[T](message: String)(thunk: => T): T = {
-    val token = startProgress(message)
+    val (task, token) = startProgress(message)
     try thunk
-    finally endProgress(token)
+    finally {
+      task.wasFinished.set(true)
+      endProgress(token)
+    }
   }
 
   def canceled(token: Token): Unit =

diff --git a/metals/src/main/scala/scala/meta/internal/worksheets/WorksheetProvider.scala b/metals/src/main/scala/scala/meta/internal/worksheets/WorksheetProvider.scala
@@ -340,12 +340,15 @@ class WorksheetProvider(
     val interruptThread = new Runnable {
       def run(): Unit = {
         if (!result.isDone()) {
-          val token = workDoneProgress.startProgress(
+          val (task, token) = workDoneProgress.startProgress(
             s"Evaluating worksheet '${path.filename}'",
             onCancel = Some(cancellable.cancel),
           )
 
-          result.asScala.onComplete(_ => workDoneProgress.endProgress(token))
+          result.asScala.onComplete { _ =>
+            task.wasFinished.set(true)
+            workDoneProgress.endProgress(token)
+          }
         }
       }
     }