diff --git a/UET/Redpoint.Uet.Patching.Runtime/ConsoleUetPatchLogging.cs b/UET/Redpoint.Uet.Patching.Runtime/ConsoleUetPatchLogging.cs new file mode 100644 index 00000000..2f4149f0 --- /dev/null +++ b/UET/Redpoint.Uet.Patching.Runtime/ConsoleUetPatchLogging.cs @@ -0,0 +1,20 @@ +namespace Redpoint.Uet.Patching.Runtime +{ + internal class ConsoleUetPatchLogging : IUetPatchLogging + { + public void LogInfo(string message) + { + Console.WriteLine($"Redpoint.Uet.Patching.Runtime [info ] {message}"); + } + + public void LogWarning(string message) + { + Console.WriteLine($"Redpoint.Uet.Patching.Runtime [warn ] {message}"); + } + + public void LogError(string message) + { + Console.WriteLine($"Redpoint.Uet.Patching.Runtime [error] {message}"); + } + } +} diff --git a/UET/Redpoint.Uet.Patching.Runtime/IUetPatchLogging.cs b/UET/Redpoint.Uet.Patching.Runtime/IUetPatchLogging.cs new file mode 100644 index 00000000..a3fb99f6 --- /dev/null +++ b/UET/Redpoint.Uet.Patching.Runtime/IUetPatchLogging.cs @@ -0,0 +1,11 @@ +namespace Redpoint.Uet.Patching.Runtime +{ + internal interface IUetPatchLogging + { + void LogInfo(string message); + + void LogWarning(string message); + + void LogError(string message); + } +} diff --git a/UET/Redpoint.Uet.Patching.Runtime/Kubernetes/IKubernetesUbaConfig.cs b/UET/Redpoint.Uet.Patching.Runtime/Kubernetes/IKubernetesUbaConfig.cs new file mode 100644 index 00000000..9d24e704 --- /dev/null +++ b/UET/Redpoint.Uet.Patching.Runtime/Kubernetes/IKubernetesUbaConfig.cs @@ -0,0 +1,14 @@ +namespace Redpoint.Uet.Patching.Runtime.Kubernetes +{ + internal interface IKubernetesUbaConfig + { + string? Namespace { get; } + string? Context { get; } + string? SmbServer { get; } + string? SmbShare { get; } + string? SmbUsername { get; } + string? SmbPassword { get; } + + public const ulong MemoryBytesPerCore = 1610612736uL; + } +} diff --git a/UET/Redpoint.Uet.Patching.Runtime/Kubernetes/KubernetesNodeState.cs b/UET/Redpoint.Uet.Patching.Runtime/Kubernetes/KubernetesNodeState.cs new file mode 100644 index 00000000..f686447f --- /dev/null +++ b/UET/Redpoint.Uet.Patching.Runtime/Kubernetes/KubernetesNodeState.cs @@ -0,0 +1,106 @@ +namespace Redpoint.Uet.Patching.Runtime.Kubernetes +{ + using System; + using System.Collections.Generic; + using System.Linq; + using k8s.Models; + + internal class KubernetesNodeState + { + public string? NodeId; + public V1Node? KubernetesNode; + public List KubernetesPods = new List(); + public readonly List AllocatedBlocks = new List(); + + public ulong MemoryTotal + { + get + { + return KubernetesNode!.Status.Capacity["memory"].ToUInt64(); + } + } + + public ulong MemoryNonUba + { + get + { + return KubernetesPods + .Where(x => x.GetLabel("uba") != "true") + .SelectMany(x => x.Spec.Containers) + .Select(x => x?.Resources?.Requests != null && x.Resources.Requests.TryGetValue("memory", out var quantity) ? quantity.ToUInt64() : 0) + .Aggregate((a, b) => a + b); + } + } + + public ulong MemoryAllocated + { + get + { + return AllocatedBlocks + .Select(x => (ulong)x.AllocatedCores * IKubernetesUbaConfig.MemoryBytesPerCore) + .DefaultIfEmpty((ulong)0) + .Aggregate((a, b) => a + b); + } + } + + public ulong MemoryAvailable + { + get + { + var memory = MemoryTotal; + memory -= MemoryNonUba; + memory -= MemoryAllocated; + return Math.Max(0, memory); + } + } + + public double CoresTotal + { + get + { + return Math.Floor(KubernetesNode!.Status.Capacity["cpu"].ToDouble()); + } + } + + public double CoresNonUba + { + get + { + return KubernetesPods + .Where(x => x.GetLabel("uba") != "true") + .SelectMany(x => x.Spec.Containers) + .Select(x => x?.Resources?.Requests != null && x.Resources.Requests.TryGetValue("cpu", out var quantity) ? quantity.ToDouble() : 0) + .Sum(); + } + } + + public double CoresAllocated + { + get + { + return AllocatedBlocks.Sum(x => x.AllocatedCores); + } + } + + public int CoresAvailable + { + get + { + var cores = CoresTotal; + cores -= CoresNonUba; + cores -= CoresAllocated; + return (int)Math.Max(0, Math.Floor(cores)); + } + } + + public int CoresAllocatable + { + get + { + var realCoresAvailable = CoresAvailable; + var memoryConstrainedCoresAvailable = MemoryAvailable / IKubernetesUbaConfig.MemoryBytesPerCore; + return Math.Min(realCoresAvailable, (int)memoryConstrainedCoresAvailable); + } + } + } +} diff --git a/UET/Redpoint.Uet.Patching.Runtime/Kubernetes/KubernetesNodeWorker.cs b/UET/Redpoint.Uet.Patching.Runtime/Kubernetes/KubernetesNodeWorker.cs new file mode 100644 index 00000000..27f3d7e3 --- /dev/null +++ b/UET/Redpoint.Uet.Patching.Runtime/Kubernetes/KubernetesNodeWorker.cs @@ -0,0 +1,13 @@ +namespace Redpoint.Uet.Patching.Runtime.Kubernetes +{ + using k8s.Models; + + internal class KubernetesNodeWorker + { + public V1Pod? KubernetesPod; + public V1Service? KubernetesService; + public int AllocatedCores; + public string? UbaHost; + public int? UbaPort; + } +} diff --git a/UET/Redpoint.Uet.Patching.Runtime/Kubernetes/KubernetesUbaConfigFromHook.cs b/UET/Redpoint.Uet.Patching.Runtime/Kubernetes/KubernetesUbaConfigFromHook.cs new file mode 100644 index 00000000..b7780b81 --- /dev/null +++ b/UET/Redpoint.Uet.Patching.Runtime/Kubernetes/KubernetesUbaConfigFromHook.cs @@ -0,0 +1,34 @@ +namespace Redpoint.Uet.Patching.Runtime.Kubernetes +{ + using System.Reflection; + + internal class KubernetesUbaConfigFromHook : IKubernetesUbaConfig + { + private readonly object _ubaKubeConfig; + + public KubernetesUbaConfigFromHook(object ubtHookObject) + { + var field = ubtHookObject.GetType().GetField("_ubaKubeConfig", BindingFlags.NonPublic | BindingFlags.Instance); + _ubaKubeConfig = field!.GetValue(ubtHookObject)!; + } + + private string? GetPropertyValue(string property) + { + return _ubaKubeConfig.GetType() + .GetProperty(property, BindingFlags.Public | BindingFlags.Instance) + ?.GetValue(_ubaKubeConfig) as string; + } + + public string? Namespace => GetPropertyValue("Namespace"); + + public string? Context => GetPropertyValue("Context"); + + public string? SmbServer => GetPropertyValue("SmbServer"); + + public string? SmbShare => GetPropertyValue("SmbShare"); + + public string? SmbUsername => GetPropertyValue("SmbUsername"); + + public string? SmbPassword => GetPropertyValue("SmbPassword"); + } +} diff --git a/UET/Redpoint.Uet.Patching.Runtime/Kubernetes/KubernetesUbaCoordinator.cs b/UET/Redpoint.Uet.Patching.Runtime/Kubernetes/KubernetesUbaCoordinator.cs new file mode 100644 index 00000000..595f6647 --- /dev/null +++ b/UET/Redpoint.Uet.Patching.Runtime/Kubernetes/KubernetesUbaCoordinator.cs @@ -0,0 +1,569 @@ +namespace Redpoint.Uet.Patching.Runtime.Kubernetes +{ + using System; + using System.Collections.Generic; + using System.Linq; + using System.Net; + using System.Reflection; + using System.Runtime.CompilerServices; + using System.Runtime.InteropServices; + using System.Text; + using System.Threading.Tasks; + using k8s; + using k8s.Autorest; + using k8s.Models; + + internal class KubernetesUbaCoordinator : IDisposable + { + private readonly object _ubtHookObject; + private readonly IKubernetesUbaConfig _ubaKubeConfig; + private readonly CancellationTokenSource _cancellationSource; + private readonly Dictionary _kubernetesNodes; + private readonly string _id; + private string? _ubaAgentRemotePath; + private string? _ubaAgentHash; + private Kubernetes? _client; + + public KubernetesUbaCoordinator(object ubtHookObject) + { + _ubtHookObject = ubtHookObject; + + // Create our config from the hook object. + _ubaKubeConfig = new KubernetesUbaConfigFromHook(ubtHookObject); + + // Set the cancellation source for cancelling the work. + _cancellationSource = new CancellationTokenSource(); + + // Initialize our dictionary that tracks Kubernetes nodes. + _kubernetesNodes = new Dictionary(); + + // Generate an ID to identify the jobs we're running. + _id = Guid.NewGuid().ToString(); + } + + private string KubernetesNamespace + { + get + { + return _ubaKubeConfig?.Namespace ?? "default"; + } + } + + public CancellationTokenSource CancellationTokenSource => _cancellationSource; + + public bool ClientIsAvailable + { + get + { + return _client != null; + } + } + + private void LogInformation(string message) + { + _ubtHookObject.GetType() + .GetMethod("LogInformation", BindingFlags.Public | BindingFlags.Instance)! + .Invoke(_ubtHookObject, new object[] + { + message + }); + } + + private void AddUbaClient(string ubaHost, int ubaPort) + { + _ubtHookObject.GetType() + .GetMethod("AddUbaClient", BindingFlags.Public | BindingFlags.Instance)! + .Invoke(_ubtHookObject, new object[] + { + ubaHost, + ubaPort + }); + } + + public void CopyAgentFileToShare(string ubaAgentExe, string ubaAgentHash) + { + LogInformation("Kubernetes UBA: Copying agent file to network share..."); + var ubaRemoteDir = Path.Combine(@$"\\{_ubaKubeConfig.SmbServer}\{_ubaKubeConfig.SmbShare}\Uba", ubaAgentExe); + var ubaRemoteFile = Path.Combine(ubaRemoteDir, "UbaAgent.exe"); + try + { + Directory.CreateDirectory(ubaRemoteDir); + File.Copy(ubaAgentExe, ubaRemoteFile); + } + catch + { + // File already copied. + } + _ubaAgentRemotePath = ubaRemoteFile; + _ubaAgentHash = ubaAgentHash; + LogInformation("Kubernetes UBA: Copied agent file to network share."); + } + + public async Task ConnectToClusterAsync() + { + LogInformation("Kubernetes UBA: Connecting to cluster..."); + var config = KubernetesClientConfiguration.BuildConfigFromConfigFile(currentContext: _ubaKubeConfig.Context); + _client = new Kubernetes(config); + await _client.ListNamespacedPodAsync(KubernetesNamespace, cancellationToken: _cancellationSource.Token); + LogInformation("Kubernetes UBA: Connected to cluster."); + } + + #region Kubernetes Helpers + + private async Task DeleteServiceAndPodAsync(V1Pod pod, CancellationToken? cancellationToken) + { + try + { + await _client.DeleteNamespacedServiceAsync(pod.Name(), pod.Namespace(), cancellationToken: cancellationToken ?? _cancellationSource?.Token ?? CancellationToken.None); + } + catch (HttpOperationException ex) when (ex.Response.StatusCode == HttpStatusCode.NotFound) + { + } + try + { + await _client.DeleteNamespacedPodAsync(pod.Name(), pod.Namespace(), cancellationToken: cancellationToken ?? _cancellationSource?.Token ?? CancellationToken.None); + } + catch (HttpOperationException ex) when (ex.Response.StatusCode == HttpStatusCode.NotFound) + { + } + if (_kubernetesNodes.TryGetValue(pod.Spec.NodeName, out var node)) + { + node.AllocatedBlocks.RemoveAll(x => x.KubernetesPod.Name() == pod.Name()); + } + } + + private async IAsyncEnumerable EnumerateNodePodAsync(string nodeName, [EnumeratorCancellation] CancellationToken cancellationToken) + { + string? continueParameter = null; + do + { + cancellationToken.ThrowIfCancellationRequested(); + var list = await _client.ListPodForAllNamespacesAsync(fieldSelector: $"spec.nodeName={nodeName}", continueParameter: continueParameter, cancellationToken: cancellationToken); + foreach (var item in list.Items) + { + yield return item; + } + continueParameter = list.Metadata.ContinueProperty; + } while (!string.IsNullOrWhiteSpace(continueParameter)); + } + + private async IAsyncEnumerable EnumerateNamespacedPodAsync(string labelSelector, [EnumeratorCancellation] CancellationToken cancellationToken) + { + string? continueParameter = null; + do + { + cancellationToken.ThrowIfCancellationRequested(); + var list = await _client.ListNamespacedPodAsync(KubernetesNamespace, labelSelector: labelSelector, continueParameter: continueParameter, cancellationToken: cancellationToken); + foreach (var item in list.Items) + { + yield return item; + } + continueParameter = list.Metadata.ContinueProperty; + } while (!string.IsNullOrWhiteSpace(continueParameter)); + } + + private async IAsyncEnumerable EnumerateNamespacedServiceAsync(string labelSelector, [EnumeratorCancellation] CancellationToken cancellationToken) + { + string? continueParameter = null; + do + { + cancellationToken.ThrowIfCancellationRequested(); + var list = await _client.ListNamespacedServiceAsync(KubernetesNamespace, labelSelector: labelSelector, continueParameter: continueParameter, cancellationToken: cancellationToken); + foreach (var item in list.Items) + { + yield return item; + } + continueParameter = list.Metadata.ContinueProperty; + } while (!string.IsNullOrWhiteSpace(continueParameter)); + } + + private async IAsyncEnumerable EnumerateNodesAsync(string labelSelector, [EnumeratorCancellation] CancellationToken cancellationToken) + { + string? continueParameter = null; + do + { + cancellationToken.ThrowIfCancellationRequested(); + var list = await _client.ListNodeAsync(labelSelector: labelSelector, continueParameter: continueParameter, cancellationToken: cancellationToken); + foreach (var item in list.Items) + { + if (item.Spec.Taints != null && ( + item.Spec.Taints.Any(x => x.Effect == "NoSchedule") || + item.Spec.Taints.Any(x => x.Effect == "NoExecute"))) + { + continue; + } + yield return item; + } + continueParameter = list.Metadata.ContinueProperty; + } while (!string.IsNullOrWhiteSpace(continueParameter)); + } + + #endregion + + public async Task CleanupFinishedKubernetesPodsGloballyAsync() + { + await foreach (var pod in EnumerateNamespacedPodAsync("uba=true", _cancellationSource.Token)) + { + if (pod.Status.Phase == "Succeeded" || pod.Status.Phase == "Failed" || pod.Status.Phase == "Unknown") + { + LogInformation($"Removing Kubernetes block: {pod.Name()} (cleanup)"); + await DeleteServiceAndPodAsync(pod, _cancellationSource.Token); + } + } + } + + public async Task CleanupFinishedKubernetesPodsLocallyAsync() + { + await foreach (var pod in EnumerateNamespacedPodAsync($"uba=true,uba.queueId={_id}", _cancellationSource.Token)) + { + if (pod.Status.Phase == "Succeeded" || pod.Status.Phase == "Failed" || pod.Status.Phase == "Unknown") + { + LogInformation($"Removing Kubernetes block: {pod.Name()}"); + await DeleteServiceAndPodAsync(pod, _cancellationSource.Token); + } + } + } + + public int GetAllocatedBlocks() + { + return _kubernetesNodes.SelectMany(x => x.Value.AllocatedBlocks).Sum(y => y.AllocatedCores); + } + + public async Task SynchroniseKubernetesNodesAsync() + { + var knownNodeNames = new HashSet(); + await foreach (var node in EnumerateNodesAsync("kubernetes.io/os=windows", _cancellationSource.Token)) + { + knownNodeNames.Add(node.Name()); + if (_kubernetesNodes.TryGetValue(node.Name(), out var existingEntry)) + { + existingEntry.KubernetesNode = node; + } + else + { + existingEntry = new KubernetesNodeState + { + NodeId = node.Name(), + KubernetesNode = node, + }; + _kubernetesNodes.Add(node.Name(), existingEntry); + } + var newPodsList = new List(); + await foreach (var pod in EnumerateNodePodAsync(node.Name(), _cancellationSource.Token)) + { + newPodsList.Add(pod); + } + existingEntry.KubernetesPods = newPodsList; + } + var currentNodeNames = new HashSet(_kubernetesNodes.Keys); + foreach (var current in currentNodeNames) + { + if (!knownNodeNames.Contains(current)) + { + _kubernetesNodes.Remove(current); + } + } + } + + public int GetMaximumBlockSize(double desiredCpus) + { + var blockSize = _kubernetesNodes + .Select(x => x.Value.CoresAllocatable) + .Where(x => x != 0) + .DefaultIfEmpty(0) + .Max(); + blockSize = Math.Min(blockSize, (int)Math.Floor(desiredCpus)); + return blockSize; + } + + public async Task TryAllocateKubernetesNodeAsync(int blockSize) + { + // Pick an available node, weighted by the available cores. + var selectedNode = _kubernetesNodes + .Select(x => x.Value) + .SelectMany(x => + { + var blocks = new List(); + LogInformation($"Kubernetes UBA: Loop: Node '{x.NodeId}' has\nCPU: {x.CoresTotal} total, {x.CoresNonUba} non-UBA, {x.CoresAllocated} allocated, {x.CoresAvailable} available, {x.CoresAllocatable} allocatable.\nMemory: {x.MemoryTotal} total, {x.MemoryNonUba} non-UBA, {x.MemoryAllocated} allocated, {x.MemoryAvailable} available.\nBlocks: {x.AllocatedBlocks.Count} allocated."); + for (var c = 0; c < x.CoresAllocatable; c += blockSize) + { + if (c + blockSize <= x.CoresAllocatable) + { + blocks.Add(x); + } + } + return blocks; + }) + .OrderBy(_ => Random.Shared.NextInt64()) + .FirstOrDefault(); + + // If we don't have any available node, break out of the loop node. + if (selectedNode == null) + { + return false; + } + + // Generate the next block ID for this node. + int highestBlockId = 0; + await foreach (var pod in EnumerateNamespacedPodAsync($"uba.nodeId={selectedNode.NodeId}", _cancellationSource.Token)) + { + var thisBlockId = int.Parse(pod.GetLabel("uba.blockId")); + highestBlockId = Math.Max(highestBlockId, thisBlockId); + } + var nextBlockId = highestBlockId + 1; + + // Create the pod and service. + var name = $"uba-{selectedNode.NodeId}-{nextBlockId}"; + LogInformation($"Allocating Kubernetes block: {name}"); + var labels = new Dictionary + { + { "uba", "true" }, + { "uba.nodeId", selectedNode.NodeId! }, + { "uba.blockId", nextBlockId.ToString() }, + { "uba.queueId", _id! }, + }; + var kubernetesPod = await _client.CreateNamespacedPodAsync( + new V1Pod + { + Metadata = new V1ObjectMeta + { + Name = name, + Labels = labels, + }, + Spec = new V1PodSpec + { + AutomountServiceAccountToken = false, + NodeSelector = new Dictionary + { + { "kubernetes.io/os", "windows" }, + }, + RestartPolicy = "Never", + TerminationGracePeriodSeconds = 0, + Volumes = new List + { + new V1Volume + { + Name = "uba-storage", + HostPath = new V1HostPathVolumeSource + { + Path = @$"C:\Uba\{name}", + } + } + }, + Containers = new List + { + new V1Container + { + Image = "mcr.microsoft.com/powershell:lts-windowsservercore-ltsc2022", + ImagePullPolicy = "IfNotPresent", + Name = "uba-agent", + Resources = new V1ResourceRequirements + { + Requests = new Dictionary + { + { "cpu", new ResourceQuantity(blockSize.ToString()) }, + // @note: We don't set 'memory' here because it can be finicky to get the container to get allocated. + }, + Limits = new Dictionary + { + { "cpu", new ResourceQuantity(blockSize.ToString()) }, + // @note: We don't set 'memory' here because it can be finicky to get the container to get allocated. + }, + }, + SecurityContext = new V1SecurityContext + { + WindowsOptions = new V1WindowsSecurityContextOptions + { + RunAsUserName = "ContainerAdministrator", + } + }, + Command = new List + { + @"C:\Program Files\PowerShell\latest\pwsh.exe", + }, + Args = new List + { + "-Command", + $@"Start-Sleep -Seconds 1; Write-Host ""Mapping network drive...""; C:\Windows\system32\net.exe use Z: \\{_ubaKubeConfig.SmbServer}\{_ubaKubeConfig.SmbShare}\Uba /USER:{_ubaKubeConfig.SmbUsername} {_ubaKubeConfig.SmbPassword}; Write-Host ""Copying UBA agent...""; Copy-Item Z:\{_ubaAgentHash}\UbaAgent.exe C:\UbaAgent.exe; Write-Host ""Running UBA agent...""; C:\UbaAgent.exe -Verbose -Listen=7000 -NoPoll -listenTimeout=120 -ProxyPort=7001 -Dir=C:\UbaData -MaxIdle=15 -MaxCpu={blockSize}; Write-Host ""UBA agent exited with exit code: $LastExitCode""; exit $LastExitCode;", + }, + Ports = new List + { + new V1ContainerPort + { + ContainerPort = 7000, + Protocol = "TCP", + }, + new V1ContainerPort + { + ContainerPort = 7001, + Protocol = "TCP", + } + }, + VolumeMounts = new List + { + new V1VolumeMount + { + Name = "uba-storage", + MountPath = @"C:\UbaData", + } + } + } + } + } + }, + KubernetesNamespace, + cancellationToken: _cancellationSource.Token); + var kubernetesService = await _client.CreateNamespacedServiceAsync( + new V1Service + { + Metadata = new V1ObjectMeta + { + Name = name, + Labels = labels, + }, + Spec = new V1ServiceSpec + { + Selector = labels, + Type = "NodePort", + Ports = new List + { + new V1ServicePort + { + Name = "uba", + Port = 7000, + TargetPort = new IntstrIntOrString("7000"), + Protocol = "TCP", + }, + new V1ServicePort + { + Name = "uba-proxy", + Port = 7001, + TargetPort = new IntstrIntOrString("7001"), + Protocol = "TCP", + }, + }, + }, + }, + KubernetesNamespace, + cancellationToken: _cancellationSource.Token); + + // Track the worker. + var worker = new KubernetesNodeWorker + { + KubernetesPod = kubernetesPod, + KubernetesService = kubernetesService, + AllocatedCores = blockSize, + }; + selectedNode.AllocatedBlocks.Add(worker); + + // In the background, wait for the worker to become ready and allocate it to UBA. + _ = Task.Run(async () => + { + var didRegister = false; + try + { + // Wait for the service to have node ports. + while (worker.UbaHost == null || worker.UbaPort == null) + { + _cancellationSource.Token.ThrowIfCancellationRequested(); + + // Refresh service status. + worker.KubernetesService = await _client.ReadNamespacedServiceAsync( + worker.KubernetesService.Name(), + worker.KubernetesService.Namespace(), + cancellationToken: _cancellationSource.Token); + + // If a port doesn't have NodePort, it's not allocated yet. + if (worker.KubernetesService.Spec.Ports.Any(x => x.NodePort == null)) + { + await Task.Delay(1000); + continue; + } + + // We should have the node port now. + worker.UbaHost = selectedNode.KubernetesNode!.Status.Addresses + .Where(x => x.Type == "InternalIP") + .Select(x => x.Address) + .First(); + worker.UbaPort = worker.KubernetesService.Spec.Ports.First(x => x.Name == "uba").NodePort!.Value; + break; + } + + // Wait for the pod to start. + var secondsElapsed = 0; + while (worker.KubernetesPod.Status.Phase == "Pending" && secondsElapsed < 30) + { + await Task.Delay(1000); + worker.KubernetesPod = await _client.ReadNamespacedPodAsync( + worker.KubernetesPod.Name(), + worker.KubernetesPod.Namespace(), + cancellationToken: _cancellationSource.Token); + secondsElapsed++; + } + if (worker.KubernetesPod.Status.Phase == "Pending") + { + // Timed out. + LogInformation($"Kubernetes timed out while allocating: {name}"); + return; + } + + // Add the worker to UBA. + LogInformation($"Kubernetes block is ready: {name} ({worker.UbaHost}:{worker.UbaPort.Value})"); + AddUbaClient(worker.UbaHost, worker.UbaPort.Value); + didRegister = true; + } + catch (Exception ex) + { + LogInformation($"Exception in Kubernetes wait loop: {ex}"); + } + finally + { + if (!didRegister) + { + // This pod/service could not be allocated. The main loop can then try to allocate again. + await DeleteServiceAndPodAsync(worker.KubernetesPod, _cancellationSource.Token); + } + } + }); + return true; + } + + public async Task CloseAsync() + { + // Remove any Kubernetes pods that are complete. + await foreach (var pod in EnumerateNamespacedPodAsync("uba=true", CancellationToken.None)) + { + if (pod.Status.Phase == "Succeeded" || pod.Status.Phase == "Failed" || pod.Status.Phase == "Unknown") + { + // _logger.LogInformation($"Removing Kubernetes block: {pod.Name()} (cleanup on close)"); + await DeleteServiceAndPodAsync(pod, CancellationToken.None); + } + } + + // Remove any Kubernetes pods owned by us. + if (!string.IsNullOrWhiteSpace(_id)) + { + await foreach (var pod in EnumerateNamespacedPodAsync($"uba.queueId={_id}", CancellationToken.None)) + { + // _logger.LogInformation($"Removing Kubernetes block: {pod.Name()} (unconditional)"); + await DeleteServiceAndPodAsync(pod, CancellationToken.None); + } + } + } + + public void Dispose() + { + CloseAsync().Wait(); + Dispose(true); + GC.SuppressFinalize(this); + } + + protected virtual void Dispose(bool disposing) + { + if (disposing) + { + _cancellationSource?.Dispose(); + } + } + } +} diff --git a/UET/Redpoint.Uet.Patching.Runtime/NullUetPatchLogging.cs b/UET/Redpoint.Uet.Patching.Runtime/NullUetPatchLogging.cs new file mode 100644 index 00000000..f3d84b8c --- /dev/null +++ b/UET/Redpoint.Uet.Patching.Runtime/NullUetPatchLogging.cs @@ -0,0 +1,17 @@ +namespace Redpoint.Uet.Patching.Runtime +{ + internal class NullUetPatchLogging : IUetPatchLogging + { + public void LogInfo(string message) + { + } + + public void LogWarning(string message) + { + } + + public void LogError(string message) + { + } + } +} diff --git a/UET/Redpoint.Uet.Patching.Runtime/Patches/BuildGraphProjectRootUetPatch.cs b/UET/Redpoint.Uet.Patching.Runtime/Patches/BuildGraphProjectRootUetPatch.cs new file mode 100644 index 00000000..9d08e73c --- /dev/null +++ b/UET/Redpoint.Uet.Patching.Runtime/Patches/BuildGraphProjectRootUetPatch.cs @@ -0,0 +1,17 @@ +namespace Redpoint.Uet.Patching.Runtime.Patches +{ + using HarmonyLib; + using System.Reflection; + + internal class BuildGraphProjectRootUetPatch : IUetPatch + { + public bool ShouldApplyPatch() + { + return Assembly.GetEntryAssembly()?.GetName()?.Name == "AutomationTool"; + } + + public void ApplyPatch(IUetPatchLogging logging, Harmony harmony) + { + } + } +} diff --git a/UET/Redpoint.Uet.Patching.Runtime/Patches/KubernetesUbaCoordinatorPatchCode.cs b/UET/Redpoint.Uet.Patching.Runtime/Patches/KubernetesUbaCoordinatorPatchCode.cs new file mode 100644 index 00000000..3118f839 --- /dev/null +++ b/UET/Redpoint.Uet.Patching.Runtime/Patches/KubernetesUbaCoordinatorPatchCode.cs @@ -0,0 +1,410 @@ +using EpicGames.Core; +using Microsoft.Extensions.Logging; +using System; +using System.Collections.Generic; +using System.Linq; +using System.Runtime.InteropServices; +using System.Threading; +using System.Threading.Tasks; +using UnrealBuildBase; +using System.Reflection; +using Timer = System.Threading.Timer; + +#nullable enable + +namespace UnrealBuildTool +{ + public static class UBAAgentCoordinatorKubernetesConstructor + { + static object Construct( + ILogger logger, + UnrealBuildAcceleratorConfig ubaConfig, + CommandLineArguments? additionalArguments = null) + { + return new UBAAgentCoordinatorKubernetes(logger, ubaConfig, additionalArguments); + } + } + + public class UnrealBuildAcceleratorKubernetesConfig + { + /// + /// Namespace to deploy Kubernetes pods into. + /// + [XmlConfigFile(Category = "Kubernetes", Name = "Namespace")] + [CommandLine("-KubernetesNamespace=")] + public string? Namespace { get; set; } + + /// + /// Context from the .kubeconfig file to use. + /// + [XmlConfigFile(Category = "Kubernetes", Name = "Context")] + [CommandLine("-KubernetesContext=")] + public string? Context { get; set; } + + /// + /// IP address of the network share to store the UBA agent on. This is used to transfer the UBA agent to containers. + /// + [XmlConfigFile(Category = "Kubernetes", Name = "SmbServer")] + [CommandLine("-KubernetesSmbServer=")] + public string? SmbServer { get; set; } + + /// + /// Share name to use to store the UBA agent. + /// + [XmlConfigFile(Category = "Kubernetes", Name = "SmbShare")] + [CommandLine("-KubernetesSmbShare=")] + public string? SmbShare { get; set; } + + /// + /// Username to use to connect to network share. + /// + [XmlConfigFile(Category = "Kubernetes", Name = "SmbUsername")] + [CommandLine("-KubernetesSmbUsername=")] + public string? SmbUsername { get; set; } + + /// + /// Password to use to connect to network share. + /// + [XmlConfigFile(Category = "Kubernetes", Name = "SmbPassword")] + [CommandLine("-KubernetesSmbPassword=")] + public string? SmbPassword { get; set; } + } + + class UBAAgentCoordinatorKubernetes : IUBAAgentCoordinator, IDisposable + { + private readonly ILogger _logger; + private readonly UnrealBuildAcceleratorConfig _ubaConfig; + private readonly UnrealBuildAcceleratorKubernetesConfig _ubaKubeConfig; + private readonly System.Type _realCoordinatorType; + + private readonly ConstructorInfo _realCoordinatorConstructor; + private readonly Dictionary _realCoordinatorProperties; + private readonly Dictionary _realCoordinatorMethods; + + private UBAExecutor? _executor; + private IDisposable? _realCoordinator; + private Timer? _timer; + private const int _timerPeriod = 5000; + private bool _forcedStop = false; + private int _quickRestartsRemaining = 10; + + public UBAAgentCoordinatorKubernetes( + ILogger logger, + UnrealBuildAcceleratorConfig ubaConfig, + CommandLineArguments? additionalArguments = null) + { + _logger = logger; + + _logger.LogInformation("Kubernetes UBA: Initializing configuration"); + _ubaConfig = ubaConfig; + _ubaKubeConfig = new UnrealBuildAcceleratorKubernetesConfig(); + XmlConfig.ApplyTo(_ubaKubeConfig); + additionalArguments?.ApplyTo(_ubaKubeConfig); + + _logger.LogInformation("Kubernetes UBA: Initializing real coordinator type"); + _realCoordinatorType = System.Runtime.Loader.AssemblyLoadContext.Default.Assemblies + .First(x => x.GetName()?.Name == "Redpoint.Uet.Patching.Runtime") + .GetType("Redpoint.Uet.Patching.Runtime.Kubernetes.KubernetesUbaCoordinator")!; + if (_realCoordinatorType == null) + { + throw new InvalidOperationException($"Unable to find type KubernetesUbaCoordinator!"); + } + + _realCoordinatorConstructor = _realCoordinatorType.GetConstructors().FirstOrDefault()!; + if (_realCoordinatorConstructor == null) + { + throw new InvalidOperationException($"Unable to find constructor in KubernetesUbaCoordinator!"); + } + + _realCoordinatorProperties = new Dictionary(); + _realCoordinatorMethods = new Dictionary(); + foreach (var property in new string[] + { + "CancellationTokenSource", + "ClientIsAvailable", + }) + { + var foundProperty = _realCoordinatorType.GetProperty(property, BindingFlags.Public | BindingFlags.Instance); + if (foundProperty == null) + { + throw new InvalidOperationException($"Unable to find property '{property}' in KubernetesUbaCoordinator!"); + } + _realCoordinatorProperties.Add(property, foundProperty); + } + foreach (var method in new string[] + { + "GetAllocatedBlocks", + "CopyAgentFileToShare", + "ConnectToClusterAsync", + "CleanupFinishedKubernetesPodsGloballyAsync", + "CleanupFinishedKubernetesPodsLocallyAsync", + "SynchroniseKubernetesNodesAsync", + "GetMaximumBlockSize", + "TryAllocateKubernetesNodeAsync", + "CloseAsync", + }) + { + var foundMethod = _realCoordinatorType.GetMethod(method, BindingFlags.Public | BindingFlags.Instance); + if (foundMethod == null) + { + throw new InvalidOperationException($"Unable to find method '{method}' in KubernetesUbaCoordinator!"); + } + _realCoordinatorMethods.Add(method, foundMethod); + } + } + + public DirectoryReference? GetUBARootDir() + { + return null; + } + + public void LogInformation(string message) + { + _logger.LogInformation(message); + } + + public void AddUbaClient(string host, int port) + { + _executor!.Server!.AddClient(host, port); + } + + private CancellationTokenSource? CancellationTokenSource + { + get + { + if (_realCoordinator != null) + { + return (CancellationTokenSource?)_realCoordinatorProperties["CancellationTokenSource"].GetValue(_realCoordinator); + } + else + { + return null; + } + } + } + + public async Task InitAsync(UBAExecutor executor) + { + try + { + _logger.LogInformation("Kubernetes UBA: InitAsync"); + + if (_ubaConfig.bDisableRemote) + { + _logger.LogInformation("Kubernetes UBA: Remoting is disabled, skipping initialization."); + return; + } + + if (string.IsNullOrWhiteSpace(_ubaKubeConfig.Namespace)) + { + _logger.LogWarning("Kubernetes UBA: Missing Kubernetes -> Namespace in BuildConfiguration.xml."); + return; + } + + if (string.IsNullOrWhiteSpace(_ubaKubeConfig.Context)) + { + _logger.LogWarning("Kubernetes UBA: Missing Kubernetes -> Context in BuildConfiguration.xml."); + return; + } + if (string.IsNullOrWhiteSpace(_ubaKubeConfig.SmbServer)) + { + _logger.LogWarning("Kubernetes UBA: Missing Kubernetes -> SmbServer in BuildConfiguration.xml."); + return; + } + if (string.IsNullOrWhiteSpace(_ubaKubeConfig.SmbShare)) + { + _logger.LogWarning("Kubernetes UBA: Missing Kubernetes -> SmbShare in BuildConfiguration.xml."); + return; + } + if (string.IsNullOrWhiteSpace(_ubaKubeConfig.SmbUsername)) + { + _logger.LogWarning("Kubernetes UBA: Missing Kubernetes -> SmbUsername in BuildConfiguration.xml."); + return; + } + if (string.IsNullOrWhiteSpace(_ubaKubeConfig.SmbPassword)) + { + _logger.LogWarning("Kubernetes UBA: Missing Kubernetes -> SmbPassword in BuildConfiguration.xml."); + return; + } + + _logger.LogInformation("Kubernetes UBA: Constructing real coordinator"); + if (_realCoordinator != null) + { + _realCoordinator.Dispose(); + } + _realCoordinator = (IDisposable)_realCoordinatorConstructor.Invoke(new object[] { this })!; + + _logger.LogInformation("Kubernetes UBA: Starting initialization"); + // Copy the UbaAgent file to the network share. + _logger.LogInformation("Kubernetes UBA: (Proxy) Copying file to network share..."); + var ubaDir = DirectoryReference.Combine(Unreal.EngineDirectory, "Binaries", "Win64", "UnrealBuildAccelerator", RuntimeInformation.ProcessArchitecture.ToString().ToLowerInvariant()); + var ubaFile = FileReference.Combine(ubaDir, "UbaAgent.exe"); + var agentHash = (await (new FileHasher()).GetDigestAsync(ubaFile)).ToString(); + _realCoordinatorMethods["CopyAgentFileToShare"].Invoke(_realCoordinator, new object[] { ubaFile.FullName, agentHash }); + + // Set up Kubernetes client and ensure we can connect to the cluster. + _logger.LogInformation("Kubernetes UBA: (Proxy) Connecting to cluster..."); + await (Task)_realCoordinatorMethods["ConnectToClusterAsync"].Invoke(_realCoordinator, new object[0])!; + + // Track the executor so we can add clients to it. + _logger.LogInformation("Kubernetes UBA: Tracking executor..."); + _executor = executor; + + _logger.LogInformation("Kubernetes UBA: Ready to run!"); + } + catch (Exception ex) + { + _logger.LogError(ex, $"Kubernetes UBA: InitAsync failed: {ex}"); + } + } + + public void Start(ImmediateActionQueue queue, Func canRunRemotely) + { + _logger.LogInformation("Kubernetes UBA: Start"); + + _timer = new Timer(async _ => + { + _timer?.Change(Timeout.Infinite, Timeout.Infinite); + + var stopping = false; + var quickRestart = false; + try + { + // If we're cancelled, stop. + if (_forcedStop || (CancellationTokenSource != null && CancellationTokenSource.IsCancellationRequested)) + { + _logger.LogInformation("Kubernetes UBA: Loop: Cancelled"); + stopping = true; + return; + } + + // If the client is unavailable, keep looping unless we're done. + if (_realCoordinator == null || + !(bool)_realCoordinatorProperties["ClientIsAvailable"].GetValue(_realCoordinator)!) + { + if (!queue.IsDone) + { + _logger.LogInformation("Kubernetes UBA: Loop: Client not available (waiting)"); + if (_quickRestartsRemaining > 0) + { + quickRestart = true; + _quickRestartsRemaining--; + } + } + else + { + _logger.LogInformation("Kubernetes UBA: Loop: Client not available (done)"); + stopping = true; + } + return; + } + + // Remove any Kubernetes pods that are complete, or failed to start. + await (Task)_realCoordinatorMethods["CleanupFinishedKubernetesPodsGloballyAsync"].Invoke(_realCoordinator, new object[0])!; + + // If we're done, return. + if (queue.IsDone) + { + _logger.LogInformation("Kubernetes UBA: Loop: Done"); + stopping = true; + return; + } + + // Synchronise Kubernetes nodes with our known node list. + await (Task)_realCoordinatorMethods["SynchroniseKubernetesNodesAsync"].Invoke(_realCoordinator, new object[0])!; + + // Determine the threshold over local allocation. + double desiredCpusThreshold = _ubaConfig.bForceBuildAllRemote ? 0 : 5; + + // Allocate cores from Kubernetes until we're satisified. + while (true) + { + // Check how many additional cores we need to allocate from the cluster. + double desiredCpus = queue.EnumerateReadyToCompileActions().Where(x => canRunRemotely(x)).Sum(x => x.Weight); + desiredCpus -= (int)_realCoordinatorMethods["GetAllocatedBlocks"].Invoke(_realCoordinator, new object[0])!; + if (desiredCpus <= desiredCpusThreshold) + { + _logger.LogInformation($"Kubernetes UBA: Loop: Skipping (desired CPU {desiredCpus} <= {desiredCpusThreshold})"); + break; + } + + // Remove any Kubernetes pods that are complete, or failed to start. + await (Task)_realCoordinatorMethods["CleanupFinishedKubernetesPodsLocallyAsync"].Invoke(_realCoordinator, new object[0])!; + + // Compute the biggest size we can allocate. + var blockSize = (int)_realCoordinatorMethods["GetMaximumBlockSize"].Invoke(_realCoordinator, new object[] { desiredCpus })!; + if (blockSize == 0) + { + break; + } + + // Try to allocate a node from Kubernetes. + var shouldContinue = await (Task)_realCoordinatorMethods["TryAllocateKubernetesNodeAsync"].Invoke(_realCoordinator, new object[] { blockSize })!; + if (!shouldContinue) + { + break; + } + } + } + catch (OperationCanceledException ex) when (CancellationTokenSource != null && ex.CancellationToken == CancellationTokenSource.Token) + { + // Expected exception. + stopping = true; + } + catch (Exception ex) + { + _logger.LogError(ex, $"Exception in Kubernetes loop: {ex}"); + } + finally + { + if (!stopping) + { + if (quickRestart) + { + _timer?.Change(_timerPeriod, Timeout.Infinite); + } + else + { + _timer?.Change(500, Timeout.Infinite); + } + } + } + }, null, 0, _timerPeriod); + } + + public void Stop() + { + CancellationTokenSource?.Cancel(); + _forcedStop = true; + } + + public async Task CloseAsync() + { + CancellationTokenSource?.Cancel(); + _forcedStop = true; + + if (_realCoordinator != null) + { + await (Task)_realCoordinatorMethods["CloseAsync"].Invoke(_realCoordinator, new object[0])!; + } + } + + public void Dispose() + { + Stop(); + CloseAsync().Wait(); + Dispose(true); + GC.SuppressFinalize(this); + } + + protected virtual void Dispose(bool disposing) + { + if (disposing) + { + CancellationTokenSource?.Dispose(); + _timer?.Dispose(); + _timer = null; + } + } + } +} \ No newline at end of file diff --git a/UET/Redpoint.Uet.Patching.Runtime/Patches/KubernetesUbaCoordinatorUetPatch.cs b/UET/Redpoint.Uet.Patching.Runtime/Patches/KubernetesUbaCoordinatorUetPatch.cs new file mode 100644 index 00000000..0b4069e6 --- /dev/null +++ b/UET/Redpoint.Uet.Patching.Runtime/Patches/KubernetesUbaCoordinatorUetPatch.cs @@ -0,0 +1,166 @@ +namespace Redpoint.Uet.Patching.Runtime.Patches +{ + using HarmonyLib; + using Microsoft.CodeAnalysis.CSharp; + using Microsoft.CodeAnalysis.Text; + using Microsoft.CodeAnalysis; + using System; + using System.Collections.Generic; + using System.Linq; + using System.Reflection; + using System.Runtime.Loader; + using System.Threading.Tasks; + using Redpoint.Uet.Patching.Runtime; + using System.Runtime.InteropServices; + using System.Collections; + + internal class KubernetesUbaCoordinatorUetPatch : IUetPatch + { + private static Assembly? _ubtUetAssembly; + + public bool ShouldApplyPatch() + { + var entryAssembly = Assembly.GetEntryAssembly(); + return + entryAssembly != null && + entryAssembly.GetName()?.Name == "UnrealBuildTool" && + entryAssembly.GetType("UnrealBuildTool.UBAExecutor") != null && + entryAssembly.GetType("UnrealBuildTool.IUBAAgentCoordinator") != null; + } + + public void ApplyPatch(IUetPatchLogging logging, Harmony harmony) + { + // Create our working directory for compilation. + var tempPath = Path.Combine(Path.GetTempPath(), $"uetstartuphook-{Environment.ProcessId}"); + Directory.CreateDirectory(tempPath); + + // Determine the location of UnrealBuildTool.dll. + var loadedUbtAssembly = AssemblyLoadContext.Default.Assemblies.First(x => x.GetName().Name == "UnrealBuildTool"); + + // Compile our Kubernetes coordinator type that we will hook into. + logging.LogInfo($"Compiling the Kubernetes UBA coordinator..."); + { + string sourceText; + using (var sourceStream = new StreamReader(Assembly.GetExecutingAssembly().GetManifestResourceStream("Redpoint.Uet.Patching.Runtime.Patches.KubernetesUbaCoordinatorPatchCode.cs")!)) + { + sourceText = sourceStream.ReadToEnd(); + } + var source = SourceText.From(sourceText); + var tree = CSharpSyntaxTree.ParseText(source, new CSharpParseOptions(LanguageVersion.Latest)); + var diagnostics = tree.GetDiagnostics(); + if (diagnostics.Any()) + { + Console.WriteLine("UET hook parse failure!"); + foreach (var diagnostic in diagnostics) + { + Console.WriteLine(diagnostic); + } + return; + } + var referencedAssemblies = new List + { + MetadataReference.CreateFromFile(typeof(object).Assembly.Location), + MetadataReference.CreateFromFile(Assembly.Load("System.Runtime").Location), + MetadataReference.CreateFromFile(Assembly.Load("System.Collections").Location), + MetadataReference.CreateFromFile(Assembly.Load("System.Linq").Location), + MetadataReference.CreateFromFile(Assembly.Load("EpicGames.Core").Location), + MetadataReference.CreateFromFile(Assembly.Load("EpicGames.UBA").Location), + MetadataReference.CreateFromFile(Assembly.Load("EpicGames.Build").Location), + MetadataReference.CreateFromFile(Assembly.Load("EpicGames.IoHash").Location), + MetadataReference.CreateFromFile(Assembly.Load("Microsoft.Extensions.Logging").Location), + MetadataReference.CreateFromFile(Assembly.Load("Microsoft.Extensions.Logging.Abstractions").Location), + MetadataReference.CreateFromFile(typeof(Task).Assembly.Location), + MetadataReference.CreateFromFile(typeof(Console).Assembly.Location), + MetadataReference.CreateFromFile(typeof(Timer).Assembly.Location), + MetadataReference.CreateFromFile(typeof(Type).Assembly.Location), + MetadataReference.CreateFromFile(typeof(RuntimeInformation).Assembly.Location), + MetadataReference.CreateFromFile(loadedUbtAssembly.Location) + }; + var compilationOptions = new CSharpCompilationOptions( + outputKind: OutputKind.DynamicallyLinkedLibrary, + optimizationLevel: OptimizationLevel.Debug, + warningLevel: 4, + assemblyIdentityComparer: DesktopAssemblyIdentityComparer.Default, + reportSuppressedDiagnostics: true); + var compilation = CSharpCompilation.Create( + // @note: Very intentional! Allows this assembly to access the internals. + assemblyName: "UnrealBuildTool.Tests", + syntaxTrees: new List { tree }, + references: referencedAssemblies, + options: compilationOptions); + using var ubtUetAssemblyMemory = new MemoryStream(); + var emitResult = compilation.Emit(ubtUetAssemblyMemory); + if (!emitResult.Success) + { + Console.WriteLine("UET hook compilation failure!"); + foreach (var diagnostic in emitResult.Diagnostics) + { + Console.WriteLine(diagnostic); + } + return; + } + ubtUetAssemblyMemory.Seek(0, SeekOrigin.Begin); + + logging.LogInfo($"Loading the Kubernetes UBA coordinator assembly..."); + _ubtUetAssembly = Assembly.Load(ubtUetAssemblyMemory.ToArray()); + + logging.LogInfo($"Loaded the Kubernetes UBA coordinator assembly!"); + } + + // Use Harmony to patch the constructor of UBAExecutor. + logging.LogInfo($"Adding the Kubernetes coordinator to UBAExecutor constructor..."); + { + var constructor = loadedUbtAssembly.GetType("UnrealBuildTool.UBAExecutor")?.GetDeclaredConstructors(false).FirstOrDefault(); + if (constructor == null) + { + throw new InvalidOperationException("Unable to find UnrealBuildTool.UBAExecutor constructor!"); + } + var postfix = GetType().GetMethod(nameof(UbaConstructorInjection), AccessTools.all); + harmony.Patch(constructor, null, new HarmonyMethod(postfix)); + } + + // Use Harmony to patch the constructor of UBAExecutor. + logging.LogInfo($"Patching XmlConfig.FindConfigurableTypes..."); + { + var findConfigurableTypes = loadedUbtAssembly.GetType("UnrealBuildTool.XmlConfig") + ?.GetMethod("FindConfigurableTypes", BindingFlags.NonPublic | BindingFlags.Static); + if (findConfigurableTypes == null) + { + throw new InvalidOperationException("Unable to find UnrealBuildTool.UBAExecutor constructor!"); + } + var postfix = GetType().GetMethod(nameof(XmlFindConfigurableTypesInjection), AccessTools.all); + harmony.Patch(findConfigurableTypes, null, new HarmonyMethod(postfix)); + } + } + + public static void UbaConstructorInjection(ref object __instance, object logger, object additionalArguments) + { + var coordinators = (IList)__instance.GetType() + .GetField("_agentCoordinators", BindingFlags.NonPublic | BindingFlags.Instance)! + .GetValue(__instance)!; + + // Remove other coordinators. + coordinators.Clear(); + + // Add our Kubernetes coordinator. + var config = __instance.GetType() + .GetProperty("UBAConfig", BindingFlags.Public | BindingFlags.Instance)! + .GetValue(__instance); + var coordinator = _ubtUetAssembly!.GetType("UnrealBuildTool.UBAAgentCoordinatorKubernetesConstructor")! + .GetMethod("Construct", BindingFlags.Static | BindingFlags.NonPublic)! + .Invoke(null, new object[] { logger, config!, additionalArguments }); + coordinators.Add(coordinator); + } + + public static void XmlFindConfigurableTypesInjection(ref List __result) + { + Console.WriteLine("Handling XmlConfig.FindConfigurableTypes"); + var type = _ubtUetAssembly!.GetType("UnrealBuildTool.UnrealBuildAcceleratorKubernetesConfig"); + if (type == null) + { + throw new InvalidOperationException("Unable to find type 'UnrealBuildTool.UnrealBuildAcceleratorKubernetesConfig'!"); + } + __result.Insert(0, type); + } + } +} diff --git a/UET/Redpoint.Uet.Patching.Runtime/Patches/UetPatch.cs b/UET/Redpoint.Uet.Patching.Runtime/Patches/UetPatch.cs new file mode 100644 index 00000000..f2a3212e --- /dev/null +++ b/UET/Redpoint.Uet.Patching.Runtime/Patches/UetPatch.cs @@ -0,0 +1,12 @@ +namespace Redpoint.Uet.Patching.Runtime.Patches +{ + using HarmonyLib; + using Redpoint.Uet.Patching.Runtime; + + internal interface IUetPatch + { + bool ShouldApplyPatch(); + + void ApplyPatch(IUetPatchLogging logging, Harmony harmony); + } +} diff --git a/UET/Redpoint.Uet.Patching.Runtime/Redpoint.Uet.Patching.Runtime.csproj b/UET/Redpoint.Uet.Patching.Runtime/Redpoint.Uet.Patching.Runtime.csproj new file mode 100644 index 00000000..59db15d5 --- /dev/null +++ b/UET/Redpoint.Uet.Patching.Runtime/Redpoint.Uet.Patching.Runtime.csproj @@ -0,0 +1,25 @@ + + + + net6.0 + enable + enable + true + + + + + + + + + + + + + + + + + + diff --git a/UET/Redpoint.Uet.Patching.Runtime/StartupHook.cs b/UET/Redpoint.Uet.Patching.Runtime/StartupHook.cs new file mode 100644 index 00000000..60f317a6 --- /dev/null +++ b/UET/Redpoint.Uet.Patching.Runtime/StartupHook.cs @@ -0,0 +1,23 @@ +using System.Reflection; +using System.Runtime.Loader; + +internal class StartupHook +{ + public static void Initialize() + { + // Set our assembly resolver for Mono.Cecil and other things. + var ourDirectoryPath = new FileInfo(Assembly.GetExecutingAssembly().Location).DirectoryName!; + AssemblyLoadContext.Default.Resolving += (AssemblyLoadContext context, AssemblyName name) => + { + var targetAssembly = Path.Combine(ourDirectoryPath, name.Name + ".dll"); + if (File.Exists(targetAssembly)) + { + return Assembly.LoadFrom(targetAssembly); + } + return null; + }; + + // Then call our real code. + Redpoint.Uet.Patching.Runtime.UetStartupHook.Initialize(); + } +} diff --git a/UET/Redpoint.Uet.Patching.Runtime/UetStartupHook.cs b/UET/Redpoint.Uet.Patching.Runtime/UetStartupHook.cs new file mode 100644 index 00000000..c7885433 --- /dev/null +++ b/UET/Redpoint.Uet.Patching.Runtime/UetStartupHook.cs @@ -0,0 +1,50 @@ +namespace Redpoint.Uet.Patching.Runtime +{ + using HarmonyLib; + using Microsoft.CodeAnalysis; + using Microsoft.CodeAnalysis.CSharp; + using Microsoft.CodeAnalysis.Text; + using Redpoint.Uet.Patching.Runtime.Patches; + using System; + using System.Reflection; + using System.Runtime.Loader; + + public static class UetStartupHook + { + public static void Initialize() + { + // Prevent MSBuild from re-using nodes, since we're injected into all .NET processes (including MSBuild). + Environment.SetEnvironmentVariable("MSBUILDDISABLENODEREUSE", "1"); + + // Our list of patches. + var patches = new IUetPatch[] + { + new KubernetesUbaCoordinatorUetPatch(), + new BuildGraphProjectRootUetPatch(), + }; + + // Determine if we have any patches to apply. If we have none, we're done. + if (!patches.Any(x => x.ShouldApplyPatch())) + { + return; + } + + // Create our logging instance. + IUetPatchLogging logging = Environment.GetEnvironmentVariable("UET_RUNTIME_PATCHING_ENABLE_LOGGING") == "1" + ? new ConsoleUetPatchLogging() + : new NullUetPatchLogging(); + + // Create our Harmony instance. + var harmony = new Harmony("games.redpoint.uet"); + + // Go through our patches and apply the ones that are relevant. + foreach (var patch in patches) + { + if (patch.ShouldApplyPatch()) + { + patch.ApplyPatch(logging, harmony); + } + } + } + } +} diff --git a/UET/UET.sln b/UET/UET.sln index b0206e30..5253a9f3 100644 --- a/UET/UET.sln +++ b/UET/UET.sln @@ -315,6 +315,8 @@ Project("{2150E333-8FDC-42A3-9474-1A3956D46DE8}") = "Redpoint.Uba", "Redpoint.Ub EndProject Project("{FAE04EC0-301F-11D3-BF4B-00C04F79EFBC}") = "Redpoint.Uba", "Redpoint.Uba\Redpoint.Uba.csproj", "{2AAD5971-F06A-4B7A-9E11-538BEC82F247}" EndProject +Project("{9A19103F-16F7-4668-BE54-9A1E7A4F7556}") = "Redpoint.Uet.Patching.Runtime", "Redpoint.Uet.Patching.Runtime\Redpoint.Uet.Patching.Runtime.csproj", "{C501B004-9ACA-467D-A15C-21DD9F42EF98}" +EndProject Global GlobalSection(SolutionConfigurationPlatforms) = preSolution Debug|Any CPU = Debug|Any CPU @@ -857,6 +859,18 @@ Global {2AAD5971-F06A-4B7A-9E11-538BEC82F247}.Debug|Any CPU.Build.0 = Debug|Any CPU {2AAD5971-F06A-4B7A-9E11-538BEC82F247}.Release|Any CPU.ActiveCfg = Release|Any CPU {2AAD5971-F06A-4B7A-9E11-538BEC82F247}.Release|Any CPU.Build.0 = Release|Any CPU + {2A500D98-9903-4659-87D7-E02498CFA64B}.Debug|Any CPU.ActiveCfg = Debug|Any CPU + {2A500D98-9903-4659-87D7-E02498CFA64B}.Debug|Any CPU.Build.0 = Debug|Any CPU + {2A500D98-9903-4659-87D7-E02498CFA64B}.Release|Any CPU.ActiveCfg = Release|Any CPU + {2A500D98-9903-4659-87D7-E02498CFA64B}.Release|Any CPU.Build.0 = Release|Any CPU + {46E05BD4-D74A-46E8-9A16-268E7D126B17}.Debug|Any CPU.ActiveCfg = Debug|Any CPU + {46E05BD4-D74A-46E8-9A16-268E7D126B17}.Debug|Any CPU.Build.0 = Debug|Any CPU + {46E05BD4-D74A-46E8-9A16-268E7D126B17}.Release|Any CPU.ActiveCfg = Release|Any CPU + {46E05BD4-D74A-46E8-9A16-268E7D126B17}.Release|Any CPU.Build.0 = Release|Any CPU + {C501B004-9ACA-467D-A15C-21DD9F42EF98}.Debug|Any CPU.ActiveCfg = Debug|Any CPU + {C501B004-9ACA-467D-A15C-21DD9F42EF98}.Debug|Any CPU.Build.0 = Debug|Any CPU + {C501B004-9ACA-467D-A15C-21DD9F42EF98}.Release|Any CPU.ActiveCfg = Release|Any CPU + {C501B004-9ACA-467D-A15C-21DD9F42EF98}.Release|Any CPU.Build.0 = Release|Any CPU EndGlobalSection GlobalSection(SolutionProperties) = preSolution HideSolutionNode = FALSE @@ -967,6 +981,9 @@ Global {7CDD0F44-C516-47E9-99E5-3885AD24636F} = {A10A6C63-109E-4825-AA79-81B0AE279A76} {D19B5289-C0A3-4025-834D-FA8AB1C19A32} = {1AE4AFCD-0F49-4CEA-8439-F1AAA2CDD183} {2AAD5971-F06A-4B7A-9E11-538BEC82F247} = {586E33AB-CC82-4ADF-92F3-3B287A5AEEAC} + {2A500D98-9903-4659-87D7-E02498CFA64B} = {5073FBDD-8F97-4261-A0B0-D2DC367CD643} + {46E05BD4-D74A-46E8-9A16-268E7D126B17} = {5073FBDD-8F97-4261-A0B0-D2DC367CD643} + {C501B004-9ACA-467D-A15C-21DD9F42EF98} = {1AE4AFCD-0F49-4CEA-8439-F1AAA2CDD183} EndGlobalSection GlobalSection(ExtensibilityGlobals) = postSolution SolutionGuid = {8598A278-509A-48A6-A7B3-3E3B0D1011F1}