Merge pull request #41844 from dibarbet/navbar_workspace_registration

Listen for workspace registration changes to ensure we add/remove the…

src/EditorFeatures/Core.Wpf/Suggestions/SuggestedActionsSource.cs

@@ -1012,7 +1012,7 @@ private void OnWorkspaceChanged(object sender, EventArgs e)
                 RegisterEventsToWorkspace(_registration.Workspace);
             }
 
-            private void RegisterEventsToWorkspace(Workspace workspace)
+            private void RegisterEventsToWorkspace(Workspace? workspace)
             {
                 _workspace = workspace;
 
@@ -1022,7 +1022,7 @@ private void RegisterEventsToWorkspace(Workspace workspace)
                 }
 
                 _workspace.DocumentActiveContextChanged += OnActiveContextChanged;
-                _workspaceStatusService = workspace.Services.GetService<IWorkspaceStatusService>();
+                _workspaceStatusService = _workspace.Services.GetService<IWorkspaceStatusService>();
                 if (_workspaceStatusService != null)
                 {
                     _workspaceStatusService.StatusChanged += OnWorkspaceStatusChanged;

src/EditorFeatures/Core/Extensibility/NavigationBar/INavigationBarController.cs

@@ -2,10 +2,14 @@
 // The .NET Foundation licenses this file to you under the MIT license.
 // See the LICENSE file in the project root for more information.
 
+#nullable enable
+
 namespace Microsoft.CodeAnalysis.Editor
 {
     internal interface INavigationBarController
     {
         void Disconnect();
+
+        void SetWorkspace(Workspace? newWorkspace);
     }
 }

src/EditorFeatures/Core/Implementation/NavigationBar/NavigationBarController.cs

@@ -2,6 +2,8 @@
 // The .NET Foundation licenses this file to you under the MIT license.
 // See the LICENSE file in the project root for more information.
 
+#nullable enable
+
 using System;
 using System.Collections.Generic;
 using System.Linq;
@@ -33,7 +35,6 @@ internal partial class NavigationBarController : ForegroundThreadAffinitizedObje
         private readonly ITextBuffer _subjectBuffer;
         private readonly IWaitIndicator _waitIndicator;
         private readonly IAsynchronousOperationListener _asyncListener;
-        private readonly WorkspaceRegistration _workspaceRegistration;
 
         /// <summary>
         /// If we have pushed a full list to the presenter in response to a focus event, this
@@ -43,7 +44,7 @@ internal partial class NavigationBarController : ForegroundThreadAffinitizedObje
         private VersionStamp? _versionStampOfFullListPushedToPresenter = null;
 
         private bool _disconnected = false;
-        private Workspace _workspace;
+        private Workspace? _workspace;
 
         public NavigationBarController(
             IThreadingContext threadingContext,
@@ -57,8 +58,6 @@ internal partial class NavigationBarController : ForegroundThreadAffinitizedObje
             _subjectBuffer = subjectBuffer;
             _waitIndicator = waitIndicator;
             _asyncListener = asyncListener;
-            _workspaceRegistration = Workspace.GetWorkspaceRegistration(subjectBuffer.AsTextContainer());
-            _workspaceRegistration.WorkspaceChanged += OnWorkspaceRegistrationChanged;
 
             presenter.CaretMoved += OnCaretMoved;
             presenter.ViewFocused += OnViewFocused;
@@ -76,18 +75,12 @@ internal partial class NavigationBarController : ForegroundThreadAffinitizedObje
                     null));
 
             _selectedItemInfoTask = Task.FromResult(new NavigationBarSelectedTypeAndMember(null, null));
-
-            if (_workspaceRegistration.Workspace != null)
-            {
-                ConnectToWorkspace(_workspaceRegistration.Workspace);
-            }
         }
 
-        private void OnWorkspaceRegistrationChanged(object sender, EventArgs e)
+        public void SetWorkspace(Workspace? newWorkspace)
         {
             DisconnectFromWorkspace();
 
-            var newWorkspace = _workspaceRegistration.Workspace;
             if (newWorkspace != null)
             {
                 ConnectToWorkspace(newWorkspace);
@@ -151,8 +144,6 @@ public void Disconnect()
 
             _presenter.Disconnect();
 
-            _workspaceRegistration.WorkspaceChanged -= OnWorkspaceRegistrationChanged;
-
             _disconnected = true;
 
             // Cancel off any remaining background work
@@ -171,8 +162,8 @@ private void OnWorkspaceChanged(object sender, WorkspaceChangeEventArgs args)
             // If the displayed project is being renamed, retrigger the update
             if (args.Kind == WorkspaceChangeKind.ProjectChanged && args.ProjectId != null)
             {
-                var oldProject = args.OldSolution.GetProject(args.ProjectId);
-                var newProject = args.NewSolution.GetProject(args.ProjectId);
+                var oldProject = args.OldSolution.GetRequiredProject(args.ProjectId);
+                var newProject = args.NewSolution.GetRequiredProject(args.ProjectId);
 
                 if (oldProject.Name != newProject.Name)
                 {
@@ -267,7 +258,7 @@ private void UpdateDropDownsSynchronously(CancellationToken cancellationToken)
             _versionStampOfFullListPushedToPresenter = _modelTask.Result.SemanticVersionStamp;
         }
 
-        private void GetProjectItems(out IList<NavigationBarProjectItem> projectItems, out NavigationBarProjectItem selectedProjectItem)
+        private void GetProjectItems(out IList<NavigationBarProjectItem> projectItems, out NavigationBarProjectItem? selectedProjectItem)
         {
             var documents = _subjectBuffer.CurrentSnapshot.GetRelatedDocumentsWithChanges();
             if (!documents.Any())
@@ -323,8 +314,8 @@ private void PushSelectedItemsToPresenter(NavigationBarSelectedTypeAndMember sel
             var oldLeft = selectedItems.TypeItem;
             var oldRight = selectedItems.MemberItem;
 
-            NavigationBarItem newLeft = null;
-            NavigationBarItem newRight = null;
+            NavigationBarItem? newLeft = null;
+            NavigationBarItem? newRight = null;
             var listOfLeft = new List<NavigationBarItem>();
             var listOfRight = new List<NavigationBarItem>();
 
@@ -396,7 +387,7 @@ private void ProcessItemSelectionSynchronously(NavigationBarItem item, Cancellat
                 var document = _subjectBuffer.CurrentSnapshot.GetOpenDocumentInCurrentContextWithChanges();
                 if (document != null)
                 {
-                    var languageService = document.GetLanguageService<INavigationBarItemService>();
+                    var languageService = document.GetRequiredLanguageService<INavigationBarItemService>();
 
                     NavigateToItem(item, document, _subjectBuffer.CurrentSnapshot, languageService, cancellationToken);
                 }

src/EditorFeatures/Test2/NavigationBar/NavigationBarPresenterTests.vb

@@ -302,6 +302,7 @@ End Class
 
                 Dim controllerFactory = workspace.GetService(Of INavigationBarControllerFactoryService)()
                 Dim controller = controllerFactory.CreateController(mockPresenter, document.GetTextBuffer())
+                controller.SetWorkspace(workspace)
 
                 mockPresenter.RaiseDropDownFocused()
                 Assert.Equal("VBProj", projectName)

src/VisualStudio/Core/Def/Implementation/LanguageService/AbstractLanguageService`2.VsCodeWindowManager.cs

@@ -2,15 +2,22 @@
 // The .NET Foundation licenses this file to you under the MIT license.
 // See the LICENSE file in the project root for more information.
 
+#nullable enable
+
 using System.Diagnostics;
 using System.Linq;
+using System.Threading.Tasks;
+using Microsoft.CodeAnalysis;
 using Microsoft.CodeAnalysis.Editor;
 using Microsoft.CodeAnalysis.Editor.Options;
+using Microsoft.CodeAnalysis.Editor.Shared.Utilities;
 using Microsoft.CodeAnalysis.Options;
 using Microsoft.CodeAnalysis.Shared.Extensions;
+using Microsoft.CodeAnalysis.Shared.TestHooks;
 using Microsoft.CodeAnalysis.Text;
 using Microsoft.VisualStudio.LanguageServices.Implementation.NavigationBar;
 using Microsoft.VisualStudio.TextManager.Interop;
+using Microsoft.VisualStudio.Threading;
 using Roslyn.Utilities;
 
 namespace Microsoft.VisualStudio.LanguageServices.Implementation.LanguageService
@@ -22,21 +29,71 @@ internal class VsCodeWindowManager : IVsCodeWindowManager, IVsCodeWindowEvents
             private readonly TLanguageService _languageService;
             private readonly IVsCodeWindow _codeWindow;
             private readonly ComEventSink _sink;
-            private readonly IOptionService _optionService;
+            private readonly IThreadingContext _threadingContext;
+            private readonly IAsynchronousOperationListener _asynchronousOperationListener;
 
-            private INavigationBarController _navigationBarController;
-            private IVsDropdownBarClient _dropdownBarClient;
+            private INavigationBarController? _navigationBarController;
+            private IVsDropdownBarClient? _dropdownBarClient;
+            private IOptionService? _optionService;
+            private WorkspaceRegistration? _workspaceRegistration;
 
             public VsCodeWindowManager(TLanguageService languageService, IVsCodeWindow codeWindow)
             {
                 _languageService = languageService;
                 _codeWindow = codeWindow;
 
-                var workspace = languageService.Package.ComponentModel.GetService<VisualStudioWorkspace>();
-                _optionService = workspace.Services.GetService<IOptionService>();
+                _threadingContext = languageService.Package.ComponentModel.GetService<IThreadingContext>();
+
+                var listenerProvider = languageService.Package.ComponentModel.GetService<IAsynchronousOperationListenerProvider>();
+                _asynchronousOperationListener = listenerProvider.GetListener(FeatureAttribute.NavigationBar);
 
                 _sink = ComEventSink.Advise<IVsCodeWindowEvents>(codeWindow, this);
-                _optionService.OptionChanged += OnOptionChanged;
+            }
+
+            private void OnWorkspaceRegistrationChanged(object sender, System.EventArgs e)
+            {
+                var token = _asynchronousOperationListener.BeginAsyncOperation(nameof(OnWorkspaceRegistrationChanged));
+
+                // Fire and forget to update the navbar based on the workspace registration
+                // to avoid blocking the caller and possible deadlocks workspace registration changed events under lock.
+                UpdateWorkspace().CompletesAsyncOperation(token).Forget();
+            }
+
+            private async Task UpdateWorkspace()
+            {
+                // This event may not be triggered on the main thread, but adding and removing the navbar
+                // must be done from the main thread.
+                await _threadingContext.JoinableTaskFactory.SwitchToMainThreadAsync();
+
+                // If the workspace registration is missing, addornments have been removed.
+                if (_workspaceRegistration == null)
+                {
+                    return;
+                }
+
+                // There's a new workspace, so make sure we unsubscribe from the old workspace option changes and subscribe to new.
+                UpdateOptionChangedSource(_workspaceRegistration.Workspace);
+
+                _navigationBarController?.SetWorkspace(_workspaceRegistration.Workspace);
+
+                // Trigger a check to see if the dropdown should be added / removed now that the buffer is in a different workspace.
+                AddOrRemoveDropdown();
+            }
+
+            private void UpdateOptionChangedSource(Workspace? newWorkspace)
+            {
+                if (_optionService != null)
+                {
+                    _optionService.OptionChanged -= OnOptionChanged;
+                    _optionService = null;
+                }
+
+                var optionService = newWorkspace?.Services.GetService<IOptionService>();
+                if (optionService != null)
+                {
+                    _optionService = optionService;
+                    _optionService.OptionChanged += OnOptionChanged;
+                }
             }
 
             private void SetupView(IVsTextView view)
@@ -50,17 +107,22 @@ private void TeardownView(IVsTextView view)
 
             private void OnOptionChanged(object sender, OptionChangedEventArgs e)
             {
+                // If the workspace registration is missing, addornments have been removed.
+                if (_workspaceRegistration == null)
+                {
+                    return;
+                }
+
                 if (e.Language != _languageService.RoslynLanguageName ||
                     e.Option != NavigationBarOptions.ShowNavigationBar)
                 {
                     return;
                 }
 
-                var enabled = _optionService.GetOption(NavigationBarOptions.ShowNavigationBar, _languageService.RoslynLanguageName);
-                AddOrRemoveDropdown(enabled);
+                AddOrRemoveDropdown();
             }
 
-            private void AddOrRemoveDropdown(bool enabled)
+            private void AddOrRemoveDropdown()
             {
                 if (!(_codeWindow is IVsDropdownBarManager dropdownManager))
                 {
@@ -74,24 +136,29 @@ private void AddOrRemoveDropdown(bool enabled)
 
                 // Temporary solution until the editor provides a proper way to resolve the correct navbar.
                 // Tracked in https://github.com/dotnet/roslyn/issues/40989
-                var document = _languageService.EditorAdaptersFactoryService.GetDataBuffer(buffer).AsTextContainer().GetRelatedDocuments().FirstOrDefault();
-                if (document.GetLanguageService<INavigationBarItemService>() == null)
+                var document = _languageService.EditorAdaptersFactoryService.GetDataBuffer(buffer)?.AsTextContainer().GetRelatedDocuments().FirstOrDefault();
+                if (document?.GetLanguageService<INavigationBarItemService>() == null)
                 {
+                    // Remove the existing dropdown bar if it is ours.
+                    if (IsOurDropdownBar(dropdownManager, out var _))
+                    {
+                        RemoveDropdownBar(dropdownManager);
+                    }
+
                     return;
                 }
 
-                if (enabled)
+                var enabled = _optionService?.GetOption(NavigationBarOptions.ShowNavigationBar, _languageService.RoslynLanguageName);
+                if (enabled == true)
                 {
-                    var existingDropdownBar = GetDropdownBar(dropdownManager);
-                    if (existingDropdownBar != null)
+                    if (IsOurDropdownBar(dropdownManager, out var existingDropdownBar))
                     {
-                        // Check if the existing dropdown is already one of ours, and do nothing if it is.
-                        if (_dropdownBarClient != null &&
-                            _dropdownBarClient == GetDropdownBarClient(existingDropdownBar))
-                        {
-                            return;
-                        }
+                        // The dropdown bar is already one of ours, do nothing.
+                        return;
+                    }
 
+                    if (existingDropdownBar != null)
+                    {
                         // Not ours, so remove the old one so that we can add ours.
                         RemoveDropdownBar(dropdownManager);
                     }
@@ -107,6 +174,21 @@ private void AddOrRemoveDropdown(bool enabled)
                 {
                     RemoveDropdownBar(dropdownManager);
                 }
+
+                bool IsOurDropdownBar(IVsDropdownBarManager dropdownBarManager, out IVsDropdownBar? existingDropdownBar)
+                {
+                    existingDropdownBar = GetDropdownBar(dropdownBarManager);
+                    if (existingDropdownBar != null)
+                    {
+                        if (_dropdownBarClient != null &&
+                            _dropdownBarClient == GetDropdownBarClient(existingDropdownBar))
+                        {
+                            return true;
+                        }
+                    }
+
+                    return false;
+                }
             }
 
             private static IVsDropdownBar GetDropdownBar(IVsDropdownBarManager dropdownManager)
@@ -132,6 +214,7 @@ private void AdddropdownBar(IVsDropdownBarManager dropdownManager)
                 var textBuffer = _languageService.EditorAdaptersFactoryService.GetDataBuffer(buffer);
                 var controllerFactoryService = _languageService.Package.ComponentModel.GetService<INavigationBarControllerFactoryService>();
                 var newController = controllerFactoryService.CreateController(navigationBarClient, textBuffer);
+                newController.SetWorkspace(_workspaceRegistration?.Workspace);
                 var hr = dropdownManager.AddDropdownBar(cCombos: 3, pClient: navigationBarClient);
 
                 if (ErrorHandler.Failed(hr))
@@ -174,8 +257,14 @@ public int AddAdornments()
                     SetupView(secondaryView);
                 }
 
-                var enabled = _optionService.GetOption(NavigationBarOptions.ShowNavigationBar, _languageService.RoslynLanguageName);
-                AddOrRemoveDropdown(enabled);
+                ErrorHandler.ThrowOnFailure(_codeWindow.GetBuffer(out var buffer));
+                var textContainer = _languageService.EditorAdaptersFactoryService.GetDataBuffer(buffer).AsTextContainer();
+                _workspaceRegistration = CodeAnalysis.Workspace.GetWorkspaceRegistration(textContainer);
+                _workspaceRegistration.WorkspaceChanged += OnWorkspaceRegistrationChanged;
+
+                UpdateOptionChangedSource(_workspaceRegistration.Workspace);
+
+                AddOrRemoveDropdown();
 
                 return VSConstants.S_OK;
             }
@@ -197,9 +286,23 @@ public int OnNewView(IVsTextView view)
             public int RemoveAdornments()
             {
                 _sink.Unadvise();
-                _optionService.OptionChanged -= OnOptionChanged;
 
-                AddOrRemoveDropdown(enabled: false);
+                if (_optionService != null)
+                {
+                    _optionService.OptionChanged -= OnOptionChanged;
+                    _optionService = null;
+                }
+
+                if (_workspaceRegistration != null)
+                {
+                    _workspaceRegistration.WorkspaceChanged -= OnWorkspaceRegistrationChanged;
+                    _workspaceRegistration = null;
+                }
+
+                if (_codeWindow is IVsDropdownBarManager dropdownManager)
+                {
+                    RemoveDropdownBar(dropdownManager);
+                }
 
                 return VSConstants.S_OK;
             }

src/Workspaces/Core/Portable/Workspace/WorkspaceRegistration.cs

@@ -2,6 +2,8 @@
 // The .NET Foundation licenses this file to you under the MIT license.
 // See the LICENSE file in the project root for more information.
 
+#nullable enable
+
 using System;
 
 namespace Microsoft.CodeAnalysis
@@ -14,17 +16,17 @@ internal WorkspaceRegistration()
         {
         }
 
-        public Workspace Workspace { get; private set; }
+        public Workspace? Workspace { get; private set; }
 
-        public event EventHandler WorkspaceChanged;
+        public event EventHandler? WorkspaceChanged;
 
-        internal void SetWorkspaceAndRaiseEvents(Workspace workspace)
+        internal void SetWorkspaceAndRaiseEvents(Workspace? workspace)
         {
             SetWorkspace(workspace);
             RaiseEvents();
         }
 
-        internal void SetWorkspace(Workspace workspace)
+        internal void SetWorkspace(Workspace? workspace)
         {
             Workspace = workspace;
         }