diff --git a/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/ManagedSni/SniMarsConnection.netcore.cs b/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/ManagedSni/SniMarsConnection.netcore.cs
index 240e62666e..f2648a3e92 100644
--- a/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/ManagedSni/SniMarsConnection.netcore.cs
+++ b/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()
///
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)
diff --git a/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/ManagedSni/SniMarsHandle.netcore.cs b/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/ManagedSni/SniMarsHandle.netcore.cs
index 46063f9484..c3fd515407 100644
--- a/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/ManagedSni/SniMarsHandle.netcore.cs
+++ b/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);
+ }
}
}
diff --git a/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/ManagedSni/SniNpHandle.netcore.cs b/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/ManagedSni/SniNpHandle.netcore.cs
index e7ef8ad890..c1ae085b8d 100644
--- a/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/ManagedSni/SniNpHandle.netcore.cs
+++ b/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;
}
diff --git a/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/ManagedSni/SniTcpHandle.netcore.cs b/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/ManagedSni/SniTcpHandle.netcore.cs
index 176420db87..41bfa0dae7 100644
--- a/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/ManagedSni/SniTcpHandle.netcore.cs
+++ b/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;
}
diff --git a/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/TdsParserStateObject.cs b/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/TdsParserStateObject.cs
index 6d82040ca6..9d7dd539c0 100644
--- a/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/TdsParserStateObject.cs
+++ b/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
diff --git a/src/Microsoft.Data.SqlClient/tests/ManualTests/SQL/AsyncTest/MarsReceiveTeardownTest.cs b/src/Microsoft.Data.SqlClient/tests/ManualTests/SQL/AsyncTest/MarsReceiveTeardownTest.cs
new file mode 100644
index 0000000000..f6e6274d45
--- /dev/null
+++ b/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
+{
+ ///
+ /// 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.
+ ///
+ [Trait("Set", "1")]
+ [Collection(nameof(MarsReceiveTeardownCollection))]
+ public sealed class MarsReceiveTeardownTest
+ {
+ ///
+ /// 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.
+ ///
+ [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 unobserved = new();
+ EventHandler 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);
+ }
+ }
+
+ ///
+ /// Starts two MARS sessions, terminates only their connection, and observes all
+ /// application tasks so any unobserved exception belongs to the driver's receive pump.
+ ///
+ /// Connection string for the test server.
+ /// Whether to use asynchronous command and reader APIs.
+ /// Delay before terminating the streaming connection, in milliseconds.
+ /// A task that completes after both sessions and the connection are cleaned up.
+ private static async Task TerminateStreamingConnection(string connectionString, bool async, int delay)
+ {
+ SqlConnection connection = new(connectionString);
+ SqlDataReader? reader = null;
+ Task? 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(() => terminate.ExecuteNonQueryAsync());
+ }
+ else
+ {
+ error = Assert.Throws(() => 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);
+ }
+ }
+
+ ///
+ /// Allows expected failures of commands on a broken connection, but rejects programming errors.
+ ///
+ /// The observed application-side teardown exception, if any.
+ private static void AssertExpectedTeardownError(Exception? exception)
+ {
+ Assert.True(exception is null or SqlException or InvalidOperationException or IOException,
+ $"Unexpected application-side teardown error: {exception}");
+ }
+
+ ///
+ /// Finalizes abandoned continuation tasks so their failures reach the event handler.
+ ///
+ private static void CollectReceiveContinuations()
+ {
+ GC.Collect();
+ GC.WaitForPendingFinalizers();
+ GC.Collect();
+ }
+ }
+
+ ///
+ /// Isolates the process-wide unobserved-exception handler and forced garbage collections.
+ ///
+ [CollectionDefinition(nameof(MarsReceiveTeardownCollection), DisableParallelization = true)]
+ public sealed class MarsReceiveTeardownCollection
+ {
+ }
+}
+
+#endif
diff --git a/src/Microsoft.Data.SqlClient/tests/UnitTests/Microsoft/Data/SqlClient/ManagedSni/SniReceiveTests.cs b/src/Microsoft.Data.SqlClient/tests/UnitTests/Microsoft/Data/SqlClient/ManagedSni/SniReceiveTests.cs
new file mode 100644
index 0000000000..ca6e87d3bd
--- /dev/null
+++ b/src/Microsoft.Data.SqlClient/tests/UnitTests/Microsoft/Data/SqlClient/ManagedSni/SniReceiveTests.cs
@@ -0,0 +1,493 @@
+// 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
+
+using System;
+using System.IO;
+using System.IO.Pipes;
+using System.Net;
+using System.Net.Sockets;
+using System.Threading;
+using System.Threading.Tasks;
+using Microsoft.Data.ProviderBase;
+using Microsoft.Data.SqlClient.ManagedSni;
+using Xunit;
+
+namespace Microsoft.Data.SqlClient.UnitTests.ManagedSni
+{
+ ///
+ /// Covers managed TCP and named-pipe receive completion during physical connection teardown.
+ ///
+ public sealed class SniReceiveTests
+ {
+ ///
+ /// Closing the connection invalidates an async reader's buffered-data estimate.
+ /// Its next network read must report closure rather than fail a debug assertion.
+ ///
+ [Theory]
+ [InlineData(false)]
+ [InlineData(true)]
+ public void ReadNetworkPacket_AfterClose_ReportsClosedConnection(bool broken)
+ {
+ TdsParser parser = new(MARS: true, fAsynchronous: true);
+ TdsParserStateObjectManaged stateObject = new(parser);
+ try
+ {
+ parser.State = broken ? TdsParserState.Broken : TdsParserState.Closed;
+#if DEBUG
+ stateObject._shouldHaveEnoughData = true;
+#endif
+ Assert.Throws(() => stateObject.TryReadNetworkPacket());
+ }
+ finally
+ {
+ stateObject.Dispose();
+ parser._physicalStateObj.Dispose();
+ }
+ }
+
+ ///
+ /// A receive started after disposal must report an SNI error, not throw into its caller.
+ ///
+ [Theory]
+ [InlineData(false, false)]
+ [InlineData(false, true)]
+ [InlineData(true, false)]
+ [InlineData(true, true)]
+ public void Receive_AfterDispose_ReturnsError(bool namedPipe, bool async)
+ {
+ using LocalConnection connection = new(namedPipe);
+ connection.Handle.Dispose();
+ SniPacket? packet = null;
+ try
+ {
+ uint result = async
+ ? connection.Handle.ReceiveAsync(ref packet)
+ : connection.Handle.Receive(out packet, 1000);
+
+ Assert.Equal(TdsEnums.SNI_ERROR, result);
+ Assert.Null(packet);
+ Assert.Equal(namedPipe ? SniProviders.NP_PROV : SniProviders.TCP_PROV, SniLoadHandle.LastError.provider);
+ Assert.IsType(SniLoadHandle.LastError.exception);
+ }
+ finally
+ {
+ if (packet is not null && !packet.IsInvalid)
+ {
+ connection.Handle.ReturnPacket(packet);
+ }
+ }
+ }
+
+ ///
+ /// Re-arming a successful MARS receive after disposal must finish error handling
+ /// without throwing, including when the SMUX header or payload spans receives.
+ ///
+ [Theory]
+ [InlineData(false, false)]
+ [InlineData(false, true)]
+ [InlineData(true, false)]
+ [InlineData(true, true)]
+ public void MarsReceiveComplete_AfterDispose_HandlesReceiveError(bool namedPipe, bool partialPayload)
+ {
+ using LocalConnection connection = new(namedPipe);
+ SniMarsConnection mars = new(connection.Handle);
+ SniPacket packet = connection.Handle.RentPacket(0, SniSmuxHeader.HEADER_LENGTH);
+ packet.SetAsyncIOCompletionCallback(mars.HandleReceiveComplete);
+ if (partialPayload)
+ {
+ byte[] headerBytes = new byte[SniSmuxHeader.HEADER_LENGTH];
+ new SniSmuxHeader
+ {
+ SMID = 83,
+ flags = (byte)SniSmuxFlags.SMUX_DATA,
+ length = SniSmuxHeader.HEADER_LENGTH + 1,
+ }.Write(headerBytes);
+ packet.AppendData(headerBytes, headerBytes.Length);
+ }
+
+ connection.Handle.Dispose();
+ try
+ {
+ mars.HandleReceiveComplete(packet, TdsEnums.SNI_SUCCESS);
+ Assert.IsType(SniLoadHandle.LastError.exception);
+ Assert.True(packet.IsInvalid);
+ }
+ finally
+ {
+ if (!packet.IsInvalid)
+ {
+ connection.Handle.ReturnPacket(packet);
+ }
+ }
+ }
+
+ ///
+ /// A successful transport callback may run after disposal but must not fault when
+ /// the MARS demultiplexer tries to receive the rest of its header.
+ ///
+ [Theory]
+ [InlineData(false)]
+ [InlineData(true)]
+ public async Task MarsReceiveCallback_DisposesBeforeRearm_CompletesWithoutException(bool namedPipe)
+ {
+ using LocalConnection connection = new(namedPipe);
+ SniMarsConnection mars = new(connection.Handle);
+ TaskCompletionSource completion = new(TaskCreationOptions.RunContinuationsAsynchronously);
+ connection.Handle.SetAsyncCallbacks((packet, error) =>
+ {
+ // Force the race's ordering, rather than depending on thread scheduling and GC.
+ try
+ {
+ Assert.Equal(TdsEnums.SNI_SUCCESS, error);
+ connection.Handle.Dispose();
+ mars.HandleReceiveComplete(packet, error);
+ completion.SetResult(SniLoadHandle.LastError);
+ }
+ catch (Exception ex)
+ {
+ completion.SetException(ex);
+ }
+ }, null);
+
+ Assert.Equal(TdsEnums.SNI_SUCCESS_IO_PENDING, mars.StartReceive());
+ await connection.Peer.WriteAsync(new byte[] { 83 });
+ SniError receiveError = await completion.Task.WaitAsync(TimeSpan.FromSeconds(5));
+ Assert.IsType(receiveError.exception);
+ }
+
+ ///
+ /// Healthy sync and async receives must retain their data and packet ownership semantics.
+ ///
+ [Theory]
+ [InlineData(false, false)]
+ [InlineData(false, true)]
+ [InlineData(true, false)]
+ [InlineData(true, true)]
+ public async Task Receive_Connected_ReturnsData(bool namedPipe, bool async)
+ {
+ using LocalConnection connection = new(namedPipe);
+ TaskCompletionSource<(SniPacket Packet, uint Error)> completion =
+ new(TaskCreationOptions.RunContinuationsAsynchronously);
+ connection.Handle.SetAsyncCallbacks(
+ (received, error) => completion.SetResult((received, error)), null);
+ SniPacket? packet = null;
+ Task write = connection.Peer.WriteAsync(new byte[] { 42 }).AsTask();
+ try
+ {
+ uint result = async
+ ? connection.Handle.ReceiveAsync(ref packet)
+ : connection.Handle.Receive(out packet, 1000);
+ if (async)
+ {
+ Assert.Equal(TdsEnums.SNI_SUCCESS_IO_PENDING, result);
+ (packet, result) = await completion.Task.WaitAsync(TimeSpan.FromSeconds(5));
+ }
+
+ Assert.Equal(TdsEnums.SNI_SUCCESS, result);
+ Assert.NotNull(packet);
+ byte[] data = new byte[1];
+ Assert.Equal(1, packet.TakeData(data, 0, data.Length));
+ Assert.Equal(42, data[0]);
+ await write.WaitAsync(TimeSpan.FromSeconds(5));
+ }
+ finally
+ {
+ if (packet is not null && !packet.IsInvalid)
+ {
+ connection.Handle.ReturnPacket(packet);
+ }
+ }
+ }
+
+ ///
+ /// Disposal must not wait for a pending MARS receive, and the resulting error callback
+ /// must be able to acquire the demultiplexer lock and release its packet.
+ ///
+ [Theory]
+ [InlineData(false)]
+ [InlineData(true)]
+ public async Task MarsReceive_PendingRead_DisposeCompletesAndReportsError(bool namedPipe)
+ {
+ using LocalConnection connection = new(namedPipe);
+ SniMarsConnection mars = new(connection.Handle);
+ TaskCompletionSource completion = new(TaskCreationOptions.RunContinuationsAsynchronously);
+ connection.Handle.SetAsyncCallbacks((packet, error) =>
+ {
+ try
+ {
+ mars.HandleReceiveComplete(packet, error);
+ Assert.True(packet.IsInvalid);
+ completion.SetResult(error);
+ }
+ catch (Exception ex)
+ {
+ completion.SetException(ex);
+ }
+ }, null);
+
+ Assert.Equal(TdsEnums.SNI_SUCCESS_IO_PENDING, mars.StartReceive());
+ Assert.False(completion.Task.IsCompleted);
+ await Task.Run(() =>
+ {
+ lock (mars.DemuxerSync)
+ {
+ connection.Handle.Dispose();
+ }
+ }).WaitAsync(TimeSpan.FromSeconds(5));
+ Assert.Equal(TdsEnums.SNI_ERROR, await completion.Task.WaitAsync(TimeSpan.FromSeconds(5)));
+ }
+
+ ///
+ /// An immediate re-arm failure keeps the consumed packet's callback and ownership
+ /// intact until MARS broadcasts the error; a new packet must not escape on initial failure.
+ ///
+ [Theory]
+ [InlineData(false)]
+ [InlineData(true)]
+ public void MarsReceive_AfterDispose_PreservesErrorPacket(bool namedPipe)
+ {
+ using LocalConnection connection = new(namedPipe);
+ SniMarsConnection mars = new(connection.Handle);
+ SniPacket previousPacket = connection.Handle.RentPacket(0, 1);
+ previousPacket.SetAsyncIOCompletionCallback(mars.HandleReceiveComplete);
+ SniPacket? packet = previousPacket;
+ connection.Handle.Dispose();
+ try
+ {
+ Assert.Equal(TdsEnums.SNI_ERROR, mars.ReceiveAsync(ref packet));
+ Assert.Same(previousPacket, packet);
+ Assert.False(previousPacket.IsInvalid);
+ Assert.True(previousPacket.HasAsyncIOCompletionCallback);
+ }
+ finally
+ {
+ if (!previousPacket.IsInvalid)
+ {
+ connection.Handle.ReturnPacket(previousPacket);
+ }
+ }
+
+ Assert.Equal(TdsEnums.SNI_ERROR, mars.StartReceive());
+ }
+
+ ///
+ /// Re-arm failure after a complete ACK must reach every pending parser callback
+ /// exactly once, including a task cleared during closure, without calling an idle session.
+ ///
+ [Theory]
+ [InlineData(false)]
+ [InlineData(true)]
+ public async Task MarsReceiveComplete_AfterDispose_NotifiesAllSessions(bool namedPipe)
+ {
+ using LocalConnection connection = new(namedPipe);
+ SniMarsConnection mars = new(connection.Handle);
+ TdsParser parser = new(MARS: true, fAsynchronous: true);
+ TdsParserStateObjectManaged[] callbacks =
+ {
+ new(parser),
+ new(parser),
+ new(parser),
+ new(parser),
+ };
+ SniPacket? packet = null;
+ try
+ {
+ SniMarsHandle[] sessions = new SniMarsHandle[callbacks.Length];
+ for (int i = 0; i < callbacks.Length; i++)
+ {
+ callbacks[i].TimeoutTime = long.MaxValue;
+ if (i > 0)
+ {
+ callbacks[i]._networkPacketTaskSource = new(TaskCreationOptions.RunContinuationsAsynchronously);
+ callbacks[i].IncrementPendingCallbacks();
+ }
+ byte[] syn = new byte[SniSmuxHeader.HEADER_LENGTH];
+ using CancellationTokenSource timeout = new(TimeSpan.FromSeconds(5));
+ // Consume the synchronous SYN write even when the pipe has no buffer space.
+ Task readSyn = connection.Peer.ReadExactlyAsync(syn, timeout.Token).AsTask();
+ using CancellationTokenRegistration abort = timeout.Token.Register(connection.Peer.Dispose);
+ sessions[i] = mars.CreateMarsSession(callbacks[i], async: true);
+ await readSyn;
+ SniSmuxHeader synHeader = new();
+ synHeader.Read(syn);
+ Assert.Equal((byte)SniSmuxFlags.SMUX_SYN, synHeader.flags);
+ Assert.Equal((ushort)i, synHeader.sessionId);
+ if (i > 0)
+ {
+ SniPacket? pending = null;
+ Assert.Equal(TdsEnums.SNI_SUCCESS_IO_PENDING, sessions[i].ReceiveAsync(ref pending));
+ Assert.Null(pending);
+ }
+ }
+
+ byte[] header = new byte[SniSmuxHeader.HEADER_LENGTH];
+ new SniSmuxHeader
+ {
+ SMID = 83,
+ flags = (byte)SniSmuxFlags.SMUX_ACK,
+ length = SniSmuxHeader.HEADER_LENGTH,
+ highwater = 4,
+ }.Write(header);
+ packet = connection.Handle.RentPacket(0, header.Length);
+ packet.SetAsyncIOCompletionCallback(mars.HandleReceiveComplete);
+ packet.AppendData(header, header.Length);
+
+ parser.State = TdsParserState.Broken;
+ // Model a pending read whose task was already completed and cleared by teardown.
+ callbacks[1]._networkPacketTaskSource.SetCanceled();
+ await Assert.ThrowsAsync(() => callbacks[1]._networkPacketTaskSource.Task);
+ callbacks[1]._networkPacketTaskSource = null;
+ connection.Handle.Dispose();
+ mars.HandleReceiveComplete(packet, TdsEnums.SNI_SUCCESS);
+ SniError receiveError = SniLoadHandle.LastError;
+ Assert.IsType(receiveError.exception);
+ Assert.True(packet.IsInvalid);
+
+ packet = connection.Handle.RentPacket(0, header.Length);
+ packet.SetAsyncIOCompletionCallback(mars.HandleReceiveComplete);
+ lock (mars.DemuxerSync)
+ {
+ mars.HandleReceiveError(packet);
+ }
+ Assert.True(packet.IsInvalid);
+
+ for (int i = 0; i < sessions.Length; i++)
+ {
+ if (i > 1)
+ {
+ await Assert.ThrowsAsync(() =>
+ callbacks[i]._networkPacketTaskSource.Task.WaitAsync(TimeSpan.FromSeconds(5)));
+ }
+ else
+ {
+ Assert.Null(callbacks[i]._networkPacketTaskSource);
+ }
+ Assert.Equal(TdsEnums.SNI_ERROR, sessions[i].Receive(out SniPacket? received, 1000));
+ Assert.Null(received);
+ Assert.Same(receiveError, SniLoadHandle.LastError);
+ Assert.Equal(0, callbacks[i].DecrementPendingCallbacks(release: false));
+ Assert.Equal(TdsEnums.SNI_ERROR, sessions[i].ReceiveAsync(ref received));
+ Assert.Null(received);
+ Assert.Same(receiveError, SniLoadHandle.LastError);
+ }
+ }
+ finally
+ {
+ if (packet is not null && !packet.IsInvalid)
+ {
+ connection.Handle.ReturnPacket(packet);
+ }
+ foreach (TdsParserStateObjectManaged callback in callbacks)
+ {
+ callback.Dispose();
+ }
+ parser._physicalStateObj.Dispose();
+ }
+ }
+
+ ///
+ /// Successful re-arms recycle consumed packets without clearing the new packet's
+ /// callback or data, including after the physical packet pool starts reusing objects.
+ ///
+ [Theory]
+ [InlineData(false)]
+ [InlineData(true)]
+ public async Task MarsReceive_Connected_RecyclesConsumedPackets(bool namedPipe)
+ {
+ using LocalConnection connection = new(namedPipe);
+ SniMarsConnection mars = new(connection.Handle);
+ SniPacket? packet = null;
+ try
+ {
+ for (byte value = 1; value <= 3; value++)
+ {
+ TaskCompletionSource<(SniPacket Packet, uint Error)> completion =
+ new(TaskCreationOptions.RunContinuationsAsynchronously);
+ connection.Handle.SetAsyncCallbacks(
+ (received, error) => completion.SetResult((received, error)), null);
+ SniPacket? previousPacket = packet;
+ Assert.Equal(TdsEnums.SNI_SUCCESS_IO_PENDING, mars.ReceiveAsync(ref packet));
+ Assert.NotNull(packet);
+ Assert.True(previousPacket is null || ReferenceEquals(previousPacket, packet) || previousPacket.IsInvalid);
+ Assert.True(packet.HasAsyncIOCompletionCallback);
+ await connection.Peer.WriteAsync(new byte[] { value });
+ (SniPacket received, uint error) = await completion.Task.WaitAsync(TimeSpan.FromSeconds(5));
+ Assert.Same(packet, received);
+ Assert.Equal(TdsEnums.SNI_SUCCESS, error);
+ byte[] data = new byte[1];
+ Assert.Equal(1, packet.TakeData(data, 0, data.Length));
+ Assert.Equal(value, data[0]);
+ }
+ }
+ finally
+ {
+ if (packet is not null && !packet.IsInvalid)
+ {
+ connection.Handle.ReturnPacket(packet);
+ }
+ }
+ }
+
+ ///
+ /// Establishes a real local transport without SQL Server and owns both ends.
+ ///
+ private sealed class LocalConnection : IDisposable
+ {
+ public SniHandle Handle { get; }
+
+ public Stream Peer { get; }
+
+ public LocalConnection(bool namedPipe)
+ {
+ if (namedPipe)
+ {
+ string pipeName = $"SqlClient-{Guid.NewGuid():N}";
+ NamedPipeServerStream server = new(
+ pipeName, PipeDirection.InOut, 1, PipeTransmissionMode.Byte, PipeOptions.Asynchronous);
+ Peer = server;
+ using CancellationTokenSource timeout = new(TimeSpan.FromSeconds(5));
+ Task accept = server.WaitForConnectionAsync(timeout.Token);
+ Handle = new SniNpHandle(
+ ".", pipeName, TimeoutTimer.StartNew(TimeSpan.FromSeconds(5)),
+ tlsFirst: false, hostNameInCertificate: null, serverCertificateFilename: null);
+ accept.GetAwaiter().GetResult();
+ }
+ else
+ {
+ using TcpListener listener = new(IPAddress.Loopback, 0);
+ listener.Start();
+ SQLDNSInfo? pendingDnsInfo = null;
+ Handle = new SniTcpHandle(
+ IPAddress.Loopback.ToString(),
+ ((IPEndPoint)listener.LocalEndpoint).Port,
+ TimeoutTimer.StartNew(TimeSpan.FromSeconds(5)),
+ parallel: false,
+ SqlConnectionIPAddressPreference.IPv4First,
+ cachedFQDN: nameof(SniReceiveTests),
+ ref pendingDnsInfo,
+ tlsFirst: false,
+ hostNameInCertificate: null,
+ serverCertificateFilename: null);
+ using CancellationTokenSource timeout = new(TimeSpan.FromSeconds(5));
+ Peer = new NetworkStream(
+ listener.AcceptSocketAsync(timeout.Token).AsTask().GetAwaiter().GetResult(), ownsSocket: true);
+ }
+ Assert.Equal(TdsEnums.SNI_SUCCESS, Handle.Status);
+ }
+
+ ///
+ /// Closes the client handle and its local server stream.
+ ///
+ public void Dispose()
+ {
+ Handle.Dispose();
+ Peer.Dispose();
+ }
+ }
+ }
+}
+
+#endif