adds the ability for shards to sync with the main server regarding their caches, when pulling files from main server the shard can stream it to the client directly while downloading and add info for server to report to client regarding file locations across shards

This commit is contained in:
Abelfreyja
2026-01-14 16:52:06 +09:00
parent e8c56bb3bc
commit 0eea339e41
11 changed files with 915 additions and 27 deletions

View File

@@ -16,6 +16,7 @@ public sealed class CachedFileProvider : IDisposable
private readonly FileStatisticsService _fileStatisticsService;
private readonly LightlessMetrics _metrics;
private readonly ServerTokenGenerator _generator;
private readonly IShardFileInventoryReporter _inventoryReporter;
private readonly Uri _remoteCacheSourceUri;
private readonly string _hotStoragePath;
private readonly ConcurrentDictionary<string, Task> _currentTransfers = new(StringComparer.Ordinal);
@@ -27,13 +28,15 @@ public sealed class CachedFileProvider : IDisposable
private bool _isDistributionServer;
public CachedFileProvider(IConfigurationService<StaticFilesServerConfiguration> configuration, ILogger<CachedFileProvider> logger,
FileStatisticsService fileStatisticsService, LightlessMetrics metrics, ServerTokenGenerator generator)
FileStatisticsService fileStatisticsService, LightlessMetrics metrics, ServerTokenGenerator generator,
IShardFileInventoryReporter inventoryReporter)
{
_configuration = configuration;
_logger = logger;
_fileStatisticsService = fileStatisticsService;
_metrics = metrics;
_generator = generator;
_inventoryReporter = inventoryReporter;
_remoteCacheSourceUri = configuration.GetValueOrDefault<Uri>(nameof(StaticFilesServerConfiguration.DistributionFileServerAddress), null);
_isDistributionServer = configuration.GetValueOrDefault(nameof(StaticFilesServerConfiguration.IsDistributionNode), false);
_hotStoragePath = configuration.GetValue<string>(nameof(StaticFilesServerConfiguration.CacheDirectory));
@@ -97,10 +100,11 @@ public sealed class CachedFileProvider : IDisposable
_metrics.IncGauge(MetricsAPI.GaugeFilesTotal);
_metrics.IncGauge(MetricsAPI.GaugeFilesTotalSize, FilePathUtil.GetFileInfoForHash(_hotStoragePath, hash).Length);
_inventoryReporter.ReportAdded(hash);
response.Dispose();
}
private bool TryCopyFromColdStorage(string hash, string destinationFilePath)
private bool TryCopyFromColdStorage(string hash, string destinationFilePath)
{
if (!_configuration.GetValueOrDefault(nameof(StaticFilesServerConfiguration.UseColdStorage), false)) return false;
@@ -124,6 +128,7 @@ public sealed class CachedFileProvider : IDisposable
_metrics.IncGauge(MetricsAPI.GaugeFilesTotal);
_metrics.IncGauge(MetricsAPI.GaugeFilesTotalSize, new FileInfo(destinationFilePath).Length);
_inventoryReporter.ReportAdded(hash);
return true;
}
catch (Exception ex)
@@ -134,6 +139,93 @@ public sealed class CachedFileProvider : IDisposable
return false;
}
public async Task<CachedFileStreamResult?> GetFileStreamForDirectDownloadAsync(string hash, CancellationToken ct)
{
var destinationFilePath = FilePathUtil.GetFilePath(_hotStoragePath, hash);
var existing = GetLocalFilePath(hash);
if (existing != null)
{
return new CachedFileStreamResult(existing, null, existing.Length);
}
if (TryCopyFromColdStorage(hash, destinationFilePath))
{
var coldFile = GetLocalFilePath(hash);
if (coldFile != null)
{
return new CachedFileStreamResult(coldFile, null, coldFile.Length);
}
}
if (_remoteCacheSourceUri == null)
{
return null;
}
TaskCompletionSource<bool>? completion = null;
await _downloadSemaphore.WaitAsync(ct).ConfigureAwait(false);
try
{
if (_currentTransfers.TryGetValue(hash, out var downloadTask)
&& !(downloadTask?.IsCompleted ?? true))
{
completion = null;
}
else
{
completion = new TaskCompletionSource<bool>(TaskCreationOptions.RunContinuationsAsynchronously);
_currentTransfers[hash] = completion.Task;
_metrics.IncGauge(MetricsAPI.GaugeFilesDownloadingFromCache);
}
}
finally
{
_downloadSemaphore.Release();
}
if (completion == null)
{
var waited = await DownloadAndGetLocalFileInfo(hash).ConfigureAwait(false);
if (waited == null) return null;
return new CachedFileStreamResult(waited, null, waited.Length);
}
var downloadUrl = LightlessFiles.DistributionGetFullPath(_remoteCacheSourceUri, hash);
_logger.LogInformation("Did not find {hash}, streaming from {server}", hash, downloadUrl);
using var requestMessage = new HttpRequestMessage(HttpMethod.Get, downloadUrl);
requestMessage.Headers.Authorization = new AuthenticationHeaderValue("Bearer", _generator.Token);
HttpResponseMessage response;
try
{
response = await _httpClient.SendAsync(requestMessage, HttpCompletionOption.ResponseHeadersRead, ct).ConfigureAwait(false);
response.EnsureSuccessStatusCode();
}
catch (Exception ex)
{
_logger.LogWarning(ex, "Failed to stream {url}", downloadUrl);
FinalizeStreamingDownload(hash, null, destinationFilePath, completion, false, 0);
return null;
}
var tempFileName = destinationFilePath + ".dl";
var fileStream = new FileStream(tempFileName, FileMode.Create, FileAccess.Write, FileShare.Read, bufferSize: 64 * 1024, useAsync: true);
var sourceStream = await response.Content.ReadAsStreamAsync(ct).ConfigureAwait(false);
var stream = new StreamingCacheWriteStream(
sourceStream,
fileStream,
response,
bytesWritten => FinalizeStreamingDownload(hash, tempFileName, destinationFilePath, completion, true, bytesWritten),
bytesWritten => FinalizeStreamingDownload(hash, tempFileName, destinationFilePath, completion, false, bytesWritten),
_logger);
return new CachedFileStreamResult(null, stream, response.Content.Headers.ContentLength);
}
public async Task DownloadFileWhenRequired(string hash)
{
var fi = FilePathUtil.GetFileInfoForHash(_hotStoragePath, hash);
@@ -219,4 +311,174 @@ public sealed class CachedFileProvider : IDisposable
{
return hashes.Exists(_currentTransfers.Keys.Contains);
}
}
private void FinalizeStreamingDownload(string hash, string? tempFileName, string destinationFilePath,
TaskCompletionSource<bool> completion, bool success, long bytesWritten)
{
try
{
if (success)
{
if (!string.IsNullOrEmpty(tempFileName))
{
File.Move(tempFileName, destinationFilePath, true);
}
var fi = FilePathUtil.GetFileInfoForHash(_hotStoragePath, hash);
if (fi != null)
{
_metrics.IncGauge(MetricsAPI.GaugeFilesTotal);
_metrics.IncGauge(MetricsAPI.GaugeFilesTotalSize, fi.Length);
_fileStatisticsService.LogFile(hash, fi.Length);
_inventoryReporter.ReportAdded(hash);
}
}
else if (!string.IsNullOrEmpty(tempFileName))
{
try { File.Delete(tempFileName); } catch { /* ignore */ }
}
}
catch (Exception ex)
{
_logger.LogWarning(ex, "Failed to finalize streaming download for {hash} after {bytes} bytes", hash, bytesWritten);
}
finally
{
_metrics.DecGauge(MetricsAPI.GaugeFilesDownloadingFromCache);
_currentTransfers.Remove(hash, out _);
completion.TrySetResult(success);
}
}
private sealed class StreamingCacheWriteStream : Stream
{
private readonly Stream _source;
private readonly FileStream _cacheStream;
private readonly HttpResponseMessage _response;
private readonly Action<long> _onSuccess;
private readonly Action<long> _onFailure;
private readonly ILogger _logger;
private long _bytesWritten;
private int _completed;
public StreamingCacheWriteStream(Stream source, FileStream cacheStream, HttpResponseMessage response,
Action<long> onSuccess, Action<long> onFailure, ILogger logger)
{
_source = source;
_cacheStream = cacheStream;
_response = response;
_onSuccess = onSuccess;
_onFailure = onFailure;
_logger = logger;
}
public override bool CanRead => true;
public override bool CanSeek => false;
public override bool CanWrite => false;
public override long Length => _bytesWritten;
public override long Position
{
get => _bytesWritten;
set => throw new NotSupportedException();
}
public override int Read(byte[] buffer, int offset, int count)
{
if (Volatile.Read(ref _completed) != 0)
{
return 0;
}
try
{
int bytesRead = _source.Read(buffer, offset, count);
if (bytesRead > 0)
{
_cacheStream.Write(buffer, offset, bytesRead);
_bytesWritten += bytesRead;
return bytesRead;
}
Complete(true);
return 0;
}
catch (Exception ex)
{
_logger.LogWarning(ex, "Streaming download failed while reading");
Complete(false);
throw;
}
}
public override async ValueTask<int> ReadAsync(Memory<byte> buffer, CancellationToken cancellationToken = default)
{
if (Volatile.Read(ref _completed) != 0)
{
return 0;
}
try
{
int bytesRead = await _source.ReadAsync(buffer, cancellationToken).ConfigureAwait(false);
if (bytesRead > 0)
{
await _cacheStream.WriteAsync(buffer.Slice(0, bytesRead), cancellationToken).ConfigureAwait(false);
_bytesWritten += bytesRead;
return bytesRead;
}
Complete(true);
return 0;
}
catch (Exception ex)
{
_logger.LogWarning(ex, "Streaming download failed while reading");
Complete(false);
throw;
}
}
public override void Flush()
{
// no-op
}
protected override void Dispose(bool disposing)
{
if (disposing)
{
Complete(false);
}
base.Dispose(disposing);
}
private void Complete(bool success)
{
if (Interlocked.Exchange(ref _completed, 1) != 0)
{
return;
}
try { _cacheStream.Flush(); } catch { /* ignore */ }
try { _cacheStream.Dispose(); } catch { /* ignore */ }
try { _source.Dispose(); } catch { /* ignore */ }
try { _response.Dispose(); } catch { /* ignore */ }
if (success)
{
_onSuccess(_bytesWritten);
}
else
{
_onFailure(_bytesWritten);
}
}
public override long Seek(long offset, SeekOrigin origin) => throw new NotSupportedException();
public override void SetLength(long value) => throw new NotSupportedException();
public override void Write(byte[] buffer, int offset, int count) => throw new NotSupportedException();
}
}
public readonly record struct CachedFileStreamResult(FileInfo? File, Stream? Stream, long? ContentLength);