Merge TDSParserStateObject.StateSnapshot creation and clearing (#2144)

dotnet · Sep 13, 2023 · 1e3cbfc · 1e3cbfc
1 parent 93d4398
commit 1e3cbfc
Show file tree

Hide file tree

Showing 3 changed files with 66 additions and 59 deletions.
diff --git a/...osoft.Data.SqlClient/netcore/src/Microsoft/Data/SqlClient/TdsParserStateObject.netcore.cs b/...osoft.Data.SqlClient/netcore/src/Microsoft/Data/SqlClient/TdsParserStateObject.netcore.cs
@@ -21,7 +21,6 @@ internal abstract partial class TdsParserStateObject
         private readonly WeakReference _cancellationOwner = new WeakReference(null);
 
 		// Async
-        private StateSnapshot _cachedSnapshot;
         private SnapshottedStateFlags _snapshottedState;
 
         //////////////////
@@ -1107,34 +1106,6 @@ internal bool TrySkipBytes(int num)
         // Network/Packet Reading & Processing //
         /////////////////////////////////////////
 
-        internal void SetSnapshot()
-        {
-            StateSnapshot snapshot = _snapshot;
-            if (snapshot is null)
-            {
-                snapshot = Interlocked.Exchange(ref _cachedSnapshot, null) ?? new StateSnapshot();
-            }
-            else
-            {
-                snapshot.Clear();
-            }
-            _snapshot = snapshot;
-            _snapshot.Snap(this);
-            _snapshotReplay = false;
-        }
-
-        internal void ResetSnapshot()
-        {
-            if (_snapshot != null)
-            {
-                StateSnapshot snapshot = _snapshot;
-                _snapshot = null;
-                snapshot.Clear();
-                Interlocked.CompareExchange(ref _cachedSnapshot, snapshot, null);
-            }
-            _snapshotReplay = false;
-        }
-
 #if DEBUG
         private string _lastStack;
 #endif
@@ -3291,24 +3262,13 @@ internal void Clear()
                 PacketData packet = _snapshotInBuffList;
                 _snapshotInBuffList = null;
                 _snapshotInBuffCount = 0;
-                _snapshotInBuffCurrent = 0;
-                _snapshotInBytesUsed = 0;
-                _snapshotInBytesPacket = 0;
-                _snapshotMessageStatus = 0;
-                _snapshotNullBitmapInfo = default;
-                _plpData = null;
-                _snapshotCleanupMetaData = null;
-                _snapshotCleanupAltMetaDataSetArray = null;
+
                 _state = SnapshottedStateFlags.None;
-#if DEBUG
-                _rollingPend = 0;
-                _rollingPendCount = 0;
-                _stateObj._lastStack = null;
-#endif
-                _stateObj = null;
 
                 packet.Clear();
                 _sparePacket = packet;
+
+                ClearCore();
             }
         }
     }

diff --git a/...Microsoft.Data.SqlClient/netfx/src/Microsoft/Data/SqlClient/TdsParserStateObject.netfx.cs b/...Microsoft.Data.SqlClient/netfx/src/Microsoft/Data/SqlClient/TdsParserStateObject.netfx.cs
@@ -1191,19 +1191,6 @@ internal bool TrySkipBytes(int num)
         // Network/Packet Reading & Processing //
         /////////////////////////////////////////
 
-        internal void SetSnapshot()
-        {
-            _snapshot = new StateSnapshot(this);
-            _snapshot.Snap();
-            _snapshotReplay = false;
-        }
-
-        internal void ResetSnapshot()
-        {
-            _snapshot = null;
-            _snapshotReplay = false;
-        }
-
 #if DEBUG
         StackTrace _lastStack;
 #endif
@@ -3249,10 +3236,9 @@ sealed partial class StateSnapshot
             private bool _snapshotReceivedColumnMetadata = false;
             private bool _snapshotAttentionReceived;
 
-            public StateSnapshot(TdsParserStateObject state)
+            public StateSnapshot()
             {
                 _snapshotInBuffs = new List<PacketData>();
-                _stateObj = state;
             }
 
 #if DEBUG
@@ -3315,8 +3301,9 @@ internal bool Replay()
                 return false;
             }
 
-            internal void Snap()
+            internal void Snap(TdsParserStateObject state)
             {
+                _stateObj = state;
                 _snapshotInBuffs.Clear();
                 _snapshotInBuffCurrent = 0;
                 _snapshotInBytesUsed = _stateObj._inBytesUsed;
@@ -3389,6 +3376,19 @@ internal void ResetSnapshotState()
 
                 _stateObj.AssertValidState();
             }
+
+            internal void Clear()
+            {
+                _snapshotInBuffs.Clear();
+
+                _snapshotPendingData = false;
+                _snapshotErrorTokenReceived = false;
+                _snapshotHasOpenResult = false;
+                _snapshotReceivedColumnMetadata = false;
+                _snapshotAttentionReceived = false;
+
+                ClearCore();
+            }
         }
     }
 }
diff --git a/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/TdsParserStateObject.cs b/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/TdsParserStateObject.cs
@@ -203,6 +203,7 @@ public TimeoutState(int value)
         internal bool _syncOverAsync = true;
         private bool _snapshotReplay;
         private StateSnapshot _snapshot;
+        private StateSnapshot _cachedSnapshot;
         internal ExecutionContext _executionContext;
         internal bool _asyncReadWithoutSnapshot;
 #if DEBUG
@@ -1102,6 +1103,34 @@ internal bool SetPacketSize(int size)
         }
         */
 
+        internal void SetSnapshot()
+        {
+            StateSnapshot snapshot = _snapshot;
+            if (snapshot is null)
+            {
+                snapshot = Interlocked.Exchange(ref _cachedSnapshot, null) ?? new StateSnapshot();
+            }
+            else
+            {
+                snapshot.Clear();
+            }
+            _snapshot = snapshot;
+            _snapshot.Snap(this);
+            _snapshotReplay = false;
+        }
+
+        internal void ResetSnapshot()
+        {
+            if (_snapshot != null)
+            {
+                StateSnapshot snapshot = _snapshot;
+                _snapshot = null;
+                snapshot.Clear();
+                Interlocked.CompareExchange(ref _cachedSnapshot, snapshot, null);
+            }
+            _snapshotReplay = false;
+        }
+
         sealed partial class StateSnapshot
         {
             private sealed class PLPData
@@ -1171,6 +1200,24 @@ internal void PrepareReplay()
             {
                 ResetSnapshotState();
             }
+
+            internal void ClearCore()
+            {
+                _snapshotInBuffCurrent = 0;
+                _snapshotInBytesUsed = 0;
+                _snapshotInBytesPacket = 0;
+                _snapshotMessageStatus = 0;
+                _snapshotNullBitmapInfo = default;
+                _plpData = null;
+                _snapshotCleanupMetaData = null;
+                _snapshotCleanupAltMetaDataSetArray = null;
+#if DEBUG
+                _rollingPend = 0;
+                _rollingPendCount = 0;
+                _stateObj._lastStack = null;
+#endif
+                _stateObj = null;
+            }
         }
     }
 }