wabbajack/Wabbajack.Compression.BSA/DiskSlabAllocator.cs

using System.Collections.Concurrent;
using System.IO;
using System.Threading;
using System.Threading.Tasks;
using Wabbajack.Paths.IO;

namespace Wabbajack.Compression.BSA
{
    public class DiskSlabAllocator
    {
        private readonly TemporaryFileManager _manager;
        private readonly long _maxMemorySize;
        private readonly ConcurrentBag<TemporaryPath> _files = new();
        private long _memorySize;
        private readonly ConcurrentBag<Stream> _streams = new();

        public DiskSlabAllocator(TemporaryFileManager manager, long maxMemorySize = 1024 * 1024 * 256)
        {
            _manager = manager;
            _memorySize = 0;
            _maxMemorySize = maxMemorySize;
        }

        public async Task DisposeAsync()
        {
            foreach (var s in _streams)
                await s.DisposeAsync();

            foreach (var file in _files) file.Dispose();
        }

        public Stream Allocate(long rLength)
        {
            var newSize = Interlocked.Add(ref _memorySize, rLength);
            if (newSize < _maxMemorySize)
            {
                var stream = new MemoryStream();
                _streams.Add(stream);
                return stream;
            }
            else
            {
                var tempFile = _manager.CreateFile();
                var stream = tempFile.Path.Open(FileMode.Create, FileAccess.ReadWrite, FileShare.None);
                _streams.Add(stream);
                _files.Add(tempFile);
                return stream;
            }
        }
    }
}