Fix edgeHub shutdown for renew certificate (#6037)

* fix shutdown and E2E test
Azure · Feb 7, 2022 · fcd4d00 · fcd4d00
1 parent cfeea7d
commit fcd4d00
Show file tree

Hide file tree

Showing 8 changed files with 94 additions and 18 deletions.
diff --git a/edge-hub/core/src/Microsoft.Azure.Devices.Edge.Hub.Amqp/AmqpProtocolHead.cs b/edge-hub/core/src/Microsoft.Azure.Devices.Edge.Hub.Amqp/AmqpProtocolHead.cs
@@ -116,7 +116,6 @@ public async Task CloseAsync(CancellationToken token)
 
         public void Dispose()
         {
-            this.CloseAsync(CancellationToken.None).Wait();
         }
 
         void OnAcceptTransport(TransportListener transportListener, TransportAsyncCallbackArgs args)

diff --git a/edge-hub/core/src/Microsoft.Azure.Devices.Edge.Hub.Mqtt/MqttProtocolHead.cs b/edge-hub/core/src/Microsoft.Azure.Devices.Edge.Hub.Mqtt/MqttProtocolHead.cs
@@ -50,6 +50,8 @@ public class MqttProtocolHead : IProtocolHead
 
         IChannel serverChannel;
         IEventLoopGroup eventLoopGroup;
+        IEventLoopGroup wsEventLoopGroup;
+        IEventLoopGroup parentEventLoopGroup;
 
         public MqttProtocolHead(
             ISettingsProvider settingsProvider,
@@ -106,13 +108,15 @@ public async Task CloseAsync(CancellationToken token)
         {
             try
             {
-                this.logger.LogInformation("Stopping");
+                this.logger.LogInformation("Stopping MQTT protocol head");
 
                 await (this.serverChannel?.CloseAsync() ?? TaskEx.Done);
                 await (this.eventLoopGroup?.ShutdownGracefullyAsync() ?? TaskEx.Done);
+                await (this.parentEventLoopGroup?.ShutdownGracefullyAsync() ?? TaskEx.Done);
+                await (this.wsEventLoopGroup?.ShutdownGracefullyAsync() ?? TaskEx.Done);
                 // TODO: gracefully shutdown the MultithreadEventLoopGroup in MqttWebSocketListener?
                 // TODO: this.webSocketListenerRegistry.TryUnregister("mqtts")?
-                this.logger.LogInformation("Stopped");
+                this.logger.LogInformation("Stopped MQTT protocol head");
             }
             catch (Exception ex)
             {
@@ -123,7 +127,6 @@ public async Task CloseAsync(CancellationToken token)
         public void Dispose()
         {
             this.mqttConnectionProvider.Dispose();
-            this.CloseAsync(CancellationToken.None).Wait();
         }
 
         ServerBootstrap SetupServerBootstrap()
@@ -138,11 +141,10 @@ ServerBootstrap SetupServerBootstrap()
 
             var bootstrap = new ServerBootstrap();
             // multithreaded event loop that handles the incoming connection
-            IEventLoopGroup parentEventLoopGroup = new MultithreadEventLoopGroup(parentEventLoopCount);
+            this.parentEventLoopGroup = new MultithreadEventLoopGroup(parentEventLoopCount);
             // multithreaded event loop (worker) that handles the traffic of the accepted connections
             this.eventLoopGroup = new MultithreadEventLoopGroup(threadCount);
-
-            bootstrap.Group(parentEventLoopGroup, this.eventLoopGroup)
+            bootstrap.Group(this.parentEventLoopGroup, this.eventLoopGroup)
                 .Option(ChannelOption.SoBacklog, listenBacklogSize)
                 // Allow listening socket to force bind to port if previous socket is still in TIME_WAIT
                 // Fixes "address is already in use" errors
@@ -185,14 +187,15 @@ ServerBootstrap SetupServerBootstrap()
                                     bridgeFactory));
                         }));
 
+            this.wsEventLoopGroup = new MultithreadEventLoopGroup(Environment.ProcessorCount);
             var mqttWebSocketListener = new MqttWebSocketListener(
                 settings,
                 bridgeFactory,
                 this.authenticator,
                 this.usernameParser,
                 this.clientCredentialsFactory,
                 () => this.sessionProvider,
-                new MultithreadEventLoopGroup(Environment.ProcessorCount),
+                this.wsEventLoopGroup,
                 this.byteBufferAllocator,
                 AutoRead,
                 maxInboundMessageSize,

diff --git a/...icrosoft.Azure.Devices.Edge.Hub.MqttBrokerAdapter/authentication/AuthAgentProtocolHead.cs b/...icrosoft.Azure.Devices.Edge.Hub.MqttBrokerAdapter/authentication/AuthAgentProtocolHead.cs
@@ -96,7 +96,6 @@ await hostToStop.Match(
 
         public void Dispose()
         {
-            this.DisposeAsync().ConfigureAwait(false).GetAwaiter().GetResult();
         }
 
         public async ValueTask DisposeAsync()

diff --git a/...osoft.Azure.Devices.Edge.Hub.MqttBrokerAdapter/brokerConnection/MqttBrokerProtocolHead.cs b/...osoft.Azure.Devices.Edge.Hub.MqttBrokerAdapter/brokerConnection/MqttBrokerProtocolHead.cs
@@ -73,7 +73,9 @@ public async Task CloseAsync(CancellationToken token)
             Events.Closed();
         }
 
-        public void Dispose() => this.CloseAsync(CancellationToken.None).Wait();
+        public void Dispose()
+        {
+        }
 
         static class Events
         {

diff --git a/edge-hub/core/src/Microsoft.Azure.Devices.Edge.Hub.Service/CertificateRenewal.cs b/edge-hub/core/src/Microsoft.Azure.Devices.Edge.Hub.Service/CertificateRenewal.cs
@@ -8,19 +8,21 @@ namespace Microsoft.Azure.Devices.Edge.Hub.Service
 
     public class CertificateRenewal : IDisposable
     {
-        static readonly TimeSpan MaxRenewAfter = TimeSpan.FromMilliseconds(int.MaxValue);
+        static readonly TimeSpan DefaultMaxRenewAfter = TimeSpan.FromMilliseconds(int.MaxValue);
         static readonly TimeSpan TimeBuffer = TimeSpan.FromMinutes(5);
 
+        readonly TimeSpan maxRenewAfter;
         readonly EdgeHubCertificates certificates;
         readonly ILogger logger;
         readonly Timer timer;
         readonly CancellationTokenSource cts;
 
-        public CertificateRenewal(EdgeHubCertificates certificates, ILogger logger)
+        public CertificateRenewal(EdgeHubCertificates certificates, ILogger logger, TimeSpan maxRenewAfter)
         {
             this.certificates = Preconditions.CheckNotNull(certificates, nameof(certificates));
             this.logger = Preconditions.CheckNotNull(logger, nameof(logger));
             this.cts = new CancellationTokenSource();
+            this.maxRenewAfter = maxRenewAfter;
 
             TimeSpan timeToExpire = certificates.ServerCertificate.NotAfter - DateTime.UtcNow;
             if (timeToExpire > TimeBuffer)
@@ -29,8 +31,8 @@ public CertificateRenewal(EdgeHubCertificates certificates, ILogger logger)
                 // This is the maximum value for the timer (~24 days)
                 // Math.Min unfortunately doesn't work with TimeSpans so we need to do the check manually
                 TimeSpan renewAfter = timeToExpire - (TimeBuffer / 2);
-                TimeSpan clamped = renewAfter > MaxRenewAfter
-                    ? MaxRenewAfter
+                TimeSpan clamped = renewAfter > this.maxRenewAfter
+                    ? this.maxRenewAfter
                     : renewAfter;
                 logger.LogInformation("Scheduling server certificate renewal for {0}.", DateTime.UtcNow.Add(renewAfter).ToString("o"));
                 logger.LogDebug("Scheduling server certificate renewal timer for {0} (clamped to Int32.MaxValue).", DateTime.UtcNow.Add(clamped).ToString("o"));
@@ -73,7 +75,7 @@ protected virtual void Dispose(bool disposing)
         void Callback(object _state)
         {
             TimeSpan timeToExpire = this.certificates.ServerCertificate.NotAfter - DateTime.UtcNow;
-            if (timeToExpire > TimeBuffer)
+            if (timeToExpire > TimeBuffer && this.maxRenewAfter == DefaultMaxRenewAfter)
             {
                 // Timer has expired but is not within the time window for renewal
                 // Reschedule the timer.
@@ -82,8 +84,8 @@ void Callback(object _state)
                 // This is the maximum value for the timer (~24 days)
                 // Math.Min unfortunately doesn't work with TimeSpans so we need to do the check manually
                 TimeSpan renewAfter = timeToExpire - (TimeBuffer / 2);
-                TimeSpan clamped = renewAfter > MaxRenewAfter
-                    ? MaxRenewAfter
+                TimeSpan clamped = renewAfter > this.maxRenewAfter
+                    ? this.maxRenewAfter
                     : renewAfter;
                 this.logger.LogDebug("Scheduling server certificate renewal timer for {0}.", DateTime.UtcNow.Add(clamped).ToString("o"));
                 this.timer.Change(clamped, Timeout.InfiniteTimeSpan);

diff --git a/edge-hub/core/src/Microsoft.Azure.Devices.Edge.Hub.Service/Program.cs b/edge-hub/core/src/Microsoft.Azure.Devices.Edge.Hub.Service/Program.cs
@@ -122,9 +122,12 @@ static async Task<int> MainAsync(IConfigurationRoot configuration)
             TimeSpan shutdownWaitPeriod = TimeSpan.FromSeconds(configuration.GetValue("ShutdownWaitPeriod", DefaultShutdownWaitPeriod));
             (CancellationTokenSource cts, ManualResetEventSlim completed, Option<object> handler) = ShutdownHandler.Init(shutdownWaitPeriod, logger);
 
+            double renewAfter = configuration.GetValue("ServerCertificateRenewAfterInMs", int.MaxValue);
+            renewAfter = renewAfter > int.MaxValue ? int.MaxValue : renewAfter;
+            TimeSpan maxRenewAfter = TimeSpan.FromMilliseconds(renewAfter);
             using (IProtocolHead mqttBrokerProtocolHead = await GetMqttBrokerProtocolHeadAsync(experimentalFeatures, container))
             using (IProtocolHead edgeHubProtocolHead = await GetEdgeHubProtocolHeadAsync(logger, configuration, experimentalFeatures, container, hosting))
-            using (var renewal = new CertificateRenewal(certificates, logger))
+            using (var renewal = new CertificateRenewal(certificates, logger, maxRenewAfter))
             {
                 try
                 {

diff --git a/...ore/test/Microsoft.Azure.Devices.Edge.Hub.MqttBrokerAdapter.Test/AuthAgentListenerTest.cs b/...ore/test/Microsoft.Azure.Devices.Edge.Hub.MqttBrokerAdapter.Test/AuthAgentListenerTest.cs
@@ -59,6 +59,8 @@ public async Task StartsUpAndServes()
                 dynamic response = await PostAsync(content, this.url);
 
                 Assert.Equal(200, (int)response.result);
+
+                await sut.CloseAsync(CancellationToken.None);
             }
         }
 
@@ -71,6 +73,7 @@ public async Task CannotStartTwice()
             {
                 await sut.StartAsync();
                 await Assert.ThrowsAsync<InvalidOperationException>(async () => await sut.StartAsync());
+                await sut.CloseAsync(CancellationToken.None);
             }
         }
 
@@ -90,6 +93,7 @@ public async Task DeniesNoPasswordNorCertificate()
                 dynamic response = await PostAsync(content, this.url);
 
                 Assert.Equal(403, (int)response.result);
+                await sut.CloseAsync(CancellationToken.None);
             }
         }
 
@@ -112,6 +116,8 @@ public async Task DeniesBothPasswordAndCertificate()
                 dynamic response = await PostAsync(content, this.url);
 
                 Assert.Equal(403, (int)response.result);
+
+                await sut.CloseAsync(CancellationToken.None);
             }
         }
 
@@ -132,6 +138,8 @@ public async Task DeniesBadCertificateFormat()
                 dynamic response = await PostAsync(content, this.url);
 
                 Assert.Equal(403, (int)response.result);
+
+                await sut.CloseAsync(CancellationToken.None);
             }
         }
 
@@ -152,6 +160,8 @@ public async Task DeniesNoVersion()
                 dynamic response = await PostAsync(content, this.url);
 
                 Assert.Equal(403, (int)response.result);
+
+                await sut.CloseAsync(CancellationToken.None);
             }
         }
 
@@ -173,6 +183,8 @@ public async Task DeniesBadVersion()
                 dynamic response = await PostAsync(content, this.url);
 
                 Assert.Equal(403, (int)response.result);
+
+                await sut.CloseAsync(CancellationToken.None);
             }
         }
 
@@ -199,6 +211,8 @@ public async Task AcceptsGoodTokenDeniesBadToken()
 
                 response = await PostAsync(content, this.url);
                 Assert.Equal(200, (int)response.result);
+
+                await sut.CloseAsync(CancellationToken.None);
             }
         }
 
@@ -224,6 +238,8 @@ public async Task AcceptsGoodThumbprintDeniesBadThumbprint()
 
                 response = await PostAsync(content, this.url);
                 Assert.Equal(200, (int)response.result);
+
+                await sut.CloseAsync(CancellationToken.None);
             }
         }
 
@@ -254,6 +270,8 @@ public async Task AcceptsGoodCaDeniesBadCa()
 
                 response = await PostAsync(content, this.url);
                 Assert.Equal(200, (int)response.result);
+
+                await sut.CloseAsync(CancellationToken.None);
             }
         }
 
@@ -275,6 +293,8 @@ public async Task ReturnsDeviceIdentity()
                 var response = await PostAsync(content, this.url);
                 Assert.Equal(200, (int)response.result);
                 Assert.Equal("testhub/device", (string)response.identity);
+
+                await sut.CloseAsync(CancellationToken.None);
             }
         }
 
@@ -296,6 +316,8 @@ public async Task ReturnsModuleIdentity()
                 var response = await PostAsync(content, this.url);
                 Assert.Equal(200, (int)response.result);
                 Assert.Equal("testhub/device/module", (string)response.identity);
+
+                await sut.CloseAsync(CancellationToken.None);
             }
         }
 
@@ -310,6 +332,8 @@ public async Task AcceptsRequestWithContentLength()
                 var result = await SendDirectRequest(RequestBody);
 
                 Assert.StartsWith(@"{""result"":200,", result);
+
+                await sut.CloseAsync(CancellationToken.None);
             }
         }
 
@@ -324,6 +348,8 @@ public async Task AcceptsRequestWithNoContentLength()
                 var result = await SendDirectRequest(RequestBody, withContentLength: false);
 
                 Assert.StartsWith(@"{""result"":200,", result);
+
+                await sut.CloseAsync(CancellationToken.None);
             }
         }
 
@@ -338,6 +364,8 @@ public async Task DeniesMalformedJsonRequest()
                 var result = await SendDirectRequest(NonJSONRequestBody);
 
                 Assert.StartsWith(@"{""result"":403,", result);
+
+                await sut.CloseAsync(CancellationToken.None);
             }
         }
 
@@ -379,6 +407,8 @@ public async Task StoresMetadataCorrectly()
                 var modelId = (await metadataStore.GetMetadata("device")).ModelId;
                 Assert.True(modelId.HasValue);
                 Assert.Equal(modelIdString, modelId.GetOrElse("impossibleValue"));
+
+                await sut.CloseAsync(CancellationToken.None);
             }
         }
 

diff --git a/test/Microsoft.Azure.Devices.Edge.Test/Module.cs b/test/Microsoft.Azure.Devices.Edge.Test/Module.cs
@@ -17,6 +17,44 @@ public class Module : SasManualProvisioningFixture
         const string SensorName = "tempSensor";
         const string DefaultSensorImage = "mcr.microsoft.com/azureiotedge-simulated-temperature-sensor:1.0";
 
+        [Test]
+        [Category("CentOsSafe")]
+        public async Task CertRenew()
+        {
+            CancellationToken token = this.TestToken;
+
+            EdgeDeployment deployment = await this.runtime.DeployConfigurationAsync(
+                    builder =>
+                    {
+                         builder.GetModule("$edgeHub").WithEnvironment(("ServerCertificateRenewAfterInMs", "6000"));
+                    },
+                    token,
+                    Context.Current.NestedEdge);
+
+            EdgeModule edgeHub = deployment.Modules[ModuleName.EdgeHub];
+            await edgeHub.WaitForStatusAsync(EdgeModuleStatus.Running, token);
+            EdgeModule edgeAgent = deployment.Modules[ModuleName.EdgeAgent];
+            // certificate renew should stop edgeHub and then it should be started by edgeAgent
+            await edgeAgent.WaitForReportedPropertyUpdatesAsync(
+                new
+                {
+                    properties = new
+                    {
+                        reported = new
+                        {
+                            systemModules = new
+                            {
+                                edgeHub = new
+                                {
+                                    restartCount = 1
+                                }
+                            }
+                        }
+                    }
+                },
+                token);
+        }
+
         [Test]
         [Category("CentOsSafe")]
         [Category("nestededge_isa95")]