From 966d010e8a1778bc5b768e54c4ea730084101da8 Mon Sep 17 00:00:00 2001 From: Vadim Kovalyov Date: Thu, 27 Aug 2026 23:53:37 +0000 Subject: [PATCH 1/3] Fix edgeHub subscription recovery Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../SubscriptionProcessor.cs | 254 +++++++++++------ .../SubscriptionProcessorTest.cs | 267 ++++++++++++++++++ 2 files changed, 429 insertions(+), 92 deletions(-) diff --git a/edge-hub/core/src/Microsoft.Azure.Devices.Edge.Hub.Core/SubscriptionProcessor.cs b/edge-hub/core/src/Microsoft.Azure.Devices.Edge.Hub.Core/SubscriptionProcessor.cs index 0ba5e1f70e7..85a18173970 100644 --- a/edge-hub/core/src/Microsoft.Azure.Devices.Edge.Hub.Core/SubscriptionProcessor.cs +++ b/edge-hub/core/src/Microsoft.Azure.Devices.Edge.Hub.Core/SubscriptionProcessor.cs @@ -4,9 +4,8 @@ namespace Microsoft.Azure.Devices.Edge.Hub.Core using System; using System.Collections.Concurrent; using System.Collections.Generic; - using System.Linq; + using System.Threading; using System.Threading.Tasks; - using System.Threading.Tasks.Dataflow; using Microsoft.Azure.Devices.Client.Exceptions; using Microsoft.Azure.Devices.Edge.Hub.Core.Cloud; using Microsoft.Azure.Devices.Edge.Hub.Core.Device; @@ -27,33 +26,59 @@ namespace Microsoft.Azure.Devices.Edge.Hub.Core /// Note that subscriptions are not stored in the SubscriptionProcessor - they are stored /// in the ConnectionManager. /// - public class SubscriptionProcessor : SubscriptionProcessorBase + public class SubscriptionProcessor : SubscriptionProcessorBase, IDisposable { + static readonly TimeSpan RecoveryMaxBackoff = TimeSpan.FromSeconds(30); static readonly ITransientErrorDetectionStrategy TransientErrorDetectionStrategy = new ErrorDetectionStrategy(); static readonly RetryStrategy TransientRetryStrategy = new ExponentialBackoff(2, TimeSpan.FromSeconds(2), TimeSpan.FromSeconds(10), TimeSpan.FromSeconds(4)); + static readonly ShouldRetry RecoveryShouldRetry = + new ExponentialBackoff(int.MaxValue, TimeSpan.FromSeconds(1), RecoveryMaxBackoff, TimeSpan.FromSeconds(1)).GetShouldRetry(); + readonly ConcurrentDictionary> pendingSubscriptions; - readonly ConcurrentDictionary subscriptionsBeingProcessed = new ConcurrentDictionary(); - readonly ActionBlock processSubscriptionsBlock; + readonly ConcurrentDictionary clientStates = new ConcurrentDictionary(); readonly IInvokeMethodHandler invokeMethodHandler; + readonly IDeviceConnectivityManager deviceConnectivityManager; + readonly CancellationTokenSource shutdown = new CancellationTokenSource(); + readonly CancellationToken shutdownToken; + int disposed; public SubscriptionProcessor(IConnectionManager connectionManager, IInvokeMethodHandler invokeMethodHandler, IDeviceConnectivityManager deviceConnectivityManager) : base(connectionManager) { - Preconditions.CheckNotNull(deviceConnectivityManager, nameof(deviceConnectivityManager)); + this.deviceConnectivityManager = Preconditions.CheckNotNull(deviceConnectivityManager, nameof(deviceConnectivityManager)); this.invokeMethodHandler = Preconditions.CheckNotNull(invokeMethodHandler, nameof(invokeMethodHandler)); this.pendingSubscriptions = new ConcurrentDictionary>(); - this.processSubscriptionsBlock = new ActionBlock(this.ProcessPendingSubscriptions); + this.shutdownToken = this.shutdown.Token; connectionManager.DeviceConnected += this.ClientConnectionToEdgeHubEstablished; - deviceConnectivityManager.DeviceConnected += this.CloudConnectivityEstablished; + this.deviceConnectivityManager.DeviceConnected += this.CloudConnectivityEstablished; connectionManager.CloudConnectionEstablished += this.ClientConnectionToCloudEstablished; } + public void Dispose() + { + if (Interlocked.Exchange(ref this.disposed, 1) == 0) + { + this.ConnectionManager.DeviceConnected -= this.ClientConnectionToEdgeHubEstablished; + this.deviceConnectivityManager.DeviceConnected -= this.CloudConnectivityEstablished; + this.ConnectionManager.CloudConnectionEstablished -= this.ClientConnectionToCloudEstablished; + this.shutdown.Cancel(); + this.shutdown.Dispose(); + } + } + protected override void HandleSubscriptions(string id, List<(DeviceSubscription, bool)> subscriptions) => this.AddToPendingSubscriptions(id, subscriptions); + static TimeSpan GetRecoveryDelay(int retryAttempt) + { + return RecoveryShouldRetry(retryAttempt - 1, null, out TimeSpan delay) + ? delay + : RecoveryMaxBackoff; + } + static Task ExecuteWithRetry(Func func, Action onRetry) { var transientRetryPolicy = new RetryPolicy(TransientErrorDetectionStrategy, TransientRetryStrategy); @@ -61,7 +86,7 @@ static Task ExecuteWithRetry(Func func, Action onRetry) return transientRetryPolicy.ExecuteAsync(func); } - async Task ProcessSubscriptionWithRetry(string id, Option cloudProxy, DeviceSubscription deviceSubscription, bool addSubscription) + async Task ProcessSubscriptionWithRetry(string id, ICloudProxy cloudProxy, DeviceSubscription deviceSubscription, bool addSubscription) { Events.ProcessingSubscription(id, deviceSubscription); try @@ -73,42 +98,44 @@ await ExecuteWithRetry( Metrics.AddRetryOperation(id, addSubscription ? "AddSubscription" : "RemoveSubscription"); Events.ErrorProcessingSubscription(id, deviceSubscription, addSubscription, r); }); + return true; } catch (Exception ex) { Events.ErrorProcessingSubscription(id, deviceSubscription, addSubscription, ex); + return false; } } - async Task ProcessSubscription(string id, Option cloudProxy, DeviceSubscription deviceSubscription, bool addSubscription) + async Task ProcessSubscription(string id, ICloudProxy cloudProxy, DeviceSubscription deviceSubscription, bool addSubscription) { switch (deviceSubscription) { case DeviceSubscription.C2D: if (addSubscription) { - cloudProxy.ForEach(c => c.StartListening()); + await cloudProxy.StartListening(); } else { - cloudProxy.ForEach(c => c.StopListening()); + await cloudProxy.StopListening(); } break; case DeviceSubscription.DesiredPropertyUpdates: - await cloudProxy.ForEachAsync(c => addSubscription ? c.SetupDesiredPropertyUpdatesAsync() : c.RemoveDesiredPropertyUpdatesAsync()); + await (addSubscription ? cloudProxy.SetupDesiredPropertyUpdatesAsync() : cloudProxy.RemoveDesiredPropertyUpdatesAsync()); break; case DeviceSubscription.Methods: if (addSubscription) { - await cloudProxy.ForEachAsync(c => c.SetupCallMethodAsync()); + await cloudProxy.SetupCallMethodAsync(); await this.invokeMethodHandler.ProcessInvokeMethodSubscription(id); } else { - await cloudProxy.ForEachAsync(c => c.RemoveCallMethodAsync()); + await cloudProxy.RemoveCallMethodAsync(); } break; @@ -118,7 +145,7 @@ async Task ProcessSubscription(string id, Option cloudProxy, Device // and because of that the rest of the CloudProxy implementations were built that way later. if (!addSubscription) { - await cloudProxy.ForEachAsync(c => c.RemoveTwinResponseAsync()); + await cloudProxy.RemoveTwinResponseAsync(); } break; @@ -130,119 +157,168 @@ async Task ProcessSubscription(string id, Option cloudProxy, Device } } - async void CloudConnectivityEstablished(object sender, EventArgs eventArgs) + void CloudConnectivityEstablished(object sender, EventArgs eventArgs) { Events.DeviceConnectedProcessingSubscriptions(); - async Task ProcessSubscriptionByIdentity(IIdentity identity) - { - try - { - Events.ProcessingSubscriptionsOnDeviceConnectedToCloud(identity); - await this.ProcessExistingSubscriptions(identity.Id); - } - catch (Exception e) - { - Events.ErrorProcessingSubscriptions(e, identity); - } - } - - try - { - IEnumerable tasks = this.ConnectionManager.GetConnectedClients().Select(id => ProcessSubscriptionByIdentity(id)); - await Task.WhenAll(tasks); - } - catch (Exception e) + foreach (IIdentity identity in this.ConnectionManager.GetConnectedClients()) { - Events.ErrorProcessingSubscriptions(e); + Events.ProcessingSubscriptionsOnDeviceConnectedToCloud(identity); + this.Signal(identity.Id); } } - async void ClientConnectionToCloudEstablished(object sender, IIdentity identity) + void ClientConnectionToCloudEstablished(object sender, IIdentity identity) { Events.ClientConnectedToCloudProcessingSubscriptions(identity); - try - { - await this.ProcessExistingSubscriptions(identity.Id); - } - catch (Exception e) - { - Events.ErrorProcessingSubscriptions(e, identity); - } + this.Signal(identity.Id); } - async void ClientConnectionToEdgeHubEstablished(object sender, IIdentity identity) + void ClientConnectionToEdgeHubEstablished(object sender, IIdentity identity) { Events.ClientConnectedToEdgeHubProcessingSubscriptions(identity); - try + this.Signal(identity.Id); + } + + void Signal(string id) + { + if (Volatile.Read(ref this.disposed) != 0) { - await this.ProcessExistingSubscriptions(identity.Id); + return; } - catch (Exception e) + + ClientState state = this.clientStates.GetOrAdd(id, _ => new ClientState()); + if (state.Signal()) { - Events.ErrorProcessingSubscriptions(e, identity); + _ = this.ProcessSubscriptionsAsync(id, state); } } - async Task ProcessExistingSubscriptions(string id) + async Task ProcessSubscriptionsAsync(string id, ClientState state) { - // Set a flag for an identity that temporarily disables subscription processing for other threads while we process - // the identity's subscriptions, so we don't trigger processing subscriptions multiple times for the same identity concurrently. - if (!this.subscriptionsBeingProcessed.TryAdd(id, true)) - { - // Identity subscription already being processed. Skip it. - Events.SkippingProcessingSubscription(id); - } - else + try { - try + while (true) { - Option cloudProxy = await this.ConnectionManager.GetCloudConnection(id); - Option> subscriptions = this.ConnectionManager.GetSubscriptions(id); - await subscriptions.ForEachAsync( - async s => + Events.ProcessingSubscriptions(id); + state.StartPass(); + bool retry; + try + { + Option cloudProxy = await this.ConnectionManager.GetCloudConnection(id); + if (cloudProxy.HasValue) { - foreach (KeyValuePair subscription in s) - { - await this.ProcessSubscriptionWithRetry(id, cloudProxy, subscription.Key, subscription.Value); - } - }); + retry = !await this.ApplySubscriptions(id, cloudProxy.OrDefault()); + } + else + { + Events.ProcessingSubscriptionsNoCloudProxy(id); + retry = true; + } + } + catch (Exception ex) when (!ex.IsFatal()) + { + Events.ErrorProcessingSubscriptions(ex); + retry = true; + } + + if (!retry) + { + state.ResetRetryAttempt(); + if (state.TryComplete()) + { + return; + } + + continue; + } + + if (!this.ConnectionManager.GetDeviceConnection(id).HasValue) + { + if (state.TryComplete()) + { + return; + } + + continue; + } + + await Task.Delay(GetRecoveryDelay(state.IncrementRetryAttempt()), this.shutdownToken); } - finally + } + catch (OperationCanceledException) when (this.shutdownToken.IsCancellationRequested) + { + state.Abort(); + } + catch (Exception ex) when (!ex.IsFatal()) + { + Events.ErrorProcessingSubscriptions(ex); + if (state.Abort()) { - this.subscriptionsBeingProcessed.TryRemove(id, out _); + this.Signal(id); } } } - async Task ProcessPendingSubscriptions(string id) + async Task ApplySubscriptions(string id, ICloudProxy cloudProxy) { - Events.ProcessingSubscriptions(id); + var processedSubscriptions = new Dictionary(); + bool succeeded = true; ConcurrentQueue<(DeviceSubscription, bool)> clientSubscriptionsQueue = this.GetClientSubscriptionsQueue(id); - if (!clientSubscriptionsQueue.IsEmpty) + while (clientSubscriptionsQueue.TryPeek(out (DeviceSubscription deviceSubscription, bool addSubscription) result)) { - Option cloudProxy = await this.ConnectionManager.GetCloudConnection(id); - if (!cloudProxy.HasValue) - { - Events.ProcessingSubscriptionsNoCloudProxy(id); - } + bool operationSucceeded = await this.ProcessSubscriptionWithRetry(id, cloudProxy, result.deviceSubscription, result.addSubscription); + processedSubscriptions[result.deviceSubscription] = operationSucceeded; + succeeded &= operationSucceeded; - while (clientSubscriptionsQueue.TryDequeue(out (DeviceSubscription deviceSubscription, bool addSubscription) result)) - { - await this.ProcessSubscriptionWithRetry(id, cloudProxy, result.deviceSubscription, result.addSubscription); - } + clientSubscriptionsQueue.TryDequeue(out _); } + + Option> subscriptions = this.ConnectionManager.GetSubscriptions(id); + await subscriptions.ForEachAsync( + async s => + { + foreach (KeyValuePair subscription in s) + { + if (!processedSubscriptions.TryGetValue(subscription.Key, out bool operationSucceeded) || !operationSucceeded) + { + succeeded &= await this.ProcessSubscriptionWithRetry(id, cloudProxy, subscription.Key, subscription.Value); + } + } + }); + return succeeded; } void AddToPendingSubscriptions(string id, List<(DeviceSubscription, bool)> subscriptions) { ConcurrentQueue<(DeviceSubscription, bool)> clientSubscriptionsQueue = this.GetClientSubscriptionsQueue(id); subscriptions.ForEach(s => clientSubscriptionsQueue.Enqueue(s)); - this.processSubscriptionsBlock.Post(id); + this.Signal(id); } ConcurrentQueue<(DeviceSubscription, bool)> GetClientSubscriptionsQueue(string id) => this.pendingSubscriptions.GetOrAdd(id, new ConcurrentQueue<(DeviceSubscription, bool)>()); + sealed class ClientState + { + const int Idle = 0; + const int Running = 1; + const int Pending = 2; + int status; + int retryAttempt; + + public bool Signal() => Interlocked.Exchange(ref this.status, Pending) == Idle; + + public void StartPass() => Interlocked.CompareExchange(ref this.status, Running, Pending); + + public bool TryComplete() => Interlocked.CompareExchange(ref this.status, Idle, Running) == Running; + + public int IncrementRetryAttempt() => ++this.retryAttempt; + + public void ResetRetryAttempt() => this.retryAttempt = 0; + + public bool Abort() => Interlocked.Exchange(ref this.status, Idle) == Pending; + } + class ErrorDetectionStrategy : ITransientErrorDetectionStrategy { static readonly ISet NonTransientExceptions = new HashSet @@ -268,8 +344,7 @@ enum EventIds ProcessingSubscription, DeviceConnectedToEdgeHubProcessingSubscription, ClientConnectedProcessingSubscriptions, - ProcessingSubscriptionsNoCloudProxy, - SkippingProcessingSubscription + ProcessingSubscriptionsNoCloudProxy } public static void ErrorProcessingSubscriptions(Exception ex, IIdentity identity) @@ -342,11 +417,6 @@ public static void ClientConnectedToEdgeHubProcessingSubscriptions(IIdentity ide Log.LogInformation((int)EventIds.DeviceConnectedToEdgeHubProcessingSubscription, Invariant($"Client {identity.Id} connected to edgeHub, processing existing subscriptions.")); } - public static void SkippingProcessingSubscription(string id) - { - Log.LogInformation((int)EventIds.SkippingProcessingSubscription, Invariant($"Skipping {id} for subscription processing, as it is currently being processed.")); - } - public static void ProcessingSubscriptions(string id) { Log.LogInformation((int)EventIds.ProcessingSubscription, Invariant($"Processing pending subscriptions for {id}")); diff --git a/edge-hub/core/test/Microsoft.Azure.Devices.Edge.Hub.Core.Test/SubscriptionProcessorTest.cs b/edge-hub/core/test/Microsoft.Azure.Devices.Edge.Hub.Core.Test/SubscriptionProcessorTest.cs index b9a062b4548..f063f49949b 100644 --- a/edge-hub/core/test/Microsoft.Azure.Devices.Edge.Hub.Core.Test/SubscriptionProcessorTest.cs +++ b/edge-hub/core/test/Microsoft.Azure.Devices.Edge.Hub.Core.Test/SubscriptionProcessorTest.cs @@ -570,5 +570,272 @@ public void ProcessSubscriptionsOnClientCloudConnectionEstablished() Mock.Get(invokeMethodHandler).VerifyAll(); Mock.Get(connectionManager).VerifyAll(); } + + [Fact] + public async Task PendingMethodsSubscriptionSurvivesMissingCloudProxy() + { + string id = "d1/m1"; + bool proxyAvailable = false; + var firstAttempt = new SemaphoreSlim(0); + var subscriptionApplied = new SemaphoreSlim(0, 1); + var cloudProxy = new Mock(MockBehavior.Strict); + cloudProxy.Setup(c => c.SetupCallMethodAsync()) + .Callback(() => subscriptionApplied.Release()) + .Returns(Task.CompletedTask); + var connectionManager = new Mock(); + connectionManager.Setup(c => c.AddSubscription(id, DeviceSubscription.Methods)).Returns(true); + connectionManager.Setup(c => c.GetDeviceConnection(id)).Returns(Option.Some(Mock.Of())); + connectionManager.Setup(c => c.GetCloudConnection(id)) + .Returns( + () => + { + if (!proxyAvailable) + { + firstAttempt.Release(); + return Task.FromResult(Option.None()); + } + + return Task.FromResult(Option.Some(cloudProxy.Object)); + }); + var invokeMethodHandler = Mock.Of( + h => h.ProcessInvokeMethodSubscription(id) == Task.CompletedTask); + using var subscriptionProcessor = new SubscriptionProcessor( + connectionManager.Object, + invokeMethodHandler, + Mock.Of()); + + await subscriptionProcessor.AddSubscription(id, DeviceSubscription.Methods); + Assert.True(await firstAttempt.WaitAsync(TimeSpan.FromSeconds(5))); + cloudProxy.Verify(c => c.SetupCallMethodAsync(), Times.Never); + + proxyAvailable = true; + Assert.True(await subscriptionApplied.WaitAsync(TimeSpan.FromSeconds(5))); + + cloudProxy.Verify(c => c.SetupCallMethodAsync(), Times.Once); + Mock.Get(invokeMethodHandler).Verify( + h => h.ProcessInvokeMethodSubscription(id), + Times.Once); + } + + [Fact] + public async Task PendingMethodsSubscriptionRecoversAfterCloudProxyCreationThrows() + { + string id = "d1/m1"; + var subscriptionApplied = new SemaphoreSlim(0, 1); + var cloudProxy = new Mock(MockBehavior.Strict); + cloudProxy.Setup(c => c.SetupCallMethodAsync()) + .Callback(() => subscriptionApplied.Release()) + .Returns(Task.CompletedTask); + var connectionManager = new Mock(); + connectionManager.Setup(c => c.AddSubscription(id, DeviceSubscription.Methods)).Returns(true); + connectionManager.Setup(c => c.GetDeviceConnection(id)).Returns(Option.Some(Mock.Of())); + connectionManager.SetupSequence(c => c.GetCloudConnection(id)) + .ThrowsAsync(new TimeoutException()) + .ReturnsAsync(Option.Some(cloudProxy.Object)); + using var subscriptionProcessor = new SubscriptionProcessor( + connectionManager.Object, + Mock.Of( + h => h.ProcessInvokeMethodSubscription(id) == Task.CompletedTask), + Mock.Of()); + + await subscriptionProcessor.AddSubscription(id, DeviceSubscription.Methods); + + Assert.True(await subscriptionApplied.WaitAsync(TimeSpan.FromSeconds(5))); + cloudProxy.Verify(c => c.SetupCallMethodAsync(), Times.Once); + connectionManager.Verify(c => c.GetCloudConnection(id), Times.Exactly(2)); + } + + [Fact] + public async Task TransientSubscriptionFailureStillRetries() + { + string id = "d1/m1"; + int setupAttempts = 0; + var subscriptionApplied = new SemaphoreSlim(0, 1); + var cloudProxy = new Mock(MockBehavior.Strict); + cloudProxy.Setup(c => c.SetupCallMethodAsync()) + .Returns( + () => + { + if (Interlocked.Increment(ref setupAttempts) == 1) + { + return Task.FromException(new InvalidOperationException()); + } + + subscriptionApplied.Release(); + return Task.CompletedTask; + }); + var connectionManager = new Mock(); + connectionManager.Setup(c => c.AddSubscription(id, DeviceSubscription.Methods)).Returns(true); + connectionManager.Setup(c => c.GetCloudConnection(id)).ReturnsAsync(Option.Some(cloudProxy.Object)); + using var subscriptionProcessor = new SubscriptionProcessor( + connectionManager.Object, + Mock.Of( + h => h.ProcessInvokeMethodSubscription(id) == Task.CompletedTask), + Mock.Of()); + + await subscriptionProcessor.AddSubscription(id, DeviceSubscription.Methods); + + Assert.True(await subscriptionApplied.WaitAsync(TimeSpan.FromSeconds(5))); + cloudProxy.Verify(c => c.SetupCallMethodAsync(), Times.Exactly(2)); + } + + [Fact] + public async Task PendingSubscriptionsPreserveAddRemoveOrderWhileCloudProxyIsMissing() + { + string id = "d1/m1"; + bool proxyAvailable = false; + var firstAttempt = new SemaphoreSlim(0); + var subscriptionsApplied = new SemaphoreSlim(0, 1); + var cloudProxy = new Mock(MockBehavior.Strict); + var sequence = new MockSequence(); + cloudProxy.InSequence(sequence) + .Setup(c => c.SetupCallMethodAsync()) + .Returns(Task.CompletedTask); + cloudProxy.InSequence(sequence) + .Setup(c => c.RemoveCallMethodAsync()) + .Callback(() => subscriptionsApplied.Release()) + .Returns(Task.CompletedTask); + var connectionManager = new Mock(); + connectionManager.Setup(c => c.AddSubscription(id, DeviceSubscription.Methods)).Returns(true); + connectionManager.Setup(c => c.RemoveSubscription(id, DeviceSubscription.Methods)).Returns(true); + connectionManager.Setup(c => c.GetDeviceConnection(id)).Returns(Option.Some(Mock.Of())); + connectionManager.Setup(c => c.GetCloudConnection(id)) + .Returns( + () => + { + if (!proxyAvailable) + { + firstAttempt.Release(); + return Task.FromResult(Option.None()); + } + + return Task.FromResult(Option.Some(cloudProxy.Object)); + }); + using var subscriptionProcessor = new SubscriptionProcessor( + connectionManager.Object, + Mock.Of( + h => h.ProcessInvokeMethodSubscription(id) == Task.CompletedTask), + Mock.Of()); + + await subscriptionProcessor.AddSubscription(id, DeviceSubscription.Methods); + await subscriptionProcessor.RemoveSubscription(id, DeviceSubscription.Methods); + Assert.True(await firstAttempt.WaitAsync(TimeSpan.FromSeconds(5))); + + proxyAvailable = true; + Assert.True(await subscriptionsApplied.WaitAsync(TimeSpan.FromSeconds(5))); + + cloudProxy.VerifyAll(); + } + + [Fact] + public async Task MissingCloudProxyForOneClientDoesNotDelayAnotherClient() + { + string unavailableId = "d1/m1"; + string availableId = "d1/m2"; + var unavailableProxyResult = new TaskCompletionSource>(TaskCreationOptions.RunContinuationsAsynchronously); + var subscriptionApplied = new SemaphoreSlim(0, 1); + var cloudProxy = new Mock(MockBehavior.Strict); + cloudProxy.Setup(c => c.SetupDesiredPropertyUpdatesAsync()) + .Callback(() => subscriptionApplied.Release()) + .Returns(Task.CompletedTask); + var connectionManager = new Mock(); + connectionManager.Setup(c => c.AddSubscription(It.IsAny(), It.IsAny())).Returns(true); + connectionManager.Setup(c => c.GetCloudConnection(unavailableId)) + .Returns(unavailableProxyResult.Task); + connectionManager.Setup(c => c.GetCloudConnection(availableId)) + .ReturnsAsync(Option.Some(cloudProxy.Object)); + using var subscriptionProcessor = new SubscriptionProcessor( + connectionManager.Object, + Mock.Of(), + Mock.Of()); + + await subscriptionProcessor.AddSubscription(unavailableId, DeviceSubscription.Methods); + await subscriptionProcessor.AddSubscription(availableId, DeviceSubscription.DesiredPropertyUpdates); + + Assert.True(await subscriptionApplied.WaitAsync(TimeSpan.FromSeconds(5))); + unavailableProxyResult.SetResult(Option.None()); + cloudProxy.VerifyAll(); + } + + [Fact] + public async Task RepeatedRecoveryEventsCoalesceIntoOneAdditionalPass() + { + string id = "d1/m1"; + var identity = Mock.Of(i => i.Id == id); + var firstProxyResult = new TaskCompletionSource>(TaskCreationOptions.RunContinuationsAsynchronously); + var secondPassCompleted = new SemaphoreSlim(0, 1); + int subscriptionCallCount = 0; + IReadOnlyDictionary subscriptions = + new Dictionary { [DeviceSubscription.Methods] = true }; + var cloudProxy = new Mock(); + cloudProxy.Setup(c => c.SetupCallMethodAsync()) + .Callback( + () => + { + if (Interlocked.Increment(ref subscriptionCallCount) == 2) + { + secondPassCompleted.Release(); + } + }) + .Returns(Task.CompletedTask); + var connectionManager = new Mock(); + connectionManager.Setup(c => c.GetCloudConnection(id)) + .Returns(firstProxyResult.Task); + connectionManager.Setup(c => c.GetSubscriptions(id)).Returns(Option.Some(subscriptions)); + using var subscriptionProcessor = new SubscriptionProcessor( + connectionManager.Object, + Mock.Of( + h => h.ProcessInvokeMethodSubscription(id) == Task.CompletedTask), + Mock.Of()); + + Mock.Get(connectionManager.Object).Raise(c => c.DeviceConnected += null, this, identity); + Mock.Get(connectionManager.Object).Raise(c => c.CloudConnectionEstablished += null, this, identity); + Mock.Get(connectionManager.Object).Raise(c => c.CloudConnectionEstablished += null, this, identity); + firstProxyResult.SetResult(Option.Some(cloudProxy.Object)); + + Assert.True(await secondPassCompleted.WaitAsync(TimeSpan.FromSeconds(5))); + Assert.Equal(2, Volatile.Read(ref subscriptionCallCount)); + connectionManager.Verify(c => c.GetCloudConnection(id), Times.Exactly(2)); + } + + [Fact] + public async Task ReplayRequestedDuringActiveReplayRunsAnotherPass() + { + string id = "d1/m1"; + var identity = Mock.Of(i => i.Id == id); + var firstReplayStarted = new SemaphoreSlim(0, 1); + var releaseFirstReplay = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + var secondReplayCompleted = new SemaphoreSlim(0, 1); + IReadOnlyDictionary subscriptions = + new Dictionary { [DeviceSubscription.Methods] = true }; + var cloudProxy = new Mock(MockBehavior.Strict); + cloudProxy.Setup(c => c.SetupCallMethodAsync()) + .Callback(() => firstReplayStarted.Release()) + .Returns(() => releaseFirstReplay.Task); + cloudProxy.Setup(c => c.RemoveCallMethodAsync()) + .Callback(() => secondReplayCompleted.Release()) + .Returns(Task.CompletedTask); + var connectionManager = new Mock(); + connectionManager.Setup(c => c.GetCloudConnection(id)) + .ReturnsAsync(Option.Some(cloudProxy.Object)); + connectionManager.Setup(c => c.GetSubscriptions(id)) + .Returns(() => Option.Some(subscriptions)); + using var subscriptionProcessor = new SubscriptionProcessor( + connectionManager.Object, + Mock.Of( + h => h.ProcessInvokeMethodSubscription(id) == Task.CompletedTask), + Mock.Of()); + + Mock.Get(connectionManager.Object).Raise(c => c.DeviceConnected += null, this, identity); + Assert.True(await firstReplayStarted.WaitAsync(TimeSpan.FromSeconds(5))); + + subscriptions = new Dictionary { [DeviceSubscription.Methods] = false }; + Mock.Get(connectionManager.Object).Raise(c => c.CloudConnectionEstablished += null, this, identity); + releaseFirstReplay.SetResult(true); + + Assert.True(await secondReplayCompleted.WaitAsync(TimeSpan.FromSeconds(5))); + cloudProxy.Verify(c => c.SetupCallMethodAsync(), Times.Once); + cloudProxy.Verify(c => c.RemoveCallMethodAsync(), Times.Once); + } } } From 23b4f5b34be179f3631b9a7bf20f305706324b78 Mon Sep 17 00:00:00 2001 From: Vadim Kovalyov Date: Fri, 28 Aug 2026 00:16:21 +0000 Subject: [PATCH 2/3] fix(edge-hub): refine subscription recovery Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../SubscriptionProcessor.cs | 27 ++++++---- .../routing/RoutingEdgeHub.cs | 1 + .../Program.cs | 1 + .../SubscriptionProcessorTest.cs | 54 +++++++++++++++++++ .../routing/RoutingEdgeHubTest.cs | 19 ++++++- 5 files changed, 91 insertions(+), 11 deletions(-) diff --git a/edge-hub/core/src/Microsoft.Azure.Devices.Edge.Hub.Core/SubscriptionProcessor.cs b/edge-hub/core/src/Microsoft.Azure.Devices.Edge.Hub.Core/SubscriptionProcessor.cs index 85a18173970..7ae644acdbc 100644 --- a/edge-hub/core/src/Microsoft.Azure.Devices.Edge.Hub.Core/SubscriptionProcessor.cs +++ b/edge-hub/core/src/Microsoft.Azure.Devices.Edge.Hub.Core/SubscriptionProcessor.cs @@ -79,11 +79,11 @@ static TimeSpan GetRecoveryDelay(int retryAttempt) : RecoveryMaxBackoff; } - static Task ExecuteWithRetry(Func func, Action onRetry) + static Task ExecuteWithRetry(Func func, Action onRetry, CancellationToken cancellationToken) { var transientRetryPolicy = new RetryPolicy(TransientErrorDetectionStrategy, TransientRetryStrategy); transientRetryPolicy.Retrying += (_, args) => onRetry(args); - return transientRetryPolicy.ExecuteAsync(func); + return transientRetryPolicy.ExecuteAsync(func, cancellationToken); } async Task ProcessSubscriptionWithRetry(string id, ICloudProxy cloudProxy, DeviceSubscription deviceSubscription, bool addSubscription) @@ -97,9 +97,14 @@ await ExecuteWithRetry( { Metrics.AddRetryOperation(id, addSubscription ? "AddSubscription" : "RemoveSubscription"); Events.ErrorProcessingSubscription(id, deviceSubscription, addSubscription, r); - }); + }, + this.shutdownToken); return true; } + catch (OperationCanceledException) when (this.shutdownToken.IsCancellationRequested) + { + throw; + } catch (Exception ex) { Events.ErrorProcessingSubscription(id, deviceSubscription, addSubscription, ex); @@ -205,6 +210,7 @@ async Task ProcessSubscriptionsAsync(string id, ClientState state) try { Option cloudProxy = await this.ConnectionManager.GetCloudConnection(id); + this.shutdownToken.ThrowIfCancellationRequested(); if (cloudProxy.HasValue) { retry = !await this.ApplySubscriptions(id, cloudProxy.OrDefault()); @@ -215,9 +221,9 @@ async Task ProcessSubscriptionsAsync(string id, ClientState state) retry = true; } } - catch (Exception ex) when (!ex.IsFatal()) + catch (Exception ex) when (!ex.IsFatal() && !(ex is OperationCanceledException && this.shutdownToken.IsCancellationRequested)) { - Events.ErrorProcessingSubscriptions(ex); + Events.ErrorProcessingSubscriptions(ex, id); retry = true; } @@ -251,7 +257,7 @@ async Task ProcessSubscriptionsAsync(string id, ClientState state) } catch (Exception ex) when (!ex.IsFatal()) { - Events.ErrorProcessingSubscriptions(ex); + Events.ErrorProcessingSubscriptions(ex, id); if (state.Abort()) { this.Signal(id); @@ -348,14 +354,17 @@ enum EventIds } public static void ErrorProcessingSubscriptions(Exception ex, IIdentity identity) + => ErrorProcessingSubscriptions(ex, identity.Id); + + public static void ErrorProcessingSubscriptions(Exception ex, string id) { if (ex.HasTimeoutException()) { - Log.LogDebug((int)EventIds.ErrorProcessingSubscriptions, ex, Invariant($"Timed out while processing subscriptions for client {identity.Id}. Will try again when connected.")); + Log.LogDebug((int)EventIds.ErrorProcessingSubscriptions, ex, Invariant($"Timed out while processing subscriptions for client {id}. Will try again when connected.")); } else { - Log.LogWarning((int)EventIds.ErrorProcessingSubscriptions, ex, Invariant($"Error processing subscriptions for client {identity.Id}.")); + Log.LogWarning((int)EventIds.ErrorProcessingSubscriptions, ex, Invariant($"Error processing subscriptions for client {id}.")); } } @@ -404,7 +413,7 @@ internal static void DeviceConnectedProcessingSubscriptions() internal static void ErrorProcessingSubscriptions(Exception e) { - Log.LogWarning((int)EventIds.ProcessingSubscription, e, Invariant($"Error processing subscriptions for connected clients.")); + Log.LogWarning((int)EventIds.ErrorProcessingSubscriptions, e, Invariant($"Error processing subscriptions for connected clients.")); } public static void ClientConnectedToCloudProcessingSubscriptions(IIdentity identity) diff --git a/edge-hub/core/src/Microsoft.Azure.Devices.Edge.Hub.Core/routing/RoutingEdgeHub.cs b/edge-hub/core/src/Microsoft.Azure.Devices.Edge.Hub.Core/routing/RoutingEdgeHub.cs index a7396522e81..5d97dd1c579 100644 --- a/edge-hub/core/src/Microsoft.Azure.Devices.Edge.Hub.Core/routing/RoutingEdgeHub.cs +++ b/edge-hub/core/src/Microsoft.Azure.Devices.Edge.Hub.Core/routing/RoutingEdgeHub.cs @@ -162,6 +162,7 @@ protected virtual void Dispose(bool disposing) { if (disposing) { + (this.subscriptionProcessor as IDisposable)?.Dispose(); this.router?.Dispose(); } } diff --git a/edge-hub/core/src/Microsoft.Azure.Devices.Edge.Hub.Service/Program.cs b/edge-hub/core/src/Microsoft.Azure.Devices.Edge.Hub.Service/Program.cs index 76a31c9a8a4..c3f07e2e61b 100644 --- a/edge-hub/core/src/Microsoft.Azure.Devices.Edge.Hub.Service/Program.cs +++ b/edge-hub/core/src/Microsoft.Azure.Devices.Edge.Hub.Service/Program.cs @@ -189,6 +189,7 @@ static async Task MainAsync(IConfigurationRoot configuration, ILogger logge logger.LogError($"Error stopping protocol heads: {ex.Message}"); } + edgeHub.Dispose(); await CloseDbStoreProviderAsync(container); } diff --git a/edge-hub/core/test/Microsoft.Azure.Devices.Edge.Hub.Core.Test/SubscriptionProcessorTest.cs b/edge-hub/core/test/Microsoft.Azure.Devices.Edge.Hub.Core.Test/SubscriptionProcessorTest.cs index f063f49949b..0608cc0c208 100644 --- a/edge-hub/core/test/Microsoft.Azure.Devices.Edge.Hub.Core.Test/SubscriptionProcessorTest.cs +++ b/edge-hub/core/test/Microsoft.Azure.Devices.Edge.Hub.Core.Test/SubscriptionProcessorTest.cs @@ -837,5 +837,59 @@ public async Task ReplayRequestedDuringActiveReplayRunsAnotherPass() cloudProxy.Verify(c => c.SetupCallMethodAsync(), Times.Once); cloudProxy.Verify(c => c.RemoveCallMethodAsync(), Times.Once); } + + [Fact] + public async Task ClientInRecoveryDelayDoesNotBlockAnotherClient() + { + string unavailableId = "d1/m1"; + string availableId = "d1/m2"; + var recoveryStarted = new SemaphoreSlim(0, 1); + var subscriptionApplied = new SemaphoreSlim(0, 1); + var availableCloudProxy = new Mock(); + availableCloudProxy.Setup(c => c.SetupDesiredPropertyUpdatesAsync()) + .Callback(() => subscriptionApplied.Release()) + .Returns(Task.CompletedTask); + var connectionManager = new Mock(); + connectionManager.Setup(c => c.AddSubscription(It.IsAny(), It.IsAny())).Returns(true); + connectionManager.Setup(c => c.GetCloudConnection(unavailableId)) + .Callback(() => recoveryStarted.Release()) + .ReturnsAsync(Option.None()); + connectionManager.Setup(c => c.GetDeviceConnection(unavailableId)).Returns(Option.Some(Mock.Of())); + connectionManager.Setup(c => c.GetCloudConnection(availableId)).ReturnsAsync(Option.Some(availableCloudProxy.Object)); + using var subscriptionProcessor = new SubscriptionProcessor( + connectionManager.Object, + Mock.Of(), + Mock.Of()); + + await subscriptionProcessor.AddSubscription(unavailableId, DeviceSubscription.Methods); + Assert.True(await recoveryStarted.WaitAsync(TimeSpan.FromSeconds(5))); + await subscriptionProcessor.AddSubscription(availableId, DeviceSubscription.DesiredPropertyUpdates); + + Assert.True(await subscriptionApplied.WaitAsync(TimeSpan.FromSeconds(5))); + } + + [Fact] + public async Task DisposeCancelsRecoveryDelay() + { + string id = "d1/m1"; + var recoveryStarted = new SemaphoreSlim(0, 1); + var connectionManager = new Mock(); + connectionManager.Setup(c => c.AddSubscription(id, DeviceSubscription.Methods)).Returns(true); + connectionManager.Setup(c => c.GetCloudConnection(id)) + .Callback(() => recoveryStarted.Release()) + .ReturnsAsync(Option.None()); + connectionManager.Setup(c => c.GetDeviceConnection(id)).Returns(Option.Some(Mock.Of())); + var subscriptionProcessor = new SubscriptionProcessor( + connectionManager.Object, + Mock.Of(), + Mock.Of()); + + await subscriptionProcessor.AddSubscription(id, DeviceSubscription.Methods); + Assert.True(await recoveryStarted.WaitAsync(TimeSpan.FromSeconds(5))); + subscriptionProcessor.Dispose(); + await Task.Delay(TimeSpan.FromMilliseconds(1200)); + + connectionManager.Verify(c => c.GetCloudConnection(id), Times.Once); + } } } diff --git a/edge-hub/core/test/Microsoft.Azure.Devices.Edge.Hub.Core.Test/routing/RoutingEdgeHubTest.cs b/edge-hub/core/test/Microsoft.Azure.Devices.Edge.Hub.Core.Test/routing/RoutingEdgeHubTest.cs index 3cf6f51e302..f8fd7772416 100644 --- a/edge-hub/core/test/Microsoft.Azure.Devices.Edge.Hub.Core.Test/routing/RoutingEdgeHubTest.cs +++ b/edge-hub/core/test/Microsoft.Azure.Devices.Edge.Hub.Core.Test/routing/RoutingEdgeHubTest.cs @@ -534,10 +534,25 @@ public async Task AddEdgeSystemPropertiesTest() Assert.True(clientMessage3.SystemProperties.ContainsKey(SystemProperties.EdgeMessageId)); } - static async Task GetTestEdgeHub(IConnectionManager connectionManager = null) + [Fact] + public async Task DisposeDisposesSubscriptionProcessor() + { + var subscriptionProcessor = new Mock(); + Mock disposableSubscriptionProcessor = subscriptionProcessor.As(); + RoutingEdgeHub edgeHub = await GetTestEdgeHub(subscriptionProcessor: subscriptionProcessor.Object); + + edgeHub.Dispose(); + + disposableSubscriptionProcessor.Verify(d => d.Dispose(), Times.Once); + } + + static async Task GetTestEdgeHub( + IConnectionManager connectionManager = null, + ISubscriptionProcessor subscriptionProcessor = null) { // Arrange connectionManager = connectionManager ?? Mock.Of(); + subscriptionProcessor = subscriptionProcessor ?? Mock.Of(); var endpoint = new Mock("myId"); var endpointExecutor = Mock.Of(); Mock.Get(endpointExecutor).SetupGet(ee => ee.Endpoint).Returns(() => endpoint.Object); @@ -558,7 +573,7 @@ static async Task GetTestEdgeHub(IConnectionManager connectionMa "ed1", "$edgeHub", Mock.Of(), - Mock.Of(), + subscriptionProcessor, Mock.Of()); return edgeHub; } From 629aa1c9cfe92f6b0f8040f0d2e4ea87235b610e Mon Sep 17 00:00:00 2001 From: Vadim Kovalyov Date: Thu, 3 Sep 2026 21:03:39 +0000 Subject: [PATCH 3/3] fix(edge-hub): cancel pending proxy lookup Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../SubscriptionProcessor.cs | 4 +-- .../SubscriptionProcessorTest.cs | 27 +++++++++++++++++++ 2 files changed, 29 insertions(+), 2 deletions(-) diff --git a/edge-hub/core/src/Microsoft.Azure.Devices.Edge.Hub.Core/SubscriptionProcessor.cs b/edge-hub/core/src/Microsoft.Azure.Devices.Edge.Hub.Core/SubscriptionProcessor.cs index 7ae644acdbc..4a0380cff89 100644 --- a/edge-hub/core/src/Microsoft.Azure.Devices.Edge.Hub.Core/SubscriptionProcessor.cs +++ b/edge-hub/core/src/Microsoft.Azure.Devices.Edge.Hub.Core/SubscriptionProcessor.cs @@ -209,7 +209,7 @@ async Task ProcessSubscriptionsAsync(string id, ClientState state) bool retry; try { - Option cloudProxy = await this.ConnectionManager.GetCloudConnection(id); + Option cloudProxy = await this.ConnectionManager.GetCloudConnection(id).WaitAsync(this.shutdownToken); this.shutdownToken.ThrowIfCancellationRequested(); if (cloudProxy.HasValue) { @@ -360,7 +360,7 @@ public static void ErrorProcessingSubscriptions(Exception ex, string id) { if (ex.HasTimeoutException()) { - Log.LogDebug((int)EventIds.ErrorProcessingSubscriptions, ex, Invariant($"Timed out while processing subscriptions for client {id}. Will try again when connected.")); + Log.LogDebug((int)EventIds.ErrorProcessingSubscriptions, ex, Invariant($"Timed out while processing subscriptions for client {id}. Will retry subscription recovery.")); } else { diff --git a/edge-hub/core/test/Microsoft.Azure.Devices.Edge.Hub.Core.Test/SubscriptionProcessorTest.cs b/edge-hub/core/test/Microsoft.Azure.Devices.Edge.Hub.Core.Test/SubscriptionProcessorTest.cs index 0608cc0c208..8d35447fbdd 100644 --- a/edge-hub/core/test/Microsoft.Azure.Devices.Edge.Hub.Core.Test/SubscriptionProcessorTest.cs +++ b/edge-hub/core/test/Microsoft.Azure.Devices.Edge.Hub.Core.Test/SubscriptionProcessorTest.cs @@ -891,5 +891,32 @@ public async Task DisposeCancelsRecoveryDelay() connectionManager.Verify(c => c.GetCloudConnection(id), Times.Once); } + + [Fact] + public async Task DisposeCancelsPendingCloudConnectionLookup() + { + string id = "d1/m1"; + var cloudLookupStarted = new SemaphoreSlim(0, 1); + var cloudLookup = new TaskCompletionSource>(TaskCreationOptions.RunContinuationsAsynchronously); + var cloudProxy = new Mock(); + var connectionManager = new Mock(); + connectionManager.Setup(c => c.AddSubscription(id, DeviceSubscription.Methods)).Returns(true); + connectionManager.Setup(c => c.GetCloudConnection(id)) + .Callback(() => cloudLookupStarted.Release()) + .Returns(cloudLookup.Task); + var subscriptionProcessor = new SubscriptionProcessor( + connectionManager.Object, + Mock.Of(), + Mock.Of()); + + await subscriptionProcessor.AddSubscription(id, DeviceSubscription.Methods); + Assert.True(await cloudLookupStarted.WaitAsync(TimeSpan.FromSeconds(5))); + + subscriptionProcessor.Dispose(); + cloudLookup.SetResult(Option.Some(cloudProxy.Object)); + await Task.Delay(TimeSpan.FromMilliseconds(100)); + + cloudProxy.Verify(c => c.SetupCallMethodAsync(), Times.Never); + } } }