mirror of
https://github.com/Stone-Red-Code/RemoteExec.git
synced 2026-09-04 09:06:20 +02:00
Improve load balancing and move to pull based model
This commit is contained in:
Stone_Red
@@ -2,9 +2,6 @@
|
||||
|
||||
public enum LoadBalancingStrategy
|
||||
{
|
||||
RoundRobin,
|
||||
Random,
|
||||
LeastConnections,
|
||||
LeastActiveTasks,
|
||||
ResourceAware
|
||||
ResourceAware,
|
||||
LeastBacklog
|
||||
}
|
||||
+220
-123
@@ -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)}");
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user