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 ValueTask<Stream> GetResult(CancellationToken cancellationToken)
        {
            var (sharedStream, _) = await GetResultInternal(cancellationToken);
            return sharedStream;
        }
 
        public async ValueTask<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).AsTask();
 
        async ValueTask<(MemoryStream Stream, long StreamLength)> 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 (Stream: localBuffer, StreamLength: localBuffer.Length);
                    }
 
            lock (semaphore)
            {
                var localBuffer = buffer ?? throw new ObjectDisposedException(nameof(BufferedFileStreamProvider));
                return (
                    Stream: localBuffer,
                    StreamLength: localBuffer.Length);
            }
        }
    }
}