diff --git a/RemoteExec.Client/AsyncManualResetEvent.cs b/RemoteExec.Client/AsyncManualResetEvent.cs new file mode 100644 index 0000000..aff7d68 --- /dev/null +++ b/RemoteExec.Client/AsyncManualResetEvent.cs @@ -0,0 +1,37 @@ +namespace RemoteExec.Client; + +internal class AsyncManualResetEvent +{ + private volatile TaskCompletionSource _tcs = new(TaskCreationOptions.RunContinuationsAsynchronously); + + public AsyncManualResetEvent(bool initialState) + { + if (initialState) + { + _ = _tcs.TrySetResult(true); + } + } + + public Task WaitAsync(CancellationToken ct) + { + return _tcs.Task.WaitAsync(ct); + } + + public void Set() + { + _ = _tcs.TrySetResult(true); + } + + public void Reset() + { + while (true) + { + TaskCompletionSource tcs = _tcs; + if (!tcs.Task.IsCompleted || + Interlocked.CompareExchange(ref _tcs, new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously), tcs) == tcs) + { + return; + } + } + } +} \ No newline at end of file diff --git a/RemoteExec.Client/RemoteExecutor.ConnectionManagement.cs b/RemoteExec.Client/RemoteExecutor.ConnectionManagement.cs new file mode 100644 index 0000000..f20229e --- /dev/null +++ b/RemoteExec.Client/RemoteExecutor.ConnectionManagement.cs @@ -0,0 +1,167 @@ +using Microsoft.AspNetCore.SignalR.Client; +using Microsoft.Extensions.Logging; + +using RemoteExec.Shared; + +using System.Collections.Concurrent; + +namespace RemoteExec.Client; + +public partial class RemoteExecutor +{ + public async Task StartAsync(CancellationToken cancellationToken = default) + { + // Dispose previous cancellation token source if exists + if (distributorCts != null && !distributorCts.IsCancellationRequested) + { + await distributorCts.CancelAsync(); + distributorCts.Dispose(); + } + + distributorCts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken); + + List startTasks = []; + + foreach (ServerConnection server in servers) + { + RegisterEventHandlers(server); + startTasks.Add(StartServerConnectionAsync(server, distributorCts.Token)); + } + + await Task.WhenAll(startTasks); + + // Check if at least one server is connected + if (!servers.Any(s => s.Connection.State == HubConnectionState.Connected)) + { + logger.LogWarning("No servers are currently connected. Tasks will be queued until a server becomes available."); + } + + distributorTask = Task.Run(() => DistributorLoop(distributorCts.Token), distributorCts.Token); + } + + public async Task StopAsync(CancellationToken cancellationToken = default) + { + logger.LogInformation("Stopping RemoteExecutor..."); + + await distributorCts.CancelAsync(); + + if (distributorTask != null) + { + try + { + await distributorTask; + logger.LogDebug("Distributor task completed successfully"); + } + catch (OperationCanceledException ex) + { + logger.LogError(ex, "Distributor task was canceled"); + } + } + + logger.LogDebug("Completing task channels for {ServerCount} servers", servers.Count); + foreach (ServerConnection server in servers) + { + server.TaskChannel.Writer.Complete(); + } + + List stopTasks = []; + + foreach (ServerConnection server in servers) + { + stopTasks.Add(server.Connection.StopAsync(cancellationToken)); + } + + await Task.WhenAll(stopTasks); + logger.LogInformation("RemoteExecutor stopped successfully"); + } + + private async Task StartServerConnectionAsync(ServerConnection server, CancellationToken cancellationToken) + { + try + { + await InitializeServerAsync(server, cancellationToken); + } + catch (Exception ex) + { + logger.LogWarning(ex, "Failed to connect to server: {ServerUrl}. Will retry automatically.", server.HttpClient.BaseAddress); + + // Start background reconnection attempts + _ = Task.Run(async () => await AttemptReconnectionAsync(server, cancellationToken), cancellationToken); + } + } + + private async Task AttemptReconnectionAsync(ServerConnection server, CancellationToken cancellationToken) + { + int retryDelay = 5000; // Start with 5 seconds + int maxRetryDelay = 60000; // Max 1 minute + + while (!cancellationToken.IsCancellationRequested && server.Connection.State == HubConnectionState.Disconnected) + { + try + { + await Task.Delay(retryDelay, cancellationToken); + logger.LogInformation("Attempting to reconnect to server: {ServerUrl}", server.HttpClient.BaseAddress); + + await InitializeServerAsync(server, cancellationToken); + + break; // Successfully connected + } + catch (OperationCanceledException) + { + break; + } + catch (Exception ex) + { + logger.LogWarning(ex, "Reconnection attempt failed for server: {ServerUrl}. Retrying in {Delay}ms", server.HttpClient.BaseAddress, retryDelay); + + // Exponential backoff + retryDelay = Math.Min(retryDelay * 2, maxRetryDelay); + } + } + } + + private async Task InitializeServerAsync(ServerConnection server, CancellationToken cancellationToken) + { + await server.Connection.StartAsync(cancellationToken); + logger.LogInformation("Connected to server: {ServerUrl}", server.HttpClient.BaseAddress); + + // Fetch initial metrics + server.Metrics = await server.Connection.InvokeAsync("GetMetrics", cancellationToken); + MetricsUpdated?.Invoke(this, new ServerMetricsUpdatedEventArgs(server.Connection, server.Metrics)); + + // Start task stream + await server.Connection.SendAsync("StartTaskStream", server.TaskChannel.Reader, cancellationToken); + + serverAvailableSignal.Set(); + } + + private async Task RequeuePendingTasksForServer(ServerConnection server) + { + while (server.TaskChannel.Reader.TryRead(out _)) { /* Discard stale items */ } + + if (!serverAssignedTasks.TryGetValue(server, out ConcurrentDictionary? taskDict)) + { + return; + } + + List taskIds = taskDict.Keys.ToList(); + int count = 0; + + foreach (Guid id in taskIds) + { + if (taskDict.TryRemove(id, out PendingTask? task) && pendingResults.ContainsKey(id)) + { + globalQueue.Add(task); + count++; + } + } + + // If no servers are left, block the distributor + if (!servers.Any(s => s.Connection.State == HubConnectionState.Connected)) + { + serverAvailableSignal.Reset(); + } + + logger.LogInformation("Successfully requeued {Count} tasks from {Url}", count, server.HttpClient.BaseAddress); + } +} \ No newline at end of file diff --git a/RemoteExec.Client/RemoteExecutor.EventHandlers.cs b/RemoteExec.Client/RemoteExecutor.EventHandlers.cs new file mode 100644 index 0000000..6311098 --- /dev/null +++ b/RemoteExec.Client/RemoteExecutor.EventHandlers.cs @@ -0,0 +1,97 @@ +using Microsoft.AspNetCore.SignalR.Client; +using Microsoft.Extensions.Logging; + +using RemoteExec.Shared; + +using System.Collections.Concurrent; +using System.Reflection; + +namespace RemoteExec.Client; + +public partial class RemoteExecutor +{ + private void RegisterEventHandlers(ServerConnection server) + { + // Unregister existing handlers to prevent duplicates + UnregisterEventHandlers(server); + + _ = server.Connection.On("MetricsUpdated", metrics => OnMetricsUpdated(server, metrics)); + _ = server.Connection.On("TaskResult", (taskId, result) => OnTaskResult(server, taskId, result)); + _ = server.Connection.On("RequestAssembly", async (assemblyName, requestId) => await OnRequestAssemblyAsync(server, assemblyName, requestId)); + + // Store the handler so it can be unregistered later + server.ClosedEventHandler = async (error) => await OnConnectionClosedAsync(server, error); + server.Connection.Closed += server.ClosedEventHandler; + } + + private static void UnregisterEventHandlers(ServerConnection server) + { + server.Connection.Remove("MetricsUpdated"); + server.Connection.Remove("TaskResult"); + server.Connection.Remove("RequestAssembly"); + + // Unregister the Closed event handler if it exists + if (server.ClosedEventHandler != null) + { + server.Connection.Closed -= server.ClosedEventHandler; + server.ClosedEventHandler = null; + } + } + + private void OnMetricsUpdated(ServerConnection server, ServerMetrics metrics) + { + server.Metrics = metrics; + MetricsUpdated?.Invoke(this, new ServerMetricsUpdatedEventArgs(server.Connection, metrics)); + } + + private void OnTaskResult(ServerConnection server, Guid taskId, RemoteExecutionResult result) + { + // Remove completed task from server's assigned tasks + if (serverAssignedTasks.TryGetValue(server, out ConcurrentDictionary? taskDictionary)) + { + _ = taskDictionary.TryRemove(taskId, out _); + } + + if (pendingResults.TryRemove(taskId, out TaskCompletionSource? tcs)) + { + tcs.SetResult(result); + } + } + + private async Task OnRequestAssemblyAsync(ServerConnection server, string assemblyName, Guid requestId) + { + Assembly? assembly = AppDomain.CurrentDomain.GetAssemblies().FirstOrDefault(a => a.GetName().FullName == assemblyName) ?? Assembly.Load(new AssemblyName(assemblyName)); + + if (string.IsNullOrEmpty(assembly.Location)) + { + logger.LogWarning("Cannot provide assembly {AssemblyName}: Location is null or empty", assemblyName); + return; + } + + byte[] dllBytes = await File.ReadAllBytesAsync(assembly.Location); + + ByteArrayContent content = new(dllBytes); + content.Headers.ContentType = new System.Net.Http.Headers.MediaTypeHeaderValue("application/octet-stream"); + _ = await server.HttpClient.PostAsync($"/provide-assembly?requestId={requestId}", content); + } + + private async Task OnConnectionClosedAsync(ServerConnection server, Exception? error) + { + if (error != null) + { + logger.LogWarning(error, "Server connection closed unexpectedly: {ServerUrl}", server.HttpClient.BaseAddress); + } + else + { + logger.LogInformation("Server connection closed gracefully: {ServerUrl}", server.HttpClient.BaseAddress); + } + + // Requeue all tasks assigned to this server + await RequeuePendingTasksForServer(server); + + if (!distributorCts.IsCancellationRequested) + { + _ = Task.Run(() => AttemptReconnectionAsync(server, distributorCts.Token)); + } + } +} \ No newline at end of file diff --git a/RemoteExec.Client/RemoteExecutor.InternalTypes.cs b/RemoteExec.Client/RemoteExecutor.InternalTypes.cs new file mode 100644 index 0000000..737ccdc --- /dev/null +++ b/RemoteExec.Client/RemoteExecutor.InternalTypes.cs @@ -0,0 +1,26 @@ +using Microsoft.AspNetCore.SignalR.Client; + +using RemoteExec.Shared; + +using System.Threading.Channels; + +namespace RemoteExec.Client; + +public partial class RemoteExecutor +{ + private sealed class ServerConnection(HubConnection connection, HttpClient httpClient) + { + public HubConnection Connection { get; } = connection; + public HttpClient HttpClient { get; } = httpClient; + public Channel TaskChannel { get; } = Channel.CreateUnbounded(); + public ServerMetrics? Metrics { get; set; } + public Func? ClosedEventHandler { get; set; } + } + + private sealed class PendingTask + { + public required Guid TaskId { get; init; } + public required RemoteExecutionRequest Request { get; init; } + public required DateTime EnqueuedAt { get; init; } + } +} \ No newline at end of file diff --git a/RemoteExec.Client/RemoteExecutor.TaskDistribution.cs b/RemoteExec.Client/RemoteExecutor.TaskDistribution.cs new file mode 100644 index 0000000..8a56985 --- /dev/null +++ b/RemoteExec.Client/RemoteExecutor.TaskDistribution.cs @@ -0,0 +1,102 @@ +using Microsoft.AspNetCore.SignalR.Client; +using Microsoft.Extensions.Logging; + +using RemoteExec.Shared; + +using System.Collections.Concurrent; + +namespace RemoteExec.Client; + +public partial class RemoteExecutor +{ + private async Task DistributorLoop(CancellationToken cancellationToken) + { + while (!cancellationToken.IsCancellationRequested) + { + PendingTask? pendingTask = null; + + try + { + // Select the best server based on metrics + ServerConnection? bestServer = SelectBestServer(); + + if (bestServer is null) + { + await serverAvailableSignal.WaitAsync(cancellationToken); + continue; + } + + // Check if there are tasks in the global queue + pendingTask = globalQueue.Take(cancellationToken); + + // Track task assignment + if (serverAssignedTasks.TryGetValue(bestServer, out ConcurrentDictionary? taskDictionary)) + { + _ = taskDictionary.TryAdd(pendingTask.TaskId, pendingTask); + } + + // Create task item with ID + TaskItem taskItem = new TaskItem + { + TaskId = pendingTask.TaskId, + Request = pendingTask.Request + }; + + // Push to server's channel - SignalR will stream it + await bestServer.TaskChannel.Writer.WriteAsync(taskItem, cancellationToken); + + } + catch (OperationCanceledException) + { + break; + } + catch (Exception ex) + { + logger.LogError(ex, "Error in distributor loop"); + + if (pendingTask != null) + { + globalQueue.Add(pendingTask, cancellationToken); + } + + await Task.Delay(100, cancellationToken); + } + } + } + + private ServerConnection? SelectBestServer() + { + if (servers.Count == 0) + { + return null; + } + + // Only select servers that are connected + IEnumerable connectedServers = servers.Where(s => s.Connection.State == HubConnectionState.Connected); + + if (!connectedServers.Any()) + { + return null; + } + + return loadBalancingStrategy switch + { + LoadBalancingStrategy.ResourceAware => connectedServers.MinBy(s => + { + if (s.Metrics == null) + { + return double.MaxValue; + } + + double cpuScore = s.Metrics.CpuUsage; + double activeTaskScore = s.Metrics.ActiveTasks * 10; + double backlogScore = s.TaskChannel.Reader.Count * 50; + + return cpuScore + activeTaskScore + backlogScore; + }), + LoadBalancingStrategy.LeastBacklog => connectedServers.MinBy(s => s.TaskChannel.Reader.Count), + + _ => connectedServers.FirstOrDefault() + }; + } +} \ No newline at end of file diff --git a/RemoteExec.Client/RemoteExecutor.cs b/RemoteExec.Client/RemoteExecutor.cs index abc4dff..89af28a 100644 --- a/RemoteExec.Client/RemoteExecutor.cs +++ b/RemoteExec.Client/RemoteExecutor.cs @@ -5,19 +5,19 @@ using Microsoft.Extensions.Logging.Abstractions; using RemoteExec.Shared; using System.Collections.Concurrent; -using System.Diagnostics; using System.Reflection; using System.Text.Json; -using System.Threading.Channels; namespace RemoteExec.Client; -public class RemoteExecutor : IAsyncDisposable +public partial class RemoteExecutor : IAsyncDisposable { private readonly List servers = []; private readonly BlockingCollection globalQueue = []; private readonly ConcurrentDictionary> pendingResults = new(); + private readonly ConcurrentDictionary> serverAssignedTasks = new(); + private readonly AsyncManualResetEvent serverAvailableSignal = new(false); private CancellationTokenSource distributorCts = new(); private Task? distributorTask; @@ -60,7 +60,6 @@ public class RemoteExecutor : IAsyncDisposable HubConnection connection = new HubConnectionBuilder() .WithUrl(signalRUri) - .WithAutomaticReconnect() .ConfigureLogging(logging => { _ = logging.AddProvider(new RemoteExecLoggerProvider(logger)); @@ -74,93 +73,10 @@ public class RemoteExecutor : IAsyncDisposable ServerConnection serverConnection = new ServerConnection(connection, httpClient); servers.Add(serverConnection); + serverAssignedTasks[serverConnection] = new ConcurrentDictionary(); } } - public async Task StartAsync(CancellationToken cancellationToken = default) - { - distributorCts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken); - - List startTasks = []; - - foreach (ServerConnection server in servers) - { - _ = server.Connection.On("MetricsUpdated", metrics => - { - server.Metrics = metrics; - MetricsUpdated?.Invoke(this, new ServerMetricsUpdatedEventArgs(server.Connection, metrics)); - }); - - _ = server.Connection.On("TaskResult", (taskId, result) => - { - if (pendingResults.TryRemove(taskId, out TaskCompletionSource? tcs)) - { - tcs.SetResult(result); - } - }); - - _ = server.Connection.On($"RequestAssembly", async (string assemblyName, Guid requestId) => - { - Assembly? assembly = AppDomain.CurrentDomain.GetAssemblies().FirstOrDefault(a => a.GetName().FullName == assemblyName) ?? Assembly.Load(new AssemblyName(assemblyName)); - byte[] dllBytes = await File.ReadAllBytesAsync(assembly.Location!); - - ByteArrayContent content = new(dllBytes); - content.Headers.ContentType = new System.Net.Http.Headers.MediaTypeHeaderValue("application/octet-stream"); - _ = await server.HttpClient.PostAsync($"/provide-assembly?requestId={requestId}", content); - }); - - startTasks.Add(server.Connection.StartAsync(cancellationToken) - .ContinueWith(async (task, state) => - { - ServerConnection conn = (ServerConnection)state!; - conn.Metrics = await conn.Connection.InvokeAsync("GetMetrics", cancellationToken); - MetricsUpdated?.Invoke(this, new ServerMetricsUpdatedEventArgs(conn.Connection, conn.Metrics)); - - await conn.Connection.SendAsync("StartTaskStream", conn.TaskChannel.Reader, cancellationToken); - }, server, TaskScheduler.Default).Unwrap()); - } - - await Task.WhenAll(startTasks); - - distributorTask = Task.Run(() => DistributorLoop(distributorCts.Token), distributorCts.Token); - } - - public async Task StopAsync(CancellationToken cancellationToken = default) - { - logger.LogInformation("Stopping RemoteExecutor..."); - - await distributorCts.CancelAsync(); - - if (distributorTask != null) - { - try - { - await distributorTask; - logger.LogDebug("Distributor task completed successfully"); - } - catch (OperationCanceledException ex) - { - logger.LogError(ex, "Distributor task was canceled"); - } - } - - logger.LogDebug("Completing task channels for {ServerCount} servers", servers.Count); - foreach (ServerConnection server in servers) - { - server.TaskChannel.Writer.Complete(); - } - - List stopTasks = []; - - foreach (ServerConnection server in servers) - { - stopTasks.Add(server.Connection.StopAsync(cancellationToken)); - } - - await Task.WhenAll(stopTasks); - logger.LogInformation("RemoteExecutor stopped successfully"); - } - public Dictionary GetCurrentServerMetrics() { return servers @@ -198,6 +114,11 @@ public class RemoteExecutor : IAsyncDisposable throw new InvalidOperationException("Only static methods supported"); } + if (assembly.IsDynamic) + { + throw new InvalidOperationException("Dynamic assemblies are not supported"); + } + RemoteExecutionRequest request = new RemoteExecutionRequest { AssemblyName = assembly.GetName().FullName, @@ -230,98 +151,17 @@ public class RemoteExecutor : IAsyncDisposable return result.Result; } - private async Task DistributorLoop(CancellationToken cancellationToken) - { - while (!cancellationToken.IsCancellationRequested) - { - try - { - // Check if there are tasks in the global queue - PendingTask pendingTask = globalQueue.Take(cancellationToken); - - // Select the best server based on metrics - ServerConnection? bestServer = SelectBestServer(); - - if (bestServer != null) - { - // Create task item with ID - TaskItem taskItem = new TaskItem - { - TaskId = pendingTask.TaskId, - Request = pendingTask.Request - }; - - // Push to server's channel - SignalR will stream it - await bestServer.TaskChannel.Writer.WriteAsync(taskItem, cancellationToken); - } - else - { - // No available server, re-enqueue - globalQueue.Add(pendingTask, cancellationToken); - await Task.Delay(100, cancellationToken); - } - } - catch (OperationCanceledException) - { - break; - } - catch (Exception ex) - { - Debug.WriteLine($"DistributorLoop exception: {ex}"); - await Task.Delay(100, cancellationToken); - } - } - } - - private ServerConnection? SelectBestServer() - { - if (servers.Count == 0) - { - return null; - } - - return loadBalancingStrategy switch - { - LoadBalancingStrategy.ResourceAware => servers.MinBy(s => - { - if (s.Metrics == null) - { - return double.MaxValue; - } - - double cpuScore = s.Metrics.CpuUsage; - double activeTaskScore = s.Metrics.ActiveTasks * 10; - double backlogScore = s.TaskChannel.Reader.Count * 50; - - return cpuScore + activeTaskScore + backlogScore; - }), - LoadBalancingStrategy.LeastBacklog => servers.MinBy(s => s.TaskChannel.Reader.Count), - - _ => servers.FirstOrDefault() - }; - } - - private sealed class ServerConnection(HubConnection connection, HttpClient httpClient) - { - public HubConnection Connection { get; } = connection; - public HttpClient HttpClient { get; } = httpClient; - public Channel TaskChannel { get; } = Channel.CreateUnbounded(); - public ServerMetrics? Metrics { get; set; } - } - - private sealed class PendingTask - { - public required Guid TaskId { get; init; } - public required RemoteExecutionRequest Request { get; init; } - public required DateTime EnqueuedAt { get; init; } - } - protected virtual async Task DisposeAsync(bool disposing) { if (!disposedValue) { if (disposing) { + if (!distributorCts.IsCancellationRequested) + { + await StopAsync(); + } + distributorCts.Dispose(); foreach (ServerConnection server in servers) @@ -329,10 +169,9 @@ public class RemoteExecutor : IAsyncDisposable await server.Connection.DisposeAsync(); server.HttpClient.Dispose(); } - } - // Free unmanaged resources (unmanaged objects) and override finalizer if needed - // Set large fields to null if needed + globalQueue.Dispose(); + } disposedValue = true; } @@ -343,4 +182,4 @@ public class RemoteExecutor : IAsyncDisposable await DisposeAsync(disposing: true); GC.SuppressFinalize(this); } -} +} \ No newline at end of file