Converted WorkspaceService.GetFile, TryGetFile, and EnumeratePSFiles to async

Fixed object disposal issue on ScriptFile constructor overload

Author: dkattan

src/PowerShellEditorServices/Extensions/Api/WorkspaceService.cs

@@ -5,6 +5,7 @@
 using System;
 using System.Collections.Generic;
 using System.Management.Automation.Language;
+using System.Threading.Tasks;
 
 namespace Microsoft.PowerShell.EditorServices.Extensions.Services
 {
@@ -64,22 +65,22 @@ public interface IWorkspaceService
         /// </summary>
         /// <param name="fileUri">The absolute URI of the file to get.</param>
         /// <returns>A representation of the file.</returns>
-        IEditorScriptFile GetFile(Uri fileUri);
+        Task<IEditorScriptFile> GetFile(Uri fileUri);
 
         /// <summary>
         /// Attempt to get a file within the workspace.
         /// </summary>
         /// <param name="fileUri">The absolute URI of the file to get.</param>
         /// <param name="file">The file, if it was found.</param>
         /// <returns>True if the file was found, false otherwise.</returns>
-        bool TryGetFile(Uri fileUri, out IEditorScriptFile file);
+        Task<IEditorScriptFile?> TryGetFile(Uri fileUri);
 
         /// <summary>
         /// Get all the open files in the editor workspace.
         /// The result is not kept up to date as files are opened or closed.
         /// </summary>
         /// <returns>All open files in the editor workspace.</returns>
-        IReadOnlyList<IEditorScriptFile> GetOpenedFiles();
+        Task<IReadOnlyList<IEditorScriptFile>> GetOpenedFiles();
     }
 
     internal class EditorScriptFile : IEditorScriptFile
@@ -124,30 +125,28 @@ internal WorkspaceService(
 
         public IReadOnlyList<string> ExcludedFileGlobs { get; }
 
-        public IEditorScriptFile GetFile(Uri fileUri) => GetEditorFileFromScriptFile(_workspaceService.GetFile(fileUri));
+        public async Task<IEditorScriptFile> GetFile(Uri fileUri) => await GetEditorFileFromScriptFile(await _workspaceService.GetFile(fileUri).ConfigureAwait(false)).ConfigureAwait(false);
 
-        public bool TryGetFile(Uri fileUri, out IEditorScriptFile file)
+        public async Task<IEditorScriptFile?> TryGetFile(Uri fileUri)
         {
-            if (!_workspaceService.TryGetFile(fileUri.LocalPath, out ScriptFile scriptFile))
+            if (await _workspaceService.TryGetFile(fileUri.LocalPath).ConfigureAwait(false) is not ScriptFile scriptFile)
             {
-                file = null;
-                return false;
+                return null;
             }
 
-            file = GetEditorFileFromScriptFile(scriptFile);
-            return true;
+            return await GetEditorFileFromScriptFile(scriptFile).ConfigureAwait(false);
         }
 
-        public IReadOnlyList<IEditorScriptFile> GetOpenedFiles()
+        public async Task<IReadOnlyList<IEditorScriptFile>> GetOpenedFiles()
         {
             List<IEditorScriptFile> files = new();
             foreach (ScriptFile openedFile in _workspaceService.GetOpenedFiles())
             {
-                files.Add(GetEditorFileFromScriptFile(openedFile));
+                files.Add(await GetEditorFileFromScriptFile(openedFile).ConfigureAwait(false));
             }
             return files.AsReadOnly();
         }
 
-        private static IEditorScriptFile GetEditorFileFromScriptFile(ScriptFile file) => new EditorScriptFile(file);
+        private static async Task<IEditorScriptFile> GetEditorFileFromScriptFile(ScriptFile file) => new EditorScriptFile(file);
     }
 }

src/PowerShellEditorServices/Services/Analysis/AnalysisService.cs

@@ -28,6 +28,7 @@ namespace Microsoft.PowerShell.EditorServices.Services
     /// </summary>
     internal class AnalysisService : IDisposable
     {
+        private readonly SemaphoreSlim _initializationSemaphore = new(1, 1);
         /// <summary>
         /// Reliably generate an ID for a diagnostic record to track it.
         /// </summary>
@@ -91,9 +92,6 @@ internal static string GetUniqueIdFromDiagnostic(Diagnostic diagnostic)
         private readonly int _analysisDelayMillis = 750;
 
         private readonly ConcurrentDictionary<ScriptFile, CorrectionTableEntry> _mostRecentCorrectionsByFile = new();
-
-        private Lazy<PssaCmdletAnalysisEngine> _analysisEngineLazy;
-
         private CancellationTokenSource _diagnosticsCancellationTokenSource;
 
         private readonly string _pssaModulePath;
@@ -112,14 +110,37 @@ public AnalysisService(
             _languageServer = languageServer;
             _configurationService = configurationService;
             _workspaceService = workspaceService;
-            _analysisEngineLazy = new Lazy<PssaCmdletAnalysisEngine>(InstantiateAnalysisEngine);
             _pssaModulePath = Path.Combine(hostInfo.BundledModulePath, "PSScriptAnalyzer");
         }
-
+        private PssaCmdletAnalysisEngine? _analysisEngine;
         /// <summary>
         /// The analysis engine to use for running script analysis.
         /// </summary>
-        internal PssaCmdletAnalysisEngine AnalysisEngine => _analysisEngineLazy?.Value;
+        internal PssaCmdletAnalysisEngine? AnalysisEngine
+        {
+            get
+            {
+                _initializationSemaphore.Wait();
+                try
+                {
+                    if (_analysisEngine == null)
+                    {
+                        _analysisEngine = InstantiateAnalysisEngine().GetAwaiter().GetResult();
+                        IsValueCreated = true;
+                    }
+                }
+                finally
+                {
+                    _initializationSemaphore.Release();
+                }
+                return _analysisEngine;
+            }
+            private set
+            {
+                IsValueCreated = true;
+                _analysisEngine = value;
+            }
+        }
 
         /// <summary>
         /// Sets up a script analysis run, eventually returning the result.
@@ -215,7 +236,8 @@ public async Task<string> GetCommentHelpText(string functionText, string helpLoc
         /// <returns>A thread-safe readonly dictionary of the code actions of the particular file.</returns>
         public async Task<IReadOnlyDictionary<string, IEnumerable<MarkerCorrection>>> GetMostRecentCodeActionsForFileAsync(DocumentUri uri)
         {
-            if (!_workspaceService.TryGetFile(uri, out ScriptFile file)
+            ScriptFile? file = await _workspaceService.TryGetFile(uri).ConfigureAwait(false);
+            if (file is null
                 || !_mostRecentCorrectionsByFile.TryGetValue(file, out CorrectionTableEntry corrections))
             {
                 return null;
@@ -239,7 +261,7 @@ public async Task<IReadOnlyDictionary<string, IEnumerable<MarkerCorrection>>> Ge
         /// </summary>
         /// <param name="_">The sender of the configuration update event.</param>
         /// <param name="settings">The new language server settings.</param>
-        public void OnConfigurationUpdated(object _, LanguageServerSettings settings)
+        public async Task OnConfigurationUpdated(LanguageServerSettings settings)
         {
             if (settings.ScriptAnalysis.Enable)
             {
@@ -249,24 +271,25 @@ public void OnConfigurationUpdated(object _, LanguageServerSettings settings)
 
         private void EnsureEngineSettingsCurrent()
         {
-            if (_analysisEngineLazy is null
+            if (AnalysisEngine is null
                     || (_pssaSettingsFilePath is not null
                         && !File.Exists(_pssaSettingsFilePath)))
             {
                 InitializeAnalysisEngineToCurrentSettings();
             }
         }
 
-        private void InitializeAnalysisEngineToCurrentSettings()
+        private async Task InitializeAnalysisEngineToCurrentSettings()
         {
+
             // We may be triggered after the lazy factory is set,
             // but before it's been able to instantiate
-            if (_analysisEngineLazy is null)
+            if (AnalysisEngine is null)
             {
-                _analysisEngineLazy = new Lazy<PssaCmdletAnalysisEngine>(InstantiateAnalysisEngine);
+                AnalysisEngine = await InstantiateAnalysisEngine().ConfigureAwait(false);
                 return;
             }
-            else if (!_analysisEngineLazy.IsValueCreated)
+            else if (IsValueCreated)
             {
                 return;
             }
@@ -276,15 +299,16 @@ private void InitializeAnalysisEngineToCurrentSettings()
 
             // Clear the open file markers and set the new engine factory
             ClearOpenFileMarkers();
-            _analysisEngineLazy = new Lazy<PssaCmdletAnalysisEngine>(() => RecreateAnalysisEngine(currentAnalysisEngine));
+            AnalysisEngine = await RecreateAnalysisEngine(currentAnalysisEngine).ConfigureAwait(false);
         }
 
-        internal PssaCmdletAnalysisEngine InstantiateAnalysisEngine()
+        internal async Task<PssaCmdletAnalysisEngine> InstantiateAnalysisEngine()
         {
             PssaCmdletAnalysisEngine.Builder pssaCmdletEngineBuilder = new(_loggerFactory);
 
             // If there's a settings file use that
-            if (TryFindSettingsFile(out string settingsFilePath))
+            string? settingsFilePath = await TryFindSettingsFile().ConfigureAwait(false);
+            if (!string.IsNullOrEmpty(settingsFilePath))
             {
                 _logger.LogInformation($"Configuring PSScriptAnalyzer with rules at '{settingsFilePath}'");
                 _pssaSettingsFilePath = settingsFilePath;
@@ -299,9 +323,10 @@ internal PssaCmdletAnalysisEngine InstantiateAnalysisEngine()
             return pssaCmdletEngineBuilder.Build(_pssaModulePath);
         }
 
-        private PssaCmdletAnalysisEngine RecreateAnalysisEngine(PssaCmdletAnalysisEngine oldAnalysisEngine)
+        private async Task<PssaCmdletAnalysisEngine> RecreateAnalysisEngine(PssaCmdletAnalysisEngine oldAnalysisEngine)
         {
-            if (TryFindSettingsFile(out string settingsFilePath))
+            string? settingsFilePath = await TryFindSettingsFile().ConfigureAwait(false);
+            if (!string.IsNullOrEmpty(settingsFilePath))
             {
                 _logger.LogInformation($"Recreating analysis engine with rules at '{settingsFilePath}'");
                 _pssaSettingsFilePath = settingsFilePath;
@@ -312,27 +337,27 @@ private PssaCmdletAnalysisEngine RecreateAnalysisEngine(PssaCmdletAnalysisEngine
             return oldAnalysisEngine.RecreateWithRules(s_defaultRules);
         }
 
-        private bool TryFindSettingsFile(out string settingsFilePath)
+        private async Task<string?> TryFindSettingsFile()
         {
             string configuredPath = _configurationService?.CurrentSettings.ScriptAnalysis.SettingsPath;
-
+            string? settingsFilePath;
             if (string.IsNullOrEmpty(configuredPath))
             {
                 settingsFilePath = null;
-                return false;
+                return settingsFilePath;
             }
 
-            settingsFilePath = _workspaceService?.ResolveWorkspacePath(configuredPath);
+            settingsFilePath = await (_workspaceService?.ResolveWorkspacePath(configuredPath)).ConfigureAwait(false);
 
             if (settingsFilePath is null
                 || !File.Exists(settingsFilePath))
             {
                 _logger.LogInformation($"Unable to find PSSA settings file at '{configuredPath}'. Loading default rules.");
                 settingsFilePath = null;
-                return false;
+                return settingsFilePath;
             }
 
-            return true;
+            return settingsFilePath;
         }
 
         private void ClearOpenFileMarkers()
@@ -467,19 +492,22 @@ private static Hashtable GetCommentHelpRuleSettings(string helpLocation, bool fo
 
         #region IDisposable Support
         private bool disposedValue; // To detect redundant calls
+        private bool IsValueCreated;
 
         protected virtual void Dispose(bool disposing)
         {
             if (!disposedValue)
             {
                 if (disposing)
                 {
-                    if (_analysisEngineLazy?.IsValueCreated == true)
+                    if (IsValueCreated)
                     {
-                        _analysisEngineLazy.Value.Dispose();
+                        AnalysisEngine.Dispose();
                     }
+                    _initializationSemaphore.Dispose();
 
                     _diagnosticsCancellationTokenSource?.Dispose();
+                    _analysisEngine?.Dispose();
                 }
 
                 disposedValue = true;

src/PowerShellEditorServices/Services/DebugAdapter/Handlers/BreakpointHandlers.cs

@@ -49,7 +49,8 @@ public BreakpointHandlers(
 
         public async Task<SetBreakpointsResponse> Handle(SetBreakpointsArguments request, CancellationToken cancellationToken)
         {
-            if (!_workspaceService.TryGetFile(request.Source.Path, out ScriptFile scriptFile))
+            ScriptFile scriptFile = await _workspaceService.TryGetFile(request.Source.Path).ConfigureAwait(false);
+            if (scriptFile is null)
             {
                 string message = _debugStateService.NoDebug ? string.Empty : "Source file could not be accessed, breakpoint not set.";
                 IEnumerable<Breakpoint> srcBreakpoints = request.Breakpoints

src/PowerShellEditorServices/Services/DebugAdapter/Handlers/ConfigurationDoneHandler.cs

@@ -116,8 +116,8 @@ internal async Task LaunchScriptAsync(string scriptToLaunch)
             }
             else // It's a URI to an untitled script, or a raw script.
             {
-                bool isScriptFile = _workspaceService.TryGetFile(scriptToLaunch, out ScriptFile untitledScript);
-                if (isScriptFile && BreakpointApiUtils.SupportsBreakpointApis(_runspaceContext.CurrentRunspace))
+                ScriptFile untitledScript = await _workspaceService.TryGetFile(scriptToLaunch).ConfigureAwait(false);
+                if (untitledScript is ScriptFile && BreakpointApiUtils.SupportsBreakpointApis(_runspaceContext.CurrentRunspace))
                 {
                     // Parse untitled files with their `Untitled:` URI as the filename which will
                     // cache the URI and contents within the PowerShell parser. By doing this, we
@@ -149,7 +149,7 @@ internal async Task LaunchScriptAsync(string scriptToLaunch)
                     command = PSCommandHelpers.BuildDotSourceCommandWithArguments(
                         string.Concat(
                             "{" + System.Environment.NewLine,
-                            isScriptFile ? untitledScript.Contents : scriptToLaunch,
+                            untitledScript is ScriptFile ? untitledScript.Contents : scriptToLaunch,
                             System.Environment.NewLine + "}"),
                             _debugStateService?.Arguments);
                 }

src/PowerShellEditorServices/Services/PowerShell/Handlers/GetCommentHelpHandler.cs

@@ -27,8 +27,8 @@ public GetCommentHelpHandler(
         public async Task<CommentHelpRequestResult> Handle(CommentHelpRequestParams request, CancellationToken cancellationToken)
         {
             CommentHelpRequestResult result = new();
-
-            if (!_workspaceService.TryGetFile(request.DocumentUri, out ScriptFile scriptFile))
+            ScriptFile? scriptFile = await _workspaceService.TryGetFile(request.DocumentUri).ConfigureAwait(false);
+            if (scriptFile is null)
             {
                 return result;
             }

src/PowerShellEditorServices/Services/PowerShell/Host/PsesInternalHost.cs

@@ -181,6 +181,7 @@ public PsesInternalHost(
             _pipelineThread = new Thread(Run, maxStackSize)
             {
                 Name = "PSES Pipeline Execution Thread",
+                IsBackground = false
             };
 
             if (VersionUtils.IsWindows)

src/PowerShellEditorServices/Services/Symbols/SymbolsService.cs

@@ -3,14 +3,6 @@
 
 #nullable enable
 
-using System;
-using System.Collections.Concurrent;
-using System.Collections.Generic;
-using System.Linq;
-using System.Management.Automation;
-using System.Management.Automation.Language;
-using System.Threading;
-using System.Threading.Tasks;
 using Microsoft.Extensions.Logging;
 using Microsoft.PowerShell.EditorServices.CodeLenses;
 using Microsoft.PowerShell.EditorServices.Logging;
@@ -21,6 +13,14 @@
 using Microsoft.PowerShell.EditorServices.Services.Symbols;
 using Microsoft.PowerShell.EditorServices.Services.TextDocument;
 using Microsoft.PowerShell.EditorServices.Utility;
+using System;
+using System.Collections.Concurrent;
+using System.Collections.Generic;
+using System.Linq;
+using System.Management.Automation;
+using System.Management.Automation.Language;
+using System.Threading;
+using System.Threading.Tasks;
 
 namespace Microsoft.PowerShell.EditorServices.Services
 {
@@ -333,11 +333,11 @@ internal async Task ScanWorkspacePSFiles(CancellationToken cancellationToken = d
             if (scanTask is null)
             {
                 scanTask = Task.Run(
-                    () =>
+                    async () =>
                     {
-                        foreach (string file in _workspaceService.EnumeratePSFiles())
+                        foreach (string file in await _workspaceService.EnumeratePSFiles(cancellationToken).ConfigureAwait(false))
                         {
-                            if (_workspaceService.TryGetFile(file, out ScriptFile scriptFile))
+                            if ((await _workspaceService.TryGetFile(file).ConfigureAwait(false)) is ScriptFile scriptFile)
                             {
                                 scriptFile.References.EnsureInitialized();
                             }
@@ -457,7 +457,7 @@ internal async Task ScanWorkspacePSFiles(CancellationToken cancellationToken = d
             return functionDefinitionAst as FunctionDefinitionAst;
         }
 
-        internal void OnConfigurationUpdated(object _, LanguageServerSettings e)
+        internal async Task OnConfigurationUpdated(LanguageServerSettings e)
         {
             if (e.AnalyzeOpenDocumentsOnly)
             {