Fix IBufferWriter: GetSpan should not go through GetMemory (#2259)

dotnet · May 4, 2018 · 35e606e · 35e606e
1 parent d9b413c
commit 35e606e
Show file tree

Hide file tree

Showing 7 changed files with 74 additions and 7 deletions.
diff --git a/samples/LowAllocationWebServer/LowAllocationWebServerLibrary/TcpConnectionFormatter.cs b/samples/LowAllocationWebServer/LowAllocationWebServerLibrary/TcpConnectionFormatter.cs
@@ -47,7 +47,16 @@ public Memory<byte> GetMemory(int minimumLength = 0)
             return Memory<byte>.Empty;
         }
 
-        public Span<byte> GetSpan(int minimumLength) => GetMemory(minimumLength).Span;
+        public Span<byte> GetSpan(int minimumLength)
+        {
+            if (_written < 1) throw new NotImplementedException();
+            Send();
+            _written = 0;
+
+            Span<byte> buffer = _buffer.AsSpan(ChunkPrefixSize + _written);
+            if (buffer.Length > 2) return buffer.Slice(0, buffer.Length - 2);
+            return Span<byte>.Empty;
+        }
 
         public int MaxBufferSize { get; } = Int32.MaxValue;
 

diff --git a/src/System.Text.Formatting/System/Text/Formatting/Formatters/BufferWriterFormatter.cs b/src/System.Text.Formatting/System/Text/Formatting/Formatters/BufferWriterFormatter.cs
@@ -31,6 +31,9 @@ public Memory<byte> GetMemory(int minimumLength = 0)
             return _output.GetMemory(minimumLength);
         }
 
-        public Span<byte> GetSpan(int minimumLength) => GetMemory(minimumLength).Span;
+        public Span<byte> GetSpan(int minimumLength)
+        {
+            return _output.GetSpan(minimumLength);
+        }
     }
 }
diff --git a/src/System.Text.Formatting/System/Text/Formatting/Formatters/SequenceFormatter.cs b/src/System.Text.Formatting/System/Text/Formatting/Formatters/SequenceFormatter.cs
@@ -76,7 +76,25 @@ Memory<byte> IBufferWriter<byte>.GetMemory(int minimumLength)
             return Current.Slice(_currentWrittenBytes);
         }
 
-        Span<byte> IBufferWriter<byte>.GetSpan(int minimumLength) => ((IBufferWriter<byte>)this).GetMemory(minimumLength).Span;
+        Span<byte> IBufferWriter<byte>.GetSpan(int minimumLength)
+        {
+            if (minimumLength == 0) minimumLength = 1;
+            if (minimumLength > Current.Length - _currentWrittenBytes)
+            {
+                if (NeedShift) throw new NotImplementedException("need to allocate temp array");
+
+                _previousSequencePosition = _currentSequencePosition;
+                _previousWrittenBytes = _currentWrittenBytes;
+
+                if (!_buffers.TryGet(ref _currentSequencePosition, out Memory<byte> span))
+                {
+                    throw new InvalidOperationException();
+                }
+                _currentWrittenBytes = 0;
+            }
+            // Duplicating logic from GetMemory so we can slice the span rather than slicing the memory
+            return Current.Span.Slice(_currentWrittenBytes);
+        }
 
 
         void IBufferWriter<byte>.Advance(int bytes)

diff --git a/src/System.Text.Formatting/System/Text/Formatting/Formatters/StreamFormatter.cs b/src/System.Text.Formatting/System/Text/Formatting/Formatters/StreamFormatter.cs
@@ -44,7 +44,21 @@ Memory<byte> IBufferWriter<byte>.GetMemory(int minimumLength)
             return _buffer;
         }
 
-        Span<byte> IBufferWriter<byte>.GetSpan(int minimumLength) => ((IBufferWriter<byte>) this).GetMemory(minimumLength).Span;
+        Span<byte> IBufferWriter<byte>.GetSpan(int minimumLength)
+        {
+            if (minimumLength > _buffer.Length)
+            {
+                var newSize = _buffer.Length * 2;
+                if (minimumLength != 0)
+                {
+                    newSize = minimumLength;
+                }
+                var temp = _buffer;
+                _buffer = _pool.Rent(newSize);
+                _pool.Return(temp);
+            }
+            return _buffer;
+        }
 
         // ISSUE
         // I would like to lazy write to the stream, but unfortunatelly this seems to be exclusive with this type being a struct.

diff --git a/src/System.Text.Formatting/System/Text/Formatting/Formatters/StringFormatter.cs b/src/System.Text.Formatting/System/Text/Formatting/Formatters/StringFormatter.cs
@@ -74,7 +74,19 @@ Memory<byte> IBufferWriter<byte>.GetMemory(int minimumLength)
             return _buffer.Free;
         }
 
-        Span<byte> IBufferWriter<byte>.GetSpan(int minimumLength) => ((IBufferWriter<byte>) this).GetMemory(minimumLength).Span;
+        Span<byte> IBufferWriter<byte>.GetSpan(int minimumLength)
+        {
+            if (minimumLength < 1) minimumLength = 1;
+            if (_buffer.Free.Count < minimumLength)
+            {
+                var doubleCount = _buffer.Free.Count * 2;
+                int newSize = minimumLength > doubleCount ? minimumLength : doubleCount;
+                var newArray = _pool.Rent(newSize + _buffer.Count);
+                var oldArray = _buffer.Resize(newArray);
+                _pool.Return(oldArray);
+            }
+            return _buffer.Free;
+        }
 
         void IBufferWriter<byte>.Advance(int bytes)
         {

diff --git a/tests/System.Binary.Base64.Tests/BasicUnitTests.cs b/tests/System.Binary.Base64.Tests/BasicUnitTests.cs
@@ -106,7 +106,10 @@ public Memory<byte> GetMemory(int minimumLength = 0)
             return ((Memory<byte>)_buffer).Slice(_written);
         }
 
-        public Span<byte> GetSpan(int minimumLength) => GetMemory(minimumLength).Span;
+        public Span<byte> GetSpan(int minimumLength)
+        {
+            return ((Span<byte>)_buffer).Slice(_written);
+        }
 
         public int MaxBufferSize { get; } = Int32.MaxValue;
     }

diff --git a/tests/System.Buffers.Experimental.Tests/BufferWriterTests_sequence.cs b/tests/System.Buffers.Experimental.Tests/BufferWriterTests_sequence.cs
@@ -66,7 +66,15 @@ public Memory<byte> GetMemory(int minimumLength = 0)
             return _current;
         }
 
-        public Span<byte> GetSpan(int minimumLength) => GetMemory(minimumLength).Span;
+        public Span<byte> GetSpan(int minimumLength)
+        {
+            if (minimumLength == 0) minimumLength = _current.Length + 1;
+            if (minimumLength < _current.Length) throw new InvalidOperationException();
+            var newBuffer = new byte[minimumLength];
+            _current.CopyTo(newBuffer.AsSpan());
+            _current = newBuffer;
+            return _current;
+        }
 
         public int MaxBufferSize { get; } = Int32.MaxValue;