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

Read/WriteAsync(Memory) on MemoryStream and UnmanagedMemoryStream need to delegate to one of the existing virtual methods, in case an existing stream has overridden the virtuals in order to change or augment the behavior (e.g. checking on each write to ensure the length doesn't exceed some amount).  Currently these delegate to the synchronous Read/Write(Span) methods.  The problem with that is, for exactly the case where there is a derived class, Read/Write(Span) themselves need to delegate to Read/Write(byte[]), which means they use ArrayPool and copy.  But with a {ReadOnly}Memory, we may already have access to the underlying array, in which case we're going from an array to a span and back to different rented array along with an unnecessary copy.  To address that, this commit changes the delegation to prefer Read/Write(byte[],...) if possible, falling back to Read/Write(Span) only if we couldn't get an array from the 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)