Files
UnrealEngine/Engine/Source/Programs/UnrealToolbox/Plugins/HordeAgent/HordeAgentPlugin.cs
2025-05-18 13:04:45 +08:00

577 lines
17 KiB
C#

// Copyright Epic Games, Inc. All Rights Reserved.
using Avalonia.Controls;
using EpicGames.Core;
using EpicGames.Horde;
using FluentAvalonia.UI.Controls;
using HordeAgent;
using Microsoft.Extensions.Logging;
using Microsoft.Win32;
using System.Diagnostics;
using System.IO.Pipes;
using System.Runtime.InteropServices;
using System.Text.Json;
using System.Text.Json.Serialization;
namespace UnrealToolbox.Plugins.HordeAgent
{
class HordeAgentPlugin : ToolboxPluginBase
{
record struct IdleStat(string Name, long Value, long MinValue);
readonly IToolboxPluginHost _host;
readonly IToolCatalog _toolCatalog;
readonly ILogger _logger;
readonly BackgroundTask _clientTask;
readonly BackgroundTask _tickPauseStateTask;
readonly FileReference _settingsFile;
byte[] _settingsData = Array.Empty<byte>();
HordeAgentSettings _settings = new HordeAgentSettings();
bool _pipeConnected;
public HordeAgentSettings Settings => _settings;
AgentSettingsMessage? _agentSettings;
ToolboxPluginStatus? _status;
ToolboxPluginStatus? _reportStatus; // Updated with _status when nothing is currently being installed
public override string Name => "Horde Agent";
public override IconSource Icon => new SymbolIconSource() { Symbol = Symbol.People };
public bool IsEnabled { get; private set; }
public override bool HasSettingsPage()
=> IsEnabled;
public override Control CreateSettingsPage(SettingsContext context)
=> new HordeAgentSettingsPage(context, this);
public void UpdateSettings(HordeAgentSettings settings)
{
byte[] data = JsonSerializer.SerializeToUtf8Bytes(settings, GetJsonSerializerOptions());
if (!data.SequenceEqual(_settingsData))
{
_settings = settings;
_settingsData = data;
DirectoryReference.CreateDirectory(_settingsFile.Directory);
FileReference.WriteAllBytes(_settingsFile, data);
_statusChangedEvent.Set();
}
}
void EnrollWithServer()
{
Uri? serverUrl = _agentSettings?.ServerUrl;
if (serverUrl != null)
{
Process.Start(new ProcessStartInfo(new Uri(serverUrl, "agents/registration").ToString()) { UseShellExecute = true });
}
}
public HordeAgentPlugin(IToolboxPluginHost host, ToolCatalog toolCatalog, ILogger<HordeAgentPlugin> logger)
{
_host = host;
_toolCatalog = toolCatalog;
_logger = logger;
_settingsFile = FileReference.Combine(Program.DataDir, "HordeAgent.json");
LoadSettings();
_clientTask = BackgroundTask.StartNew(StatusTaskAsync);
_tickPauseStateTask = BackgroundTask.StartNew(ctx => TickPauseStateAsync(ctx));
}
public override bool Refresh()
{
if (OperatingSystem.IsWindows())
{
bool enabled = _pipeConnected || ((Registry.GetValue("HKEY_LOCAL_MACHINE\\SOFTWARE\\Epic Games\\Horde\\Agent", "Installed", null) as int?) ?? 0) != 0;
if (enabled != IsEnabled)
{
IsEnabled = enabled;
return true;
}
}
return false;
}
static JsonSerializerOptions GetJsonSerializerOptions()
{
JsonSerializerOptions options = new JsonSerializerOptions();
options.PropertyNamingPolicy = JsonNamingPolicy.CamelCase;
options.PropertyNameCaseInsensitive = true;
options.DefaultIgnoreCondition = JsonIgnoreCondition.WhenWritingNull;
options.AllowTrailingCommas = true;
options.WriteIndented = true;
options.Converters.Add(new JsonStringEnumConverter());
return options;
}
void LoadSettings()
{
if (FileReference.Exists(_settingsFile))
{
try
{
byte[] data = FileReference.ReadAllBytes(_settingsFile);
if (!data.SequenceEqual(_settingsData))
{
_settings = JsonSerializer.Deserialize<HordeAgentSettings>(data, GetJsonSerializerOptions())!;
if (_settings.Mode == null)
{
_settings.Mode = ReadLegacyMode();
UpdateSettings(_settings);
}
_settingsData = data;
return;
}
}
catch (Exception ex)
{
_logger.LogError(ex, "Error while reading {File}: {Message}", _settingsFile, ex.Message);
}
}
else
{
DirectoryReference? programData = DirectoryReference.GetSpecialFolder(Environment.SpecialFolder.CommonApplicationData);
if (programData != null)
{
FileReference legacySettingsFile = FileReference.Combine(programData, "Epic", "Horde", "TrayApp", "settings.json");
if (FileReference.Exists(legacySettingsFile))
{
try
{
byte[] data = FileReference.ReadAllBytes(legacySettingsFile);
HordeAgentSettings settings = JsonSerializer.Deserialize<HordeAgentSettings>(data, GetJsonSerializerOptions())!;
settings.Mode ??= ReadLegacyMode();
UpdateSettings(settings);
return;
}
catch (Exception ex)
{
_logger.LogError(ex, "Error while reading {File}: {Message}", legacySettingsFile, ex.Message);
}
}
}
}
UpdateSettings(new HordeAgentSettings());
}
static AgentMode? ReadLegacyMode()
{
if (!OperatingSystem.IsWindows())
{
return null;
}
const string RegistryKey = "HKEY_CURRENT_USER\\Software\\Epic Games\\Horde\\TrayApp";
const string RegistryStatusValue = "Status";
int? status = Registry.GetValue(RegistryKey, RegistryStatusValue, null) as int?;
if (status == null)
{
return null;
}
return status.Value switch
{
0 => AgentMode.Dedicated,
1 => AgentMode.Disabled,
2 => AgentMode.Workstation,
_ => null
};
}
public override void PopulateContextMenu(NativeMenu contextMenu)
{
if (IsEnabled)
{
NativeMenuItem enrollMenuItem = new NativeMenuItem("Enroll with Server...");
enrollMenuItem.Click += (s, e) => EnrollWithServer();
contextMenu.Items.Add(enrollMenuItem);
}
}
public override async ValueTask DisposeAsync()
{
await base.DisposeAsync();
await _tickPauseStateTask.DisposeAsync();
await _clientTask.DisposeAsync();
}
private void OnOpenLogs(object? sender, EventArgs e)
{
DirectoryReference? programDataDir = DirectoryReference.GetSpecialFolder(Environment.SpecialFolder.CommonApplicationData);
if (programDataDir != null)
{
DirectoryReference logsDir = DirectoryReference.Combine(programDataDir, "Epic", "Horde", "Agent");
if (DirectoryReference.Exists(logsDir))
{
Process.Start(new ProcessStartInfo { FileName = logsDir.FullName, UseShellExecute = true });
}
}
}
void SetStatus(AgentStatusMessage status)
{
if (!IsEnabled)
{
_status = null;
}
else if (!status.Healthy)
{
string message = String.IsNullOrEmpty(status.Detail) ? "Error. Check logs." : status.Detail.Length > 100 ? status.Detail.Substring(0, 100) : status.Detail;
_status = new ToolboxPluginStatus(TrayAppPluginState.Error, message);
if (_settings.Mode != AgentMode.Disabled)
{
// Make some known error messages more friendly
if (message.Contains("actively refused it", StringComparison.OrdinalIgnoreCase))
{
message = $"Could not connect to Horde Server: {HordeOptions.GetDefaultServerUrl()?.ToString() ?? "(Not configured)"}";
}
if (message.Contains("enrollment key does not match", StringComparison.OrdinalIgnoreCase))
{
message = $"Agent registration revoked by Horde Server: {HordeOptions.GetDefaultServerUrl()?.ToString() ?? "(Not configured)"}";
}
ToolboxNotificationManager.PostNotification("Horde Agent", message);
}
}
else if (status.NumLeases > 0)
{
string message = status.NumLeases == 1 ? "Currently handling 1 lease" : $"Currently handling {status.NumLeases} leases";
_status = new ToolboxPluginStatus(TrayAppPluginState.Busy, message);
}
else if (_enabled)
{
string message = "Agent is operating normally";
_status = new ToolboxPluginStatus(TrayAppPluginState.Ok, message);
}
else
{
string message = "Agent is paused";
_status = new ToolboxPluginStatus(TrayAppPluginState.Paused, message);
}
_host.UpdateStatus();
}
public override ToolboxPluginStatus GetStatus()
{
if (_settings.Mode == AgentMode.Disabled)
{
return ToolboxPluginStatus.Default;
}
if (!_toolCatalog.Items.Any(x => x.Pending != null && !x.Pending.Failed))
{
_reportStatus = _status;
}
return _reportStatus ?? ToolboxPluginStatus.Default;
}
async Task StatusTaskAsync(CancellationToken cancellationToken)
{
SetStatus(new AgentStatusMessage(true, 0, AgentStatusMessage.Starting));
for (; ; )
{
try
{
try
{
await PollForStatusUpdatesAsync(cancellationToken);
}
finally
{
_pipeConnected = false;
}
}
catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested)
{
break;
}
catch
{
SetStatus(new AgentStatusMessage(false, 0, "Unable to connect to Agent. Check the service is running."));
await Task.Delay(TimeSpan.FromSeconds(5.0), cancellationToken);
}
}
}
#pragma warning disable IDE1006
[StructLayout(LayoutKind.Sequential)]
struct LASTINPUTINFO
{
public int cbSize;
public uint dwTime;
}
[DllImport("user32.dll")]
static extern bool GetLastInputInfo(ref LASTINPUTINFO plii);
[DllImport("kernel32.dll")]
static extern uint GetTickCount();
[StructLayout(LayoutKind.Sequential)]
struct FILETIME
{
public uint dwLowDateTime;
public uint dwHighDateTime;
public readonly ulong Total => dwLowDateTime | (ulong)dwHighDateTime << 32;
};
[StructLayout(LayoutKind.Sequential)]
struct MEMORYSTATUSEX
{
public int dwLength;
public uint dwMemoryLoad;
public ulong ullTotalPhys;
public ulong ullAvailPhys;
public ulong ullTotalPageFile;
public ulong ullAvailPageFile;
public ulong ullTotalVirtual;
public ulong ullAvailVirtual;
public ulong ullAvailExtendedVirtual;
}
[DllImport("kernel32.dll", CharSet = CharSet.Auto, SetLastError = true)]
static extern bool GlobalMemoryStatusEx(ref MEMORYSTATUSEX lpBuffer);
[DllImport("kernel32.dll", SetLastError = true)]
static extern bool GetSystemTimes(out FILETIME lpIdleTime, out FILETIME lpKernelTime, out FILETIME lpUserTime);
#pragma warning restore IDE1006
bool _enabled;
readonly AsyncEvent _statusChangedEvent = new AsyncEvent();
readonly AsyncEvent _enabledChangedEvent = new AsyncEvent();
async Task TickPauseStateAsync(CancellationToken cancellationToken)
{
await using BackgroundTask cpuStatsTask = BackgroundTask.StartNew(ctx => TickCpuStatsAsync(ctx));
await using BackgroundTask criticalProcessTask = BackgroundTask.StartNew(ctx => TickCriticalProcessAsync(ctx));
TimeSpan pollInterval = TimeSpan.FromSeconds(0.25);
Stopwatch stateChangeTimer = Stopwatch.StartNew();
while (!cancellationToken.IsCancellationRequested)
{
Task statusChangedTask = _statusChangedEvent.Task;
AgentMode mode = _settings.Mode ?? AgentMode.Disabled;
if (mode == AgentMode.Dedicated)
{
if (!_enabled)
{
_enabled = true;
_enabledChangedEvent.Set();
}
}
else if (mode == AgentMode.Disabled)
{
if (_enabled)
{
_enabled = false;
_enabledChangedEvent.Set();
}
}
DateTime utcNow = DateTime.UtcNow;
IEnumerable<IdleStat> idleStats = GetIdleStats();
bool idle = idleStats.All(x => x.Value >= x.MinValue);
if (idle == _enabled)
{
stateChangeTimer.Restart();
}
const int WakeTimeSecs = 2;
const int IdleTimeSecs = 30;
int stateChangeTime = (int)stateChangeTimer.Elapsed.TotalSeconds;
int stateChangeMaxTime = _enabled ? WakeTimeSecs : IdleTimeSecs;
// _idleForm?.TickStats(_enabled, stateChangeTime, stateChangeMaxTime, idleStats);
if (mode == AgentMode.Workstation && stateChangeTime >= stateChangeMaxTime)
{
_enabled ^= true;
_enabledChangedEvent.Set();
stateChangeTimer.Restart();
}
await Task.WhenAny(statusChangedTask, Task.Delay(pollInterval, cancellationToken));
}
}
IEnumerable<IdleStat> GetIdleStats()
{
// Check there has been no input for a while
LASTINPUTINFO lastInputInfo = new LASTINPUTINFO();
lastInputInfo.cbSize = Marshal.SizeOf<LASTINPUTINFO>();
if (GetLastInputInfo(ref lastInputInfo))
{
yield return new IdleStat("LastInputTime", (GetTickCount() - lastInputInfo.dwTime) / 1000, _settings.Idle.MinIdleTimeSecs);
}
// Check that no critical processes are running
if (_settings.Idle.CriticalProcesses.Any())
{
yield return new IdleStat("CriticalProcCount", -_idleCriticalProcessCount, 0);
}
// Only look at memory/CPU usage if we're not paused; executing jobs will increase them
if (!_enabled)
{
// Check the CPU usage doesn't exceed the limit
yield return new IdleStat("IdleCpuPct", _idleCpuPct, _settings.Idle.MinIdleCpuPct);
// Check there's enough available virtual memory
MEMORYSTATUSEX memoryStatus = new MEMORYSTATUSEX();
memoryStatus.dwLength = Marshal.SizeOf<MEMORYSTATUSEX>();
if (GlobalMemoryStatusEx(ref memoryStatus))
{
yield return new IdleStat("VirtualMemMb", (long)(memoryStatus.ullAvailPhys + memoryStatus.ullAvailPageFile) / (1024 * 1024), _settings.Idle.MinFreeVirtualMemMb);
}
}
}
int _idleCpuPct = 0;
int _idleCriticalProcessCount = 0;
async Task TickCpuStatsAsync(CancellationToken cancellationToken)
{
const int NumSamples = 10;
TimeSpan sampleInterval = TimeSpan.FromSeconds(0.2);
(ulong IdleTime, ulong TotalTime)[] samples = new (ulong IdleTime, ulong TotalTime)[NumSamples];
int sampleIdx = 0;
for (; ; )
{
if (GetSystemTimes(out FILETIME idleTime, out FILETIME kernelTime, out FILETIME userTime))
{
(ulong prevIdleTime, ulong prevTotalTime) = samples[sampleIdx];
(ulong nextIdleTime, ulong nextTotalTime) = (idleTime.Total, kernelTime.Total + userTime.Total);
samples[sampleIdx] = (nextIdleTime, nextTotalTime);
sampleIdx = (sampleIdx + 1) % NumSamples;
if (prevTotalTime > 0 && nextTotalTime > prevTotalTime)
{
_idleCpuPct = (int)((nextIdleTime - prevIdleTime) * 100 / (nextTotalTime - prevTotalTime));
}
}
await Task.Delay(sampleInterval, cancellationToken);
}
}
async Task TickCriticalProcessAsync(CancellationToken cancellationToken)
{
TimeSpan sampleInterval = TimeSpan.FromSeconds(1.0);
for (; ; )
{
try
{
if (RuntimeInformation.IsOSPlatform(OSPlatform.Windows) && _settings.Idle.CriticalProcesses.Any())
{
IEnumerable<int> hordeProcessIds = Process.GetProcessesByName("HordeAgent").Select(x => x.Id);
IEnumerable<Process> criticalProcesses = _settings.Idle.CriticalProcesses
.Select(x => Path.GetFileNameWithoutExtension(x).ToUpperInvariant())
.Distinct()
.SelectMany(x => Process.GetProcessesByName(x))
.Where(x => !x.HasExited);
// Ignore processes that are descendants of HordeAgent
if (hordeProcessIds.Any() && criticalProcesses.Any())
{
criticalProcesses = criticalProcesses
.Where(x => !ProcessUtils.GetAncestorProcesses(x)
.Select(x => x.Id).Intersect(hordeProcessIds).Any())
.Where(x => !x.HasExited);
}
_idleCriticalProcessCount = criticalProcesses.Count();
}
}
catch (InvalidOperationException)
{
// If a process stops running Process.Id will throw an exception
}
await Task.Delay(sampleInterval, cancellationToken);
}
}
async Task PollForStatusUpdatesAsync(CancellationToken cancellationToken)
{
AgentMessageBuffer message = new AgentMessageBuffer();
using (NamedPipeClientStream pipeClient = new NamedPipeClientStream(".", AgentMessagePipe.PipeName, PipeDirection.InOut))
{
SetStatus(new AgentStatusMessage(true, 0, "Connecting to agent..."));
await pipeClient.ConnectAsync(cancellationToken);
_pipeConnected = true;
SetStatus(new AgentStatusMessage(true, 0, "Waiting for status update."));
for (; ; )
{
Task idleChangeTask = _enabledChangedEvent.Task;
bool enabled = _enabled;
message.Set(AgentMessageType.SetEnabledRequest, new AgentEnabledMessage(enabled));
await message.SendAsync(pipeClient, cancellationToken);
if (_agentSettings == null)
{
message.Set(AgentMessageType.GetSettingsRequest);
await message.SendAsync(pipeClient, cancellationToken);
if (!await message.TryReadAsync(pipeClient, cancellationToken))
{
break;
}
if (message.Type == AgentMessageType.GetSettingsResponse)
{
_agentSettings = message.Parse<AgentSettingsMessage>();
}
}
message.Set(AgentMessageType.SetSettingsRequest, new AgentSetSettingsRequest(_settings.Cpu.CpuCount, _settings.Cpu.CpuMultiplier));
await message.SendAsync(pipeClient, cancellationToken);
message.Set(AgentMessageType.GetStatusRequest);
await message.SendAsync(pipeClient, cancellationToken);
if (!await message.TryReadAsync(pipeClient, cancellationToken))
{
break;
}
switch (message.Type)
{
case AgentMessageType.GetStatusResponse:
AgentStatusMessage status = message.Parse<AgentStatusMessage>();
SetStatus(status);
break;
}
await Task.WhenAny(idleChangeTask, Task.Delay(TimeSpan.FromSeconds(5.0), cancellationToken));
}
}
}
}
}