Fix managed SNI MARS receive teardown race

src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/ManagedSni/SniMarsConnection.netcore.cs

@@ -130,22 +130,35 @@ public uint ReceiveAsync(ref SniPacket packet)
         {
             using (SqlClientSNIEventScope.Create(nameof(SniMarsConnection)))
             {
-                if (packet != null)
+                lock (DemuxerSync)
                 {
-                    ReturnPacket(packet);
+                    SniPacket previousPacket = packet;
+                    packet = null;
+                    try
+                    {
+                        var response = _lowerHandle.ReceiveAsync(ref packet);
+                        if (response != TdsEnums.SNI_SUCCESS_IO_PENDING && packet == null)
+                        {
+                            // Immediate transport errors release the new packet. Keep the consumed
+                            // packet alive so HandleReceiveError can notify all sessions and release it.
+                            packet = previousPacket;
+                            previousPacket = null;
+                        }
 #if DEBUG
-                    SqlClientEventSource.Log.TrySNITraceEvent(nameof(SniMarsConnection), EventType.INFO, "MARS Session Id {0}, Packet {1} returned", args0: ConnectionId, args1: packet?._id);
+                        SqlClientEventSource.Log.TrySNITraceEvent(nameof(SniMarsConnection), EventType.INFO, "MARS Session Id {0}, Received new packet {1}", args0: ConnectionId, args1: packet?._id);
 #endif
-                    packet = null;
-                }
-
-                lock (DemuxerSync)
-                {
-                    var response = _lowerHandle.ReceiveAsync(ref packet);
+                        return response;
+                    }
+                    finally
+                    {
+                        if (previousPacket != null)
+                        {
+                            ReturnPacket(previousPacket);
 #if DEBUG
-                    SqlClientEventSource.Log.TrySNITraceEvent(nameof(SniMarsConnection), EventType.INFO, "MARS Session Id {0}, Received new packet {1}", args0: ConnectionId, args1: packet?._id);
+                            SqlClientEventSource.Log.TrySNITraceEvent(nameof(SniMarsConnection), EventType.INFO, "MARS Session Id {0}, Packet {1} returned", args0: ConnectionId, args1: previousPacket._id);
 #endif
-                    return response;
+                        }
+                    }
                 }
             }
         }
@@ -170,7 +183,14 @@ public uint CheckConnection()
         /// </summary>
         public void HandleReceiveError(SniPacket packet)
         {
-            Debug.Assert(Monitor.IsEntered(this), "HandleReceiveError was called without being locked.");
+            Debug.Assert(Monitor.IsEntered(DemuxerSync), "HandleReceiveError was called without being locked.");
+            if (_dataBytesLeft > 0 && _currentPacket != null)
+            {
+                ReturnPacket(_currentPacket);
+                _currentPacket = null;
+                _dataBytesLeft = 0;
+            }
+
             foreach (SniMarsHandle handle in _sessions.Values)
             {
                 if (packet.HasAsyncIOCompletionCallback)

src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/ManagedSni/SniMarsHandle.netcore.cs

@@ -345,14 +345,24 @@ public void HandleReceiveError(SniPacket packet)
                 // which should handle ownership of the packet because the individual mars handles are not aware of
                 // each other and cannot know if they are the last one in the list and that it is safe to return the packet
 
+                bool notifyAsyncReceive;
                 lock (_receivedPacketQueue)
                 {
                     _connectionError = SniLoadHandle.LastError;
                     SqlClientEventSource.Log.TrySNITraceEvent(nameof(SniMarsHandle), EventType.ERR, "MARS Session Id {0}, _connectionError to be handled: {1}", args0: ConnectionId, args1: _connectionError);
                     _packetEvent.Set();
+                    notifyAsyncReceive = _asyncReceives > 0;
+                    if (notifyAsyncReceive)
+                    {
+                        _asyncReceives--;
+                    }
                 }
 
-                ((TdsParserStateObject)_callbackObject).ReadAsyncCallback(PacketHandle.FromManagedPacket(packet), 1);
+                // Sync and idle sessions observe _connectionError without an async callback.
+                if (notifyAsyncReceive)
+                {
+                    ((TdsParserStateObject)_callbackObject).ReadAsyncCallback(PacketHandle.FromManagedPacket(packet), 1);
+                }
             }
         }
 

src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/ManagedSni/SniNpHandle.netcore.cs

@@ -195,7 +195,7 @@ public override uint Receive(out SniPacket packet, int timeout)
                     try
                     {
                         packet = RentPacket(headerSize: 0, dataSize: _bufferSize);
-                        packet.ReadFromStream(_stream);
+                        packet.ReadFromStream(_stream ?? throw new ObjectDisposedException(nameof(SniNpHandle)));
                         SqlClientEventSource.Log.TrySNITraceEvent(nameof(SniNpHandle), EventType.INFO, "Connection Id {0}, Rented and read packet, dataLeft {1}", args0: _connectionId, args1: packet?.DataLeft);
 
                         if (packet.Length == 0)
@@ -235,7 +235,9 @@ public override uint ReceiveAsync(ref SniPacket packet)
                 packet.SetAsyncIOCompletionCallback(_receiveCallback);
                 try
                 {
-                    packet.ReadFromStreamAsync(_stream);
+                    // Capture once: Dispose can clear the field while a MARS receive is re-armed.
+                    Stream stream = _stream ?? throw new ObjectDisposedException(nameof(SniNpHandle));
+                    packet.ReadFromStreamAsync(stream);
                     SqlClientEventSource.Log.TrySNITraceEvent(nameof(SniNpHandle), EventType.INFO, "Connection Id {0}, Rented and read packet asynchronously, dataLeft {1}", args0: _connectionId, args1: packet?.DataLeft);
                     return TdsEnums.SNI_SUCCESS_IO_PENDING;
                 }

src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/ManagedSni/SniTcpHandle.netcore.cs

@@ -890,6 +890,11 @@ public override uint Receive(out SniPacket packet, int timeoutInMilliseconds)
                     return ReportTcpSNIError(0, SniCommon.ConnOpenFailedError, Strings.SNI_ERROR_10);
                 }
 
+                if (_stream == null)
+                {
+                    return ReportTcpSNIError(new ObjectDisposedException(nameof(SniTcpHandle)));
+                }
+
                 try
                 {
                     // TODO: convert these to async versions that accept a cancellation token
@@ -1022,7 +1027,9 @@ public override uint ReceiveAsync(ref SniPacket packet)
             packet.SetAsyncIOCompletionCallback(_receiveCallback);
             try
             {
-                packet.ReadFromStreamAsync(_stream);
+                // Capture once: Dispose can clear the field while a MARS receive is re-armed.
+                Stream stream = _stream ?? throw new ObjectDisposedException(nameof(SniTcpHandle));
+                packet.ReadFromStreamAsync(stream);
                 SqlClientEventSource.Log.TrySNITraceEvent(nameof(SniTcpHandle), EventType.INFO, "Connection Id {0}, Data received from stream asynchronously", args0: _connectionId);
                 return TdsEnums.SNI_SUCCESS_IO_PENDING;
             }

src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/TdsParserStateObject.cs

@@ -3408,7 +3408,10 @@ internal void SendAttention(bool mustTakeWriteLock = false, bool asyncClose = fa
         internal TdsOperationStatus TryReadNetworkPacket()
         {
 #if DEBUG
-            Debug.Assert(!_shouldHaveEnoughData || _attentionSent, "Caller said there should be enough data, but we are currently reading a packet");
+            // Teardown can invalidate the reader's buffered-data estimate before this call.
+            Debug.Assert(!_shouldHaveEnoughData || _attentionSent ||
+                _parser.State == TdsParserState.Closed || _parser.State == TdsParserState.Broken,
+                "Caller said there should be enough data, but we are currently reading a packet");
 #endif
             TdsOperationStatus result = TdsOperationStatus.InvalidData;
             if (_snapshot != null)
@@ -3635,7 +3638,14 @@ public void ReadAsyncCallback(IntPtr key, PacketHandle packet, uint error)
             try
             {
 #if NET
-                Debug.Assert((packet.Type == 0 && PartialPacketContainsCompletePacket()) || (CheckPacket(packet, source) && source != null), "AsyncResult null on callback");
+                // Closing a MARS connection can complete and clear a session's task before
+                // its pending error callback arrives. It must still retire the callback count.
+                Debug.Assert(
+                    (packet.Type == 0 && PartialPacketContainsCompletePacket()) ||
+                    (CheckPacket(packet, source) && source != null) ||
+                    (source == null && error != 0 && _parser.MARSOn &&
+                        (_parser.State == TdsParserState.Closed || _parser.State == TdsParserState.Broken)),
+                    "AsyncResult null on callback");
 #else
                 Debug.Assert(CheckPacket(packet, source), "AsyncResult null on callback");
 #endif

src/Microsoft.Data.SqlClient/tests/ManualTests/SQL/AsyncTest/MarsReceiveTeardownTest.cs

@@ -0,0 +1,187 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+// See the LICENSE file in the project root for more information.
+
+#if NET
+
+#nullable enable
+
+using System;
+using System.Collections.Concurrent;
+using System.IO;
+using System.Threading.Tasks;
+using Xunit;
+
+namespace Microsoft.Data.SqlClient.ManualTesting.Tests
+{
+    /// <summary>
+    /// Exercises GH#4679 with a streaming MARS session and a connection-terminating error
+    /// on a second session. Requires a test SQL Server and a sysadmin login; each case
+    /// writes a severity-20 error to the server log.
+    /// </summary>
+    [Trait("Set", "1")]
+    [Collection(nameof(MarsReceiveTeardownCollection))]
+    public sealed class MarsReceiveTeardownTest
+    {
+        /// <summary>
+        /// Physical teardown must report normal command errors rather than faulting an
+        /// unobserved MARS receive continuation, with or without connection pooling.
+        /// Each delay is a separate case rather than a long-running stress loop.
+        /// </summary>
+        [ConditionalTheory(typeof(DataTestUtility), nameof(DataTestUtility.AreConnStringsSetup),
+            nameof(DataTestUtility.IsNotAzureServer), nameof(DataTestUtility.IsUsingManagedSNI),
+            nameof(DataTestUtility.IsSysAdmin))]
+        [InlineData(false, false, 1)]
+        [InlineData(false, true, 1)]
+        [InlineData(true, false, 1)]
+        [InlineData(true, true, 1)]
+        [InlineData(false, false, 29)]
+        [InlineData(false, true, 29)]
+        [InlineData(true, false, 29)]
+        [InlineData(true, true, 29)]
+        public async Task FatalError_WhileMarsReaderStreams_DoesNotFaultReceivePump(bool async, bool pooling, int delay)
+        {
+            string connectionString = new SqlConnectionStringBuilder(DataTestUtility.TCPConnectionString)
+            {
+                MultipleActiveResultSets = true,
+                Pooling = pooling,
+                ConnectRetryCount = 0,
+                ApplicationName = nameof(MarsReceiveTeardownTest),
+            }.ConnectionString;
+            ConcurrentQueue<Exception> unobserved = new();
+            EventHandler<UnobservedTaskExceptionEventArgs> handler = (_, args) =>
+            {
+                foreach (Exception exception in args.Exception.Flatten().InnerExceptions)
+                {
+                    if (exception.StackTrace?.Contains("ManagedSni.SniMarsConnection") == true)
+                    {
+                        unobserved.Enqueue(exception);
+                        args.SetObserved();
+                    }
+                }
+            };
+
+            TaskScheduler.UnobservedTaskException += handler;
+            try
+            {
+                await TerminateStreamingConnection(connectionString, async, delay);
+
+                for (int attempt = 0; attempt < 10; attempt++)
+                {
+                    CollectReceiveContinuations();
+                    await Task.Delay(50);
+                }
+
+                Assert.Empty(unobserved);
+            }
+            finally
+            {
+                TaskScheduler.UnobservedTaskException -= handler;
+                using SqlConnection poolKey = new(connectionString);
+                SqlConnection.ClearPool(poolKey);
+            }
+        }
+
+        /// <summary>
+        /// Starts two MARS sessions, terminates only their connection, and observes all
+        /// application tasks so any unobserved exception belongs to the driver's receive pump.
+        /// </summary>
+        /// <param name="connectionString">Connection string for the test server.</param>
+        /// <param name="async">Whether to use asynchronous command and reader APIs.</param>
+        /// <param name="delay">Delay before terminating the streaming connection, in milliseconds.</param>
+        /// <returns>A task that completes after both sessions and the connection are cleaned up.</returns>
+        private static async Task TerminateStreamingConnection(string connectionString, bool async, int delay)
+        {
+            SqlConnection connection = new(connectionString);
+            SqlDataReader? reader = null;
+            Task<Exception?>? consumer = null;
+            using SqlCommand stream = new(
+                "SELECT TOP (200000) a.object_id, b.name, REPLICATE('x', 200) AS pad " +
+                "FROM sys.all_objects a CROSS JOIN sys.all_columns b", connection);
+            stream.CommandTimeout = 60;
+            try
+            {
+                if (async)
+                {
+                    await connection.OpenAsync();
+                }
+                else
+                {
+                    connection.Open();
+                }
+
+                reader = async ? await stream.ExecuteReaderAsync() : stream.ExecuteReader();
+                Assert.True(async ? await reader.ReadAsync() : reader.Read());
+                consumer = Task.Run(() => Record.ExceptionAsync(async () =>
+                {
+                    while (async ? await reader.ReadAsync() : reader.Read())
+                    {
+                    }
+                }));
+                await Task.Delay(delay);
+
+                using SqlCommand terminate = new(
+                    "RAISERROR('SqlClient MARS receive teardown regression', 20, 1) WITH LOG", connection);
+                SqlException error;
+                if (async)
+                {
+                    error = await Assert.ThrowsAsync<SqlException>(() => terminate.ExecuteNonQueryAsync());
+                }
+                else
+                {
+                    error = Assert.Throws<SqlException>(() => terminate.ExecuteNonQuery());
+                }
+                Assert.True(error.Class >= 20, "The second session must terminate the physical connection.");
+            }
+            finally
+            {
+                Exception? closeError;
+                Exception? readerError;
+                try
+                {
+                    Exception? consumeError = consumer is null
+                        ? null
+                        : await consumer.WaitAsync(TimeSpan.FromSeconds(65));
+                    AssertExpectedTeardownError(consumeError);
+                }
+                finally
+                {
+                    closeError = Record.Exception(connection.Dispose);
+                    readerError = reader is null ? null : Record.Exception(reader.Dispose);
+                }
+                AssertExpectedTeardownError(closeError);
+                AssertExpectedTeardownError(readerError);
+            }
+        }
+
+        /// <summary>
+        /// Allows expected failures of commands on a broken connection, but rejects programming errors.
+        /// </summary>
+        /// <param name="exception">The observed application-side teardown exception, if any.</param>
+        private static void AssertExpectedTeardownError(Exception? exception)
+        {
+            Assert.True(exception is null or SqlException or InvalidOperationException or IOException,
+                $"Unexpected application-side teardown error: {exception}");
+        }
+
+        /// <summary>
+        /// Finalizes abandoned continuation tasks so their failures reach the event handler.
+        /// </summary>
+        private static void CollectReceiveContinuations()
+        {
+            GC.Collect();
+            GC.WaitForPendingFinalizers();
+            GC.Collect();
+        }
+    }
+
+    /// <summary>
+    /// Isolates the process-wide unobserved-exception handler and forced garbage collections.
+    /// </summary>
+    [CollectionDefinition(nameof(MarsReceiveTeardownCollection), DisableParallelization = true)]
+    public sealed class MarsReceiveTeardownCollection
+    {
+    }
+}
+
+#endif