src/VisualStudio/Core/Def/ProjectSystem/OpenTextBufferProvider.cs - main - githubmirror/roslyn

src/VisualStudio/Core/Def/ProjectSystem/OpenTextBufferProvider.cs
347 строк16 KB
Jason Malinowski
Switch the fetch of the Running Document Table to the UI thread again (#81566)
09 дек 2025, 05:27
09 дек 2025, 05:276a49297
// Licensed to the .NET Foundation under one or more agreements.
// The .NET Foundation licenses this file to you under the MIT license.
// See the LICENSE file in the project root for more information.

using System;
using System.Collections.Generic;
using System.Collections.Immutable;
using System.ComponentModel.Composition;
using System.Diagnostics.CodeAnalysis;
using System.Threading.Tasks;
using Microsoft.CodeAnalysis.Editor.Shared.Extensions;
using Microsoft.CodeAnalysis.Editor.Shared.Utilities;
using Microsoft.CodeAnalysis.ErrorReporting;
using Microsoft.CodeAnalysis.Host.Mef;
using Microsoft.CodeAnalysis.PooledObjects;
using Microsoft.CodeAnalysis.Shared.TestHooks;
using Microsoft.VisualStudio.Editor;
using Microsoft.VisualStudio.Shell;
using Microsoft.VisualStudio.Shell.Interop;
using Microsoft.VisualStudio.Text;
using Microsoft.VisualStudio.Threading;

namespace Microsoft.VisualStudio.LanguageServices.Implementation.ProjectSystem;

/// <summary>
/// A class that provides access to the currently open list of files in Visual Studio.
/// </summary>
/// <remarks>
/// You are able to ask for the text buffer for a document on any thread; events are raised on the UI thread, and any method that provides a <see cref="IVsHierarchy"/> must be used on the UI thread.
/// Individual methods are documented for which threading contracts they expect.
/// </remarks>
[Export(typeof(OpenTextBufferProvider))]
internal sealed class OpenTextBufferProvider : IVsRunningDocTableEvents3, IDisposable
{
    private bool _isDisposed = false;
    private readonly IThreadingContext _threadingContext;

    /// <summary>
    /// A simple object for asserting when we're on the UI thread.
    /// </summary>
    private readonly IVsEditorAdaptersFactoryService _editorAdaptersFactoryService;
    private readonly Lazy<IVsRunningDocumentTable4> _runningDocumentTable;

    private ImmutableArray<IOpenTextBufferEventListener> _listeners = [];

    /// <summary>
    /// The map from monikers to open text buffers; because we can only fetch the text buffer on the UI thread, all updates to this must be done from the UI thread.
    /// </summary>
    private ImmutableDictionary<string, ITextBuffer> _monikerToTextBufferMap = ImmutableDictionary<string, ITextBuffer>.Empty.WithComparers(StringComparer.OrdinalIgnoreCase);

    private uint _runningDocumentTableEventsCookie;

    [ImportingConstructor]
    [Obsolete(MefConstruction.ImportingConstructorMessage, error: true)]
    public OpenTextBufferProvider(
        IThreadingContext threadingContext,
        IVsEditorAdaptersFactoryService editorAdaptersFactoryService,
        [Import(typeof(SVsServiceProvider))] IServiceProvider serviceProvider,
        IAsynchronousOperationListenerProvider listenerProvider)
    {
        _threadingContext = threadingContext;
        _editorAdaptersFactoryService = editorAdaptersFactoryService;

        _runningDocumentTable = new(() =>
        {
            // The hope was fetching the running document table and calling Advise() can be fully free-threaded, but at the moment (in 18.3) that's not entirely the case.
            // As of this writing, this Running Document Table object has an internal managed code implementation; when we call Advise() that native implementation forwards
            // to the managed implementation, and might send along not only our listener, but any other listeners that had been registered with the native side but not the managed side.
            // In this case, the other handlers get a QueryInterface called, which is expected to be on the UI thread. Until this is fixed, we'll do this on the UI thread.
            // See https://devdiv.visualstudio.com/DevDiv/_workitems/edit/2578713 for the investigation.

            _threadingContext.ThrowIfNotOnUIThread();
            var runningDocumentTable = (IVsRunningDocumentTable)serviceProvider.GetService(typeof(SVsRunningDocumentTable));
            runningDocumentTable.AdviseRunningDocTableEvents(this, out _runningDocumentTableEventsCookie);

            return (IVsRunningDocumentTable4)runningDocumentTable;
        });

        // We also need to check for any documents that might have been open before we subscribed. That we do have to do on the UI thread.
        var listener = listenerProvider.GetListener(FeatureAttribute.Workspace);
        var asyncToken = listener.BeginAsyncOperation(nameof(CheckForExistingOpenDocumentsAsync));

        CheckForExistingOpenDocumentsAsync().CompletesAsyncOperation(asyncToken);
    }

    private void RaiseEventForEachListener(Action<IOpenTextBufferEventListener> action)
    {
        _threadingContext.ThrowIfNotOnUIThread();

        foreach (var listener in _listeners)
        {
            try
            {
                action(listener);
            }
            catch (Exception e) when (FatalError.ReportAndCatch(e, ErrorSeverity.Critical))
            {
                // We'll catch the exception; this way if one listener is broken, we don't end up breaking other features that might no longer get events. Any exceptions would get caught by the
                // RunningDocumentTable itself which wouldn't report them in a useful way regardless.
            }
        }
    }

    private async Task CheckForExistingOpenDocumentsAsync()
    {
        // Yield the thread, so the calling constructor can proceed immediately
        await Task.Yield().ConfigureAwait(false);

        // Ensure we obtain the RDT before transitioning to the main thread
        // TODO: add this back once we move the RDT acquisition and event advising off the UI thread again.
        // _ = _runningDocumentTable.Value;

        // Now switch to the main thread since document enumeration does require it
        await _threadingContext.JoinableTaskFactory.SwitchToMainThreadAsync();

        foreach (var (filePath, textBuffer, hierarchy) in EnumerateDocumentSet())
        {
            // We might or might not have seen this file be opened if it was opened between when we subscribed to the running document table and when
            // we got scheduled to the UI thread.
            if (!_monikerToTextBufferMap.ContainsKey(filePath))
            {
                _monikerToTextBufferMap = _monikerToTextBufferMap.Add(filePath, textBuffer);
                RaiseEventForEachListener(l => l.OnOpenDocument(filePath, textBuffer, hierarchy));
            }
        }
    }

    public void AddListener(IOpenTextBufferEventListener listener) => ImmutableInterlocked.Update(ref _listeners, static (array, listener) => array.Add(listener), listener);
    public void RemoveListener(IOpenTextBufferEventListener listener) => ImmutableInterlocked.Update(ref _listeners, static (array, listener) => array.Remove(listener), listener);

    public int OnAfterFirstDocumentLock(uint docCookie, uint dwRDTLockType, uint dwReadLocksRemaining, uint dwEditLocksRemaining)
        => VSConstants.E_NOTIMPL;

    public int OnBeforeLastDocumentUnlock(uint docCookie, uint dwRDTLockType, uint dwReadLocksRemaining, uint dwEditLocksRemaining)
    {
        if (dwReadLocksRemaining + dwEditLocksRemaining == 0)
        {
            _threadingContext.ThrowIfNotOnUIThread();
            if (_runningDocumentTable.Value.IsDocumentInitialized(docCookie))
            {
                var moniker = _runningDocumentTable.Value.GetDocumentMoniker(docCookie);
                _monikerToTextBufferMap = _monikerToTextBufferMap.Remove(moniker);

                RaiseEventForEachListener(l => l.OnCloseDocument(moniker));
            }
        }

        return VSConstants.S_OK;
    }

    public int OnAfterSave(uint docCookie)
    {
        if (_runningDocumentTable.Value.IsDocumentInitialized(docCookie))
        {
            var moniker = _runningDocumentTable.Value.GetDocumentMoniker(docCookie);
            RaiseEventForEachListener(l => l.OnSaveDocument(moniker));
        }

        return VSConstants.S_OK;
    }

    public int OnAfterAttributeChange(uint docCookie, uint grfAttribs)
        => VSConstants.E_NOTIMPL;

    public int OnAfterAttributeChangeEx(uint docCookie, uint grfAttribs, IVsHierarchy pHierOld, uint itemidOld, string pszMkDocumentOld, IVsHierarchy pHierNew, uint itemidNew, string pszMkDocumentNew)
    {
        // Did we rename?
        if ((grfAttribs & (uint)__VSRDTATTRIB.RDTA_MkDocument) != 0)
        {
            _threadingContext.ThrowIfNotOnUIThread();
            if (_runningDocumentTable.Value.IsDocumentInitialized(docCookie))
            {
                // We should already have a text buffer for this one
                if (_monikerToTextBufferMap.TryGetValue(pszMkDocumentOld, out var textBuffer))
                {
                    _monikerToTextBufferMap = _monikerToTextBufferMap.Remove(pszMkDocumentOld).Add(pszMkDocumentNew, textBuffer);
                }
                else
                {
                    // Odd we don't have one, but fetch it now
                    if (TryGetBufferFromRunningDocumentTable(docCookie, out textBuffer))
                    {
                        _monikerToTextBufferMap = _monikerToTextBufferMap.Add(pszMkDocumentNew, textBuffer);
                    }
                }

                // Only raise an event if we had a text buffer; otherwise this is a rename of something else and we don't need to report it
                if (textBuffer != null)
                {
                    RaiseEventForEachListener(l => l.OnRenameDocument(newMoniker: pszMkDocumentNew, oldMoniker: pszMkDocumentOld, textBuffer: textBuffer));
                }
            }
        }

        // Either RDTA_DocDataReloaded or RDTA_DocumentInitialized will be triggered if there's a lazy load and the document is now available.
        // See https://devdiv.visualstudio.com/DevDiv/_workitems/edit/937712 for a scenario where we do need the RDTA_DocumentInitialized check.
        // We still check for RDTA_DocDataReloaded because the RDT will mark something as initialized as soon as there is something in the doc data,
        // but that might still not be associated with an ITextBuffer.
        if ((grfAttribs & ((uint)__VSRDTATTRIB.RDTA_DocDataReloaded | (uint)__VSRDTATTRIB3.RDTA_DocumentInitialized)) != 0)
        {
            _threadingContext.ThrowIfNotOnUIThread();
            if (_runningDocumentTable.Value.IsDocumentInitialized(docCookie) && TryGetMoniker(docCookie, out var moniker) && TryGetBufferFromRunningDocumentTable(docCookie, out var buffer))
            {
                _monikerToTextBufferMap = _monikerToTextBufferMap.Add(moniker, buffer);
                _runningDocumentTable.Value.GetDocumentHierarchyItem(docCookie, out var hierarchy, out _);

                RaiseEventForEachListener(l => l.OnOpenDocument(moniker, buffer, hierarchy));
            }
        }

        if ((grfAttribs & (uint)__VSRDTATTRIB.RDTA_Hierarchy) != 0)
        {
            _threadingContext.ThrowIfNotOnUIThread();
            if (_runningDocumentTable.Value.IsDocumentInitialized(docCookie) && TryGetMoniker(docCookie, out var moniker))
            {
                _runningDocumentTable.Value.GetDocumentHierarchyItem(docCookie, out var hierarchy, out _);

                RaiseEventForEachListener(l => l.OnRefreshDocumentContext(moniker, hierarchy));
            }
        }

        return VSConstants.S_OK;
    }

    public int OnBeforeDocumentWindowShow(uint docCookie, int fFirstShow, IVsWindowFrame pFrame)
    {
        // Doc data reloaded is not triggered for the underlying aspx.cs file when changes are made to the aspx file, so catch it here.
        if (fFirstShow != 0 && _runningDocumentTable.Value.IsDocumentInitialized(docCookie) && TryGetMoniker(docCookie, out var moniker))
        {
            // If we hadn't already raised an event for this, do it now
            if (!_monikerToTextBufferMap.ContainsKey(moniker) && TryGetBufferFromRunningDocumentTable(docCookie, out var buffer))
            {
                _monikerToTextBufferMap = _monikerToTextBufferMap.Add(moniker, buffer);
                _runningDocumentTable.Value.GetDocumentHierarchyItem(docCookie, out var hierarchy, out _);

                RaiseEventForEachListener(l => l.OnOpenDocument(moniker, buffer, hierarchy));
            }

            RaiseEventForEachListener(l => l.OnDocumentOpenedIntoWindowFrame(moniker, pFrame));
        }

        return VSConstants.S_OK;
    }

    public int OnAfterDocumentWindowHide(uint docCookie, IVsWindowFrame pFrame)
        => VSConstants.E_NOTIMPL;

    public int OnBeforeSave(uint docCookie)
        => VSConstants.E_NOTIMPL;

    /// <summary>
    /// Attempts to get a text buffer from the specified file path. May be called on any thread.
    /// </summary>
    /// <param name="filePath">The file path to retrieve the text buffer for.</param>
    /// <param name="textBuffer">The buffer if the file is open and initialized.</param>
    /// <returns>true if the buffer was found with a non null value.</returns>
    public bool TryGetBufferFromFilePath(string filePath, [NotNullWhen(true)] out ITextBuffer? textBuffer)
    {
        return _monikerToTextBufferMap.TryGetValue(filePath, out textBuffer);
    }

    /// <summary>
    /// Checks if a file is open. May be called on any thread.
    /// </summary>
    public bool IsFileOpen(string filePath)
    {
        return _monikerToTextBufferMap.ContainsKey(filePath);
    }

    /// <summary>
    /// Fetches the <see cref="IVsHierarchy"/> for a document. Must be called on the UI thread.
    /// </summary>
    public IVsHierarchy? GetDocumentHierarchy(string filePath)
    {
        _threadingContext.ThrowIfNotOnUIThread();

        if (!_runningDocumentTable.Value.IsFileOpen(filePath))
        {
            return null;
        }

        var cookie = _runningDocumentTable.Value.GetDocumentCookie(filePath);
        _runningDocumentTable.Value.GetDocumentHierarchyItem(cookie, out var hierarchy, out _);
        return hierarchy;
    }

    /// <summary>
    /// Enumerates the running document table to retrieve all initialized files. Must be called on the UI thread, since this returns <see cref="IVsHierarchy"/> objects.
    /// </summary>
    public IEnumerable<(string filePath, ITextBuffer textBuffer, IVsHierarchy hierarchy)> EnumerateDocumentSet()
    {
        _threadingContext.ThrowIfNotOnUIThread();

        var documents = ArrayBuilder<(string, ITextBuffer, IVsHierarchy)>.GetInstance();
        foreach (var cookie in GetInitializedRunningDocumentTableCookies())
        {
            if (TryGetMoniker(cookie, out var moniker) && TryGetBufferFromRunningDocumentTable(cookie, out var buffer))
            {
                _runningDocumentTable.Value.GetDocumentHierarchyItem(cookie, out var hierarchy, out _);
                documents.Add((moniker, buffer, hierarchy));
            }
        }

        return documents.ToArray();
    }

    private IEnumerable<uint> GetInitializedRunningDocumentTableCookies()
    {
        foreach (var cookie in _runningDocumentTable.Value.GetRunningDocuments())
        {
            if (_runningDocumentTable.Value.IsDocumentInitialized(cookie))
            {
                yield return cookie;
            }
        }
    }

    private bool TryGetMoniker(uint docCookie, out string moniker)
    {
        moniker = _runningDocumentTable.Value.GetDocumentMoniker(docCookie);
        return !string.IsNullOrEmpty(moniker);
    }

    private bool TryGetBufferFromRunningDocumentTable(uint docCookie, [NotNullWhen(true)] out ITextBuffer? textBuffer)
    {
        _threadingContext.ThrowIfNotOnUIThread();
        return _runningDocumentTable.Value.TryGetBuffer(_editorAdaptersFactoryService, docCookie, out textBuffer);
    }

    public void Dispose()
    {
        if (_isDisposed)
        {
            return;
        }

        // Unsubscribe from the RDT, assuming we ever subscribed in the first case
        if (_runningDocumentTable.IsValueCreated)
        {
            var runningDocumentTableForEvents = (IVsRunningDocumentTable)_runningDocumentTable.Value;
            runningDocumentTableForEvents.UnadviseRunningDocTableEvents(_runningDocumentTableEventsCookie);
            _runningDocumentTableEventsCookie = 0;
        }

        _isDisposed = true;
    }
}