Add timeout of 10 minutes on ADE related HTTP calls

Author: Kamran Khan

src/ResourceManager/Compute/Commands.Compute/Extension/AzureDiskEncryption/AzureDiskEncryptionExtensionConstants.cs

@@ -39,5 +39,6 @@ public static class AzureDiskEncryptionExtensionConstants
         public const string osTypeLinux = "Linux";
         public const string osTypeWindows = "Windows";
         public const string defaultKeyEncryptionAlgorithm = "RSA-OAEP";
+        public const int httpTimeoutMinutes = 10;
     }
 }

src/ResourceManager/Compute/Commands.Compute/Extension/AzureDiskEncryption/DisableAzureDiskEncryption.cs

@@ -22,6 +22,7 @@
 using System.Collections;
 using System.Globalization;
 using System.Management.Automation;
+using System.Threading.Tasks;
 
 namespace Microsoft.Azure.Commands.Compute.Extension.AzureDiskEncryption
 {
@@ -199,10 +200,21 @@ private AzureOperationResponse<VirtualMachine> UpdateVmEncryptionSettings()
                 Tags = vmParameters.Tags
             };
 
-            return this.ComputeClient.ComputeManagementClient.VirtualMachines.CreateOrUpdateWithHttpMessagesAsync(
+            var httpTask = this.ComputeClient.ComputeManagementClient.VirtualMachines.CreateOrUpdateWithHttpMessagesAsync(
                 this.ResourceGroupName,
                 vmParameters.Name,
-                parameters).GetAwaiter().GetResult();
+                parameters);
+
+            if (!httpTask.Wait(new TimeSpan(0, AzureDiskEncryptionExtensionConstants.httpTimeoutMinutes, 0)))
+            {
+                string errorMessage = string.Format(CultureInfo.CurrentUICulture, "Extension operation timed out");
+                ThrowTerminatingError(new ErrorRecord(new ApplicationException(errorMessage),
+                                                      errorMessage,
+                                                      ErrorCategory.InvalidResult,
+                                                      null));
+            }
+
+            return httpTask.Result;
         }
 
         public override void ExecuteCmdlet()

src/ResourceManager/Compute/Commands.Compute/Extension/AzureDiskEncryption/GetAzureDiskEncryptionStatus.cs

@@ -368,11 +368,20 @@ public override void ExecuteCmdlet()
                             this.ResourceGroupName, VMName).Body;
                         VirtualMachineExtension parameters = GetVmExtensionParameters(virtualMachineResponse, osType);
 
-                        this.VirtualMachineExtensionClient.CreateOrUpdateWithHttpMessagesAsync(
+                        var httpTask = this.VirtualMachineExtensionClient.CreateOrUpdateWithHttpMessagesAsync(
                             this.ResourceGroupName,
                             this.VMName,
                             this.Name,
-                            parameters).GetAwaiter().GetResult();
+                            parameters);
+
+                        if (!httpTask.Wait(new TimeSpan(0, AzureDiskEncryptionExtensionConstants.httpTimeoutMinutes, 0)))
+                        {
+                            string errorMessage = string.Format(CultureInfo.CurrentUICulture, "Extension operation timed out");
+                            ThrowTerminatingError(new ErrorRecord(new ApplicationException(errorMessage),
+                                                                  errorMessage,
+                                                                  ErrorCategory.InvalidResult,
+                                                                  null));
+                        }
 
                         Dictionary<string, string> encryptionStatusParsed = null;
                         try

src/ResourceManager/Compute/Commands.Compute/Extension/AzureDiskEncryption/SetAzureDiskEncryptionExtension.cs

@@ -23,6 +23,7 @@
 using System.Collections;
 using System.Globalization;
 using System.Management.Automation;
+using System.Threading.Tasks;
 
 namespace Microsoft.Azure.Commands.Compute.Extension.AzureDiskEncryption
 {
@@ -299,10 +300,21 @@ private AzureOperationResponse<VirtualMachine> UpdateVmEncryptionSettings()
                 Location = vmParameters.Location,
                 Tags = vmParameters.Tags
             };
-            return this.ComputeClient.ComputeManagementClient.VirtualMachines.CreateOrUpdateWithHttpMessagesAsync(
+            var httpTask = this.ComputeClient.ComputeManagementClient.VirtualMachines.CreateOrUpdateWithHttpMessagesAsync(
                 this.ResourceGroupName,
                 vmParameters.Name,
-                parameters).GetAwaiter().GetResult();
+                parameters);
+
+            if (!httpTask.Wait(new TimeSpan(0, AzureDiskEncryptionExtensionConstants.httpTimeoutMinutes, 0)))
+            {
+                string errorMessage = string.Format(CultureInfo.CurrentUICulture, "Extension operation timed out");
+                ThrowTerminatingError(new ErrorRecord(new ApplicationException(errorMessage),
+                                                      errorMessage,
+                                                      ErrorCategory.InvalidResult,
+                                                      null));
+            }
+
+            return httpTask.Result;
         }
 
         private Hashtable GetExtensionPublicSettings()
@@ -437,11 +449,20 @@ public override void ExecuteCmdlet()
 
                     VirtualMachineExtension parameters = GetVmExtensionParameters(virtualMachineResponse);
 
-                    this.VirtualMachineExtensionClient.CreateOrUpdateWithHttpMessagesAsync(
+                    var httpTask = this.VirtualMachineExtensionClient.CreateOrUpdateWithHttpMessagesAsync(
                         this.ResourceGroupName,
                         this.VMName,
                         this.Name,
-                        parameters).GetAwaiter().GetResult();
+                        parameters);
+
+                    if (!httpTask.Wait(new TimeSpan(0, AzureDiskEncryptionExtensionConstants.httpTimeoutMinutes, 0)))
+                    {
+                        string errorMessage = string.Format(CultureInfo.CurrentUICulture, "Extension operation timed out");
+                        ThrowTerminatingError(new ErrorRecord(new ApplicationException(errorMessage),
+                                                              errorMessage,
+                                                              ErrorCategory.InvalidResult,
+                                                              null));
+                    }
 
                     var op = UpdateVmEncryptionSettings();
                     var result = Mapper.Map<PSAzureOperationResponse>(op);