Skip to content

Commit

Permalink
Rent byte buffers in SqlSequentialTextReader (#2356)
Browse files Browse the repository at this point in the history
  • Loading branch information
TrayanZapryanov authored Apr 8, 2024
1 parent 9811c65 commit 1bf025a
Showing 1 changed file with 27 additions and 18 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@
// See the LICENSE file in the project root for more information.

using System;
using System.Buffers;
using System.Buffers.Binary;
using System.Diagnostics;
using System.Text;
Expand All @@ -18,7 +19,8 @@ sealed internal class SqlSequentialTextReader : System.IO.TextReader
private readonly int _columnIndex; // The index of out column in the table
private readonly Encoding _encoding; // Encoding for this character stream
private readonly Decoder _decoder; // Decoder based on the encoding (NOTE: Decoders are stateful as they are designed to process streams of data)
private byte[] _leftOverBytes; // Bytes leftover from the last Read() operation - this can be null if there were no bytes leftover (Possible optimization: re-use the same array?)
private byte[] _leftOverBytes; // Bytes leftover from the last Read() operation - this can be null if there were no bytes leftover
private int _leftOverByteBufferUsed; //Number of bytes used from _leftOverBytes buffer - will be zero in case of null buffer
private int _peekedChar; // The last character that we peeked at (or -1 if we haven't peeked at anything)
private Task _currentTask; // The current async task
private readonly CancellationTokenSource _disposalTokenSource; // Used to indicate that a cancellation is requested due to disposal
Expand Down Expand Up @@ -357,23 +359,18 @@ private byte[] PrepareByteBuffer(int numberOfChars, out int byteBufferUsed)

if (_leftOverBytes != null)
{
// If we have more leftover bytes than we need for this conversion, then just re-use the leftover buffer
if (_leftOverBytes.Length > byteBufferSize)
{
byteBuffer = _leftOverBytes;
byteBufferUsed = byteBuffer.Length;
}
else
{
// Otherwise, copy over the leftover buffer
byteBuffer = new byte[byteBufferSize];
Buffer.BlockCopy(_leftOverBytes, 0, byteBuffer, 0, _leftOverBytes.Length);
byteBufferUsed = _leftOverBytes.Length;
}
// Copy over the leftover buffer
byteBuffer = ArrayPool<byte>.Shared.Rent(byteBufferSize);
Buffer.BlockCopy(_leftOverBytes, 0, byteBuffer, 0, _leftOverByteBufferUsed);
byteBufferUsed = _leftOverByteBufferUsed;
//return _leftOverBytes and clean _leftOverBytes reference
ArrayPool<byte>.Shared.Return(_leftOverBytes);
_leftOverBytes = null;
_leftOverByteBufferUsed = 0;
}
else
{
byteBuffer = new byte[byteBufferSize];
byteBuffer = ArrayPool<byte>.Shared.Rent(byteBufferSize);
byteBufferUsed = 0;
}
}
Expand Down Expand Up @@ -402,14 +399,26 @@ private int DecodeBytesToChars(byte[] inBuffer, int inBufferCount, char[] outBuf
// completed may be false and there is no spare bytes if the Decoder has stored bytes to use later
if ((!completed) && (bytesUsed < inBufferCount))
{
_leftOverBytes = new byte[inBufferCount - bytesUsed];
Buffer.BlockCopy(inBuffer, bytesUsed, _leftOverBytes, 0, _leftOverBytes.Length);
_leftOverByteBufferUsed = inBufferCount - bytesUsed;
_leftOverBytes = ArrayPool<byte>.Shared.Rent(_leftOverByteBufferUsed);

Buffer.BlockCopy(inBuffer, bytesUsed, _leftOverBytes, 0, _leftOverByteBufferUsed);
}
else
{
// If Convert() sets completed to true, then it must have used all of the bytes we gave it
Debug.Assert(bytesUsed >= inBufferCount, "Converted completed, but not all bytes were used");
_leftOverBytes = null;
if (_leftOverBytes != null)
{
ArrayPool<byte>.Shared.Return(_leftOverBytes);
_leftOverBytes = null;
_leftOverByteBufferUsed = 0;
}
}

if (inBuffer.Length > 0)
{
ArrayPool<byte>.Shared.Return(inBuffer);
}

Debug.Assert(((_reader == null) || (_reader.ColumnDataBytesRemaining() > 0) || (!completed) || (_leftOverBytes == null)), "Stream has run out of data and the decoder finished, but there are leftover bytes");
Expand Down

0 comments on commit 1bf025a

Please sign in to comment.