Merge pull request dotnet#33415 from stephentoub/compressionstreams

Add {Deflate\GZip\Brotli}Stream.DisposeAsync overrides

Author: stephentoub

src/System.IO.Compression.Brotli/ref/System.IO.Compression.Brotli.cs

@@ -36,6 +36,7 @@ public BrotliStream(System.IO.Stream stream, System.IO.Compression.CompressionMo
         public override long Length { get { throw null; } }
         public override long Position { get { throw null; } set { } }
         protected override void Dispose(bool disposing) { }
+        public override System.Threading.Tasks.ValueTask DisposeAsync() { throw null; }
         public override void Flush() { }
         public override IAsyncResult BeginRead(byte[] buffer, int offset, int count, AsyncCallback asyncCallback, object asyncState) { throw null; }
         public override int EndRead(IAsyncResult asyncResult) { throw null; }

src/System.IO.Compression.Brotli/src/System/IO/Compression/BrotliStream.cs

@@ -5,6 +5,7 @@
 using System.Diagnostics;
 using System.Runtime.CompilerServices;
 using System.Threading;
+using System.Threading.Tasks;
 
 namespace System.IO.Compression
 {
@@ -74,6 +75,31 @@ protected override void Dispose(bool disposing)
             }
         }
 
+        public override async ValueTask DisposeAsync()
+        {
+            try
+            {
+                if (_stream != null)
+                {
+                    if (_mode == CompressionMode.Compress)
+                    {
+                        await WriteAsyncMemoryCore(ReadOnlyMemory<byte>.Empty, CancellationToken.None, isFinalBlock: true).ConfigureAwait(false);
+                    }
+
+                    if (!_leaveOpen)
+                    {
+                        await _stream.DisposeAsync().ConfigureAwait(false);
+                    }
+                }
+            }
+            finally
+            {
+                _stream = null;
+                _encoder.Dispose();
+                _decoder.Dispose();
+            }
+        }
+
         private static void ValidateParameters(byte[] array, int offset, int count)
         {
             if (array == null)

src/System.IO.Compression.Brotli/src/System/IO/Compression/enc/BrotliStream.Compress.cs

@@ -81,7 +81,7 @@ public override Task WriteAsync(byte[] buffer, int offset, int count, Cancellati
                 WriteAsyncMemoryCore(buffer, cancellationToken));
         }
 
-        private async Task WriteAsyncMemoryCore(ReadOnlyMemory<byte> buffer, CancellationToken cancellationToken)
+        private async Task WriteAsyncMemoryCore(ReadOnlyMemory<byte> buffer, CancellationToken cancellationToken, bool isFinalBlock = false)
         {
             AsyncOperationStarting();
             try
@@ -92,7 +92,7 @@ private async Task WriteAsyncMemoryCore(ReadOnlyMemory<byte> buffer, Cancellatio
                     Memory<byte> output = new Memory<byte>(_buffer);
                     int bytesConsumed = 0;
                     int bytesWritten = 0;
-                    lastResult = _encoder.Compress(buffer, output, out bytesConsumed, out bytesWritten, isFinalBlock: false);
+                    lastResult = _encoder.Compress(buffer, output, out bytesConsumed, out bytesWritten, isFinalBlock);
                     if (lastResult == OperationStatus.InvalidData)
                         throw new InvalidOperationException(SR.BrotliStream_Compress_InvalidData);
                     if (bytesConsumed > 0)

src/System.IO.Compression.Brotli/tests/CompressionStreamUnitTests.Brotli.cs

@@ -33,6 +33,27 @@ public void Precancellation()
             }
         }
 
+        [Theory]
+        [InlineData(false)]
+        [InlineData(true)]
+        public async Task DisposeAsync_Flushes(bool leaveOpen)
+        {
+            var ms = new MemoryStream();
+            var bs = new BrotliStream(ms, CompressionMode.Compress, leaveOpen);
+            bs.WriteByte(1);
+            Assert.Equal(0, ms.Position);
+            await bs.DisposeAsync();
+            Assert.InRange(ms.ToArray().Length, 1, int.MaxValue);
+            if (leaveOpen)
+            {
+                Assert.InRange(ms.Position, 1, int.MaxValue);
+            }
+            else
+            {
+                Assert.Throws<ObjectDisposedException>(() => ms.Position);
+            }
+        }
+
         [Fact]
         [OuterLoop("Test takes ~6 seconds to run")]
         public override void FlushAsync_DuringWriteAsync() { base.FlushAsync_DuringWriteAsync(); }

src/System.IO.Compression/ref/System.IO.Compression.cs

@@ -32,6 +32,7 @@ public DeflateStream(System.IO.Stream stream, System.IO.Compression.CompressionM
         public override long Position { get { throw null; } set { } }
         public override System.Threading.Tasks.Task CopyToAsync(System.IO.Stream destination, int bufferSize, System.Threading.CancellationToken cancellationToken) { throw null; }
         protected override void Dispose(bool disposing) { }
+        public override System.Threading.Tasks.ValueTask DisposeAsync() { throw null; }
         public override void Flush() { }
         public override IAsyncResult BeginRead(byte[] buffer, int offset, int count, AsyncCallback asyncCallback, object asyncState) { throw null; }
         public override int EndRead(IAsyncResult asyncResult) { throw null; }
@@ -66,6 +67,7 @@ public GZipStream(System.IO.Stream stream, System.IO.Compression.CompressionMode
         public override void CopyTo(System.IO.Stream destination, int bufferSize) { }
         public override System.Threading.Tasks.Task CopyToAsync(System.IO.Stream destination, int bufferSize, System.Threading.CancellationToken cancellationToken) { throw null; }
         protected override void Dispose(bool disposing) { }
+        public override System.Threading.Tasks.ValueTask DisposeAsync() { throw null; }
         public override void Flush() { }
         public override IAsyncResult BeginRead(byte[] array, int offset, int count, AsyncCallback asyncCallback, object asyncState) { throw null; }
         public override int EndRead(IAsyncResult asyncResult) { throw null; }

src/System.IO.Compression/src/System/IO/Compression/DeflateZLib/DeflateStream.cs

@@ -4,7 +4,6 @@
 
 using System.Buffers;
 using System.Diagnostics;
-using System.Runtime.CompilerServices;
 using System.Runtime.InteropServices;
 using System.Threading;
 using System.Threading.Tasks;
@@ -628,6 +627,53 @@ private void PurgeBuffers(bool disposing)
             }
         }
 
+        private async Task PurgeBuffersAsync()
+        {
+            // Same logic as PurgeBuffers, except with async counterparts.
+
+            if (_stream == null)
+                return;
+
+            if (_mode != CompressionMode.Compress)
+                return;
+
+            // Some deflaters (e.g. ZLib) write more than zero bytes for zero byte inputs.
+            // This round-trips and we should be ok with this, but our legacy managed deflater
+            // always wrote zero output for zero input and upstack code (e.g. ZipArchiveEntry)
+            // took dependencies on it. Thus, make sure to only "flush" when we actually had
+            // some input:
+            if (_wroteBytes)
+            {
+                // Compress any bytes left
+                await WriteDeflaterOutputAsync(default).ConfigureAwait(false);
+
+                // Pull out any bytes left inside deflater:
+                bool finished;
+                do
+                {
+                    int compressedBytes;
+                    finished = _deflater.Finish(_buffer, out compressedBytes);
+
+                    if (compressedBytes > 0)
+                        await _stream.WriteAsync(new ReadOnlyMemory<byte>(_buffer, 0, compressedBytes)).ConfigureAwait(false);
+                } while (!finished);
+            }
+            else
+            {
+                // In case of zero length buffer, we still need to clean up the native created stream before
+                // the object get disposed because eventually ZLibNative.ReleaseHandle will get called during
+                // the dispose operation and although it frees the stream but it return error code because the
+                // stream state was still marked as in use. The symptoms of this problem will not be seen except
+                // if running any diagnostic tools which check for disposing safe handle objects
+                bool finished;
+                do
+                {
+                    int compressedBytes;
+                    finished = _deflater.Finish(_buffer, out compressedBytes);
+                } while (!finished);
+            }
+        }
+
         protected override void Dispose(bool disposing)
         {
             try
@@ -674,6 +720,58 @@ protected override void Dispose(bool disposing)
             }
         }
 
+        public override ValueTask DisposeAsync()
+        {
+            return GetType() == typeof(DeflateStream) ?
+                DisposeAsyncCore() :
+                base.DisposeAsync();
+        }
+
+        private async ValueTask DisposeAsyncCore()
+        {
+            // Same logic as Dispose(true), except with async counterparts.
+            try
+            {
+                await PurgeBuffersAsync().ConfigureAwait(false);
+            }
+            finally
+            {
+                // Close the underlying stream even if PurgeBuffers threw.
+                // Stream.Close() may throw here (may or may not be due to the same error).
+                // In this case, we still need to clean up internal resources, hence the inner finally blocks.
+                Stream stream = _stream;
+                _stream = null;
+                try
+                {
+                    if (!_leaveOpen && stream != null)
+                        await stream.DisposeAsync().ConfigureAwait(false);
+                }
+                finally
+                {
+                    try
+                    {
+                        _deflater?.Dispose();
+                        _inflater?.Dispose();
+                    }
+                    finally
+                    {
+                        _deflater = null;
+                        _inflater = null;
+
+                        byte[] buffer = _buffer;
+                        if (buffer != null)
+                        {
+                            _buffer = null;
+                            if (!AsyncOperationIsActive)
+                            {
+                                ArrayPool<byte>.Shared.Return(buffer);
+                            }
+                        }
+                    }
+                }
+            }
+        }
+
         public override IAsyncResult BeginWrite(byte[] array, int offset, int count, AsyncCallback asyncCallback, object asyncState) =>
             TaskToApm.Begin(WriteAsync(array, offset, count, CancellationToken.None), asyncCallback, asyncState);
 

src/System.IO.Compression/src/System/IO/Compression/GZipStream.cs

@@ -150,6 +150,18 @@ protected override void Dispose(bool disposing)
             }
         }
 
+        public override ValueTask DisposeAsync()
+        {
+            if (GetType() != typeof(GZipStream))
+            {
+                return base.DisposeAsync();
+            }
+
+            DeflateStream ds = _deflateStream;
+            _deflateStream = null;
+            return ds.DisposeAsync();
+        }
+
         public Stream BaseStream => _deflateStream?.BaseStream;
 
         public override Task<int> ReadAsync(byte[] array, int offset, int count, CancellationToken cancellationToken)

src/System.IO.Compression/tests/CompressionStreamUnitTests.Deflate.cs

@@ -77,6 +77,31 @@ public void DerivedStream_ReadWriteSpan_UsesReadWriteArray()
             }
         }
 
+        [Theory]
+        [InlineData(false, false)]
+        [InlineData(false, true)]
+        [InlineData(true, false)]
+        [InlineData(true, false)]
+        public async Task DisposeAsync_Flushes(bool derived, bool leaveOpen)
+        {
+            var ms = new MemoryStream();
+            var ds = derived ?
+                new DerivedDeflateStream(ms, CompressionMode.Compress, leaveOpen) :
+                new DeflateStream(ms, CompressionMode.Compress, leaveOpen);
+            ds.WriteByte(1);
+            Assert.Equal(0, ms.Position);
+            await ds.DisposeAsync();
+            Assert.InRange(ms.ToArray().Length, 1, int.MaxValue);
+            if (leaveOpen)
+            {
+                Assert.InRange(ms.Position, 1, int.MaxValue);
+            }
+            else
+            {
+                Assert.Throws<ObjectDisposedException>(() => ms.Position);
+            }
+        }
+
         private sealed class DerivedDeflateStream : DeflateStream
         {
             public bool ReadArrayInvoked = false, WriteArrayInvoked = false;

src/System.IO.Compression/tests/CompressionStreamUnitTests.Gzip.cs

@@ -193,6 +193,36 @@ public void DerivedStream_ReadWriteSpan_UsesReadWriteArray()
             }
         }
 
+        [Theory]
+        [InlineData(false, false)]
+        [InlineData(false, true)]
+        [InlineData(true, false)]
+        [InlineData(true, false)]
+        public async Task DisposeAsync_Flushes(bool derived, bool leaveOpen)
+        {
+            var ms = new MemoryStream();
+            var gs = derived ?
+                new DerivedGZipStream(ms, CompressionMode.Compress, leaveOpen) :
+                new GZipStream(ms, CompressionMode.Compress, leaveOpen);
+            gs.WriteByte(1);
+            await gs.FlushAsync();
+
+            long pos = ms.Position;
+            gs.WriteByte(1);
+            Assert.Equal(pos, ms.Position);
+
+            await gs.DisposeAsync();
+            Assert.InRange(ms.ToArray().Length, pos + 1, int.MaxValue);
+            if (leaveOpen)
+            {
+                Assert.InRange(ms.Position, pos + 1, int.MaxValue);
+            }
+            else
+            {
+                Assert.Throws<ObjectDisposedException>(() => ms.Position);
+            }
+        }
+
         private sealed class DerivedGZipStream : GZipStream
         {
             public bool ReadArrayInvoked = false, WriteArrayInvoked = false;