Fix edge agent connection (#3172) (#3173)

* Recreate connection after every exception
Azure · Jul 13, 2020 · 23ffb26 · 23ffb26
1 parent 0569489
commit 23ffb26
Show file tree

Hide file tree

Showing 3 changed files with 26 additions and 7 deletions.
diff --git a/edge-agent/src/Microsoft.Azure.Devices.Edge.Agent.IoTHub/EdgeAgentConnection.cs b/edge-agent/src/Microsoft.Azure.Devices.Edge.Agent.IoTHub/EdgeAgentConnection.cs
@@ -264,7 +264,7 @@ async Task<Twin> GetTwinFunc()
             {
                 Events.ErrorGettingTwin(e);
 
-                if (!retrying && moduleClient != null && !(e is TimeoutException))
+                if (!retrying && moduleClient != null)
                 {
                     try
                     {

diff --git a/...-agent/src/Microsoft.Azure.Devices.Edge.Agent.IoTHub/sdkClient/WrappingSdkModuleClient.cs b/...-agent/src/Microsoft.Azure.Devices.Edge.Agent.IoTHub/sdkClient/WrappingSdkModuleClient.cs
@@ -1,6 +1,7 @@
 // Copyright (c) Microsoft. All rights reserved.
 namespace Microsoft.Azure.Devices.Edge.Agent.IoTHub.SdkClient
 {
+    using System;
     using System.Threading;
     using System.Threading.Tasks;
     using Microsoft.Azure.Devices.Client;
@@ -14,7 +15,18 @@ public class WrappingSdkModuleClient : ISdkModuleClient
         public WrappingSdkModuleClient(ModuleClient sdkModuleClient)
             => this.sdkModuleClient = Preconditions.CheckNotNull(sdkModuleClient, nameof(sdkModuleClient));
 
-        public Task OpenAsync() => this.sdkModuleClient.OpenAsync();
+        public Task OpenAsync()
+        {
+            try
+            {
+                return this.sdkModuleClient.OpenAsync();
+            }
+            catch (Exception)
+            {
+                this.sdkModuleClient?.Dispose();
+                throw;
+            }
+        }
 
         public void SetConnectionStatusChangesHandler(ConnectionStatusChangesHandler statusChangesHandler)
             => this.sdkModuleClient.SetConnectionStatusChangesHandler(statusChangesHandler);
@@ -49,6 +61,10 @@ public Task UpdateReportedPropertiesAsync(TwinCollection reportedProperties)
         ////    return await EdgeClientWebSocket.Connect(deviceStreamRequest.Url, deviceStreamRequest.AuthorizationToken, cancellationToken);
         ////}
 
-        public Task CloseAsync() => this.sdkModuleClient.CloseAsync();
+        public Task CloseAsync()
+        {
+            this.sdkModuleClient.Dispose();
+            return Task.CompletedTask;
+        }
     }
 }
diff --git a/edge-agent/test/Microsoft.Azure.Devices.Edge.Agent.IoTHub.Test/EdgeAgentConnectionTest.cs b/edge-agent/test/Microsoft.Azure.Devices.Edge.Agent.IoTHub.Test/EdgeAgentConnectionTest.cs
@@ -1525,9 +1525,11 @@ public async Task GetTwinRetryLogicGetsNewClient()
             }
         }
 
-        [Fact]
+        [Theory]
         [Unit]
-        public async Task GetDeploymentConfigInfoAsync_CreateNewModuleClientWhenGetTwinThrowsMoreThanRetryCount()
+        [InlineData(typeof(InvalidOperationException))]
+        [InlineData(typeof(TimeoutException))]
+        public async Task GetDeploymentConfigInfoAsync_CreateNewModuleClientWhenGetTwinThrowsMoreThanRetryCount(Type thrownException)
         {
             // Arrange
             var moduleClient = new Mock<IModuleClient>();
@@ -1580,9 +1582,10 @@ public async Task GetDeploymentConfigInfoAsync_CreateNewModuleClientWhenGetTwinT
                 }
             };
 
+            var ex = Activator.CreateInstance(thrownException, "msg str") as Exception;
             moduleClient.SetupSequence(d => d.GetTwinAsync())
-                .ThrowsAsync(new InvalidOperationException())
-                .ThrowsAsync(new InvalidOperationException())
+                .ThrowsAsync(ex)
+                .ThrowsAsync(ex)
                 .ReturnsAsync(twin);
             moduleClient.Setup(d => d.SetDesiredPropertyUpdateCallbackAsync(It.IsAny<DesiredPropertyUpdateCallback>()))
                 .Returns(Task.CompletedTask);