Clear buffers before returning them to the ArrayPool

Snappier/Internal/ByteArrayPoolMemoryOwner.cs

@@ -45,6 +45,9 @@ public void Dispose()
             byte[]? innerArray = _innerArray;
             if (innerArray is not null)
             {
+                // Clear the used portion of the array before returning it to the pool
+                Memory.Span.Clear();
+
                 _innerArray = null;
                 Memory = default;
                 ArrayPool<byte>.Shared.Return(innerArray);

Snappier/Internal/Helpers.cs

@@ -1,8 +1,10 @@
 using System;
+using System.Buffers;
 using System.Buffers.Binary;
 using System.Diagnostics;
 using System.Runtime.CompilerServices;
 using System.Runtime.InteropServices;
+
 #if NET6_0_OR_GREATER
 using System.Numerics;
 using System.Runtime.Intrinsics;
@@ -45,6 +47,23 @@ public static int MaxCompressedLength(int sourceBytes)
             return 32 + sourceBytes + sourceBytes / 6 + 1;
         }
 
+        // Constant for MaxCompressedLength when passed Constants.BlockSize, keep this in sync with the above method
+        public const int MaxBlockCompressedLength = (int)(32 + Constants.BlockSize + Constants.BlockSize / 6 + 1);
+
+        /// <summary>
+        /// Clears the array and returns it to the pool, clearing only the used portion of the array.
+        /// This is a minor performance optimization to avoid clearing the entire array.
+        /// </summary>
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        public static void ClearAndReturn(byte[] array, int usedLength)
+        {
+            Debug.Assert(array is not null);
+            Debug.Assert(usedLength >= 0);
+
+            array.AsSpan(0, usedLength).Clear();
+            ArrayPool<byte>.Shared.Return(array);
+        }
+
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
         public static bool LeftShiftOverflows(byte value, int shift)
         {

Snappier/Internal/SnappyCompressor.cs

@@ -60,21 +60,18 @@ public bool TryCompress(ReadOnlySpan<byte> input, Span<byte> output, out int byt
                     // compress to a temporary buffer and copy the compressed data to the output span.
 
                     byte[] scratch = ArrayPool<byte>.Shared.Rent(maxOutput);
-                    try
+                    written = CompressFragment(fragment, scratch.AsSpan(), hashTable);
+                    if (output.Length < written)
                     {
-                        written = CompressFragment(fragment, scratch.AsSpan(), hashTable);
-                        if (output.Length < written)
-                        {
-                            bytesWritten = 0;
-                            return false;
-                        }
-
-                        scratch.AsSpan(0, written).CopyTo(output);
-                    }
-                    finally
-                    {
-                        ArrayPool<byte>.Shared.Return(scratch);
+                        Helpers.ClearAndReturn(scratch, written);
+                        bytesWritten = 0;
+                        return false;
                     }
+
+                    Span<byte> writtenScratch = scratch.AsSpan(0, written);
+                    writtenScratch.CopyTo(output);
+                    writtenScratch.Clear();
+                    ArrayPool<byte>.Shared.Return(scratch);
                 }
 
                 output = output.Slice(written);
@@ -132,19 +129,17 @@ public void Compress(ReadOnlySequence<byte> input, IBufferWriter<byte> bufferWri
 
                     int fragmentLength = (int)fragment.Length;
                     byte[] scratch = ArrayPool<byte>.Shared.Rent(fragmentLength);
-                    try
-                    {
-                        fragment.CopyTo(scratch);
 
-                        CompressFragment(scratch.AsSpan(0, fragmentLength), bufferWriter);
+                    Span<byte> usedScratch = scratch.AsSpan(0, fragmentLength);
+                    fragment.CopyTo(usedScratch);
 
-                        // Advance the length of the entire fragment
-                        input = input.Slice(position);
-                    }
-                    finally
-                    {
-                        ArrayPool<byte>.Shared.Return(scratch);
-                    }
+                    CompressFragment(usedScratch, bufferWriter);
+
+                    usedScratch.Clear();
+                    ArrayPool<byte>.Shared.Return(scratch);
+
+                    // Advance the length of the entire fragment
+                    input = input.Slice(position);
                 }
             }
         }

Snappier/Internal/SnappyDecompressor.cs

@@ -526,7 +526,8 @@ private int? ExpectedLength
                 {
                     if (_lookbackBufferArray is not null)
                     {
-                        ArrayPool<byte>.Shared.Return(_lookbackBufferArray);
+                        // Clear the used portion of the lookback buffer before returning
+                        Helpers.ClearAndReturn(_lookbackBufferArray, _lookbackPosition);
                     }
 
                     if (BufferWriter is not null)
@@ -731,7 +732,9 @@ public void Dispose()
         {
             if (_lookbackBufferArray is not null)
             {
-                ArrayPool<byte>.Shared.Return(_lookbackBufferArray);
+                // Clear the used portion of the lookback buffer before returning
+                Helpers.ClearAndReturn(_lookbackBufferArray, _lookbackPosition);
+
                 _lookbackBufferArray = null;
                 _lookbackBuffer = default;
             }

Snappier/Internal/SnappyStreamCompressor.cs

@@ -276,7 +276,7 @@ private void EnsureBuffer()
         {
             // Allocate enough room for the stream header and block headers
             _outputBuffer ??=
-                ArrayPool<byte>.Shared.Rent(Helpers.MaxCompressedLength((int) Constants.BlockSize) + 8 + SnappyHeader.Length);
+                ArrayPool<byte>.Shared.Rent(Helpers.MaxBlockCompressedLength + 8 + SnappyHeader.Length);
 
             // Allocate enough room for the stream header and block headers
             _inputBuffer ??= ArrayPool<byte>.Shared.Rent((int) Constants.BlockSize);
@@ -289,12 +289,12 @@ public void Dispose()
 
             if (_outputBuffer is not null)
             {
-                ArrayPool<byte>.Shared.Return(_outputBuffer);
+                ArrayPool<byte>.Shared.Return(_outputBuffer, clearArray: true);
                 _outputBuffer = null;
             }
             if (_inputBuffer is not null)
             {
-                ArrayPool<byte>.Shared.Return(_inputBuffer);
+                ArrayPool<byte>.Shared.Return(_inputBuffer, clearArray: true);
                 _inputBuffer = null;
             }
         }

Snappier/Snappy.cs

@@ -102,21 +102,16 @@ public static IMemoryOwner<byte> CompressToMemory(ReadOnlySpan<byte> input)
         {
             byte[] buffer = ArrayPool<byte>.Shared.Rent(GetMaxCompressedLength(input.Length));
 
-            try
+            if (!TryCompress(input, buffer, out int length))
             {
-                if (!TryCompress(input, buffer, out int length))
-                {
-                    // Should be unreachable since we're allocating a buffer of the correct size.
-                    ThrowHelper.ThrowInvalidOperationException();
-                }
+                // The amount of data written is unknown, so clear the entire buffer when returning
+                ArrayPool<byte>.Shared.Return(buffer, clearArray: true);
 
-                return new ByteArrayPoolMemoryOwner(buffer, length);
-            }
-            catch
-            {
-                ArrayPool<byte>.Shared.Return(buffer);
-                throw;
+                // Should be unreachable since we're allocating a buffer of the correct size.
+                ThrowHelper.ThrowInvalidOperationException();
             }
+
+            return new ByteArrayPoolMemoryOwner(buffer, length);
         }
 
         /// <summary>

Snappier/SnappyStream.cs

@@ -494,7 +494,7 @@ protected override void Dispose(bool disposing)
                             _buffer = null;
                             if (!AsyncOperationIsActive)
                             {
-                                ArrayPool<byte>.Shared.Return(buffer);
+                                ArrayPool<byte>.Shared.Return(buffer, clearArray: true);
                             }
                         }
 
@@ -545,7 +545,7 @@ public override async ValueTask DisposeAsync()
                             _buffer = null;
                             if (!AsyncOperationIsActive)
                             {
-                                ArrayPool<byte>.Shared.Return(buffer);
+                                ArrayPool<byte>.Shared.Return(buffer, clearArray: true);
                             }
                         }
                     }