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..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 @@ -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; @@ -21,9 +22,11 @@ class ClientTokenCloudConnection : CloudConnection, IClientTokenCloudConnection static readonly TimeSpan TokenRetryWaitTime = TimeSpan.FromSeconds(20); readonly AsyncLock identityUpdateLock = new AsyncLock(); + readonly Func tokenRetryDelay; bool callbacksEnabled = true; - Option> tokenGetter; + TaskCompletionSource tokenGetter; + int tokenUpdatesCanceled; Option cloudProxy; ClientTokenCloudConnection( @@ -38,7 +41,8 @@ class ClientTokenCloudConnection : CloudConnection, IClientTokenCloudConnection TimeSpan operationTimeout, TimeSpan cloudConnectionHangingTimeout, string productInfo, - Option modelId) + Option modelId, + Func tokenRetryDelay) : base( identity, connectionStatusChangedHandler, @@ -53,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, @@ -70,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, @@ -84,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); @@ -92,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)); @@ -123,20 +156,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 +170,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 +199,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) { @@ -188,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); @@ -201,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) @@ -216,42 +263,61 @@ async Task GetNewToken(string currentToken) return token; } - else + + Events.TokenNotUsable(this.Identity, token); + + bool tokenGetterPublished = false; + this.ThrowIfTokenUpdatesCanceled(cancelPublishedWaiter: false); + TaskCompletionSource tcs = Volatile.Read(ref this.tokenGetter); + if (tcs == null) { - Events.TokenNotUsable(this.Identity, token); + Events.SafeCreateNewToken(this.Identity.Id); + var taskCompletionSource = new TaskCompletionSource(); + tcs = Interlocked.CompareExchange( + ref this.tokenGetter, + taskCompletionSource, + null); + if (tcs == null) + { + tcs = taskCompletionSource; + tokenGetterPublished = true; + } } - bool newTokenGetterCreated = false; - // 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; - }); - - // If a new tokenGetter was created, then invoke the connection status changed handler - if (newTokenGetterCreated) + this.ThrowIfTokenUpdatesCanceled(cancelPublishedWaiter: true); + + if (tokenGetterPublished) { - // If retrying, wait for some time. if (retrying) { - await Task.Delay(TokenRetryWaitTime); + 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 06a3273b534..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 @@ -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; @@ -21,6 +23,12 @@ 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); + + // 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(); @@ -30,6 +38,8 @@ public class ConnectionManager : IConnectionManager readonly IIdentityProvider identityProvider; readonly IDeviceConnectivityManager connectivityManager; readonly bool closeCloudConnectionOnDeviceDisconnect; + readonly TimeSpan cloudConnectionRetryInterval; + readonly Func getMonotonicTime; public ConnectionManager( ICloudConnectionProvider cloudConnectionProvider, @@ -38,12 +48,37 @@ public ConnectionManager( IDeviceConnectivityManager connectivityManager, int maxClients = DefaultMaxClients, bool closeCloudConnectionOnDeviceDisconnect = true) + : this( + cloudConnectionProvider, + credentialsCache, + identityProvider, + connectivityManager, + maxClients, + closeCloudConnectionOnDeviceDisconnect, + DefaultCloudConnectionRetryInterval, + () => MonotonicClock.Elapsed) + { + } + + internal ConnectionManager( + ICloudConnectionProvider cloudConnectionProvider, + ICredentialsCache credentialsCache, + IIdentityProvider identityProvider, + IDeviceConnectivityManager connectivityManager, + int maxClients, + bool closeCloudConnectionOnDeviceDisconnect, + TimeSpan cloudConnectionRetryInterval, + Func getMonotonicTime) { 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.getMonotonicTime = Preconditions.CheckNotNull(getMonotonicTime, nameof(getMonotonicTime)); this.connectivityManager.DeviceDisconnected += (o, args) => this.HandleDeviceCloudConnectionDisconnected(); this.closeCloudConnectionOnDeviceDisconnect = closeCloudConnectionOnDeviceDisconnect; } @@ -78,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; } @@ -235,7 +270,8 @@ 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)); Events.GetCloudConnection(credentials.Identity, cloudConnectionTry); Try cloudProxyTry = GetCloudProxyFromCloudConnection(cloudConnectionTry, credentials.Identity); return cloudProxyTry.Success @@ -248,7 +284,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); @@ -257,8 +296,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: throttleReconnect, + preserveConnection: throttleReconnect); } Events.RemoveDeviceConnection(id); @@ -267,7 +307,7 @@ await device.CloudConnection.Filter(cp => cp.IsActive) } Task> CreateOrUpdateCloudConnection(ConnectedDevice device, IClientCredentials credentials) => - device.CloudConnection.Map( + device.CloudConnectionForUpdate.Map( async c => { try @@ -319,26 +359,26 @@ 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, 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); + await this.RemoveDeviceConnection(device, removeCloudConnection: true, throttleReconnect: true); this.CloudConnectionLost?.Invoke(this, device.Identity); } break; case CloudConnectionStatus.DisconnectedTokenExpired: - await this.RemoveDeviceConnection(device, true); + await this.RemoveDeviceConnection(device, removeCloudConnection: true, throttleReconnect: true); Events.InvokingCloudConnectionLostEvent(device.Identity); this.CloudConnectionLost?.Invoke(this, device.Identity); break; @@ -357,18 +397,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); } } } @@ -388,7 +431,11 @@ ConnectedDevice CreateOrUpdateConnectedDevice(IIdentity identity) return this.devices.AddOrUpdate( deviceId, id => this.CreateNewConnectedDevice(identity), - (id, cd) => new ConnectedDevice(identity, cd.CloudConnection, cd.DeviceConnection)); + (id, cd) => + { + cd.UpdateIdentity(identity); + return cd; + }); } ConnectedDevice CreateNewConnectedDevice(IIdentity identity) @@ -400,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); + return new ConnectedDevice(identity, this.cloudConnectionRetryInterval, this.getMonotonicTime); } } @@ -422,39 +469,93 @@ 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(); + readonly object cloudConnectionStateLock = new object(); readonly AsyncLock cloudConnectionLock = new AsyncLock(); - Option>> cloudConnectionCreateTask = Option.None>>(); - - public ConnectedDevice(IIdentity identity) - : this(identity, Option.None(), Option.None()) + readonly AsyncLock cloudConnectionRemovalLock = new AsyncLock(); + readonly TimeSpan cloudConnectionRetryInterval; + readonly Func getMonotonicTime; + ICloudConnection cloudConnection; + ICloudConnection preservedCloudConnection; + DeviceConnection deviceConnection; + IIdentity identity; + Task> cloudConnectionCreateTask; + long cloudConnectionCreateGeneration; + long cloudConnectionGeneration; + TimeSpan? cloudConnectionCreateCompletedTime; + bool shouldThrottleCloudConnectionCreation; + + public ConnectedDevice(IIdentity identity, TimeSpan cloudConnectionRetryInterval, Func getMonotonicTime) { + this.identity = identity; + this.cloudConnectionRetryInterval = cloudConnectionRetryInterval; + this.getMonotonicTime = getMonotonicTime; } - public ConnectedDevice(IIdentity identity, Option cloudProxy, Option deviceConnection) + public IIdentity Identity => Volatile.Read(ref this.identity); + + public Option CloudConnection { - this.Identity = identity; - this.CloudConnection = cloudProxy; - this.DeviceConnection = deviceConnection; + get + { + ICloudConnection currentCloudConnection = Volatile.Read(ref this.cloudConnection); + return currentCloudConnection != null + ? Option.Some(currentCloudConnection) + : Option.None(); + } } - public IIdentity Identity { get; } - - 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(); } } @@ -462,13 +563,43 @@ public async Task> CreateOrUpdateCloudConnection( Func>> cloudConnectionUpdater) { Preconditions.CheckNotNull(cloudConnectionUpdater, nameof(cloudConnectionUpdater)); + 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()) { + long connectionGeneration = Volatile.Read(ref this.cloudConnectionGeneration); Try newCloudConnection = await cloudConnectionUpdater(this); - if (newCloudConnection.Success) + bool invalidated; + lock (this.cloudConnectionStateLock) + { + 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.RecordFailedCloudConnectionAttemptCore( + Task.FromResult(newCloudConnection), + connectionGeneration); + } + } + } + + if (invalidated) { - this.CloudConnection = Option.Some(newCloudConnection.Value); + return await this.DiscardInvalidatedCloudConnection(newCloudConnection); } return newCloudConnection; @@ -480,36 +611,364 @@ public async Task> GetOrCreateCloudConnection( { 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 Try.Success(activeCloudConnection.OrDefault()); + } + } + + TaskCompletionSource> createTaskSource = null; + bool replacesExistingConnection = false; + lock (this.cloudConnectionStateLock) + { + connectionGeneration = Volatile.Read(ref this.cloudConnectionGeneration); + if (!this.IsCloudConnectionGenerationValid(connectionGeneration)) { - 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 () => - { - return await this.cloudConnectionCreateTask.Filter(c => !c.IsCompleted) - .GetOrElse( - async () => - { - Task> createTask = cloudConnectionCreator(this); - this.cloudConnectionCreateTask = Option.Some(createTask); - Try cloudConnectionResult = await createTask; - this.CloudConnection = cloudConnectionResult.Ok(); - return cloudConnectionResult; - }); - }); - } - }); - }); + return Try.Failure( + new EdgeHubConnectionException($"Cloud connection for device {this.Identity.Id} is being removed.")); + } + + bool reuseCreateTask = false; + 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) + { + 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); + 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; + } + } + + 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)) + { + await CancelTokenUpdateAndCloseAsync(displacedPreservedConnection); + } + + return cloudConnectionResult; + } + finally + { + lock (this.cloudConnectionStateLock) + { + if (this.IsCloudConnectionGenerationValid(connectionGeneration)) + { + this.cloudConnectionCreateCompletedTime = this.getMonotonicTime(); + if (createTask == null || createTask.IsFaulted || createTask.IsCanceled) + { + this.shouldThrottleCloudConnectionCreation = true; + } + } + } + } + } + + // 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()) + { + } + } + + bool CloudConnectionRetryIntervalElapsed() + { + if (!this.cloudConnectionCreateCompletedTime.HasValue) + { + return true; + } + + TimeSpan elapsed = this.getMonotonicTime() - this.cloudConnectionCreateCompletedTime.Value; + return elapsed < TimeSpan.Zero || elapsed >= this.cloudConnectionRetryInterval; + } + + void RecordFailedCloudConnectionAttemptCore( + Task> createTask, + long connectionGeneration) + { + if (!this.IsCloudConnectionGenerationValid(connectionGeneration)) + { + return; + } + + Volatile.Write(ref this.cloudConnectionCreateGeneration, connectionGeneration); + Volatile.Write(ref this.cloudConnectionCreateTask, createTask); + this.cloudConnectionCreateCompletedTime = this.getMonotonicTime(); + this.shouldThrottleCloudConnectionCreation = true; + } + + void ResetCloudConnectionRetryStateCore(long connectionGeneration) + { + if (!this.IsCloudConnectionGenerationValid(connectionGeneration)) + { + return; + } + + Volatile.Write(ref this.cloudConnectionCreateGeneration, connectionGeneration); + Volatile.Write(ref this.cloudConnectionCreateTask, null); + this.cloudConnectionCreateCompletedTime = null; + 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) + { + this.BeginCloudConnectionRemoval(); + 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 = this.NextStableCloudConnectionGeneration(); + Volatile.Write(ref this.cloudConnectionCreateGeneration, stableGeneration); + this.cloudConnectionCreateCompletedTime = throttleReconnect + ? this.getMonotonicTime() + : (TimeSpan?)null; + 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); + 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) => + IsStableGeneration(connectionGeneration) + && 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); + await CancelTokenUpdateAndCloseAsync(cloudConnection.Value); + } + + return Try.Failure( + new EdgeHubConnectionException($"Cloud connection attempt for device {this.Identity.Id} was invalidated.")); } } @@ -550,7 +1009,8 @@ enum EventIds HandlingConnectionStatusChangedHandler, CloudConnectionLostClosingClient, CloudConnectionLostClosingAllClients, - GettingCloudConnectionForDeviceSubscriptions + GettingCloudConnectionForDeviceSubscriptions, + ReusingRecentCloudConnectionAttempt } public static void NewCloudConnection(IIdentity identity, Try cloudConnection) @@ -626,6 +1086,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..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 @@ -206,6 +206,111 @@ 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 retryDelayEntered = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + var releaseRetryDelay = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + 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(), + () => + { + retryDelayEntered.TrySetResult(true); + return releaseRetryDelay.Task; + }); + Assert.NotNull(tokenProvider); + + Task getTokenTask = tokenProvider.GetTokenAsync(Option.None()); + Assert.Single(receivedStatuses); + await cloudConnection.UpdateTokenAsync(credentials); + await retryDelayEntered.Task; + Assert.True(cloudConnection.HasPendingTokenUpdate); + + cloudConnection.CancelTokenUpdate(); + releaseRetryDelay.TrySetResult(true); + + await Assert.ThrowsAnyAsync(() => getTokenTask); + 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 b17368a4e7b..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 @@ -572,6 +572,402 @@ public async Task GetOrCreateCloudProxyTest() cloudProxyProviderMock.Verify(c => c.Connect(It.IsAny(), It.IsAny>()), Times.Exactly(2)); } + [Fact] + [Unit] + public async Task GetCloudConnectionThrottlesRepeatedFailedCreation() + { + const string DeviceId = "device1"; + const int OperationCount = 200; + TimeSpan retryInterval = TimeSpan.FromSeconds(5); + TimeSpan monotonicTime = TimeSpan.Zero; + var cloudConnectionProvider = new Mock(); + cloudConnectionProvider + .Setup(c => c.Connect(It.Is(i => i.Id == DeviceId), It.IsAny>())) + .ReturnsAsync(Try.Failure(new TimeoutException())); + + var connectionManager = new ConnectionManager( + cloudConnectionProvider.Object, + Mock.Of(), + GetIdentityProvider(), + Mock.Of(), + maxClients: 101, + closeCloudConnectionOnDeviceDisconnect: true, + retryInterval, + () => monotonicTime); + + 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); + + monotonicTime += retryInterval; + 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); + TimeSpan monotonicTime = TimeSpan.Zero; + 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(), + maxClients: 101, + closeCloudConnectionOnDeviceDisconnect: true, + retryInterval, + () => monotonicTime); + + Option cloudProxy = await connectionManager.GetCloudConnection(DeviceId); + Assert.True(cloudProxy.HasValue); + initialConnectionActive = false; + + IMessage message = Mock.Of(); + Task[] operations = Enumerable.Range(0, OperationCount) + .Select( + _ => Assert.ThrowsAsync( + () => cloudProxy.OrDefault().SendMessageAsync(message))) + .ToArray(); + await Task.WhenAll(operations); + + cloudConnectionProvider.Verify( + c => c.Connect(It.IsAny(), It.IsAny>()), + Times.Exactly(2)); + + monotonicTime += retryInterval; + 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); + TimeSpan monotonicTime = TimeSpan.Zero; + 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(), + maxClients: 101, + closeCloudConnectionOnDeviceDisconnect: true, + retryInterval, + () => monotonicTime); + + 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 76c23e44eb0..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 @@ -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; @@ -93,7 +94,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 +205,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 +300,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 +362,24 @@ async Task GetCloudProxy(IConnectionManager cm) Assert.Equal(expectedMessageIds, receivedMessageIds); } + static ConnectionManager CreateConnectionManagerWithoutCloudConnectionRetryDelay( + ICloudConnectionProvider connectionProvider, + ICredentialsCache credentialsCache, + IIdentityProvider identityProvider, + IDeviceConnectivityManager deviceConnectivityManager) + { + var monotonicClock = Stopwatch.StartNew(); + return new ConnectionManager( + connectionProvider, + credentialsCache, + identityProvider, + deviceConnectivityManager, + maxClients: 101, + closeCloudConnectionOnDeviceDisconnect: true, + TimeSpan.Zero, + () => monotonicClock.Elapsed); + } + static async Task RunSendMessages(ICloudProxy cloudProxy, IEnumerable messages, int batchSize = 1) { if (batchSize == 1)