Improve load balancing and move to pull based model

RemoteExec.Client/RemoteExecutor.cs

@@ -1,92 +1,135 @@
 using Microsoft.AspNetCore.SignalR.Client;
+using Microsoft.Extensions.Logging;
 
 using RemoteExec.Shared;
 
 using System.Collections.Concurrent;
-using System.Diagnostics.CodeAnalysis;
+using System.Diagnostics;
 using System.Reflection;
 using System.Text.Json;
 using System.Threading.Channels;
 
 namespace RemoteExec.Client;
 
-public class RemoteExecutor
+public class RemoteExecutor : IDisposable
 {
-    private readonly List<HubConnection> connections = [];
-    private readonly ConcurrentDictionary<HubConnection, ServerMetrics> serverMetrics = new();
-    private int _currentConnectionIndex = 0;
-    private readonly Lock @lock = new Lock();
+    private readonly List<ServerConnection> servers = [];
+    private readonly BlockingCollection<PendingTask> globalQueue = [];
+    private readonly ConcurrentDictionary<Guid, TaskCompletionSource<RemoteExecutionResult>> pendingResults = new();
+    private CancellationTokenSource distributorCts = new();
+    private Task? distributorTask;
     private readonly LoadBalancingStrategy loadBalancingStrategy;
 
+    private bool disposedValue;
+
     public event EventHandler<ServerMetricsUpdatedEventArgs>? MetricsUpdated;
 
-    public RemoteExecutor(string url) : this([url], LoadBalancingStrategy.RoundRobin)
+    public RemoteExecutor(string url) : this([url], LoadBalancingStrategy.ResourceAware)
     {
     }
 
-    public RemoteExecutor(string[] urls, LoadBalancingStrategy loadBalancingStrategy = LoadBalancingStrategy.RoundRobin)
+    public RemoteExecutor(string[] urls, LoadBalancingStrategy loadBalancingStrategy = LoadBalancingStrategy.ResourceAware)
     {
         this.loadBalancingStrategy = loadBalancingStrategy;
 
         foreach (string url in urls)
         {
+            Uri baseUri = new(url);
+            Uri signalRUri = new(baseUri, "/remote");
+
             HubConnection connection = new HubConnectionBuilder()
-                .WithUrl(url)
+                .WithUrl(signalRUri)
                 .WithAutomaticReconnect()
+                .ConfigureLogging(logging =>
+                {
+                    _ = logging.AddProvider(new RemoteExecLoggerProvider());
+                })
                 .Build();
 
-            connections.Add(connection);
+            HttpClient httpClient = new()
+            {
+                BaseAddress = baseUri
+            };
+
+            ServerConnection serverConnection = new ServerConnection(connection, httpClient);
+            servers.Add(serverConnection);
         }
     }
 
     public async Task StartAsync(CancellationToken cancellationToken = default)
     {
+        distributorCts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
+
         List<Task> startTasks = [];
 
-        foreach (HubConnection connection in connections)
+        foreach (ServerConnection server in servers)
         {
-            _ = connection.On<ServerMetrics>("MetricsUpdated", metrics =>
+            _ = server.Connection.On<ServerMetrics>("MetricsUpdated", metrics =>
             {
-                serverMetrics[connection] = metrics;
-                MetricsUpdated?.Invoke(this, new ServerMetricsUpdatedEventArgs(connection, metrics));
+                server.Metrics = metrics;
+                MetricsUpdated?.Invoke(this, new ServerMetricsUpdatedEventArgs(server.Connection, metrics));
             });
 
-            _ = connection.On($"RequestAssembly", async (string assemblyName, Guid requestId) =>
+            _ = 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!);
 
-                Channel<byte> channel = Channel.CreateUnbounded<byte>();
-
-                foreach (byte b in dllBytes)
-                {
-                    await channel.Writer.WriteAsync(b);
-                }
-
-                channel.Writer.Complete();
-
-                await connection.InvokeAsync("ProvideAssembly", requestId, channel.Reader);
+                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(connection.StartAsync(cancellationToken)
+            startTasks.Add(server.Connection.StartAsync(cancellationToken)
                 .ContinueWith(async (task, state) =>
                 {
-                    HubConnection conn = (HubConnection)state!;
-                    serverMetrics[conn] = await conn.InvokeAsync<ServerMetrics>("GetMetrics", cancellationToken);
-                    MetricsUpdated?.Invoke(this, new ServerMetricsUpdatedEventArgs(conn, serverMetrics[conn]));
-                }, connection, TaskScheduler.Default).Unwrap());
+                    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)
     {
+        await distributorCts.CancelAsync();
+
+        if (distributorTask != null)
+        {
+            try
+            {
+                await distributorTask;
+            }
+            catch (OperationCanceledException)
+            {
+                // Expected
+            }
+        }
+
+        foreach (ServerConnection server in servers)
+        {
+            server.TaskChannel.Writer.Complete();
+        }
+
         List<Task> stopTasks = [];
 
-        foreach (HubConnection connection in connections)
+        foreach (ServerConnection server in servers)
         {
-            stopTasks.Add(connection.StopAsync(cancellationToken));
+            stopTasks.Add(server.Connection.StopAsync(cancellationToken));
         }
 
         await Task.WhenAll(stopTasks);
@@ -94,52 +137,15 @@ public class RemoteExecutor
 
     public Dictionary<string, ServerMetrics> GetCurrentServerMetrics()
     {
-        Dictionary<string, ServerMetrics> metrics = [];
-
-        foreach (HubConnection connection in connections)
-        {
-            if (serverMetrics.TryGetValue(connection, out ServerMetrics? newServerMetrics))
-            {
-                metrics[newServerMetrics.ServerId] = newServerMetrics;
-            }
-        }
-
-        return metrics;
+        return servers
+            .Select(server => server.Metrics)
+            .Where(metrics => metrics != null)
+            .ToDictionary(metrics => metrics!.ServerId, metrics => metrics!);
     }
 
-    public bool TryExecute<TDelegate, TResult>(TDelegate del, out TResult? result, params object[] args) where TDelegate : Delegate
+    public async Task<TResult> Execute<TDelegate, TResult>(TDelegate del, params object[] args) where TDelegate : Delegate
     {
-        try
-        {
-            result = Execute<TDelegate, TResult>(del, args);
-            return true;
-        }
-        catch
-        {
-            result = default;
-            return false;
-        }
-    }
-
-    public bool TryExecute<TDelegate, TResult>(TDelegate del, out TResult? result, [NotNullWhen(false)] out Exception? exception, params object[] args) where TDelegate : Delegate
-    {
-        try
-        {
-            result = Execute<TDelegate, TResult>(del, args);
-            exception = null;
-            return true;
-        }
-        catch (Exception ex)
-        {
-            result = default;
-            exception = ex;
-            return false;
-        }
-    }
-
-    public TResult Execute<TDelegate, TResult>(TDelegate del, params object[] args) where TDelegate : Delegate
-    {
-        object? execResult = Execute(del, args);
+        object? execResult = await Execute(del, args);
 
         if (execResult is TResult typedResult)
         {
@@ -155,7 +161,7 @@ public class RemoteExecutor
         }
     }
 
-    public object? Execute<T>(T del, params object[] args) where T : Delegate
+    public async Task<object?> Execute<T>(T del, params object[] args) where T : Delegate
     {
         MethodInfo method = del.Method;
         Type declaringType = method.DeclaringType!;
@@ -175,12 +181,20 @@ public class RemoteExecutor
             Arguments = args
         };
 
-        HubConnection connection = GetNextConnection();
+        Guid taskId = Guid.NewGuid();
+        TaskCompletionSource<RemoteExecutionResult> tcs = new();
+        pendingResults[taskId] = tcs;
 
-        RemoteExecutionResult result = connection
-            .InvokeAsync<RemoteExecutionResult>("Execute", request)
-            .GetAwaiter()
-            .GetResult();
+        PendingTask pendingTask = new PendingTask
+        {
+            TaskId = taskId,
+            Request = request,
+            EnqueuedAt = DateTime.UtcNow
+        };
+
+        globalQueue.Add(pendingTask);
+
+        RemoteExecutionResult result = await tcs.Task;
 
         if (result.Exception != null)
         {
@@ -190,62 +204,145 @@ public class RemoteExecutor
         return result.Result;
     }
 
-    private HubConnection GetNextConnection()
+    private async Task DistributorLoop(CancellationToken cancellationToken)
     {
-        if (connections.Count == 0)
+        while (!cancellationToken.IsCancellationRequested)
         {
-            throw new InvalidOperationException("No connections available");
+            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.RoundRobin => GetRoundRobinConnection(),
-            LoadBalancingStrategy.Random => GetRandomConnection(),
-            LoadBalancingStrategy.LeastConnections => GetLeastConnections(),
-            LoadBalancingStrategy.LeastActiveTasks => GetLeastActiveTasksConnections(),
-            LoadBalancingStrategy.ResourceAware => GetResourceAwareConnections(),
-            _ => throw new NotSupportedException($"Load balancing strategy {loadBalancingStrategy} is not supported")
+            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 HubConnection GetRoundRobinConnection()
+    private sealed class ServerConnection(HubConnection connection, HttpClient httpClient)
     {
-        lock (@lock)
+        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; }
+    }
+
+    // S2930: Dispose distributorCts when no longer needed
+    protected virtual void Dispose(bool disposing)
+    {
+        if (!disposedValue)
         {
-            HubConnection connection = connections[_currentConnectionIndex];
-            _currentConnectionIndex = (_currentConnectionIndex + 1) % connections.Count;
-            return connection;
+            if (disposing)
+            {
+                distributorCts.Dispose();
+                // Dispose managed state (managed objects) here if needed
+            }
+
+            // Free unmanaged resources (unmanaged objects) and override finalizer if needed
+            // Set large fields to null if needed
+
+            disposedValue = true;
         }
     }
 
-    private HubConnection GetRandomConnection()
+    public void Dispose()
     {
-        int index = Random.Shared.Next(connections.Count);
-        return connections[index];
-    }
-
-    private HubConnection GetLeastConnections()
-    {
-        // Purely looks at how many clients are connected to the Hub
-        return connections.OrderBy(c => serverMetrics.TryGetValue(c, out ServerMetrics? m) ? m.ActiveConnections : 0).First();
-    }
-
-    private HubConnection GetLeastActiveTasksConnections()
-    {
-        return connections.OrderBy(c => serverMetrics.TryGetValue(c, out ServerMetrics? m) ? m.ActiveTasks : 0).First();
-    }
-
-    private HubConnection GetResourceAwareConnections()
-    {
-        return connections.OrderBy(c =>
-        {
-            if (!serverMetrics.TryGetValue(c, out ServerMetrics? m))
-            {
-                return 0;
-            }
-
-            // Simple heuristic: CPU percentage + (Memory in MB / 1024)
-            return m.CpuUsage + (m.TotalMemoryUsage / 1024 / 1024 / 100);
-        }).First();
+        // Do not change this code. Put cleanup code in 'Dispose(bool disposing)' method
+        Dispose(disposing: true);
+        GC.SuppressFinalize(this);
     }
 }
+
+internal class RemoteExecLoggerProvider : ILoggerProvider
+{
+    public ILogger CreateLogger(string categoryName)
+    {
+        return new RemoteExecLogger();
+    }
+
+    public void Dispose()
+    {
+        throw new NotImplementedException();
+    }
+}
+
+internal class RemoteExecLogger : ILogger
+{
+    public IDisposable? BeginScope<TState>(TState state) where TState : notnull
+    {
+        return null;
+    }
+
+    public bool IsEnabled(LogLevel logLevel)
+    {
+        return true;
+    }
+
+    public void Log<TState>(LogLevel logLevel, EventId eventId, TState state, Exception? exception, Func<TState, Exception?, string> formatter)
+    {
+        Console.WriteLine($"[{logLevel}] {formatter(state, exception)}");
+    }
+}