Merge pull request #317 from Washi1337/feature/rentable-memory-stream…

…-writers

Rentable BinaryStreamWriters that write to temporary MemoryStream

src/AsmResolver.DotNet/Builder/Metadata/Blob/BlobStreamBuffer.cs

@@ -17,6 +17,8 @@ public class BlobStreamBuffer : IMetadataStreamBuffer
         private readonly BinaryStreamWriter _writer;
         private readonly Dictionary<byte[], uint> _blobs = new(ByteArrayEqualityComparer.Instance);
 
+        private readonly MemoryStreamWriterPool _blobWriterPool = new();
+
         /// <summary>
         /// Creates a new blob stream buffer with the default blob stream name.
         /// </summary>
@@ -116,11 +118,11 @@ public uint GetBlobIndex(ITypeCodedIndexProvider provider, BlobSignature? signat
                 return 0u;
 
             // Serialize blob.
-            using var stream = new MemoryStream();
-            var writer = new BinaryStreamWriter(stream);
-            signature.Write(new BlobSerializationContext(writer, provider, errorListener));
 
-            return GetBlobIndex(stream.ToArray());
+            using var rentedWriter = _blobWriterPool.Rent();
+            signature.Write(new BlobSerializationContext(rentedWriter.Writer, provider, errorListener));
+
+            return GetBlobIndex(rentedWriter.GetData());
         }
 
         /// <summary>

src/AsmResolver.DotNet/Code/Cil/CilMethodBodySerializer.cs

@@ -15,6 +15,8 @@ namespace AsmResolver.DotNet.Code.Cil
     /// </summary>
     public class CilMethodBodySerializer : IMethodBodySerializer
     {
+        private readonly MemoryStreamWriterPool _writerPool = new();
+
         /// <summary>
         /// Gets or sets the value of an override switch indicating whether the max stack should always be recalculated
         /// or should always be preserved.
@@ -88,7 +90,7 @@ public ISegmentReference SerializeMethodBody(MethodBodySerializationContext cont
             }
 
             // Serialize CIL stream.
-            var code = BuildRawCodeStream(context, body);
+            byte[] code = BuildRawCodeStream(context, body);
 
             // Build method body.
             var rawBody = body.IsFat
@@ -98,8 +100,7 @@ public ISegmentReference SerializeMethodBody(MethodBodySerializationContext cont
             return rawBody.ToReference();
         }
 
-        private static CilRawMethodBody BuildTinyMethodBody(byte[] code) =>
-            new CilRawTinyMethodBody(code);
+        private static CilRawMethodBody BuildTinyMethodBody(byte[] code) => new CilRawTinyMethodBody(code);
 
         private CilRawMethodBody BuildFatMethodBody(MethodBodySerializationContext context, CilMethodBody body, byte[] code)
         {
@@ -137,35 +138,33 @@ private CilRawMethodBody BuildFatMethodBody(MethodBodySerializationContext conte
             return fatBody;
         }
 
-        private static byte[] BuildRawCodeStream(MethodBodySerializationContext context, CilMethodBody body)
+        private byte[] BuildRawCodeStream(MethodBodySerializationContext context, CilMethodBody body)
         {
-            using var codeStream = new MemoryStream();
             var bag = context.ErrorListener;
 
-            var writer = new BinaryStreamWriter(codeStream);
+            using var rentedWriter = _writerPool.Rent();
             var assembler = new CilAssembler(
-                writer,
+                rentedWriter.Writer,
                 new CilOperandBuilder(context.TokenProvider, bag),
                 body.Owner.SafeToString,
                 bag);
 
             assembler.WriteInstructions(body.Instructions);
 
-            return codeStream.ToArray();
+            return rentedWriter.GetData();
         }
 
         private byte[] SerializeExceptionHandlers(MethodBodySerializationContext context, IList<CilExceptionHandler> exceptionHandlers, bool needsFatFormat)
         {
-            using var sectionStream = new MemoryStream();
-            var writer = new BinaryStreamWriter(sectionStream);
+            using var rentedWriter = _writerPool.Rent();
 
             for (int i = 0; i < exceptionHandlers.Count; i++)
             {
                 var handler = exceptionHandlers[i];
-                WriteExceptionHandler(context, writer, handler, needsFatFormat);
+                WriteExceptionHandler(context, rentedWriter.Writer, handler, needsFatFormat);
             }
 
-            return sectionStream.ToArray();
+            return rentedWriter.GetData();
         }
 
         private void WriteExceptionHandler(MethodBodySerializationContext context, IBinaryStreamWriter writer, CilExceptionHandler handler, bool useFatFormat)

src/AsmResolver/IO/BinaryStreamWriter.cs

@@ -8,107 +8,113 @@ namespace AsmResolver.IO
     /// </summary>
     public class BinaryStreamWriter : IBinaryStreamWriter
     {
-        private readonly Stream _stream;
-
         /// <summary>
         /// Creates a new binary stream writer using the provided output stream.
         /// </summary>
         /// <param name="stream">The stream to write to.</param>
         public BinaryStreamWriter(Stream stream)
         {
-            _stream = stream ?? throw new ArgumentNullException(nameof(stream));
+            BaseStream = stream ?? throw new ArgumentNullException(nameof(stream));
+        }
+
+        /// <summary>
+        /// Gets the stream this writer writes to.
+        /// </summary>
+        public Stream BaseStream
+        {
+            get;
         }
 
         /// <inheritdoc />
         public ulong Offset
         {
-            get => (uint) _stream.Position;
+            get => (uint) BaseStream.Position;
             set
             {
                 // Check if position actually changed before actually setting. If we don't do this, this can cause
                 // performance issues on some systems. See https://github.com/Washi1337/AsmResolver/issues/232
-                if (_stream.Position != (long) value)
-                    _stream.Position = (long) value;
+                if (BaseStream.Position != (long) value)
+                    BaseStream.Position = (long) value;
             }
         }
 
         /// <inheritdoc />
-        public uint Length => (uint) _stream.Length;
+        public uint Length => (uint) BaseStream.Length;
 
         /// <inheritdoc />
         public void WriteBytes(byte[] buffer, int startIndex, int count)
         {
-            _stream.Write(buffer, startIndex, count);
+            BaseStream.Write(buffer, startIndex, count);
         }
 
         /// <inheritdoc />
         public void WriteByte(byte value)
         {
-            _stream.WriteByte(value);
+            BaseStream.WriteByte(value);
         }
 
         /// <inheritdoc />
         public void WriteUInt16(ushort value)
         {
-            _stream.WriteByte((byte) (value & 0xFF));
-            _stream.WriteByte((byte) ((value >> 8) & 0xFF));
+            BaseStream.WriteByte((byte) (value & 0xFF));
+            BaseStream.WriteByte((byte) ((value >> 8) & 0xFF));
         }
 
         /// <inheritdoc />
         public void WriteUInt32(uint value)
         {
-            _stream.WriteByte((byte) (value & 0xFF));
-            _stream.WriteByte((byte) ((value >> 8) & 0xFF));
-            _stream.WriteByte((byte) ((value >> 16) & 0xFF));
-            _stream.WriteByte((byte) ((value >> 24) & 0xFF));
+            BaseStream.WriteByte((byte) (value & 0xFF));
+            BaseStream.WriteByte((byte) ((value >> 8) & 0xFF));
+            BaseStream.WriteByte((byte) ((value >> 16) & 0xFF));
+            BaseStream.WriteByte((byte) ((value >> 24) & 0xFF));
         }
 
         /// <inheritdoc />
         public void WriteUInt64(ulong value)
         {
-            _stream.WriteByte((byte) (value & 0xFF));
-            _stream.WriteByte((byte) ((value >> 8) & 0xFF));
-            _stream.WriteByte((byte) ((value >> 16) & 0xFF));
-            _stream.WriteByte((byte) ((value >> 24) & 0xFF));
-            _stream.WriteByte((byte) ((value >> 32) & 0xFF));
-            _stream.WriteByte((byte) ((value >> 40) & 0xFF));
-            _stream.WriteByte((byte) ((value >> 48) & 0xFF));
-            _stream.WriteByte((byte) ((value >> 56) & 0xFF));
+            BaseStream.WriteByte((byte) (value & 0xFF));
+            BaseStream.WriteByte((byte) ((value >> 8) & 0xFF));
+            BaseStream.WriteByte((byte) ((value >> 16) & 0xFF));
+            BaseStream.WriteByte((byte) ((value >> 24) & 0xFF));
+            BaseStream.WriteByte((byte) ((value >> 32) & 0xFF));
+            BaseStream.WriteByte((byte) ((value >> 40) & 0xFF));
+            BaseStream.WriteByte((byte) ((value >> 48) & 0xFF));
+            BaseStream.WriteByte((byte) ((value >> 56) & 0xFF));
         }
 
         /// <inheritdoc />
         public void WriteSByte(sbyte value)
         {
-            _stream.WriteByte(unchecked((byte) value));
+            BaseStream.WriteByte(unchecked((byte) value));
         }
 
         /// <inheritdoc />
         public void WriteInt16(short value)
         {
-            _stream.WriteByte((byte) (value & 0xFF));
-            _stream.WriteByte((byte) ((value >> 8) & 0xFF));
+            BaseStream.WriteByte((byte) (value & 0xFF));
+            BaseStream.WriteByte((byte) ((value >> 8) & 0xFF));
         }
 
         /// <inheritdoc />
         public void WriteInt32(int value)
         {
-            _stream.WriteByte((byte) (value & 0xFF));
-            _stream.WriteByte((byte) ((value >> 8) & 0xFF));
-            _stream.WriteByte((byte) ((value >> 16) & 0xFF));
-            _stream.WriteByte((byte) ((value >> 24) & 0xFF));
+            BaseStream.WriteByte((byte) (value & 0xFF));
+            BaseStream.WriteByte((byte) ((value >> 8) & 0xFF));
+            BaseStream.WriteByte((byte) ((value >> 16) & 0xFF));
+            BaseStream.WriteByte((byte) ((value >> 24) & 0xFF));
         }
 
         /// <inheritdoc />
         public void WriteInt64(long value)
         {
-            _stream.WriteByte((byte) (value & 0xFF));
-            _stream.WriteByte((byte) ((value >> 8) & 0xFF));
-            _stream.WriteByte((byte) ((value >> 16) & 0xFF));
-            _stream.WriteByte((byte) ((value >> 24) & 0xFF));
-            _stream.WriteByte((byte) ((value >> 32) & 0xFF));
-            _stream.WriteByte((byte) ((value >> 40) & 0xFF));
-            _stream.WriteByte((byte) ((value >> 48) & 0xFF));
-            _stream.WriteByte((byte) ((value >> 56) & 0xFF));
+            BaseStream.WriteByte((byte) (value & 0xFF));
+            BaseStream.WriteByte((byte) ((value >> 8) & 0xFF));
+            BaseStream.WriteByte((byte) ((value >> 16) & 0xFF));
+            BaseStream.WriteByte((byte) ((value >> 24) & 0xFF));
+            BaseStream.WriteByte((byte) ((value >> 32) & 0xFF));
+            BaseStream.WriteByte((byte) ((value >> 40) & 0xFF));
+            BaseStream.WriteByte((byte) ((value >> 48) & 0xFF));
+            BaseStream.WriteByte((byte) ((value >> 56) & 0xFF));
         }
 
         /// <inheritdoc />

src/AsmResolver/IO/MemoryStreamWriterPool.cs

@@ -0,0 +1,87 @@
+using System;
+using System.Collections.Concurrent;
+using System.IO;
+
+namespace AsmResolver.IO
+{
+    /// <summary>
+    /// Provides a pool of reusable instances of <see cref="BinaryStreamWriter"/> that are meant to be used for
+    /// constructing byte arrays.
+    /// </summary>
+    /// <remarks>
+    /// This class is thread-safe. All threads are allowed to rent and return writers from this pool simultaneously.
+    /// </remarks>
+    public class MemoryStreamWriterPool
+    {
+        private readonly ConcurrentQueue<BinaryStreamWriter> _writers = new();
+
+        /// <summary>
+        /// Rents a single binary stream writer.
+        /// </summary>
+        /// <returns>The writer.</returns>
+        public RentedWriter Rent()
+        {
+            if (!_writers.TryDequeue(out var writer))
+                writer = new BinaryStreamWriter(new MemoryStream());
+
+            writer.BaseStream.SetLength(0);
+            return new RentedWriter(this, writer);
+        }
+
+        private void Return(BinaryStreamWriter writer) => _writers.Enqueue(writer);
+
+        /// <summary>
+        /// Represents a single instance of a <see cref="BinaryStreamWriter"/> that is rented by a writer pool.
+        /// </summary>
+        public ref struct RentedWriter
+        {
+            private bool _isDisposed = false;
+            private readonly BinaryStreamWriter _writer;
+
+            internal RentedWriter(MemoryStreamWriterPool pool, BinaryStreamWriter writer)
+            {
+                Pool = pool;
+                _writer = writer;
+            }
+
+            /// <summary>
+            /// Gets the pool the writer was rented from.
+            /// </summary>
+            public MemoryStreamWriterPool Pool
+            {
+                get;
+            }
+
+            /// <summary>
+            /// Gets the writer instance.
+            /// </summary>
+            public BinaryStreamWriter Writer
+            {
+                get
+                {
+                    if (_isDisposed)
+                        throw new ObjectDisposedException(nameof(Writer));
+                    return _writer;
+                }
+            }
+
+            /// <summary>
+            /// Gets the data that was written to the temporary stream.
+            /// </summary>
+            /// <returns></returns>
+            public byte[] GetData() => ((MemoryStream) Writer.BaseStream).ToArray();
+
+            /// <summary>
+            /// Returns the stream writer to the pool.
+            /// </summary>
+            public void Dispose()
+            {
+                if (_isDisposed)
+                    return;
+
+                Pool.Return(Writer);
+                _isDisposed = true;
+            }
+        }
+    }
+}