eh

This reverts commit 845e4d4877.

src/modules/cmdpal/Microsoft.CmdPal.UI.ViewModels/TopLevelCommandManager.cs

@@ -2,8 +2,10 @@
 // The Microsoft Corporation licenses this file to you under the MIT license.
 // See the LICENSE file in the project root for more information.
 
+using System.Collections.Concurrent;
 using System.Collections.Immutable;
 using System.Collections.ObjectModel;
+using System.Diagnostics;
 using CommunityToolkit.Mvvm.ComponentModel;
 using CommunityToolkit.Mvvm.Input;
 using CommunityToolkit.Mvvm.Messaging;
@@ -53,53 +55,44 @@ public partial class TopLevelCommandManager : ObservableObject,
         {
             CommandProviderWrapper wrapper = new(provider, _taskScheduler);
             _builtInCommands.Add(wrapper);
-            await LoadTopLevelCommandsFromProvider(wrapper);
+            var commands = await LoadTopLevelCommandsFromProvider(wrapper);
+            lock (TopLevelCommands)
+            {
+                foreach (var c in commands)
+                {
+                    TopLevelCommands.Add(c);
+                }
+            }
         }
 
         return true;
     }
 
     // May be called from a background thread
-    private async Task LoadTopLevelCommandsFromProvider(CommandProviderWrapper commandProvider)
+    private async Task<IEnumerable<TopLevelViewModel>> LoadTopLevelCommandsFromProvider(CommandProviderWrapper commandProvider)
     {
         WeakReference<IPageContext> weakSelf = new(this);
 
         await commandProvider.LoadTopLevelCommands(_serviceProvider, weakSelf);
 
         var settings = _serviceProvider.GetService<SettingsModel>()!;
-        var makeAndAdd = (ICommandItem? i, bool fallback) =>
+
+        List<TopLevelViewModel> commands = [];
+
+        foreach (var item in commandProvider.TopLevelItems)
         {
-            var commandItemViewModel = new CommandItemViewModel(new(i), weakSelf);
-            var topLevelViewModel = new TopLevelViewModel(commandItemViewModel, fallback, commandProvider.ExtensionHost, commandProvider.ProviderId, settings, _serviceProvider);
+            commands.Add(item);
+        }
 
-            lock (TopLevelCommands)
-            {
-                TopLevelCommands.Add(topLevelViewModel);
-            }
-        };
-
-        await Task.Factory.StartNew(
-            () =>
-            {
-                lock (TopLevelCommands)
-                {
-                    foreach (var item in commandProvider.TopLevelItems)
-                    {
-                        TopLevelCommands.Add(item);
-                    }
-
-                    foreach (var item in commandProvider.FallbackItems)
-                    {
-                        TopLevelCommands.Add(item);
-                    }
-                }
-            },
-            CancellationToken.None,
-            TaskCreationOptions.None,
-            _taskScheduler);
+        foreach (var item in commandProvider.FallbackItems)
+        {
+            commands.Add(item);
+        }
 
         commandProvider.CommandsChanged -= CommandProvider_CommandsChanged;
         commandProvider.CommandsChanged += CommandProvider_CommandsChanged;
+
+        return commands;
     }
 
     // By all accounts, we're already on a background thread (the COM call
@@ -214,7 +207,7 @@ public partial class TopLevelCommandManager : ObservableObject,
         _extensionCommandProviders.Clear();
         if (extensions != null)
         {
-            await StartExtensionsAndGetCommands(extensions);
+            StartExtensionsAndGetCommands(extensions);
         }
 
         extensionService.OnExtensionAdded += ExtensionService_OnExtensionAdded;
@@ -228,36 +221,94 @@ public partial class TopLevelCommandManager : ObservableObject,
     private void ExtensionService_OnExtensionAdded(IExtensionService sender, IEnumerable<IExtensionWrapper> extensions)
     {
         // When we get an extension install event, hop off to a BG thread
-        _ = Task.Run(async () =>
+        _ = Task.Run(() =>
         {
             // for each newly installed extension, start it and get commands
             // from it. One single package might have more than one
             // IExtensionWrapper in it.
-            await StartExtensionsAndGetCommands(extensions);
+            StartExtensionsAndGetCommands(extensions);
         });
     }
 
-    private async Task StartExtensionsAndGetCommands(IEnumerable<IExtensionWrapper> extensions)
+    private void StartExtensionsAndGetCommands(IEnumerable<IExtensionWrapper> extensions)
     {
-        // TODO This most definitely needs a lock
-        foreach (var extension in extensions)
-        {
-            Logger.LogDebug($"Starting {extension.PackageFullName}");
-            try
-            {
-                // start it ...
-                await extension.StartExtensionAsync();
+        var timer = new Stopwatch();
+        timer.Start();
 
-                // ... and fetch the command provider from it.
-                CommandProviderWrapper wrapper = new(extension, _taskScheduler);
-                _extensionCommandProviders.Add(wrapper);
-                await LoadTopLevelCommandsFromProvider(wrapper);
-            }
-            catch (Exception ex)
+        // Start all extensions in parallel using Parallel.ForEach
+        var wrappers = new ConcurrentBag<CommandProviderWrapper>();
+        Parallel.ForEach(extensions, extension =>
+        {
+            var wrapper = StartExtensionWithTimeoutAsync(extension).GetAwaiter().GetResult();
+            if (wrapper != null)
             {
-                Logger.LogError(ex.ToString());
+                wrappers.Add(wrapper);
+            }
+        });
+
+        foreach (var wrapper in wrappers)
+        {
+            _extensionCommandProviders.Add(wrapper);
+        }
+
+        // Load the commands from the providers in parallel
+        var commandSets = new ConcurrentBag<IEnumerable<TopLevelViewModel>>();
+
+        Parallel.ForEach(wrappers, wrapper =>
+        {
+            var commands = LoadCommandsWithTimeoutAsync(wrapper).GetAwaiter().GetResult();
+            if (commands != null)
+            {
+                commandSets.Add(commands);
+            }
+        });
+
+        lock (TopLevelCommands)
+        {
+            foreach (var commands in commandSets)
+            {
+                foreach (var c in commands)
+                {
+                    TopLevelCommands.Add(c);
+                }
             }
         }
+
+        timer.Stop();
+        Logger.LogDebug($"Loading extensions took {timer.ElapsedMilliseconds} ms");
+    }
+
+    private async Task<CommandProviderWrapper?> StartExtensionWithTimeoutAsync(IExtensionWrapper extension)
+    {
+        Logger.LogDebug($"Starting {extension.PackageFullName}");
+        try
+        {
+            await extension.StartExtensionAsync().WaitAsync(TimeSpan.FromSeconds(10));
+            return new CommandProviderWrapper(extension, _taskScheduler);
+        }
+        catch (Exception ex)
+        {
+            Logger.LogError($"Failed to start extension {extension.PackageFullName}: {ex}");
+            return null; // Return null for failed extensions
+        }
+    }
+
+    private async Task<IEnumerable<TopLevelViewModel>?> LoadCommandsWithTimeoutAsync(CommandProviderWrapper wrapper)
+    {
+        try
+        {
+            return await LoadTopLevelCommandsFromProvider(wrapper!).WaitAsync(TimeSpan.FromSeconds(10));
+        }
+        catch (TimeoutException)
+        {
+            Logger.LogError($"Loading commands from {wrapper!.ExtensionHost?.Extension?.PackageFullName} timed out");
+        }
+        catch (Exception ex)
+        {
+            Logger.LogError($"Failed to load commands for extension {wrapper!.ExtensionHost?.Extension?.PackageFullName}: {ex}");
+        }
+
+        return null;
     }
 
     private void ExtensionService_OnExtensionRemoved(IExtensionService sender, IEnumerable<IExtensionWrapper> extensions)