Serialize command info lookups on a single dedicated runspace

Co-authored-by: jessehouwing <4173387+jessehouwing@users.noreply.github.com>

Author: Copilot

Engine/CommandInfoCache.cs

@@ -23,7 +23,17 @@ internal class CommandInfoCache : IDisposable
         private const int MaxLookupAttempts = 3;
 
         private readonly ConcurrentDictionary<CommandLookupKey, Lazy<CommandInfo>> _commandInfoCache;
-        private readonly RunspacePool _runspacePool;
+
+        /// <summary>
+        /// Guards all access to <see cref="_runspace"/> so that only one thread at a time drives the
+        /// PowerShell engine. The engine is not thread safe, so concurrent lookups can fail transiently,
+        /// see https://github.com/PowerShell/PowerShell/issues/4003.
+        /// A monitor is used rather than a semaphore because it is re-entrant, which avoids a deadlock
+        /// should a lookup ever end up calling back into the cache on the same thread.
+        /// </summary>
+        private readonly object _runspaceLock = new object();
+
+        private readonly Runspace _runspace;
         private bool disposed = false;
 
         /// <summary>
@@ -32,11 +42,13 @@ internal class CommandInfoCache : IDisposable
         public CommandInfoCache()
         {
             _commandInfoCache = new ConcurrentDictionary<CommandLookupKey, Lazy<CommandInfo>>();
-            _runspacePool = RunspaceFactory.CreateRunspacePool(1, 10);
-            _runspacePool.Open();
+            // A single runspace rather than a pool: all lookups are serialized on it, so that the
+            // PowerShell engine is never driven concurrently.
+            _runspace = RunspaceFactory.CreateRunspace();
+            _runspace.Open();
         }
 
-        /// <summary>Dispose the runspace pool</summary>
+        /// <summary>Dispose the runspace</summary>
         public void Dispose()
         {
             Dispose(true);
@@ -52,7 +64,14 @@ protected virtual void Dispose(bool disposing)
 
             if ( disposing )
             {
-                _runspacePool.Dispose();
+                // Take the lock so that the runspace is not disposed while a lookup is in flight.
+                lock (_runspaceLock)
+                {
+                    disposed = true;
+                    _runspace.Dispose();
+                }
+
+                return;
             }
 
             disposed = true;
@@ -123,41 +142,51 @@ private CommandInfo GetCommandInfoInternal(string cmdName, CommandTypes? command
 
             for (int attempt = 1; ; attempt++)
             {
-                using (var ps = System.Management.Automation.PowerShell.Create())
+                // Serialize all use of the PowerShell engine. Only cache misses reach this point;
+                // lookups that are already cached are served without taking the lock.
+                lock (_runspaceLock)
                 {
-                    ps.RunspacePool = _runspacePool;
-
-                    ps.AddCommand("Get-Command")
-                        .AddParameter("Name", actualCmdName)
-                        .AddParameter("ErrorAction", "SilentlyContinue");
-
-                    if (commandType != null)
+                    if (disposed)
                     {
-                        ps.AddParameter("CommandType", commandType);
+                        return null;
                     }
 
-                    if (!string.IsNullOrEmpty(moduleName))
+                    using (var ps = System.Management.Automation.PowerShell.Create())
                     {
-                        ps.AddParameter("Module", moduleName);
-                    }
+                        ps.Runspace = _runspace;
 
-                    try
-                    {
-                        return ps.Invoke<CommandInfo>()
-                            .FirstOrDefault();
-                    }
-                    // 'Get-Command' is invoked with 'SilentlyContinue', so a CommandNotFoundException can only
-                    // mean that the engine failed to resolve 'Get-Command' itself in the pooled runspace.
-                    // That happens intermittently because the PowerShell engine is not thread safe, see
-                    // https://github.com/PowerShell/PowerShell/issues/4003 and
-                    // https://github.com/PowerShell/PSScriptAnalyzer/issues/2205
-                    // Retrying usually succeeds, but rather than failing the whole analysis when it does not,
-                    // treat the command as unresolvable.
-                    catch (CommandNotFoundException)
-                    {
-                        if (attempt >= MaxLookupAttempts)
+                        ps.AddCommand("Get-Command")
+                            .AddParameter("Name", actualCmdName)
+                            .AddParameter("ErrorAction", "SilentlyContinue");
+
+                        if (commandType != null)
+                        {
+                            ps.AddParameter("CommandType", commandType);
+                        }
+
+                        if (!string.IsNullOrEmpty(moduleName))
+                        {
+                            ps.AddParameter("Module", moduleName);
+                        }
+
+                        try
+                        {
+                            return ps.Invoke<CommandInfo>()
+                                .FirstOrDefault();
+                        }
+                        // 'Get-Command' is invoked with 'SilentlyContinue', so a CommandNotFoundException can only
+                        // mean that the engine failed to resolve 'Get-Command' itself in the runspace.
+                        // That happened intermittently when lookups ran concurrently because the PowerShell engine
+                        // is not thread safe, see https://github.com/PowerShell/PowerShell/issues/4003 and
+                        // https://github.com/PowerShell/PSScriptAnalyzer/issues/2205
+                        // Lookups are serialized now, so this should no longer occur, but the retry is kept as a
+                        // safety net for hosts that drive the engine from other threads at the same time.
+                        catch (CommandNotFoundException)
                         {
-                            return null;
+                            if (attempt >= MaxLookupAttempts)
+                            {
+                                return null;
+                            }
                         }
                     }
                 }

Tests/Engine/CommandInfoCacheConcurrency.tests.ps1

@@ -0,0 +1,59 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+Describe "Concurrent command lookups" {
+    BeforeAll {
+        # The concurrency driver is written in C# so that the lookups really do run on separate
+        # threads. Invoking a PowerShell script block on a thread pool thread would introduce
+        # runspace affinity problems of its own and would not test the command info cache.
+        $analyzerAssembly = [Microsoft.Windows.PowerShell.ScriptAnalyzer.Helper].Assembly.Location
+        Add-Type -IgnoreWarnings -WarningAction SilentlyContinue -ReferencedAssemblies $analyzerAssembly, ([System.Management.Automation.PSObject].Assembly.Location) -TypeDefinition @'
+using System.Threading.Tasks;
+using Microsoft.Windows.PowerShell.ScriptAnalyzer;
+
+public static class ConcurrentCommandLookup
+{
+    public static string[] Lookup(string[] commandNames)
+    {
+        var helper = Helper.Instance;
+        var tasks = new Task<string>[commandNames.Length];
+        for (int i = 0; i < commandNames.Length; i++)
+        {
+            string name = commandNames[i];
+            tasks[i] = Task.Run(() =>
+            {
+                var commandInfo = helper.GetCommandInfo(name);
+                return commandInfo == null ? null : commandInfo.Name;
+            });
+        }
+
+        Task.WaitAll(tasks);
+
+        var results = new string[tasks.Length];
+        for (int i = 0; i < tasks.Length; i++)
+        {
+            results[i] = tasks[i].Result;
+        }
+
+        return results;
+    }
+}
+'@
+    }
+
+    It "resolves commands from several threads without failing" {
+        $commandNames = @(
+            'Get-ChildItem', 'Where-Object', 'ForEach-Object', 'Get-Content', 'Write-Output',
+            'Test-Path', 'Get-Command', 'Select-Object', 'Sort-Object', 'Measure-Object'
+        ) * 4
+
+        # A lookup that hits the thread safety problem throws, which fails the test.
+        $results = [ConcurrentCommandLookup]::Lookup($commandNames)
+
+        $results.Count | Should -Be $commandNames.Count
+        # A failed lookup returns null, so every entry must name the command that was requested.
+        for ($i = 0; $i -lt $commandNames.Count; $i++) {
+            $results[$i] | Should -BeExactly $commandNames[$i]
+        }
+    }
+}