Refactor RemoteExecutor

RemoteExec.Client/AsyncManualResetEvent.cs

@@ -0,0 +1,37 @@
+namespace RemoteExec.Client;
+
+internal class AsyncManualResetEvent
+{
+    private volatile TaskCompletionSource<bool> _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<bool> tcs = _tcs;
+            if (!tcs.Task.IsCompleted ||
+                Interlocked.CompareExchange(ref _tcs, new TaskCompletionSource<bool>(TaskCreationOptions.RunContinuationsAsynchronously), tcs) == tcs)
+            {
+                return;
+            }
+        }
+    }
+}

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<Task> 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<Task> 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<ServerMetrics>("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<Guid, PendingTask>? taskDict))
+        {
+            return;
+        }
+
+        List<Guid> 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);
+    }
+}

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<ServerMetrics>("MetricsUpdated", metrics => OnMetricsUpdated(server, metrics));
+        _ = server.Connection.On<Guid, RemoteExecutionResult>("TaskResult", (taskId, result) => OnTaskResult(server, taskId, result));
+        _ = server.Connection.On<string, Guid>("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<Guid, PendingTask>? taskDictionary))
+        {
+            _ = taskDictionary.TryRemove(taskId, out _);
+        }
+
+        if (pendingResults.TryRemove(taskId, out TaskCompletionSource<RemoteExecutionResult>? 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));
+        }
+    }
+}

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<TaskItem> TaskChannel { get; } = Channel.CreateUnbounded<TaskItem>();
+        public ServerMetrics? Metrics { get; set; }
+        public Func<Exception?, Task>? 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; }
+    }
+}

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<Guid, PendingTask>? 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<ServerConnection> 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()
+        };
+    }
+}

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<ServerConnection> servers = [];
     private readonly BlockingCollection<PendingTask> globalQueue = [];
     private readonly ConcurrentDictionary<Guid, TaskCompletionSource<RemoteExecutionResult>> pendingResults = new();
+    private readonly ConcurrentDictionary<ServerConnection, ConcurrentDictionary<Guid, PendingTask>> 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<Guid, PendingTask>();
         }
     }
 
-    public async Task StartAsync(CancellationToken cancellationToken = default)
-    {
-        distributorCts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
-
-        List<Task> startTasks = [];
-
-        foreach (ServerConnection server in servers)
-        {
-            _ = server.Connection.On<ServerMetrics>("MetricsUpdated", metrics =>
-            {
-                server.Metrics = metrics;
-                MetricsUpdated?.Invoke(this, new ServerMetricsUpdatedEventArgs(server.Connection, metrics));
-            });
-
-            _ = server.Connection.On<Guid, RemoteExecutionResult>("TaskResult", (taskId, result) =>
-            {
-                if (pendingResults.TryRemove(taskId, out TaskCompletionSource<RemoteExecutionResult>? 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<ServerMetrics>("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<Task> stopTasks = [];
-
-        foreach (ServerConnection server in servers)
-        {
-            stopTasks.Add(server.Connection.StopAsync(cancellationToken));
-        }
-
-        await Task.WhenAll(stopTasks);
-        logger.LogInformation("RemoteExecutor stopped successfully");
-    }
-
     public Dictionary<string, ServerMetrics> 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<TaskItem> TaskChannel { get; } = Channel.CreateUnbounded<TaskItem>();
-        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);
     }
-}
+}