Skip to content

More changes to how we load gladstone extensions. #77987

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 3 commits into from
Apr 3, 2025
Merged

More changes to how we load gladstone extensions. #77987

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
16a0b29
CHange how locking works
CyrusNajmabadi Apr 3, 2025
ae8e77d
Simplify
CyrusNajmabadi Apr 3, 2025
e359cc2
Docs
CyrusNajmabadi Apr 3, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
243 changes: 123 additions & 120 deletions src/Features/Core/Portable/Extensions/ExtensionMessageHandlerService.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -47,24 +47,26 @@ internal sealed class ExtensionMessageHandlerService(
private readonly SolutionServices _solutionServices = solutionServices;
private readonly IExtensionMessageHandlerFactory _customMessageHandlerFactory = customMessageHandlerFactory;

// Core design: To make things lightweight, and to avoid locking, all work is computed and cached in simple
// immutable dictionaries. These dictionaries are populated on demand, but contain data that can be recomputed
// safely if missing.
/// <summary>
/// Lock for <see cref="_folderPathToExtensionFolder"/>, <see cref="_cachedDocumentHandlers"/>, and <see
/// cref="_cachedWorkspaceHandlers"/>.
/// </summary>
private readonly SemaphoreSlim _gate = new(initialCount: 1);

/// <summary>
/// Extensions assembly load contexts and loaded handlers, indexed by extension folder path.
/// </summary>
private ImmutableDictionary<string, AsyncLazy<IExtensionFolder>> _folderPathToExtensionFolder = ImmutableDictionary<string, AsyncLazy<IExtensionFolder>>.Empty;
private readonly Dictionary<string, AsyncLazy<ExtensionFolder>> _folderPathToExtensionFolder = new();

/// <summary>
/// Cached handlers of document-related messages, indexed by handler message name.
/// </summary>
private ImmutableDictionary<string, AsyncLazy<ImmutableArray<IExtensionMessageHandlerWrapper<Document>>>> _cachedDocumentHandlers = ImmutableDictionary<string, AsyncLazy<ImmutableArray<IExtensionMessageHandlerWrapper<Document>>>>.Empty;
private readonly Dictionary<string, AsyncLazy<ImmutableArray<IExtensionMessageHandlerWrapper<Document>>>> _cachedDocumentHandlers = new();

/// <summary>
/// Cached handlers of non-document-related messages, indexed by handler message name.
/// </summary>
private ImmutableDictionary<string, AsyncLazy<ImmutableArray<IExtensionMessageHandlerWrapper<Solution>>>> _cachedWorkspaceHandlers = ImmutableDictionary<string, AsyncLazy<ImmutableArray<IExtensionMessageHandlerWrapper<Solution>>>>.Empty;
private readonly Dictionary<string, AsyncLazy<ImmutableArray<IExtensionMessageHandlerWrapper<Solution>>>> _cachedWorkspaceHandlers = new();

private async ValueTask<TResult> ExecuteInRemoteOrCurrentProcessAsync<TResult>(
Solution? solution,
Expand Down Expand Up @@ -116,18 +118,22 @@ public async ValueTask<RegisterExtensionResponse> RegisterExtensionInCurrentProc
var assemblyFolderPath = Path.GetDirectoryName(assemblyFilePath)
?? throw new InvalidOperationException($"Unable to get the directory name for {assemblyFilePath}.");

var lazy = ImmutableInterlocked.GetOrAdd(
ref _folderPathToExtensionFolder,
assemblyFolderPath,
static (assemblyFolderPath, @this) => AsyncLazy.Create(
cancellationToken => ExtensionFolder.Create(@this, assemblyFolderPath, cancellationToken)),
this);
AsyncLazy<ExtensionFolder> lazyExtensionFolder;
using (await _gate.DisposableWaitAsync(cancellationToken).ConfigureAwait(false))
{
lazyExtensionFolder = _folderPathToExtensionFolder.GetOrAdd(
assemblyFolderPath,
static (assemblyFolderPath, @this) => AsyncLazy.Create(
cancellationToken => ExtensionFolder.Create(@this, assemblyFolderPath, cancellationToken)),
this);

var extensionFolder = await lazy.GetValueAsync(cancellationToken).ConfigureAwait(false);
var assemblyHandlers = await extensionFolder.RegisterAssemblyAsync(assemblyFilePath, cancellationToken).ConfigureAwait(false);
// After registering, clear out the cached handler names. They will be recomputed the next time we need them.
ClearCachedHandlers();
}

// After registering, clear out the cached handler names. They will be recomputed the next time we need them.
ClearCachedHandlers();
var extensionFolder = await lazyExtensionFolder.GetValueAsync(cancellationToken).ConfigureAwait(false);
var lazyAssemblyHandlers = extensionFolder.RegisterAssembly(assemblyFilePath);
var assemblyHandlers = await lazyAssemblyHandlers.GetValueAsync(cancellationToken).ConfigureAwait(false);

return new(
[.. assemblyHandlers.WorkspaceMessageHandlers.Keys],
Expand All @@ -152,17 +158,23 @@ private async ValueTask<VoidResult> UnregisterExtensionInCurrentProcessAsync(
var assemblyFolderPath = Path.GetDirectoryName(assemblyFilePath)
?? throw new InvalidOperationException($"Unable to get the directory name for {assemblyFilePath}.");

if (_folderPathToExtensionFolder.TryGetValue(assemblyFolderPath, out var lazy))
// Note: unregistering is slightly expensive as we do everything under a lock, to ensure that we have a
// consistent view of the world. This is fine as we don't expect this to be called very often.
using (await _gate.DisposableWaitAsync(cancellationToken).ConfigureAwait(false))
{
var extensionFolder = await lazy.GetValueAsync(cancellationToken).ConfigureAwait(false);
// Unregister this particular assembly file from teh assembly folder. If it was the last extension within
// this folder, we can remove the registration for the extension entirely.
if (await extensionFolder.UnregisterAssemblyAsync(assemblyFilePath, cancellationToken).ConfigureAwait(false))
_folderPathToExtensionFolder = _folderPathToExtensionFolder.Remove(assemblyFolderPath);
if (_folderPathToExtensionFolder.TryGetValue(assemblyFolderPath, out var lazyExtensionFolder))
{
var extensionFolder = await lazyExtensionFolder.GetValueAsync(cancellationToken).ConfigureAwait(false);
// Unregister this particular assembly file from teh assembly folder. If it was the last extension within
// this folder, we can remove the registration for the extension entirely.
if (extensionFolder.UnregisterAssembly(assemblyFilePath))
_folderPathToExtensionFolder.Remove(assemblyFolderPath);
}

// After unregistering, clear out the cached handler names. They will be recomputed the next time we need them.
ClearCachedHandlers();
}

// After unregistering, clear out the cached handler names. They will be recomputed the next time we need them.
ClearCachedHandlers();
return default;
}

Expand All @@ -177,15 +189,19 @@ await ExecuteInRemoteOrCurrentProcessAsync(

private ValueTask<VoidResult> ResetInCurrentProcessAsync()
{
_folderPathToExtensionFolder = ImmutableDictionary<string, AsyncLazy<IExtensionFolder>>.Empty;
ClearCachedHandlers();
return default;
lock (_gate)
{
_folderPathToExtensionFolder.Clear();
ClearCachedHandlers();
return default;
}
}

private void ClearCachedHandlers()
{
_cachedWorkspaceHandlers = ImmutableDictionary<string, AsyncLazy<ImmutableArray<IExtensionMessageHandlerWrapper<Solution>>>>.Empty;
_cachedDocumentHandlers = ImmutableDictionary<string, AsyncLazy<ImmutableArray<IExtensionMessageHandlerWrapper<Document>>>>.Empty;
Contract.ThrowIfTrue(!Monitor.IsEntered(_gate));
_cachedWorkspaceHandlers.Clear();
_cachedDocumentHandlers.Clear();
}

public async ValueTask<string> HandleExtensionWorkspaceMessageAsync(Solution solution, string messageName, string jsonMessage, CancellationToken cancellationToken)
Expand All @@ -210,18 +226,21 @@ public async ValueTask<string> HandleExtensionDocumentMessageAsync(Document docu

private async ValueTask<string> HandleExtensionMessageInCurrentProcessAsync<TArgument>(
TArgument executeArgument, bool isSolution, string messageName, string jsonMessage,
ImmutableDictionary<string, AsyncLazy<ImmutableArray<IExtensionMessageHandlerWrapper<TArgument>>>> cachedHandlers,
Dictionary<string, AsyncLazy<ImmutableArray<IExtensionMessageHandlerWrapper<TArgument>>>> cachedHandlers,
CancellationToken cancellationToken)
{
var lazy = ImmutableInterlocked.GetOrAdd(
ref cachedHandlers,
messageName,
static (messageName, arg) => AsyncLazy.Create(
static (arg, cancellationToken) => ComputeHandlersAsync<TArgument>(arg.@this, arg.messageName, arg.isSolution, cancellationToken),
(messageName, arg.@this, arg.isSolution)),
(@this: this, executeArgument, isSolution));

var handlers = await lazy.GetValueAsync(cancellationToken).ConfigureAwait(false);
AsyncLazy<ImmutableArray<IExtensionMessageHandlerWrapper<TArgument>>> lazyHandlers;
using (await _gate.DisposableWaitAsync(cancellationToken).ConfigureAwait(false))
{
lazyHandlers = cachedHandlers.GetOrAdd(
messageName,
static (messageName, arg) => AsyncLazy.Create(
static (arg, cancellationToken) => ComputeHandlersAsync<TArgument>(arg.@this, arg.messageName, arg.isSolution, cancellationToken),
(messageName, arg.@this, arg.isSolution)),
(messageName, @this: this, isSolution));
}

var handlers = await lazyHandlers.GetValueAsync(cancellationToken).ConfigureAwait(false);
if (handlers.Length == 0)
throw new InvalidOperationException($"No handler found for message {messageName}.");

Expand Down Expand Up @@ -269,58 +288,62 @@ private static async Task<ImmutableArray<IExtensionMessageHandlerWrapper<TResult
return result.ToImmutable();
}

private interface IExtensionFolder
private abstract class ExtensionFolder
{
ValueTask<AssemblyHandlers> RegisterAssemblyAsync(string assemblyFilePath, CancellationToken cancellationToken);

/// <summary>
/// Unregisters this assembly path from this extension folder. If this was the last registered path, then this
/// will return true so that this folder can be unloaded.
/// Mapping from assembly file path to the handlers it contains. Used as its own lock when mutating.
/// </summary>
ValueTask<bool> UnregisterAssemblyAsync(string assemblyFilePath, CancellationToken cancellationToken);
private readonly Dictionary<string, AsyncLazy<AssemblyHandlers>> _assemblyFilePathToHandlers = new();

ValueTask AddHandlersAsync<TResult>(string messageName, bool isSolution, ArrayBuilder<IExtensionMessageHandlerWrapper<TResult>> result, CancellationToken cancellationToken);
}
protected abstract AssemblyHandlers CreateAssemblyHandlers(string assemblyFilePath, CancellationToken cancellationToken);

/// <summary>
/// Trivial placeholder impl of <see cref="IExtensionFolder"/> when we fail for some reason to even process the
/// folder we are told contains extensions.
/// </summary>
private sealed class TrivialExtensionFolder : IExtensionFolder
{
public static readonly TrivialExtensionFolder Instance = new();
public AsyncLazy<AssemblyHandlers> RegisterAssembly(string assemblyFilePath)
{
lock (_assemblyFilePathToHandlers)
{
return _assemblyFilePathToHandlers.GetOrAdd(
assemblyFilePath,
static (assemblyFilePath, @this) => AsyncLazy.Create(
static (args, cancellationToken) => args.@this.CreateAssemblyHandlers(args.assemblyFilePath, cancellationToken),
(assemblyFilePath, @this)),
this);
}
}

/// <summary>
/// No lock needed as registratin/unregistration must happen serially.
/// Unregisters this assembly path from this extension folder. If this was the last registered path, then this
/// will return true so that this folder can be unloaded.
/// </summary>
private readonly List<string> _registeredFilePaths = [];

public ValueTask<AssemblyHandlers> RegisterAssemblyAsync(string assemblyFilePath, CancellationToken cancellationToken)
public bool UnregisterAssembly(string assemblyFilePath)
{
_registeredFilePaths.Add(assemblyFilePath);
return new(AssemblyHandlers.Empty);
lock (_assemblyFilePathToHandlers)
{
_assemblyFilePathToHandlers.Remove(assemblyFilePath);
return _assemblyFilePathToHandlers.Count == 0;
}
}

public ValueTask<bool> UnregisterAssemblyAsync(string assemblyFilePath, CancellationToken cancellationToken)
public async ValueTask AddHandlersAsync<TResult>(string messageName, bool isSolution, ArrayBuilder<IExtensionMessageHandlerWrapper<TResult>> result, CancellationToken cancellationToken)
{
_registeredFilePaths.Remove(assemblyFilePath);
return new(_registeredFilePaths.Count == 0);
}

public ValueTask AddHandlersAsync<TResult>(string messageName, bool isSolution, ArrayBuilder<IExtensionMessageHandlerWrapper<TResult>> result, CancellationToken cancellationToken)
=> default;
}

private sealed class ExtensionFolder(
ExtensionMessageHandlerService extensionMessageHandlerService,
IAnalyzerAssemblyLoaderInternal analyzerAssemblyLoader) : IExtensionFolder
{
private readonly ExtensionMessageHandlerService _extensionMessageHandlerService = extensionMessageHandlerService;
private readonly IAnalyzerAssemblyLoaderInternal _analyzerAssemblyLoader = analyzerAssemblyLoader;
foreach (var (_, lazyHandlers) in _assemblyFilePathToHandlers)
{
cancellationToken.ThrowIfCancellationRequested();

private ImmutableDictionary<string, AsyncLazy<AssemblyHandlers>> _assemblyFilePathToHandlers = ImmutableDictionary<string, AsyncLazy<AssemblyHandlers>>.Empty;
var handlers = await lazyHandlers.GetValueAsync(cancellationToken).ConfigureAwait(false);
if (isSolution)
{
if (handlers.WorkspaceMessageHandlers.TryGetValue(messageName, out var handler))
result.Add((IExtensionMessageHandlerWrapper<TResult>)handler);
}
else
{
if (handlers.DocumentMessageHandlers.TryGetValue(messageName, out var handler))
result.Add((IExtensionMessageHandlerWrapper<TResult>)handler);
}
}
}

public static IExtensionFolder Create(
public static ExtensionFolder Create(
ExtensionMessageHandlerService extensionMessageHandlerService,
string assemblyFolderPath,
CancellationToken cancellationToken)
Expand Down Expand Up @@ -348,36 +371,42 @@ public static IExtensionFolder Create(
analyzerAssemblyLoader.AddDependencyLocation(dll);
}

return new ExtensionFolder(extensionMessageHandlerService, analyzerAssemblyLoader);
return new ShadowCopyExtensionFolder(extensionMessageHandlerService, analyzerAssemblyLoader);
}
catch (Exception ex) when (FatalError.ReportAndCatch(ex, ErrorSeverity.Critical))
{
// TODO: Log this exception so the client knows something went wrong.
return new TrivialExtensionFolder();
}
}
}

public async ValueTask<AssemblyHandlers> RegisterAssemblyAsync(
string assemblyFilePath, CancellationToken cancellationToken)
{
var lazy = ImmutableInterlocked.GetOrAdd(
ref _assemblyFilePathToHandlers,
assemblyFilePath,
static (assemblyFilePath, @this) => AsyncLazy.Create(
static (args, cancellationToken) => CreateAssemblyHandlers(args.@this, args.assemblyFilePath, cancellationToken),
(assemblyFilePath, @this)),
this);
/// <summary>
/// Trivial placeholder impl of <see cref="ExtensionFolder"/> when we fail for some reason to even process the
/// folder we are told contains extensions.
/// </summary>
private sealed class TrivialExtensionFolder : ExtensionFolder
{
protected override AssemblyHandlers CreateAssemblyHandlers(string assemblyFilePath, CancellationToken cancellationToken)
=> AssemblyHandlers.Empty;
}

return await lazy.GetValueAsync(cancellationToken).ConfigureAwait(false);
}
/// <summary>
/// Standard impl of <see cref="ExtensionFolder"/> that uses a shadow copy loader to load extensions.
/// </summary>
private sealed class ShadowCopyExtensionFolder(
ExtensionMessageHandlerService extensionMessageHandlerService,
IAnalyzerAssemblyLoaderInternal analyzerAssemblyLoader) : ExtensionFolder
{
private readonly ExtensionMessageHandlerService _extensionMessageHandlerService = extensionMessageHandlerService;
private readonly IAnalyzerAssemblyLoaderInternal _analyzerAssemblyLoader = analyzerAssemblyLoader;

private static AssemblyHandlers CreateAssemblyHandlers(
ExtensionFolder @this, string assemblyFilePath, CancellationToken cancellationToken)
protected override AssemblyHandlers CreateAssemblyHandlers(string assemblyFilePath, CancellationToken cancellationToken)
{
try
{
var assembly = @this._analyzerAssemblyLoader.LoadFromPath(assemblyFilePath);
var factory = @this._extensionMessageHandlerService._customMessageHandlerFactory;
var assembly = _analyzerAssemblyLoader.LoadFromPath(assemblyFilePath);
var factory = _extensionMessageHandlerService._customMessageHandlerFactory;

var messageWorkspaceHandlers = factory
.CreateWorkspaceMessageHandlers(assembly, extensionIdentifier: assemblyFilePath, cancellationToken)
Expand All @@ -401,32 +430,6 @@ private static AssemblyHandlers CreateAssemblyHandlers(
return AssemblyHandlers.Empty;
}
}

public async ValueTask AddHandlersAsync<TResult>(string messageName, bool isSolution, ArrayBuilder<IExtensionMessageHandlerWrapper<TResult>> result, CancellationToken cancellationToken)
{
foreach (var (_, lazy) in _assemblyFilePathToHandlers)
{
cancellationToken.ThrowIfCancellationRequested();

var handlers = await lazy.GetValueAsync(cancellationToken).ConfigureAwait(false);
if (isSolution)
{
if (handlers.WorkspaceMessageHandlers.TryGetValue(messageName, out var handler))
result.Add((IExtensionMessageHandlerWrapper<TResult>)handler);
}
else
{
if (handlers.DocumentMessageHandlers.TryGetValue(messageName, out var handler))
result.Add((IExtensionMessageHandlerWrapper<TResult>)handler);
}
}
}

public ValueTask<bool> UnregisterAssemblyAsync(string assemblyFilePath, CancellationToken cancellationToken)
{
_assemblyFilePathToHandlers = _assemblyFilePathToHandlers.Remove(assemblyFilePath);
return new(_assemblyFilePathToHandlers.IsEmpty);
}
}

private sealed class AssemblyHandlers
Expand Down
Loading