git.stella-ops.org/src/Plugin/StellaOps.Plugin.Host/Loading/AssemblyPluginLoader.cs

using Microsoft.Extensions.Logging;
using Microsoft.Extensions.Options;
using StellaOps.Plugin.Abstractions;
using StellaOps.Plugin.Abstractions.Manifest;
using System.Collections.Concurrent;
using System.Reflection;

namespace StellaOps.Plugin.Host.Loading;

/// <summary>
/// Loads plugins from assemblies with optional isolation using AssemblyLoadContext.
/// </summary>
public sealed class AssemblyPluginLoader : IHostPluginLoader
{
    private readonly ConcurrentDictionary<string, PluginLoadContextReference> _loadContexts = new();
    private readonly ILogger<AssemblyPluginLoader> _logger;
    private readonly PluginHostOptions _options;

    /// <summary>
    /// Creates a new assembly plugin loader.
    /// </summary>
    /// <param name="options">Plugin host options.</param>
    /// <param name="logger">Logger instance.</param>
    public AssemblyPluginLoader(
        IOptions<PluginHostOptions> options,
        ILogger<AssemblyPluginLoader> logger)
    {
        _options = options.Value;
        _logger = logger;
    }

    /// <inheritdoc />
    public Task<PluginAssemblyLoadResult> LoadAsync(
        PluginManifest manifest,
        PluginTrustLevel trustLevel,
        CancellationToken ct)
    {
        var assemblyPath = ResolveAssemblyPath(manifest);
        var pluginId = manifest.Info.Id;

        _logger.LogDebug("Loading plugin assembly from {Path}", assemblyPath);

        // Determine if we should isolate this plugin
        var shouldIsolate = _options.EnableAssemblyIsolation && trustLevel != PluginTrustLevel.BuiltIn;

        PluginAssemblyLoadContext? loadContext = null;
        Assembly assembly;

        if (shouldIsolate)
        {
            // Create isolated load context
            loadContext = new PluginAssemblyLoadContext(
                pluginId,
                assemblyPath,
                isCollectible: true);

            _loadContexts[pluginId] = new PluginLoadContextReference(loadContext, pluginId);

            try
            {
                assembly = loadContext.LoadFromAssemblyPath(assemblyPath);
            }
            catch (Exception ex)
            {
                // Cleanup on failure
                _loadContexts.TryRemove(pluginId, out _);
                loadContext.Unload();
                throw new PluginLoadException(pluginId, $"Failed to load assembly: {ex.Message}", ex);
            }
        }
        else
        {
            // Load into default context (for built-in plugins)
            assembly = Assembly.LoadFrom(assemblyPath);
        }

        try
        {
            // Find the entry point type
            var entryPointType = assembly.GetType(manifest.EntryPoint)
                ?? throw new PluginLoadException(pluginId, $"Entry point type '{manifest.EntryPoint}' not found");

            // Verify it implements IPlugin
            if (!typeof(IPlugin).IsAssignableFrom(entryPointType))
                throw new PluginLoadException(pluginId, $"Entry point type '{manifest.EntryPoint}' does not implement IPlugin");

            // Create instance
            var instance = Activator.CreateInstance(entryPointType) as IPlugin
                ?? throw new PluginLoadException(pluginId, $"Failed to create instance of '{manifest.EntryPoint}'");

            _logger.LogDebug("Loaded plugin {PluginId} (isolated: {Isolated})", pluginId, shouldIsolate);

            return Task.FromResult(new PluginAssemblyLoadResult(instance, assembly, loadContext));
        }
        catch (Exception ex) when (ex is not PluginLoadException)
        {
            // Cleanup on failure
            if (loadContext != null)
            {
                _loadContexts.TryRemove(pluginId, out _);
                loadContext.Unload();
            }
            throw new PluginLoadException(pluginId, $"Failed to instantiate plugin: {ex.Message}", ex);
        }
    }

    /// <inheritdoc />
    public Task<PluginAssemblyLoadResult> LoadAsync<T>(
        PluginTrustLevel trustLevel,
        CancellationToken ct) where T : class, IPlugin, new()
    {
        var instance = new T();
        var assembly = typeof(T).Assembly;

        // Embedded plugins are loaded in the default context
        return Task.FromResult(new PluginAssemblyLoadResult(instance, assembly, null));
    }

    /// <inheritdoc />
    public async Task UnloadAsync(string pluginId, CancellationToken ct)
    {
        if (!_loadContexts.TryRemove(pluginId, out var contextRef))
        {
            _logger.LogDebug("Plugin {PluginId} was not isolated, no unload needed", pluginId);
            return;
        }

        contextRef.Unload();

        // Wait for GC to collect the assemblies
        const int maxAttempts = 10;
        const int delayMs = 100;

        for (int i = 0; i < maxAttempts && !contextRef.IsCollected; i++)
        {
            ct.ThrowIfCancellationRequested();

            GC.Collect();
            GC.WaitForPendingFinalizers();
            await Task.Delay(delayMs, ct);
        }

        if (!contextRef.IsCollected)
        {
            _logger.LogWarning(
                "Plugin {PluginId} load context still alive after unload - possible memory leak",
                pluginId);
        }
        else
        {
            _logger.LogDebug("Plugin {PluginId} unloaded and collected", pluginId);
        }
    }

    /// <inheritdoc />
    public bool IsLoaded(string pluginId)
    {
        return _loadContexts.ContainsKey(pluginId);
    }

    private static string ResolveAssemblyPath(PluginManifest manifest)
    {
        if (string.IsNullOrEmpty(manifest.AssemblyPath))
        {
            throw new PluginLoadException(
                manifest.Info.Id,
                "Assembly path not specified in manifest");
        }

        if (!File.Exists(manifest.AssemblyPath))
        {
            throw new PluginLoadException(
                manifest.Info.Id,
                $"Assembly not found: {manifest.AssemblyPath}");
        }

        return manifest.AssemblyPath;
    }
}

/// <summary>
/// Exception thrown when plugin loading fails.
/// </summary>
public class PluginLoadException : Exception
{
    /// <summary>
    /// The plugin ID that failed to load.
    /// </summary>
    public string PluginId { get; }

    /// <summary>
    /// Creates a new plugin load exception.
    /// </summary>
    /// <param name="pluginId">The plugin ID.</param>
    /// <param name="message">The error message.</param>
    public PluginLoadException(string pluginId, string message)
        : base($"Failed to load plugin '{pluginId}': {message}")
    {
        PluginId = pluginId;
    }

    /// <summary>
    /// Creates a new plugin load exception with inner exception.
    /// </summary>
    /// <param name="pluginId">The plugin ID.</param>
    /// <param name="message">The error message.</param>
    /// <param name="inner">The inner exception.</param>
    public PluginLoadException(string pluginId, string message, Exception inner)
        : base($"Failed to load plugin '{pluginId}': {message}", inner)
    {
        PluginId = pluginId;
    }
}