Merging some methods of TdsParserStateObject (note: ports part of dot…

…net#1060 from netcore to common file).
panoskj · Mar 18, 2022 · 417a6da · 417a6da
1 parent 72d8e54
commit 417a6da
Show file tree

Hide file tree

Showing 4 changed files with 73 additions and 81 deletions.
diff --git a/src/Microsoft.Data.SqlClient/netcore/src/Microsoft/Data/SqlClient/TdsParserStateObject.cs b/src/Microsoft.Data.SqlClient/netcore/src/Microsoft/Data/SqlClient/TdsParserStateObject.cs
@@ -300,51 +300,6 @@ internal void StartSession(object cancellationOwner)
             _cancellationOwner.Target = cancellationOwner;
         }
 
-        internal void ThrowExceptionAndWarning(bool callerHasConnectionLock = false, bool asyncClose = false)
-        {
-            _parser.ThrowExceptionAndWarning(this, callerHasConnectionLock, asyncClose);
-        }
-
-        ////////////////////////////////////////////
-        // TDS Packet/buffer manipulation methods //
-        ////////////////////////////////////////////
-
-        internal Task ExecuteFlush()
-        {
-            lock (this)
-            {
-                if (_cancelled && 1 == _outputPacketNumber)
-                {
-                    ResetBuffer();
-                    _cancelled = false;
-                    throw SQL.OperationCancelled();
-                }
-                else
-                {
-                    Task writePacketTask = WritePacket(TdsEnums.HARDFLUSH);
-                    if (writePacketTask == null)
-                    {
-                        HasPendingData = true;
-                        _messageStatus = 0;
-                        return null;
-                    }
-                    else
-                    {
-                        return AsyncHelper.CreateContinuationTaskWithState(
-                            task: writePacketTask, 
-                            state: this,
-                            onSuccess: static (object state) => 
-                            {
-                                TdsParserStateObject stateObject = (TdsParserStateObject)state;
-                                stateObject.HasPendingData = true;
-                                stateObject._messageStatus = 0; 
-                            }
-                        );
-                    }
-                }
-            }
-        }
-
         // Processes the tds header that is present in the buffer
         internal bool TryProcessHeader()
         {

diff --git a/src/Microsoft.Data.SqlClient/netfx/src/Microsoft/Data/SqlClient/SqlUtil.cs b/src/Microsoft.Data.SqlClient/netfx/src/Microsoft/Data/SqlClient/SqlUtil.cs
@@ -46,6 +46,28 @@ internal static Task CreateContinuationTask(Task task, Action onSuccess, SqlInte
             }
         }
 
+        internal static Task CreateContinuationTaskWithState(Task task, object state, Action<object> onSuccess, Action<Exception, object> onFailure = null)
+        {
+            if (task == null)
+            {
+                onSuccess(state);
+                return null;
+            }
+            else
+            {
+                var completion = new TaskCompletionSource<object>();
+                ContinueTaskWithState(task, completion, state,
+                    onSuccess: (object continueState) =>
+                    {
+                        onSuccess(continueState);
+                        completion.SetResult(null);
+                    },
+                    onFailure: onFailure
+                );
+                return completion.Task;
+            }
+        }
+
         internal static Task CreateContinuationTask<T1, T2>(Task task, Action<T1, T2> onSuccess, T1 arg1, T2 arg2, SqlInternalConnectionTds connectionToDoom = null, Action<Exception> onFailure = null)
         {
             return CreateContinuationTask(task, () => onSuccess(arg1, arg2), connectionToDoom, onFailure);

diff --git a/src/Microsoft.Data.SqlClient/netfx/src/Microsoft/Data/SqlClient/TdsParserStateObject.cs b/src/Microsoft.Data.SqlClient/netfx/src/Microsoft/Data/SqlClient/TdsParserStateObject.cs
@@ -111,6 +111,12 @@ internal bool HasOpenResult
             set => _hasOpenResult = value;
         }
 
+        internal bool HasPendingData
+        {
+            get => _pendingData;
+            set => _pendingData = value;
+        }
+
         internal UInt32 Status
         {
             get
@@ -434,42 +440,6 @@ internal void StartSession(int objectID)
             _allowObjectID = objectID;
         }
 
-        internal void ThrowExceptionAndWarning(bool callerHasConnectionLock = false, bool asyncClose = false)
-        {
-            _parser.ThrowExceptionAndWarning(this, callerHasConnectionLock, asyncClose);
-        }
-
-        ////////////////////////////////////////////
-        // TDS Packet/buffer manipulation methods //
-        ////////////////////////////////////////////
-
-        internal Task ExecuteFlush()
-        {
-            lock (this)
-            {
-                if (_cancelled && 1 == _outputPacketNumber)
-                {
-                    ResetBuffer();
-                    _cancelled = false;
-                    throw SQL.OperationCancelled();
-                }
-                else
-                {
-                    Task writePacketTask = WritePacket(TdsEnums.HARDFLUSH);
-                    if (writePacketTask == null)
-                    {
-                        _pendingData = true;
-                        _messageStatus = 0;
-                        return null;
-                    }
-                    else
-                    {
-                        return AsyncHelper.CreateContinuationTask(writePacketTask, () => { _pendingData = true; _messageStatus = 0; });
-                    }
-                }
-            }
-        }
-
         // Processes the tds header that is present in the buffer
         internal bool TryProcessHeader()
         {

diff --git a/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/TdsParserStateObject.cs b/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/TdsParserStateObject.cs
@@ -701,5 +701,50 @@ internal void SetTimeoutMilliseconds(long timeout)
                 _timeoutTime = 0;
             }
         }
+
+        internal void ThrowExceptionAndWarning(bool callerHasConnectionLock = false, bool asyncClose = false)
+        {
+            _parser.ThrowExceptionAndWarning(this, callerHasConnectionLock, asyncClose);
+        }
+
+        ////////////////////////////////////////////
+        // TDS Packet/buffer manipulation methods //
+        ////////////////////////////////////////////
+
+        internal Task ExecuteFlush()
+        {
+            lock (this)
+            {
+                if (_cancelled && 1 == _outputPacketNumber)
+                {
+                    ResetBuffer();
+                    _cancelled = false;
+                    throw SQL.OperationCancelled();
+                }
+                else
+                {
+                    Task writePacketTask = WritePacket(TdsEnums.HARDFLUSH);
+                    if (writePacketTask == null)
+                    {
+                        HasPendingData = true;
+                        _messageStatus = 0;
+                        return null;
+                    }
+                    else
+                    {
+                        return AsyncHelper.CreateContinuationTaskWithState(
+                            task: writePacketTask,
+                            state: this,
+                            onSuccess: static (object state) =>
+                            {
+                                TdsParserStateObject stateObject = (TdsParserStateObject)state;
+                                stateObject.HasPendingData = true;
+                                stateObject._messageStatus = 0;
+                            }
+                        );
+                    }
+                }
+            }
+        }
     }
 }