BufferedFileStreamProvider.cs

Go to the documentation of this file.

using System;
using System.IO;
using System.Threading;
using System.Threading.Tasks;
 
using Tgstation.Server.Host.Utils;
 
namespace Tgstation.Server.Host.IO
{
    public sealed class BufferedFileStreamProvider : ISeekableFileStreamProvider
    {
        public bool Disposed => buffer == null;
 
        readonly Stream input;
 
        readonly SemaphoreSlim semaphore;
 
        volatile MemoryStream buffer;
 
        volatile bool buffered;
 
        public BufferedFileStreamProvider(Stream input)
        {
            this.input = input ?? throw new ArgumentNullException(nameof(input));
 
            semaphore = new SemaphoreSlim(1);
            try
            {
                buffer = new MemoryStream();
            }
            catch
            {
                semaphore.Dispose();
                throw;
            }
        }
 
        public async ValueTask DisposeAsync()
        {
            MemoryStream localBuffer;
            lock (semaphore)
            {
                localBuffer = buffer;
                if (localBuffer == null)
                    return;
 
                // important to drop the reference so it can properly GC
                // The implementation of MemoryStream doesn't fucking do this for some reason
                buffer = null;
                buffered = true;
            }
 
            var bufferDispose = localBuffer.DisposeAsync();
            semaphore.Dispose();
            await bufferDispose;
        }
 
        public async Task<Stream> GetResult(CancellationToken cancellationToken)
        {
            var (sharedStream, _) = await GetResultInternal(cancellationToken);
            return sharedStream;
        }
 
        public async Task<MemoryStream> GetOwnedResult(CancellationToken cancellationToken)
        {
            var (sharedStream, length) = await GetResultInternal(cancellationToken);
            return new MemoryStream(sharedStream.GetBuffer(), 0, (int)length, false, true);
        }
 
        public Task EnsureBuffered(CancellationToken cancellationToken) => GetResultInternal(cancellationToken);
 
        async Task<(MemoryStream, long)> GetResultInternal(CancellationToken cancellationToken)
        {
            if (!buffered)
                using (await SemaphoreSlimContext.Lock(semaphore, cancellationToken))
                    if (!buffered)
                    {
                        MemoryStream localBuffer;
                        lock (semaphore)
                            localBuffer = buffer ?? throw new ObjectDisposedException(nameof(BufferedFileStreamProvider));
 
                        await input.CopyToAsync(localBuffer, cancellationToken);
                        localBuffer.Seek(0, SeekOrigin.Begin);
                        buffered = true;
                        return (localBuffer, localBuffer.Length);
                    }
 
            lock (semaphore)
            {
                var localBuffer = buffer ?? throw new ObjectDisposedException(nameof(BufferedFileStreamProvider));
                return (
                    localBuffer,
                    localBuffer.Length);
            }
        }
    }
}