Reduce memory usage when reading payloads that span multiple packets.

Author: bgrainger

src/MySqlConnector/Protocol/Serialization/ArraySegmentHolder.cs

@@ -0,0 +1,24 @@
+using System;
+
+namespace MySql.Data.Protocol.Serialization
+{
+	/// <summary>
+	/// <see cref="ArraySegmentHolder{T}"/> is a class that holds an instance of <see cref="ArraySegment{T}"/>.
+	/// Its primary difference from <see cref="ArraySegment{T}"/> is that it's a reference type, so mutations
+	/// to this object are visible to other objects that hold a reference to it.
+	/// </summary>
+	internal sealed class ArraySegmentHolder<T>
+	{
+		public ArraySegment<T> ArraySegment { get; set; }
+
+		public T[] Array => ArraySegment.Array;
+		public int Offset => ArraySegment.Offset;
+		public int Count => ArraySegment.Count;
+
+		public void Clear()
+		{
+			if (ArraySegment.Count > 0)
+				ArraySegment = new ArraySegment<T>(ArraySegment.Array, 0, 0);
+		}
+	}
+}

src/MySqlConnector/Protocol/Serialization/CompressedPayloadHandler.cs

@@ -30,8 +30,8 @@ public IByteHandler ByteHandler
 			set => throw new NotSupportedException();
 		}
 
-		public ValueTask<ArraySegment<byte>> ReadPayloadAsync(ProtocolErrorBehavior protocolErrorBehavior, IOBehavior ioBehavior) =>
-			ProtocolUtility.ReadPayloadAsync(m_bufferedByteReader, new CompressedByteHandler(this, protocolErrorBehavior), () => -1, default(ArraySegment<byte>), protocolErrorBehavior, ioBehavior);
+		public ValueTask<ArraySegment<byte>> ReadPayloadAsync(ArraySegmentHolder<byte> cache, ProtocolErrorBehavior protocolErrorBehavior, IOBehavior ioBehavior) =>
+			ProtocolUtility.ReadPayloadAsync(m_bufferedByteReader, new CompressedByteHandler(this, protocolErrorBehavior), () => -1, cache, protocolErrorBehavior, ioBehavior);
 
 		public ValueTask<int> WritePayloadAsync(ArraySegment<byte> payload, IOBehavior ioBehavior)
 		{

src/MySqlConnector/Protocol/Serialization/IPayloadHandler.cs

@@ -19,12 +19,15 @@ internal interface IPayloadHandler
 		/// <summary>
 		/// Reads the next payload.
 		/// </summary>
+		/// <param name="cache">An <see cref="ArraySegmentHolder{Byte}"/> that will cache any buffers allocated during this
+		/// read. (To disable caching, pass <code>new ArraySegmentHolder&lt;byte&gt;</code> so the cache will be garbage-collected
+		/// when this method returns.)</param>
 		/// <param name="protocolErrorBehavior">The <see cref="ProtocolErrorBehavior"/> to use if there is a protocol error.</param>
 		/// <param name="ioBehavior">The <see cref="IOBehavior"/> to use when reading data.</param>
 		/// <returns>An <see cref="ArraySegment{Byte}"/> containing the data that was read. This
 		/// <see cref="ArraySegment{Byte}"/> will be valid to read from until the next time <see cref="ReadPayloadAsync"/> or
 		/// <see cref="WritePayloadAsync"/> is called.</returns>
-		ValueTask<ArraySegment<byte>> ReadPayloadAsync(ProtocolErrorBehavior protocolErrorBehavior, IOBehavior ioBehavior);
+		ValueTask<ArraySegment<byte>> ReadPayloadAsync(ArraySegmentHolder<byte> cache, ProtocolErrorBehavior protocolErrorBehavior, IOBehavior ioBehavior);
 
 		/// <summary>
 		/// Writes a payload.

src/MySqlConnector/Protocol/Serialization/ProtocolUtility.cs

@@ -58,13 +58,19 @@ private static ValueTask<Packet> CreatePacketFromPayload(ArraySegment<byte> payl
 				protocolErrorBehavior == ProtocolErrorBehavior.Throw ? ValueTaskExtensions.FromException<Packet>(new EndOfStreamException()) :
 				default(ValueTask<Packet>);
 
-		public static ValueTask<ArraySegment<byte>> ReadPayloadAsync(BufferedByteReader bufferedByteReader, IByteHandler byteHandler, Func<int> getNextSequenceNumber, ArraySegment<byte> previousPayloads, ProtocolErrorBehavior protocolErrorBehavior, IOBehavior ioBehavior)
+		public static ValueTask<ArraySegment<byte>> ReadPayloadAsync(BufferedByteReader bufferedByteReader, IByteHandler byteHandler, Func<int> getNextSequenceNumber, ArraySegmentHolder<byte> cache, ProtocolErrorBehavior protocolErrorBehavior, IOBehavior ioBehavior)
+		{
+			cache.Clear();
+			return DoReadPayloadAsync(bufferedByteReader, byteHandler, getNextSequenceNumber, cache, protocolErrorBehavior, ioBehavior);
+		}
+
+		private static ValueTask<ArraySegment<byte>> DoReadPayloadAsync(BufferedByteReader bufferedByteReader, IByteHandler byteHandler, Func<int> getNextSequenceNumber, ArraySegmentHolder<byte> previousPayloads, ProtocolErrorBehavior protocolErrorBehavior, IOBehavior ioBehavior)
 		{
 			var readPacketTask = ReadPacketAsync(bufferedByteReader, byteHandler, getNextSequenceNumber, protocolErrorBehavior, ioBehavior);
 			while (readPacketTask.IsCompleted)
 			{
 				ValueTask<ArraySegment<byte>> result;
-				if (HasReadPayload(ref previousPayloads, readPacketTask.Result, protocolErrorBehavior, out result))
+				if (HasReadPayload(previousPayloads, readPacketTask.Result, protocolErrorBehavior, out result))
 					return result;
 
 				readPacketTask = ReadPacketAsync(bufferedByteReader, byteHandler, getNextSequenceNumber, protocolErrorBehavior, ioBehavior);
@@ -73,40 +79,40 @@ public static ValueTask<ArraySegment<byte>> ReadPayloadAsync(BufferedByteReader
 			return AddContinuation(readPacketTask, bufferedByteReader, byteHandler, getNextSequenceNumber, previousPayloads, protocolErrorBehavior, ioBehavior);
 
 			// NOTE: use a local function (with no captures) to defer creation of lambda objects
-			ValueTask<ArraySegment<byte>> AddContinuation(ValueTask<Packet> readPacketTask_, BufferedByteReader bufferedByteReader_, IByteHandler byteHandler_, Func<int> getNextSequenceNumber_, ArraySegment<byte> previousPayloads_, ProtocolErrorBehavior protocolErrorBehavior_, IOBehavior ioBehavior_)
+			ValueTask<ArraySegment<byte>> AddContinuation(ValueTask<Packet> readPacketTask_, BufferedByteReader bufferedByteReader_, IByteHandler byteHandler_, Func<int> getNextSequenceNumber_, ArraySegmentHolder<byte> previousPayloads_, ProtocolErrorBehavior protocolErrorBehavior_, IOBehavior ioBehavior_)
 			{
 				return readPacketTask_.ContinueWith(packet =>
-					HasReadPayload(ref previousPayloads_, packet, protocolErrorBehavior_, out var result_) ? result_ :
-						ReadPayloadAsync(bufferedByteReader_, byteHandler_, getNextSequenceNumber_, previousPayloads_, protocolErrorBehavior_, ioBehavior_));
+					HasReadPayload(previousPayloads_, packet, protocolErrorBehavior_, out var result_) ? result_ :
+						DoReadPayloadAsync(bufferedByteReader_, byteHandler_, getNextSequenceNumber_, previousPayloads_, protocolErrorBehavior_, ioBehavior_));
 			}
 		}
 
-		private static bool HasReadPayload(ref ArraySegment<byte> previousPayloads, Packet packet, ProtocolErrorBehavior protocolErrorBehavior, out ValueTask<ArraySegment<byte>> result)
+		private static bool HasReadPayload(ArraySegmentHolder<byte> previousPayloads, Packet packet, ProtocolErrorBehavior protocolErrorBehavior, out ValueTask<ArraySegment<byte>> result)
 		{
 			if (packet == null && protocolErrorBehavior == ProtocolErrorBehavior.Ignore)
 			{
 				result = default(ValueTask<ArraySegment<byte>>);
 				return true;
 			}
 
-			var previousPayloadsArray = previousPayloads.Array;
-			if (previousPayloadsArray == null && packet.Contents.Count < MaxPacketSize)
+			if (previousPayloads.Count == 0 && packet.Contents.Count < MaxPacketSize)
 			{
 				result = new ValueTask<ArraySegment<byte>>(packet.Contents);
 				return true;
 			}
 
+			var previousPayloadsArray = previousPayloads.Array;
 			if (previousPayloadsArray == null)
 				previousPayloadsArray = new byte[ProtocolUtility.MaxPacketSize + 1];
 			else if (previousPayloads.Offset + previousPayloads.Count + packet.Contents.Count > previousPayloadsArray.Length)
 				Array.Resize(ref previousPayloadsArray, previousPayloadsArray.Length * 2);
 
 			Buffer.BlockCopy(packet.Contents.Array, packet.Contents.Offset, previousPayloadsArray, previousPayloads.Offset + previousPayloads.Count, packet.Contents.Count);
-			previousPayloads = new ArraySegment<byte>(previousPayloadsArray, previousPayloads.Offset, previousPayloads.Count + packet.Contents.Count);
+			previousPayloads.ArraySegment = new ArraySegment<byte>(previousPayloadsArray, previousPayloads.Offset, previousPayloads.Count + packet.Contents.Count);
 
 			if (packet.Contents.Count < ProtocolUtility.MaxPacketSize)
 			{
-				result = new ValueTask<ArraySegment<byte>>(previousPayloads);
+				result = new ValueTask<ArraySegment<byte>>(previousPayloads.ArraySegment);
 				return true;
 			}
 

src/MySqlConnector/Protocol/Serialization/StandardPayloadHandler.cs

@@ -26,8 +26,8 @@ public IByteHandler ByteHandler
 			}
 		}
 
-		public ValueTask<ArraySegment<byte>> ReadPayloadAsync(ProtocolErrorBehavior protocolErrorBehavior, IOBehavior ioBehavior) =>
-			ProtocolUtility.ReadPayloadAsync(m_bufferedByteReader, m_byteHandler, m_getNextSequenceNumber, default(ArraySegment<byte>), protocolErrorBehavior, ioBehavior);
+		public ValueTask<ArraySegment<byte>> ReadPayloadAsync(ArraySegmentHolder<byte> cache, ProtocolErrorBehavior protocolErrorBehavior, IOBehavior ioBehavior) =>
+			ProtocolUtility.ReadPayloadAsync(m_bufferedByteReader, m_byteHandler, m_getNextSequenceNumber, cache, protocolErrorBehavior, ioBehavior);
 
 		public ValueTask<int> WritePayloadAsync(ArraySegment<byte> payload, IOBehavior ioBehavior) =>
 			ProtocolUtility.WritePayloadAsync(m_byteHandler, m_getNextSequenceNumber, payload, ioBehavior);

src/MySqlConnector/Serialization/MySqlSession.cs

@@ -24,6 +24,7 @@ public MySqlSession()
 		public MySqlSession(ConnectionPool pool, int poolGeneration)
 		{
 			m_lock = new object();
+			m_payloadCache = new ArraySegmentHolder<byte>();
 			CreatedUtc = DateTime.UtcNow;
 			Pool = pool;
 			PoolGeneration = poolGeneration;
@@ -167,7 +168,7 @@ public async Task DisposeAsync(IOBehavior ioBehavior, CancellationToken cancella
 				{
 					m_payloadHandler.StartNewConversation();
 					await m_payloadHandler.WritePayloadAsync(QuitPayload.Create(), ioBehavior).ConfigureAwait(false);
-					await m_payloadHandler.ReadPayloadAsync(ProtocolErrorBehavior.Ignore, ioBehavior).ConfigureAwait(false);
+					await m_payloadHandler.ReadPayloadAsync(m_payloadCache, ProtocolErrorBehavior.Ignore, ioBehavior).ConfigureAwait(false);
 				}
 				catch (IOException)
 				{
@@ -321,7 +322,7 @@ public ValueTask<PayloadData> ReceiveReplyAsync(IOBehavior ioBehavior, Cancellat
 			try
 			{
 				VerifyConnected();
-				task = m_payloadHandler.ReadPayloadAsync(ProtocolErrorBehavior.Throw, ioBehavior);
+				task = m_payloadHandler.ReadPayloadAsync(m_payloadCache, ProtocolErrorBehavior.Throw, ioBehavior);
 			}
 			catch (Exception ex)
 			{
@@ -690,6 +691,7 @@ private enum State
 		}
 
 		readonly object m_lock;
+		readonly ArraySegmentHolder<byte> m_payloadCache;
 		State m_state;
 		string m_hostname = "";
 		TcpClient m_tcpClient;