Add timeout to workload api operations (#1307) (#1353)

* Add default operation timeout to Workload api client * Update WorkloadClientVersioned.cs
Azure · Jun 18, 2019 · a1b77bf · a1b77bf
1 parent c6b0ba9
commit a1b77bf
Show file tree

Hide file tree

Showing 4 changed files with 69 additions and 8 deletions.
diff --git a/edge-util/src/Microsoft.Azure.Devices.Edge.Util/edged/WorkloadClientVersioned.cs b/edge-util/src/Microsoft.Azure.Devices.Edge.Util/edged/WorkloadClientVersioned.cs
@@ -8,18 +8,28 @@ namespace Microsoft.Azure.Devices.Edge.Util.Edged
 
     abstract class WorkloadClientVersioned
     {
+        static readonly TimeSpan DefaultOperationTimeout = TimeSpan.FromMinutes(5);
+
         static readonly RetryStrategy TransientRetryStrategy =
             new ExponentialBackoff(retryCount: 2, minBackoff: TimeSpan.FromSeconds(1), maxBackoff: TimeSpan.FromSeconds(3), deltaBackoff: TimeSpan.FromSeconds(2));
 
         readonly ITransientErrorDetectionStrategy transientErrorDetectionStrategy;
-
-        protected WorkloadClientVersioned(Uri serverUri, ApiVersion apiVersion, string moduleId, string moduleGenerationId, ITransientErrorDetectionStrategy transientErrorDetectionStrategy)
+        readonly TimeSpan operationTimeout;
+
+        protected WorkloadClientVersioned(
+            Uri serverUri,
+            ApiVersion apiVersion,
+            string moduleId,
+            string moduleGenerationId,
+            ITransientErrorDetectionStrategy transientErrorDetectionStrategy,
+            Option<TimeSpan> operationTimeout)
         {
             this.WorkloadUri = Preconditions.CheckNotNull(serverUri, nameof(serverUri));
             this.Version = Preconditions.CheckNotNull(apiVersion, nameof(apiVersion));
             this.ModuleId = Preconditions.CheckNonWhiteSpace(moduleId, nameof(moduleId));
             this.ModuleGenerationId = Preconditions.CheckNonWhiteSpace(moduleGenerationId, nameof(moduleGenerationId));
             this.transientErrorDetectionStrategy = transientErrorDetectionStrategy;
+            this.operationTimeout = operationTimeout.GetOrElse(DefaultOperationTimeout);
         }
 
         protected Uri WorkloadUri { get; }
@@ -40,14 +50,13 @@ protected WorkloadClientVersioned(Uri serverUri, ApiVersion apiVersion, string m
 
         public abstract Task<string> SignAsync(string keyId, string algorithm, string data);
 
-        protected abstract void HandleException(Exception ex, string operation);
-
-        protected async Task<T> Execute<T>(Func<Task<T>> func, string operation)
+        protected internal async Task<T> Execute<T>(Func<Task<T>> func, string operation)
         {
             try
             {
                 Events.ExecutingOperation(operation, this.WorkloadUri.ToString());
-                T result = await ExecuteWithRetry(func, r => Events.RetryingOperation(operation, this.WorkloadUri.ToString(), r), this.transientErrorDetectionStrategy);
+                T result = await ExecuteWithRetry(func, r => Events.RetryingOperation(operation, this.WorkloadUri.ToString(), r), this.transientErrorDetectionStrategy)
+                    .TimeoutAfter(this.operationTimeout);
                 Events.SuccessfullyExecutedOperation(operation, this.WorkloadUri.ToString());
                 return result;
             }
@@ -59,6 +68,8 @@ protected async Task<T> Execute<T>(Func<Task<T>> func, string operation)
             }
         }
 
+        protected abstract void HandleException(Exception ex, string operation);
+
         static Task<T> ExecuteWithRetry<T>(Func<Task<T>> func, Action<RetryingEventArgs> onRetry, ITransientErrorDetectionStrategy transientErrorDetection)
         {
             var transientRetryPolicy = new RetryPolicy(transientErrorDetection, TransientRetryStrategy);

diff --git a/edge-util/src/Microsoft.Azure.Devices.Edge.Util/edged/version_2018_06_28/WorkloadClient.cs b/edge-util/src/Microsoft.Azure.Devices.Edge.Util/edged/version_2018_06_28/WorkloadClient.cs
@@ -15,7 +15,12 @@ class WorkloadClient : WorkloadClientVersioned
             new ExponentialBackoff(retryCount: 3, minBackoff: TimeSpan.FromSeconds(2), maxBackoff: TimeSpan.FromSeconds(30), deltaBackoff: TimeSpan.FromSeconds(3));
 
         public WorkloadClient(Uri serverUri, ApiVersion apiVersion, string moduleId, string moduleGenerationId)
-            : base(serverUri, apiVersion, moduleId, moduleGenerationId, new ErrorDetectionStrategy())
+            : this(serverUri, apiVersion, moduleId, moduleGenerationId, Option.None<TimeSpan>())
+        {
+        }
+
+        internal WorkloadClient(Uri serverUri, ApiVersion apiVersion, string moduleId, string moduleGenerationId, Option<TimeSpan> operationTimeout)
+            : base(serverUri, apiVersion, moduleId, moduleGenerationId, new ErrorDetectionStrategy(), operationTimeout)
         {
         }
 

diff --git a/edge-util/src/Microsoft.Azure.Devices.Edge.Util/edged/version_2019_01_30/WorkloadClient.cs b/edge-util/src/Microsoft.Azure.Devices.Edge.Util/edged/version_2019_01_30/WorkloadClient.cs
@@ -12,7 +12,12 @@ namespace Microsoft.Azure.Devices.Edge.Util.Edged.Version_2019_01_30
     class WorkloadClient : WorkloadClientVersioned
     {
         public WorkloadClient(Uri serverUri, ApiVersion apiVersion, string moduleId, string moduleGenerationId)
-            : base(serverUri, apiVersion, moduleId, moduleGenerationId, new ErrorDetectionStrategy())
+            : this(serverUri, apiVersion, moduleId, moduleGenerationId, Option.None<TimeSpan>())
+        {
+        }
+
+        internal WorkloadClient(Uri serverUri, ApiVersion apiVersion, string moduleId, string moduleGenerationId, Option<TimeSpan> operationTimeout)
+            : base(serverUri, apiVersion, moduleId, moduleGenerationId, new ErrorDetectionStrategy(), operationTimeout)
         {
         }
 

diff --git a/edge-util/test/Microsoft.Azure.Devices.Edge.Util.Test/edged/WorkloadClientTest.cs b/edge-util/test/Microsoft.Azure.Devices.Edge.Util.Test/edged/WorkloadClientTest.cs
@@ -144,5 +144,45 @@ public async Task GetTrustBundleAsyncShouldReturnCert(string workloadApiVersion,
             string response = await workloadClient.GetTrustBundleAsync();
             Assert.Equal(new X509Certificate2(Encoding.UTF8.GetBytes(TestCertificateHelper.CertificatePem)), new X509Certificate2(Encoding.UTF8.GetBytes(response)));
         }
+
+        [Fact]
+        public async Task ExecuteTimeoutTest_Version_2018_06_28()
+        {
+            // Arrange
+            var client = new Util.Edged.Version_2018_06_28.WorkloadClient(this.serverUri, ApiVersion.Version20180628, "m1", Guid.NewGuid().ToString(), Option.Some(TimeSpan.FromSeconds(10)));
+
+            async Task<int> LongOperation()
+            {
+                await Task.Delay(TimeSpan.FromHours(1));
+                return 10;
+            }
+
+            // Act
+            Task assertTask = Assert.ThrowsAsync<TimeoutException>(() => client.Execute<int>(LongOperation, "Dummy"));
+            Task delayTask = Task.Delay(TimeSpan.FromSeconds(20));
+
+            Task completedTask = await Task.WhenAny(assertTask, delayTask);
+            Assert.Equal(assertTask, completedTask);
+        }
+
+        [Fact]
+        public async Task ExecuteTimeoutTest_Version_2019_01_30()
+        {
+            // Arrange
+            var client = new Util.Edged.Version_2019_01_30.WorkloadClient(this.serverUri, ApiVersion.Version20190130, "m1", Guid.NewGuid().ToString(), Option.Some(TimeSpan.FromSeconds(10)));
+
+            async Task<int> LongOperation()
+            {
+                await Task.Delay(TimeSpan.FromHours(1));
+                return 10;
+            }
+
+            // Act
+            Task assertTask = Assert.ThrowsAsync<TimeoutException>(() => client.Execute<int>(LongOperation, "Dummy"));
+            Task delayTask = Task.Delay(TimeSpan.FromSeconds(20));
+
+            Task completedTask = await Task.WhenAny(assertTask, delayTask);
+            Assert.Equal(assertTask, completedTask);
+        }
     }
 }