Optimize SqlClient SNIPacket async paths

src/System.Data.SqlClient/src/System/Data/SqlClient/SNI/SNIPacket.cs

@@ -247,37 +247,50 @@ public void Reset()
         /// <param name="callback">Completion callback</param>
         public void ReadFromStreamAsync(Stream stream, SNIAsyncCallback callback)
         {
-            bool error = false;
-
-            stream.ReadAsync(_data, 0, _capacity, CancellationToken.None).ContinueWith(t =>
+            // Treat local function as a static and pass all params otherwise as async will allocate
+            async Task ReadFromStreamAsync(SNIPacket packet, SNIAsyncCallback cb, ValueTask<int> valueTask)
             {
-                Exception e = t.Exception?.InnerException;
-                if (e != null)
+                bool error = false;
+                try
                 {
-                    SNILoadHandle.SingletonInstance.LastError = new SNIError(SNIProviders.TCP_PROV, SNICommon.InternalExceptionError, e);
-                    error = true;
-                }
-                else
-                {
-                    _length = t.Result;
-
-                    if (_length == 0)
+                    packet._length = await valueTask.ConfigureAwait(false);
+                    if (packet._length == 0)
                     {
                         SNILoadHandle.SingletonInstance.LastError = new SNIError(SNIProviders.TCP_PROV, 0, SNICommon.ConnTerminatedError, string.Empty);
                         error = true;
                     }
                 }
+                catch (Exception ex)
+                {
+                    SNILoadHandle.SingletonInstance.LastError = new SNIError(SNIProviders.TCP_PROV, SNICommon.InternalExceptionError, ex);
+                    error = true;
+                }
 
                 if (error)
                 {
-                    Release();
+                    packet.Release();
                 }
 
-                callback(this, error ? TdsEnums.SNI_ERROR : TdsEnums.SNI_SUCCESS);
-            },
-            CancellationToken.None,
-            TaskContinuationOptions.DenyChildAttach,
-            TaskScheduler.Default);
+                cb(packet, error ? TdsEnums.SNI_ERROR : TdsEnums.SNI_SUCCESS);
+            }
+
+            ValueTask<int> vt = stream.ReadAsync(new Memory<byte>(_data, 0, _capacity), CancellationToken.None);
+
+            if (vt.IsCompletedSuccessfully)
+            {
+                _length = vt.Result;
+                // Zero length to go via async local function as is error condition
+                if (_length > 0)
+                {
+                    callback(this, TdsEnums.SNI_SUCCESS);
+
+                    // Completed
+                    return;
+                }
+            }
+
+            // Not complete or error call the async local function to complete
+            _ = ReadFromStreamAsync(this, callback, vt);
         }
 
         /// <summary>
@@ -302,24 +315,50 @@ public void WriteToStream(Stream stream)
         /// Write data to a stream asynchronously
         /// </summary>
         /// <param name="stream">Stream to write to</param>
-        public async void WriteToStreamAsync(Stream stream, SNIAsyncCallback callback, SNIProviders provider, bool disposeAfterWriteAsync = false)
+        public void WriteToStreamAsync(Stream stream, SNIAsyncCallback callback, SNIProviders provider, bool disposeAfterWriteAsync = false)
         {
-            uint status = TdsEnums.SNI_SUCCESS;
-            try
+            // Treat local function as a static and pass all params otherwise as async will allocate
+            async Task WriteToStreamAsync(SNIPacket packet, SNIAsyncCallback cb, SNIProviders providers, bool disposeAfter, ValueTask valueTask)
             {
-                await stream.WriteAsync(_data, 0, _length, CancellationToken.None).ConfigureAwait(false);
-            }
-            catch (Exception e)
-            {
-                SNILoadHandle.SingletonInstance.LastError = new SNIError(provider, SNICommon.InternalExceptionError, e);
-                status = TdsEnums.SNI_ERROR;
+                uint status = TdsEnums.SNI_SUCCESS;
+                try
+                {
+                    await valueTask.ConfigureAwait(false);
+                }
+                catch (Exception e)
+                {
+                    SNILoadHandle.SingletonInstance.LastError = new SNIError(providers, SNICommon.InternalExceptionError, e);
+                    status = TdsEnums.SNI_ERROR;
+                }
+
+                cb(packet, status);
+
+                if (disposeAfter)
+                {
+                    packet.Dispose();
+                }
             }
-            callback(this, status);
 
-            if (disposeAfterWriteAsync)
+            ValueTask vt = stream.WriteAsync(new Memory<byte>(_data, 0, _length), CancellationToken.None);
+
+            if (vt.IsCompletedSuccessfully)
             {
-                Dispose();
+                // Read the result to register as complete for the ValueTask
+                vt.GetAwaiter().GetResult();
+
+                callback(this, TdsEnums.SNI_SUCCESS);
+
+                if (disposeAfterWriteAsync)
+                {
+                    Dispose();
+                }
+
+                // Completed
+                return;
             }
+
+            // Not complete or error call the async local function to complete
+            _ = WriteToStreamAsync(this, callback, provider, disposeAfterWriteAsync, vt);
         }
 
         /// <summary>