Files
UnrealEngine/Engine/Source/Programs/UnrealBuildTool/Artifacts/HordeStorageArtifactCache.cs
2025-05-18 13:04:45 +08:00

536 lines
16 KiB
C#

// Copyright Epic Games, Inc. All Rights Reserved.
using System;
using System.Collections.Generic;
using System.IO;
using System.Linq;
using System.Threading;
using System.Threading.Tasks;
using EpicGames.Core;
using EpicGames.Horde.Storage;
using EpicGames.Horde.Storage.Bundles;
using EpicGames.Horde.Storage.Nodes;
using Microsoft.Extensions.Logging;
namespace UnrealBuildTool.Artifacts
{
/// <summary>
/// Horde specific artifact action structure that also contains the file nodes for the outputs
/// </summary>
readonly struct HordeArtifactAction
{
/// <summary>
/// Artifact action
/// </summary>
public readonly ArtifactAction ArtifactAction;
/// <summary>
/// Collection of output file references. There should be exactly the same number
/// of file references as outputs in the action
/// </summary>
public readonly IHashedBlobRef<ChunkedDataNode>[] OutputRefs;
/// <summary>
/// Construct a new horde artifact number
/// </summary>
/// <param name="artifactAction">Artifact action</param>
/// <exception cref="ArgumentException"></exception>
public HordeArtifactAction(ArtifactAction artifactAction)
{
ArtifactAction = artifactAction;
OutputRefs = new IHashedBlobRef<ChunkedDataNode>[ArtifactAction.Outputs.Length];
}
/// <summary>
/// Construct a new artifact action from the reader
/// </summary>
/// <param name="reader">Source reader</param>
public HordeArtifactAction(IBlobReader reader)
{
ArtifactAction = reader.ReadArtifactAction();
OutputRefs = reader.ReadVariableLengthArray(() => reader.ReadBlobRef<ChunkedDataNode>());
}
/// <summary>
/// Serialize the artifact action
/// </summary>
/// <param name="writer">Destination writer</param>
public void Serialize(IBlobWriter writer)
{
writer.WriteArtifactAction(ArtifactAction);
writer.WriteVariableLengthArray(OutputRefs, x => writer.WriteBlobRef(x));
}
/// <summary>
/// Write all the files to disk
/// </summary>
/// <param name="writer">Destination writer</param>
/// <param name="cancellationToken">Cancellation token</param>
/// <returns>Task</returns>
public async Task WriteFilesAsync(IBlobWriter writer, CancellationToken cancellationToken)
{
LeafChunkedDataNodeOptions leafOptions = new(512 * 1024, 1 * 1024 * 1024, 2 * 1024 * 1024);
InteriorChunkedDataNodeOptions interiorOptions = new(1, 10, 20);
using LeafChunkedDataWriter fileWriter = new(writer, leafOptions);
int index = 0;
foreach (ArtifactFile artifact in ArtifactAction.Outputs)
{
string outputName = artifact.GetFullPath(ArtifactAction.DirectoryMapping);
using FileStream stream = new(outputName, FileMode.Open, FileAccess.Read, FileShare.Read);
LeafChunkedData leafChunkedData = await fileWriter.CreateAsync(stream, leafOptions.TargetSize, cancellationToken);
ChunkedData chunkedData = await InteriorChunkedDataNode.CreateTreeAsync(leafChunkedData, interiorOptions, writer, cancellationToken);
OutputRefs[index++] = chunkedData.Root;
}
}
}
/// <summary>
/// Series of helper methods for serialization
/// </summary>
static class HordeArtifactReaderWriterExtensions
{
/// <summary>
/// Read a horde artifact action
/// </summary>
/// <param name="reader">Source reader</param>
/// <returns>Created artifact action</returns>
public static HordeArtifactAction ReadHordeArtifactAction(this IBlobReader reader)
{
return new HordeArtifactAction(reader);
}
/// <summary>
/// Write a horde artifact action
/// </summary>
/// <param name="writer">Destination writer</param>
/// <param name="artifactAction">Artifact action to write</param>
public static void WriteHordeArtifactAction(this IBlobWriter writer, HordeArtifactAction artifactAction)
{
artifactAction.Serialize(writer);
}
}
/// <summary>
/// Horde node that represents a collection of action nodes
/// </summary>
[BlobConverter(typeof(ArtifactActionCollectionNodeConverter))]
class ArtifactActionCollectionNode
{
/// <summary>
/// Collection of actions
/// </summary>
public Dictionary<IoHash, HordeArtifactAction> ArtifactActions = new();
/// <summary>
/// Construct a new collection
/// </summary>
public ArtifactActionCollectionNode()
{
}
}
class ArtifactActionCollectionNodeConverter : BlobConverter<ArtifactActionCollectionNode>
{
static readonly BlobType s_blobType = new BlobType("{E8DBCD77-4CAE-861D-0758-7FB733256ED2}", 1);
public override ArtifactActionCollectionNode Read(IBlobReader reader, BlobSerializerOptions options)
{
ArtifactActionCollectionNode node = new ArtifactActionCollectionNode();
node.ArtifactActions = reader.ReadDictionary<IoHash, HordeArtifactAction>(() => reader.ReadIoHash(), () => reader.ReadHordeArtifactAction());
return node;
}
public override BlobType Write(IBlobWriter writer, ArtifactActionCollectionNode value, BlobSerializerOptions options)
{
writer.WriteDictionary<IoHash, HordeArtifactAction>(value.ArtifactActions, (x) => writer.WriteIoHash(x), (x) => writer.WriteHordeArtifactAction(x));
return s_blobType;
}
}
/// <summary>
/// Class for managing artifacts using horde storage
/// </summary>
public sealed class HordeStorageArtifactCache : IArtifactCache, IDisposable
{
/// <summary>
/// Defines the theoretical max number of pending actions to write
/// </summary>
const int MaxPendingSize = 128;
/// <summary>
/// Underlying storage object
/// </summary>
private IStorageNamespace? _store = null;
/// <summary>
/// Task used to wait on ready state
/// </summary>
private Task<ArtifactCacheState>? _readyTask = null;
/// <summary>
/// Ready state
/// </summary>
private int _state = (int)ArtifactCacheState.Pending;
/// <summary>
/// Collection of actions waiting to be written
/// </summary>
private readonly List<ArtifactAction> _pendingWrites;
/// <summary>
/// Task for any pending flush
/// </summary>
private Task? _pendingWritesFlushTask = null;
/// <summary>
/// Controls access to shared data structures
/// </summary>
private readonly SemaphoreSlim _semaphore = new(1);
/// <summary>
/// Test to see if the cache is ready
/// </summary>
public ArtifactCacheState State
{
get => (ArtifactCacheState)Interlocked.Add(ref _state, 0);
private set => Interlocked.Exchange(ref _state, (int)value);
}
/// <inheritdoc/>
public void Dispose()
{
_pendingWritesFlushTask?.Dispose();
_pendingWritesFlushTask = null;
_semaphore.Dispose();
}
/// <summary>
/// Create a memory only cache
/// </summary>
public static IArtifactCache CreateMemoryCache(ILogger logger)
{
HordeStorageArtifactCache cache = new(BundleStorageNamespace.CreateInMemory(logger))
{
State = ArtifactCacheState.Available
};
return cache;
}
/// <summary>
/// Create a file based cache
/// </summary>
/// <param name="directory">Destination directory</param>
/// <param name="memoryMappedFileCache">Cache for memory mapped files</param>
/// <param name="logger">Logging object</param>
/// <param name="cleanDirectory">If true, clean the directory</param>
public static IArtifactCache CreateFileCache(DirectoryReference directory, MemoryMappedFileCache memoryMappedFileCache, ILogger logger, bool cleanDirectory)
{
HordeStorageArtifactCache cache = new(null);
cache._readyTask = Task.Run(() => cache.InitFileCache(directory, memoryMappedFileCache, logger, cleanDirectory));
return cache;
}
/// <summary>
/// Constructor
/// </summary>
/// <param name="storage">Storage object to use</param>
private HordeStorageArtifactCache(IStorageNamespace? storage)
{
_store = storage;
_pendingWrites = new(MaxPendingSize);
}
/// <inheritdoc/>
public Task<ArtifactCacheState> WaitForReadyAsync()
{
return _readyTask ?? Task.FromResult<ArtifactCacheState>(State);
}
/// <inheritdoc/>
public async Task<ArtifactAction[]> QueryArtifactActionsAsync(IoHash[] partialKeys, CancellationToken cancellationToken)
{
if (State != ArtifactCacheState.Available || _store == null)
{
return Array.Empty<ArtifactAction>();
}
List<ArtifactAction> artifactActions = new();
await _semaphore.WaitAsync(cancellationToken);
try
{
foreach (IoHash key in partialKeys)
{
lock (_pendingWrites)
{
artifactActions.AddRange(_pendingWrites.Where(x => x.Key == key));
}
ArtifactActionCollectionNode? node = await _store.TryReadRefTargetAsync<ArtifactActionCollectionNode>(GetRefName(key), default, cancellationToken: cancellationToken);
if (node != null)
{
foreach (HordeArtifactAction artifactAction in node.ArtifactActions.Values)
{
artifactActions.Add(artifactAction.ArtifactAction);
}
}
}
}
finally
{
_semaphore.Release();
}
return artifactActions.ToArray();
}
/// <inheritdoc/>
public async Task<bool[]?> QueryArtifactOutputsAsync(ArtifactAction[] artifactActions, CancellationToken cancellationToken)
{
if (State != ArtifactCacheState.Available || _store == null)
{
return null;
}
bool[] output = new bool[artifactActions.Length];
Array.Fill(output, false);
for (int index = 0; index < artifactActions.Length; index++)
{
output[index] = false;
ArtifactAction artifactAction = artifactActions[index];
ArtifactActionCollectionNode? node = await _store.TryReadRefTargetAsync<ArtifactActionCollectionNode>(GetRefName(artifactAction.Key), default, cancellationToken: cancellationToken);
if (node != null)
{
if (node.ArtifactActions.TryGetValue(artifactAction.ActionKey, out HordeArtifactAction hordeArtifactAction))
{
output[index] = true;
int refIndex = 0;
foreach (IHashedBlobRef<ChunkedDataNode> artifactRef in hordeArtifactAction.OutputRefs)
{
if (artifactRef == null)
{
output[index] = false;
break;
}
try
{
string outputName = hordeArtifactAction.ArtifactAction.Outputs[refIndex++].GetFullPath(artifactAction.DirectoryMapping);
using FileStream stream = new(outputName, FileMode.Create, FileAccess.Write, FileShare.ReadWrite);
await artifactRef.CopyToStreamAsync(stream, cancellationToken);
}
catch (Exception)
{
output[index] = false;
break;
}
}
if (!output[index])
{
foreach (ArtifactFile artifact in hordeArtifactAction.ArtifactAction.Outputs)
{
string outputName = artifact.GetFullPath(artifactAction.DirectoryMapping);
if (File.Exists(outputName))
{
try
{
File.Delete(outputName);
}
catch (Exception)
{
}
}
}
}
}
}
}
return output;
}
/// <inheritdoc/>
public async Task SaveArtifactActionsAsync(ArtifactAction[] artifactActions, CancellationToken cancellationToken)
{
if (State != ArtifactCacheState.Available || _store == null)
{
return;
}
lock (_pendingWrites)
{
_pendingWrites.AddRange(artifactActions);
}
Task? task = FlushChangesInternalAsync(false, cancellationToken);
if (task != null)
{
await task;
}
}
/// <inheritdoc/>
public async Task FlushChangesAsync(CancellationToken cancellationToken)
{
if (State != ArtifactCacheState.Available || _store == null)
{
return;
}
Task? task = FlushChangesInternalAsync(true, cancellationToken);
if (task != null)
{
await task;
}
}
/// <summary>
/// Optionally flush all pending writes
/// </summary>
/// <param name="force">If true, force a flush</param>
/// <param name="cancellationToken">Cancellation token</param>
private Task? FlushChangesInternalAsync(bool force, CancellationToken cancellationToken)
{
Task? pendingFlushTask = null;
lock (_pendingWrites)
{
// If any prior flush task has completed, then forget it
if (_pendingWritesFlushTask != null && _pendingWritesFlushTask.IsCompleted)
{
_pendingWritesFlushTask = null;
}
// We start a new flush under the following condition
//
// 1) Actions must be pending
// 2) Create a new task if force is specified
// 3) -OR- Create a new task if there is no current task and we have reached the limit
if (_pendingWrites.Count > 0 && (force || (_pendingWrites.Count >= MaxPendingSize && _pendingWritesFlushTask == null)))
{
ArtifactAction[] artifactActionsToFlush = _pendingWrites.ToArray();
Task? priorTask = _pendingWritesFlushTask;
_pendingWrites.Clear();
async Task action()
{
// When forcing, we might have a prior flush task in progress. Wait for it to complete
if (priorTask != null)
{
await priorTask;
}
// Block reading while we update the actions
await _semaphore.WaitAsync(cancellationToken);
try
{
List<Task> tasks = CommitArtifactActions(artifactActionsToFlush, cancellationToken);
await Task.WhenAll(tasks);
}
finally
{
_semaphore.Release();
}
}
pendingFlushTask = _pendingWritesFlushTask = new(() => action().Wait(), cancellationToken);
}
}
// Start the task outside of the lock
pendingFlushTask?.Start();
return pendingFlushTask;
}
/// <summary>
/// Add a group of artifact actions to a new or existing source
/// </summary>
/// <param name="artifactActions">New artifact actions to add</param>
/// <param name="cancellationToken">Token to be used to cancel operations</param>
/// <returns>List of tasks</returns>
private List<Task> CommitArtifactActions(ArtifactAction[] artifactActions, CancellationToken cancellationToken)
{
List<Task> tasks = new();
if (artifactActions.Length == 0)
{
return tasks;
}
// Loop through the artifact actions
foreach (ArtifactAction artifactAction in artifactActions)
{
// Create the task to write the files
tasks.Add(Task.Run(async () =>
{
RefName refName = GetRefName(artifactAction.Key);
// Locate the destination collection for this key
ArtifactActionCollectionNode? node = await _store!.TryReadRefTargetAsync<ArtifactActionCollectionNode>(refName, default, cancellationToken: cancellationToken);
node ??= new ArtifactActionCollectionNode();
// Update the artifact action collection
HordeArtifactAction hordeArtifactAction = new(artifactAction);
node.ArtifactActions[artifactAction.ActionKey] = hordeArtifactAction;
// Save the artifact action file
await using IBlobWriter writer = _store!.CreateBlobWriter();
await hordeArtifactAction.WriteFilesAsync(writer, cancellationToken);
IHashedBlobRef<ArtifactActionCollectionNode> nodeRef = await writer.WriteBlobAsync(node);
await writer.FlushAsync();
// Save the collection
await _store.AddRefAsync(refName, nodeRef, cancellationToken: cancellationToken);
}, cancellationToken));
}
return tasks;
}
/// <summary>
/// Initialize a file based cache
/// </summary>
/// <param name="directory">Destination directory</param>
/// <param name="memoryMappedFileCache">Cache for memory mapped files</param>
/// <param name="logger">Logger</param>
/// <param name="cleanDirectory">If true, clean the directory</param>
/// <returns>Cache state</returns>
private ArtifactCacheState InitFileCache(DirectoryReference directory, MemoryMappedFileCache memoryMappedFileCache, ILogger logger, bool cleanDirectory)
{
try
{
if (cleanDirectory)
{
// Clear the output directory
try
{
Directory.Delete(directory.FullName, true);
}
catch (Exception)
{ }
}
Directory.CreateDirectory(directory.FullName);
_store = BundleStorageNamespace.CreateFromDirectory(directory, BundleCache.None, memoryMappedFileCache, logger);
State = ArtifactCacheState.Available;
return State;
}
catch (Exception)
{
State = ArtifactCacheState.Unavailable;
throw;
}
}
/// <summary>
/// Return the ref name for horde storage given a artitfact action collection key
/// </summary>
/// <param name="key">Artifact action collection key</param>
/// <returns>The reference name</returns>
private static RefName GetRefName(IoHash key)
{
return new RefName($"action_artifact_v2_{key}");
}
}
}