[Event Hubs Client] AMQP Cancellation Support (Azure#25654)

The focus of these changes is to take advantage of the support for
cancellation tokens that was added to the AMQP transport library, preserving
existing timeout behavior by specifying an operation timeout for the link.

Author: jsquire

sdk/eventhub/Azure.Messaging.EventHubs/CHANGELOG.md

@@ -4,6 +4,8 @@
 
 ### Features Added
 
+- Support for cancellation tokens has been improved for AMQP operations, enabling earlier detection of cancellation requests without needing to wait for the configured timeout to elapse.
+
 ### Breaking Changes
 
 ### Bugs Fixed

sdk/eventhub/Azure.Messaging.EventHubs/src/Amqp/AmqpClient.cs

@@ -99,6 +99,7 @@ internal class AmqpClient : TransportClient
         ///
         /// <param name="host">The fully qualified host name for the Event Hubs namespace.  This is likely to be similar to <c>{yournamespace}.servicebus.windows.net</c>.</param>
         /// <param name="eventHubName">The name of the specific Event Hub to connect the client to.</param>
+        /// <param name="operationTimeout">The amount of time to allow for an AMQP operation using the link to complete before attempting to cancel it.</param>
         /// <param name="credential">The Azure managed identity credential to use for authorization.  Access controls may be specified by the Event Hubs namespace or the requested Event Hub, depending on Azure configuration.</param>
         /// <param name="clientOptions">A set of options to apply when configuring the client.</param>
         ///
@@ -113,8 +114,9 @@ internal class AmqpClient : TransportClient
         ///
         public AmqpClient(string host,
                           string eventHubName,
+                          TimeSpan operationTimeout,
                           EventHubTokenCredential credential,
-                          EventHubConnectionOptions clientOptions) : this(host, eventHubName, credential, clientOptions, null, null)
+                          EventHubConnectionOptions clientOptions) : this(host, eventHubName, operationTimeout, credential, clientOptions, null, null)
         {
         }
 
@@ -124,6 +126,7 @@ public AmqpClient(string host,
         ///
         /// <param name="host">The fully qualified host name for the Event Hubs namespace.  This is likely to be similar to <c>{yournamespace}.servicebus.windows.net</c>.</param>
         /// <param name="eventHubName">The name of the specific Event Hub to connect the client to.</param>
+        /// <param name="operationTimeout">The amount of time to allow for an AMQP operation using the link to complete before attempting to cancel it.</param>
         /// <param name="credential">The Azure managed identity credential to use for authorization.  Access controls may be specified by the Event Hubs namespace or the requested Event Hub, depending on Azure configuration.</param>
         /// <param name="clientOptions">A set of options to apply when configuring the client.</param>
         /// <param name="connectionScope">The optional scope to use for AMQP connection management.  If <c>null</c>, a new scope will be created.</param>
@@ -140,13 +143,15 @@ public AmqpClient(string host,
         ///
         protected AmqpClient(string host,
                              string eventHubName,
+                             TimeSpan operationTimeout,
                              EventHubTokenCredential credential,
                              EventHubConnectionOptions clientOptions,
                              AmqpConnectionScope connectionScope,
                              AmqpMessageConverter messageConverter)
         {
             Argument.AssertNotNullOrEmpty(host, nameof(host));
             Argument.AssertNotNullOrEmpty(eventHubName, nameof(eventHubName));
+            Argument.AssertNotNegative(operationTimeout, nameof(operationTimeout));
             Argument.AssertNotNull(credential, nameof(credential));
             Argument.AssertNotNull(clientOptions, nameof(clientOptions));
 
@@ -189,7 +194,7 @@ protected AmqpClient(string host,
                     clientOptions.CertificateValidationCallback);
 
                 ManagementLink = new FaultTolerantAmqpObject<RequestResponseAmqpLink>(
-                    timeout => ConnectionScope.OpenManagementLinkAsync(timeout, CancellationToken.None),
+                    linkTimeout => ConnectionScope.OpenManagementLinkAsync(operationTimeout, linkTimeout, CancellationToken.None),
                     link =>
                     {
                         link.Session?.SafeClose();
@@ -238,17 +243,13 @@ public override async Task<EventHubProperties> GetPropertiesAsync(EventHubsRetry
 
                         var token = await AcquireAccessTokenAsync(cancellationToken).ConfigureAwait(false);
                         using AmqpMessage request = MessageConverter.CreateEventHubPropertiesRequest(EventHubName, token);
-                        cancellationToken.ThrowIfCancellationRequested<TaskCanceledException>();
 
                         RequestResponseAmqpLink link = await ManagementLink.GetOrCreateAsync(UseMinimum(ConnectionScope.SessionTimeout, tryTimeout.CalculateRemaining(stopWatch.GetElapsedTime()))).ConfigureAwait(false);
                         cancellationToken.ThrowIfCancellationRequested<TaskCanceledException>();
 
-                        // Send the request and wait for the response.
-
-                        using AmqpMessage response = await link.RequestAsync(request, tryTimeout.CalculateRemaining(stopWatch.GetElapsedTime())).ConfigureAwait(false);
-                        cancellationToken.ThrowIfCancellationRequested<TaskCanceledException>();
+                        // Send the request and process the response.
 
-                        // Process the response.
+                        using AmqpMessage response = await link.RequestAsync(request, cancellationToken).ConfigureAwait(false);
 
                         AmqpError.ThrowIfErrorResponse(response, EventHubName);
                         return MessageConverter.CreateEventHubPropertiesFromResponse(response);
@@ -338,17 +339,13 @@ public override async Task<PartitionProperties> GetPartitionPropertiesAsync(stri
 
                         token = await AcquireAccessTokenAsync(cancellationToken).ConfigureAwait(false);
                         using AmqpMessage request = MessageConverter.CreatePartitionPropertiesRequest(EventHubName, partitionId, token);
-                        cancellationToken.ThrowIfCancellationRequested<TaskCanceledException>();
 
                         link = await ManagementLink.GetOrCreateAsync(UseMinimum(ConnectionScope.SessionTimeout, tryTimeout.CalculateRemaining(stopWatch.GetElapsedTime()))).ConfigureAwait(false);
                         cancellationToken.ThrowIfCancellationRequested<TaskCanceledException>();
 
-                        // Send the request and wait for the response.
+                        // Send the request and process the response.
 
-                        using AmqpMessage response = await link.RequestAsync(request, tryTimeout.CalculateRemaining(stopWatch.GetElapsedTime())).ConfigureAwait(false);
-                        cancellationToken.ThrowIfCancellationRequested<TaskCanceledException>();
-
-                        // Process the response.
+                        using AmqpMessage response = await link.RequestAsync(request, cancellationToken).ConfigureAwait(false);
 
                         AmqpError.ThrowIfErrorResponse(response, EventHubName);
                         return MessageConverter.CreatePartitionPropertiesFromResponse(response);
@@ -506,6 +503,7 @@ public override async Task CloseAsync(CancellationToken cancellationToken)
                 return;
             }
 
+            cancellationToken.ThrowIfCancellationRequested<TaskCanceledException>();
             _closed = true;
 
             var clientId = GetHashCode().ToString(CultureInfo.InvariantCulture);
@@ -514,11 +512,9 @@ public override async Task CloseAsync(CancellationToken cancellationToken)
             try
             {
                 EventHubsEventSource.Log.ClientCloseStart(clientType, EventHubName, clientId);
-                cancellationToken.ThrowIfCancellationRequested<TaskCanceledException>();
 
                 if (ManagementLink?.TryGetOpenedObject(out var _) == true)
                 {
-                    cancellationToken.ThrowIfCancellationRequested<TaskCanceledException>();
                     await ManagementLink.CloseAsync().ConfigureAwait(false);
                 }