TryGet for ReadOnlysequence (#27229)

Author: pakrym

src/System.IO.Pipelines/ref/System.IO.Pipelines.cs

@@ -82,7 +82,6 @@ protected PipeWriter() { }
         public abstract System.IO.Pipelines.PipeAwaiter<System.IO.Pipelines.FlushResult> FlushAsync(System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken));
         public abstract System.Memory<byte> GetMemory(int minimumLength = 0);
         public abstract System.Span<byte> GetSpan(int minimumLength = 0);
-        public abstract int MaxBufferSize { get; }
         public abstract void OnReaderCompleted(System.Action<System.Exception, object> callback, object state);
         public virtual System.IO.Pipelines.PipeAwaiter<System.IO.Pipelines.FlushResult> WriteAsync(System.ReadOnlyMemory<byte> source, System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) { throw null; }
     }

src/System.IO.Pipelines/src/System/IO/Pipelines/Pipe.DefaultPipeWriter.cs

@@ -30,11 +30,9 @@ public DefaultPipeWriter(Pipe pipe)
 
             public override void Advance(int bytes) => _pipe.Advance(bytes);
 
-            public override Memory<byte> GetMemory(int minimumLength = 0) => _pipe.GetMemory(minimumLength);
+            public override Memory<byte> GetMemory(int sizeHint = 0) => _pipe.GetMemory(sizeHint);
 
-            public override Span<byte> GetSpan(int minimumLength = 0) => _pipe.GetMemory(minimumLength).Span;
-
-            public override int MaxBufferSize => _pipe._pool.MaxBufferSize;
+            public override Span<byte> GetSpan(int sizeHint = 0) => _pipe.GetMemory(sizeHint).Span;
 
             public bool IsCompleted => _pipe.IsFlushAsyncCompleted;
 

src/System.IO.Pipelines/src/System/IO/Pipelines/Pipe.cs

@@ -111,30 +111,29 @@ private void ResetState()
             _length = 0;
         }
 
-        internal Memory<byte> GetMemory(int minimumSize)
+        internal Memory<byte> GetMemory(int sizeHint)
         {
             if (_writerCompletion.IsCompleted)
             {
                 ThrowHelper.ThrowInvalidOperationException_NoWritingAllowed();
             }
 
-            if (minimumSize < 0)
+            if (sizeHint < 0)
             {
                 ThrowHelper.ThrowArgumentOutOfRangeException(ExceptionArgument.minimumSize);
             }
 
             lock (_sync)
             {
-                BufferSegment segment = _writingHead ?? AllocateWriteHeadUnsynchronized(minimumSize);
+                BufferSegment segment = _writingHead ?? AllocateWriteHeadUnsynchronized(sizeHint);
 
                 int bytesLeftInBuffer = segment.WritableBytes;
 
                 // If inadequate bytes left or if the segment is readonly
-                if (bytesLeftInBuffer == 0 || bytesLeftInBuffer < minimumSize || segment.ReadOnly)
+                if (bytesLeftInBuffer == 0 || bytesLeftInBuffer < sizeHint || segment.ReadOnly)
                 {
                     BufferSegment nextSegment = CreateSegmentUnsynchronized();
-
-                    nextSegment.SetMemory(_pool.Rent(Math.Max(_minimumSegmentSize, minimumSize)));
+                    nextSegment.SetMemory(_pool.Rent(GetSegmentSize(sizeHint)));
 
                     segment.SetNext(nextSegment);
 
@@ -145,7 +144,7 @@ internal Memory<byte> GetMemory(int minimumSize)
             return _writingHead.AvailableMemory.Slice(_writingHead.End, _writingHead.WritableBytes);
         }
 
-        private BufferSegment AllocateWriteHeadUnsynchronized(int count)
+        private BufferSegment AllocateWriteHeadUnsynchronized(int sizeHint)
         {
             BufferSegment segment = null;
 
@@ -154,7 +153,7 @@ private BufferSegment AllocateWriteHeadUnsynchronized(int count)
                 // Try to return the tail so the calling code can append to it
                 int remaining = _commitHead.WritableBytes;
 
-                if (count <= remaining && remaining > 0)
+                if (sizeHint <= remaining && remaining > 0)
                 {
                     // Free tail space of the right amount, use that
                     segment = _commitHead;
@@ -165,7 +164,7 @@ private BufferSegment AllocateWriteHeadUnsynchronized(int count)
             {
                 // No free tail space, allocate a new segment
                 segment = CreateSegmentUnsynchronized();
-                segment.SetMemory(_pool.Rent(Math.Max(_minimumSegmentSize, count)));
+                segment.SetMemory(_pool.Rent(GetSegmentSize(sizeHint)));
             }
 
             if (_commitHead == null)
@@ -186,6 +185,15 @@ private BufferSegment AllocateWriteHeadUnsynchronized(int count)
             return segment;
         }
 
+        private int GetSegmentSize(int sizeHint)
+        {
+            // First we need to handle case where hint is smaller than minimum segment size
+            var adjustedToMinimumSize = Math.Max(_minimumSegmentSize, sizeHint);
+            // After that adjust it to fit into pools max buffer size
+            var adjustedToMaximumSize = Math.Min(_pool.MaxBufferSize, adjustedToMinimumSize);
+            return adjustedToMaximumSize;
+        }
+
         private BufferSegment CreateSegmentUnsynchronized()
         {
             if (_pooledSegmentCount > 0)
@@ -338,26 +346,26 @@ internal void AdvanceReader(SequencePosition consumed, SequencePosition examined
             lock (_sync)
             {
                 var examinedEverything = false;
-                if (examined.Segment == _commitHead)
+                if (examined.GetObject() == _commitHead)
                 {
-                    examinedEverything = _commitHead != null ? examined.Index == _commitHeadIndex - _commitHead.Start : examined.Index == 0;
+                    examinedEverything = _commitHead != null ? examined.GetInteger() == _commitHeadIndex - _commitHead.Start : examined.GetInteger() == 0;
                 }
 
-                if (consumed.Segment != null)
+                if (consumed.GetObject() != null)
                 {
                     if (_readHead == null)
                     {
                         ThrowHelper.ThrowInvalidOperationException_AdvanceToInvalidCursor();
                         return;
                     }
 
-                    var consumedSegment = (BufferSegment)consumed.Segment;
+                    var consumedSegment = (BufferSegment)consumed.GetObject();
 
                     returnStart = _readHead;
                     returnEnd = consumedSegment;
 
                     // Check if we crossed _maximumSizeLow and complete backpressure
-                    long consumedBytes = new ReadOnlySequence<byte>(returnStart, _readHeadIndex, consumedSegment, consumed.Index).Length;
+                    long consumedBytes = new ReadOnlySequence<byte>(returnStart, _readHeadIndex, consumedSegment, consumed.GetInteger()).Length;
                     long oldLength = _length;
                     _length -= consumedBytes;
 
@@ -370,7 +378,7 @@ internal void AdvanceReader(SequencePosition consumed, SequencePosition examined
                     // Check if we consumed entire last segment
                     // if we are going to return commit head we need to check that there is no writing operation that
                     // might be using tailspace
-                    if (consumed.Index == returnEnd.Length && _writingHead != returnEnd)
+                    if (consumed.GetInteger() == returnEnd.Length && _writingHead != returnEnd)
                     {
                         BufferSegment nextBlock = returnEnd.NextSegment;
                         if (_commitHead == returnEnd)
@@ -386,7 +394,7 @@ internal void AdvanceReader(SequencePosition consumed, SequencePosition examined
                     else
                     {
                         _readHead = consumedSegment;
-                        _readHeadIndex = consumed.Index;
+                        _readHeadIndex = consumed.GetInteger();
                     }
                 }
 

src/System.IO.Pipelines/src/System/IO/Pipelines/PipeWriter.cs

@@ -37,13 +37,10 @@ public abstract class PipeWriter : IBufferWriter<byte>
         public abstract void Advance(int bytes);
 
         /// <inheritdoc />
-        public abstract Memory<byte> GetMemory(int minimumLength = 0);
+        public abstract Memory<byte> GetMemory(int sizeHint = 0);
 
         /// <inheritdoc />
-        public abstract Span<byte> GetSpan(int minimumLength = 0);
-
-        /// <inheritdoc />
-        public abstract int MaxBufferSize { get; }
+        public abstract Span<byte> GetSpan(int sizeHint = 0);
 
         /// <summary>
         /// Writes <paramref name="source"/> to the pipe and makes data accessible to <see cref="PipeReader"/>

src/System.IO.Pipelines/tests/PipeReaderWriterFacts.cs

@@ -233,7 +233,7 @@ public async Task HelloWorldAcrossTwoBlocks()
             //     block 1       ->    block2
             // [padding..hello]  ->  [  world   ]
             PipeWriter writeBuffer = _pipe.Writer;
-            var blockSize = _pipe.Writer.GetMemory(0).Length;
+            var blockSize = _pipe.Writer.GetMemory().Length;
 
             byte[] paddingBytes = Enumerable.Repeat((byte)'a', blockSize - 5).ToArray();
             byte[] bytes = Encoding.ASCII.GetBytes("Hello World");

src/System.IO.Pipelines/tests/PipeWriterTests.cs

@@ -197,5 +197,13 @@ public void ThrowsOnAdvanceWithNoMemory()
             var exception = Assert.Throws<InvalidOperationException>(() => buffer.Advance(1));
             Assert.Equal("No writing operation. Make sure GetMemory() was called.", exception.Message);
         }
+
+        [Fact]
+        public void GetMemory_AdjustsToPoolMaxBufferSize()
+        {
+            PipeWriter buffer = Pipe.Writer;
+            var memory = buffer.GetMemory(int.MaxValue);
+            Assert.True(memory.Length >= 4096);
+        }
     }
 }

src/System.IO.Pipelines/tests/TestMemoryPool.cs

@@ -25,7 +25,7 @@ protected override void Dispose(bool disposing)
             _disposed = true;
         }
 
-        public override int MaxBufferSize => _pool.MaxBufferSize;
+        public override int MaxBufferSize => 4096;
 
         internal void CheckDisposed()
         {

src/System.Memory/ref/System.Memory.cs

@@ -165,8 +165,10 @@ public ref partial struct Enumerator
     {
         private readonly object _dummy;
         public SequencePosition(object segment, int index) { throw null; }
-        public int Index { get { throw null; } }
-        public object Segment { get { throw null; } }
+        [System.ComponentModel.EditorBrowsableAttribute((System.ComponentModel.EditorBrowsableState)(1))]
+        public int GetInteger() { throw null; }
+        [System.ComponentModel.EditorBrowsableAttribute((System.ComponentModel.EditorBrowsableState)(1))]
+        public object GetObject() { throw null; }
         [System.ComponentModel.EditorBrowsableAttribute((System.ComponentModel.EditorBrowsableState)(1))]
         public override bool Equals(object obj) { throw null; }
         public bool Equals(System.SequencePosition position) { throw null; }
@@ -219,9 +221,8 @@ namespace System.Buffers
     public partial interface IBufferWriter<T>
     {
         void Advance(int count);
-        System.Memory<T> GetMemory(int minimumLength = 0);
-        System.Span<T> GetSpan(int minimumLength = 0);
-        int MaxBufferSize { get; }
+        System.Memory<T> GetMemory(int sizeHint = 0);
+        System.Span<T> GetSpan(int sizeHint = 0);
     }
     public partial interface IMemoryList<T>
     {
@@ -512,4 +513,12 @@ public static bool TryGetOwnedMemory<T, TOwner>(ReadOnlyMemory<T> readOnlyMemory
         public static bool TryGetOwnedMemory<T, TOwner>(ReadOnlyMemory<T> readOnlyMemory, out TOwner ownedMemory, out int index, out int length)
             where TOwner : System.Buffers.OwnedMemory<T> { throw null; }
     }
+
+    public static partial class SequenceMarshal
+    {
+        public static bool TryGetArray<T>(System.Buffers.ReadOnlySequence<T> sequence, out System.ArraySegment<T> arraySegment) { throw null; }
+        public static bool TryGetMemoryList<T>(System.Buffers.ReadOnlySequence<T> sequence, out System.Buffers.IMemoryList<T> startSegment, out int startIndex, out System.Buffers.IMemoryList<T> endSegment, out int endIndex) { throw null; }
+        public static bool TryGetOwnedMemory<T>(System.Buffers.ReadOnlySequence<T> sequence, out System.Buffers.OwnedMemory<T> ownedMemory, out int start, out int length) { throw null; }
+        public static bool TryGetReadOnlyMemory<T>(System.Buffers.ReadOnlySequence<T> sequence, out System.ReadOnlyMemory<T> readOnlyMemory) { throw null; }
+    }
 }

src/System.Memory/src/System.Memory.csproj

@@ -45,6 +45,7 @@
     <Compile Include="System\Buffers\Binary\WriterLittleEndian.cs" />
     <Compile Include="System\Buffers\Text\Base64Decoder.cs" />
     <Compile Include="System\Buffers\Text\Base64Encoder.cs" />
+    <Compile Include="System\Runtime\InteropServices\SequenceMarshal.cs" />
   </ItemGroup>
   <!-- Utf8 Formatter/Parser -->
   <ItemGroup>
@@ -159,4 +160,4 @@
     <ProjectReference Include="..\..\System.Numerics.Vectors\src\System.Numerics.Vectors.csproj" />
   </ItemGroup>
   <Import Project="$([MSBuild]::GetDirectoryNameOfFileAbove($(MSBuildThisFileDirectory), dir.targets))\dir.targets" />
-</Project>
+</Project>

src/System.Memory/src/System/Buffers/BuffersExtensions.cs

@@ -70,20 +70,24 @@ public static void Write<T>(this IBufferWriter<T> bufferWriter, ReadOnlySpan<T>
                 return;
             }
 
-            while (source.Length > 0)
-            {
-                int writeSize = destination.Length;
-
-                if (destination.Length == 0)
-                {
-                    writeSize = Math.Min(source.Length, bufferWriter.MaxBufferSize);
-                    destination = bufferWriter.GetSpan(writeSize);
-                }
+            WriteMultiSegment(bufferWriter, source, destination);
+        }
 
+        private static void WriteMultiSegment<T>(IBufferWriter<T> bufferWriter, ReadOnlySpan<T> source, Span<T> destination)
+        {
+            while (true)
+            {
+                int writeSize = Math.Min(destination.Length, source.Length);
                 source.Slice(0, writeSize).CopyTo(destination);
                 bufferWriter.Advance(writeSize);
                 source = source.Slice(writeSize);
-                destination = default;
+                if (source.Length > 0)
+                {
+                    destination = bufferWriter.GetSpan(source.Length);
+                    continue;
+                }
+
+                return;
             }
         }
     }