diff --git a/RemoteExec.Server/Configuration/ExecutionConfiguration.cs b/RemoteExec.Server/Configuration/ExecutionConfiguration.cs index f188205..f7a9e3e 100644 --- a/RemoteExec.Server/Configuration/ExecutionConfiguration.cs +++ b/RemoteExec.Server/Configuration/ExecutionConfiguration.cs @@ -11,4 +11,9 @@ public class ExecutionConfiguration /// Timeout for assembly loading requests in seconds. Default is 30 seconds. /// public int AssemblyLoadTimeoutSeconds { get; set; } = 30; + + /// + /// Type of execution environment to use. Default is "AssemblyLoadContext". + /// + public string ExecutionEnvironment { get; set; } = "AssemblyLoadContext"; } \ No newline at end of file diff --git a/RemoteExec.Server/Hubs/RemoteExecutionHub.cs b/RemoteExec.Server/Hubs/RemoteExecutionHub.cs index ee796ad..6555e1d 100644 --- a/RemoteExec.Server/Hubs/RemoteExecutionHub.cs +++ b/RemoteExec.Server/Hubs/RemoteExecutionHub.cs @@ -2,12 +2,11 @@ using Microsoft.Extensions.Options; using RemoteExec.Server.Configuration; +using RemoteExec.Server.Services; using RemoteExec.Shared; using System.Collections.Concurrent; using System.Diagnostics; -using System.Reflection; -using System.Text.Json; namespace RemoteExec.Server.Hubs; @@ -16,12 +15,12 @@ namespace RemoteExec.Server.Hubs; /// public class RemoteExecutionHub : Hub { - private static readonly ConcurrentDictionary connections = new(); + private static readonly ConcurrentDictionary connections = new(); private static readonly ConcurrentDictionary> pendingAssemblyRequests = new(); // Track pending assembly requests per connection to avoid duplicate requests - private static readonly ConcurrentDictionary>>> pendingAssemblyRequestsByConnection = new(); + private static readonly ConcurrentDictionary>>> pendingAssemblyRequestsByConnection = new(); private static ServerMetrics? lastMetrics; private static DateTime lastMetricsTimestamp; @@ -31,11 +30,13 @@ public class RemoteExecutionHub : Hub private static int maxConcurrentTasks; private static SemaphoreSlim taskSemaphore = null!; + private static string? executionEnvironmentName; private static int assemblyLoadTimeoutSeconds; private static double cpuDifferenceThreshold; private static long memoryDifferenceThreshold; private readonly ILogger logger; + private readonly IEnumerable executionEnvironments; /// /// Initializes a new instance of the class. @@ -43,15 +44,18 @@ public class RemoteExecutionHub : Hub /// The logger instance. /// The execution configuration options. /// The metrics configuration options. - public RemoteExecutionHub(ILogger logger, IOptions executionOptions, IOptions metricsOptions) + /// The available execution environments. + public RemoteExecutionHub(ILogger logger, IOptions executionOptions, IOptions metricsOptions, IEnumerable executionEnvironments) { this.logger = logger; + this.executionEnvironments = executionEnvironments; // Initialize static configuration values once if (taskSemaphore is null) { maxConcurrentTasks = executionOptions.Value.MaxConcurrentTasks ?? (Environment.ProcessorCount * 2); taskSemaphore = new SemaphoreSlim(maxConcurrentTasks, maxConcurrentTasks); + executionEnvironmentName = executionOptions.Value.ExecutionEnvironment; assemblyLoadTimeoutSeconds = executionOptions.Value.AssemblyLoadTimeoutSeconds; cpuDifferenceThreshold = metricsOptions.Value.CpuDifferenceThreshold; memoryDifferenceThreshold = metricsOptions.Value.MemoryDifferenceThreshold; @@ -59,30 +63,44 @@ public class RemoteExecutionHub : Hub } /// - public override Task OnConnectedAsync() + public override async Task OnConnectedAsync() { - RemoteJobAssemblyLoadContext assemblyLoadContext = new RemoteJobAssemblyLoadContext($"RemoteJob_{Guid.NewGuid()}"); + ExecutionEnvironment? executionEnvironment = executionEnvironments.FirstOrDefault(env => env.Name.Equals(executionEnvironmentName, StringComparison.OrdinalIgnoreCase)); - _ = connections.TryAdd(Context.ConnectionId, assemblyLoadContext); - _ = pendingAssemblyRequestsByConnection.TryAdd(Context.ConnectionId, new ConcurrentDictionary>>()); + if (executionEnvironment is null) + { + logger.LogError("Execution environment '{ExecutionEnvironment}' not found for connection {ConnectionId}", executionEnvironmentName, Context.ConnectionId); + throw new InvalidOperationException($"Execution environment '{executionEnvironmentName}' not found"); + } + + // Capture Context and Clients to avoid accessing disposed Hub instance + HubCallerContext capturedContext = Context; + IHubCallerClients capturedClients = Clients; + + executionEnvironment.RequestAssembly += async (sender, e) => + { + byte[] assemblyBytes = await RequestAssemblyBytesAsync(e.Value, capturedContext, capturedClients); + e.SetCompleted(assemblyBytes); + }; + + await executionEnvironment.PrepareEnvironmentAsync(Context.ConnectionAborted); + + _ = connections.TryAdd(Context.ConnectionId, executionEnvironment); + _ = pendingAssemblyRequestsByConnection.TryAdd(Context.ConnectionId, new()); logger.LogInformation("Connection {ConnectionId} established", Context.ConnectionId); - - return base.OnConnectedAsync(); } /// - public override Task OnDisconnectedAsync(Exception? exception) + public override async Task OnDisconnectedAsync(Exception? exception) { - if (connections.TryRemove(Context.ConnectionId, out RemoteJobAssemblyLoadContext? assemblyLoadContext)) + if (connections.TryRemove(Context.ConnectionId, out ExecutionEnvironment? executionEnvironment)) { - assemblyLoadContext.Unload(); + await executionEnvironment.CleanupEnvironmentAsync(CancellationToken.None); logger.LogInformation("Connection {ConnectionId} disconnected", Context.ConnectionId); } _ = pendingAssemblyRequestsByConnection.TryRemove(Context.ConnectionId, out _); - - return base.OnDisconnectedAsync(exception); } /// @@ -134,111 +152,23 @@ public class RemoteExecutionHub : Hub } } - /// - /// Executes a single remote method request. - /// - /// The execution request. - /// The execution result. - public async Task Execute(RemoteExecutionRequest req) - { - return await ExecuteTask(req); - } - - private async Task ExecuteTask(RemoteExecutionRequest req) + private async Task ExecuteTask(RemoteExecutionRequest request) { _ = Interlocked.Increment(ref activeTasks); try { - if (!connections.TryGetValue(Context.ConnectionId, out RemoteJobAssemblyLoadContext? assemblyLoadContext)) + if (!connections.TryGetValue(Context.ConnectionId, out ExecutionEnvironment? executionEnvironment)) { - logger.LogError("Connection {ConnectionId} not found", Context.ConnectionId); + logger.LogError("Connection {ConnectionId} not found for executing method {Method} in type {Type}", Context.ConnectionId, request.MethodName, request.TypeName); throw new InvalidOperationException("Connection not found"); } - // Check if assembly is already loaded in the context - Assembly? assembly = assemblyLoadContext.Assemblies.FirstOrDefault(a => a.GetName().FullName == req.AssemblyName); - - // If not loaded, request and load it into the context - assembly ??= await LoadAssemblyAsync(req.AssemblyName, assemblyLoadContext); - - Type type = assembly.GetType(req.TypeName, throwOnError: true)!; - - Type[] argTypes = req.ArgumentTypes - .Select(Type.GetType) - .ToArray()!; - - MethodInfo? method = type.GetMethod( - req.MethodName, - BindingFlags.Static | BindingFlags.Public | BindingFlags.NonPublic, - binder: null, - argTypes, - modifiers: null) ?? throw new MissingMethodException(req.TypeName, req.MethodName); - - // Pre-load all referenced assemblies to avoid triggering Resolving event during Invoke - await PreLoadReferencedAssembliesAsync(assemblyLoadContext, assembly); - - ParameterInfo[] parameters = method.GetParameters(); - - if (parameters.Length != req.Arguments.Length) - { - logger.LogError("Argument count mismatch for method {Method} in type {Type} for connection {ConnectionId}", req.MethodName, req.TypeName, Context.ConnectionId); - throw new ArgumentException("Argument count mismatch"); - } - - object?[] invokeArgs = new object?[req.Arguments.Length]; - - for (int i = 0; i < invokeArgs.Length; i++) - { - Type targetType = parameters[i].ParameterType; - object arg = req.Arguments[i]; - - if (arg is JsonElement je) - { - // Deserialize the JSON element into the expected CLR type - invokeArgs[i] = JsonSerializer.Deserialize(je.GetRawText(), targetType); - } - else if (arg == null) - { - invokeArgs[i] = null; - } - else if (!targetType.IsInstanceOfType(arg)) - { - // Fallback for simple primitive conversions - invokeArgs[i] = Convert.ChangeType(arg, targetType); - } - else - { - invokeArgs[i] = arg; - } - } - - object? result = method.Invoke(null, invokeArgs); - - if (result is Task taskResult) - { - await taskResult.ConfigureAwait(false); - Type returnType = method.ReturnType; - if (returnType.IsGenericType && returnType.GetGenericTypeDefinition() == typeof(Task<>)) - { - PropertyInfo resultProperty = returnType.GetProperty("Result")!; - result = resultProperty.GetValue(taskResult); - } - else - { - // For non-generic Task, result is null - result = null; - } - } - - return new RemoteExecutionResult - { - Result = result - }; + return await executionEnvironment.ExecuteTaskAsync(request); } catch (Exception ex) { - logger.LogError(ex, "Error executing remote method {Method} in type {Type} for connection {ConnectionId}", req.MethodName, req.TypeName, Context.ConnectionId); + logger.LogError(ex, "Error executing remote method {Method} in type {Type} for connection {ConnectionId}", request.MethodName, request.TypeName, Context.ConnectionId); return new RemoteExecutionResult { @@ -251,12 +181,7 @@ public class RemoteExecutionHub : Hub } } - /// - /// Provides assembly bytes to fulfill a pending assembly request. - /// - /// The unique identifier for the assembly request. - /// The assembly binary data. - public static async Task ProvideAssembly(Guid requestId, byte[] assemblyBytes) + internal static async Task ProvideAssembly(Guid requestId, byte[] assemblyBytes) { if (pendingAssemblyRequests.TryRemove(requestId, out TaskCompletionSource? tcs)) { @@ -328,20 +253,20 @@ public class RemoteExecutionHub : Hub }; } - private async Task LoadAssemblyAsync(string assemblyName, RemoteJobAssemblyLoadContext assemblyLoadContext) + private async Task RequestAssemblyBytesAsync(string assemblyName, HubCallerContext context, IHubCallerClients hubCallerClients) { try { - if (!pendingAssemblyRequestsByConnection.TryGetValue(Context.ConnectionId, out ConcurrentDictionary>>? connectionPendingRequests)) + if (!pendingAssemblyRequestsByConnection.TryGetValue(context.ConnectionId, out ConcurrentDictionary>>? connectionPendingRequests)) { throw new InvalidOperationException("Connection not found"); } // Use Lazy> pattern to ensure only one request is made // The Lazy.Value is only evaluated once, even if multiple threads access it simultaneously - Lazy> lazyTask = connectionPendingRequests.GetOrAdd(assemblyName, key => + Lazy> lazyTask = connectionPendingRequests.GetOrAdd(assemblyName, key => { - return new Lazy>(() => Task.Run(async () => + return new Lazy>(() => Task.Run(async () => { try { @@ -350,16 +275,21 @@ public class RemoteExecutionHub : Hub _ = pendingAssemblyRequests.TryAdd(guid, tcs); - await Clients.Caller.SendAsync("RequestAssembly", key, guid); + logger.LogInformation("Requesting assembly {Assembly} with RequestId {RequestId} for connection {ConnectionId}", key, guid, context.ConnectionId); - byte[] assemblyBytes = await tcs.Task.WaitAsync(TimeSpan.FromSeconds(assemblyLoadTimeoutSeconds)); + await hubCallerClients.Caller.SendAsync("RequestAssembly", key, guid); - using MemoryStream ms = new MemoryStream(assemblyBytes); - return assemblyLoadContext.LoadFromStream(ms); + return await tcs.Task.WaitAsync(TimeSpan.FromSeconds(assemblyLoadTimeoutSeconds)); } finally { - _ = connectionPendingRequests.TryRemove(key, out _); + _ = Task.Run(async () => + { + // WORKAROUND: Delay to avoid race condition where the same assembly is requested again because the assembly hasn't been loaded yet + // An alternative would be to only remove the request after the client disconnects or after a longer timeout + await Task.Delay(1000); + _ = connectionPendingRequests.TryRemove(key, out _); + }); } })); }); @@ -373,38 +303,4 @@ public class RemoteExecutionHub : Hub throw; } } - - private async Task PreLoadReferencedAssembliesAsync(RemoteJobAssemblyLoadContext assemblyLoadContext, Assembly assembly) - { - AssemblyName[] referencedAssemblies = assembly.GetReferencedAssemblies(); - - foreach (AssemblyName referencedAssembly in referencedAssemblies) - { - try - { - // Try to load from the assembly load context first - Assembly? loadedAssembly = assemblyLoadContext.Assemblies.FirstOrDefault(a => a.GetName().FullName == referencedAssembly.FullName); - - if (loadedAssembly != null) - { - continue; // Already loaded in the context - } - - // Try to load from default context (BCL assemblies) - try - { - _ = assemblyLoadContext.LoadFromAssemblyName(referencedAssembly); - continue; // Successfully loaded from default context - } - catch - { - _ = await LoadAssemblyAsync(referencedAssembly.FullName!, assemblyLoadContext); - } - } - catch (Exception ex) - { - logger.LogWarning(ex, "Could not pre-load referenced assembly {Assembly}", referencedAssembly.FullName); - } - } - } } \ No newline at end of file diff --git a/RemoteExec.Server/Program.cs b/RemoteExec.Server/Program.cs index 8500f3c..a90f54f 100644 --- a/RemoteExec.Server/Program.cs +++ b/RemoteExec.Server/Program.cs @@ -14,6 +14,7 @@ builder.Services.AddHealthChecks(); builder.Services.Configure(builder.Configuration.GetSection("Authentication")); builder.Services.Configure(builder.Configuration.GetSection("Execution")); builder.Services.Configure(builder.Configuration.GetSection("Metrics")); +builder.Services.AddScoped(); builder.Services.AddHostedService(); diff --git a/RemoteExec.Server/Services/AssemblyLoadContextExecutionEnvironment.cs b/RemoteExec.Server/Services/AssemblyLoadContextExecutionEnvironment.cs new file mode 100644 index 0000000..97eb921 --- /dev/null +++ b/RemoteExec.Server/Services/AssemblyLoadContextExecutionEnvironment.cs @@ -0,0 +1,115 @@ +using RemoteExec.Server.Utilities; +using RemoteExec.Shared; + +using System.Reflection; +using System.Text.Json; + +namespace RemoteExec.Server.Services; + +public class AssemblyLoadContextExecutionEnvironment : ExecutionEnvironment +{ + public override string Name => "AssemblyLoadContext"; + + private RemoteJobAssemblyLoadContext? assemblyLoadContext; + + public override async Task ExecuteTaskAsync(RemoteExecutionRequest request) + { + if (assemblyLoadContext is null) + { + throw new InvalidOperationException("The execution environment has not been prepared."); + } + + // Check if assembly is already loaded in the context + Assembly? assembly = assemblyLoadContext.Assemblies.FirstOrDefault(a => a.GetName().FullName == request.AssemblyName); + + // If not loaded, request and load it into the context + assembly ??= assemblyLoadContext.LoadFromBytes(await RequestAssemblyAsync(request.AssemblyName)); + + Type type = assembly.GetType(request.TypeName, throwOnError: true)!; + + Type[] argTypes = request.ArgumentTypes + .Select(Type.GetType) + .ToArray()!; + + MethodInfo? method = type.GetMethod( + request.MethodName, + BindingFlags.Static | BindingFlags.Public | BindingFlags.NonPublic, + binder: null, + argTypes, + modifiers: null) ?? throw new MissingMethodException(request.TypeName, request.MethodName); + + // Pre-load all referenced assemblies to avoid triggering Resolving event during Invoke + await AssemblyUtilities.PreLoadReferencedAssembliesAsync(assemblyLoadContext, assembly, RequestAssemblyAsync); + + ParameterInfo[] parameters = method.GetParameters(); + + if (parameters.Length != request.Arguments.Length) + { + throw new ArgumentException($"Argument count mismatch: expected {parameters.Length}, received {request.Arguments.Length}"); + } + + object?[] invokeArgs = new object?[request.Arguments.Length]; + + for (int i = 0; i < invokeArgs.Length; i++) + { + Type targetType = parameters[i].ParameterType; + object arg = request.Arguments[i]; + + if (arg is JsonElement je) + { + // Deserialize the JSON element into the expected CLR type + invokeArgs[i] = JsonSerializer.Deserialize(je.GetRawText(), targetType); + } + else if (arg == null) + { + invokeArgs[i] = null; + } + else if (!targetType.IsInstanceOfType(arg)) + { + // Fallback for simple primitive conversions + invokeArgs[i] = Convert.ChangeType(arg, targetType); + } + else + { + invokeArgs[i] = arg; + } + } + + object? result = method.Invoke(null, invokeArgs); + + if (result is Task taskResult) + { + await taskResult.ConfigureAwait(false); + Type returnType = method.ReturnType; + if (returnType.IsGenericType && returnType.GetGenericTypeDefinition() == typeof(Task<>)) + { + PropertyInfo resultProperty = returnType.GetProperty("Result")!; + result = resultProperty.GetValue(taskResult); + } + else + { + // For non-generic Task, result is null + result = null; + } + } + + return new RemoteExecutionResult + { + Result = result + }; + } + + public override Task PrepareEnvironmentAsync(CancellationToken cancellationToken) + { + assemblyLoadContext = new RemoteJobAssemblyLoadContext($"RemoteJob_{Guid.NewGuid()}"); + + return Task.CompletedTask; + } + + public override Task CleanupEnvironmentAsync(CancellationToken cancellationToken) + { + assemblyLoadContext?.Unload(); + assemblyLoadContext = null; + return Task.CompletedTask; + } +} diff --git a/RemoteExec.Server/Services/ExecutionEnvironment.cs b/RemoteExec.Server/Services/ExecutionEnvironment.cs new file mode 100644 index 0000000..41899df --- /dev/null +++ b/RemoteExec.Server/Services/ExecutionEnvironment.cs @@ -0,0 +1,24 @@ +using RemoteExec.Server.Utilities; +using RemoteExec.Shared; + +namespace RemoteExec.Server.Services; + +public abstract class ExecutionEnvironment +{ + public event EventHandler>? RequestAssembly; + + public abstract string Name { get; } + + public abstract Task PrepareEnvironmentAsync(CancellationToken cancellationToken); + + public abstract Task ExecuteTaskAsync(RemoteExecutionRequest request); + + public abstract Task CleanupEnvironmentAsync(CancellationToken cancellationToken); + + protected async Task RequestAssemblyAsync(string assemblyName) + { + CompletableEventArgs args = new CompletableEventArgs(assemblyName); + RequestAssembly?.Invoke(this, args); + return await args.WaitAsync(); + } +} diff --git a/RemoteExec.Server/Utilities/AssemblyUtilities.cs b/RemoteExec.Server/Utilities/AssemblyUtilities.cs new file mode 100644 index 0000000..c9b792b --- /dev/null +++ b/RemoteExec.Server/Utilities/AssemblyUtilities.cs @@ -0,0 +1,40 @@ +using System.Reflection; + +namespace RemoteExec.Server.Utilities; + +public static class AssemblyUtilities +{ + public static async Task PreLoadReferencedAssembliesAsync(RemoteJobAssemblyLoadContext assemblyLoadContext, Assembly assembly, Func> requestAssembly) + { + AssemblyName[] referencedAssemblies = assembly.GetReferencedAssemblies(); + + foreach (AssemblyName referencedAssembly in referencedAssemblies) + { + // Try to load from the assembly load context first + Assembly? loadedAssembly = assemblyLoadContext.Assemblies.FirstOrDefault(a => a.GetName().FullName == referencedAssembly.FullName); + + if (loadedAssembly != null) + { + continue; // Already loaded in the context + } + + // Try to load from default context (BCL assemblies) + try + { + _ = assemblyLoadContext.LoadFromAssemblyName(referencedAssembly); + continue; // Successfully loaded from default context + } + catch + { + byte[] assemblyBytes = await requestAssembly(referencedAssembly.FullName!); + _ = assemblyLoadContext.LoadFromBytes(assemblyBytes); + } + } + } + + public static Assembly LoadFromBytes(this RemoteJobAssemblyLoadContext assemblyLoadContext, byte[] assemblyBytes) + { + using MemoryStream ms = new(assemblyBytes); + return assemblyLoadContext.LoadFromStream(ms); + } +} diff --git a/RemoteExec.Server/Utilities/CompletableEventArgs.cs b/RemoteExec.Server/Utilities/CompletableEventArgs.cs new file mode 100644 index 0000000..4aa8422 --- /dev/null +++ b/RemoteExec.Server/Utilities/CompletableEventArgs.cs @@ -0,0 +1,57 @@ +namespace RemoteExec.Server.Utilities; + +public class CompletableEventArgs : EventArgs +{ + private readonly TaskCompletionSource tcs = new(TaskCreationOptions.RunContinuationsAsynchronously); + + public void SetCompleted() + { + _ = tcs.TrySetResult(true); + } + + public Task WaitAsync(CancellationToken cancellationToken = default) + { + if (cancellationToken.CanBeCanceled) + { + _ = cancellationToken.Register(() => tcs.TrySetCanceled(cancellationToken)); + } + return tcs.Task; + } +} + +public class CompletableEventArgs : EventArgs +{ + private readonly TaskCompletionSource tcs = new(TaskCreationOptions.RunContinuationsAsynchronously); + public void SetCompleted(T result) + { + _ = tcs.TrySetResult(result); + } + public Task WaitAsync(CancellationToken cancellationToken = default) + { + if (cancellationToken.CanBeCanceled) + { + _ = cancellationToken.Register(() => tcs.TrySetCanceled(cancellationToken)); + } + return tcs.Task; + } +} + +public class CompletableEventArgs(TValue value) +{ + public TValue Value { get; } = value; + + private readonly TaskCompletionSource tcs = new(TaskCreationOptions.RunContinuationsAsynchronously); + public void SetCompleted(TResult result) + { + _ = tcs.TrySetResult(result); + } + public Task WaitAsync(CancellationToken cancellationToken = default) + { + if (cancellationToken.CanBeCanceled) + { + _ = cancellationToken.Register(() => tcs.TrySetCanceled(cancellationToken)); + } + return tcs.Task; + } +} +