Gently stops fleet of virtual machines (#50)

* Gently stops fleet of virtual machines * Adds info text when stopping * Removes activeMachineNames
shapehq · Sep 22, 2023 · ad92a4d · ad92a4d
1 parent 7a3c234
commit ad92a4d
Show file tree

Hide file tree

Showing 7 changed files with 93 additions and 75 deletions.
diff --git a/Packages/MenuBar/Sources/MenuBarItem/Internal/FleetMenuBarItem.swift b/Packages/MenuBar/Sources/MenuBarItem/Internal/FleetMenuBarItem.swift
@@ -1,42 +1,49 @@
 import SwiftUI
 
 struct FleetMenuBarItem: View {
+    enum Action {
+        case start
+        case stop
+    }
+
     let hasSelectedVirtualMachine: Bool
     let isFleetStarted: Bool
+    let isStoppingFleet: Bool
     let isEditorStarted: Bool
-    let startsSingleVirtualMachine: Bool
-    let onSelect: () -> Void
+    let onSelect: (Action) -> Void
 
     var body: some View {
-        if isFleetStarted {
+        if isStoppingFleet {
+            Button {} label: {
+                HStack {
+                    Image(systemName: "stop.fill")
+                    Text(L10n.MenuBarItem.VirtualMachines.stopping)
+                }
+            }.disabled(true)
+            Button {} label: {
+                Text(L10n.MenuBarItem.VirtualMachines.stoppingInfo)
+            }.disabled(true)
+        } else if isFleetStarted {
             Button {
-                onSelect()
+                onSelect(.stop)
             } label: {
                 HStack {
                     Image(systemName: "stop.fill")
-                    if startsSingleVirtualMachine {
-                        Text(L10n.MenuBarItem.VirtualMachines.Stop.singularis)
-                    } else {
-                        Text(L10n.MenuBarItem.VirtualMachines.Stop.pluralis)
-                    }
+                    Text(L10n.MenuBarItem.VirtualMachines.stop)
                 }
             }
         } else if hasSelectedVirtualMachine {
             Button {
-                onSelect()
+                onSelect(.start)
             } label: {
                 HStack {
                     Image(systemName: "play.fill")
-                    if startsSingleVirtualMachine {
-                        Text(L10n.MenuBarItem.VirtualMachines.Start.singularis)
-                    } else {
-                        Text(L10n.MenuBarItem.VirtualMachines.Start.pluralis)
-                    }
+                    Text(L10n.MenuBarItem.VirtualMachines.start)
                 }
             }.disabled(isEditorStarted)
         } else {
             Button {
-                onSelect()
+                onSelect(.start)
             } label: {
                 HStack {
                     Image(systemName: "desktopcomputer")

diff --git a/Packages/MenuBar/Sources/MenuBarItem/Internal/L10n.swift b/Packages/MenuBar/Sources/MenuBarItem/Internal/L10n.swift
@@ -36,20 +36,16 @@ internal enum L10n {
       }
     }
     internal enum VirtualMachines {
+      /// Start
+      internal static let start = L10n.tr("Localizable", "menu_bar_item.virtual_machines.start", fallback: "Start")
+      /// Stop
+      internal static let stop = L10n.tr("Localizable", "menu_bar_item.virtual_machines.stop", fallback: "Stop")
+      /// Stopping...
+      internal static let stopping = L10n.tr("Localizable", "menu_bar_item.virtual_machines.stopping", fallback: "Stopping...")
+      /// Stops when virtual machines have terminated.
+      internal static let stoppingInfo = L10n.tr("Localizable", "menu_bar_item.virtual_machines.stopping_info", fallback: "Stops when virtual machines have terminated.")
       /// Select Virtual Machine...
       internal static let unavailable = L10n.tr("Localizable", "menu_bar_item.virtual_machines.unavailable", fallback: "Select Virtual Machine...")
-      internal enum Start {
-        /// Start
-        internal static let pluralis = L10n.tr("Localizable", "menu_bar_item.virtual_machines.start.pluralis", fallback: "Start")
-        /// Start
-        internal static let singularis = L10n.tr("Localizable", "menu_bar_item.virtual_machines.start.singularis", fallback: "Start")
-      }
-      internal enum Stop {
-        /// Stop
-        internal static let pluralis = L10n.tr("Localizable", "menu_bar_item.virtual_machines.stop.pluralis", fallback: "Stop")
-        /// Stop
-        internal static let singularis = L10n.tr("Localizable", "menu_bar_item.virtual_machines.stop.singularis", fallback: "Stop")
-      }
     }
   }
 }

diff --git a/Packages/MenuBar/Sources/MenuBarItem/Internal/VirtualMachinesMenuContent.swift b/Packages/MenuBar/Sources/MenuBarItem/Internal/VirtualMachinesMenuContent.swift
@@ -1,23 +1,30 @@
-import SettingsStore
 import SwiftUI
 
 struct VirtualMachinesMenuContent: View {
     @StateObject private var viewModel: VirtualMachinesMenuContentViewModel
-    @ObservedObject private var settingsStore: SettingsStore
 
     init(viewModel: VirtualMachinesMenuContentViewModel) {
         _viewModel = StateObject(wrappedValue: viewModel)
-        _settingsStore = ObservedObject(wrappedValue: viewModel.settingsStore)
     }
 
     var body: some View {
         FleetMenuBarItem(
             hasSelectedVirtualMachine: viewModel.hasSelectedVirtualMachine,
             isFleetStarted: viewModel.isFleetStarted,
-            isEditorStarted: viewModel.isEditorStarted,
-            startsSingleVirtualMachine: settingsStore.numberOfVirtualMachines == 1,
-            onSelect: viewModel.presentFleet
-        )
+            isStoppingFleet: viewModel.isStoppingFleet,
+            isEditorStarted: viewModel.isEditorStarted
+        ) { action in
+            switch action {
+            case .start:
+                if viewModel.hasSelectedVirtualMachine {
+                    viewModel.startFleet()
+                } else {
+                    viewModel.presentSettings()
+                }
+            case .stop:
+                viewModel.stopFleet()
+            }
+        }
         Divider()
         EditorMenuBarItem(
             isEditorStarted: viewModel.isEditorStarted,

diff --git a/Packages/MenuBar/Sources/MenuBarItem/Internal/VirtualMachinesMenuContentViewModel.swift b/Packages/MenuBar/Sources/MenuBarItem/Internal/VirtualMachinesMenuContentViewModel.swift
@@ -10,6 +10,7 @@ final class VirtualMachinesMenuContentViewModel: ObservableObject {
     let settingsStore: SettingsStore
     @Published private(set) var hasSelectedVirtualMachine: Bool
     @Published private(set) var isFleetStarted = false
+    @Published private(set) var isStoppingFleet = false
     @Published private(set) var isEditorStarted = false
     var isEditorMenuBarItemEnabled: Bool {
         return !isFleetStarted && !isEditorStarted && hasSelectedVirtualMachine
@@ -35,20 +36,34 @@ final class VirtualMachinesMenuContentViewModel: ObservableObject {
         self.settingsPresenter = settingsPresenter
         self.hasSelectedVirtualMachine = settingsStore.virtualMachine != .unknown
         settingsStore.onChange.map { $0.virtualMachine != .unknown }.assign(to: \.hasSelectedVirtualMachine, on: self).store(in: &cancellables)
-        fleet.isStarted.assign(to: \.isFleetStarted, on: self).store(in: &cancellables)
+        fleet.isStarted.receive(on: DispatchQueue.main).assign(to: \.isFleetStarted, on: self).store(in: &cancellables)
+        fleet.isStopping.receive(on: DispatchQueue.main).assign(to: \.isStoppingFleet, on: self).store(in: &cancellables)
         editorService.isStarted.assign(to: \.isEditorStarted, on: self).store(in: &cancellables)
     }
 
-    func presentFleet() {
+    func startFleet() {
+        guard !isFleetStarted && hasSelectedVirtualMachine else {
+            return
+        }
+        do {
+            try fleet.start(numberOfMachines: settingsStore.numberOfVirtualMachines)
+        } catch {
+            #if DEBUG
+            print(error)
+            #endif
+        }
+    }
+
+    func stopFleet() {
         if isFleetStarted {
-            stopFleet()
-        } else if hasSelectedVirtualMachine {
-            startFleet()
-        } else {
-            settingsPresenter.presentSettings()
+            fleet.stop()
         }
     }
 
+    func presentSettings() {
+        settingsPresenter.presentSettings()
+    }
+
     func startEditor() {
         if !isEditorStarted {
             editorService.start()
@@ -64,24 +79,3 @@ final class VirtualMachinesMenuContentViewModel: ObservableObject {
         }
     }
 }
-
-private extension VirtualMachinesMenuContentViewModel {
-    private func startFleet() {
-        guard !isFleetStarted && hasSelectedVirtualMachine else {
-            return
-        }
-        do {
-            try fleet.start(numberOfMachines: settingsStore.numberOfVirtualMachines)
-        } catch {
-            #if DEBUG
-            print(error)
-            #endif
-        }
-    }
-
-    private func stopFleet() {
-        if isFleetStarted {
-            fleet.stop()
-        }
-    }
-}
diff --git a/Packages/MenuBar/Sources/MenuBarItem/Supporting files/Localizable.strings b/Packages/MenuBar/Sources/MenuBarItem/Supporting files/Localizable.strings
@@ -1,7 +1,7 @@
-"menu_bar_item.virtual_machines.start.singularis" = "Start";
-"menu_bar_item.virtual_machines.start.pluralis" = "Start";
-"menu_bar_item.virtual_machines.stop.singularis" = "Stop";
-"menu_bar_item.virtual_machines.stop.pluralis" = "Stop";
+"menu_bar_item.virtual_machines.start" = "Start";
+"menu_bar_item.virtual_machines.stop" = "Stop";
+"menu_bar_item.virtual_machines.stopping" = "Stopping...";
+"menu_bar_item.virtual_machines.stopping_info" = "Stops when virtual machines have terminated.";
 "menu_bar_item.virtual_machines.unavailable" = "Select Virtual Machine...";
 "menu_bar_item.editor.edit_virtual_machine.start" = "Edit Virtual Machine";
 "menu_bar_item.editor.edit_virtual_machine.editing" = "Editing...";

diff --git a/Packages/VirtualMachine/Sources/VirtualMachineFleet/VirtualMachineFleet.swift b/Packages/VirtualMachine/Sources/VirtualMachineFleet/VirtualMachineFleet.swift
@@ -2,6 +2,8 @@ import Combine
 
 public protocol VirtualMachineFleet {
     var isStarted: AnyPublisher<Bool, Never> { get }
+    var isStopping: AnyPublisher<Bool, Never> { get }
     func start(numberOfMachines: Int) throws
+    func stopImmediately()
     func stop()
 }
diff --git a/Packages/VirtualMachine/Sources/VirtualMachineFleetLive/VirtualMachineFleetLive.swift b/Packages/VirtualMachine/Sources/VirtualMachineFleetLive/VirtualMachineFleetLive.swift
@@ -7,15 +7,18 @@ import VirtualMachineFleet
 
 public final class VirtualMachineFleetLive: VirtualMachineFleet {
     public let isStarted: AnyPublisher<Bool, Never>
+    public let isStopping: AnyPublisher<Bool, Never>
 
     private let logger = Logger(category: "VirtualMachineFleetLive")
     private let virtualMachineFactory: VirtualMachineFactory
-    private var activeTasks: [Task<(), Never>] = []
+    private var activeTasks: [String: Task<(), Never>] = [:]
     private let _isStarted = CurrentValueSubject<Bool, Never>(false)
+    private let _isStopping = CurrentValueSubject<Bool, Never>(false)
 
     public init(virtualMachineFactory: VirtualMachineFactory) {
         self.virtualMachineFactory = virtualMachineFactory
         self.isStarted = _isStarted.eraseToAnyPublisher()
+        self.isStopping = _isStopping.eraseToAnyPublisher()
     }
 
     public func start(numberOfMachines: Int) throws {
@@ -30,15 +33,17 @@ public final class VirtualMachineFleetLive: VirtualMachineFleet {
         }
     }
 
-    public func stop() {
-        guard _isStarted.value else {
-            return
-        }
+    public func stopImmediately() {
         _isStarted.value = false
-        for task in activeTasks {
+        _isStopping.value = false
+        for (_, task) in activeTasks {
             task.cancel()
         }
-        activeTasks = []
+        activeTasks = [:]
+    }
+
+    public func stop() {
+        _isStopping.value = true
     }
 }
 
@@ -48,6 +53,9 @@ private extension VirtualMachineFleetLive {
             while !Task.isCancelled {
                 do {
                     try await runVirtualMachine(named: name)
+                    if _isStopping.value {
+                        activeTasks[name]?.cancel()
+                    }
                 } catch {
                     // Ignore the error and try again until the task is cancelled.
                     // The error should have been logged using OSLog so we know what is going on in case we need to debug.
@@ -56,8 +64,12 @@ private extension VirtualMachineFleetLive {
                 }
             }
             logger.info("Task running virtual machine named \(name, privacy: .public) was cancelled.")
+            activeTasks.removeValue(forKey: name)
+            if activeTasks.isEmpty {
+                stopImmediately()
+            }
         }
-        activeTasks.append(task)
+        activeTasks[name] = task
     }
 
     private func runVirtualMachine(named name: String) async throws {