Skip to content

Commit 9fcc8cc

Browse files
DorothySun216SeanFeldman
DorothySun216
and
SeanFeldman
authored
Enable a way to Unregister Message Handler and Session Handler (#14021)
* add UnregisterMessageHandler method * Update sdk/servicebus/Microsoft.Azure.ServiceBus/src/Core/IReceiverClient.cs Co-authored-by: Sean Feldman <[email protected]> * Update sdk/servicebus/Microsoft.Azure.ServiceBus/src/Core/MessageReceiver.cs Co-authored-by: Sean Feldman <[email protected]> * Update the unregister method to be async and await for inflight operations to finish * Update sdk/servicebus/Microsoft.Azure.ServiceBus/src/SubscriptionClient.cs Co-authored-by: Sean Feldman <[email protected]> * Update sdk/servicebus/Microsoft.Azure.ServiceBus/src/Core/MessageReceiver.cs Co-authored-by: Sean Feldman <[email protected]> * Update sdk/servicebus/Microsoft.Azure.ServiceBus/src/QueueClient.cs Co-authored-by: Sean Feldman <[email protected]> * Change name to have async suffix and add to existing onMessageQueueTests * Add UnregisterSessionHandlerAsync and corresponding tests * nit * nit * Add a new cancellation type to not cancel inflight message handling operations when unregister is called. * Add another type of cancellation token to session handler path * nit * Add a timeout parameter to unregister functions and add according unit tests * nit * cancel runningTaskCancellationTokenSource after unregister is done * change public API * update the API header * update the API definition * fix spacing * fix ApproveAzureServiceBus CIT test Co-authored-by: Sean Feldman <[email protected]>
1 parent b15911e commit 9fcc8cc

17 files changed

+634
-16
lines changed

‎sdk/servicebus/Microsoft.Azure.ServiceBus/src/Core/IReceiverClient.cs‎

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -63,6 +63,14 @@ public interface IReceiverClient : IClientEntity
6363
/// <remarks>Enable prefetch to speed up the receive rate.</remarks>
6464
void RegisterMessageHandler(Func<Message, CancellationToken, Task> handler, MessageHandlerOptions messageHandlerOptions);
6565

66+
/// <summary>
67+
/// Unregister message handler from the receiver if there is an active message handler registered. This operation waits for the completion
68+
/// of inflight receive and message handling operations to finish and unregisters future receives on the message handler which previously
69+
/// registered.
70+
/// <param name="inflightMessageHandlerTasksWaitTimeout"> is the waitTimeout for inflight message handling tasks.
71+
/// </summary>
72+
Task UnregisterMessageHandlerAsync(TimeSpan inflightMessageHandlerTasksWaitTimeout);
73+
6674
/// <summary>
6775
/// Completes a <see cref="Message"/> using its lock token. This will delete the message from the queue.
6876
/// </summary>
@@ -115,4 +123,4 @@ public interface IReceiverClient : IClientEntity
115123
/// </remarks>
116124
Task DeadLetterAsync(string lockToken, string deadLetterReason, string deadLetterErrorDescription = null);
117125
}
118-
}
126+
}

‎sdk/servicebus/Microsoft.Azure.ServiceBus/src/Core/MessageReceiver.cs‎

Lines changed: 61 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -54,7 +54,11 @@ public class MessageReceiver : ClientEntity, IMessageReceiver
5454
int prefetchCount;
5555
long lastPeekedSequenceNumber;
5656
MessageReceivePump receivePump;
57+
// Cancellation token to cancel the message pump. Once this is fired, all future message handling operations registered by application will be
58+
// cancelled.
5759
CancellationTokenSource receivePumpCancellationTokenSource;
60+
// Cancellation token to cancel the inflight message handling operations registered by application in the message pump.
61+
CancellationTokenSource runningTaskCancellationTokenSource;
5862

5963
/// <summary>
6064
/// Creates a new MessageReceiver from a <see cref="ServiceBusConnectionStringBuilder"/>.
@@ -899,6 +903,51 @@ public void RegisterMessageHandler(Func<Message, CancellationToken, Task> handle
899903
this.OnMessageHandler(messageHandlerOptions, handler);
900904
}
901905

906+
/// <summary>
907+
/// Unregister message handler from the receiver if there is an active message handler registered. This operation waits for the completion
908+
/// of inflight receive and message handling operations to finish and unregisters future receives on the message handler which previously
909+
/// registered.
910+
/// <param name="inflightMessageHandlerTasksWaitTimeout"> is the waitTimeout for inflight message handling tasks.
911+
/// </summary>
912+
public async Task UnregisterMessageHandlerAsync(TimeSpan inflightMessageHandlerTasksWaitTimeout)
913+
{
914+
this.ThrowIfClosed();
915+
916+
if (inflightMessageHandlerTasksWaitTimeout <= TimeSpan.Zero)
917+
{
918+
throw Fx.Exception.ArgumentOutOfRange(nameof(inflightMessageHandlerTasksWaitTimeout), inflightMessageHandlerTasksWaitTimeout, Resources.TimeoutMustBePositiveNonZero.FormatForUser(nameof(inflightMessageHandlerTasksWaitTimeout), inflightMessageHandlerTasksWaitTimeout));
919+
}
920+
921+
MessagingEventSource.Log.UnregisterMessageHandlerStart(this.ClientId);
922+
lock (this.messageReceivePumpSyncLock)
923+
{
924+
if (this.receivePump == null || this.receivePumpCancellationTokenSource.IsCancellationRequested)
925+
{
926+
// Silently return if handler has already been unregistered.
927+
return;
928+
}
929+
930+
this.receivePumpCancellationTokenSource.Cancel();
931+
this.receivePumpCancellationTokenSource.Dispose();
932+
}
933+
934+
Stopwatch stopWatch = Stopwatch.StartNew();
935+
while (this.receivePump != null
936+
&& stopWatch.Elapsed < inflightMessageHandlerTasksWaitTimeout
937+
&& this.receivePump.maxConcurrentCallsSemaphoreSlim.CurrentCount < this.receivePump.registerHandlerOptions.MaxConcurrentCalls)
938+
{
939+
await Task.Delay(10).ConfigureAwait(false);
940+
}
941+
942+
lock (this.messageReceivePumpSyncLock)
943+
{
944+
this.runningTaskCancellationTokenSource.Cancel();
945+
this.runningTaskCancellationTokenSource.Dispose();
946+
this.receivePump = null;
947+
}
948+
MessagingEventSource.Log.UnregisterMessageHandlerStop(this.ClientId);
949+
}
950+
902951
/// <summary>
903952
/// Registers a <see cref="ServiceBusPlugin"/> to be used with this receiver.
904953
/// </summary>
@@ -1003,6 +1052,9 @@ protected override async Task OnClosingAsync()
10031052
{
10041053
this.receivePumpCancellationTokenSource.Cancel();
10051054
this.receivePumpCancellationTokenSource.Dispose();
1055+
// For back-compatibility
1056+
this.runningTaskCancellationTokenSource.Cancel();
1057+
this.runningTaskCancellationTokenSource.Dispose();
10061058
this.receivePump = null;
10071059
}
10081060
}
@@ -1279,7 +1331,13 @@ protected virtual void OnMessageHandler(
12791331
}
12801332

12811333
this.receivePumpCancellationTokenSource = new CancellationTokenSource();
1282-
this.receivePump = new MessageReceivePump(this, registerHandlerOptions, callback, this.ServiceBusConnection.Endpoint, this.receivePumpCancellationTokenSource.Token);
1334+
1335+
if (this.runningTaskCancellationTokenSource == null)
1336+
{
1337+
this.runningTaskCancellationTokenSource = new CancellationTokenSource();
1338+
}
1339+
1340+
this.receivePump = new MessageReceivePump(this, registerHandlerOptions, callback, this.ServiceBusConnection.Endpoint, this.receivePumpCancellationTokenSource.Token, this.runningTaskCancellationTokenSource.Token);
12831341
}
12841342

12851343
try
@@ -1295,6 +1353,8 @@ protected virtual void OnMessageHandler(
12951353
{
12961354
this.receivePumpCancellationTokenSource.Cancel();
12971355
this.receivePumpCancellationTokenSource.Dispose();
1356+
this.runningTaskCancellationTokenSource.Cancel();
1357+
this.runningTaskCancellationTokenSource.Dispose();
12981358
this.receivePump = null;
12991359
}
13001360
}

‎sdk/servicebus/Microsoft.Azure.ServiceBus/src/IQueueClient.cs‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -77,5 +77,13 @@ public interface IQueueClient : IReceiverClient, ISenderClient
7777
/// <param name="sessionHandlerOptions">Options used to configure the settings of the session pump.</param>
7878
/// <remarks>Enable prefetch to speed up the receive rate. </remarks>
7979
void RegisterSessionHandler(Func<IMessageSession, Message, CancellationToken, Task> handler, SessionHandlerOptions sessionHandlerOptions);
80+
81+
/// <summary>
82+
/// Unregister session handler from the receiver if there is an active session handler registered. This operation waits for the completion
83+
/// of inflight receive and session handling operations to finish and unregisters future receives on the session handler which previously
84+
/// registered.
85+
/// <param name="inflightSessionHandlerTasksWaitTimeout"> is the waitTimeout for inflight session handling tasks.
86+
/// </summary>
87+
Task UnregisterSessionHandlerAsync(TimeSpan inflightSessionHandlerTasksWaitTimeout);
8088
}
8189
}

‎sdk/servicebus/Microsoft.Azure.ServiceBus/src/ISubscriptionClient.cs‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -114,5 +114,13 @@ public interface ISubscriptionClient : IReceiverClient
114114
/// <param name="sessionHandlerOptions">Options used to configure the settings of the session pump.</param>
115115
/// <remarks>Enable prefetch to speed up the receive rate. </remarks>
116116
void RegisterSessionHandler(Func<IMessageSession, Message, CancellationToken, Task> handler, SessionHandlerOptions sessionHandlerOptions);
117+
118+
/// <summary>
119+
/// Unregister session handler from the receiver if there is an active session handler registered. This operation waits for the completion
120+
/// of inflight receive and session handling operations to finish and unregisters future receives on the session handler which previously
121+
/// registered.
122+
/// <param name="inflightSessionHandlerTasksWaitTimeout"> is the waitTimeout for inflight session handling tasks.
123+
/// </summary>
124+
Task UnregisterSessionHandlerAsync(TimeSpan inflightSessionHandlerTasksWaitTimeout);
117125
}
118126
}

‎sdk/servicebus/Microsoft.Azure.ServiceBus/src/MessageReceivePump.cs‎

Lines changed: 7 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -12,25 +12,28 @@ namespace Microsoft.Azure.ServiceBus
1212

1313
sealed class MessageReceivePump
1414
{
15+
public readonly SemaphoreSlim maxConcurrentCallsSemaphoreSlim;
16+
public readonly MessageHandlerOptions registerHandlerOptions;
1517
readonly Func<Message, CancellationToken, Task> onMessageCallback;
1618
readonly string endpoint;
17-
readonly MessageHandlerOptions registerHandlerOptions;
1819
readonly IMessageReceiver messageReceiver;
1920
readonly CancellationToken pumpCancellationToken;
20-
readonly SemaphoreSlim maxConcurrentCallsSemaphoreSlim;
21+
readonly CancellationToken runningTaskCancellationToken;
2122
readonly ServiceBusDiagnosticSource diagnosticSource;
2223

2324
public MessageReceivePump(IMessageReceiver messageReceiver,
2425
MessageHandlerOptions registerHandlerOptions,
2526
Func<Message, CancellationToken, Task> callback,
2627
Uri endpoint,
27-
CancellationToken pumpCancellationToken)
28+
CancellationToken pumpCancellationToken,
29+
CancellationToken runningTaskCancellationToken)
2830
{
2931
this.messageReceiver = messageReceiver ?? throw new ArgumentNullException(nameof(messageReceiver));
3032
this.registerHandlerOptions = registerHandlerOptions;
3133
this.onMessageCallback = callback;
3234
this.endpoint = endpoint.Authority;
3335
this.pumpCancellationToken = pumpCancellationToken;
36+
this.runningTaskCancellationToken = runningTaskCancellationToken;
3437
this.maxConcurrentCallsSemaphoreSlim = new SemaphoreSlim(this.registerHandlerOptions.MaxConcurrentCalls);
3538
this.diagnosticSource = new ServiceBusDiagnosticSource(messageReceiver.Path, endpoint);
3639
}
@@ -163,7 +166,7 @@ async Task MessageDispatchTask(Message message)
163166
try
164167
{
165168
MessagingEventSource.Log.MessageReceiverPumpUserCallbackStart(this.messageReceiver.ClientId, message);
166-
await this.onMessageCallback(message, this.pumpCancellationToken).ConfigureAwait(false);
169+
await this.onMessageCallback(message, this.runningTaskCancellationToken).ConfigureAwait(false);
167170

168171
MessagingEventSource.Log.MessageReceiverPumpUserCallbackStop(this.messageReceiver.ClientId, message);
169172
}

‎sdk/servicebus/Microsoft.Azure.ServiceBus/src/MessagingEventSource.cs‎

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1380,6 +1380,42 @@ public void ManagementSerializationException(string objectName, string details =
13801380
this.WriteEvent(117, objectName, details);
13811381
}
13821382
}
1383+
1384+
[Event(118, Level = EventLevel.Informational, Message = "{0}: Unregister MessageHandler start.")]
1385+
public void UnregisterMessageHandlerStart(string clientId)
1386+
{
1387+
if (this.IsEnabled())
1388+
{
1389+
this.WriteEvent(118, clientId);
1390+
}
1391+
}
1392+
1393+
[Event(119, Level = EventLevel.Informational, Message = "{0}: Unregister MessageHandler done.")]
1394+
public void UnregisterMessageHandlerStop(string clientId)
1395+
{
1396+
if (this.IsEnabled())
1397+
{
1398+
this.WriteEvent(119, clientId);
1399+
}
1400+
}
1401+
1402+
[Event(120, Level = EventLevel.Informational, Message = "{0}: Unregister SessionHandler start.")]
1403+
public void UnregisterSessionHandlerStart(string clientId)
1404+
{
1405+
if (this.IsEnabled())
1406+
{
1407+
this.WriteEvent(120, clientId);
1408+
}
1409+
}
1410+
1411+
[Event(121, Level = EventLevel.Informational, Message = "{0}: Unregister SessionHandler done.")]
1412+
public void UnregisterSessionHandlerStop(string clientId)
1413+
{
1414+
if (this.IsEnabled())
1415+
{
1416+
this.WriteEvent(121, clientId);
1417+
}
1418+
}
13831419
}
13841420

13851421
internal static class TraceHelper

‎sdk/servicebus/Microsoft.Azure.ServiceBus/src/QueueClient.cs‎

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -446,6 +446,18 @@ public void RegisterMessageHandler(Func<Message, CancellationToken, Task> handle
446446
this.InnerReceiver.RegisterMessageHandler(handler, messageHandlerOptions);
447447
}
448448

449+
/// <summary>
450+
/// Unregister message handler from the receiver if there is an active message handler registered. This operation waits for the completion
451+
/// of inflight receive and message handling operations to finish and unregisters future receives on the message handler which previously
452+
/// registered.
453+
/// <param name="inflightMessageHandlerTasksWaitTimeout"> is the waitTimeout for inflight message handling tasks.
454+
/// </summary>
455+
public async Task UnregisterMessageHandlerAsync(TimeSpan inflightMessageHandlerTasksWaitTimeout)
456+
{
457+
this.ThrowIfClosed();
458+
await this.InnerReceiver.UnregisterMessageHandlerAsync(inflightMessageHandlerTasksWaitTimeout).ConfigureAwait(false);
459+
}
460+
449461
/// <summary>
450462
/// Receive session messages continuously from the queue. Registers a message handler and begins a new thread to receive session-messages.
451463
/// This handler(<see cref="Func{IMessageSession, Message, CancellationToken, Task}"/>) is awaited on every time a new message is received by the queue client.
@@ -476,6 +488,18 @@ public void RegisterSessionHandler(Func<IMessageSession, Message, CancellationTo
476488
this.SessionPumpHost.OnSessionHandler(handler, sessionHandlerOptions);
477489
}
478490

491+
/// <summary>
492+
/// Unregister session handler from the receiver if there is an active session handler registered. This operation waits for the completion
493+
/// of inflight receive and session handling operations to finish and unregisters future receives on the session handler which previously
494+
/// registered.
495+
/// <param name="inflightSessionHandlerTasksWaitTimeout"> is the waitTimeout for inflight session handling tasks.
496+
/// </summary>
497+
public async Task UnregisterSessionHandlerAsync(TimeSpan inflightSessionHandlerTasksWaitTimeout)
498+
{
499+
this.ThrowIfClosed();
500+
await this.SessionPumpHost.UnregisterSessionHandlerAsync(inflightSessionHandlerTasksWaitTimeout).ConfigureAwait(false);
501+
}
502+
479503
/// <summary>
480504
/// Schedules a message to appear on Service Bus at a later time.
481505
/// </summary>

‎sdk/servicebus/Microsoft.Azure.ServiceBus/src/SessionPumpHost.cs‎

Lines changed: 53 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,9 @@
33

44
namespace Microsoft.Azure.ServiceBus
55
{
6+
using Microsoft.Azure.ServiceBus.Primitives;
67
using System;
8+
using System.Diagnostics;
79
using System.Threading;
810
using System.Threading.Tasks;
911

@@ -12,6 +14,7 @@ internal sealed class SessionPumpHost
1214
readonly object syncLock;
1315
SessionReceivePump sessionReceivePump;
1416
CancellationTokenSource sessionPumpCancellationTokenSource;
17+
CancellationTokenSource runningTaskCancellationTokenSource;
1518
readonly Uri endpoint;
1619

1720
public SessionPumpHost(string clientId, ReceiveMode receiveMode, ISessionClient sessionClient, Uri endpoint)
@@ -35,6 +38,9 @@ public void Close()
3538
{
3639
this.sessionPumpCancellationTokenSource?.Cancel();
3740
this.sessionPumpCancellationTokenSource?.Dispose();
41+
// For back-compatibility
42+
this.runningTaskCancellationTokenSource?.Cancel();
43+
this.runningTaskCancellationTokenSource?.Dispose();
3844
this.sessionReceivePump = null;
3945
}
4046
}
@@ -53,14 +59,22 @@ public void OnSessionHandler(
5359
}
5460

5561
this.sessionPumpCancellationTokenSource = new CancellationTokenSource();
62+
63+
// Running task cancellation token source can be reused if previously UnregisterSessionHandlerAsync was called
64+
if (this.runningTaskCancellationTokenSource == null)
65+
{
66+
this.runningTaskCancellationTokenSource = new CancellationTokenSource();
67+
}
68+
5669
this.sessionReceivePump = new SessionReceivePump(
5770
this.ClientId,
5871
this.SessionClient,
5972
this.ReceiveMode,
6073
sessionHandlerOptions,
6174
callback,
6275
this.endpoint,
63-
this.sessionPumpCancellationTokenSource.Token);
76+
this.sessionPumpCancellationTokenSource.Token,
77+
this.runningTaskCancellationTokenSource.Token);
6478
}
6579

6680
try
@@ -82,5 +96,43 @@ public void OnSessionHandler(
8296

8397
MessagingEventSource.Log.RegisterOnSessionHandlerStop(this.ClientId);
8498
}
99+
100+
public async Task UnregisterSessionHandlerAsync(TimeSpan inflightSessionHandlerTasksWaitTimeout)
101+
{
102+
if (inflightSessionHandlerTasksWaitTimeout <= TimeSpan.Zero)
103+
{
104+
throw Fx.Exception.ArgumentOutOfRange(nameof(inflightSessionHandlerTasksWaitTimeout), inflightSessionHandlerTasksWaitTimeout, Resources.TimeoutMustBePositiveNonZero.FormatForUser(nameof(inflightSessionHandlerTasksWaitTimeout), inflightSessionHandlerTasksWaitTimeout));
105+
}
106+
107+
MessagingEventSource.Log.UnregisterSessionHandlerStart(this.ClientId);
108+
lock (this.syncLock)
109+
{
110+
if (this.sessionReceivePump == null || this.sessionPumpCancellationTokenSource.IsCancellationRequested)
111+
{
112+
// Silently return if handler has already been unregistered.
113+
return;
114+
}
115+
116+
this.sessionPumpCancellationTokenSource.Cancel();
117+
this.sessionPumpCancellationTokenSource.Dispose();
118+
}
119+
120+
Stopwatch stopWatch = Stopwatch.StartNew();
121+
while (this.sessionReceivePump != null
122+
&& stopWatch.Elapsed < inflightSessionHandlerTasksWaitTimeout
123+
&& (this.sessionReceivePump.maxConcurrentSessionsSemaphoreSlim.CurrentCount <
124+
this.sessionReceivePump.sessionHandlerOptions.MaxConcurrentSessions
125+
|| this.sessionReceivePump.maxPendingAcceptSessionsSemaphoreSlim.CurrentCount <
126+
this.sessionReceivePump.sessionHandlerOptions.MaxConcurrentAcceptSessionCalls))
127+
{
128+
await Task.Delay(10).ConfigureAwait(false);
129+
}
130+
131+
lock (this.sessionPumpCancellationTokenSource)
132+
{
133+
this.sessionReceivePump = null;
134+
}
135+
MessagingEventSource.Log.UnregisterSessionHandlerStop(this.ClientId);
136+
}
85137
}
86138
}

0 commit comments

Comments
 (0)