Merge pull request #1347 from adamhathcock/adam/deflate-allocations

Add pooling for arrays for deflate

Author: adamhathcock

src/SharpCompress/Compressors/Deflate/CRC32.cs

@@ -33,6 +33,7 @@
 // ------------------------------------------------------------------
 
 using System;
+using System.Buffers;
 using System.IO;
 
 namespace SharpCompress.Compressors.Deflate;
@@ -120,22 +121,29 @@ public uint GetCrc32AndCopy(Stream input, Stream? output)
         {
             //UInt32 crc32Result;
             //crc32Result = 0xFFFFFFFF;
-            var buffer = new byte[BUFFER_SIZE];
-            var readSize = BUFFER_SIZE;
-
-            TotalBytesRead = 0;
-            var count = input.Read(buffer, 0, readSize);
-            output?.Write(buffer, 0, count);
-            TotalBytesRead += count;
-            while (count > 0)
+            var buffer = ArrayPool<byte>.Shared.Rent(BUFFER_SIZE);
+            try
             {
-                SlurpBlock(buffer, 0, count);
-                count = input.Read(buffer, 0, readSize);
+                var readSize = BUFFER_SIZE;
+
+                TotalBytesRead = 0;
+                var count = input.Read(buffer, 0, readSize);
                 output?.Write(buffer, 0, count);
                 TotalBytesRead += count;
-            }
+                while (count > 0)
+                {
+                    SlurpBlock(buffer, 0, count);
+                    count = input.Read(buffer, 0, readSize);
+                    output?.Write(buffer, 0, count);
+                    TotalBytesRead += count;
+                }
 
-            return ~runningCrc32Result;
+                return ~runningCrc32Result;
+            }
+            finally
+            {
+                ArrayPool<byte>.Shared.Return(buffer, clearArray: true);
+            }
         }
     }
 

src/SharpCompress/Compressors/Deflate/DeflateManager.cs

@@ -69,6 +69,7 @@
 // -----------------------------------------------------------------------
 
 using System;
+using System.Buffers;
 using SharpCompress.Algorithms;
 using SharpCompress.Common;
 
@@ -1726,9 +1727,12 @@ CompressionStrategy strategy
         hash_mask = hash_size - 1;
         hash_shift = ((hash_bits + MIN_MATCH - 1) / MIN_MATCH);
 
-        window = new byte[w_size * 2];
-        prev = new short[w_size];
-        head = new short[hash_size];
+        window = ArrayPool<byte>.Shared.Rent(w_size * 2);
+        prev = ArrayPool<short>.Shared.Rent(w_size);
+        head = ArrayPool<short>.Shared.Rent(hash_size);
+        Array.Clear(window, 0, w_size * 2);
+        Array.Clear(prev, 0, w_size);
+        Array.Clear(head, 0, hash_size);
 
         // for memLevel==8, this will be 16384, 16k
         lit_bufsize = 1 << (memLevel + 6);
@@ -1737,7 +1741,8 @@ CompressionStrategy strategy
         // the output distance codes, and the output length codes (aka tree).
         // orig comment: This works just fine since the average
         // output size for (length,distance) codes is <= 24 bits.
-        pending = new byte[lit_bufsize * 4];
+        pending = ArrayPool<byte>.Shared.Rent(lit_bufsize * 4);
+        Array.Clear(pending, 0, lit_bufsize * 4);
         _distanceOffset = lit_bufsize;
         _lengthOffset = (1 + 2) * lit_bufsize;
 
@@ -1776,20 +1781,46 @@ internal void Reset()
 
     internal int End()
     {
+        var result = ZlibConstants.Z_OK;
         if (status != INIT_STATE && status != BUSY_STATE && status != FINISH_STATE)
         {
-            return ZlibConstants.Z_STREAM_ERROR;
+            result = ZlibConstants.Z_STREAM_ERROR;
+        }
+        else if (status == BUSY_STATE)
+        {
+            result = ZlibConstants.Z_DATA_ERROR;
         }
 
         // Deallocate in reverse order of allocations:
-        pending = null;
-        head = null;
-        prev = null;
-        window = null;
+        ReturnBuffers();
 
         // free
         // dstate=null;
-        return status == BUSY_STATE ? ZlibConstants.Z_DATA_ERROR : ZlibConstants.Z_OK;
+        return result;
+    }
+
+    private void ReturnBuffers()
+    {
+        if (pending is not null)
+        {
+            ArrayPool<byte>.Shared.Return(pending, clearArray: true);
+            pending = null;
+        }
+        if (head is not null)
+        {
+            ArrayPool<short>.Shared.Return(head, clearArray: true);
+            head = null;
+        }
+        if (prev is not null)
+        {
+            ArrayPool<short>.Shared.Return(prev, clearArray: true);
+            prev = null;
+        }
+        if (window is not null)
+        {
+            ArrayPool<byte>.Shared.Return(window, clearArray: true);
+            window = null;
+        }
     }
 
     private void SetDeflater() =>

src/SharpCompress/Compressors/Deflate/Inflate.cs

@@ -124,7 +124,8 @@ internal sealed class InflateBlocks
     internal InflateBlocks(ZlibCodec codec, object checkfn, int w)
     {
         _codec = codec;
-        hufts = new int[MANY * 3];
+        hufts = ArrayPool<int>.Shared.Rent(MANY * 3);
+        Array.Clear(hufts, 0, MANY * 3);
         window = MemoryPool<byte>.Shared.Rent(w);
         end = w;
         this.checkfn = checkfn;
@@ -719,7 +720,11 @@ internal void Free()
         Reset();
         window?.Dispose();
         window = null;
-        hufts = null;
+        if (hufts is not null)
+        {
+            ArrayPool<int>.Shared.Return(hufts, clearArray: true);
+            hufts = null;
+        }
     }
 
     internal void SetDictionary(byte[] d, int start, int n)

src/SharpCompress/Compressors/Deflate/ZlibBaseStream.cs

@@ -27,6 +27,7 @@
 // ------------------------------------------------------------------
 
 using System;
+using System.Buffers;
 using System.Buffers.Binary;
 using System.Collections.Generic;
 using System.IO;
@@ -136,7 +137,7 @@ private ZlibCodec z
         }
     }
 
-    private byte[] workingBuffer => _workingBuffer ??= new byte[_bufferSize];
+    private byte[] workingBuffer => _workingBuffer ??= ArrayPool<byte>.Shared.Rent(_bufferSize);
 
     public override void Write(byte[] buffer, int offset, int count)
     {
@@ -541,6 +542,7 @@ protected override void Dispose(bool disposing)
             finally
             {
                 end();
+                ReturnWorkingBuffer();
                 if (!_leaveOpen)
                 {
                     _stream?.Dispose();
@@ -575,6 +577,7 @@ public async ValueTask DisposeAsync()
         finally
         {
             end();
+            ReturnWorkingBuffer();
             if (_stream != null)
             {
                 if (!_leaveOpen)
@@ -593,6 +596,17 @@ public async ValueTask DisposeAsync()
         }
     }
 
+    private void ReturnWorkingBuffer()
+    {
+        if (_workingBuffer is null)
+        {
+            return;
+        }
+
+        ArrayPool<byte>.Shared.Return(_workingBuffer, clearArray: true);
+        _workingBuffer = null;
+    }
+
     public override void Flush()
     {
         // Only flush the underlying stream when in write mode

src/SharpCompress/Compressors/Deflate/ZlibCodec.cs

@@ -612,10 +612,9 @@ public int EndDeflate()
             throw new ZlibException("No Deflate State!");
         }
 
-        // TODO: dinoch Tue, 03 Nov 2009  15:39 (test this)
-        //int ret = dstate.End();
+        _ = dstate.End();
         dstate = null;
-        return ZlibConstants.Z_OK; //ret;
+        return ZlibConstants.Z_OK;
     }
 
     /// <summary>