feat(provisioning-device, prov-amqp, prov-mqtt, prov-https): Add support for timespan timeouts to provisioning device client

e2e/test/provisioning/ProvisioningE2ETests.cs

@@ -4,6 +4,7 @@
 using System;
 using System.Collections.Generic;
 using System.Net;
+using System.Net.Sockets;
 using System.Security.Cryptography;
 using System.Security.Cryptography.X509Certificates;
 using System.Text;
@@ -432,6 +433,62 @@ public async Task ProvisioningDeviceClient_ValidRegistrationId_MqttWsWithProxy_S
             await ProvisioningDeviceClient_ValidRegistrationId_Register_Ok(Client.TransportType.Mqtt, AttestationMechanismType.SymmetricKey, EnrollmentType.Individual, true, s_proxyServerAddress).ConfigureAwait(false);
         }
 
+        [LoggedTestMethod]
+        public async Task ProvisioningDeviceClient_ValidRegistrationId_TimeSpanTimeoutRespected_Mqtt()
+        {
+            try
+            {
+                await ProvisioningDeviceClient_ValidRegistrationId_Register_Ok(Client.TransportType.Mqtt_Tcp_Only, AttestationMechanismType.SymmetricKey, EnrollmentType.Individual, TimeSpan.Zero).ConfigureAwait(false);
+            }
+            catch (OperationCanceledException)
+            {
+                return; // expected exception was thrown, so exit the test
+            }
+
+            throw new AssertFailedException("Expected an OperationCanceledException to be thrown since the timeout was set to TimeSpan.Zero");
+        }
+
+        [LoggedTestMethod]
+        public async Task ProvisioningDeviceClient_ValidRegistrationId_TimeSpanTimeoutRespected_Https()
+        {
+            try
+            {
+                await ProvisioningDeviceClient_ValidRegistrationId_Register_Ok(Client.TransportType.Http1, AttestationMechanismType.SymmetricKey, EnrollmentType.Individual, TimeSpan.Zero).ConfigureAwait(false);
+            }
+            catch (OperationCanceledException)
+            {
+                return; // expected exception was thrown, so exit the test
+            }
+
+            throw new AssertFailedException("Expected an OperationCanceledException to be thrown since the timeout was set to TimeSpan.Zero");
+        }
+
+        [LoggedTestMethod]
+        public async Task ProvisioningDeviceClient_ValidRegistrationId_TimeSpanTimeoutRespected_Amqps()
+        {
+            try
+            {
+                await ProvisioningDeviceClient_ValidRegistrationId_Register_Ok(Client.TransportType.Amqp_Tcp_Only, AttestationMechanismType.SymmetricKey, EnrollmentType.Individual, TimeSpan.Zero).ConfigureAwait(false);
+            }
+            catch (ProvisioningTransportException ex) when (ex.InnerException is SocketException && ((SocketException) ex.InnerException).SocketErrorCode == SocketError.TimedOut)
+            {
+                // The expected exception is a bit different in AMQP compared to MQTT/HTTPS
+                return; // expected exception was thrown, so exit the test
+            }
+
+            throw new AssertFailedException("Expected an OperationCanceledException to be thrown since the timeout was set to TimeSpan.Zero");
+        }
+
+        public async Task ProvisioningDeviceClient_ValidRegistrationId_Register_Ok(
+            Client.TransportType transportType,
+            AttestationMechanismType attestationType,
+            EnrollmentType? enrollmentType,
+            TimeSpan timeout)
+        {
+            //Default reprovisioning settings: Hashed allocation, no reprovision policy, hub names, or custom allocation policy
+            await ProvisioningDeviceClientValidRegistrationIdRegisterOkAsync(transportType, attestationType, enrollmentType, false, null, AllocationPolicy.Hashed, null, null, null, timeout, s_proxyServerAddress).ConfigureAwait(false);
+        }
+
         public async Task ProvisioningDeviceClient_ValidRegistrationId_Register_Ok(
             Client.TransportType transportType,
             AttestationMechanismType attestationType,
@@ -440,7 +497,7 @@ public async Task ProvisioningDeviceClient_ValidRegistrationId_Register_Ok(
             string proxyServerAddress = null)
         {
             //Default reprovisioning settings: Hashed allocation, no reprovision policy, hub names, or custom allocation policy
-            await ProvisioningDeviceClientValidRegistrationIdRegisterOkAsync(transportType, attestationType, enrollmentType, setCustomProxy, null, AllocationPolicy.Hashed, null, null, null, s_proxyServerAddress).ConfigureAwait(false);
+            await ProvisioningDeviceClientValidRegistrationIdRegisterOkAsync(transportType, attestationType, enrollmentType, setCustomProxy, null, AllocationPolicy.Hashed, null, null, null, TimeSpan.MaxValue, proxyServerAddress).ConfigureAwait(false);
         }
 
         public async Task ProvisioningDeviceClient_ValidRegistrationId_Register_Ok(
@@ -463,7 +520,8 @@ await ProvisioningDeviceClientValidRegistrationIdRegisterOkAsync(
                     null,
                     iothubs,
                     capabilities,
-                    s_proxyServerAddress)
+                    TimeSpan.MaxValue,
+                    proxyServerAddress)
                 .ConfigureAwait(false);
         }
 
@@ -477,6 +535,7 @@ private async Task ProvisioningDeviceClientValidRegistrationIdRegisterOkAsync(
             CustomAllocationDefinition customAllocationDefinition,
             ICollection<string> iothubs,
             DeviceCapabilities deviceCapabilities,
+            TimeSpan timeout,
             string proxyServerAddress = null)
         {
             string groupId = _idPrefix + AttestationTypeToString(attestationType) + "-" + Guid.NewGuid();
@@ -516,7 +575,14 @@ private async Task ProvisioningDeviceClientValidRegistrationIdRegisterOkAsync(
             {
                 try
                 {
-                    result = await provClient.RegisterAsync(cts.Token).ConfigureAwait(false);
+                    if (timeout != TimeSpan.MaxValue)
+                    {
+                        result = await provClient.RegisterAsync(timeout).ConfigureAwait(false);
+                    }
+                    else
+                    { 
+                        result = await provClient.RegisterAsync(cts.Token).ConfigureAwait(false);
+                    }
                     break;
                 }
                 // Catching all ProvisioningTransportException as the status code is not the same for Mqtt, Amqp and Http.

provisioning/device/src/ProvisioningDeviceClient.cs

@@ -3,6 +3,7 @@
 
 using Microsoft.Azure.Devices.Provisioning.Client.Transport;
 using Microsoft.Azure.Devices.Shared;
+using System;
 using System.Threading;
 using System.Threading.Tasks;
 
@@ -63,52 +64,87 @@ private ProvisioningDeviceClient(
         /// <summary>
         /// Registers the current device using the Device Provisioning Service and assigns it to an IoT Hub.
         /// </summary>
+        /// <param name="timeout">The maximum amount of time to allow this operation to run for before timing out.</param>
+        /// <remarks>
+        /// Due to the AMQP library used by this library uses not accepting cancellation tokens, this overload and <see cref="RegisterAsync(ProvisioningRegistrationAdditionalData, TimeSpan)"/>
+        /// are the only overloads for this method that allow for a specified timeout to be respected in the middle of an AMQP operation such as opening
+        /// the AMQP connection. MQTT and HTTPS connections do not share that same limitation, though.
+        /// </remarks>
         /// <returns>The registration result.</returns>
-        public Task<DeviceRegistrationResult> RegisterAsync()
+        public Task<DeviceRegistrationResult> RegisterAsync(TimeSpan timeout)
         {
-            return RegisterAsync(CancellationToken.None);
+            return RegisterAsync(null, timeout);
         }
 
         /// <summary>
         /// Registers the current device using the Device Provisioning Service and assigns it to an IoT Hub.
         /// </summary>
-        /// <param name="data">The optional additional data.</param>
+        /// <param name="data">
+        /// The optional additional data that is passed through to the custom allocation policy webhook if 
+        /// a custom allocation policy webhook is setup for this enrollment.
+        /// </param>
+        /// <param name="timeout">The maximum amount of time to allow this operation to run for before timing out.</param>
+        /// <remarks>
+        /// Due to the AMQP library used by this library uses not accepting cancellation tokens, this overload and <see cref="RegisterAsync(TimeSpan)"/>
+        /// are the only overloads for this method that allow for a specified timeout to be respected in the middle of an AMQP operation such as opening
+        /// the AMQP connection. MQTT and HTTPS connections do not share that same limitation, though.
+        /// </remarks>
         /// <returns>The registration result.</returns>
-        public Task<DeviceRegistrationResult> RegisterAsync(ProvisioningRegistrationAdditionalData data)
+        public Task<DeviceRegistrationResult> RegisterAsync(ProvisioningRegistrationAdditionalData data, TimeSpan timeout)
         {
-            return RegisterAsync(data, CancellationToken.None);
+            Logging.RegisterAsync(this, _globalDeviceEndpoint, _idScope, _transport, _security);
+
+            var request = new ProvisioningTransportRegisterMessage(_globalDeviceEndpoint, _idScope, _security, data?.JsonData)
+            {
+                ProductInfo = ProductInfo,
+            };
+
+            return _transport.RegisterAsync(request, timeout);
         }
 
         /// <summary>
         /// Registers the current device using the Device Provisioning Service and assigns it to an IoT Hub.
         /// </summary>
         /// <param name="cancellationToken">The cancellation token.</param>
+        /// <remarks>
+        /// Due to the AMQP library used by this library uses not accepting cancellation tokens, the provided cancellation token will only be checked 
+        /// for cancellation in between AMQP operations, and not during. In order to have a timeout for this operation that is checked during AMQP operations
+        /// (such as opening the connection), you must use <see cref="RegisterAsync(TimeSpan)"/> instead. MQTT and HTTPS connections do not have the same 
+        /// behavior as AMQP connections in this regard. MQTT and HTTPS connections will check this cancellation token for cancellation during their protocol level operations.
+        /// </remarks>
         /// <returns>The registration result.</returns>
-        public Task<DeviceRegistrationResult> RegisterAsync(CancellationToken cancellationToken)
+        public Task<DeviceRegistrationResult> RegisterAsync(CancellationToken cancellationToken = default)
         {
             Logging.RegisterAsync(this, _globalDeviceEndpoint, _idScope, _transport, _security);
 
-            var request = new ProvisioningTransportRegisterMessage(_globalDeviceEndpoint, _idScope, _security)
-            {
-                ProductInfo = ProductInfo
-            };
-            return _transport.RegisterAsync(request, cancellationToken);
+            return RegisterAsync(null, cancellationToken);
         }
 
         /// <summary>
         /// Registers the current device using the Device Provisioning Service and assigns it to an IoT Hub.
         /// </summary>
-        /// <param name="data">The custom content.</param>
+        /// <param name="data">
+        /// The optional additional data that is passed through to the custom allocation policy webhook if 
+        /// a custom allocation policy webhook is setup for this enrollment.
+        /// </param>
         /// <param name="cancellationToken">The cancellation token.</param>
+        /// <remarks>
+        /// Due to the AMQP library used by this library uses not accepting cancellation tokens, the provided cancellation token will only be checked 
+        /// for cancellation in between AMQP operations, and not during. In order to have a timeout for this operation that is checked during AMQP operations
+        /// (such as opening the connection), you must use <see cref="RegisterAsync(ProvisioningRegistrationAdditionalData, TimeSpan)">this overload</see> instead. 
+        /// MQTT and HTTPS connections do not have the same behavior as AMQP connections in this regard. MQTT and HTTPS connections will check this cancellation 
+        /// token for cancellation during their protocol level operations.
+        /// </remarks>
         /// <returns>The registration result.</returns>
-        public Task<DeviceRegistrationResult> RegisterAsync(ProvisioningRegistrationAdditionalData data, CancellationToken cancellationToken)
+        public Task<DeviceRegistrationResult> RegisterAsync(ProvisioningRegistrationAdditionalData data, CancellationToken cancellationToken = default)
         {
             Logging.RegisterAsync(this, _globalDeviceEndpoint, _idScope, _transport, _security);
 
             var request = new ProvisioningTransportRegisterMessage(_globalDeviceEndpoint, _idScope, _security, data?.JsonData)
             {
                 ProductInfo = ProductInfo,
             };
+
             return _transport.RegisterAsync(request, cancellationToken);
         }
     }

provisioning/device/src/ProvisioningTransportHandler.cs

@@ -85,6 +85,19 @@ public virtual Task<DeviceRegistrationResult> RegisterAsync(
             return _innerHandler.RegisterAsync(message, cancellationToken);
         }
 
+        /// <summary>
+        /// Registers a device described by the message.
+        /// </summary>
+        /// <param name="message">The provisioning message.</param>
+        /// <param name="timeout">The maximum amount of time to allow this operation to run for before timing out.</param>
+        /// <returns>The registration result.</returns>
+        public virtual Task<DeviceRegistrationResult> RegisterAsync(
+            ProvisioningTransportRegisterMessage message,
+            TimeSpan timeout)
+        {
+            return _innerHandler.RegisterAsync(message, timeout);
+        }
+
         /// <summary>
         /// Releases the unmanaged resources and disposes of the managed resources used by the invoker.
         /// </summary>