Reduce StreamWriter allocation (#44495)

* Lazily-allocate StreamWriter's byte[] buffer

The byte[] buffer is used just when flushing to the underlying stream, to store the intermediate bytes as generated from the characters by the encoding.  Currently the byte[] buffer is allocated in the constructor, but for relatively small strings written/flushed synchronously (such as those often used with Console.Write), we can avoid the buffer entirely via stack allocation (in Flush); for all other cases, we can allocate it on demand.

(I considered using ArrayPool to rent/return in each Flush{Async} call, but opted not to in this change.  Someone could investigate that subsequently.)

* Undo use of static async helpers

In .NET Framework, there was a non-trivial penalty due to StreamWriter deriving from MarshalByRefObject and the impact that had on async methods accessing lifted locals on the state machine (it made them significantly more expensive).  The workaround was to make the async methods be statics and pass in all of the relevant instance state needed.  That workaround is no longer necessary on core, where MarshalByRefObjects are nops, leaving behind a workaround that is not just clunkier, but actually makes things more expensive because more state needs to be stored on the state machine objects (all of the arguments, which can instead just be accessed off of `this`).  Undo that whole change.

* Address PR feedback

diff --git a/src/libraries/System.Private.CoreLib/src/System/IO/StreamWriter.cs b/src/libraries/System.Private.CoreLib/src/System/IO/StreamWriter.cs
index 8b88aa0..ed71b29 100644 (file)
--- a/src/libraries/System.Private.CoreLib/src/System/IO/StreamWriter.cs
+++ b/src/libraries/System.Private.CoreLib/src/System/IO/StreamWriter.cs
@@ -32,7 +32,7 @@ namespace System.IO
         private readonly Stream _stream;
         private readonly Encoding _encoding;
         private readonly Encoder _encoder;
-        private readonly byte[] _byteBuffer;
+        private byte[]? _byteBuffer;
         private readonly char[] _charBuffer;
         private int _charPos;
         private int _charLen;
@@ -123,7 +123,6 @@ namespace System.IO
             }
 
             _charBuffer = new char[bufferSize];
-            _byteBuffer = new byte[_encoding.GetMaxByteCount(bufferSize)];
             _charLen = bufferSize;
             // If we're appending to a Stream that already has data, don't write
             // the preamble.
@@ -275,15 +274,28 @@ namespace System.IO
                 }
             }
 
-            int count = _encoder.GetBytes(_charBuffer, 0, _charPos, _byteBuffer, 0, flushEncoder);
+            // For sufficiently small char data being flushed, try to encode to the stack.
+            // For anything else, fall back to allocating the byte[] buffer.
+            Span<byte> byteBuffer = stackalloc byte[0];
+            if (_byteBuffer is not null)
+            {
+                byteBuffer = _byteBuffer;
+            }
+            else
+            {
+                int maxBytesForCharPos = _encoding.GetMaxByteCount(_charPos);
+                byteBuffer = maxBytesForCharPos <= 1024 ? // arbitrary threshold
+                    stackalloc byte[1024] :
+                    (_byteBuffer = new byte[_encoding.GetMaxByteCount(_charBuffer.Length)]);
+            }
+
+            int count = _encoder.GetBytes(new ReadOnlySpan<char>(_charBuffer, 0, _charPos), byteBuffer, flushEncoder);
             _charPos = 0;
             if (count > 0)
             {
-                _stream.Write(_byteBuffer, 0, count);
+                _stream.Write(byteBuffer.Slice(0, count));
             }
-            // By definition, calling Flush should flush the stream, but this is
-            // only necessary if we passed in true for flushStream.  The Web
-            // Services guys have some perf tests where flushing needlessly hurts.
+
             if (flushStream)
             {
                 _stream.Flush();
@@ -613,53 +625,38 @@ namespace System.IO
             ThrowIfDisposed();
             CheckAsyncTaskInProgress();
 
-            Task task = WriteAsyncInternal(this, value, _charBuffer, _charPos, _charLen, CoreNewLine, _autoFlush, appendNewLine: false);
+            Task task = WriteAsyncInternal(value, appendNewLine: false);
             _asyncWriteTask = task;
 
             return task;
         }
 
-        // We pass in private instance fields of this MarshalByRefObject-derived type as local params
-        // to ensure performant access inside the state machine that corresponds this async method.
-        // Fields that are written to must be assigned at the end of the method *and* before instance invocations.
-        private static async Task WriteAsyncInternal(StreamWriter _this, char value,
-                                                     char[] charBuffer, int charPos, int charLen, char[] coreNewLine,
-                                                     bool autoFlush, bool appendNewLine)
+        private async Task WriteAsyncInternal(char value, bool appendNewLine)
         {
-            if (charPos == charLen)
+            if (_charPos == _charLen)
             {
-                await _this.FlushAsyncInternal(false, false, charBuffer, charPos).ConfigureAwait(false);
-                Debug.Assert(_this._charPos == 0);
-                charPos = 0;
+                await FlushAsyncInternal(flushStream: false, flushEncoder: false).ConfigureAwait(false);
             }
 
-            charBuffer[charPos] = value;
-            charPos++;
+            _charBuffer[_charPos++] = value;
 
             if (appendNewLine)
             {
-                for (int i = 0; i < coreNewLine.Length; i++)   // Expect 2 iterations, no point calling BlockCopy
+                for (int i = 0; i < CoreNewLine.Length; i++) // Generally 1 (\n) or 2 (\r\n) iterations
                 {
-                    if (charPos == charLen)
+                    if (_charPos == _charLen)
                     {
-                        await _this.FlushAsyncInternal(false, false, charBuffer, charPos).ConfigureAwait(false);
-                        Debug.Assert(_this._charPos == 0);
-                        charPos = 0;
+                        await FlushAsyncInternal(flushStream: false, flushEncoder: false).ConfigureAwait(false);
                     }
 
-                    charBuffer[charPos] = coreNewLine[i];
-                    charPos++;
+                    _charBuffer[_charPos++] = CoreNewLine[i];
                 }
             }
 
-            if (autoFlush)
+            if (_autoFlush)
             {
-                await _this.FlushAsyncInternal(true, false, charBuffer, charPos).ConfigureAwait(false);
-                Debug.Assert(_this._charPos == 0);
-                charPos = 0;
+                await FlushAsyncInternal(flushStream: true, flushEncoder: false).ConfigureAwait(false);
             }
-
-            _this._charPos = charPos;
         }
 
         public override Task WriteAsync(string? value)
@@ -678,7 +675,7 @@ namespace System.IO
                 ThrowIfDisposed();
                 CheckAsyncTaskInProgress();
 
-                Task task = WriteAsyncInternal(this, value, _charBuffer, _charPos, _charLen, CoreNewLine, _autoFlush, appendNewLine: false);
+                Task task = WriteAsyncInternal(value.AsMemory(), appendNewLine: false, default);
                 _asyncWriteTask = task;
 
                 return task;
@@ -689,68 +686,6 @@ namespace System.IO
             }
         }
 
-        // We pass in private instance fields of this MarshalByRefObject-derived type as local params
-        // to ensure performant access inside the state machine that corresponds this async method.
-        // Fields that are written to must be assigned at the end of the method *and* before instance invocations.
-        private static async Task WriteAsyncInternal(StreamWriter _this, string value,
-                                                     char[] charBuffer, int charPos, int charLen, char[] coreNewLine,
-                                                     bool autoFlush, bool appendNewLine)
-        {
-            Debug.Assert(value != null);
-
-            int count = value.Length;
-            int index = 0;
-
-            while (count > 0)
-            {
-                if (charPos == charLen)
-                {
-                    await _this.FlushAsyncInternal(false, false, charBuffer, charPos).ConfigureAwait(false);
-                    Debug.Assert(_this._charPos == 0);
-                    charPos = 0;
-                }
-
-                int n = charLen - charPos;
-                if (n > count)
-                {
-                    n = count;
-                }
-
-                Debug.Assert(n > 0, "StreamWriter::Write(String) isn't making progress!  This is most likely a race condition in user code.");
-
-                value.CopyTo(index, charBuffer, charPos, n);
-
-                charPos += n;
-                index += n;
-                count -= n;
-            }
-
-            if (appendNewLine)
-            {
-                for (int i = 0; i < coreNewLine.Length; i++)   // Expect 2 iterations, no point calling BlockCopy
-                {
-                    if (charPos == charLen)
-                    {
-                        await _this.FlushAsyncInternal(false, false, charBuffer, charPos).ConfigureAwait(false);
-                        Debug.Assert(_this._charPos == 0);
-                        charPos = 0;
-                    }
-
-                    charBuffer[charPos] = coreNewLine[i];
-                    charPos++;
-                }
-            }
-
-            if (autoFlush)
-            {
-                await _this.FlushAsyncInternal(true, false, charBuffer, charPos).ConfigureAwait(false);
-                Debug.Assert(_this._charPos == 0);
-                charPos = 0;
-            }
-
-            _this._charPos = charPos;
-        }
-
         public override Task WriteAsync(char[] buffer, int index, int count)
         {
             if (buffer == null)
@@ -782,7 +717,7 @@ namespace System.IO
             ThrowIfDisposed();
             CheckAsyncTaskInProgress();
 
-            Task task = WriteAsyncInternal(this, new ReadOnlyMemory<char>(buffer, index, count), _charBuffer, _charPos, _charLen, CoreNewLine, _autoFlush, appendNewLine: false, cancellationToken: default);
+            Task task = WriteAsyncInternal(new ReadOnlyMemory<char>(buffer, index, count), appendNewLine: false, cancellationToken: default);
             _asyncWriteTask = task;
 
             return task;
@@ -804,60 +739,46 @@ namespace System.IO
                 return Task.FromCanceled(cancellationToken);
             }
 
-            Task task = WriteAsyncInternal(this, buffer, _charBuffer, _charPos, _charLen, CoreNewLine, _autoFlush, appendNewLine: false, cancellationToken: cancellationToken);
+            Task task = WriteAsyncInternal(buffer, appendNewLine: false, cancellationToken: cancellationToken);
             _asyncWriteTask = task;
             return task;
         }
 
-        // We pass in private instance fields of this MarshalByRefObject-derived type as local params
-        // to ensure performant access inside the state machine that corresponds this async method.
-        // Fields that are written to must be assigned at the end of the method *and* before instance invocations.
-        private static async Task WriteAsyncInternal(StreamWriter _this, ReadOnlyMemory<char> source,
-                                                     char[] charBuffer, int charPos, int charLen, char[] coreNewLine,
-                                                     bool autoFlush, bool appendNewLine, CancellationToken cancellationToken)
+        private async Task WriteAsyncInternal(ReadOnlyMemory<char> source, bool appendNewLine, CancellationToken cancellationToken)
         {
             int copied = 0;
             while (copied < source.Length)
             {
-                if (charPos == charLen)
+                if (_charPos == _charLen)
                 {
-                    await _this.FlushAsyncInternal(false, false, charBuffer, charPos, cancellationToken).ConfigureAwait(false);
-                    Debug.Assert(_this._charPos == 0);
-                    charPos = 0;
+                    await FlushAsyncInternal(flushStream: false, flushEncoder: false, cancellationToken).ConfigureAwait(false);
                 }
 
-                int n = Math.Min(charLen - charPos, source.Length - copied);
+                int n = Math.Min(_charLen - _charPos, source.Length - copied);
                 Debug.Assert(n > 0, "StreamWriter::Write(char[], int, int) isn't making progress!  This is most likely a race condition in user code.");
 
-                source.Span.Slice(copied, n).CopyTo(new Span<char>(charBuffer, charPos, n));
-                charPos += n;
+                source.Span.Slice(copied, n).CopyTo(new Span<char>(_charBuffer, _charPos, n));
+                _charPos += n;
                 copied += n;
             }
 
             if (appendNewLine)
             {
-                for (int i = 0; i < coreNewLine.Length; i++)   // Expect 2 iterations, no point calling BlockCopy
+                for (int i = 0; i < CoreNewLine.Length; i++) // Generally 1 (\n) or 2 (\r\n) iterations
                 {
-                    if (charPos == charLen)
+                    if (_charPos == _charLen)
                     {
-                        await _this.FlushAsyncInternal(false, false, charBuffer, charPos, cancellationToken).ConfigureAwait(false);
-                        Debug.Assert(_this._charPos == 0);
-                        charPos = 0;
+                        await FlushAsyncInternal(flushStream: false, flushEncoder: false, cancellationToken).ConfigureAwait(false);
                     }
 
-                    charBuffer[charPos] = coreNewLine[i];
-                    charPos++;
+                    _charBuffer[_charPos++] = CoreNewLine[i];
                 }
             }
 
-            if (autoFlush)
+            if (_autoFlush)
             {
-                await _this.FlushAsyncInternal(true, false, charBuffer, charPos, cancellationToken).ConfigureAwait(false);
-                Debug.Assert(_this._charPos == 0);
-                charPos = 0;
+                await FlushAsyncInternal(flushStream: true, flushEncoder: false, cancellationToken).ConfigureAwait(false);
             }
-
-            _this._charPos = charPos;
         }
 
         public override Task WriteLineAsync()
@@ -874,13 +795,12 @@ namespace System.IO
             ThrowIfDisposed();
             CheckAsyncTaskInProgress();
 
-            Task task = WriteAsyncInternal(this, ReadOnlyMemory<char>.Empty, _charBuffer, _charPos, _charLen, CoreNewLine, _autoFlush, appendNewLine: true, cancellationToken: default);
+            Task task = WriteAsyncInternal(ReadOnlyMemory<char>.Empty, appendNewLine: true, cancellationToken: default);
             _asyncWriteTask = task;
 
             return task;
         }
 
-
         public override Task WriteLineAsync(char value)
         {
             // If we have been inherited into a subclass, the following implementation could be incorrect
@@ -895,13 +815,12 @@ namespace System.IO
             ThrowIfDisposed();
             CheckAsyncTaskInProgress();
 
-            Task task = WriteAsyncInternal(this, value, _charBuffer, _charPos, _charLen, CoreNewLine, _autoFlush, appendNewLine: true);
+            Task task = WriteAsyncInternal(value, appendNewLine: true);
             _asyncWriteTask = task;
 
             return task;
         }
 
-
         public override Task WriteLineAsync(string? value)
         {
             if (value == null)
@@ -921,13 +840,12 @@ namespace System.IO
             ThrowIfDisposed();
             CheckAsyncTaskInProgress();
 
-            Task task = WriteAsyncInternal(this, value, _charBuffer, _charPos, _charLen, CoreNewLine, _autoFlush, appendNewLine: true);
+            Task task = WriteAsyncInternal(value.AsMemory(), appendNewLine: true, default);
             _asyncWriteTask = task;
 
             return task;
         }
 
-
         public override Task WriteLineAsync(char[] buffer, int index, int count)
         {
             if (buffer == null)
@@ -959,7 +877,7 @@ namespace System.IO
             ThrowIfDisposed();
             CheckAsyncTaskInProgress();
 
-            Task task = WriteAsyncInternal(this, new ReadOnlyMemory<char>(buffer, index, count), _charBuffer, _charPos, _charLen, CoreNewLine, _autoFlush, appendNewLine: true, cancellationToken: default);
+            Task task = WriteAsyncInternal(new ReadOnlyMemory<char>(buffer, index, count), appendNewLine: true, cancellationToken: default);
             _asyncWriteTask = task;
 
             return task;
@@ -980,13 +898,12 @@ namespace System.IO
                 return Task.FromCanceled(cancellationToken);
             }
 
-            Task task = WriteAsyncInternal(this, buffer, _charBuffer, _charPos, _charLen, CoreNewLine, _autoFlush, appendNewLine: true, cancellationToken: cancellationToken);
+            Task task = WriteAsyncInternal(buffer, appendNewLine: true, cancellationToken: cancellationToken);
             _asyncWriteTask = task;
 
             return task;
         }
 
-
         public override Task FlushAsync()
         {
             // If we have been inherited into a subclass, the following implementation could be incorrect
@@ -1005,14 +922,13 @@ namespace System.IO
             ThrowIfDisposed();
             CheckAsyncTaskInProgress();
 
-            Task task = FlushAsyncInternal(true, true, _charBuffer, _charPos);
+            Task task = FlushAsyncInternal(flushStream: true, flushEncoder: true);
             _asyncWriteTask = task;
 
             return task;
         }
 
-        private Task FlushAsyncInternal(bool flushStream, bool flushEncoder,
-                                        char[] sCharBuffer, int sCharPos, CancellationToken cancellationToken = default)
+        private Task FlushAsyncInternal(bool flushStream, bool flushEncoder, CancellationToken cancellationToken = default)
         {
             if (cancellationToken.IsCancellationRequested)
             {
@@ -1020,47 +936,43 @@ namespace System.IO
             }
 
             // Perf boost for Flush on non-dirty writers.
-            if (sCharPos == 0 && !flushStream && !flushEncoder)
+            if (_charPos == 0 && !flushStream && !flushEncoder)
             {
                 return Task.CompletedTask;
             }
 
-            Task flushTask = FlushAsyncInternal(this, flushStream, flushEncoder, sCharBuffer, sCharPos, _haveWrittenPreamble,
-                                                _encoding, _encoder, _byteBuffer, _stream, cancellationToken);
+            Task flushTask = Core(flushStream, flushEncoder, cancellationToken);
 
             _charPos = 0;
             return flushTask;
-        }
 
-
-        // We pass in private instance fields of this MarshalByRefObject-derived type as local params
-        // to ensure performant access inside the state machine that corresponds this async method.
-        private static async Task FlushAsyncInternal(StreamWriter _this, bool flushStream, bool flushEncoder,
-                                                     char[] charBuffer, int charPos, bool haveWrittenPreamble,
-                                                     Encoding encoding, Encoder encoder, byte[] byteBuffer, Stream stream, CancellationToken cancellationToken)
-        {
-            if (!haveWrittenPreamble)
+            async Task Core(bool flushStream, bool flushEncoder, CancellationToken cancellationToken)
             {
-                _this._haveWrittenPreamble = true;
-                byte[] preamble = encoding.GetPreamble();
-                if (preamble.Length > 0)
+                if (!_haveWrittenPreamble)
                 {
-                    await stream.WriteAsync(new ReadOnlyMemory<byte>(preamble), cancellationToken).ConfigureAwait(false);
+                    _haveWrittenPreamble = true;
+                    byte[] preamble = _encoding.GetPreamble();
+                    if (preamble.Length > 0)
+                    {
+                        await _stream.WriteAsync(new ReadOnlyMemory<byte>(preamble), cancellationToken).ConfigureAwait(false);
+                    }
                 }
-            }
 
-            int count = encoder.GetBytes(charBuffer, 0, charPos, byteBuffer, 0, flushEncoder);
-            if (count > 0)
-            {
-                await stream.WriteAsync(new ReadOnlyMemory<byte>(byteBuffer, 0, count), cancellationToken).ConfigureAwait(false);
-            }
+                byte[] byteBuffer = _byteBuffer ??= new byte[_encoding.GetMaxByteCount(_charBuffer.Length)];
 
-            // By definition, calling Flush should flush the stream, but this is
-            // only necessary if we passed in true for flushStream.  The Web
-            // Services guys have some perf tests where flushing needlessly hurts.
-            if (flushStream)
-            {
-                await stream.FlushAsync(cancellationToken).ConfigureAwait(false);
+                int count = _encoder.GetBytes(new ReadOnlySpan<char>(_charBuffer, 0, _charPos), byteBuffer, flushEncoder);
+                if (count > 0)
+                {
+                    await _stream.WriteAsync(new ReadOnlyMemory<byte>(byteBuffer, 0, count), cancellationToken).ConfigureAwait(false);
+                }
+
+                // By definition, calling Flush should flush the stream, but this is
+                // only necessary if we passed in true for flushStream.  The Web
+                // Services guys have some perf tests where flushing needlessly hurts.
+                if (flushStream)
+                {
+                    await _stream.FlushAsync(cancellationToken).ConfigureAwait(false);
+                }
             }
         }
 
@@ -1073,5 +985,5 @@ namespace System.IO
 
             void ThrowObjectDisposedException() => throw new ObjectDisposedException(GetType().Name, SR.ObjectDisposed_WriterClosed);
         }
-    }  // class StreamWriter
-}  // namespace
+    }
+}