RequestFileStreamProvider.cs

Go to the documentation of this file.

using System;
using System.IO;
using System.Net.Http;
using System.Threading;
using System.Threading.Tasks;
 
using Tgstation.Server.Common.Http;
 
namespace Tgstation.Server.Host.IO
{
    sealed class RequestFileStreamProvider : IFileStreamProvider
    {
        readonly IHttpClient httpClient;
 
        readonly HttpRequestMessage requestMessage;
 
        readonly CancellationTokenSource downloadCts;
 
        Task<CachedResponseStream>? downloadTask;
 
        volatile bool disposed;
 
        public RequestFileStreamProvider(IHttpClient httpClient, HttpRequestMessage requestMessage)
        {
            this.httpClient = httpClient ?? throw new ArgumentNullException(nameof(httpClient));
            this.requestMessage = requestMessage ?? throw new ArgumentNullException(nameof(requestMessage));
 
            downloadCts = new CancellationTokenSource();
        }
 
        public async ValueTask DisposeAsync()
        {
            Task<CachedResponseStream>? localDownloadTask;
            lock (downloadCts)
            {
                if (disposed)
                    return;
 
                disposed = true;
 
                localDownloadTask = downloadTask;
                downloadTask = null;
            }
 
            downloadCts.Cancel();
            downloadCts.Dispose();
 
            if (localDownloadTask != null)
            {
                CachedResponseStream result;
                try
                {
                    result = await localDownloadTask;
                }
                catch
                {
                    // Unsightly yes, but, if we're here, that means someone called GetResult. So, either they handled the error or decided to ignore it
                    return;
                }
 
                await result.DisposeAsync();
            }
 
            requestMessage.Dispose();
            httpClient.Dispose();
        }
 
        public async ValueTask<Stream> GetResult(CancellationToken cancellationToken)
        {
            if (disposed)
                throw new ObjectDisposedException(nameof(RequestFileStreamProvider));
 
            Task<CachedResponseStream> localTask;
            using (cancellationToken.Register(() => downloadCts.Cancel()))
            {
                lock (downloadCts)
                {
                    downloadTask ??= InitiateDownload(downloadCts.Token);
                    localTask = downloadTask;
                }
 
                return await localTask;
            }
        }
 
        async Task<CachedResponseStream> InitiateDownload(CancellationToken cancellationToken)
        {
            var response = await httpClient.SendAsync(requestMessage, HttpCompletionOption.ResponseHeadersRead, cancellationToken);
            try
            {
                response.EnsureSuccessStatusCode();
                return await CachedResponseStream.Create(response);
            }
            catch
            {
                response.Dispose();
                throw;
            }
        }
    }
}