112 lines
4.4 KiB
C#
112 lines
4.4 KiB
C#
// Copyright Epic Games, Inc. All Rights Reserved.
|
|
|
|
using System;
|
|
using System.Collections.Generic;
|
|
using System.Net;
|
|
using System.Net.Sockets;
|
|
using System.Threading;
|
|
using System.Threading.Tasks;
|
|
using EpicGames.Core;
|
|
using EpicGames.Horde.Compute.Transports;
|
|
using Microsoft.Extensions.Logging;
|
|
|
|
namespace EpicGames.Horde.Compute.Clients
|
|
{
|
|
/// <summary>
|
|
/// Implementation of <see cref="IComputeClient"/> which marshals data over a loopback connection to a method running on a background task in the same process.
|
|
/// </summary>
|
|
public sealed class LocalComputeClient : IComputeClient, IAsyncDisposable
|
|
{
|
|
private static readonly ClusterId s_cluster = new ("_local");
|
|
|
|
class LeaseImpl : IComputeLease
|
|
{
|
|
public ClusterId Cluster { get; } = s_cluster;
|
|
public IReadOnlyList<string> Properties { get; } = new List<string>();
|
|
public IReadOnlyDictionary<string, int> AssignedResources => new Dictionary<string, int>();
|
|
public RemoteComputeSocket Socket => _socket;
|
|
public string Ip => "127.0.0.1";
|
|
public ConnectionMode ConnectionMode => ConnectionMode.Direct;
|
|
public IReadOnlyDictionary<string, ConnectionMetadataPort> Ports => new Dictionary<string, ConnectionMetadataPort>();
|
|
|
|
readonly RemoteComputeSocket _socket;
|
|
|
|
public LeaseImpl(RemoteComputeSocket socket) => _socket = socket;
|
|
|
|
/// <inheritdoc/>
|
|
public ValueTask DisposeAsync() => _socket.DisposeAsync();
|
|
|
|
/// <inheritdoc/>
|
|
public ValueTask CloseAsync(CancellationToken cancellationToken) => _socket.CloseAsync(cancellationToken);
|
|
}
|
|
|
|
readonly BackgroundTask _listenerTask;
|
|
readonly Socket _listener;
|
|
readonly Socket _socket;
|
|
readonly bool _executeInProcess;
|
|
|
|
/// <summary>
|
|
/// Constructor
|
|
/// </summary>
|
|
/// <param name="port">Port to connect on</param>
|
|
/// <param name="sandboxDir">Sandbox directory for the worker</param>
|
|
/// <param name="executeInProcess">Whether to run external assemblies in-process. Useful for debugging.</param>
|
|
/// <param name="logger">Logger for diagnostic output</param>
|
|
public LocalComputeClient(int port, DirectoryReference sandboxDir, bool executeInProcess, ILogger logger)
|
|
{
|
|
_executeInProcess = executeInProcess;
|
|
|
|
_listener = new Socket(SocketType.Stream, ProtocolType.IP);
|
|
_listener.Bind(new IPEndPoint(IPAddress.Loopback, port));
|
|
_listener.Listen();
|
|
|
|
_listenerTask = BackgroundTask.StartNew(ctx => RunListenerAsync(_listener, sandboxDir, _executeInProcess, logger, ctx));
|
|
|
|
_socket = new Socket(SocketType.Stream, ProtocolType.IP);
|
|
_socket.Connect(IPAddress.Loopback, port);
|
|
}
|
|
|
|
/// <inheritdoc/>
|
|
public async ValueTask DisposeAsync()
|
|
{
|
|
_socket.Dispose();
|
|
await _listenerTask.DisposeAsync();
|
|
_listener.Dispose();
|
|
}
|
|
|
|
/// <summary>
|
|
/// Sets up the loopback listener and calls the server method
|
|
/// </summary>
|
|
static async Task RunListenerAsync(Socket listener, DirectoryReference sandboxDir, bool executeInProcess, ILogger logger, CancellationToken cancellationToken)
|
|
{
|
|
using Socket tcpSocket = await listener.AcceptAsync(cancellationToken);
|
|
await using TcpTransport tcpTransport = new(tcpSocket);
|
|
await using RemoteComputeSocket socket = new(tcpTransport, ComputeProtocol.Latest, logger);
|
|
AgentMessageHandler worker = new(sandboxDir, null, executeInProcess, null, null, logger);
|
|
await worker.RunAsync(socket, cancellationToken);
|
|
await socket.CloseAsync(cancellationToken);
|
|
}
|
|
|
|
/// <inheritdoc/>
|
|
public Task<ClusterId> GetClusterAsync(Requirements? requirements, string? requestId, ConnectionMetadataRequest? connection, ILogger logger, CancellationToken cancellationToken = default)
|
|
{
|
|
return Task.FromResult(s_cluster);
|
|
}
|
|
|
|
/// <inheritdoc/>
|
|
public Task<IComputeLease?> TryAssignWorkerAsync(ClusterId? clusterId, Requirements? requirements, string? requestId, ConnectionMetadataRequest? connection, ILogger logger, CancellationToken cancellationToken)
|
|
{
|
|
#pragma warning disable CA2000 // Dispose objects before losing scope
|
|
RemoteComputeSocket socket = new RemoteComputeSocket(new TcpTransport(_socket), ComputeProtocol.Latest, logger);
|
|
return Task.FromResult<IComputeLease?>(new LeaseImpl(socket));
|
|
#pragma warning restore CA2000 // Dispose objects before losing scope
|
|
}
|
|
|
|
/// <inheritdoc/>
|
|
public Task DeclareResourceNeedsAsync(ClusterId clusterId, string pool, Dictionary<string, int> resourceNeeds, CancellationToken cancellationToken = default)
|
|
{
|
|
return Task.CompletedTask;
|
|
}
|
|
}
|
|
}
|