Implement CopyToAsync in the FileBufferingReadStream

- overrride Span and Memory overloads and implement array overloads in terms of those overloads.
- Implemented CopyToAsync (but not CopyTo)
- Added tests

Fixes #24032

Author: davidfowl

src/Http/WebUtilities/src/FileBufferingReadStream.cs

@@ -208,39 +208,41 @@ private Stream CreateTempFile()
                 FileOptions.Asynchronous | FileOptions.DeleteOnClose | FileOptions.SequentialScan);
         }
 
-        public override int Read(byte[] buffer, int offset, int count)
+        public override int Read(Span<byte> buffer)
         {
             ThrowIfDisposed();
-            if (_buffer.Position < _buffer.Length || _completelyBuffered)
+
+            if (_completelyBuffered)
             {
                 // Just read from the buffer
-                return _buffer.Read(buffer, offset, (int)Math.Min(count, _buffer.Length - _buffer.Position));
+                return _buffer.Read(buffer);
             }
 
-            int read = _inner.Read(buffer, offset, count);
+            var read = _inner.Read(buffer);
 
             if (_bufferLimit.HasValue && _bufferLimit - read < _buffer.Length)
             {
-                Dispose();
                 throw new IOException("Buffer limit exceeded.");
             }
 
-            if (_inMemory && _buffer.Length + read > _memoryThreshold)
+            // We're about to go over the threshold, switch to a file
+            if (_inMemory && _memoryThreshold - read < _buffer.Length)
             {
                 _inMemory = false;
                 var oldBuffer = _buffer;
                 _buffer = CreateTempFile();
                 if (_rentedBuffer == null)
                 {
+                    // Copy data from the in memory buffer to the file stream using a pooled buffer
                     oldBuffer.Position = 0;
                     var rentedBuffer = _bytePool.Rent(Math.Min((int)oldBuffer.Length, _maxRentedBufferSize));
                     try
                     {
-                        var copyRead = oldBuffer.Read(rentedBuffer, 0, rentedBuffer.Length);
+                        var copyRead = oldBuffer.Read(rentedBuffer);
                         while (copyRead > 0)
                         {
-                            _buffer.Write(rentedBuffer, 0, copyRead);
-                            copyRead = oldBuffer.Read(rentedBuffer, 0, rentedBuffer.Length);
+                            _buffer.Write(rentedBuffer.AsSpan(0, copyRead));
+                            copyRead = oldBuffer.Read(rentedBuffer);
                         }
                     }
                     finally
@@ -250,15 +252,15 @@ public override int Read(byte[] buffer, int offset, int count)
                 }
                 else
                 {
-                    _buffer.Write(_rentedBuffer, 0, (int)oldBuffer.Length);
+                    _buffer.Write(_rentedBuffer.AsSpan(0, (int)oldBuffer.Length));
                     _bytePool.Return(_rentedBuffer);
                     _rentedBuffer = null;
                 }
             }
 
             if (read > 0)
             {
-                _buffer.Write(buffer, offset, read);
+                _buffer.Write(buffer.Slice(0, read));
             }
             else
             {
@@ -268,24 +270,34 @@ public override int Read(byte[] buffer, int offset, int count)
             return read;
         }
 
-        public override async Task<int> ReadAsync(byte[] buffer, int offset, int count, CancellationToken cancellationToken)
+        public override int Read(byte[] buffer, int offset, int count)
+        {
+            return Read(buffer.AsSpan(offset, count));
+        }
+
+        public override Task<int> ReadAsync(byte[] buffer, int offset, int count, CancellationToken cancellationToken)
+        {
+            return ReadAsync(buffer.AsMemory(offset, count), cancellationToken).AsTask();
+        }
+
+        public override async ValueTask<int> ReadAsync(Memory<byte> buffer, CancellationToken cancellationToken = default)
         {
             ThrowIfDisposed();
-            if (_buffer.Position < _buffer.Length || _completelyBuffered)
+
+            if (_completelyBuffered)
             {
                 // Just read from the buffer
-                return await _buffer.ReadAsync(buffer, offset, (int)Math.Min(count, _buffer.Length - _buffer.Position), cancellationToken);
+                return await _buffer.ReadAsync(buffer, cancellationToken);
             }
 
-            int read = await _inner.ReadAsync(buffer, offset, count, cancellationToken);
+            var read = await _inner.ReadAsync(buffer, cancellationToken);
 
             if (_bufferLimit.HasValue && _bufferLimit - read < _buffer.Length)
             {
-                Dispose();
                 throw new IOException("Buffer limit exceeded.");
             }
 
-            if (_inMemory && _buffer.Length + read > _memoryThreshold)
+            if (_inMemory && _memoryThreshold - read < _buffer.Length)
             {
                 _inMemory = false;
                 var oldBuffer = _buffer;
@@ -297,11 +309,11 @@ public override async Task<int> ReadAsync(byte[] buffer, int offset, int count,
                     try
                     {
                         // oldBuffer is a MemoryStream, no need to do async reads.
-                        var copyRead = oldBuffer.Read(rentedBuffer, 0, rentedBuffer.Length);
+                        var copyRead = oldBuffer.Read(rentedBuffer);
                         while (copyRead > 0)
                         {
-                            await _buffer.WriteAsync(rentedBuffer, 0, copyRead, cancellationToken);
-                            copyRead = oldBuffer.Read(rentedBuffer, 0, rentedBuffer.Length);
+                            await _buffer.WriteAsync(rentedBuffer.AsMemory(0, copyRead), cancellationToken);
+                            copyRead = oldBuffer.Read(rentedBuffer);
                         }
                     }
                     finally
@@ -311,15 +323,15 @@ public override async Task<int> ReadAsync(byte[] buffer, int offset, int count,
                 }
                 else
                 {
-                    await _buffer.WriteAsync(_rentedBuffer, 0, (int)oldBuffer.Length, cancellationToken);
+                    await _buffer.WriteAsync(_rentedBuffer.AsMemory(0, (int)oldBuffer.Length), cancellationToken);
                     _bytePool.Return(_rentedBuffer);
                     _rentedBuffer = null;
                 }
             }
 
             if (read > 0)
             {
-                await _buffer.WriteAsync(buffer, offset, read, cancellationToken);
+                await _buffer.WriteAsync(buffer.Slice(0, read), cancellationToken);
             }
             else
             {
@@ -349,6 +361,39 @@ public override void Flush()
             throw new NotSupportedException();
         }
 
+        public override Task CopyToAsync(Stream destination, int bufferSize, CancellationToken cancellationToken)
+        {
+            // If we're completed buffered then copy from the underlying source
+            if (_completelyBuffered)
+            {
+                return _buffer.CopyToAsync(destination, bufferSize, cancellationToken);
+            }
+
+            async Task CopyToAsyncImpl()
+            {
+                // At least a 4K buffer
+                byte[] buffer = _bytePool.Rent(Math.Min(bufferSize, 4096));
+                try
+                {
+                    while (true)
+                    {
+                        int bytesRead = await ReadAsync(buffer, cancellationToken);
+                        if (bytesRead == 0)
+                        {
+                            break;
+                        }
+                        await destination.WriteAsync(buffer.AsMemory(0, bytesRead), cancellationToken);
+                    }
+                }
+                finally
+                {
+                    _bytePool.Return(buffer);
+                }
+            }
+
+            return CopyToAsyncImpl();
+        }
+
         protected override void Dispose(bool disposing)
         {
             if (!_disposed)

src/Http/WebUtilities/test/FileBufferingReadStreamTests.cs

@@ -4,6 +4,7 @@
 using System;
 using System.Buffers;
 using System.IO;
+using System.Linq;
 using System.Text;
 using System.Threading.Tasks;
 using Moq;
@@ -157,7 +158,6 @@ public void FileBufferingReadStream_SyncReadWithOnDiskLimit_EnforcesLimit()
                 Assert.Equal("Buffer limit exceeded.", exception.Message);
                 Assert.False(stream.InMemory);
                 Assert.NotNull(stream.TempFileName);
-                Assert.False(File.Exists(tempFileName));
             }
 
             Assert.False(File.Exists(tempFileName));
@@ -287,7 +287,6 @@ public async Task FileBufferingReadStream_AsyncReadWithOnDiskLimit_EnforcesLimit
                 Assert.Equal("Buffer limit exceeded.", exception.Message);
                 Assert.False(stream.InMemory);
                 Assert.NotNull(stream.TempFileName);
-                Assert.False(File.Exists(tempFileName));
             }
 
             Assert.False(File.Exists(tempFileName));
@@ -351,6 +350,78 @@ public async Task FileBufferingReadStream_UsingMemoryStream_RentsAndReturnsRente
             Assert.False(File.Exists(tempFileName));
         }
 
+        [Fact]
+        public async Task CopyToAsyncWorks()
+        {
+            var data = Enumerable.Range(0, 1024).Select(b => (byte)b).Reverse().ToArray();
+            var inner = new MemoryStream(data);
+
+            using var stream = new FileBufferingReadStream(inner, 1024 * 1024, bufferLimit: null, GetCurrentDirectory());
+
+            var withoutBufferMs = new MemoryStream();
+            await stream.CopyToAsync(withoutBufferMs);
+
+            var withBufferMs = new MemoryStream();
+            stream.Position = 0;
+            await stream.CopyToAsync(withBufferMs);
+
+            Assert.Equal(data, withoutBufferMs.ToArray());
+            Assert.Equal(data, withBufferMs.ToArray());
+        }
+
+        [Fact]
+        public async Task CopyToAsyncWorksWithFileThreshold()
+        {
+            var data = Enumerable.Range(0, 1024).Select(b => (byte)b).Reverse().ToArray();
+            var inner = new MemoryStream(data);
+
+            using var stream = new FileBufferingReadStream(inner, 100, bufferLimit: null, GetCurrentDirectory());
+
+            var withoutBufferMs = new MemoryStream();
+            await stream.CopyToAsync(withoutBufferMs);
+
+            var withBufferMs = new MemoryStream();
+            stream.Position = 0;
+            await stream.CopyToAsync(withBufferMs);
+
+            Assert.Equal(data, withoutBufferMs.ToArray());
+            Assert.Equal(data, withBufferMs.ToArray());
+        }
+
+        [Fact]
+        public async Task ReadAsyncThenCopyToAsyncWorks()
+        {
+            var data = Enumerable.Range(0, 1024).Select(b => (byte)b).ToArray();
+            var inner = new MemoryStream(data);
+
+            using var stream = new FileBufferingReadStream(inner, 1024 * 1024, bufferLimit: null, GetCurrentDirectory());
+
+            var withoutBufferMs = new MemoryStream();
+            var buffer = new byte[100];
+            await stream.ReadAsync(buffer);
+            await stream.CopyToAsync(withoutBufferMs);
+
+            Assert.Equal(data.AsMemory(0, 100).ToArray(), buffer);
+            Assert.Equal(data.AsMemory(100).ToArray(), withoutBufferMs.ToArray());
+        }
+
+        [Fact]
+        public async Task ReadThenCopyToAsyncWorks()
+        {
+            var data = Enumerable.Range(0, 1024).Select(b => (byte)b).ToArray();
+            var inner = new MemoryStream(data);
+
+            using var stream = new FileBufferingReadStream(inner, 1024 * 1024, bufferLimit: null, GetCurrentDirectory());
+
+            var withoutBufferMs = new MemoryStream();
+            var buffer = new byte[100];
+            stream.Read(buffer);
+            await stream.CopyToAsync(withoutBufferMs);
+
+            Assert.Equal(data.AsMemory(0, 100).ToArray(), buffer);
+            Assert.Equal(data.AsMemory(100).ToArray(), withoutBufferMs.ToArray());
+        }
+
         private static string GetCurrentDirectory()
         {
             return AppContext.BaseDirectory;