357 lines
9.3 KiB
C#
357 lines
9.3 KiB
C#
// Copyright Epic Games, Inc. All Rights Reserved.
|
|
|
|
using System;
|
|
using System.IO;
|
|
using System.IO.MemoryMappedFiles;
|
|
using System.Threading;
|
|
using System.Threading.Channels;
|
|
using System.Threading.Tasks;
|
|
using EpicGames.Core;
|
|
using EpicGames.Horde.Storage.Nodes;
|
|
using EpicGames.Horde.Utilities;
|
|
using Microsoft.Extensions.Logging;
|
|
|
|
namespace EpicGames.Horde.Storage
|
|
{
|
|
/// <summary>
|
|
/// Represents an output file that write requests may be issued against
|
|
/// </summary>
|
|
class OutputFile
|
|
{
|
|
/// <summary>
|
|
/// The relative output path
|
|
/// </summary>
|
|
public string Path { get; }
|
|
|
|
/// <summary>
|
|
/// File metadata
|
|
/// </summary>
|
|
public FileInfo FileInfo { get; }
|
|
|
|
/// <summary>
|
|
/// File entry with metadata about the target file
|
|
/// </summary>
|
|
public FileEntry FileEntry { get; }
|
|
|
|
bool _createdFile;
|
|
int _remainingChunks;
|
|
|
|
/// <summary>
|
|
/// Constructor
|
|
/// </summary>
|
|
public OutputFile(string path, FileInfo fileInfo, FileEntry fileEntry)
|
|
{
|
|
Path = path;
|
|
FileInfo = fileInfo;
|
|
FileEntry = fileEntry;
|
|
}
|
|
|
|
public int IncrementRemaining() => Interlocked.Increment(ref _remainingChunks);
|
|
public int DecrementRemaining() => Interlocked.Decrement(ref _remainingChunks);
|
|
|
|
/// <summary>
|
|
/// Opens the file for writing, setting its length on the first run
|
|
/// </summary>
|
|
public FileStream OpenStream()
|
|
{
|
|
lock (FileEntry)
|
|
{
|
|
if (!_createdFile)
|
|
{
|
|
if (FileInfo.Exists)
|
|
{
|
|
if (FileInfo.LinkTarget != null)
|
|
{
|
|
FileInfo.Delete();
|
|
}
|
|
else if (FileInfo.IsReadOnly)
|
|
{
|
|
FileInfo.IsReadOnly = false;
|
|
}
|
|
}
|
|
else
|
|
{
|
|
FileInfo.Directory?.Create();
|
|
}
|
|
}
|
|
|
|
FileStream? stream = null;
|
|
try
|
|
{
|
|
try
|
|
{
|
|
stream = FileInfo.Open(FileMode.OpenOrCreate, FileAccess.ReadWrite, FileShare.ReadWrite);
|
|
}
|
|
catch (IOException ex)
|
|
{
|
|
string? lockInfo = FileUtils.GetFileLockInfo(FileInfo.FullName);
|
|
if (lockInfo == null)
|
|
{
|
|
throw;
|
|
}
|
|
else
|
|
{
|
|
throw new WrappedFileOrDirectoryException(ex, $"{ex.Message}\n{lockInfo}");
|
|
}
|
|
}
|
|
|
|
if (!_createdFile)
|
|
{
|
|
stream.SetLength(FileEntry.Length);
|
|
_createdFile = true;
|
|
}
|
|
return stream;
|
|
}
|
|
catch
|
|
{
|
|
stream?.Dispose();
|
|
throw;
|
|
}
|
|
}
|
|
}
|
|
|
|
/// <summary>
|
|
/// Callback for a file having been fully written
|
|
/// </summary>
|
|
public virtual ValueTask CompleteAsync(CancellationToken cancellationToken)
|
|
=> default;
|
|
}
|
|
|
|
/// <summary>
|
|
/// Request to output a block of data A chunk which needs to be written to an output file
|
|
/// </summary>
|
|
record class OutputChunk(OutputFile File, long Offset, ReadOnlyMemory<byte> Data, object Source, IDisposable? DataOwner) : IDisposable
|
|
{
|
|
/// <inheritdoc/>
|
|
public void Dispose()
|
|
=> DataOwner?.Dispose();
|
|
}
|
|
|
|
/// <summary>
|
|
/// Batches file write requests
|
|
/// </summary>
|
|
sealed class BatchOutputWriter : IDisposable
|
|
{
|
|
readonly Channel<OutputChunk[]> _channel;
|
|
readonly ILogger _logger;
|
|
|
|
int _writtenFiles;
|
|
long _writtenBytes;
|
|
|
|
/// <summary>
|
|
/// Number of files that have been written
|
|
/// </summary>
|
|
public int WrittenFiles
|
|
=> Interlocked.CompareExchange(ref _writtenFiles, 0, 0);
|
|
|
|
/// <summary>
|
|
/// Total number of bytes that have been written
|
|
/// </summary>
|
|
public long WrittenBytes
|
|
=> Interlocked.CompareExchange(ref _writtenBytes, 0, 0);
|
|
|
|
/// <summary>
|
|
/// Number of writes to execute sequentially vs in parallel
|
|
/// </summary>
|
|
public int WriteBatchSize { get; set; } = 64;
|
|
|
|
/// <summary>
|
|
/// If false, disables output to disk. Useful for performance testing.
|
|
/// </summary>
|
|
public bool EnableOutput { get; set; } = true;
|
|
|
|
/// <summary>
|
|
/// If true, hashes output files after writing to verify their integrity
|
|
/// </summary>
|
|
public bool VerifyOutput { get; set; }
|
|
|
|
/// <summary>
|
|
/// If true, outputs verbose information to the log
|
|
/// </summary>
|
|
public bool VerboseOutput { get; set; }
|
|
|
|
/// <summary>
|
|
/// Sink for write requests
|
|
/// </summary>
|
|
public ChannelWriter<OutputChunk[]> RequestWriter
|
|
=> _channel.Writer;
|
|
|
|
/// <summary>
|
|
/// Constructor
|
|
/// </summary>
|
|
public BatchOutputWriter(ILogger logger)
|
|
{
|
|
_channel = Channel.CreateBounded<OutputChunk[]>(new BoundedChannelOptions(128));
|
|
_logger = logger;
|
|
}
|
|
|
|
/// <summary>
|
|
/// Dispose remaining items in the queue
|
|
/// </summary>
|
|
public void Dispose()
|
|
{
|
|
while (_channel.Reader.TryRead(out OutputChunk[]? chunks))
|
|
{
|
|
foreach (OutputChunk chunk in chunks)
|
|
{
|
|
chunk.Dispose();
|
|
}
|
|
}
|
|
}
|
|
|
|
/// <summary>
|
|
/// Adds tasks for the writer to an async pipeline
|
|
/// </summary>
|
|
/// <param name="pipeline">Pipeline instance</param>
|
|
/// <param name="numWriteTasks">Number of parallel writes</param>
|
|
public Task[] AddToPipeline(AsyncPipeline pipeline, int numWriteTasks)
|
|
{
|
|
return pipeline.AddTasks(numWriteTasks, ProcessAsync);
|
|
}
|
|
|
|
/// <summary>
|
|
/// Processes requests from the given input channel
|
|
/// </summary>
|
|
/// <param name="cancellationToken">Cancellation token for the operation</param>
|
|
public async Task ProcessAsync(CancellationToken cancellationToken)
|
|
{
|
|
await foreach (OutputChunk[] chunks in _channel.Reader.ReadAllAsync(cancellationToken))
|
|
{
|
|
try
|
|
{
|
|
await WriteChunksAsync(chunks, cancellationToken);
|
|
}
|
|
catch (Exception ex) when (ex is not OperationCanceledException)
|
|
{
|
|
_logger.LogError(ex, "Error writing to disk: {Message}", ex.Message);
|
|
throw;
|
|
}
|
|
finally
|
|
{
|
|
foreach (OutputChunk chunk in chunks)
|
|
{
|
|
chunk.Dispose();
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
async Task WriteChunksAsync(OutputChunk[] chunks, CancellationToken cancellationToken)
|
|
{
|
|
for (int requestIdx = 0; requestIdx < chunks.Length;)
|
|
{
|
|
OutputFile outputFile = chunks[requestIdx].File;
|
|
|
|
int maxChunkIdx = requestIdx + 1;
|
|
while (maxChunkIdx < chunks.Length && chunks[maxChunkIdx].File == outputFile)
|
|
{
|
|
maxChunkIdx++;
|
|
}
|
|
|
|
try
|
|
{
|
|
await ExtractChunksToFileAsync(outputFile, new ArraySegment<OutputChunk>(chunks, requestIdx, maxChunkIdx - requestIdx), cancellationToken);
|
|
}
|
|
catch (OperationCanceledException)
|
|
{
|
|
throw;
|
|
}
|
|
catch (Exception ex)
|
|
{
|
|
throw new StorageException($"Unable to extract {outputFile.FileInfo.FullName}: {ex.Message}", ex);
|
|
}
|
|
|
|
requestIdx = maxChunkIdx;
|
|
}
|
|
}
|
|
|
|
async Task ExtractChunksToFileAsync(OutputFile outputFile, ArraySegment<OutputChunk> chunks, CancellationToken cancellationToken)
|
|
{
|
|
using PlatformFileLock? fileLock = await PlatformFileLock.CreateAsync(cancellationToken);
|
|
|
|
// Open the file for the current chunk
|
|
int remainingChunks = 0;
|
|
await using (FileStream stream = outputFile.OpenStream())
|
|
{
|
|
if (outputFile.FileEntry.Length == 0 || !EnableOutput)
|
|
{
|
|
// If this file is empty, don't write anything and just move to the next chunk
|
|
for (int idx = 0; idx < chunks.Count; idx++)
|
|
{
|
|
remainingChunks = outputFile.DecrementRemaining();
|
|
}
|
|
}
|
|
else
|
|
{
|
|
// Process as many chunks as we can for this file
|
|
using MemoryMappedFile memoryMappedFile = MemoryMappedFile.CreateFromFile(stream, null, outputFile.FileEntry.Length, MemoryMappedFileAccess.ReadWrite, HandleInheritability.None, false);
|
|
using MemoryMappedView memoryMappedView = CreateMemoryMappedView(memoryMappedFile, outputFile);
|
|
|
|
for (int chunkIdx = 0; chunkIdx < chunks.Count; chunkIdx++)
|
|
{
|
|
OutputChunk chunk = chunks[chunkIdx];
|
|
cancellationToken.ThrowIfCancellationRequested();
|
|
|
|
if (VerboseOutput)
|
|
{
|
|
_logger.LogInformation("Writing {File},{Offset} ({Source})", outputFile.Path, chunk.Offset, chunk.Source);
|
|
}
|
|
|
|
// Write this chunk
|
|
ReadOnlyMemory<byte> memory = chunk.Data;
|
|
memory.CopyTo(memoryMappedView!.GetMemory(chunk.Offset, memory.Length));
|
|
|
|
// Update the stats
|
|
remainingChunks = outputFile.DecrementRemaining();
|
|
Interlocked.Add(ref _writtenBytes, memory.Length);
|
|
}
|
|
}
|
|
}
|
|
|
|
// Set correct permissions on the output file
|
|
if (remainingChunks == 0)
|
|
{
|
|
if (VerifyOutput)
|
|
{
|
|
IoHash computedHash;
|
|
using (Stream stream = outputFile.OpenStream())
|
|
{
|
|
computedHash = await IoHash.ComputeAsync(stream, cancellationToken);
|
|
}
|
|
if (computedHash != outputFile.FileEntry.Hash)
|
|
{
|
|
throw new InvalidDataException($"File {outputFile.Path} does not have correct hash after extract; expected {outputFile.FileEntry.Hash}, got {computedHash}.");
|
|
}
|
|
_logger.LogDebug("Verified hash of {Path}: {Hash}", outputFile.Path, computedHash);
|
|
}
|
|
|
|
outputFile.FileInfo.Refresh();
|
|
FileEntry.SetPermissions(outputFile.FileInfo!, outputFile.FileEntry.Flags);
|
|
|
|
if ((outputFile.FileEntry.Flags & FileEntryFlags.HasModTime) != 0)
|
|
{
|
|
outputFile.FileInfo.LastWriteTimeUtc = outputFile.FileEntry.ModTime;
|
|
}
|
|
|
|
Interlocked.Increment(ref _writtenFiles);
|
|
|
|
await outputFile.CompleteAsync(cancellationToken);
|
|
}
|
|
}
|
|
|
|
static MemoryMappedView CreateMemoryMappedView(MemoryMappedFile memoryMappedFile, OutputFile file)
|
|
{
|
|
try
|
|
{
|
|
MemoryMappedView view = new MemoryMappedView(memoryMappedFile, 0, file.FileEntry.Length);
|
|
return view;
|
|
}
|
|
catch (Exception ex)
|
|
{
|
|
file.FileInfo.Refresh();
|
|
throw new StorageException($"Unable to create mapped view of file (entry length: {file.FileEntry.Length}, file length: {file.FileInfo.Length}): {ex.Message}", ex);
|
|
}
|
|
}
|
|
}
|
|
}
|