Remove device connection will always remove cloud connection (#3100)

Add environment variable to control closing cloud connection when close device connection in connection manager
Azure · Jun 16, 2020 · fb33e62 · fb33e62
1 parent cca8049
commit fb33e62
Show file tree

Hide file tree

Showing 5 changed files with 96 additions and 7 deletions.
diff --git a/edge-hub/src/Microsoft.Azure.Devices.Edge.Hub.Core/ConnectionManager.cs b/edge-hub/src/Microsoft.Azure.Devices.Edge.Hub.Core/ConnectionManager.cs
@@ -29,13 +29,15 @@ public class ConnectionManager : IConnectionManager
         readonly ICredentialsCache credentialsCache;
         readonly IIdentityProvider identityProvider;
         readonly IDeviceConnectivityManager connectivityManager;
+        readonly bool closeCloudConnectionWhenCloseDeviceConnection;
 
         public ConnectionManager(
             ICloudConnectionProvider cloudConnectionProvider,
             ICredentialsCache credentialsCache,
             IIdentityProvider identityProvider,
             IDeviceConnectivityManager connectivityManager,
-            int maxClients = DefaultMaxClients)
+            int maxClients = DefaultMaxClients,
+            bool closeCloudConnectionWhenCloseDeviceConnection = true)
         {
             this.cloudConnectionProvider = Preconditions.CheckNotNull(cloudConnectionProvider, nameof(cloudConnectionProvider));
             this.maxClients = Preconditions.CheckRange(maxClients, 1, nameof(maxClients));
@@ -44,6 +46,7 @@ public ConnectionManager(
             this.connectivityManager = Preconditions.CheckNotNull(connectivityManager, nameof(connectivityManager));
             this.connectivityManager.DeviceDisconnected += (o, args) => this.HandleDeviceCloudConnectionDisconnected();
             Util.Metrics.MetricsV0.RegisterGaugeCallback(() => MetricsV0.SetConnectedClientCountGauge(this));
+            this.closeCloudConnectionWhenCloseDeviceConnection = closeCloudConnectionWhenCloseDeviceConnection;
         }
 
         public event EventHandler<IIdentity> CloudConnectionEstablished;
@@ -75,7 +78,7 @@ await currentDeviceConnection
         public Task RemoveDeviceConnection(string id)
         {
             return this.devices.TryGetValue(Preconditions.CheckNonWhiteSpace(id, nameof(id)), out ConnectedDevice device)
-                ? this.RemoveDeviceConnection(device, false)
+                ? this.RemoveDeviceConnection(device, this.closeCloudConnectionWhenCloseDeviceConnection)
                 : Task.CompletedTask;
         }
 
@@ -178,6 +181,7 @@ static Try<ICloudProxy> GetCloudProxyFromCloudConnection(Try<ICloudConnection> c
 
         async Task RemoveDeviceConnection(ConnectedDevice device, bool removeCloudConnection)
         {
+            Events.RemovingDeviceConnection(device.Identity.Id, removeCloudConnection);
             await device.DeviceConnection.Filter(dp => dp.IsActive)
                 .ForEachAsync(dp => dp.CloseAsync(new EdgeHubConnectionException($"Connection closed for device {device.Identity.Id}.")));
 
@@ -250,7 +254,7 @@ await clientCredentials.ForEachAsync(
                                 }
                                 else
                                 {
-                                    await this.RemoveDeviceConnection(device, false);
+                                    await this.RemoveDeviceConnection(device, this.closeCloudConnectionWhenCloseDeviceConnection);
                                 }
                             });
                     }
@@ -466,6 +470,7 @@ enum EventIds
             {
                 CreateNewCloudConnection = IdStart,
                 NewDeviceConnection,
+                RemovingDeviceConnection,
                 RemoveDeviceConnection,
                 CreateNewCloudConnectionError,
                 ObtainedCloudConnection,
@@ -495,6 +500,11 @@ public static void NewDeviceConnection(IIdentity identity)
                 Log.LogInformation((int)EventIds.NewDeviceConnection, Invariant($"New device connection for device {identity.Id}"));
             }
 
+            public static void RemovingDeviceConnection(string id, bool removeCloudConnection)
+            {
+                Log.LogInformation((int)EventIds.RemovingDeviceConnection, Invariant($"Removing device connection for device {id} with removeCloudConnection flag '{removeCloudConnection}'."));
+            }
+
             public static void RemoveDeviceConnection(string id)
             {
                 Log.LogInformation((int)EventIds.RemoveDeviceConnection, Invariant($"Device connection removed for device {id}"));

diff --git a/edge-hub/src/Microsoft.Azure.Devices.Edge.Hub.Service/DependencyManager.cs b/edge-hub/src/Microsoft.Azure.Devices.Edge.Hub.Service/DependencyManager.cs
@@ -151,6 +151,7 @@ void RegisterRoutingModule(
             bool encryptTwinStore = this.configuration.GetValue("EncryptTwinStore", false);
             int configUpdateFrequencySecs = this.configuration.GetValue("ConfigRefreshFrequencySecs", 3600);
             TimeSpan configUpdateFrequency = TimeSpan.FromSeconds(configUpdateFrequencySecs);
+            bool closeCloudConnectionWhenCloseDeviceConnection = this.configuration.GetValue("CloseCloudConnectionWhenCloseDeviceConnection", true);
 
             builder.RegisterModule(
                 new RoutingModule(
@@ -179,7 +180,8 @@ void RegisterRoutingModule(
                     upstreamFanOutFactor,
                     encryptTwinStore,
                     configUpdateFrequency,
-                    experimentalFeatures));
+                    experimentalFeatures,
+                    closeCloudConnectionWhenCloseDeviceConnection));
         }
 
         void RegisterCommonModule(

diff --git a/edge-hub/src/Microsoft.Azure.Devices.Edge.Hub.Service/modules/RoutingModule.cs b/edge-hub/src/Microsoft.Azure.Devices.Edge.Hub.Service/modules/RoutingModule.cs
@@ -53,6 +53,7 @@ public class RoutingModule : Module
         readonly bool encryptTwinStore;
         readonly TimeSpan configUpdateFrequency;
         readonly ExperimentalFeatures experimentalFeatures;
+        readonly bool closeCloudConnectionWhenCloseDeviceConnection;
 
         public RoutingModule(
             string iotHubName,
@@ -80,7 +81,8 @@ public RoutingModule(
             int upstreamFanOutFactor,
             bool encryptTwinStore,
             TimeSpan configUpdateFrequency,
-            ExperimentalFeatures experimentalFeatures)
+            ExperimentalFeatures experimentalFeatures,
+            bool closeCloudConnectionWhenCloseDeviceConnection)
         {
             this.iotHubName = Preconditions.CheckNonWhiteSpace(iotHubName, nameof(iotHubName));
             this.edgeDeviceId = Preconditions.CheckNonWhiteSpace(edgeDeviceId, nameof(edgeDeviceId));
@@ -108,6 +110,7 @@ public RoutingModule(
             this.encryptTwinStore = encryptTwinStore;
             this.configUpdateFrequency = configUpdateFrequency;
             this.experimentalFeatures = experimentalFeatures;
+            this.closeCloudConnectionWhenCloseDeviceConnection = closeCloudConnectionWhenCloseDeviceConnection;
         }
 
         protected override void Load(ContainerBuilder builder)
@@ -255,7 +258,8 @@ protected override void Load(ContainerBuilder builder)
                             credentialsCache,
                             identityProvider,
                             deviceConnectivityManager,
-                            this.maxConnectedClients);
+                            this.maxConnectedClients,
+                            this.closeCloudConnectionWhenCloseDeviceConnection);
                         return connectionManager;
                     })
                 .As<Task<IConnectionManager>>()

diff --git a/edge-hub/test/Microsoft.Azure.Devices.Edge.Hub.Core.Test/ConnectionManagerTest.cs b/edge-hub/test/Microsoft.Azure.Devices.Edge.Hub.Core.Test/ConnectionManagerTest.cs
@@ -935,6 +935,78 @@ public async Task GetMultipleCloudProxiesTest()
             Mock.Get(client2).Verify(cp => cp.CloseAsync(), Times.Never);
         }
 
+        [Unit]
+        [Theory]
+        [InlineData(true)]
+        [InlineData(false)]
+        public async Task TestCloseDeviceAndCloudConnection(bool closeCloudConnectionWhenCloseDeviceConnection)
+        {
+            // Arrange
+            string edgeDeviceId = "edgeDevice";
+            string module1Id = "module1";
+            string iotHub = "foo.azure-devices.net";
+            string token = TokenHelper.CreateSasToken(iotHub);
+            var module1Credentials = new TokenCredentials(new ModuleIdentity(iotHub, edgeDeviceId, module1Id), token, DummyProductInfo, true);
+
+            IClient client1 = GetDeviceClient();
+            var messageConverterProvider = Mock.Of<IMessageConverterProvider>();
+            var deviceClientProvider = new Mock<IClientProvider>();
+            deviceClientProvider.Setup(d => d.Create(It.IsAny<IIdentity>(), It.IsAny<ITokenProvider>(), It.IsAny<ITransportSettings[]>())).Returns(client1);
+
+            ICredentialsCache credentialsCache = new CredentialsCache(new NullCredentialsCache());
+            await credentialsCache.Add(module1Credentials);
+            var productInfoStore = Mock.Of<IProductInfoStore>();
+            var cloudConnectionProvider = new CloudConnectionProvider(
+                messageConverterProvider,
+                1,
+                deviceClientProvider.Object,
+                Option.None<UpstreamProtocol>(),
+                Mock.Of<ITokenProvider>(),
+                Mock.Of<IDeviceScopeIdentitiesCache>(),
+                credentialsCache,
+                new ModuleIdentity(iotHub, edgeDeviceId, "$edgeHub"),
+                TimeSpan.FromMinutes(60),
+                true,
+                TimeSpan.FromSeconds(20),
+                false,
+                Option.None<IWebProxy>(),
+                productInfoStore);
+            cloudConnectionProvider.BindEdgeHub(Mock.Of<IEdgeHub>());
+            var deviceConnectivityManager = Mock.Of<IDeviceConnectivityManager>();
+
+            var module1Identity = Mock.Of<IModuleIdentity>(m => m.Id == module1Credentials.Identity.Id);
+            var moduleProxy1 = Mock.Of<IDeviceProxy>(m => m.IsActive);
+            IConnectionManager connectionManager = new ConnectionManager(
+                cloudConnectionProvider,
+                credentialsCache,
+                GetIdentityProvider(),
+                deviceConnectivityManager,
+                closeCloudConnectionWhenCloseDeviceConnection: closeCloudConnectionWhenCloseDeviceConnection);
+            await connectionManager.AddDeviceConnection(module1Identity, moduleProxy1);
+
+            // Act
+            Option<ICloudProxy> getCloudProxyTask = await connectionManager.GetCloudConnection(module1Credentials.Identity.Id);
+
+            // Assert
+            Assert.True(getCloudProxyTask.HasValue);
+            Assert.True(getCloudProxyTask.OrDefault().IsActive);
+
+            // Act
+            await connectionManager.RemoveDeviceConnection(module1Credentials.Identity.Id);
+
+            // Assert
+            if (closeCloudConnectionWhenCloseDeviceConnection)
+            {
+                Assert.False(getCloudProxyTask.OrDefault().IsActive);
+                Mock.Get(client1).Verify(cp => cp.CloseAsync(), Times.Once);
+            }
+            else
+            {
+                Assert.True(getCloudProxyTask.OrDefault().IsActive);
+                Mock.Get(client1).Verify(cp => cp.CloseAsync(), Times.Never);
+            }
+        }
+
         static ICloudConnection GetCloudConnectionMock()
         {
             ICloudProxy cloudProxyMock = GetCloudProxyMock();

diff --git a/edge-hub/test/Microsoft.Azure.Devices.Edge.Hub.E2E.Test/DependencyManager.cs b/edge-hub/test/Microsoft.Azure.Devices.Edge.Hub.E2E.Test/DependencyManager.cs
@@ -153,7 +153,8 @@ public void Register(ContainerBuilder builder)
                     10,
                     false,
                     TimeSpan.FromHours(1),
-                    experimentalFeatures));
+                    experimentalFeatures,
+                    true));
 
             builder.RegisterModule(new HttpModule());
             builder.RegisterModule(new MqttModule(mqttSettingsConfiguration.Object, topics, this.serverCertificate, false, false, false));