Don't close session when closing link if cross entity transactions are enabled (Azure#28351)

* Don't close session when closing link if cross entity transactions are enabled

* sender

* flaky tests

Author: JoshLove-msft

sdk/servicebus/Azure.Messaging.ServiceBus/src/Amqp/AmqpConnectionScope.cs

@@ -1330,5 +1330,23 @@ private static void ValidateTransport(ServiceBusTransportType transport)
                 throw new ArgumentException(string.Format(CultureInfo.CurrentCulture, Resources.UnknownConnectionType, transport), nameof(transport));
             }
         }
+
+        internal void CloseLink(AmqpLink link)
+        {
+            if (!_useSingleSession)
+            {
+                link.Session?.SafeClose();
+            }
+            link.SafeClose();
+        }
+
+        internal void CloseLink(RequestResponseAmqpLink link)
+        {
+            if (!_useSingleSession)
+            {
+                link.Session?.SafeClose();
+            }
+            link.SafeClose();
+        }
     }
 }

sdk/servicebus/Azure.Messaging.ServiceBus/src/Amqp/AmqpReceiver.cs

@@ -167,11 +167,11 @@ public AmqpReceiver(
                         // it is okay to register the user provided cancellationToken from the AcceptNextSessionAsync call in
                         // the fault tolerant object because session receivers are never reconnected.
                         cancellationToken: cancellationToken),
-                link => CloseLink(link));
+                link => _connectionScope.CloseLink(link));
 
             _managementLink = new FaultTolerantAmqpObject<RequestResponseAmqpLink>(
                 timeout => OpenManagementLinkAsync(timeout),
-                link => CloseLink(link));
+                link => _connectionScope.CloseLink(link));
         }
 
         private async Task<RequestResponseAmqpLink> OpenManagementLinkAsync(
@@ -254,18 +254,6 @@ private async Task<ReceivingAmqpLink> OpenReceiverLinkAsync(
             }
         }
 
-        private static void CloseLink(ReceivingAmqpLink link)
-        {
-            link.Session?.SafeClose();
-            link.SafeClose();
-        }
-
-        private static void CloseLink(RequestResponseAmqpLink link)
-        {
-            link.Session?.SafeClose();
-            link.SafeClose();
-        }
-
         /// <summary>
         /// Receives a list of <see cref="ServiceBusReceivedMessage" /> from the entity using <see cref="ServiceBusReceiveMode"/> mode.
         /// </summary>
@@ -1325,10 +1313,10 @@ public override async Task CloseAsync(CancellationToken cancellationToken)
         private void OnReceiverLinkClosed(object receiver, EventArgs e)
         {
             var receivingAmqpLink = (ReceivingAmqpLink)receiver;
-            if (_isSessionReceiver && receivingAmqpLink != null)
+            if (receivingAmqpLink != null)
             {
                 Exception exception = receivingAmqpLink.GetInnerException();
-                if (((exception is ServiceBusException sbException) && sbException.Reason != ServiceBusFailureReason.SessionLockLost) ||
+                if (_isSessionReceiver && ((exception is ServiceBusException sbException) && sbException.Reason != ServiceBusFailureReason.SessionLockLost) ||
                     !(exception is ServiceBusException))
                 {
                     exception = new ServiceBusException(

sdk/servicebus/Azure.Messaging.ServiceBus/src/Amqp/AmqpSender.cs

@@ -116,19 +116,11 @@ public AmqpSender(
 
             _sendLink = new FaultTolerantAmqpObject<SendingAmqpLink>(
                 timeout => CreateLinkAndEnsureSenderStateAsync(timeout, CancellationToken.None),
-                link =>
-                {
-                    link.Session?.SafeClose();
-                    link.SafeClose();
-                });
+                link => _connectionScope.CloseLink(link));
 
             _managementLink = new FaultTolerantAmqpObject<RequestResponseAmqpLink>(
                 timeout => OpenManagementLinkAsync(timeout),
-                link =>
-                {
-                    link.Session?.SafeClose();
-                    link.SafeClose();
-                });
+                link => _connectionScope.CloseLink(link));
         }
 
         private async Task<RequestResponseAmqpLink> OpenManagementLinkAsync(

sdk/servicebus/Azure.Messaging.ServiceBus/tests/Client/ServiceBusClientLiveTests.cs

@@ -265,7 +265,7 @@ public async Task MetricsAreUpdatedCorrectly()
                 await client.DisposeAsync();
                 // The close frame does not come back from the service before the DisposeAsync
                 // call is returned.
-                await Task.Delay(500);
+                await Task.Delay(1000);
                 metrics = client.GetTransportMetrics();
                 Assert.Greater(metrics.LastConnectionClose, thirdOpen);
                 Assert.Greater(metrics.LastHeartBeat, firstHeartBeat);

sdk/servicebus/Azure.Messaging.ServiceBus/tests/Diagnostics/EventSourceLiveTests.cs

@@ -487,5 +487,49 @@ public void LogsMessageEvents()
                 ServiceBusEventSource.PartitionKeyValueOverwritten,
                 e => e.Payload.Contains("sessionId1") && e.Payload.Contains("sessionId2") && e.Payload.Contains("messageId"));
         }
+
+        [Test]
+        public async Task ClosingSendLinkDoesNotCloseSessionWithCrossEntityTransactionEnabled()
+        {
+            await using (var scope = await ServiceBusScope.CreateWithQueue(enablePartitioning: false, enableSession: false))
+            {
+                await using var client = new ServiceBusClient(TestEnvironment.ServiceBusConnectionString,
+                    new ServiceBusClientOptions { EnableCrossEntityTransactions = true });
+                var sender = client.CreateSender(scope.QueueName);
+                var receiver = client.CreateReceiver(scope.QueueName);
+                await sender.SendMessageAsync(ServiceBusTestUtilities.GetMessage());
+                var message = await receiver.ReceiveMessageAsync();
+                Assert.IsNotNull(message);
+                await sender.CloseAsync();
+
+                // link closed event is fired asynchronously, so add a small delay
+                await Task.Delay(TimeSpan.FromSeconds(5));
+
+                _listener.SingleEventById(ServiceBusEventSource.SendLinkClosedEvent, e => e.Payload.Contains(sender.Identifier));
+                Assert.False(_listener.EventsById(ServiceBusEventSource.ReceiveLinkClosedEvent).Any(e => e.Payload.Contains(receiver.Identifier)));
+            }
+        }
+
+        [Test]
+        public async Task ClosingReceiveLinkDoesNotCloseSessionWithCrossEntityTransactionEnabled()
+        {
+            await using (var scope = await ServiceBusScope.CreateWithQueue(enablePartitioning: false, enableSession: false))
+            {
+                await using var client = new ServiceBusClient(TestEnvironment.ServiceBusConnectionString,
+                    new ServiceBusClientOptions { EnableCrossEntityTransactions = true });
+                var sender = client.CreateSender(scope.QueueName);
+                var receiver = client.CreateReceiver(scope.QueueName);
+                await sender.SendMessageAsync(ServiceBusTestUtilities.GetMessage());
+                var message = await receiver.ReceiveMessageAsync();
+                Assert.IsNotNull(message);
+                await receiver.CloseAsync();
+
+                // link closed event is fired asynchronously, so add a small delay
+                await Task.Delay(TimeSpan.FromSeconds(5));
+
+                _listener.SingleEventById(ServiceBusEventSource.ReceiveLinkClosedEvent, e => e.Payload.Contains(receiver.Identifier));
+                Assert.False(_listener.EventsById(ServiceBusEventSource.SendLinkClosedEvent).Any(e => e.Payload.Contains(sender.Identifier)));
+            }
+        }
     }
 }