Increase lock granularity for CommandInfo cache (#1166)

* Redo command info cache changes

* Fix comment typo

Co-Authored-By: rjmholt <[email protected]>

* Trigger CI

Engine/CommandInfoCache.cs

@@ -0,0 +1,119 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+using System;
+using System.Collections.Concurrent;
+using System.Management.Automation;
+using System.Linq;
+
+namespace Microsoft.Windows.PowerShell.ScriptAnalyzer
+{
+    /// <summary>
+    /// Provides threadsafe caching around CommandInfo lookups with `Get-Command -Name ...`.
+    /// </summary>
+    internal class CommandInfoCache
+    {
+        private readonly ConcurrentDictionary<CommandLookupKey, Lazy<CommandInfo>> _commandInfoCache;
+
+        private readonly Helper _helperInstance;
+
+        /// <summary>
+        /// Create a fresh command info cache instance.
+        /// </summary>
+        public CommandInfoCache(Helper pssaHelperInstance)
+        {
+            _commandInfoCache = new ConcurrentDictionary<CommandLookupKey, Lazy<CommandInfo>>();
+            _helperInstance = pssaHelperInstance;
+        }
+
+        /// <summary>
+        /// Retrieve a command info object about a command.
+        /// </summary>
+        /// <param name="commandName">Name of the command to get a commandinfo object for.</param>
+        /// <param name="aliasName">The alias of the command to be used in the cache key. If not given, uses the command name.</param>
+        /// <param name="commandTypes">What types of command are needed. If omitted, all types are retrieved.</param>
+        /// <returns></returns>
+        public CommandInfo GetCommandInfo(string commandName, string aliasName = null, CommandTypes? commandTypes = null)
+        {
+            if (string.IsNullOrWhiteSpace(commandName))
+            {
+                return null;
+            }
+
+            // If alias name is given, we store the entry under that, but search with the command name
+            var key = new CommandLookupKey(aliasName ?? commandName, commandTypes);
+
+            // Atomically either use PowerShell to query a command info object, or fetch it from the cache
+            return _commandInfoCache.GetOrAdd(key, new Lazy<CommandInfo>(() => GetCommandInfoInternal(commandName, commandTypes))).Value;
+        }
+
+        /// <summary>
+        /// Retrieve a command info object about a command.
+        /// </summary>
+        /// <param name="commandName">Name of the command to get a commandinfo object for.</param>
+        /// <param name="commandTypes">What types of command are needed. If omitted, all types are retrieved.</param>
+        /// <returns></returns>
+        [Obsolete("Alias lookup is expensive and should not be relied upon for command lookup")]
+        public CommandInfo GetCommandInfoLegacy(string commandOrAliasName, CommandTypes? commandTypes = null)
+        {
+            string commandName = _helperInstance.GetCmdletNameFromAlias(commandOrAliasName);
+
+            return string.IsNullOrEmpty(commandName)
+                ? GetCommandInfo(commandOrAliasName, commandTypes: commandTypes)
+                : GetCommandInfo(commandName, aliasName: commandOrAliasName, commandTypes: commandTypes);
+        }
+
+        /// <summary>
+        /// Get a CommandInfo object of the given command name
+        /// </summary>
+        /// <returns>Returns null if command does not exists</returns>
+        private static CommandInfo GetCommandInfoInternal(string cmdName, CommandTypes? commandType)
+        {
+            using (var ps = System.Management.Automation.PowerShell.Create())
+            {
+                ps.AddCommand("Get-Command")
+                    .AddParameter("Name", cmdName)
+                    .AddParameter("ErrorAction", "SilentlyContinue");
+
+                if (commandType != null)
+                {
+                    ps.AddParameter("CommandType", commandType);
+                }
+
+                return ps.Invoke<CommandInfo>()
+                    .FirstOrDefault();
+            }
+        }
+
+        private struct CommandLookupKey : IEquatable<CommandLookupKey>
+        {
+            private readonly string Name;
+
+            private readonly CommandTypes CommandTypes;
+
+            internal CommandLookupKey(string name, CommandTypes? commandTypes)
+            {
+                Name = name;
+                CommandTypes = commandTypes ?? CommandTypes.All;
+            }
+
+            public bool Equals(CommandLookupKey other)
+            {
+                return CommandTypes == other.CommandTypes
+                    && Name.Equals(other.Name, StringComparison.OrdinalIgnoreCase);
+            }
+
+            public override int GetHashCode()
+            {
+                // Algorithm from https://stackoverflow.com/questions/1646807/quick-and-simple-hash-code-combinations
+                unchecked
+                {
+                    int hash = 17;
+                    hash = hash * 31 + Name.ToUpperInvariant().GetHashCode();
+                    hash = hash * 31 + CommandTypes.GetHashCode();
+                    return hash;
+                }
+            }
+        }
+    }
+}

Engine/Helper.cs

@@ -24,11 +24,11 @@ public class Helper
 
         private CommandInvocationIntrinsics invokeCommand;
         private IOutputWriter outputWriter;
-        private Object getCommandLock = new object();
         private readonly static Version minSupportedPSVersion = new Version(3, 0);
         private Dictionary<string, Dictionary<string, object>> ruleArguments;
         private PSVersionTable psVersionTable;
-        private Dictionary<CommandLookupKey, CommandInfo> commandInfoCache;
+
+        private readonly Lazy<CommandInfoCache> _commandInfoCacheLazy;
 
         #endregion
 
@@ -100,14 +100,20 @@ internal set
         private string[] functionScopes = new string[] { "global:", "local:", "script:", "private:"};
 
         private string[] variableScopes = new string[] { "global:", "local:", "script:", "private:", "variable:", ":"};
+
+        /// <summary>
+        /// Store of command info objects for commands. Memoizes results.
+        /// </summary>
+        private CommandInfoCache CommandInfoCache => _commandInfoCacheLazy.Value;
+
         #endregion
 
         /// <summary>
         /// Initializes the Helper class.
         /// </summary>
         private Helper()
         {
-
+            _commandInfoCacheLazy = new Lazy<CommandInfoCache>(() => new CommandInfoCache(pssaHelperInstance: this));
         }
 
         /// <summary>
@@ -123,7 +129,7 @@ private Helper()
         /// </param>
         public Helper(
             CommandInvocationIntrinsics invokeCommand,
-            IOutputWriter outputWriter)
+            IOutputWriter outputWriter) : this()
         {
             this.invokeCommand = invokeCommand;
             this.outputWriter = outputWriter;
@@ -140,10 +146,6 @@ public void Initialize()
             KeywordBlockDictionary = new Dictionary<String, List<Tuple<int, int>>>(StringComparer.OrdinalIgnoreCase);
             VariableAnalysisDictionary = new Dictionary<Ast, VariableAnalysis>();
             ruleArguments = new Dictionary<string, Dictionary<string, object>>(StringComparer.OrdinalIgnoreCase);
-            if (commandInfoCache == null)
-            {
-                commandInfoCache = new Dictionary<CommandLookupKey, CommandInfo>();
-            }
 
             IEnumerable<CommandInfo> aliases = this.invokeCommand.GetCommands("*", CommandTypes.Alias, true);
 
@@ -677,7 +679,6 @@ private CommandInfo GetCommandInfoInternal(string cmdName, CommandTypes? command
         }
 
         /// <summary>
-
         ///  Legacy method, new callers should use <see cref="GetCommandInfo"/> instead.
         ///  Given a command's name, checks whether it exists. It does not use the passed in CommandTypes parameter, which is a bug.
         ///  But existing method callers are already depending on this behaviour and therefore this could not be simply fixed.
@@ -689,30 +690,7 @@ private CommandInfo GetCommandInfoInternal(string cmdName, CommandTypes? command
         [Obsolete]
         public CommandInfo GetCommandInfoLegacy(string name, CommandTypes? commandType = null)
         {
-            if (string.IsNullOrWhiteSpace(name))
-            {
-                return null;
-            }
-
-            // check if it is an alias
-            string cmdletName = Helper.Instance.GetCmdletNameFromAlias(name);
-            if (string.IsNullOrWhiteSpace(cmdletName))
-            {
-                cmdletName = name;
-            }
-
-            var key = new CommandLookupKey(name, commandType);
-            lock (getCommandLock)
-            {
-                if (commandInfoCache.ContainsKey(key))
-                {
-                    return commandInfoCache[key];
-                }
-
-                var commandInfo = GetCommandInfoInternal(cmdletName, commandType);
-                commandInfoCache.Add(key, commandInfo);
-                return commandInfo;
-            }
+            return CommandInfoCache.GetCommandInfoLegacy(commandOrAliasName: name, commandTypes: commandType);
         }
 
         /// <summary>
@@ -723,23 +701,7 @@ public CommandInfo GetCommandInfoLegacy(string name, CommandTypes? commandType =
         /// <returns></returns>
         public CommandInfo GetCommandInfo(string name, CommandTypes? commandType = null)
         {
-            if (string.IsNullOrWhiteSpace(name))
-            {
-                return null;
-            }
-
-            var key = new CommandLookupKey(name, commandType);
-            lock (getCommandLock)
-            {
-                if (commandInfoCache.ContainsKey(key))
-                {
-                    return commandInfoCache[key];
-                }
-
-                var commandInfo = GetCommandInfoInternal(name, commandType);
-                commandInfoCache.Add(key, commandInfo);
-                return commandInfo;
-            }
+            return CommandInfoCache.GetCommandInfo(name, commandTypes: commandType);
         }
 
         /// <summary>