From 9ff5a4f86dad3b54715f28056ad9cd02ad372e14 Mon Sep 17 00:00:00 2001 From: yophilav Date: Thu, 6 Aug 2026 17:22:07 +0000 Subject: [PATCH 1/4] Throttle repeated cloud proxy rebuilds Reuse recently completed retry connection attempts per device so queued operations cannot trigger a sequential rebuild stampede when replacements immediately become inactive. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 0b428cdd-2248-4f74-87b1-cd02dc9ad50a --- .../ConnectionManager.cs | 117 +++++++++++++++--- .../ConnectionManagerTest.cs | 56 +++++++++ .../RetryingCloudProxyTest.cs | 33 ++++- 3 files changed, 184 insertions(+), 22 deletions(-) diff --git a/edge-hub/core/src/Microsoft.Azure.Devices.Edge.Hub.Core/ConnectionManager.cs b/edge-hub/core/src/Microsoft.Azure.Devices.Edge.Hub.Core/ConnectionManager.cs index 06a3273b534..1b6f9939fc9 100644 --- a/edge-hub/core/src/Microsoft.Azure.Devices.Edge.Hub.Core/ConnectionManager.cs +++ b/edge-hub/core/src/Microsoft.Azure.Devices.Edge.Hub.Core/ConnectionManager.cs @@ -21,6 +21,7 @@ namespace Microsoft.Azure.Devices.Edge.Hub.Core public class ConnectionManager : IConnectionManager { const int DefaultMaxClients = 101; // 100 Clients + 1 Edgehub + static readonly TimeSpan DefaultCloudConnectionRetryInterval = TimeSpan.FromSeconds(5); readonly object deviceConnLock = new object(); readonly AsyncReaderWriterLock connectToCloudLock = new AsyncReaderWriterLock(); readonly ConcurrentDictionary devices = new ConcurrentDictionary(); @@ -30,6 +31,8 @@ public class ConnectionManager : IConnectionManager readonly IIdentityProvider identityProvider; readonly IDeviceConnectivityManager connectivityManager; readonly bool closeCloudConnectionOnDeviceDisconnect; + readonly TimeSpan cloudConnectionRetryInterval; + readonly ISystemTime systemTime; public ConnectionManager( ICloudConnectionProvider cloudConnectionProvider, @@ -38,12 +41,37 @@ public ConnectionManager( IDeviceConnectivityManager connectivityManager, int maxClients = DefaultMaxClients, bool closeCloudConnectionOnDeviceDisconnect = true) + : this( + cloudConnectionProvider, + credentialsCache, + identityProvider, + connectivityManager, + maxClients, + closeCloudConnectionOnDeviceDisconnect, + DefaultCloudConnectionRetryInterval, + SystemTime.Instance) + { + } + + internal ConnectionManager( + ICloudConnectionProvider cloudConnectionProvider, + ICredentialsCache credentialsCache, + IIdentityProvider identityProvider, + IDeviceConnectivityManager connectivityManager, + int maxClients, + bool closeCloudConnectionOnDeviceDisconnect, + TimeSpan cloudConnectionRetryInterval, + ISystemTime systemTime) { this.cloudConnectionProvider = Preconditions.CheckNotNull(cloudConnectionProvider, nameof(cloudConnectionProvider)); this.maxClients = Preconditions.CheckRange(maxClients, 1, nameof(maxClients)); this.credentialsCache = Preconditions.CheckNotNull(credentialsCache, nameof(credentialsCache)); this.identityProvider = Preconditions.CheckNotNull(identityProvider, nameof(identityProvider)); this.connectivityManager = Preconditions.CheckNotNull(connectivityManager, nameof(connectivityManager)); + this.cloudConnectionRetryInterval = cloudConnectionRetryInterval >= TimeSpan.Zero + ? cloudConnectionRetryInterval + : throw new ArgumentOutOfRangeException(nameof(cloudConnectionRetryInterval)); + this.systemTime = Preconditions.CheckNotNull(systemTime, nameof(systemTime)); this.connectivityManager.DeviceDisconnected += (o, args) => this.HandleDeviceCloudConnectionDisconnected(); this.closeCloudConnectionOnDeviceDisconnect = closeCloudConnectionOnDeviceDisconnect; } @@ -91,27 +119,28 @@ public Option GetDeviceConnection(string id) public async Task> GetCloudConnection(string id) { - Try cloudProxyTry = await this.TryGetCloudConnectionInternal(id); + Try cloudProxyTry = await this.TryGetCloudConnectionInternal(id, false); return cloudProxyTry .Ok() - .Map(c => (ICloudProxy)new RetryingCloudProxy(id, () => this.TryGetCloudConnectionInternal(id), c)); + .Map(c => (ICloudProxy)new RetryingCloudProxy(id, () => this.TryGetCloudConnectionInternal(id, true), c)); } public async Task> TryGetCloudConnection(string id) { - Try cloudProxyTry = await this.TryGetCloudConnectionInternal(id); + Try cloudProxyTry = await this.TryGetCloudConnectionInternal(id, false); return cloudProxyTry.Success - ? Try.Success((ICloudProxy)new RetryingCloudProxy(id, () => this.TryGetCloudConnectionInternal(id), cloudProxyTry.Value)) + ? Try.Success((ICloudProxy)new RetryingCloudProxy(id, () => this.TryGetCloudConnectionInternal(id, true), cloudProxyTry.Value)) : cloudProxyTry; } - async Task> TryGetCloudConnectionInternal(string id) + async Task> TryGetCloudConnectionInternal(string id, bool isRetry) { IIdentity identity = this.identityProvider.Create(Preconditions.CheckNonWhiteSpace(id, nameof(id))); ConnectedDevice device = this.GetOrCreateConnectedDevice(identity); Try cloudConnectionTry = await device.GetOrCreateCloudConnection( - c => this.ConnectToCloud(c.Identity, this.CloudConnectionStatusChangedHandler)); + c => this.ConnectToCloud(c.Identity, this.CloudConnectionStatusChangedHandler), + isRetry); Events.GetCloudConnection(device.Identity, cloudConnectionTry); Try cloudProxyTry = GetCloudProxyFromCloudConnection(cloudConnectionTry, device.Identity); @@ -221,7 +250,7 @@ public async Task> CreateCloudConnectionAsync(IClientCredential Events.NewCloudConnection(credentials.Identity, newCloudConnection); Try cloudProxyTry = GetCloudProxyFromCloudConnection(newCloudConnection, credentials.Identity); return cloudProxyTry.Success - ? Try.Success((ICloudProxy)new RetryingCloudProxy(credentials.Identity.Id, () => this.TryGetCloudConnectionInternal(credentials.Identity.Id), cloudProxyTry.Value)) + ? Try.Success((ICloudProxy)new RetryingCloudProxy(credentials.Identity.Id, () => this.TryGetCloudConnectionInternal(credentials.Identity.Id, true), cloudProxyTry.Value)) : cloudProxyTry; } @@ -235,11 +264,13 @@ public async Task> GetOrCreateCloudConnectionAsync(IClientCrede // instance to this.devices and return that. ConnectedDevice device = this.GetOrCreateConnectedDevice(credentials.Identity); - Try cloudConnectionTry = await device.GetOrCreateCloudConnection((c) => this.CreateOrUpdateCloudConnection(c, credentials)); + Try cloudConnectionTry = await device.GetOrCreateCloudConnection( + c => this.CreateOrUpdateCloudConnection(c, credentials), + false); Events.GetCloudConnection(credentials.Identity, cloudConnectionTry); Try cloudProxyTry = GetCloudProxyFromCloudConnection(cloudConnectionTry, credentials.Identity); return cloudProxyTry.Success - ? Try.Success((ICloudProxy)new RetryingCloudProxy(credentials.Identity.Id, () => this.TryGetCloudConnectionInternal(credentials.Identity.Id), cloudProxyTry.Value)) + ? Try.Success((ICloudProxy)new RetryingCloudProxy(credentials.Identity.Id, () => this.TryGetCloudConnectionInternal(credentials.Identity.Id, true), cloudProxyTry.Value)) : cloudProxyTry; } @@ -388,7 +419,12 @@ ConnectedDevice CreateOrUpdateConnectedDevice(IIdentity identity) return this.devices.AddOrUpdate( deviceId, id => this.CreateNewConnectedDevice(identity), - (id, cd) => new ConnectedDevice(identity, cd.CloudConnection, cd.DeviceConnection)); + (id, cd) => new ConnectedDevice( + identity, + cd.CloudConnection, + cd.DeviceConnection, + this.cloudConnectionRetryInterval, + this.systemTime)); } ConnectedDevice CreateNewConnectedDevice(IIdentity identity) @@ -400,7 +436,7 @@ ConnectedDevice CreateNewConnectedDevice(IIdentity identity) throw new EdgeHubConnectionException($"Edge hub already has maximum allowed clients ({this.maxClients - 1}) connected."); } - return new ConnectedDevice(identity); + return new ConnectedDevice(identity, this.cloudConnectionRetryInterval, this.systemTime); } } @@ -426,18 +462,34 @@ class ConnectedDevice // so using traditional locking mechanism for those. readonly object deviceProxyLock = new object(); readonly AsyncLock cloudConnectionLock = new AsyncLock(); + readonly TimeSpan cloudConnectionRetryInterval; + readonly ISystemTime systemTime; Option>> cloudConnectionCreateTask = Option.None>>(); - - public ConnectedDevice(IIdentity identity) - : this(identity, Option.None(), Option.None()) + Option cloudConnectionCreateCompletedTime = Option.None(); + bool cloudConnectionCreateWasRetry; + + public ConnectedDevice(IIdentity identity, TimeSpan cloudConnectionRetryInterval, ISystemTime systemTime) + : this( + identity, + Option.None(), + Option.None(), + cloudConnectionRetryInterval, + systemTime) { } - public ConnectedDevice(IIdentity identity, Option cloudProxy, Option deviceConnection) + public ConnectedDevice( + IIdentity identity, + Option cloudProxy, + Option deviceConnection, + TimeSpan cloudConnectionRetryInterval, + ISystemTime systemTime) { this.Identity = identity; this.CloudConnection = cloudProxy; this.DeviceConnection = deviceConnection; + this.cloudConnectionRetryInterval = cloudConnectionRetryInterval; + this.systemTime = systemTime; } public IIdentity Identity { get; } @@ -465,7 +517,19 @@ public async Task> CreateOrUpdateCloudConnection( // Lock in case multiple connections are created to the cloud for the same device at the same time using (await this.cloudConnectionLock.LockAsync()) { - Try newCloudConnection = await cloudConnectionUpdater(this); + Task> updateTask = cloudConnectionUpdater(this); + this.cloudConnectionCreateTask = Option.Some(updateTask); + this.cloudConnectionCreateWasRetry = false; + Try newCloudConnection; + try + { + newCloudConnection = await updateTask; + } + finally + { + this.cloudConnectionCreateCompletedTime = Option.Some(this.systemTime.UtcNow); + } + if (newCloudConnection.Success) { this.CloudConnection = Option.Some(newCloudConnection.Value); @@ -476,7 +540,8 @@ public async Task> CreateOrUpdateCloudConnection( } public async Task> GetOrCreateCloudConnection( - Func>> cloudConnectionCreator) + Func>> cloudConnectionCreator, + bool isRetry) { Preconditions.CheckNotNull(cloudConnectionCreator, nameof(cloudConnectionCreator)); @@ -496,13 +561,27 @@ public async Task> GetOrCreateCloudConnection( .GetOrElse( async () => { - return await this.cloudConnectionCreateTask.Filter(c => !c.IsCompleted) + bool retryIntervalElapsed = this.cloudConnectionCreateCompletedTime + .Map(t => this.systemTime.UtcNow - t >= this.cloudConnectionRetryInterval) + .GetOrElse(true); + return await this.cloudConnectionCreateTask.Filter( + c => !c.IsCompleted || this.cloudConnectionCreateWasRetry && !retryIntervalElapsed) .GetOrElse( async () => { Task> createTask = cloudConnectionCreator(this); this.cloudConnectionCreateTask = Option.Some(createTask); - Try cloudConnectionResult = await createTask; + this.cloudConnectionCreateWasRetry = isRetry; + Try cloudConnectionResult; + try + { + cloudConnectionResult = await createTask; + } + finally + { + this.cloudConnectionCreateCompletedTime = Option.Some(this.systemTime.UtcNow); + } + this.CloudConnection = cloudConnectionResult.Ok(); return cloudConnectionResult; }); diff --git a/edge-hub/core/test/Microsoft.Azure.Devices.Edge.Hub.Core.Test/ConnectionManagerTest.cs b/edge-hub/core/test/Microsoft.Azure.Devices.Edge.Hub.Core.Test/ConnectionManagerTest.cs index b17368a4e7b..0c3db3f98f9 100644 --- a/edge-hub/core/test/Microsoft.Azure.Devices.Edge.Hub.Core.Test/ConnectionManagerTest.cs +++ b/edge-hub/core/test/Microsoft.Azure.Devices.Edge.Hub.Core.Test/ConnectionManagerTest.cs @@ -572,6 +572,62 @@ public async Task GetOrCreateCloudProxyTest() cloudProxyProviderMock.Verify(c => c.Connect(It.IsAny(), It.IsAny>()), Times.Exactly(2)); } + [Fact] + [Unit] + public async Task GetCloudConnectionReusesRecentInactiveConnectionCreation() + { + const string DeviceId = "device1"; + const int OperationCount = 200; + TimeSpan retryInterval = TimeSpan.FromSeconds(5); + DateTime now = new DateTime(2026, 8, 6, 0, 0, 0, DateTimeKind.Utc); + var systemTime = new Mock(); + systemTime.SetupGet(t => t.UtcNow).Returns(() => now); + + var cloudProxy = Mock.Of(p => !p.IsActive); + Mock.Get(cloudProxy) + .Setup(p => p.SendMessageAsync(It.IsAny())) + .ThrowsAsync(new ObjectDisposedException("cloud proxy")); + var cloudConnection = Mock.Of( + c => !c.IsActive && c.CloudProxy == Option.Some(cloudProxy)); + var cloudConnectionProvider = new Mock(); + cloudConnectionProvider + .Setup(c => c.Connect(It.Is(i => i.Id == DeviceId), It.IsAny>())) + .ReturnsAsync(Try.Success(cloudConnection)); + + var connectionManager = new ConnectionManager( + cloudConnectionProvider.Object, + Mock.Of(), + GetIdentityProvider(), + Mock.Of(), + 101, + true, + retryInterval, + systemTime.Object); + + Option initialCloudProxy = await connectionManager.GetCloudConnection(DeviceId); + Assert.True(initialCloudProxy.HasValue); + + IMessage message = Mock.Of(); + Task[] operations = Enumerable.Range(0, OperationCount) + .Select( + _ => Assert.ThrowsAsync( + () => initialCloudProxy.OrDefault().SendMessageAsync(message))) + .ToArray(); + await Task.WhenAll(operations); + + cloudConnectionProvider.Verify( + c => c.Connect(It.IsAny(), It.IsAny>()), + Times.Exactly(2)); + + now += retryInterval; + await Assert.ThrowsAsync( + () => initialCloudProxy.OrDefault().SendMessageAsync(message)); + + cloudConnectionProvider.Verify( + c => c.Connect(It.IsAny(), It.IsAny>()), + Times.Exactly(3)); + } + [Fact] [Unit] public async Task CreateCloudProxyTest() diff --git a/edge-hub/core/test/Microsoft.Azure.Devices.Edge.Hub.Core.Test/RetryingCloudProxyTest.cs b/edge-hub/core/test/Microsoft.Azure.Devices.Edge.Hub.Core.Test/RetryingCloudProxyTest.cs index 76c23e44eb0..d3040285651 100644 --- a/edge-hub/core/test/Microsoft.Azure.Devices.Edge.Hub.Core.Test/RetryingCloudProxyTest.cs +++ b/edge-hub/core/test/Microsoft.Azure.Devices.Edge.Hub.Core.Test/RetryingCloudProxyTest.cs @@ -93,7 +93,11 @@ public async Task TestSendMessages() connectionProvider.BindEdgeHub(edgeHub.Object); var deviceConnectivityManager = Mock.Of(); - var connectionManager = new ConnectionManager(connectionProvider, credentialsCache.Object, identityProvider.Object, deviceConnectivityManager); + var connectionManager = CreateConnectionManagerWithoutCloudConnectionRetryDelay( + connectionProvider, + credentialsCache.Object, + identityProvider.Object, + deviceConnectivityManager); var messagesToSend = new List(); for (int i = 0; i < 10; i++) { @@ -200,7 +204,11 @@ public async Task TestGetTwin() connectionProvider.BindEdgeHub(edgeHub.Object); var deviceConnectivityManager = Mock.Of(); - var connectionManager = new ConnectionManager(connectionProvider, credentialsCache.Object, identityProvider.Object, deviceConnectivityManager); + var connectionManager = CreateConnectionManagerWithoutCloudConnectionRetryDelay( + connectionProvider, + credentialsCache.Object, + identityProvider.Object, + deviceConnectivityManager); // Act Option cloudProxyOption = await connectionManager.GetCloudConnection(Id); @@ -291,7 +299,11 @@ public async Task TestMultipleOperations() connectionProvider.BindEdgeHub(edgeHub.Object); var deviceConnectivityManager = Mock.Of(); - var connectionManager = new ConnectionManager(connectionProvider, credentialsCache.Object, identityProvider.Object, deviceConnectivityManager); + var connectionManager = CreateConnectionManagerWithoutCloudConnectionRetryDelay( + connectionProvider, + credentialsCache.Object, + identityProvider.Object, + deviceConnectivityManager); async Task GetCloudProxy(IConnectionManager cm) { @@ -349,6 +361,21 @@ async Task GetCloudProxy(IConnectionManager cm) Assert.Equal(expectedMessageIds, receivedMessageIds); } + static ConnectionManager CreateConnectionManagerWithoutCloudConnectionRetryDelay( + ICloudConnectionProvider connectionProvider, + ICredentialsCache credentialsCache, + IIdentityProvider identityProvider, + IDeviceConnectivityManager deviceConnectivityManager) => + new ConnectionManager( + connectionProvider, + credentialsCache, + identityProvider, + deviceConnectivityManager, + 101, + true, + TimeSpan.Zero, + SystemTime.Instance); + static async Task RunSendMessages(ICloudProxy cloudProxy, IEnumerable messages, int batchSize = 1) { if (batchSize == 1) From 9c310bbef9517461e9fa00765a6bdb664a4a839a Mon Sep 17 00:00:00 2001 From: yophilav Date: Fri, 7 Aug 2026 00:56:40 +0000 Subject: [PATCH 2/4] Harden cloud proxy retry cooldown Coalesce failed per-device rebuild attempts, fence concurrent removals, and preserve token handoff without exposing stale connections to normal lookup. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 0b428cdd-2248-4f74-87b1-cd02dc9ad50a --- .../ClientTokenCloudConnection.cs | 92 ++- .../ConnectionManager.cs | 589 ++++++++++++++---- .../cloud/IClientTokenCloudConnection.cs | 4 + .../ClientTokenCloudConnectionTest.cs | 98 +++ .../ConnectionManagerTest.cs | 383 +++++++++++- .../RetryingCloudProxyTest.cs | 3 +- 6 files changed, 1009 insertions(+), 160 deletions(-) diff --git a/edge-hub/core/src/Microsoft.Azure.Devices.Edge.Hub.CloudProxy/ClientTokenCloudConnection.cs b/edge-hub/core/src/Microsoft.Azure.Devices.Edge.Hub.CloudProxy/ClientTokenCloudConnection.cs index cca2f6496b7..0dee76de0ff 100644 --- a/edge-hub/core/src/Microsoft.Azure.Devices.Edge.Hub.CloudProxy/ClientTokenCloudConnection.cs +++ b/edge-hub/core/src/Microsoft.Azure.Devices.Edge.Hub.CloudProxy/ClientTokenCloudConnection.cs @@ -2,6 +2,7 @@ namespace Microsoft.Azure.Devices.Edge.Hub.CloudProxy { using System; + using System.Threading; using System.Threading.Tasks; using Microsoft.Azure.Devices.Client; using Microsoft.Azure.Devices.Edge.Hub.Core; @@ -23,7 +24,8 @@ class ClientTokenCloudConnection : CloudConnection, IClientTokenCloudConnection readonly AsyncLock identityUpdateLock = new AsyncLock(); bool callbacksEnabled = true; - Option> tokenGetter; + TaskCompletionSource tokenGetter; + int tokenUpdatesCanceled; Option cloudProxy; ClientTokenCloudConnection( @@ -123,20 +125,9 @@ public async Task UpdateTokenAsync(ITokenCredentials newTokenCreden // If the Identity has a token, and we have a tokenGetter, that means // the connection is waiting for a new token. So give it the token and // complete the tokenGetter - if (this.tokenGetter.HasValue) + if (Volatile.Read(ref this.tokenGetter) != null) { - if (TokenHelper.IsTokenExpired(this.Identity.IotHubHostname, newTokenCredentials.Token)) - { - throw new InvalidOperationException($"Token for client {this.Identity.Id} is expired"); - } - - this.tokenGetter.ForEach( - tg => - { - // First reset the token getter and then set the result. - this.tokenGetter = Option.None>(); - tg.SetResult(newTokenCredentials.Token); - }); + this.CompleteTokenGetter(newTokenCredentials); return cp; } else @@ -148,8 +139,12 @@ public async Task UpdateTokenAsync(ITokenCredentials newTokenCreden return newCloudProxy; } }) - // No existing cloud proxy, so just create a new one. - .GetOrElse(() => this.CreateNewCloudProxyAsync(tokenProvider)); + .GetOrElse( + async () => + { + this.CompleteTokenGetter(newTokenCredentials); + return await this.CreateNewCloudProxyAsync(tokenProvider); + }); // Set Identity only after successfully opening cloud proxy // That way, if a we have one existing connection for a deviceA, @@ -173,6 +168,26 @@ public async Task UpdateTokenAsync(ITokenCredentials newTokenCreden protected override Option GetCloudProxy() => this.cloudProxy; + public bool HasPendingTokenUpdate => Volatile.Read(ref this.tokenGetter) != null; + + public void CancelTokenUpdate() + { + Volatile.Write(ref this.tokenUpdatesCanceled, 1); + Interlocked.Exchange(ref this.tokenGetter, null)?.TrySetCanceled(); + } + + void CompleteTokenGetter(ITokenCredentials newTokenCredentials) + { + TaskCompletionSource currentTokenGetter = Volatile.Read(ref this.tokenGetter); + if (currentTokenGetter != null + && TokenHelper.IsTokenExpired(this.Identity.IotHubHostname, newTokenCredentials.Token)) + { + throw new InvalidOperationException($"Token for client {this.Identity.Id} is expired"); + } + + Interlocked.Exchange(ref this.tokenGetter, null)?.TrySetResult(newTokenCredentials.Token); + } + // Checks if the token expires too soon static bool IsTokenUsable(string hostname, string token) { @@ -222,17 +237,35 @@ async Task GetNewToken(string currentToken) } bool newTokenGetterCreated = false; + if (Volatile.Read(ref this.tokenUpdatesCanceled) != 0) + { + throw new OperationCanceledException( + $"Token updates for client {this.Identity.Id} have been canceled."); + } + // No need to lock here as the lock is being held by the refresher. - TaskCompletionSource tcs = this.tokenGetter - .GetOrElse( - () => - { - Events.SafeCreateNewToken(this.Identity.Id); - var taskCompletionSource = new TaskCompletionSource(); - this.tokenGetter = Option.Some(taskCompletionSource); - newTokenGetterCreated = true; - return taskCompletionSource; - }); + TaskCompletionSource tcs = Volatile.Read(ref this.tokenGetter); + if (tcs == null) + { + Events.SafeCreateNewToken(this.Identity.Id); + var taskCompletionSource = new TaskCompletionSource(); + tcs = Interlocked.CompareExchange( + ref this.tokenGetter, + taskCompletionSource, + null); + if (tcs == null) + { + tcs = taskCompletionSource; + newTokenGetterCreated = true; + } + } + + if (Volatile.Read(ref this.tokenUpdatesCanceled) != 0) + { + this.CancelTokenUpdate(); + throw new OperationCanceledException( + $"Token updates for client {this.Identity.Id} have been canceled."); + } // If a new tokenGetter was created, then invoke the connection status changed handler if (newTokenGetterCreated) @@ -243,6 +276,13 @@ async Task GetNewToken(string currentToken) await Task.Delay(TokenRetryWaitTime); } + if (Volatile.Read(ref this.tokenUpdatesCanceled) != 0) + { + this.CancelTokenUpdate(); + throw new OperationCanceledException( + $"Token updates for client {this.Identity.Id} have been canceled."); + } + this.ConnectionStatusChangedHandler(this.Identity.Id, CloudConnectionStatus.TokenNearExpiry); } diff --git a/edge-hub/core/src/Microsoft.Azure.Devices.Edge.Hub.Core/ConnectionManager.cs b/edge-hub/core/src/Microsoft.Azure.Devices.Edge.Hub.Core/ConnectionManager.cs index 1b6f9939fc9..3a19b15df73 100644 --- a/edge-hub/core/src/Microsoft.Azure.Devices.Edge.Hub.Core/ConnectionManager.cs +++ b/edge-hub/core/src/Microsoft.Azure.Devices.Edge.Hub.Core/ConnectionManager.cs @@ -5,7 +5,9 @@ namespace Microsoft.Azure.Devices.Edge.Hub.Core using System.Collections.Concurrent; using System.Collections.Generic; using System.Collections.ObjectModel; + using System.Diagnostics; using System.Linq; + using System.Threading; using System.Threading.Tasks; using App.Metrics; using App.Metrics.Gauge; @@ -32,7 +34,7 @@ public class ConnectionManager : IConnectionManager readonly IDeviceConnectivityManager connectivityManager; readonly bool closeCloudConnectionOnDeviceDisconnect; readonly TimeSpan cloudConnectionRetryInterval; - readonly ISystemTime systemTime; + readonly Func getTimestamp; public ConnectionManager( ICloudConnectionProvider cloudConnectionProvider, @@ -49,7 +51,7 @@ public ConnectionManager( maxClients, closeCloudConnectionOnDeviceDisconnect, DefaultCloudConnectionRetryInterval, - SystemTime.Instance) + Stopwatch.GetTimestamp) { } @@ -61,7 +63,7 @@ internal ConnectionManager( int maxClients, bool closeCloudConnectionOnDeviceDisconnect, TimeSpan cloudConnectionRetryInterval, - ISystemTime systemTime) + Func getTimestamp) { this.cloudConnectionProvider = Preconditions.CheckNotNull(cloudConnectionProvider, nameof(cloudConnectionProvider)); this.maxClients = Preconditions.CheckRange(maxClients, 1, nameof(maxClients)); @@ -71,7 +73,7 @@ internal ConnectionManager( this.cloudConnectionRetryInterval = cloudConnectionRetryInterval >= TimeSpan.Zero ? cloudConnectionRetryInterval : throw new ArgumentOutOfRangeException(nameof(cloudConnectionRetryInterval)); - this.systemTime = Preconditions.CheckNotNull(systemTime, nameof(systemTime)); + this.getTimestamp = Preconditions.CheckNotNull(getTimestamp, nameof(getTimestamp)); this.connectivityManager.DeviceDisconnected += (o, args) => this.HandleDeviceCloudConnectionDisconnected(); this.closeCloudConnectionOnDeviceDisconnect = closeCloudConnectionOnDeviceDisconnect; } @@ -119,28 +121,27 @@ public Option GetDeviceConnection(string id) public async Task> GetCloudConnection(string id) { - Try cloudProxyTry = await this.TryGetCloudConnectionInternal(id, false); + Try cloudProxyTry = await this.TryGetCloudConnectionInternal(id); return cloudProxyTry .Ok() - .Map(c => (ICloudProxy)new RetryingCloudProxy(id, () => this.TryGetCloudConnectionInternal(id, true), c)); + .Map(c => (ICloudProxy)new RetryingCloudProxy(id, () => this.TryGetCloudConnectionInternal(id), c)); } public async Task> TryGetCloudConnection(string id) { - Try cloudProxyTry = await this.TryGetCloudConnectionInternal(id, false); + Try cloudProxyTry = await this.TryGetCloudConnectionInternal(id); return cloudProxyTry.Success - ? Try.Success((ICloudProxy)new RetryingCloudProxy(id, () => this.TryGetCloudConnectionInternal(id, true), cloudProxyTry.Value)) + ? Try.Success((ICloudProxy)new RetryingCloudProxy(id, () => this.TryGetCloudConnectionInternal(id), cloudProxyTry.Value)) : cloudProxyTry; } - async Task> TryGetCloudConnectionInternal(string id, bool isRetry) + async Task> TryGetCloudConnectionInternal(string id) { IIdentity identity = this.identityProvider.Create(Preconditions.CheckNonWhiteSpace(id, nameof(id))); ConnectedDevice device = this.GetOrCreateConnectedDevice(identity); Try cloudConnectionTry = await device.GetOrCreateCloudConnection( - c => this.ConnectToCloud(c.Identity, this.CloudConnectionStatusChangedHandler), - isRetry); + c => this.ConnectToCloud(c.Identity, this.CloudConnectionStatusChangedHandler)); Events.GetCloudConnection(device.Identity, cloudConnectionTry); Try cloudProxyTry = GetCloudProxyFromCloudConnection(cloudConnectionTry, device.Identity); @@ -250,7 +251,7 @@ public async Task> CreateCloudConnectionAsync(IClientCredential Events.NewCloudConnection(credentials.Identity, newCloudConnection); Try cloudProxyTry = GetCloudProxyFromCloudConnection(newCloudConnection, credentials.Identity); return cloudProxyTry.Success - ? Try.Success((ICloudProxy)new RetryingCloudProxy(credentials.Identity.Id, () => this.TryGetCloudConnectionInternal(credentials.Identity.Id, true), cloudProxyTry.Value)) + ? Try.Success((ICloudProxy)new RetryingCloudProxy(credentials.Identity.Id, () => this.TryGetCloudConnectionInternal(credentials.Identity.Id), cloudProxyTry.Value)) : cloudProxyTry; } @@ -265,12 +266,11 @@ public async Task> GetOrCreateCloudConnectionAsync(IClientCrede ConnectedDevice device = this.GetOrCreateConnectedDevice(credentials.Identity); Try cloudConnectionTry = await device.GetOrCreateCloudConnection( - c => this.CreateOrUpdateCloudConnection(c, credentials), - false); + c => this.CreateOrUpdateCloudConnection(c, credentials)); Events.GetCloudConnection(credentials.Identity, cloudConnectionTry); Try cloudProxyTry = GetCloudProxyFromCloudConnection(cloudConnectionTry, credentials.Identity); return cloudProxyTry.Success - ? Try.Success((ICloudProxy)new RetryingCloudProxy(credentials.Identity.Id, () => this.TryGetCloudConnectionInternal(credentials.Identity.Id, true), cloudProxyTry.Value)) + ? Try.Success((ICloudProxy)new RetryingCloudProxy(credentials.Identity.Id, () => this.TryGetCloudConnectionInternal(credentials.Identity.Id), cloudProxyTry.Value)) : cloudProxyTry; } @@ -279,7 +279,10 @@ static Try GetCloudProxyFromCloudConnection(Try c .GetOrElse(() => Try.Failure(new EdgeHubConnectionException($"Unable to get cloud proxy for device {identity.Id}"))) : Try.Failure(cloudConnection.Exception); - async Task RemoveDeviceConnection(ConnectedDevice device, bool removeCloudConnection) + async Task RemoveDeviceConnection( + ConnectedDevice device, + bool removeCloudConnection, + bool throttleReconnect = false) { var id = device.Identity.Id; Events.RemovingDeviceConnection(id, removeCloudConnection); @@ -288,8 +291,9 @@ await device.DeviceConnection.Filter(dp => dp.IsActive) if (removeCloudConnection) { - await device.CloudConnection.Filter(cp => cp.IsActive) - .ForEachAsync(cp => cp.CloseAsync()); + await device.RemoveCloudConnection( + throttleReconnect, + preserveConnection: throttleReconnect); } Events.RemoveDeviceConnection(id); @@ -298,7 +302,7 @@ await device.CloudConnection.Filter(cp => cp.IsActive) } Task> CreateOrUpdateCloudConnection(ConnectedDevice device, IClientCredentials credentials) => - device.CloudConnection.Map( + device.CloudConnectionForUpdate.Map( async c => { try @@ -350,7 +354,7 @@ await clientCredentials.ForEachAsync( Try cloudConnectionTry = await device.CreateOrUpdateCloudConnection(c => this.CreateOrUpdateCloudConnection(c, tokenCredentials)); if (!cloudConnectionTry.Success) { - await this.RemoveDeviceConnection(device, true); + await this.RemoveDeviceConnection(device, true, true); this.CloudConnectionLost?.Invoke(this, device.Identity); } } @@ -362,14 +366,14 @@ await clientCredentials.ForEachAsync( } else { - await this.RemoveDeviceConnection(device, true); + await this.RemoveDeviceConnection(device, true, true); this.CloudConnectionLost?.Invoke(this, device.Identity); } break; case CloudConnectionStatus.DisconnectedTokenExpired: - await this.RemoveDeviceConnection(device, true); + await this.RemoveDeviceConnection(device, true, true); Events.InvokingCloudConnectionLostEvent(device.Identity); this.CloudConnectionLost?.Invoke(this, device.Identity); break; @@ -388,18 +392,21 @@ await clientCredentials.ForEachAsync( async void HandleDeviceCloudConnectionDisconnected() { + KeyValuePair[] snapshot; using (await this.connectToCloudLock.WriterLockAsync()) { - KeyValuePair[] snapshot = this.devices.ToArray(); + snapshot = this.devices.ToArray(); Events.CloudConnectionLostClosingAllClients(); foreach (var item in snapshot) { - await item.Value.CloudConnection.Filter(cp => cp.IsActive).ForEachAsync( - cp => - { - Events.CloudConnectionLostClosingClient(item.Value.Identity); - return cp.CloseAsync(); - }); + if (item.Value.CloudConnection.Filter(cp => cp.IsActive).HasValue) + { + Events.CloudConnectionLostClosingClient(item.Value.Identity); + } + + await item.Value.RemoveCloudConnection( + throttleReconnect: false, + preserveConnection: false); } } } @@ -419,12 +426,11 @@ ConnectedDevice CreateOrUpdateConnectedDevice(IIdentity identity) return this.devices.AddOrUpdate( deviceId, id => this.CreateNewConnectedDevice(identity), - (id, cd) => new ConnectedDevice( - identity, - cd.CloudConnection, - cd.DeviceConnection, - this.cloudConnectionRetryInterval, - this.systemTime)); + (id, cd) => + { + cd.UpdateIdentity(identity); + return cd; + }); } ConnectedDevice CreateNewConnectedDevice(IIdentity identity) @@ -436,7 +442,7 @@ ConnectedDevice CreateNewConnectedDevice(IIdentity identity) throw new EdgeHubConnectionException($"Edge hub already has maximum allowed clients ({this.maxClients - 1}) connected."); } - return new ConnectedDevice(identity, this.cloudConnectionRetryInterval, this.systemTime); + return new ConnectedDevice(identity, this.cloudConnectionRetryInterval, this.getTimestamp); } } @@ -461,20 +467,29 @@ class ConnectedDevice // Device Proxy methods are sync coming from the Protocol gateway, // so using traditional locking mechanism for those. readonly object deviceProxyLock = new object(); + readonly object cloudConnectionStateLock = new object(); readonly AsyncLock cloudConnectionLock = new AsyncLock(); + readonly AsyncLock cloudConnectionRemovalLock = new AsyncLock(); readonly TimeSpan cloudConnectionRetryInterval; - readonly ISystemTime systemTime; - Option>> cloudConnectionCreateTask = Option.None>>(); - Option cloudConnectionCreateCompletedTime = Option.None(); - bool cloudConnectionCreateWasRetry; - - public ConnectedDevice(IIdentity identity, TimeSpan cloudConnectionRetryInterval, ISystemTime systemTime) + readonly Func getTimestamp; + ICloudConnection cloudConnection; + ICloudConnection preservedCloudConnection; + DeviceConnection deviceConnection; + IIdentity identity; + Task> cloudConnectionCreateTask; + long cloudConnectionCreateGeneration; + long cloudConnectionCreateCompletedTimestamp; + long cloudConnectionGeneration; + bool hasCloudConnectionCreateCompletedTimestamp; + bool shouldThrottleCloudConnectionCreation; + + public ConnectedDevice(IIdentity identity, TimeSpan cloudConnectionRetryInterval, Func getTimestamp) : this( identity, Option.None(), Option.None(), cloudConnectionRetryInterval, - systemTime) + getTimestamp) { } @@ -483,30 +498,71 @@ public ConnectedDevice( Option cloudProxy, Option deviceConnection, TimeSpan cloudConnectionRetryInterval, - ISystemTime systemTime) + Func getTimestamp) { - this.Identity = identity; - this.CloudConnection = cloudProxy; - this.DeviceConnection = deviceConnection; + this.identity = identity; + this.cloudConnection = cloudProxy.OrDefault(); + this.deviceConnection = deviceConnection.OrDefault(); this.cloudConnectionRetryInterval = cloudConnectionRetryInterval; - this.systemTime = systemTime; + this.getTimestamp = getTimestamp; } - public IIdentity Identity { get; } + public IIdentity Identity => Volatile.Read(ref this.identity); + + public Option CloudConnection + { + get + { + ICloudConnection currentCloudConnection = Volatile.Read(ref this.cloudConnection); + return currentCloudConnection != null + ? Option.Some(currentCloudConnection) + : Option.None(); + } + } - public Option CloudConnection { get; private set; } + public Option CloudConnectionForUpdate + { + get + { + ICloudConnection currentCloudConnection = Volatile.Read(ref this.cloudConnection) + ?? Volatile.Read(ref this.preservedCloudConnection); + return currentCloudConnection != null + ? Option.Some(currentCloudConnection) + : Option.None(); + } + } // ReSharper disable once MemberHidesStaticFromOuterClass - public Option DeviceConnection { get; private set; } + public Option DeviceConnection + { + get + { + DeviceConnection currentDeviceConnection = Volatile.Read(ref this.deviceConnection); + return currentDeviceConnection != null + ? Option.Some(currentDeviceConnection) + : Option.None(); + } + } + + public void UpdateIdentity(IIdentity identity) + { + Volatile.Write(ref this.identity, Preconditions.CheckNotNull(identity, nameof(identity))); + } public Option AddDeviceConnection(IDeviceProxy deviceProxy) { Preconditions.CheckNotNull(deviceProxy, nameof(deviceProxy)); lock (this.deviceProxyLock) { - Option currentValue = this.DeviceConnection; - this.DeviceConnection = Option.Some(new DeviceConnection(deviceProxy, new ConcurrentDictionary())); - return currentValue; + DeviceConnection newDeviceConnection = new DeviceConnection( + deviceProxy, + new ConcurrentDictionary()); + DeviceConnection currentDeviceConnection = Interlocked.Exchange( + ref this.deviceConnection, + newDeviceConnection); + return currentDeviceConnection != null + ? Option.Some(currentDeviceConnection) + : Option.None(); } } @@ -514,25 +570,44 @@ public async Task> CreateOrUpdateCloudConnection( Func>> cloudConnectionUpdater) { Preconditions.CheckNotNull(cloudConnectionUpdater, nameof(cloudConnectionUpdater)); + await this.WaitForCloudConnectionRemoval(); // Lock in case multiple connections are created to the cloud for the same device at the same time using (await this.cloudConnectionLock.LockAsync()) { - Task> updateTask = cloudConnectionUpdater(this); - this.cloudConnectionCreateTask = Option.Some(updateTask); - this.cloudConnectionCreateWasRetry = false; - Try newCloudConnection; - try - { - newCloudConnection = await updateTask; - } - finally + long connectionGeneration = Volatile.Read(ref this.cloudConnectionGeneration); + Try newCloudConnection = await cloudConnectionUpdater(this); + bool invalidated; + lock (this.cloudConnectionStateLock) { - this.cloudConnectionCreateCompletedTime = Option.Some(this.systemTime.UtcNow); + invalidated = !this.IsCloudConnectionGenerationValid(connectionGeneration); + if (!invalidated) + { + if (newCloudConnection.Success) + { + Volatile.Write(ref this.cloudConnection, newCloudConnection.Value); + Interlocked.CompareExchange( + ref this.preservedCloudConnection, + null, + newCloudConnection.Value); + } + + if (newCloudConnection.Success && newCloudConnection.Value.IsActive) + { + this.ResetCloudConnectionRetryStateCore(connectionGeneration); + } + else + { + this.RecordCloudConnectionAttemptCore( + Task.FromResult(newCloudConnection), + true, + connectionGeneration); + } + } } - if (newCloudConnection.Success) + if (invalidated) { - this.CloudConnection = Option.Some(newCloudConnection.Value); + return await this.DiscardInvalidatedCloudConnection(newCloudConnection); } return newCloudConnection; @@ -540,55 +615,337 @@ public async Task> CreateOrUpdateCloudConnection( } public async Task> GetOrCreateCloudConnection( - Func>> cloudConnectionCreator, - bool isRetry) + Func>> cloudConnectionCreator) { Preconditions.CheckNotNull(cloudConnectionCreator, nameof(cloudConnectionCreator)); - return await this.CloudConnection.Filter(cp => cp.IsActive) - .Map(c => Task.FromResult(Try.Success(c))) - .GetOrElse( - async () => + long connectionGeneration = Volatile.Read(ref this.cloudConnectionGeneration); + Task> createTask = Volatile.Read(ref this.cloudConnectionCreateTask); + if (createTask != null + && !createTask.IsCompleted + && this.IsCloudConnectionGenerationValid(connectionGeneration) + && connectionGeneration == Volatile.Read(ref this.cloudConnectionCreateGeneration)) + { + return await createTask; + } + + using (await this.cloudConnectionLock.LockAsync()) + { + lock (this.cloudConnectionStateLock) + { + connectionGeneration = Volatile.Read(ref this.cloudConnectionGeneration); + Option activeCloudConnection = this.CloudConnection.Filter(cp => cp.IsActive); + if (this.IsCloudConnectionGenerationValid(connectionGeneration) + && activeCloudConnection.HasValue) { - return await this.cloudConnectionCreateTask.Filter(c => !c.IsCompleted) - .GetOrElse( - async () => - { - using (await this.cloudConnectionLock.LockAsync()) - { - return await this.CloudConnection.Filter(cp => cp.IsActive) - .Map(c => Task.FromResult(Try.Success(c))) - .GetOrElse( - async () => - { - bool retryIntervalElapsed = this.cloudConnectionCreateCompletedTime - .Map(t => this.systemTime.UtcNow - t >= this.cloudConnectionRetryInterval) - .GetOrElse(true); - return await this.cloudConnectionCreateTask.Filter( - c => !c.IsCompleted || this.cloudConnectionCreateWasRetry && !retryIntervalElapsed) - .GetOrElse( - async () => - { - Task> createTask = cloudConnectionCreator(this); - this.cloudConnectionCreateTask = Option.Some(createTask); - this.cloudConnectionCreateWasRetry = isRetry; - Try cloudConnectionResult; - try - { - cloudConnectionResult = await createTask; - } - finally - { - this.cloudConnectionCreateCompletedTime = Option.Some(this.systemTime.UtcNow); - } - - this.CloudConnection = cloudConnectionResult.Ok(); - return cloudConnectionResult; - }); - }); - } - }); - }); + return Try.Success(activeCloudConnection.OrDefault()); + } + } + + bool reuseCreateTask = false; + lock (this.cloudConnectionStateLock) + { + connectionGeneration = Volatile.Read(ref this.cloudConnectionGeneration); + if (!this.IsCloudConnectionGenerationValid(connectionGeneration)) + { + return Try.Failure( + new EdgeHubConnectionException($"Cloud connection for device {this.Identity.Id} is being removed.")); + } + + createTask = this.cloudConnectionCreateTask; + if (createTask != null && connectionGeneration == this.cloudConnectionCreateGeneration) + { + reuseCreateTask = !createTask.IsCompleted + || this.shouldThrottleCloudConnectionCreation + && !this.CloudConnectionRetryIntervalElapsed(); + if (reuseCreateTask && createTask.IsCompleted) + { + Events.ReusingRecentCloudConnectionAttempt(this.Identity, this.cloudConnectionRetryInterval); + } + } + + if (!reuseCreateTask) + { + bool replacesExistingConnection = this.CloudConnectionForUpdate.HasValue; + Volatile.Write(ref this.cloudConnectionCreateGeneration, connectionGeneration); + createTask = this.CreateCloudConnection( + cloudConnectionCreator, + replacesExistingConnection, + connectionGeneration); + Volatile.Write(ref this.cloudConnectionCreateTask, createTask); + } + } + + return await createTask; + } + } + + async Task> CreateCloudConnection( + Func>> cloudConnectionCreator, + bool replacesExistingConnection, + long connectionGeneration) + { + Task> createTask = null; + try + { + createTask = cloudConnectionCreator(this); + Try cloudConnectionResult = await createTask; + bool invalidated; + ICloudConnection displacedPreservedConnection = null; + lock (this.cloudConnectionStateLock) + { + invalidated = !this.IsCloudConnectionGenerationValid(connectionGeneration); + if (!invalidated) + { + Volatile.Write( + ref this.cloudConnection, + cloudConnectionResult.Success ? cloudConnectionResult.Value : null); + if (cloudConnectionResult.Success) + { + displacedPreservedConnection = Interlocked.Exchange( + ref this.preservedCloudConnection, + null); + } + + this.shouldThrottleCloudConnectionCreation = replacesExistingConnection + || !cloudConnectionResult.Success + || !cloudConnectionResult.Value.IsActive; + } + } + + if (invalidated) + { + return await this.DiscardInvalidatedCloudConnection(cloudConnectionResult); + } + + if (displacedPreservedConnection != null + && !ReferenceEquals(displacedPreservedConnection, cloudConnectionResult.Value)) + { + if (displacedPreservedConnection is IClientTokenCloudConnection clientTokenCloudConnection) + { + clientTokenCloudConnection.CancelTokenUpdate(); + } + + await displacedPreservedConnection.CloseAsync(); + } + + return cloudConnectionResult; + } + finally + { + lock (this.cloudConnectionStateLock) + { + if (this.IsCloudConnectionGenerationValid(connectionGeneration)) + { + this.cloudConnectionCreateCompletedTimestamp = this.getTimestamp(); + this.hasCloudConnectionCreateCompletedTimestamp = true; + if (createTask == null || createTask.IsFaulted || createTask.IsCanceled) + { + this.shouldThrottleCloudConnectionCreation = true; + } + } + } + } + } + + async Task WaitForCloudConnectionRemoval() + { + using (await this.cloudConnectionRemovalLock.LockAsync()) + { + } + } + + bool CloudConnectionRetryIntervalElapsed() + { + if (!this.hasCloudConnectionCreateCompletedTimestamp) + { + return true; + } + + long elapsedTimestamp = this.getTimestamp() - this.cloudConnectionCreateCompletedTimestamp; + return elapsedTimestamp < 0 + || elapsedTimestamp >= this.cloudConnectionRetryInterval.TotalSeconds * Stopwatch.Frequency; + } + + void RecordCloudConnectionAttemptCore( + Task> createTask, + bool shouldThrottle, + long connectionGeneration) + { + if (!this.IsCloudConnectionGenerationValid(connectionGeneration)) + { + return; + } + + Volatile.Write(ref this.cloudConnectionCreateGeneration, connectionGeneration); + Volatile.Write(ref this.cloudConnectionCreateTask, createTask); + this.cloudConnectionCreateCompletedTimestamp = this.getTimestamp(); + this.hasCloudConnectionCreateCompletedTimestamp = true; + this.shouldThrottleCloudConnectionCreation = shouldThrottle; + } + + void ResetCloudConnectionRetryStateCore(long connectionGeneration) + { + if (!this.IsCloudConnectionGenerationValid(connectionGeneration)) + { + return; + } + + Volatile.Write(ref this.cloudConnectionCreateGeneration, connectionGeneration); + Volatile.Write(ref this.cloudConnectionCreateTask, null); + this.cloudConnectionCreateCompletedTimestamp = 0; + this.hasCloudConnectionCreateCompletedTimestamp = false; + this.shouldThrottleCloudConnectionCreation = false; + } + + public async Task RemoveCloudConnection( + bool throttleReconnect, + bool preserveConnection) + { + using (await this.cloudConnectionRemovalLock.LockAsync()) + { + var removedCloudConnections = new List(2); + var supersededCloudConnections = new List(1); + lock (this.cloudConnectionStateLock) + { + Interlocked.Increment(ref this.cloudConnectionGeneration); + ICloudConnection activeCloudConnection = + Interlocked.Exchange(ref this.cloudConnection, null); + if (activeCloudConnection != null) + { + removedCloudConnections.Add(activeCloudConnection); + } + + if (preserveConnection) + { + if (activeCloudConnection != null) + { + ICloudConnection previousPreservedConnection = + Interlocked.Exchange( + ref this.preservedCloudConnection, + activeCloudConnection); + if (previousPreservedConnection != null + && !ReferenceEquals(previousPreservedConnection, activeCloudConnection)) + { + supersededCloudConnections.Add(previousPreservedConnection); + } + } + } + else + { + ICloudConnection preservedConnection = + Interlocked.Exchange(ref this.preservedCloudConnection, null); + if (preservedConnection != null + && !ReferenceEquals(preservedConnection, activeCloudConnection)) + { + removedCloudConnections.Add(preservedConnection); + } + } + } + + try + { + foreach (ICloudConnection supersededCloudConnection in supersededCloudConnections) + { + if (supersededCloudConnection is IClientTokenCloudConnection clientTokenCloudConnection) + { + clientTokenCloudConnection.CancelTokenUpdate(); + } + + if (supersededCloudConnection.IsActive) + { + await supersededCloudConnection.CloseAsync(); + } + } + + foreach (ICloudConnection removedCloudConnection in removedCloudConnections) + { + bool preserveForTokenUpdate = false; + if (removedCloudConnection is IClientTokenCloudConnection clientTokenCloudConnection) + { + if (!preserveConnection) + { + clientTokenCloudConnection.CancelTokenUpdate(); + } + else if (clientTokenCloudConnection.HasPendingTokenUpdate) + { + preserveForTokenUpdate = true; + } + else + { + clientTokenCloudConnection.CancelTokenUpdate(); + } + } + + if (preserveForTokenUpdate) + { + continue; + } + + if (preserveConnection) + { + Interlocked.CompareExchange( + ref this.preservedCloudConnection, + null, + removedCloudConnection); + } + + if (removedCloudConnection.IsActive) + { + await removedCloudConnection.CloseAsync(); + } + } + } + finally + { + lock (this.cloudConnectionStateLock) + { + long stableGeneration = Volatile.Read(ref this.cloudConnectionGeneration) + 1; + Volatile.Write(ref this.cloudConnectionCreateGeneration, stableGeneration); + this.cloudConnectionCreateCompletedTimestamp = this.getTimestamp(); + this.hasCloudConnectionCreateCompletedTimestamp = throttleReconnect; + this.shouldThrottleCloudConnectionCreation = throttleReconnect; + Volatile.Write( + ref this.cloudConnectionCreateTask, + throttleReconnect + ? Task.FromResult( + Try.Failure( + new EdgeHubConnectionException( + $"Cloud connection for device {this.Identity.Id} was removed after a cloud failure."))) + : null); + Interlocked.Increment(ref this.cloudConnectionGeneration); + } + } + } + } + + bool IsCloudConnectionGenerationValid(long connectionGeneration) => + (connectionGeneration & 1) == 0 + && connectionGeneration == Volatile.Read(ref this.cloudConnectionGeneration); + + async Task> DiscardInvalidatedCloudConnection( + Try cloudConnection) + { + if (cloudConnection.Success) + { + Interlocked.CompareExchange( + ref this.cloudConnection, + null, + cloudConnection.Value); + Interlocked.CompareExchange( + ref this.preservedCloudConnection, + null, + cloudConnection.Value); + if (cloudConnection.Value is IClientTokenCloudConnection clientTokenCloudConnection) + { + clientTokenCloudConnection.CancelTokenUpdate(); + } + + await cloudConnection.Value.CloseAsync(); + } + + return Try.Failure( + new EdgeHubConnectionException($"Cloud connection attempt for device {this.Identity.Id} was invalidated.")); } } @@ -629,7 +986,8 @@ enum EventIds HandlingConnectionStatusChangedHandler, CloudConnectionLostClosingClient, CloudConnectionLostClosingAllClients, - GettingCloudConnectionForDeviceSubscriptions + GettingCloudConnectionForDeviceSubscriptions, + ReusingRecentCloudConnectionAttempt } public static void NewCloudConnection(IIdentity identity, Try cloudConnection) @@ -705,6 +1063,13 @@ public static void GettingCloudConnectionForDeviceSubscriptions() { Log.LogDebug((int)EventIds.GettingCloudConnectionForDeviceSubscriptions, $"Device has subscriptions. Trying to get cloud connection."); } + + public static void ReusingRecentCloudConnectionAttempt(IIdentity identity, TimeSpan retryInterval) + { + Log.LogDebug( + (int)EventIds.ReusingRecentCloudConnectionAttempt, + Invariant($"Reusing recent cloud connection attempt for device {identity.Id} during {retryInterval} retry interval.")); + } } static class MetricsV0 diff --git a/edge-hub/core/src/Microsoft.Azure.Devices.Edge.Hub.Core/cloud/IClientTokenCloudConnection.cs b/edge-hub/core/src/Microsoft.Azure.Devices.Edge.Hub.Core/cloud/IClientTokenCloudConnection.cs index 6f4250e37bb..572ecfce03a 100644 --- a/edge-hub/core/src/Microsoft.Azure.Devices.Edge.Hub.Core/cloud/IClientTokenCloudConnection.cs +++ b/edge-hub/core/src/Microsoft.Azure.Devices.Edge.Hub.Core/cloud/IClientTokenCloudConnection.cs @@ -6,6 +6,10 @@ namespace Microsoft.Azure.Devices.Edge.Hub.Core.Cloud public interface IClientTokenCloudConnection : ICloudConnection { + bool HasPendingTokenUpdate { get; } + + void CancelTokenUpdate(); + Task UpdateTokenAsync(ITokenCredentials tokenCredentials); } } diff --git a/edge-hub/core/test/Microsoft.Azure.Devices.Edge.Hub.CloudProxy.Test/ClientTokenCloudConnectionTest.cs b/edge-hub/core/test/Microsoft.Azure.Devices.Edge.Hub.CloudProxy.Test/ClientTokenCloudConnectionTest.cs index 819071c40e2..0d05a463681 100644 --- a/edge-hub/core/test/Microsoft.Azure.Devices.Edge.Hub.CloudProxy.Test/ClientTokenCloudConnectionTest.cs +++ b/edge-hub/core/test/Microsoft.Azure.Devices.Edge.Hub.CloudProxy.Test/ClientTokenCloudConnectionTest.cs @@ -206,6 +206,104 @@ ITokenCredentials GetClientCredentialsWithNonExpiringToken() Assert.Equal(getTokenTask.Result, clientCredentialsWithExpiringToken2.Token); } + [Fact] + [Unit] + public async Task CanceledConnectionRejectsFutureTokenRequests() + { + string iothubHostName = "test.azure-devices.net"; + string token = TokenHelper.CreateSasToken(iothubHostName, DateTime.UtcNow.AddMinutes(3)); + var credentials = new TokenCredentials( + new DeviceIdentity(iothubHostName, "device1"), + token, + string.Empty, + Option.None(), + Option.None(), + false); + ITokenProvider tokenProvider = null; + IClientProvider clientProvider = + GetMockDeviceClientProviderWithToken((_, provider, _, _) => tokenProvider = provider); + var transportSettings = + new ITransportSettings[] { new AmqpTransportSettings(TransportType.Amqp_Tcp_Only) }; + var messageConverterProvider = new MessageConverterProvider( + new Dictionary + { + [typeof(TwinCollection)] = Mock.Of() + }); + ClientTokenCloudConnection cloudConnection = await ClientTokenCloudConnection.Create( + credentials, + (_, __) => { }, + transportSettings, + messageConverterProvider, + clientProvider, + Mock.Of(), + TimeSpan.FromMinutes(60), + true, + TimeSpan.FromSeconds(20), + TimeSpan.FromSeconds(50), + DummyProductInfo, + Option.None()); + + cloudConnection.CancelTokenUpdate(); + + Assert.NotNull(tokenProvider); + await Assert.ThrowsAsync( + () => tokenProvider.GetTokenAsync(Option.None())); + Assert.False(cloudConnection.HasPendingTokenUpdate); + } + + [Fact] + [Unit] + public async Task CancellationDuringRetrySuppressesTokenStatusCallback() + { + string iothubHostName = "test.azure-devices.net"; + string token = TokenHelper.CreateSasToken(iothubHostName, DateTime.UtcNow.AddMinutes(3)); + var credentials = new TokenCredentials( + new DeviceIdentity(iothubHostName, "device1"), + token, + string.Empty, + Option.None(), + Option.None(), + false); + ITokenProvider tokenProvider = null; + IClientProvider clientProvider = + GetMockDeviceClientProviderWithToken((_, provider, _, _) => tokenProvider = provider); + var transportSettings = + new ITransportSettings[] { new AmqpTransportSettings(TransportType.Amqp_Tcp_Only) }; + var receivedStatuses = new List(); + var messageConverterProvider = new MessageConverterProvider( + new Dictionary + { + [typeof(TwinCollection)] = Mock.Of() + }); + ClientTokenCloudConnection cloudConnection = await ClientTokenCloudConnection.Create( + credentials, + (_, status) => receivedStatuses.Add(status), + transportSettings, + messageConverterProvider, + clientProvider, + Mock.Of(), + TimeSpan.FromMinutes(60), + true, + TimeSpan.FromSeconds(20), + TimeSpan.FromSeconds(50), + DummyProductInfo, + Option.None()); + Assert.NotNull(tokenProvider); + + Task getTokenTask = tokenProvider.GetTokenAsync(Option.None()); + Assert.Single(receivedStatuses); + await cloudConnection.UpdateTokenAsync(credentials); + await Task.Delay(TimeSpan.FromSeconds(1)); + Assert.True(cloudConnection.HasPendingTokenUpdate); + + cloudConnection.CancelTokenUpdate(); + + await Assert.ThrowsAnyAsync(() => getTokenTask); + await Task.Delay(TimeSpan.FromSeconds(20)); + Assert.Single(receivedStatuses); + Assert.False(cloudConnection.HasPendingTokenUpdate); + } + [Fact] [Unit] public async Task RefreshTokenWithRetryTest() diff --git a/edge-hub/core/test/Microsoft.Azure.Devices.Edge.Hub.Core.Test/ConnectionManagerTest.cs b/edge-hub/core/test/Microsoft.Azure.Devices.Edge.Hub.Core.Test/ConnectionManagerTest.cs index 0c3db3f98f9..2dbf1110e3d 100644 --- a/edge-hub/core/test/Microsoft.Azure.Devices.Edge.Hub.Core.Test/ConnectionManagerTest.cs +++ b/edge-hub/core/test/Microsoft.Azure.Devices.Edge.Hub.Core.Test/ConnectionManagerTest.cs @@ -3,6 +3,7 @@ namespace Microsoft.Azure.Devices.Edge.Hub.Core.Test { using System; using System.Collections.Generic; + using System.Diagnostics; using System.Linq; using System.Net; using System.Threading.Tasks; @@ -574,25 +575,16 @@ public async Task GetOrCreateCloudProxyTest() [Fact] [Unit] - public async Task GetCloudConnectionReusesRecentInactiveConnectionCreation() + public async Task GetCloudConnectionThrottlesRepeatedFailedCreation() { const string DeviceId = "device1"; const int OperationCount = 200; TimeSpan retryInterval = TimeSpan.FromSeconds(5); - DateTime now = new DateTime(2026, 8, 6, 0, 0, 0, DateTimeKind.Utc); - var systemTime = new Mock(); - systemTime.SetupGet(t => t.UtcNow).Returns(() => now); - - var cloudProxy = Mock.Of(p => !p.IsActive); - Mock.Get(cloudProxy) - .Setup(p => p.SendMessageAsync(It.IsAny())) - .ThrowsAsync(new ObjectDisposedException("cloud proxy")); - var cloudConnection = Mock.Of( - c => !c.IsActive && c.CloudProxy == Option.Some(cloudProxy)); + long timestamp = 0; var cloudConnectionProvider = new Mock(); cloudConnectionProvider .Setup(c => c.Connect(It.Is(i => i.Id == DeviceId), It.IsAny>())) - .ReturnsAsync(Try.Success(cloudConnection)); + .ReturnsAsync(Try.Failure(new TimeoutException())); var connectionManager = new ConnectionManager( cloudConnectionProvider.Object, @@ -602,16 +594,92 @@ public async Task GetCloudConnectionReusesRecentInactiveConnectionCreation() 101, true, retryInterval, - systemTime.Object); + () => timestamp); - Option initialCloudProxy = await connectionManager.GetCloudConnection(DeviceId); - Assert.True(initialCloudProxy.HasValue); + Task>[] operations = Enumerable.Range(0, OperationCount) + .Select(_ => connectionManager.GetCloudConnection(DeviceId)) + .ToArray(); + Option[] results = await Task.WhenAll(operations); + + Assert.All(results, result => Assert.False(result.HasValue)); + cloudConnectionProvider.Verify( + c => c.Connect(It.IsAny(), It.IsAny>()), + Times.Once); + + timestamp += (long)(retryInterval.TotalSeconds * Stopwatch.Frequency); + Option resultAfterRetryInterval = await connectionManager.GetCloudConnection(DeviceId); + + Assert.False(resultAfterRetryInterval.HasValue); + cloudConnectionProvider.Verify( + c => c.Connect(It.IsAny(), It.IsAny>()), + Times.Exactly(2)); + } + + [Fact] + [Unit] + public async Task GetCloudConnectionThrottlesImmediatelyInactiveReplacements() + { + const string DeviceId = "device1"; + const int OperationCount = 200; + TimeSpan retryInterval = TimeSpan.FromSeconds(5); + long timestamp = 0; + bool initialConnectionActive = true; + + var initialCloudProxy = new Mock(); + initialCloudProxy.SetupGet(p => p.IsActive).Returns(() => initialConnectionActive); + var initialCloudConnection = new Mock(); + initialCloudConnection.SetupGet(c => c.IsActive).Returns(() => initialConnectionActive); + initialCloudConnection.SetupGet(c => c.CloudProxy).Returns( + () => initialConnectionActive + ? Option.Some(initialCloudProxy.Object) + : Option.None()); + + Try CreateFailingReplacement() + { + bool isActive = true; + var cloudProxy = new Mock(); + cloudProxy.SetupGet(p => p.IsActive).Returns(() => isActive); + cloudProxy + .Setup(p => p.SendMessageAsync(It.IsAny())) + .Callback(() => isActive = false) + .ThrowsAsync(new ObjectDisposedException("cloud proxy")); + var cloudConnection = new Mock(); + cloudConnection.SetupGet(c => c.IsActive).Returns(() => isActive); + cloudConnection.SetupGet(c => c.CloudProxy).Returns( + () => isActive + ? Option.Some(cloudProxy.Object) + : Option.None()); + return Try.Success(cloudConnection.Object); + } + + int connectionCount = 0; + var cloudConnectionProvider = new Mock(); + cloudConnectionProvider + .Setup(c => c.Connect(It.Is(i => i.Id == DeviceId), It.IsAny>())) + .ReturnsAsync( + () => ++connectionCount == 1 + ? Try.Success(initialCloudConnection.Object) + : CreateFailingReplacement()); + + var connectionManager = new ConnectionManager( + cloudConnectionProvider.Object, + Mock.Of(), + GetIdentityProvider(), + Mock.Of(), + 101, + true, + retryInterval, + () => timestamp); + + Option cloudProxy = await connectionManager.GetCloudConnection(DeviceId); + Assert.True(cloudProxy.HasValue); + initialConnectionActive = false; IMessage message = Mock.Of(); - Task[] operations = Enumerable.Range(0, OperationCount) + Task[] operations = Enumerable.Range(0, OperationCount) .Select( - _ => Assert.ThrowsAsync( - () => initialCloudProxy.OrDefault().SendMessageAsync(message))) + _ => Assert.ThrowsAsync( + () => cloudProxy.OrDefault().SendMessageAsync(message))) .ToArray(); await Task.WhenAll(operations); @@ -619,15 +687,288 @@ public async Task GetCloudConnectionReusesRecentInactiveConnectionCreation() c => c.Connect(It.IsAny(), It.IsAny>()), Times.Exactly(2)); - now += retryInterval; - await Assert.ThrowsAsync( - () => initialCloudProxy.OrDefault().SendMessageAsync(message)); + timestamp += (long)(retryInterval.TotalSeconds * Stopwatch.Frequency); + await Assert.ThrowsAsync( + () => cloudProxy.OrDefault().SendMessageAsync(message)); cloudConnectionProvider.Verify( c => c.Connect(It.IsAny(), It.IsAny>()), Times.Exactly(3)); } + [Fact] + [Unit] + public async Task RemoveDeviceConnectionInvalidatesInFlightCloudConnection() + { + const string DeviceId = "device1"; + bool isActive = true; + var cloudProxy = Mock.Of(p => p.IsActive); + var cloudConnection = new Mock(); + cloudConnection.SetupGet(c => c.IsActive).Returns(() => isActive); + cloudConnection.SetupGet(c => c.CloudProxy).Returns(() => Option.Some(cloudProxy)); + cloudConnection + .Setup(c => c.CloseAsync()) + .Callback(() => isActive = false) + .ReturnsAsync(true); + + var connectionSource = new TaskCompletionSource>( + TaskCreationOptions.RunContinuationsAsynchronously); + var cloudConnectionProvider = new Mock(); + cloudConnectionProvider + .Setup(c => c.Connect(It.Is(i => i.Id == DeviceId), It.IsAny>())) + .Returns(connectionSource.Task); + var connectionManager = new ConnectionManager( + cloudConnectionProvider.Object, + Mock.Of(), + GetIdentityProvider(), + Mock.Of()); + + Task> getCloudConnection = connectionManager.GetCloudConnection(DeviceId); + cloudConnectionProvider.Verify( + c => c.Connect(It.IsAny(), It.IsAny>()), + Times.Once); + + await connectionManager.RemoveDeviceConnection(DeviceId); + connectionSource.SetResult(Try.Success(cloudConnection.Object)); + Option result = await getCloudConnection; + + Assert.False(result.HasValue); + Assert.False(isActive); + cloudConnection.Verify(c => c.CloseAsync(), Times.Once); + } + + [Fact] + [Unit] + public async Task CreateCloudConnectionWaitsForRemovalToComplete() + { + const string DeviceId = "device1"; + var identity = new DeviceIdentity(IotHubHostName, DeviceId); + var credentials = new TokenCredentials( + identity, + DummyToken, + DummyProductInfo, + Option.None(), + Option.None(), + false); + bool firstConnectionActive = true; + var firstCloudProxy = Mock.Of(p => p.IsActive); + var closeStarted = new TaskCompletionSource( + TaskCreationOptions.RunContinuationsAsynchronously); + var closeConnection = new TaskCompletionSource( + TaskCreationOptions.RunContinuationsAsynchronously); + var firstCloudConnection = new Mock(); + firstCloudConnection.SetupGet(c => c.IsActive).Returns(() => firstConnectionActive); + firstCloudConnection.SetupGet(c => c.CloudProxy).Returns(() => Option.Some(firstCloudProxy)); + firstCloudConnection + .Setup(c => c.CloseAsync()) + .Callback(() => closeStarted.SetResult(true)) + .Returns(closeConnection.Task); + + var secondCloudProxy = Mock.Of(p => p.IsActive); + var secondCloudConnection = Mock.Of( + c => c.IsActive && c.CloudProxy == Option.Some(secondCloudProxy)); + var cloudConnectionProvider = new Mock(); + cloudConnectionProvider + .SetupSequence(c => c.Connect(It.IsAny(), It.IsAny>())) + .ReturnsAsync(Try.Success(firstCloudConnection.Object)) + .ReturnsAsync(Try.Success(secondCloudConnection)); + var connectionManager = new ConnectionManager( + cloudConnectionProvider.Object, + Mock.Of(), + GetIdentityProvider(), + Mock.Of()); + + Try firstResult = await connectionManager.CreateCloudConnectionAsync(credentials); + Assert.True(firstResult.Success); + + Task removeConnection = connectionManager.RemoveDeviceConnection(DeviceId); + await closeStarted.Task; + Option resultDuringRemoval = await connectionManager.GetCloudConnection(DeviceId); + Task> createConnection = connectionManager.CreateCloudConnectionAsync(credentials); + + Assert.False(resultDuringRemoval.HasValue); + cloudConnectionProvider.Verify( + c => c.Connect(It.IsAny(), It.IsAny>()), + Times.Once); + + firstConnectionActive = false; + closeConnection.SetResult(true); + await removeConnection; + Try secondResult = await createConnection; + + Assert.True(secondResult.Success); + cloudConnectionProvider.Verify( + c => c.Connect(It.IsAny(), It.IsAny>()), + Times.Exactly(2)); + } + + [Fact] + [Unit] + public async Task TokenExpiredRemovalPreservesConnectionOnlyForTokenUpdate() + { + const string DeviceId = "device1"; + const int OperationCount = 200; + TimeSpan retryInterval = TimeSpan.FromSeconds(5); + long timestamp = 0; + Action statusChangedHandler = null; + var connectionRemoved = new TaskCompletionSource( + TaskCreationOptions.RunContinuationsAsynchronously); + var cloudProxy = Mock.Of(p => p.IsActive); + var updatedCloudProxy = Mock.Of(p => p.IsActive); + bool tokenUpdated = false; + var cloudConnection = new Mock(); + cloudConnection.SetupGet(c => c.IsActive).Returns(true); + cloudConnection + .SetupGet(c => c.CloudProxy) + .Returns(() => Option.Some(tokenUpdated ? updatedCloudProxy : cloudProxy)); + cloudConnection + .SetupGet(c => c.HasPendingTokenUpdate) + .Callback(() => connectionRemoved.TrySetResult(true)) + .Returns(true); + cloudConnection + .Setup(c => c.CloseAsync()) + .ReturnsAsync(true); + cloudConnection + .Setup(c => c.UpdateTokenAsync(It.IsAny())) + .Callback(() => tokenUpdated = true) + .ReturnsAsync(updatedCloudProxy); + + int connectionCount = 0; + var cloudConnectionProvider = new Mock(); + cloudConnectionProvider + .Setup(c => c.Connect(It.Is(i => i.Id == DeviceId), It.IsAny>())) + .Callback>((_, handler) => statusChangedHandler = handler) + .ReturnsAsync( + () => ++connectionCount == 1 + ? Try.Success(cloudConnection.Object as ICloudConnection) + : Try.Failure(new TimeoutException())); + + var connectionManager = new ConnectionManager( + cloudConnectionProvider.Object, + Mock.Of(), + GetIdentityProvider(), + Mock.Of(), + 101, + true, + retryInterval, + () => timestamp); + + Option initialCloudProxy = await connectionManager.GetCloudConnection(DeviceId); + Assert.True(initialCloudProxy.HasValue); + Assert.NotNull(statusChangedHandler); + + statusChangedHandler(DeviceId, CloudConnectionStatus.DisconnectedTokenExpired); + await connectionRemoved.Task; + + Task>[] operations = Enumerable.Range(0, OperationCount) + .Select(_ => connectionManager.GetCloudConnection(DeviceId)) + .ToArray(); + Option[] results = await Task.WhenAll(operations); + + Assert.All(results, result => Assert.False(result.HasValue)); + cloudConnectionProvider.Verify( + c => c.Connect(It.IsAny(), It.IsAny>()), + Times.Once); + cloudConnection.Verify(c => c.CloseAsync(), Times.Never); + + var updatedCredentials = new TokenCredentials( + new DeviceIdentity(IotHubHostName, DeviceId), + DummyToken, + DummyProductInfo, + Option.None(), + Option.None(), + true); + Try updatedConnection = + await connectionManager.CreateCloudConnectionAsync(updatedCredentials); + + Assert.True(updatedConnection.Success); + Assert.Same(updatedCloudProxy, ((RetryingCloudProxy)updatedConnection.Value).InnerCloudProxy); + Option activeConnection = await connectionManager.GetCloudConnection(DeviceId); + Assert.True(activeConnection.HasValue); + Assert.Same(updatedCloudProxy, ((RetryingCloudProxy)activeConnection.OrDefault()).InnerCloudProxy); + cloudConnection.Verify(c => c.UpdateTokenAsync(updatedCredentials), Times.Once); + cloudConnectionProvider.Verify( + c => c.Connect(It.IsAny(), It.IsAny>()), + Times.Once); + } + + [Fact] + [Unit] + public async Task TokenExpiredRemovalCancelsConnectionWithoutPendingUpdate() + { + const string DeviceId = "device1"; + bool tokenUpdateCanceled = false; + Action statusChangedHandler = null; + var connectionClosed = new TaskCompletionSource( + TaskCreationOptions.RunContinuationsAsynchronously); + var cloudProxy = Mock.Of(p => p.IsActive); + var replacementCloudProxy = Mock.Of(p => p.IsActive); + var replacementCloudConnection = Mock.Of( + c => c.IsActive && c.CloudProxy == Option.Some(replacementCloudProxy)); + var cloudConnection = new Mock(); + cloudConnection.SetupGet(c => c.IsActive).Returns(true); + cloudConnection.SetupGet(c => c.CloudProxy).Returns(Option.Some(cloudProxy)); + cloudConnection.SetupGet(c => c.HasPendingTokenUpdate).Returns(false); + cloudConnection + .Setup(c => c.CancelTokenUpdate()) + .Callback(() => tokenUpdateCanceled = true); + cloudConnection + .Setup(c => c.CloseAsync()) + .Callback( + () => + { + Assert.True(tokenUpdateCanceled); + connectionClosed.TrySetResult(true); + }) + .ReturnsAsync(true); + var cloudConnectionProvider = new Mock(); + cloudConnectionProvider + .Setup(c => c.Connect(It.Is(i => i.Id == DeviceId), It.IsAny>())) + .Callback>((_, handler) => statusChangedHandler = handler) + .ReturnsAsync(Try.Success(cloudConnection.Object as ICloudConnection)); + cloudConnectionProvider + .Setup(c => c.Connect(It.Is(credentials => credentials.Identity.Id == DeviceId), It.IsAny>())) + .ReturnsAsync(Try.Success(replacementCloudConnection)); + var connectionManager = new ConnectionManager( + cloudConnectionProvider.Object, + Mock.Of(), + GetIdentityProvider(), + Mock.Of()); + + Option initialConnection = await connectionManager.GetCloudConnection(DeviceId); + Assert.True(initialConnection.HasValue); + Assert.NotNull(statusChangedHandler); + + statusChangedHandler(DeviceId, CloudConnectionStatus.DisconnectedTokenExpired); + await connectionClosed.Task; + + cloudConnection.Verify(c => c.CancelTokenUpdate(), Times.Once); + cloudConnection.Verify(c => c.CloseAsync(), Times.Once); + Option removedConnection = await connectionManager.GetCloudConnection(DeviceId); + Assert.False(removedConnection.HasValue); + + var replacementCredentials = new TokenCredentials( + new DeviceIdentity(IotHubHostName, DeviceId), + DummyToken, + DummyProductInfo, + Option.None(), + Option.None(), + true); + Try replacementConnection = + await connectionManager.CreateCloudConnectionAsync(replacementCredentials); + + Assert.True(replacementConnection.Success); + Assert.Same( + replacementCloudProxy, + ((RetryingCloudProxy)replacementConnection.Value).InnerCloudProxy); + cloudConnection.Verify( + c => c.UpdateTokenAsync(It.IsAny()), + Times.Never); + cloudConnectionProvider.Verify( + c => c.Connect(replacementCredentials, It.IsAny>()), + Times.Once); + } + [Fact] [Unit] public async Task CreateCloudProxyTest() diff --git a/edge-hub/core/test/Microsoft.Azure.Devices.Edge.Hub.Core.Test/RetryingCloudProxyTest.cs b/edge-hub/core/test/Microsoft.Azure.Devices.Edge.Hub.Core.Test/RetryingCloudProxyTest.cs index d3040285651..efdb9db7600 100644 --- a/edge-hub/core/test/Microsoft.Azure.Devices.Edge.Hub.Core.Test/RetryingCloudProxyTest.cs +++ b/edge-hub/core/test/Microsoft.Azure.Devices.Edge.Hub.Core.Test/RetryingCloudProxyTest.cs @@ -3,6 +3,7 @@ namespace Microsoft.Azure.Devices.Edge.Hub.Core.Test { using System; using System.Collections.Generic; + using System.Diagnostics; using System.Linq; using System.Net; using System.Threading; @@ -374,7 +375,7 @@ static ConnectionManager CreateConnectionManagerWithoutCloudConnectionRetryDelay 101, true, TimeSpan.Zero, - SystemTime.Instance); + Stopwatch.GetTimestamp); static async Task RunSendMessages(ICloudProxy cloudProxy, IEnumerable messages, int batchSize = 1) { From 901853a6fdead076a545760c05ea58ed7f2e7d09 Mon Sep 17 00:00:00 2001 From: yophilav Date: Fri, 7 Aug 2026 17:05:24 +0000 Subject: [PATCH 3/4] Fix cloud connection precedence Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 0b428cdd-2248-4f74-87b1-cd02dc9ad50a --- .../ConnectionManager.cs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/edge-hub/core/src/Microsoft.Azure.Devices.Edge.Hub.Core/ConnectionManager.cs b/edge-hub/core/src/Microsoft.Azure.Devices.Edge.Hub.Core/ConnectionManager.cs index 3a19b15df73..770d1e94ab1 100644 --- a/edge-hub/core/src/Microsoft.Azure.Devices.Edge.Hub.Core/ConnectionManager.cs +++ b/edge-hub/core/src/Microsoft.Azure.Devices.Edge.Hub.Core/ConnectionManager.cs @@ -656,8 +656,8 @@ public async Task> GetOrCreateCloudConnection( if (createTask != null && connectionGeneration == this.cloudConnectionCreateGeneration) { reuseCreateTask = !createTask.IsCompleted - || this.shouldThrottleCloudConnectionCreation - && !this.CloudConnectionRetryIntervalElapsed(); + || (this.shouldThrottleCloudConnectionCreation + && !this.CloudConnectionRetryIntervalElapsed()); if (reuseCreateTask && createTask.IsCompleted) { Events.ReusingRecentCloudConnectionAttempt(this.Identity, this.cloudConnectionRetryInterval); From 12403807051bce1f6286f5562771fd8298fcafe8 Mon Sep 17 00:00:00 2001 From: yophilav Date: Fri, 7 Aug 2026 20:42:36 +0000 Subject: [PATCH 4/4] Simplify cloud connection retry state Make retry timing and generation invariants explicit, avoid starting asynchronous creation under the state lock, and replace fixed-delay token tests with deterministic synchronization. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 0b428cdd-2248-4f74-87b1-cd02dc9ad50a --- .../ClientTokenCloudConnection.cs | 124 ++++++++----- .../ConnectionManager.cs | 171 ++++++++++-------- .../ClientTokenCloudConnectionTest.cs | 13 +- .../ConnectionManagerTest.cs | 29 ++- .../RetryingCloudProxyTest.cs | 13 +- 5 files changed, 204 insertions(+), 146 deletions(-) diff --git a/edge-hub/core/src/Microsoft.Azure.Devices.Edge.Hub.CloudProxy/ClientTokenCloudConnection.cs b/edge-hub/core/src/Microsoft.Azure.Devices.Edge.Hub.CloudProxy/ClientTokenCloudConnection.cs index 0dee76de0ff..e4cfabb84df 100644 --- a/edge-hub/core/src/Microsoft.Azure.Devices.Edge.Hub.CloudProxy/ClientTokenCloudConnection.cs +++ b/edge-hub/core/src/Microsoft.Azure.Devices.Edge.Hub.CloudProxy/ClientTokenCloudConnection.cs @@ -22,6 +22,7 @@ class ClientTokenCloudConnection : CloudConnection, IClientTokenCloudConnection static readonly TimeSpan TokenRetryWaitTime = TimeSpan.FromSeconds(20); readonly AsyncLock identityUpdateLock = new AsyncLock(); + readonly Func tokenRetryDelay; bool callbacksEnabled = true; TaskCompletionSource tokenGetter; @@ -40,7 +41,8 @@ class ClientTokenCloudConnection : CloudConnection, IClientTokenCloudConnection TimeSpan operationTimeout, TimeSpan cloudConnectionHangingTimeout, string productInfo, - Option modelId) + Option modelId, + Func tokenRetryDelay) : base( identity, connectionStatusChangedHandler, @@ -55,11 +57,12 @@ class ClientTokenCloudConnection : CloudConnection, IClientTokenCloudConnection productInfo, modelId) { + this.tokenRetryDelay = tokenRetryDelay; } protected override bool CallbacksEnabled => this.callbacksEnabled; - public static async Task Create( + public static Task Create( ITokenCredentials tokenCredentials, Action connectionStatusChangedHandler, ITransportSettings[] transportSettings, @@ -72,8 +75,38 @@ public static async Task Create( TimeSpan cloudConnectionHangingTimeout, string productInfo, Option modelId) + => Create( + tokenCredentials, + connectionStatusChangedHandler, + transportSettings, + messageConverterProvider, + clientProvider, + cloudListener, + idleTimeout, + closeOnIdleTimeout, + operationTimeout, + cloudConnectionHangingTimeout, + productInfo, + modelId, + () => Task.Delay(TokenRetryWaitTime)); + + internal static async Task Create( + ITokenCredentials tokenCredentials, + Action connectionStatusChangedHandler, + ITransportSettings[] transportSettings, + IMessageConverterProvider messageConverterProvider, + IClientProvider clientProvider, + ICloudListener cloudListener, + TimeSpan idleTimeout, + bool closeOnIdleTimeout, + TimeSpan operationTimeout, + TimeSpan cloudConnectionHangingTimeout, + string productInfo, + Option modelId, + Func tokenRetryDelay) { Preconditions.CheckNotNull(tokenCredentials, nameof(tokenCredentials)); + Preconditions.CheckNotNull(tokenRetryDelay, nameof(tokenRetryDelay)); var cloudConnection = new ClientTokenCloudConnection( tokenCredentials.Identity, connectionStatusChangedHandler, @@ -86,7 +119,8 @@ public static async Task Create( operationTimeout, cloudConnectionHangingTimeout, productInfo, - modelId); + modelId, + tokenRetryDelay); ITokenProvider tokenProvider = new ClientTokenBasedTokenProvider(tokenCredentials, cloudConnection); ICloudProxy cloudProxy = await cloudConnection.CreateNewCloudProxyAsync(tokenProvider); cloudConnection.cloudProxy = Option.Some(cloudProxy); @@ -94,17 +128,14 @@ public static async Task Create( } /// - /// This method does the following - - /// 1. Updates the Identity to be used for the cloud connection - /// 2. Updates the cloud proxy - - /// i. If there is an existing device client and - /// a. If is waiting for an updated token, and the Identity has a token, - /// then it uses that to give it to the waiting client authentication method. - /// b. If not, then it creates a new cloud proxy (and device client) and closes the existing one - /// ii. Else, if there is no cloud proxy, then opens a device client and creates a cloud proxy. + /// Applies new token credentials to the cloud connection. /// + /// + /// A pending token request reuses the existing proxy. Otherwise, a replacement proxy opens + /// before the existing proxy closes so invalid credentials do not disrupt an active connection. + /// /// New token credentials. - /// task of ICloudProxy interface. + /// The active cloud proxy. public async Task UpdateTokenAsync(ITokenCredentials newTokenCredentials) { Preconditions.CheckNotNull(newTokenCredentials, nameof(newTokenCredentials)); @@ -203,12 +234,15 @@ static bool IsTokenUsable(string hostname, string token) } /// - /// If the existing Identity has a usable token, then use it. - /// Else, generate a notification of token being near expiry and return a task that - /// can be completed later. - /// Keep retrying till we get a usable token. - /// Note - Don't use this.Identity in this method, as it may not have been set yet! + /// Returns the supplied token when usable; otherwise, requests replacement tokens until one is usable. /// + /// + /// Token validation uses the supplied value because the connection identity may not yet contain + /// the latest credentials. + /// + /// The token to validate first. + /// A usable token. + /// Token updates have been canceled. async Task GetNewToken(string currentToken) { Events.GetNewToken(this.Identity.Id); @@ -216,8 +250,6 @@ async Task GetNewToken(string currentToken) string token = currentToken; while (true) { - // We have to catch UnauthorizedAccessException, because on IsTokenUsable, we call parse from - // Device Client and it throws if the token is expired. if (IsTokenUsable(this.Identity.IotHubHostname, token)) { if (retrying) @@ -231,19 +263,11 @@ async Task GetNewToken(string currentToken) return token; } - else - { - Events.TokenNotUsable(this.Identity, token); - } - bool newTokenGetterCreated = false; - if (Volatile.Read(ref this.tokenUpdatesCanceled) != 0) - { - throw new OperationCanceledException( - $"Token updates for client {this.Identity.Id} have been canceled."); - } + Events.TokenNotUsable(this.Identity, token); - // No need to lock here as the lock is being held by the refresher. + bool tokenGetterPublished = false; + this.ThrowIfTokenUpdatesCanceled(cancelPublishedWaiter: false); TaskCompletionSource tcs = Volatile.Read(ref this.tokenGetter); if (tcs == null) { @@ -256,42 +280,44 @@ async Task GetNewToken(string currentToken) if (tcs == null) { tcs = taskCompletionSource; - newTokenGetterCreated = true; + tokenGetterPublished = true; } } - if (Volatile.Read(ref this.tokenUpdatesCanceled) != 0) - { - this.CancelTokenUpdate(); - throw new OperationCanceledException( - $"Token updates for client {this.Identity.Id} have been canceled."); - } + this.ThrowIfTokenUpdatesCanceled(cancelPublishedWaiter: true); - // If a new tokenGetter was created, then invoke the connection status changed handler - if (newTokenGetterCreated) + if (tokenGetterPublished) { - // If retrying, wait for some time. if (retrying) { - await Task.Delay(TokenRetryWaitTime); - } - - if (Volatile.Read(ref this.tokenUpdatesCanceled) != 0) - { - this.CancelTokenUpdate(); - throw new OperationCanceledException( - $"Token updates for client {this.Identity.Id} have been canceled."); + await this.tokenRetryDelay(); } + this.ThrowIfTokenUpdatesCanceled(cancelPublishedWaiter: true); this.ConnectionStatusChangedHandler(this.Identity.Id, CloudConnectionStatus.TokenNearExpiry); } retrying = true; - // this.tokenGetter will be reset when this task returns. token = await tcs.Task; } } + void ThrowIfTokenUpdatesCanceled(bool cancelPublishedWaiter) + { + if (Volatile.Read(ref this.tokenUpdatesCanceled) == 0) + { + return; + } + + if (cancelPublishedWaiter) + { + Interlocked.Exchange(ref this.tokenGetter, null)?.TrySetCanceled(); + } + + throw new OperationCanceledException( + $"Token updates for client {this.Identity.Id} have been canceled."); + } + class ClientTokenBasedTokenProvider : ITokenProvider { readonly ClientTokenCloudConnection cloudConnection; diff --git a/edge-hub/core/src/Microsoft.Azure.Devices.Edge.Hub.Core/ConnectionManager.cs b/edge-hub/core/src/Microsoft.Azure.Devices.Edge.Hub.Core/ConnectionManager.cs index 770d1e94ab1..d75dd4d0a3c 100644 --- a/edge-hub/core/src/Microsoft.Azure.Devices.Edge.Hub.Core/ConnectionManager.cs +++ b/edge-hub/core/src/Microsoft.Azure.Devices.Edge.Hub.Core/ConnectionManager.cs @@ -24,6 +24,11 @@ public class ConnectionManager : IConnectionManager { const int DefaultMaxClients = 101; // 100 Clients + 1 Edgehub static readonly TimeSpan DefaultCloudConnectionRetryInterval = TimeSpan.FromSeconds(5); + + // Retry throttling must not be skewed by wall-clock jumps, so it is measured against a + // process-wide monotonic clock. + static readonly Stopwatch MonotonicClock = Stopwatch.StartNew(); + readonly object deviceConnLock = new object(); readonly AsyncReaderWriterLock connectToCloudLock = new AsyncReaderWriterLock(); readonly ConcurrentDictionary devices = new ConcurrentDictionary(); @@ -34,7 +39,7 @@ public class ConnectionManager : IConnectionManager readonly IDeviceConnectivityManager connectivityManager; readonly bool closeCloudConnectionOnDeviceDisconnect; readonly TimeSpan cloudConnectionRetryInterval; - readonly Func getTimestamp; + readonly Func getMonotonicTime; public ConnectionManager( ICloudConnectionProvider cloudConnectionProvider, @@ -51,7 +56,7 @@ public ConnectionManager( maxClients, closeCloudConnectionOnDeviceDisconnect, DefaultCloudConnectionRetryInterval, - Stopwatch.GetTimestamp) + () => MonotonicClock.Elapsed) { } @@ -63,7 +68,7 @@ internal ConnectionManager( int maxClients, bool closeCloudConnectionOnDeviceDisconnect, TimeSpan cloudConnectionRetryInterval, - Func getTimestamp) + Func getMonotonicTime) { this.cloudConnectionProvider = Preconditions.CheckNotNull(cloudConnectionProvider, nameof(cloudConnectionProvider)); this.maxClients = Preconditions.CheckRange(maxClients, 1, nameof(maxClients)); @@ -73,7 +78,7 @@ internal ConnectionManager( this.cloudConnectionRetryInterval = cloudConnectionRetryInterval >= TimeSpan.Zero ? cloudConnectionRetryInterval : throw new ArgumentOutOfRangeException(nameof(cloudConnectionRetryInterval)); - this.getTimestamp = Preconditions.CheckNotNull(getTimestamp, nameof(getTimestamp)); + this.getMonotonicTime = Preconditions.CheckNotNull(getMonotonicTime, nameof(getMonotonicTime)); this.connectivityManager.DeviceDisconnected += (o, args) => this.HandleDeviceCloudConnectionDisconnected(); this.closeCloudConnectionOnDeviceDisconnect = closeCloudConnectionOnDeviceDisconnect; } @@ -108,7 +113,7 @@ await currentDeviceConnection public Task RemoveDeviceConnection(string id) { return this.devices.TryGetValue(Preconditions.CheckNonWhiteSpace(id, nameof(id)), out ConnectedDevice device) - ? this.RemoveDeviceConnection(device, this.closeCloudConnectionOnDeviceDisconnect) + ? this.RemoveDeviceConnection(device, removeCloudConnection: this.closeCloudConnectionOnDeviceDisconnect) : Task.CompletedTask; } @@ -292,7 +297,7 @@ await device.DeviceConnection.Filter(dp => dp.IsActive) if (removeCloudConnection) { await device.RemoveCloudConnection( - throttleReconnect, + throttleReconnect: throttleReconnect, preserveConnection: throttleReconnect); } @@ -354,26 +359,26 @@ await clientCredentials.ForEachAsync( Try cloudConnectionTry = await device.CreateOrUpdateCloudConnection(c => this.CreateOrUpdateCloudConnection(c, tokenCredentials)); if (!cloudConnectionTry.Success) { - await this.RemoveDeviceConnection(device, true, true); + await this.RemoveDeviceConnection(device, removeCloudConnection: true, throttleReconnect: true); this.CloudConnectionLost?.Invoke(this, device.Identity); } } else { - await this.RemoveDeviceConnection(device, this.closeCloudConnectionOnDeviceDisconnect); + await this.RemoveDeviceConnection(device, removeCloudConnection: this.closeCloudConnectionOnDeviceDisconnect); } }); } else { - await this.RemoveDeviceConnection(device, true, true); + await this.RemoveDeviceConnection(device, removeCloudConnection: true, throttleReconnect: true); this.CloudConnectionLost?.Invoke(this, device.Identity); } break; case CloudConnectionStatus.DisconnectedTokenExpired: - await this.RemoveDeviceConnection(device, true, true); + await this.RemoveDeviceConnection(device, removeCloudConnection: true, throttleReconnect: true); Events.InvokingCloudConnectionLostEvent(device.Identity); this.CloudConnectionLost?.Invoke(this, device.Identity); break; @@ -442,7 +447,7 @@ ConnectedDevice CreateNewConnectedDevice(IIdentity identity) throw new EdgeHubConnectionException($"Edge hub already has maximum allowed clients ({this.maxClients - 1}) connected."); } - return new ConnectedDevice(identity, this.cloudConnectionRetryInterval, this.getTimestamp); + return new ConnectedDevice(identity, this.cloudConnectionRetryInterval, this.getMonotonicTime); } } @@ -464,6 +469,12 @@ async Task> ConnectToCloud(IClientCredentials credentials, class ConnectedDevice { + // Generation parity is the removal flag: even means the cloud connection state is stable, + // odd means a removal is in progress. A connection attempt captures the generation it + // started under and may only publish its result while that generation is still current, + // so any interleaving removal discards it. + const long RemovalInProgressBit = 1; + // Device Proxy methods are sync coming from the Protocol gateway, // so using traditional locking mechanism for those. readonly object deviceProxyLock = new object(); @@ -471,40 +482,22 @@ class ConnectedDevice readonly AsyncLock cloudConnectionLock = new AsyncLock(); readonly AsyncLock cloudConnectionRemovalLock = new AsyncLock(); readonly TimeSpan cloudConnectionRetryInterval; - readonly Func getTimestamp; + readonly Func getMonotonicTime; ICloudConnection cloudConnection; ICloudConnection preservedCloudConnection; DeviceConnection deviceConnection; IIdentity identity; Task> cloudConnectionCreateTask; long cloudConnectionCreateGeneration; - long cloudConnectionCreateCompletedTimestamp; long cloudConnectionGeneration; - bool hasCloudConnectionCreateCompletedTimestamp; + TimeSpan? cloudConnectionCreateCompletedTime; bool shouldThrottleCloudConnectionCreation; - public ConnectedDevice(IIdentity identity, TimeSpan cloudConnectionRetryInterval, Func getTimestamp) - : this( - identity, - Option.None(), - Option.None(), - cloudConnectionRetryInterval, - getTimestamp) - { - } - - public ConnectedDevice( - IIdentity identity, - Option cloudProxy, - Option deviceConnection, - TimeSpan cloudConnectionRetryInterval, - Func getTimestamp) + public ConnectedDevice(IIdentity identity, TimeSpan cloudConnectionRetryInterval, Func getMonotonicTime) { this.identity = identity; - this.cloudConnection = cloudProxy.OrDefault(); - this.deviceConnection = deviceConnection.OrDefault(); this.cloudConnectionRetryInterval = cloudConnectionRetryInterval; - this.getTimestamp = getTimestamp; + this.getMonotonicTime = getMonotonicTime; } public IIdentity Identity => Volatile.Read(ref this.identity); @@ -570,7 +563,7 @@ public async Task> CreateOrUpdateCloudConnection( Func>> cloudConnectionUpdater) { Preconditions.CheckNotNull(cloudConnectionUpdater, nameof(cloudConnectionUpdater)); - await this.WaitForCloudConnectionRemoval(); + await this.WaitForPendingCloudConnectionRemoval(); // Lock in case multiple connections are created to the cloud for the same device at the same time using (await this.cloudConnectionLock.LockAsync()) { @@ -597,9 +590,8 @@ public async Task> CreateOrUpdateCloudConnection( } else { - this.RecordCloudConnectionAttemptCore( + this.RecordFailedCloudConnectionAttemptCore( Task.FromResult(newCloudConnection), - true, connectionGeneration); } } @@ -642,7 +634,8 @@ public async Task> GetOrCreateCloudConnection( } } - bool reuseCreateTask = false; + TaskCompletionSource> createTaskSource = null; + bool replacesExistingConnection = false; lock (this.cloudConnectionStateLock) { connectionGeneration = Volatile.Read(ref this.cloudConnectionGeneration); @@ -652,6 +645,7 @@ public async Task> GetOrCreateCloudConnection( new EdgeHubConnectionException($"Cloud connection for device {this.Identity.Id} is being removed.")); } + bool reuseCreateTask = false; createTask = this.cloudConnectionCreateTask; if (createTask != null && connectionGeneration == this.cloudConnectionCreateGeneration) { @@ -666,16 +660,33 @@ public async Task> GetOrCreateCloudConnection( if (!reuseCreateTask) { - bool replacesExistingConnection = this.CloudConnectionForUpdate.HasValue; + replacesExistingConnection = this.CloudConnectionForUpdate.HasValue; + // Publish the placeholder under the lock so concurrent callers join this attempt, + // but run the caller-supplied creator only after the lock is released. + createTaskSource = new TaskCompletionSource>( + TaskCreationOptions.RunContinuationsAsynchronously); + createTask = createTaskSource.Task; Volatile.Write(ref this.cloudConnectionCreateGeneration, connectionGeneration); - createTask = this.CreateCloudConnection( - cloudConnectionCreator, - replacesExistingConnection, - connectionGeneration); Volatile.Write(ref this.cloudConnectionCreateTask, createTask); } } + if (createTaskSource != null) + { + try + { + createTaskSource.SetResult( + await this.CreateCloudConnection( + cloudConnectionCreator, + replacesExistingConnection, + connectionGeneration)); + } + catch (Exception ex) + { + createTaskSource.SetException(ex); + } + } + return await createTask; } } @@ -721,12 +732,7 @@ async Task> CreateCloudConnection( if (displacedPreservedConnection != null && !ReferenceEquals(displacedPreservedConnection, cloudConnectionResult.Value)) { - if (displacedPreservedConnection is IClientTokenCloudConnection clientTokenCloudConnection) - { - clientTokenCloudConnection.CancelTokenUpdate(); - } - - await displacedPreservedConnection.CloseAsync(); + await CancelTokenUpdateAndCloseAsync(displacedPreservedConnection); } return cloudConnectionResult; @@ -737,8 +743,7 @@ async Task> CreateCloudConnection( { if (this.IsCloudConnectionGenerationValid(connectionGeneration)) { - this.cloudConnectionCreateCompletedTimestamp = this.getTimestamp(); - this.hasCloudConnectionCreateCompletedTimestamp = true; + this.cloudConnectionCreateCompletedTime = this.getMonotonicTime(); if (createTask == null || createTask.IsFaulted || createTask.IsCanceled) { this.shouldThrottleCloudConnectionCreation = true; @@ -748,7 +753,11 @@ async Task> CreateCloudConnection( } } - async Task WaitForCloudConnectionRemoval() + // Barrier: acquiring and immediately releasing the removal lock guarantees that a removal + // already in flight has published its final generation and throttle state before the caller + // reads them. The lock is deliberately not held across the caller's work, because a cloud + // callback raised during that work can itself trigger a removal and would deadlock on it. + async Task WaitForPendingCloudConnectionRemoval() { using (await this.cloudConnectionRemovalLock.LockAsync()) { @@ -757,19 +766,17 @@ async Task WaitForCloudConnectionRemoval() bool CloudConnectionRetryIntervalElapsed() { - if (!this.hasCloudConnectionCreateCompletedTimestamp) + if (!this.cloudConnectionCreateCompletedTime.HasValue) { return true; } - long elapsedTimestamp = this.getTimestamp() - this.cloudConnectionCreateCompletedTimestamp; - return elapsedTimestamp < 0 - || elapsedTimestamp >= this.cloudConnectionRetryInterval.TotalSeconds * Stopwatch.Frequency; + TimeSpan elapsed = this.getMonotonicTime() - this.cloudConnectionCreateCompletedTime.Value; + return elapsed < TimeSpan.Zero || elapsed >= this.cloudConnectionRetryInterval; } - void RecordCloudConnectionAttemptCore( + void RecordFailedCloudConnectionAttemptCore( Task> createTask, - bool shouldThrottle, long connectionGeneration) { if (!this.IsCloudConnectionGenerationValid(connectionGeneration)) @@ -779,9 +786,8 @@ void RecordCloudConnectionAttemptCore( Volatile.Write(ref this.cloudConnectionCreateGeneration, connectionGeneration); Volatile.Write(ref this.cloudConnectionCreateTask, createTask); - this.cloudConnectionCreateCompletedTimestamp = this.getTimestamp(); - this.hasCloudConnectionCreateCompletedTimestamp = true; - this.shouldThrottleCloudConnectionCreation = shouldThrottle; + this.cloudConnectionCreateCompletedTime = this.getMonotonicTime(); + this.shouldThrottleCloudConnectionCreation = true; } void ResetCloudConnectionRetryStateCore(long connectionGeneration) @@ -793,8 +799,7 @@ void ResetCloudConnectionRetryStateCore(long connectionGeneration) Volatile.Write(ref this.cloudConnectionCreateGeneration, connectionGeneration); Volatile.Write(ref this.cloudConnectionCreateTask, null); - this.cloudConnectionCreateCompletedTimestamp = 0; - this.hasCloudConnectionCreateCompletedTimestamp = false; + this.cloudConnectionCreateCompletedTime = null; this.shouldThrottleCloudConnectionCreation = false; } @@ -808,7 +813,7 @@ public async Task RemoveCloudConnection( var supersededCloudConnections = new List(1); lock (this.cloudConnectionStateLock) { - Interlocked.Increment(ref this.cloudConnectionGeneration); + this.BeginCloudConnectionRemoval(); ICloudConnection activeCloudConnection = Interlocked.Exchange(ref this.cloudConnection, null); if (activeCloudConnection != null) @@ -900,10 +905,11 @@ public async Task RemoveCloudConnection( { lock (this.cloudConnectionStateLock) { - long stableGeneration = Volatile.Read(ref this.cloudConnectionGeneration) + 1; + long stableGeneration = this.NextStableCloudConnectionGeneration(); Volatile.Write(ref this.cloudConnectionCreateGeneration, stableGeneration); - this.cloudConnectionCreateCompletedTimestamp = this.getTimestamp(); - this.hasCloudConnectionCreateCompletedTimestamp = throttleReconnect; + this.cloudConnectionCreateCompletedTime = throttleReconnect + ? this.getMonotonicTime() + : (TimeSpan?)null; this.shouldThrottleCloudConnectionCreation = throttleReconnect; Volatile.Write( ref this.cloudConnectionCreateTask, @@ -913,14 +919,36 @@ public async Task RemoveCloudConnection( new EdgeHubConnectionException( $"Cloud connection for device {this.Identity.Id} was removed after a cloud failure."))) : null); - Interlocked.Increment(ref this.cloudConnectionGeneration); + this.CompleteCloudConnectionRemoval(); } } } } + static bool IsStableGeneration(long generation) => (generation & RemovalInProgressBit) == 0; + + // Cancelling first is what makes the close final: an in-flight token update would otherwise + // hand back a live client for a connection the caller has already discarded. + static async Task CancelTokenUpdateAndCloseAsync(ICloudConnection cloudConnection) + { + if (cloudConnection is IClientTokenCloudConnection clientTokenCloudConnection) + { + clientTokenCloudConnection.CancelTokenUpdate(); + } + + await cloudConnection.CloseAsync(); + } + + void BeginCloudConnectionRemoval() => Interlocked.Increment(ref this.cloudConnectionGeneration); + + void CompleteCloudConnectionRemoval() => Interlocked.Increment(ref this.cloudConnectionGeneration); + + // Only valid while a removal is in progress: the generation is odd, so the next increment + // restores the stable generation that post-removal connection attempts will run under. + long NextStableCloudConnectionGeneration() => Volatile.Read(ref this.cloudConnectionGeneration) + 1; + bool IsCloudConnectionGenerationValid(long connectionGeneration) => - (connectionGeneration & 1) == 0 + IsStableGeneration(connectionGeneration) && connectionGeneration == Volatile.Read(ref this.cloudConnectionGeneration); async Task> DiscardInvalidatedCloudConnection( @@ -936,12 +964,7 @@ async Task> DiscardInvalidatedCloudConnection( ref this.preservedCloudConnection, null, cloudConnection.Value); - if (cloudConnection.Value is IClientTokenCloudConnection clientTokenCloudConnection) - { - clientTokenCloudConnection.CancelTokenUpdate(); - } - - await cloudConnection.Value.CloseAsync(); + await CancelTokenUpdateAndCloseAsync(cloudConnection.Value); } return Try.Failure( diff --git a/edge-hub/core/test/Microsoft.Azure.Devices.Edge.Hub.CloudProxy.Test/ClientTokenCloudConnectionTest.cs b/edge-hub/core/test/Microsoft.Azure.Devices.Edge.Hub.CloudProxy.Test/ClientTokenCloudConnectionTest.cs index 0d05a463681..920c653875c 100644 --- a/edge-hub/core/test/Microsoft.Azure.Devices.Edge.Hub.CloudProxy.Test/ClientTokenCloudConnectionTest.cs +++ b/edge-hub/core/test/Microsoft.Azure.Devices.Edge.Hub.CloudProxy.Test/ClientTokenCloudConnectionTest.cs @@ -270,6 +270,8 @@ public async Task CancellationDuringRetrySuppressesTokenStatusCallback() var transportSettings = new ITransportSettings[] { new AmqpTransportSettings(TransportType.Amqp_Tcp_Only) }; var receivedStatuses = new List(); + var retryDelayEntered = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + var releaseRetryDelay = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); var messageConverterProvider = new MessageConverterProvider( new Dictionary { @@ -287,19 +289,24 @@ public async Task CancellationDuringRetrySuppressesTokenStatusCallback() TimeSpan.FromSeconds(20), TimeSpan.FromSeconds(50), DummyProductInfo, - Option.None()); + Option.None(), + () => + { + retryDelayEntered.TrySetResult(true); + return releaseRetryDelay.Task; + }); Assert.NotNull(tokenProvider); Task getTokenTask = tokenProvider.GetTokenAsync(Option.None()); Assert.Single(receivedStatuses); await cloudConnection.UpdateTokenAsync(credentials); - await Task.Delay(TimeSpan.FromSeconds(1)); + await retryDelayEntered.Task; Assert.True(cloudConnection.HasPendingTokenUpdate); cloudConnection.CancelTokenUpdate(); + releaseRetryDelay.TrySetResult(true); await Assert.ThrowsAnyAsync(() => getTokenTask); - await Task.Delay(TimeSpan.FromSeconds(20)); Assert.Single(receivedStatuses); Assert.False(cloudConnection.HasPendingTokenUpdate); } diff --git a/edge-hub/core/test/Microsoft.Azure.Devices.Edge.Hub.Core.Test/ConnectionManagerTest.cs b/edge-hub/core/test/Microsoft.Azure.Devices.Edge.Hub.Core.Test/ConnectionManagerTest.cs index 2dbf1110e3d..82f64affe76 100644 --- a/edge-hub/core/test/Microsoft.Azure.Devices.Edge.Hub.Core.Test/ConnectionManagerTest.cs +++ b/edge-hub/core/test/Microsoft.Azure.Devices.Edge.Hub.Core.Test/ConnectionManagerTest.cs @@ -3,7 +3,6 @@ namespace Microsoft.Azure.Devices.Edge.Hub.Core.Test { using System; using System.Collections.Generic; - using System.Diagnostics; using System.Linq; using System.Net; using System.Threading.Tasks; @@ -580,7 +579,7 @@ public async Task GetCloudConnectionThrottlesRepeatedFailedCreation() const string DeviceId = "device1"; const int OperationCount = 200; TimeSpan retryInterval = TimeSpan.FromSeconds(5); - long timestamp = 0; + TimeSpan monotonicTime = TimeSpan.Zero; var cloudConnectionProvider = new Mock(); cloudConnectionProvider .Setup(c => c.Connect(It.Is(i => i.Id == DeviceId), It.IsAny>())) @@ -591,10 +590,10 @@ public async Task GetCloudConnectionThrottlesRepeatedFailedCreation() Mock.Of(), GetIdentityProvider(), Mock.Of(), - 101, - true, + maxClients: 101, + closeCloudConnectionOnDeviceDisconnect: true, retryInterval, - () => timestamp); + () => monotonicTime); Task>[] operations = Enumerable.Range(0, OperationCount) .Select(_ => connectionManager.GetCloudConnection(DeviceId)) @@ -606,7 +605,7 @@ public async Task GetCloudConnectionThrottlesRepeatedFailedCreation() c => c.Connect(It.IsAny(), It.IsAny>()), Times.Once); - timestamp += (long)(retryInterval.TotalSeconds * Stopwatch.Frequency); + monotonicTime += retryInterval; Option resultAfterRetryInterval = await connectionManager.GetCloudConnection(DeviceId); Assert.False(resultAfterRetryInterval.HasValue); @@ -622,7 +621,7 @@ public async Task GetCloudConnectionThrottlesImmediatelyInactiveReplacements() const string DeviceId = "device1"; const int OperationCount = 200; TimeSpan retryInterval = TimeSpan.FromSeconds(5); - long timestamp = 0; + TimeSpan monotonicTime = TimeSpan.Zero; bool initialConnectionActive = true; var initialCloudProxy = new Mock(); @@ -666,10 +665,10 @@ Try CreateFailingReplacement() Mock.Of(), GetIdentityProvider(), Mock.Of(), - 101, - true, + maxClients: 101, + closeCloudConnectionOnDeviceDisconnect: true, retryInterval, - () => timestamp); + () => monotonicTime); Option cloudProxy = await connectionManager.GetCloudConnection(DeviceId); Assert.True(cloudProxy.HasValue); @@ -687,7 +686,7 @@ Try CreateFailingReplacement() c => c.Connect(It.IsAny(), It.IsAny>()), Times.Exactly(2)); - timestamp += (long)(retryInterval.TotalSeconds * Stopwatch.Frequency); + monotonicTime += retryInterval; await Assert.ThrowsAsync( () => cloudProxy.OrDefault().SendMessageAsync(message)); @@ -809,7 +808,7 @@ public async Task TokenExpiredRemovalPreservesConnectionOnlyForTokenUpdate() const string DeviceId = "device1"; const int OperationCount = 200; TimeSpan retryInterval = TimeSpan.FromSeconds(5); - long timestamp = 0; + TimeSpan monotonicTime = TimeSpan.Zero; Action statusChangedHandler = null; var connectionRemoved = new TaskCompletionSource( TaskCreationOptions.RunContinuationsAsynchronously); @@ -848,10 +847,10 @@ public async Task TokenExpiredRemovalPreservesConnectionOnlyForTokenUpdate() Mock.Of(), GetIdentityProvider(), Mock.Of(), - 101, - true, + maxClients: 101, + closeCloudConnectionOnDeviceDisconnect: true, retryInterval, - () => timestamp); + () => monotonicTime); Option initialCloudProxy = await connectionManager.GetCloudConnection(DeviceId); Assert.True(initialCloudProxy.HasValue); diff --git a/edge-hub/core/test/Microsoft.Azure.Devices.Edge.Hub.Core.Test/RetryingCloudProxyTest.cs b/edge-hub/core/test/Microsoft.Azure.Devices.Edge.Hub.Core.Test/RetryingCloudProxyTest.cs index efdb9db7600..03d38aa67b2 100644 --- a/edge-hub/core/test/Microsoft.Azure.Devices.Edge.Hub.Core.Test/RetryingCloudProxyTest.cs +++ b/edge-hub/core/test/Microsoft.Azure.Devices.Edge.Hub.Core.Test/RetryingCloudProxyTest.cs @@ -366,16 +366,19 @@ static ConnectionManager CreateConnectionManagerWithoutCloudConnectionRetryDelay ICloudConnectionProvider connectionProvider, ICredentialsCache credentialsCache, IIdentityProvider identityProvider, - IDeviceConnectivityManager deviceConnectivityManager) => - new ConnectionManager( + IDeviceConnectivityManager deviceConnectivityManager) + { + var monotonicClock = Stopwatch.StartNew(); + return new ConnectionManager( connectionProvider, credentialsCache, identityProvider, deviceConnectivityManager, - 101, - true, + maxClients: 101, + closeCloudConnectionOnDeviceDisconnect: true, TimeSpan.Zero, - Stopwatch.GetTimestamp); + () => monotonicClock.Elapsed); + } static async Task RunSendMessages(ICloudProxy cloudProxy, IEnumerable messages, int batchSize = 1) {