From 3a2f4f8459359e651617e881d359a017cffe035e Mon Sep 17 00:00:00 2001
From: Khan Winter <35942988+thecoolwinter@users.noreply.github.com>
Date: Fri, 2 Aug 2024 22:05:29 -0500
Subject: [PATCH] Fix Task Renew Race Condition, Stop Button Listeners

---
 .../Notificaitons/TaskNotificationView.swift  |  5 +++-
 .../CodeEditWindowController+Toolbar.swift    | 30 +++++++++++++++----
 CodeEdit/Features/Tasks/TaskManager.swift     | 19 +++++++-----
 .../DebugUtility/TaskOutputActionsView.swift  |  4 ++-
 4 files changed, 43 insertions(+), 15 deletions(-)

diff --git a/CodeEdit/Features/ActivityViewer/Notificaitons/TaskNotificationView.swift b/CodeEdit/Features/ActivityViewer/Notificaitons/TaskNotificationView.swift
index c1d00561e..d86df90d9 100644
--- a/CodeEdit/Features/ActivityViewer/Notificaitons/TaskNotificationView.swift
+++ b/CodeEdit/Features/ActivityViewer/Notificaitons/TaskNotificationView.swift
@@ -18,7 +18,10 @@ struct TaskNotificationView: View {
                 HStack {
                     Text(notification.title)
                         .font(.subheadline)
-                        .transition(.asymmetric(insertion: .move(edge: .top), removal: .move(edge: .bottom)).combined(with: .opacity))
+                        .transition(
+                            .asymmetric(insertion: .move(edge: .top), removal: .move(edge: .bottom))
+                            .combined(with: .opacity)
+                        )
                         .id("NotificationTitle" + notification.title)
 
                     if notification.isLoading {
diff --git a/CodeEdit/Features/Documents/Controllers/CodeEditWindowController+Toolbar.swift b/CodeEdit/Features/Documents/Controllers/CodeEditWindowController+Toolbar.swift
index 1c26077d8..c489267c0 100644
--- a/CodeEdit/Features/Documents/Controllers/CodeEditWindowController+Toolbar.swift
+++ b/CodeEdit/Features/Documents/Controllers/CodeEditWindowController+Toolbar.swift
@@ -7,6 +7,7 @@
 
 import AppKit
 import SwiftUI
+import Combine
 
 extension CodeEditWindowController {
     internal func setupToolbar() {
@@ -215,11 +216,15 @@ struct StopToolbarButton: View {
     // TODO: try to get this from the environment
     @ObservedObject var taskManager: TaskManager
 
+    /// Tracks the current selected task's status. Updated by `updateStatusListener`
+    @State private var currentSelectedStatus: CETaskStatus?
+    /// The listener that listens to the active task's status publisher. Is updated frequently as the active task
+    /// changes.
+    @State private var statusListener: AnyCancellable?
+
     var body: some View {
         HStack {
-            if let selectedTaskID = taskManager.selectedTask?.id,
-               let activeTask = taskManager.activeTasks[selectedTaskID],
-               activeTask.status == .running {
+            if let currentSelectedStatus, currentSelectedStatus == .running {
                     Button {
                         taskManager.terminateActiveTask()
                     } label: {
@@ -237,10 +242,23 @@ struct StopToolbarButton: View {
         .frame(width: 38, height: 22)
         .animation(
             .easeInOut(duration: 0.3),
-            value: taskManager.activeTasks[taskManager.selectedTask?.id ?? UUID()]?.status
+            value: currentSelectedStatus
         )
-        .onChange(of: taskManager.activeTasks[taskManager.selectedTask?.id ?? UUID()]?.status) { newValue in
-            print(newValue)
+        .onChange(of: taskManager.selectedTaskID) { _ in updateStatusListener() }
+        .onChange(of: taskManager.activeTasks) { _ in updateStatusListener() }
+        .onAppear(perform: updateStatusListener)
+        .onDisappear {
+            statusListener?.cancel()
+        }
+    }
+
+    /// Update the ``statusListener`` to listen to a potentially new active task.
+    private func updateStatusListener() {
+        statusListener?.cancel()
+        currentSelectedStatus = taskManager.activeTasks[taskManager.selectedTaskID ?? UUID()]?.status
+        guard let id = taskManager.selectedTaskID else { return }
+        statusListener = taskManager.activeTasks[id]?.$status.sink { newValue in
+            currentSelectedStatus = newValue
         }
     }
 }
diff --git a/CodeEdit/Features/Tasks/TaskManager.swift b/CodeEdit/Features/Tasks/TaskManager.swift
index 37069f71c..0036c06f6 100644
--- a/CodeEdit/Features/Tasks/TaskManager.swift
+++ b/CodeEdit/Features/Tasks/TaskManager.swift
@@ -60,22 +60,27 @@ class TaskManager: ObservableObject {
     func executeActiveTask() {
         let task = workspaceSettings.tasks.first { $0.id == selectedTaskID }
         guard let task else { return }
-        runTask(task: task)
+        Task {
+            await runTask(task: task)
+        }
     }
 
-    func runTask(task: CETask) {
+    func runTask(task: CETask) async {
         // A process can only be started once, that means we have to renew the Process and Pipe
-        // but don't initialise a new object.
+        // but don't initialize a new object.
         if let activeTask = activeTasks[task.id] {
             activeTask.renew()
+            // Wait until the task is no longer running.
+            // The termination handler is asynchronous, so we avoid a race condition using this.
+            while activeTask.status != .notRunning {
+                await Task.yield()
+            }
             activeTask.run()
         } else {
             let runningTask = CEActiveTask(task: task)
             runningTask.run()
-            Task {
-                await MainActor.run {
-                    activeTasks[task.id] = runningTask
-                }
+            await MainActor.run {
+                activeTasks[task.id] = runningTask
             }
         }
     }
diff --git a/CodeEdit/Features/UtilityArea/DebugUtility/TaskOutputActionsView.swift b/CodeEdit/Features/UtilityArea/DebugUtility/TaskOutputActionsView.swift
index 1e7b4df49..2dbf78d5b 100644
--- a/CodeEdit/Features/UtilityArea/DebugUtility/TaskOutputActionsView.swift
+++ b/CodeEdit/Features/UtilityArea/DebugUtility/TaskOutputActionsView.swift
@@ -18,7 +18,9 @@ struct TaskOutputActionsView: View {
             Spacer()
 
             Button {
-                taskManager.runTask(task: activeTask.task)
+                Task {
+                    await taskManager.runTask(task: activeTask.task)
+                }
             } label: {
                 Image(systemName: "memories")
                     .foregroundStyle(.green)