FileUploadProvider.cs

Go to the documentation of this file.

using System;
using System.IO;
using System.Threading;
using System.Threading.Tasks;
 
using Microsoft.AspNetCore.WebUtilities;
 
using Tgstation.Server.Api.Models;
using Tgstation.Server.Api.Models.Response;
using Tgstation.Server.Host.IO;
 
namespace Tgstation.Server.Host.Transfer
{
    sealed class FileUploadProvider : IFileUploadTicket
    {
        public FileTicketResponse Ticket { get; }
 
        readonly CancellationTokenSource ticketExpiryCts;
 
        readonly TaskCompletionSource<Stream?> streamTcs;
 
        readonly TaskCompletionSource completionTcs;
 
        readonly FileUploadStreamKind streamKind;
 
        ErrorMessageResponse? errorMessage;
 
        public FileUploadProvider(FileTicketResponse ticket, FileUploadStreamKind streamKind)
        {
            Ticket = ticket ?? throw new ArgumentNullException(nameof(ticket));
 
            ticketExpiryCts = new CancellationTokenSource();
            streamTcs = new TaskCompletionSource<Stream?>();
            completionTcs = new TaskCompletionSource();
            this.streamKind = streamKind;
        }
 
        public ValueTask DisposeAsync()
        {
            ticketExpiryCts.Dispose();
            completionTcs.TrySetResult();
            return ValueTask.CompletedTask;
        }
 
        public void Expire()
        {
            if (!completionTcs.Task.IsCompleted)
                ticketExpiryCts.Cancel();
        }
 
        public async ValueTask<ErrorMessageResponse?> Completion(Stream stream, CancellationToken cancellationToken)
        {
            ArgumentNullException.ThrowIfNull(stream);
 
            if (ticketExpiryCts.IsCancellationRequested)
                return new ErrorMessageResponse(ErrorCode.ResourceNotPresent);
 
            Stream? bufferedStream = null;
            try
            {
                switch (streamKind)
                {
                    case FileUploadStreamKind.ForSynchronousIO:
                        // big reads, we should buffer to disk
                        // NOTE: We do this here ONLY because we can't use async methods in a synchronous context
                        // Ideally we should be doing 0 buffering here, but this is a compromise
                        bufferedStream = new FileBufferingReadStream(stream, DefaultIOManager.DefaultBufferSize);
                        await bufferedStream.DrainAsync(cancellationToken);
                        break;
                    case FileUploadStreamKind.None:
                        break;
                    default:
                        throw new InvalidOperationException($"Invalid FileUploadStreamKind: {streamKind}");
                }
            }
            catch
            {
                if (bufferedStream != null)
                    await bufferedStream.DisposeAsync();
 
                throw;
            }
 
            await using (bufferedStream)
            {
                streamTcs.TrySetResult(bufferedStream ?? stream);
 
                await completionTcs.Task.WaitAsync(cancellationToken);
                return errorMessage;
            }
        }
 
        public void SetError(ErrorCode errorCode, string? additionalData)
        {
            if (errorMessage != null)
                throw new InvalidOperationException("Error already set!");
 
            errorMessage = new ErrorMessageResponse(errorCode)
            {
                AdditionalData = additionalData,
            };
            completionTcs.TrySetResult();
        }
 
        public async ValueTask<Stream> GetResult(CancellationToken cancellationToken)
            => await ((IFileUploadTicket)this).GetResult(cancellationToken)
            ?? throw new InvalidOperationException("Upload ticket expired!");
 
        async ValueTask<Stream?> IFileUploadTicket.GetResult(CancellationToken cancellationToken)
        {
            using (cancellationToken.Register(() => streamTcs.TrySetCanceled(cancellationToken)))
            using (ticketExpiryCts.Token.Register(() => streamTcs.TrySetResult(null)))
                return await streamTcs.Task;
        }
    }
}