Flatten | SqlInternalConnectionTds and SqlInternalConnection

src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/Connection/CachedContexts.cs

@@ -41,6 +41,11 @@ internal class CachedContexts
         /// </summary>
         private SqlDataReader.ReadAsyncCallContext? _dataReaderReadAsyncContext;
 
+        /// <summary>
+        /// Stores a data reader snapshot.
+        /// </summary>
+        private SqlDataReader.Snapshot? _dataReaderSnapshot;
+
         #endregion
 
         #region Access Methods
@@ -80,6 +85,13 @@ internal class CachedContexts
         internal SqlDataReader.IsDBNullAsyncCallContext? ClearDataReaderIsDbNullContext() =>
             Interlocked.Exchange(ref _dataReaderIsDbNullContext, null);
 
+        /// <summary>
+        /// Removes and returns the cached data reader snapshot.
+        /// </summary>
+        /// <returns>The previously cached snapshot or null when empty.</returns>
+        internal SqlDataReader.Snapshot? ClearDataReaderSnapshot() =>
+            Interlocked.Exchange(ref _dataReaderSnapshot, null);
+
         /// <summary>
         /// Attempts to cache the provided ExecuteNonQueryAsync context.
         /// </summary>
@@ -130,6 +142,16 @@ internal bool TrySetDataReaderReadAsyncContext(SqlDataReader.ReadAsyncCallContex
         internal bool TrySetDataReaderIsDbNullContext(SqlDataReader.IsDBNullAsyncCallContext value) =>
             TrySetContext(value, ref _dataReaderIsDbNullContext);
 
+        /// <summary>
+        /// Attempts to cache the provided data reader snapshot context.
+        /// </summary>
+        /// <param name="value">Context instance to store.</param>
+        /// <returns>
+        /// True when the snapshot is cached; false if an existing snapshot is preserved.
+        /// </returns>
+        internal bool TrySetDataReaderSnapshot(SqlDataReader.Snapshot value) =>
+            TrySetContext(value, ref _dataReaderSnapshot);
+
         #endregion
 
         private static bool TrySetContext<TContext>(TContext value, ref TContext? location)

src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlDataReader.cs

@@ -4101,10 +4101,11 @@ internal TdsOperationStatus TryReadColumnInternal(int i, bool readHeaderOnly = f
                 {
                     // reset snapshot to save memory use.  We can safely do that here because all SqlDataReader values are stable.
                     // The retry logic can use the current values to get back to the right state.
-                    if (_connection?.InnerConnection is SqlInternalConnectionTds sqlInternalConnection && sqlInternalConnection.CachedDataReaderSnapshot is null)
+                    if (_connection?.InnerConnection is SqlInternalConnectionTds sqlInternalConnection)
                     {
-                        sqlInternalConnection.CachedDataReaderSnapshot = _snapshot;
+                        sqlInternalConnection.CachedContexts.TrySetDataReaderSnapshot(_snapshot);
                     }
+
                     _snapshot = null;
                     PrepareAsyncInvocation(useSnapshot: true);
                 }
@@ -5085,10 +5086,11 @@ private static Task<bool> ReadAsyncExecute(Task task, object state)
                     if (!hasReadRowToken)
                     {
                         hasReadRowToken = true;
-                        if (reader.Connection?.InnerConnection is SqlInternalConnectionTds sqlInternalConnection && sqlInternalConnection.CachedDataReaderSnapshot is null)
+                        if (reader.Connection?.InnerConnection is SqlInternalConnectionTds sqlInternalConnection)
                         {
-                            sqlInternalConnection.CachedDataReaderSnapshot = reader._snapshot;
+                            sqlInternalConnection.CachedContexts.TrySetDataReaderSnapshot(reader._snapshot);
                         }
+
                         reader._snapshot = null;
                         reader.PrepareAsyncInvocation(useSnapshot: true);
                     }
@@ -5794,7 +5796,7 @@ private void PrepareAsyncInvocation(bool useSnapshot)
                 {
                     if (_connection?.InnerConnection is SqlInternalConnectionTds sqlInternalConnection)
                     {
-                        _snapshot = Interlocked.Exchange(ref sqlInternalConnection.CachedDataReaderSnapshot, null) ?? new Snapshot();
+                        _snapshot = sqlInternalConnection.CachedContexts.ClearDataReaderSnapshot() ?? new Snapshot();
                     }
                     else
                     {
@@ -5872,10 +5874,11 @@ private void CleanupAfterAsyncInvocationInternal(TdsParserStateObject stateObj,
             stateObj._permitReplayStackTraceToDiffer = false;
 #endif
 
-            if (_connection?.InnerConnection is SqlInternalConnectionTds sqlInternalConnection && sqlInternalConnection.CachedDataReaderSnapshot is null)
+            if (_connection?.InnerConnection is SqlInternalConnectionTds sqlInternalConnection)
             {
-                sqlInternalConnection.CachedDataReaderSnapshot = _snapshot;
+                sqlInternalConnection.CachedContexts.TrySetDataReaderSnapshot(_snapshot);
             }
+
             // We are setting this to null inside the if-statement because stateObj==null means that the reader hasn't been initialized or has been closed (either way _snapshot should already be null)
             _snapshot = null;
         }
@@ -5914,10 +5917,11 @@ private void SwitchToAsyncWithoutSnapshot()
             Debug.Assert(_snapshot != null, "Should currently have a snapshot");
             Debug.Assert(_stateObj != null && !_stateObj._asyncReadWithoutSnapshot, "Already in async without snapshot");
 
-            if (_connection?.InnerConnection is SqlInternalConnectionTds sqlInternalConnection && sqlInternalConnection.CachedDataReaderSnapshot is null)
+            if (_connection?.InnerConnection is SqlInternalConnectionTds sqlInternalConnection)
             {
-                sqlInternalConnection.CachedDataReaderSnapshot = _snapshot;
+                sqlInternalConnection.CachedContexts.TrySetDataReaderSnapshot(_snapshot);
             }
+
             _snapshot = null;
             _stateObj.ResetSnapshot();
             _stateObj._asyncReadWithoutSnapshot = true;

src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlInternalConnection.cs

@@ -28,8 +28,6 @@ internal abstract class SqlInternalConnection : DbConnectionInternal
         /// </summary>
         protected static readonly Guid s_globalTransactionTMID = new("1c742caf-6680-40ea-9c26-6b6846079764");
 
-        internal SqlDataReader.Snapshot CachedDataReaderSnapshot;
-
         /// <summary>
         /// Constructs a new SqlInternalConnection object using the provided connection options.
         /// </summary>