Merge pull request #13976 from stephentoub/fix_memorystream_delegation

Change delegation in {Unmanaged}MemoryStream.Read/WriteAsync(Memory)

2 files changed, 52 insertions, 4 deletions

diff --git a/src/mscorlib/shared/System/IO/UnmanagedMemoryStream.cs b/src/mscorlib/shared/System/IO/UnmanagedMemoryStream.cs
index b899951ba7..76e6beb7a6 100644
--- a/src/mscorlib/shared/System/IO/UnmanagedMemoryStream.cs
+++ b/src/mscorlib/shared/System/IO/UnmanagedMemoryStream.cs
@@ -484,7 +484,22 @@ namespace System.IO
 
             try
             {
-                return new ValueTask<int>(Read(destination.Span));
+                // ReadAsync(Memory<byte>,...) needs to delegate to an existing virtual to do the work, in case an existing derived type
+                // has changed or augmented the logic associated with reads.  If the Memory wraps an array, we could delegate to
+                // ReadAsync(byte[], ...), but that would defeat part of the purpose, as ReadAsync(byte[], ...) often needs to allocate
+                // a Task<int> for the return value, so we want to delegate to one of the synchronous methods.  We could always
+                // delegate to the Read(Span<byte>) method, and that's the most efficient solution when dealing with a concrete
+                // UnmanagedMemoryStream, but if we're dealing with a type derived from UnmanagedMemoryStream, Read(Span<byte>) will end up delegating
+                // to Read(byte[], ...), which requires it to get a byte[] from ArrayPool and copy the data.  So, we special-case the
+                // very common case of the Memory<byte> wrapping an array: if it does, we delegate to Read(byte[], ...) with it,
+                // as that will be efficient in both cases, and we fall back to Read(Span<byte>) if the Memory<byte> wrapped something
+                // else; if this is a concrete UnmanagedMemoryStream, that'll be efficient, and only in the case where the Memory<byte> wrapped
+                // something other than an array and this is an UnmanagedMemoryStream-derived type that doesn't override Read(Span<byte>) will
+                // it then fall back to doing the ArrayPool/copy behavior.
+                return new ValueTask<int>(
+                    destination.TryGetArray(out ArraySegment<byte> destinationArray) ?
+                        Read(destinationArray.Array, destinationArray.Offset, destinationArray.Count) :
+                        Read(destination.Span));
             }
             catch (Exception ex)
             {
@@ -764,7 +779,16 @@ namespace System.IO
 
             try
             {
-                Write(source.Span);
+                // See corresponding comment in ReadAsync for why we don't just always use Write(ReadOnlySpan<byte>).
+                // Unlike ReadAsync, we could delegate to WriteAsync(byte[], ...) here, but we don't for consistency.
+                if (source.DangerousTryGetArray(out ArraySegment<byte> sourceArray))
+                {
+                    Write(sourceArray.Array, sourceArray.Offset, sourceArray.Count);
+                }
+                else
+                {
+                    Write(source.Span);
+                }
                 return Task.CompletedTask;
             }
             catch (Exception ex)
diff --git a/src/mscorlib/src/System/IO/MemoryStream.cs b/src/mscorlib/src/System/IO/MemoryStream.cs
index 2ac7c07db4..7d9ac901c0 100644
--- a/src/mscorlib/src/System/IO/MemoryStream.cs
+++ b/src/mscorlib/src/System/IO/MemoryStream.cs
@@ -464,7 +464,22 @@ namespace System.IO
 
             try
             {
-                return new ValueTask<int>(Read(destination.Span));
+                // ReadAsync(Memory<byte>,...) needs to delegate to an existing virtual to do the work, in case an existing derived type
+                // has changed or augmented the logic associated with reads.  If the Memory wraps an array, we could delegate to
+                // ReadAsync(byte[], ...), but that would defeat part of the purpose, as ReadAsync(byte[], ...) often needs to allocate
+                // a Task<int> for the return value, so we want to delegate to one of the synchronous methods.  We could always
+                // delegate to the Read(Span<byte>) method, and that's the most efficient solution when dealing with a concrete
+                // MemoryStream, but if we're dealing with a type derived from MemoryStream, Read(Span<byte>) will end up delegating
+                // to Read(byte[], ...), which requires it to get a byte[] from ArrayPool and copy the data.  So, we special-case the
+                // very common case of the Memory<byte> wrapping an array: if it does, we delegate to Read(byte[], ...) with it,
+                // as that will be efficient in both cases, and we fall back to Read(Span<byte>) if the Memory<byte> wrapped something
+                // else; if this is a concrete MemoryStream, that'll be efficient, and only in the case where the Memory<byte> wrapped
+                // something other than an array and this is a MemoryStream-derived type that doesn't override Read(Span<byte>) will
+                // it then fall back to doing the ArrayPool/copy behavior.
+                return new ValueTask<int>(
+                    destination.TryGetArray(out ArraySegment<byte> destinationArray) ?
+                        Read(destinationArray.Array, destinationArray.Offset, destinationArray.Count) :
+                        Read(destination.Span));
             }
             catch (OperationCanceledException oce)
             {
@@ -771,7 +786,16 @@ namespace System.IO
 
             try
             {
-                Write(source.Span);
+                // See corresponding comment in ReadAsync for why we don't just always use Write(ReadOnlySpan<byte>).
+                // Unlike ReadAsync, we could delegate to WriteAsync(byte[], ...) here, but we don't for consistency.
+                if (source.DangerousTryGetArray(out ArraySegment<byte> sourceArray))
+                {
+                    Write(sourceArray.Array, sourceArray.Offset, sourceArray.Count);
+                }
+                else
+                {
+                    Write(source.Span);
+                }
                 return Task.CompletedTask;
             }
             catch (OperationCanceledException oce)