From b13294343fa8faf0756d3b23fabf689cf0f0b1a5 Mon Sep 17 00:00:00 2001
From: Edward Neal <55035479+edwardneal@users.noreply.github.com>
Date: Wed, 2 Apr 2025 20:46:11 +0100
Subject: [PATCH] Port PR 499

---
 .../Microsoft/Data/SqlClient/SqlDataReader.cs | 62 ++++++++++++++++---
 .../Data/SqlClient/SqlInternalConnection.cs   |  3 +-
 2 files changed, 53 insertions(+), 12 deletions(-)

diff --git a/src/Microsoft.Data.SqlClient/netfx/src/Microsoft/Data/SqlClient/SqlDataReader.cs b/src/Microsoft.Data.SqlClient/netfx/src/Microsoft/Data/SqlClient/SqlDataReader.cs
index 645ba1edc7..71a2b49613 100644
--- a/src/Microsoft.Data.SqlClient/netfx/src/Microsoft/Data/SqlClient/SqlDataReader.cs
+++ b/src/Microsoft.Data.SqlClient/netfx/src/Microsoft/Data/SqlClient/SqlDataReader.cs
@@ -30,7 +30,7 @@ namespace Microsoft.Data.SqlClient
     /// <include file='../../../../../../../doc/snippets/Microsoft.Data.SqlClient/SqlDataReader.xml' path='docs/members[@name="SqlDataReader"]/SqlDataReader/*' />
     public class SqlDataReader : DbDataReader, IDataReader, IDbColumnSchemaGenerator
     {
-        private enum ALTROWSTATUS
+        internal enum ALTROWSTATUS
         {
             Null = 0,           // default and after Done
             AltRow,             // after calling NextResult and the first AltRow is available for read
@@ -97,9 +97,6 @@ internal class SharedState
         private SqlSequentialStream _currentStream;
         private SqlSequentialTextReader _currentTextReader;
 
-        private IsDBNullAsyncCallContext _cachedIsDBNullContext;
-        private ReadAsyncCallContext _cachedReadAsyncContext;
-
         internal SqlDataReader(SqlCommand command, CommandBehavior behavior)
         {
             SqlConnection.VerifyExecutePermission();
@@ -4387,6 +4384,10 @@ internal TdsOperationStatus TryReadColumnInternal(int i, bool readHeaderOnly/* =
                 {
                     // 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 SqlInternalConnection sqlInternalConnection && sqlInternalConnection.CachedDataReaderSnapshot is null)
+                    {
+                        sqlInternalConnection.CachedDataReaderSnapshot = _snapshot;
+                    }
                     _snapshot = null;
                     PrepareAsyncInvocation(useSnapshot: true);
                 }
@@ -5318,7 +5319,15 @@ public override Task<bool> ReadAsync(CancellationToken cancellationToken)
                     return source.Task;
                 }
 
-                var context = Interlocked.Exchange(ref _cachedReadAsyncContext, null) ?? new ReadAsyncCallContext();
+                ReadAsyncCallContext context = null;
+                if (_connection?.InnerConnection is SqlInternalConnection sqlInternalConnection)
+                {
+                    context = Interlocked.Exchange(ref sqlInternalConnection.CachedDataReaderReadAsyncContext, null);
+                }
+                if (context is null)
+                {
+                    context = new ReadAsyncCallContext();
+                }
 
                 Debug.Assert(context.Reader == null && context.Source == null && context.Disposable == default, "cached ReadAsyncCallContext was not properly disposed");
 
@@ -5358,6 +5367,10 @@ private static Task<bool> ReadAsyncExecute(Task task, object state)
                     if (!hasReadRowToken)
                     {
                         hasReadRowToken = true;
+                        if (reader.Connection?.InnerConnection is SqlInternalConnection sqlInternalConnection && sqlInternalConnection.CachedDataReaderSnapshot is null)
+                        {
+                            sqlInternalConnection.CachedDataReaderSnapshot = reader._snapshot;
+                        }
                         reader._snapshot = null;
                         reader.PrepareAsyncInvocation(useSnapshot: true);
                     }
@@ -5377,7 +5390,10 @@ private static Task<bool> ReadAsyncExecute(Task task, object state)
 
         private void SetCachedReadAsyncCallContext(ReadAsyncCallContext instance)
         {
-            Interlocked.CompareExchange(ref _cachedReadAsyncContext, instance, null);
+            if (_connection?.InnerConnection is SqlInternalConnection sqlInternalConnection)
+            {
+                Interlocked.CompareExchange(ref sqlInternalConnection.CachedDataReaderReadAsyncContext, instance, null);
+            }
         }
 
         /// <include file='../../../../../../../doc/snippets/Microsoft.Data.SqlClient/SqlDataReader.xml' path='docs/members[@name="SqlDataReader"]/IsDBNullAsync/*' />
@@ -5479,7 +5495,15 @@ override public Task<bool> IsDBNullAsync(int i, CancellationToken cancellationTo
                         registrationHolder.Set(cancellationToken.Register(SqlCommand.s_cancelIgnoreFailure, _command));
                     }
 
-                    IsDBNullAsyncCallContext context = Interlocked.Exchange(ref _cachedIsDBNullContext, null) ?? new IsDBNullAsyncCallContext();
+                    IsDBNullAsyncCallContext context = null;
+                    if (_connection?.InnerConnection is SqlInternalConnection sqlInternalConnection)
+                    {
+                        context = Interlocked.Exchange(ref sqlInternalConnection.CachedDataReaderIsDBNullContext, null);
+                    }
+                    if (context is null)
+                    {
+                        context = new IsDBNullAsyncCallContext();
+                    }
 
                     Debug.Assert(context.Reader == null && context.Source == null && context.Disposable == default, "cached ISDBNullAsync context not properly disposed");
 
@@ -5517,7 +5541,10 @@ private static Task<bool> IsDBNullAsyncExecute(Task task, object state)
 
         private void SetCachedIDBNullAsyncCallContext(IsDBNullAsyncCallContext instance)
         {
-            Interlocked.CompareExchange(ref _cachedIsDBNullContext, instance, null);
+            if (_connection?.InnerConnection is SqlInternalConnection sqlInternalConnection)
+            {
+                Interlocked.CompareExchange(ref sqlInternalConnection.CachedDataReaderIsDBNullContext, instance, null);
+            }
         }
 
         /// <include file='../../../../../../../doc/snippets/Microsoft.Data.SqlClient/SqlDataReader.xml' path='docs/members[@name="SqlDataReader"]/GetFieldValueAsync/*' />
@@ -6019,7 +6046,7 @@ private void CompleteAsyncCall<T>(Task<T> task, SqlDataReaderBaseAsyncCallContex
             }
         }
 
-        private sealed class Snapshot
+        internal sealed class Snapshot
         {
             public bool _dataReady;
             public bool _haltRead;
@@ -6051,7 +6078,14 @@ private void PrepareAsyncInvocation(bool useSnapshot)
 
                 if (_snapshot == null)
                 {
-                    _snapshot = new Snapshot();
+                    if (_connection?.InnerConnection is SqlInternalConnection sqlInternalConnection)
+                    {
+                        _snapshot = Interlocked.Exchange(ref sqlInternalConnection.CachedDataReaderSnapshot, null) ?? new Snapshot();
+                    }
+                    else
+                    {
+                        _snapshot = new Snapshot();
+                    }
 
                     _snapshot._dataReady = _sharedState._dataReady;
                     _snapshot._haltRead = _haltRead;
@@ -6124,6 +6158,10 @@ private void CleanupAfterAsyncInvocationInternal(TdsParserStateObject stateObj,
             stateObj._permitReplayStackTraceToDiffer = false;
 #endif
 
+            if (_connection?.InnerConnection is SqlInternalConnection sqlInternalConnection && sqlInternalConnection.CachedDataReaderSnapshot is null)
+            {
+                sqlInternalConnection.CachedDataReaderSnapshot = _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;
         }
@@ -6162,6 +6200,10 @@ 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 SqlInternalConnection sqlInternalConnection && sqlInternalConnection.CachedDataReaderSnapshot is null)
+            {
+                sqlInternalConnection.CachedDataReaderSnapshot = _snapshot;
+            }
             _snapshot = null;
             _stateObj.ResetSnapshot();
             _stateObj._asyncReadWithoutSnapshot = true;
diff --git a/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlInternalConnection.cs b/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlInternalConnection.cs
index d8950280ed..a39c83d831 100644
--- a/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlInternalConnection.cs
+++ b/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlInternalConnection.cs
@@ -31,11 +31,10 @@ internal abstract class SqlInternalConnection : DbConnectionInternal
         internal SqlCommand.ExecuteReaderAsyncCallContext CachedCommandExecuteReaderAsyncContext;
         internal SqlCommand.ExecuteNonQueryAsyncCallContext CachedCommandExecuteNonQueryAsyncContext;
         internal SqlCommand.ExecuteXmlReaderAsyncCallContext CachedCommandExecuteXmlReaderAsyncContext;
-
+#endif
         internal SqlDataReader.Snapshot CachedDataReaderSnapshot;
         internal SqlDataReader.IsDBNullAsyncCallContext CachedDataReaderIsDBNullContext;
         internal SqlDataReader.ReadAsyncCallContext CachedDataReaderReadAsyncContext;
-#endif
 
         // if connection is not open: null
         // if connection is open: currently active database