Update PSConfiguration.ReadValueFromFile to make it faster and more memory efficient (#10839)

src/System.Management.Automation/engine/PSConfiguration.cs

@@ -33,30 +33,47 @@ namespace System.Management.Automation.Configuration
     /// Reads from and writes to the JSON configuration files.
     /// The config values were originally stored in the Windows registry.
     /// </summary>
+    /// <remarks>
+    /// The config file access APIs are designed to avoid hitting the disk as much as possible.
+    /// - For the first read request targeting a config file, the config data is read from the file and then cached as a 'JObject' instance;
+    ///     * the first read request happens very early during the startup of 'pwsh'.
+    /// - For the subsequent read requests targeting the same config file, they will then work with that 'JObject' instance;
+    /// - For the write request targeting a config file, the cached config data corresponding to that config file will be refreshed after the write operation is successfully done.
+    ///
+    /// To summarize the expected behavior:
+    /// Once a 'pwsh' process starts up -
+    ///   1. any changes made to the config file from outside this 'pwsh' process is not guaranteed to be seen by it (most likely won't be seen).
+    ///   2. any changes to the config file by this 'pwsh' process via the config file access APIs will be seen by it, if it chooses to read those changes afterwards.
+    /// </remarks>
     internal sealed class PowerShellConfig
     {
+        private const string configFileName = "powershell.config.json";
+
         // Provide a singleton
-        private static readonly PowerShellConfig s_instance = new PowerShellConfig();
-        internal static PowerShellConfig Instance => s_instance;
+        internal static readonly PowerShellConfig Instance = new PowerShellConfig();
 
         // The json file containing system-wide configuration settings.
-        // When passed as a pwsh command-line option,
-        // overrides the system wide configuration file.
+        // When passed as a pwsh command-line option, overrides the system wide configuration file.
         private string systemWideConfigFile;
         private string systemWideConfigDirectory;
 
         // The json file containing the per-user configuration settings.
-        private string perUserConfigFile;
-        private string perUserConfigDirectory;
+        private readonly string perUserConfigFile;
+        private readonly string perUserConfigDirectory;
 
-        private const string configFileName = "powershell.config.json";
+        // Note: JObject and JsonSerializer are thread safe.
+        // Root Json objects corresponding to the configuration file for 'AllUsers' and 'CurrentUser' respectively.
+        // They are used as a cache to avoid hitting the disk for every read operation.
+        private readonly JObject[] configRoots;
+        private readonly JObject emptyConfig;
+        private readonly JsonSerializer serializer;
 
         /// <summary>
         /// Lock used to enable multiple concurrent readers and singular write locks within a single process.
         /// TODO: This solution only works for IO from a single process.
         ///       A more robust solution is needed to enable ReaderWriterLockSlim behavior between processes.
         /// </summary>
-        private ReaderWriterLockSlim fileLock = new ReaderWriterLockSlim();
+        private readonly ReaderWriterLockSlim fileLock;
 
         private PowerShellConfig()
         {
@@ -65,11 +82,16 @@ namespace System.Management.Automation.Configuration
             systemWideConfigFile = Path.Combine(systemWideConfigDirectory, configFileName);
 
             // Sets the per-user configuration directory
-            // Note: This directory may or may not exist depending upon the
-            // execution scenario. Writes will attempt to create the directory
-            // if it does not already exist.
+            // Note: This directory may or may not exist depending upon the execution scenario.
+            // Writes will attempt to create the directory if it does not already exist.
             perUserConfigDirectory = Platform.ConfigDirectory;
             perUserConfigFile = Path.Combine(perUserConfigDirectory, configFileName);
+
+            emptyConfig = new JObject();
+            configRoots = new JObject[2];
+            serializer = JsonSerializer.Create(new JsonSerializerSettings { TypeNameHandling = TypeNameHandling.None, MaxDepth = 10 });
+
+            fileLock = new ReaderWriterLockSlim();
         }
 
         private string GetConfigFilePath(ConfigScope scope)
@@ -201,15 +223,11 @@ namespace System.Management.Automation.Configuration
         /// </summary>
         internal string[] GetExperimentalFeatures()
         {
-            string[] features = Array.Empty<string>();
-            if (File.Exists(perUserConfigFile))
-            {
-                features = ReadValueFromFile<string[]>(ConfigScope.CurrentUser, "ExperimentalFeatures", Array.Empty<string>());
-            }
+            string[] features = ReadValueFromFile(ConfigScope.CurrentUser, "ExperimentalFeatures", Array.Empty<string>());
 
             if (features.Length == 0)
             {
-                features = ReadValueFromFile<string[]>(ConfigScope.AllUsers, "ExperimentalFeatures", Array.Empty<string>());
+                features = ReadValueFromFile(ConfigScope.AllUsers, "ExperimentalFeatures", Array.Empty<string>());
             }
 
             return features;
@@ -382,42 +400,52 @@ namespace System.Management.Automation.Configuration
         /// <param name="scope">The ConfigScope of the configuration file to update.</param>
         /// <param name="key">The string key of the value.</param>
         /// <param name="defaultValue">The default value to return if the key is not present.</param>
-        /// <param name="readImpl"></param>
-        private T ReadValueFromFile<T>(ConfigScope scope, string key, T defaultValue = default(T),
-                                       Func<JToken, JsonSerializer, T, T> readImpl = null)
+        private T ReadValueFromFile<T>(ConfigScope scope, string key, T defaultValue = default)
         {
             string fileName = GetConfigFilePath(scope);
-            if (!File.Exists(fileName)) { return defaultValue; }
+            JObject configData = configRoots[(int)scope];
 
-            // Open file for reading, but allow multiple readers
-            fileLock.EnterReadLock();
-            try
+            if (configData == null)
             {
-                // The config file can be locked by another process
-                // so we wait some milliseconds in 'WaitForFile()' for recovery before stop current process.
-                using (var readerStream = WaitForFile(fileName, FileMode.Open, FileAccess.Read, FileShare.ReadWrite))
-                using (var streamReader = new StreamReader(readerStream))
-                using (var jsonReader = new JsonTextReader(streamReader))
+                if (File.Exists(fileName))
                 {
-                    var settings = new JsonSerializerSettings() { TypeNameHandling = TypeNameHandling.None, MaxDepth = 10 };
-                    var serializer = JsonSerializer.Create(settings);
-
-                    var configData = serializer.Deserialize<JObject>(jsonReader);
-                    if (configData != null && configData.TryGetValue(key, StringComparison.OrdinalIgnoreCase, out JToken jToken))
+                    try
                     {
-                        return readImpl != null ? readImpl(jToken, serializer, defaultValue) : jToken.ToObject<T>(serializer);
+                        // Open file for reading, but allow multiple readers
+                        fileLock.EnterReadLock();
+
+                        using var stream = OpenFileStreamWithRetry(fileName, FileMode.Open, FileAccess.Read, FileShare.ReadWrite);
+                        using var jsonReader = new JsonTextReader(new StreamReader(stream));
+
+                        configData = serializer.Deserialize<JObject>(jsonReader) ?? emptyConfig;
+                    }
+                    finally
+                    {
+                        fileLock.ExitReadLock();
                     }
                 }
+                else
+                {
+                    configData = emptyConfig;
+                }
+
+                // Set the configuration cache.
+                JObject originalValue = Interlocked.CompareExchange(ref configRoots[(int)scope], configData, null);
+                if (originalValue != null)
+                {
+                    configData = originalValue;
+                }
             }
-            finally
+
+            if (configData != emptyConfig && configData.TryGetValue(key, StringComparison.OrdinalIgnoreCase, out JToken jToken))
             {
-                fileLock.ExitReadLock();
+                return jToken.ToObject<T>(serializer);
             }
 
             return defaultValue;
         }
 
-        private FileStream WaitForFile(string fullPath, FileMode mode, FileAccess access, FileShare share)
+        private static FileStream OpenFileStreamWithRetry(string fullPath, FileMode mode, FileAccess access, FileShare share)
         {
             const int MaxTries = 5;
             for (int numTries = 0; numTries < MaxTries; numTries++)
@@ -437,7 +465,7 @@ namespace System.Management.Automation.Configuration
                 }
             }
 
-            throw new IOException(nameof(WaitForFile));
+            throw new IOException(nameof(OpenFileStreamWithRetry));
         }
 
         /// <summary>
@@ -450,94 +478,91 @@ namespace System.Management.Automation.Configuration
         /// <param name="addValue">Whether the key-value pair should be added to or removed from the file.</param>
         private void UpdateValueInFile<T>(ConfigScope scope, string key, T value, bool addValue)
         {
-            string fileName = GetConfigFilePath(scope);
-            fileLock.EnterWriteLock();
             try
             {
-                // Since multiple properties can be in a single file, replacement
-                // is required instead of overwrite if a file already exists.
-                // Handling the read and write operations within a single FileStream
-                // prevents other processes from reading or writing the file while
-                // the update is in progress. It also locks out readers during write
-                // operations.
-                using (FileStream fs = WaitForFile(fileName, FileMode.OpenOrCreate, FileAccess.ReadWrite, FileShare.None))
+                string fileName = GetConfigFilePath(scope);
+                fileLock.EnterWriteLock();
+
+                // Since multiple properties can be in a single file, replacement is required instead of overwrite if a file already exists.
+                // Handling the read and write operations within a single FileStream prevents other processes from reading or writing the file while
+                // the update is in progress. It also locks out readers during write operations.
+
+                JObject jsonObject = null;
+                using FileStream fs = OpenFileStreamWithRetry(fileName, FileMode.OpenOrCreate, FileAccess.ReadWrite, FileShare.None);
+
+                // UTF8, BOM detection, and bufferSize are the same as the basic stream constructor.
+                // The most important parameter here is the last one, which keeps underlying stream open after StreamReader is disposed
+                // so that it can be reused for the subsequent write operation.
+                using (StreamReader streamRdr = new StreamReader(fs, Encoding.UTF8, detectEncodingFromByteOrderMarks: true, bufferSize: 1024, leaveOpen: true))
+                using (JsonTextReader jsonReader = new JsonTextReader(streamRdr))
                 {
-                    JObject jsonObject = null;
-
-                    // UTF8, BOM detection, and bufferSize are the same as the basic stream constructor.
-                    // The most important parameter here is the last one, which keeps the StreamReader
-                    // (and FileStream) open during Dispose so that it can be reused for the write
-                    // operation.
-                    using (StreamReader streamRdr = new StreamReader(fs, Encoding.UTF8, true, 1024, true))
-                    using (JsonTextReader jsonReader = new JsonTextReader(streamRdr))
+                    // Safely determines whether there is content to read from the file
+                    bool isReadSuccess = jsonReader.Read();
+                    if (isReadSuccess)
                     {
-                        // Safely determines whether there is content to read from the file
-                        bool isReadSuccess = jsonReader.Read();
-                        if (isReadSuccess)
-                        {
-                            // Read the stream into a root JObject for manipulation
-                            jsonObject = (JObject)JToken.ReadFrom(jsonReader);
-                            JProperty propertyToModify = jsonObject.Property(key);
+                        // Read the stream into a root JObject for manipulation
+                        jsonObject = serializer.Deserialize<JObject>(jsonReader);
+                        JProperty propertyToModify = jsonObject.Property(key);
 
-                            if (propertyToModify == null)
+                        if (propertyToModify == null)
+                        {
+                            // The property doesn't exist, so add it
+                            if (addValue)
                             {
-                                // The property doesn't exist, so add it
-                                if (addValue)
-                                {
-                                    jsonObject.Add(new JProperty(key, value));
-                                }
-                                // else the property doesn't exist so there is nothing to remove
-                            }
-                            // The property exists
-                            else
-                            {
-                                if (addValue)
-                                {
-                                    propertyToModify.Replace(new JProperty(key, value));
-                                }
-                                else
-                                {
-                                    propertyToModify.Remove();
-                                }
+                                jsonObject.Add(new JProperty(key, value));
                             }
+                            // else the property doesn't exist so there is nothing to remove
                         }
                         else
                         {
-                            // The file doesn't already exist and we want to write to it
-                            // or it exists with no content.
-                            // A new file will be created that contains only this value.
-                            // If the file doesn't exist and a we don't want to write to it, no
-                            // action is necessary.
+                            // The property exists
                             if (addValue)
                             {
-                                jsonObject = new JObject(new JProperty(key, value));
+                                propertyToModify.Replace(new JProperty(key, value));
                             }
                             else
                             {
-                                return;
+                                propertyToModify.Remove();
                             }
                         }
                     }
-
-                    // Reset the stream position to the beginning so that the
-                    // changes to the file can be written to disk
-                    fs.Seek(0, SeekOrigin.Begin);
-
-                    // Update the file with new content
-                    using (StreamWriter streamWriter = new StreamWriter(fs))
-                    using (JsonTextWriter jsonWriter = new JsonTextWriter(streamWriter))
+                    else
                     {
-                        // The entire document exists within the root JObject.
-                        // I just need to write that object to produce the document.
-                        jsonObject.WriteTo(jsonWriter);
-
-                        // This trims the file if the file shrank. If the file grew,
-                        // it is a no-op. The purpose is to trim extraneous characters
-                        // from the file stream when the resultant JObject is smaller
-                        // than the input JObject.
-                        fs.SetLength(fs.Position);
+                        // The file doesn't already exist and we want to write to it or it exists with no content.
+                        // A new file will be created that contains only this value.
+                        // If the file doesn't exist and a we don't want to write to it, no action is needed.
+                        if (addValue)
+                        {
+                            jsonObject = new JObject(new JProperty(key, value));
+                        }
+                        else
+                        {
+                            return;
+                        }
                     }
                 }
+
+                // Reset the stream position to the beginning so that the
+                // changes to the file can be written to disk
+                fs.Seek(0, SeekOrigin.Begin);
+
+                // Update the file with new content
+                using (StreamWriter streamWriter = new StreamWriter(fs))
+                using (JsonTextWriter jsonWriter = new JsonTextWriter(streamWriter))
+                {
+                    // The entire document exists within the root JObject.
+                    // I just need to write that object to produce the document.
+                    jsonObject.WriteTo(jsonWriter);
+
+                    // This trims the file if the file shrank. If the file grew,
+                    // it is a no-op. The purpose is to trim extraneous characters
+                    // from the file stream when the resultant JObject is smaller
+                    // than the input JObject.
+                    fs.SetLength(fs.Position);
+                }
+
+                // Refresh the configuration cache.
+                Interlocked.Exchange(ref configRoots[(int)scope], jsonObject);
             }
             finally
             {
@@ -554,13 +579,8 @@ namespace System.Management.Automation.Configuration
         /// <param name="value">The value to write.</param>
         private void WriteValueToFile<T>(ConfigScope scope, string key, T value)
         {
-            // Defaults to system wide.
-            if (ConfigScope.CurrentUser == scope)
+            if (ConfigScope.CurrentUser == scope && !Directory.Exists(perUserConfigDirectory))
             {
-                // Exceptions are not caught so that they will propagate to the
-                // host for display to the user.
-                // CreateDirectory will succeed if the directory already exists
-                // so there is no reason to check Directory.Exists().
                 Directory.CreateDirectory(perUserConfigDirectory);
             }
 

src/TypeCatalogGen/TypeCatalogGen.cs

@@ -415,7 +415,7 @@ DUPLICATE key '{fullName}' from '{strongAssemblyName}' (IsObsolete? {isTypeObsol
         /// </summary>
         private static void WritePowerShellAssemblyLoadContextPartialClass(string targetFilePath, Dictionary<string, TypeMetadata> typeNameToAssemblyMap)
         {
-            const string SourceFormat = "            typeCatalog[\"{0}\"] = \"{1}\";";
+            const string SourceFormat = "{2}                {{\"{0}\", \"{1}\"}},";
             const string SourceHead = @"//
 // This file is auto-generated by TypeCatalogGen.exe during build of Microsoft.PowerShell.CoreCLR.AssemblyLoadContext.dll.
 // This file will be compiled into Microsoft.PowerShell.CoreCLR.AssemblyLoadContext.dll.
@@ -425,7 +425,6 @@ DUPLICATE key '{fullName}' from '{strongAssemblyName}' (IsObsolete? {isTypeObsol
 // catalog based on the reference assemblies of .NET Core.
 //
 using System.Collections.Generic;
-using System.Runtime.Loader;
 
 namespace System.Management.Automation
 {{
@@ -433,10 +432,9 @@ namespace System.Management.Automation
     {{
         private Dictionary<string, string> InitializeTypeCatalog()
         {{
-            Dictionary<string, string> typeCatalog = new Dictionary<string, string>({0}, StringComparer.OrdinalIgnoreCase);
-";
+            return new Dictionary<string, string>({0}, StringComparer.OrdinalIgnoreCase) {{";
             const string SourceEnd = @"
-            return typeCatalog;
+            };
         }
     }
 }
@@ -445,7 +443,7 @@ namespace System.Management.Automation
             StringBuilder sourceCode = new StringBuilder(string.Format(CultureInfo.InvariantCulture, SourceHead, typeNameToAssemblyMap.Count));
             foreach (KeyValuePair<string, TypeMetadata> pair in typeNameToAssemblyMap)
             {
-                sourceCode.AppendLine(string.Format(CultureInfo.InvariantCulture, SourceFormat, pair.Key, pair.Value.AssemblyName));
+                sourceCode.Append(string.Format(CultureInfo.InvariantCulture, SourceFormat, pair.Key, pair.Value.AssemblyName, Environment.NewLine));
             }
 
             sourceCode.Append(SourceEnd);

test/xUnit/csharp/test_PSConfiguration.cs

@@ -360,6 +360,15 @@ namespace PSTests.Sequential
             File.Create(fileName).Dispose();
         }
 
+        internal void ForceReadingFromFile()
+        {
+            // Reset the cached roots.
+            FieldInfo roots = typeof(PowerShellConfig).GetField("configRoots", BindingFlags.NonPublic | BindingFlags.Instance);
+            JObject[] value = (JObject[])roots.GetValue(PowerShellConfig.Instance);
+            value[0] = null;
+            value[1] = null;
+        }
+
         #endregion
     }
 
@@ -376,6 +385,8 @@ namespace PSTests.Sequential
         public void PowerShellConfig_GetPowerShellPolicies_BothConfigFilesNotEmpty()
         {
             fixture.SetupConfigFile1();
+            fixture.ForceReadingFromFile();
+
             var sysPolicies = PowerShellConfig.Instance.GetPowerShellPolicies(ConfigScope.AllUsers);
             var userPolicies = PowerShellConfig.Instance.GetPowerShellPolicies(ConfigScope.CurrentUser);
 
@@ -390,6 +401,8 @@ namespace PSTests.Sequential
         public void PowerShellConfig_GetPowerShellPolicies_EmptyUserConfig()
         {
             fixture.SetupConfigFile2();
+            fixture.ForceReadingFromFile();
+
             var sysPolicies = PowerShellConfig.Instance.GetPowerShellPolicies(ConfigScope.AllUsers);
             var userPolicies = PowerShellConfig.Instance.GetPowerShellPolicies(ConfigScope.CurrentUser);
 
@@ -403,6 +416,8 @@ namespace PSTests.Sequential
         public void PowerShellConfig_GetPowerShellPolicies_EmptySystemConfig()
         {
             fixture.SetupConfigFile3();
+            fixture.ForceReadingFromFile();
+
             var sysPolicies = PowerShellConfig.Instance.GetPowerShellPolicies(ConfigScope.AllUsers);
             var userPolicies = PowerShellConfig.Instance.GetPowerShellPolicies(ConfigScope.CurrentUser);
 
@@ -416,6 +431,8 @@ namespace PSTests.Sequential
         public void PowerShellConfig_GetPowerShellPolicies_BothConfigFilesEmpty()
         {
             fixture.SetupConfigFile4();
+            fixture.ForceReadingFromFile();
+
             var sysPolicies = PowerShellConfig.Instance.GetPowerShellPolicies(ConfigScope.AllUsers);
             var userPolicies = PowerShellConfig.Instance.GetPowerShellPolicies(ConfigScope.CurrentUser);
 
@@ -427,6 +444,8 @@ namespace PSTests.Sequential
         public void PowerShellConfig_GetPowerShellPolicies_BothConfigFilesNotExist()
         {
             fixture.CleanupConfigFiles();
+            fixture.ForceReadingFromFile();
+
             var sysPolicies = PowerShellConfig.Instance.GetPowerShellPolicies(ConfigScope.AllUsers);
             var userPolicies = PowerShellConfig.Instance.GetPowerShellPolicies(ConfigScope.CurrentUser);
 
@@ -438,6 +457,7 @@ namespace PSTests.Sequential
         public void Utils_GetPolicySetting_BothConfigFilesNotEmpty()
         {
             fixture.SetupConfigFile1();
+            fixture.ForceReadingFromFile();
 
             ScriptExecution scriptExecution;
             scriptExecution = Utils.GetPolicySetting<ScriptExecution>(Utils.SystemWideOnlyConfig);
@@ -536,6 +556,7 @@ namespace PSTests.Sequential
         public void Utils_GetPolicySetting_EmptyUserConfig()
         {
             fixture.SetupConfigFile2();
+            fixture.ForceReadingFromFile();
 
             // The CurrentUser config is empty
             ScriptExecution scriptExecution;
@@ -634,6 +655,7 @@ namespace PSTests.Sequential
         public void Utils_GetPolicySetting_EmptySystemConfig()
         {
             fixture.SetupConfigFile3();
+            fixture.ForceReadingFromFile();
 
             // The SystemWide config is empty
             ScriptExecution scriptExecution;
@@ -733,6 +755,7 @@ namespace PSTests.Sequential
         public void Utils_GetPolicySetting_BothConfigFilesEmpty()
         {
             fixture.SetupConfigFile4();
+            fixture.ForceReadingFromFile();
 
             // Both config files are empty
             ScriptExecution scriptExecution;
@@ -832,6 +855,7 @@ namespace PSTests.Sequential
         public void Utils_GetPolicySetting_BothConfigFilesNotExist()
         {
             fixture.CleanupConfigFiles();
+            fixture.ForceReadingFromFile();
 
             // Both config files don't exist
             ScriptExecution scriptExecution;