diff --git a/Snappier/Internal/ByteArrayPoolMemoryOwner.cs b/Snappier/Internal/ByteArrayPoolMemoryOwner.cs index 2cf5904..3490fe2 100644 --- a/Snappier/Internal/ByteArrayPoolMemoryOwner.cs +++ b/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.Shared.Return(innerArray); diff --git a/Snappier/Internal/Helpers.cs b/Snappier/Internal/Helpers.cs index 9163cc7..fe50228 100644 --- a/Snappier/Internal/Helpers.cs +++ b/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); + + /// + /// 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. + /// + [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.Shared.Return(array); + } + [MethodImpl(MethodImplOptions.AggressiveInlining)] public static bool LeftShiftOverflows(byte value, int shift) { diff --git a/Snappier/Internal/SnappyCompressor.cs b/Snappier/Internal/SnappyCompressor.cs index 4086a2a..8bebfc2 100644 --- a/Snappier/Internal/SnappyCompressor.cs +++ b/Snappier/Internal/SnappyCompressor.cs @@ -60,21 +60,18 @@ public bool TryCompress(ReadOnlySpan input, Span output, out int byt // compress to a temporary buffer and copy the compressed data to the output span. byte[] scratch = ArrayPool.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.Shared.Return(scratch); + Helpers.ClearAndReturn(scratch, written); + bytesWritten = 0; + return false; } + + Span writtenScratch = scratch.AsSpan(0, written); + writtenScratch.CopyTo(output); + writtenScratch.Clear(); + ArrayPool.Shared.Return(scratch); } output = output.Slice(written); @@ -132,19 +129,17 @@ public void Compress(ReadOnlySequence input, IBufferWriter bufferWri int fragmentLength = (int)fragment.Length; byte[] scratch = ArrayPool.Shared.Rent(fragmentLength); - try - { - fragment.CopyTo(scratch); - CompressFragment(scratch.AsSpan(0, fragmentLength), bufferWriter); + Span usedScratch = scratch.AsSpan(0, fragmentLength); + fragment.CopyTo(usedScratch); - // Advance the length of the entire fragment - input = input.Slice(position); - } - finally - { - ArrayPool.Shared.Return(scratch); - } + CompressFragment(usedScratch, bufferWriter); + + usedScratch.Clear(); + ArrayPool.Shared.Return(scratch); + + // Advance the length of the entire fragment + input = input.Slice(position); } } } diff --git a/Snappier/Internal/SnappyDecompressor.cs b/Snappier/Internal/SnappyDecompressor.cs index 9c40cb4..0083b6e 100644 --- a/Snappier/Internal/SnappyDecompressor.cs +++ b/Snappier/Internal/SnappyDecompressor.cs @@ -526,7 +526,8 @@ private int? ExpectedLength { if (_lookbackBufferArray is not null) { - ArrayPool.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.Shared.Return(_lookbackBufferArray); + // Clear the used portion of the lookback buffer before returning + Helpers.ClearAndReturn(_lookbackBufferArray, _lookbackPosition); + _lookbackBufferArray = null; _lookbackBuffer = default; } diff --git a/Snappier/Internal/SnappyStreamCompressor.cs b/Snappier/Internal/SnappyStreamCompressor.cs index 2e46fce..275cd7a 100644 --- a/Snappier/Internal/SnappyStreamCompressor.cs +++ b/Snappier/Internal/SnappyStreamCompressor.cs @@ -276,7 +276,7 @@ private void EnsureBuffer() { // Allocate enough room for the stream header and block headers _outputBuffer ??= - ArrayPool.Shared.Rent(Helpers.MaxCompressedLength((int) Constants.BlockSize) + 8 + SnappyHeader.Length); + ArrayPool.Shared.Rent(Helpers.MaxBlockCompressedLength + 8 + SnappyHeader.Length); // Allocate enough room for the stream header and block headers _inputBuffer ??= ArrayPool.Shared.Rent((int) Constants.BlockSize); @@ -289,12 +289,12 @@ public void Dispose() if (_outputBuffer is not null) { - ArrayPool.Shared.Return(_outputBuffer); + ArrayPool.Shared.Return(_outputBuffer, clearArray: true); _outputBuffer = null; } if (_inputBuffer is not null) { - ArrayPool.Shared.Return(_inputBuffer); + ArrayPool.Shared.Return(_inputBuffer, clearArray: true); _inputBuffer = null; } } diff --git a/Snappier/Snappy.cs b/Snappier/Snappy.cs index 378a461..375ea57 100644 --- a/Snappier/Snappy.cs +++ b/Snappier/Snappy.cs @@ -102,21 +102,16 @@ public static IMemoryOwner CompressToMemory(ReadOnlySpan input) { byte[] buffer = ArrayPool.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.Shared.Return(buffer, clearArray: true); - return new ByteArrayPoolMemoryOwner(buffer, length); - } - catch - { - ArrayPool.Shared.Return(buffer); - throw; + // Should be unreachable since we're allocating a buffer of the correct size. + ThrowHelper.ThrowInvalidOperationException(); } + + return new ByteArrayPoolMemoryOwner(buffer, length); } /// diff --git a/Snappier/SnappyStream.cs b/Snappier/SnappyStream.cs index 8e8a8f8..85711e3 100644 --- a/Snappier/SnappyStream.cs +++ b/Snappier/SnappyStream.cs @@ -494,7 +494,7 @@ protected override void Dispose(bool disposing) _buffer = null; if (!AsyncOperationIsActive) { - ArrayPool.Shared.Return(buffer); + ArrayPool.Shared.Return(buffer, clearArray: true); } } @@ -545,7 +545,7 @@ public override async ValueTask DisposeAsync() _buffer = null; if (!AsyncOperationIsActive) { - ArrayPool.Shared.Return(buffer); + ArrayPool.Shared.Return(buffer, clearArray: true); } } }