Add BrotliStream.DisposeAsync override

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(); }