Skip to content

Don't pass cancelled token to TryExecute when draining #38063

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 5 commits into from
Aug 10, 2023
Merged

Don't pass cancelled token to TryExecute when draining #38063

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
1bc9d18
Don't pass cancelled token to TryExecute when draining
JoshLove-msft Aug 7, 2023
f687e0f
comment
JoshLove-msft Aug 8, 2023
018830e
Use graceful shutdown in Dispose
JoshLove-msft Aug 8, 2023
db286aa
Update to respect drain mode
JoshLove-msft Aug 9, 2023
c6731cd
Update sdk/servicebus/Microsoft.Azure.WebJobs.Extensions.ServiceBus/s…
JoshLove-msft Aug 9, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
34 changes: 20 additions & 14 deletions ...vicebus/Microsoft.Azure.WebJobs.Extensions.ServiceBus/src/Listeners/ServiceBusListener.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,8 @@ internal sealed class ServiceBusListener : IListener, IScaleMonitorProvider, ITa
private readonly string _entityPath;
private readonly bool _isSessionsEnabled;
private readonly bool _autoCompleteMessages;
private readonly CancellationTokenSource _cancellationTokenSource;
private readonly CancellationTokenSource _stoppingCancellationTokenSource;
private readonly CancellationTokenSource _disposingCancellationTokenSource;
private readonly ServiceBusOptions _serviceBusOptions;
private readonly bool _singleDispatch;
private readonly ILogger<ServiceBusListener> _logger;
Expand Down Expand Up @@ -85,7 +86,8 @@ public ServiceBusListener(
_isSessionsEnabled = isSessionsEnabled;
_autoCompleteMessages = autoCompleteMessages;
_triggerExecutor = triggerExecutor;
_cancellationTokenSource = new CancellationTokenSource();
_stoppingCancellationTokenSource = new CancellationTokenSource();
_disposingCancellationTokenSource = new CancellationTokenSource();
_logger = loggerFactory.CreateLogger<ServiceBusListener>();
_functionId = functionId;

Expand Down Expand Up @@ -195,7 +197,7 @@ public async Task StartAsync(CancellationToken cancellationToken)
}
else
{
_batchLoop = RunBatchReceiveLoopAsync(_cancellationTokenSource);
_batchLoop = RunBatchReceiveLoopAsync(_stoppingCancellationTokenSource);
}
}
catch
Expand Down Expand Up @@ -226,7 +228,7 @@ public async Task StopAsync(CancellationToken cancellationToken)
}

// This will also cancel the background monitoring task through the linked cancellation token source.
_cancellationTokenSource.Cancel();
_stoppingCancellationTokenSource.Cancel();

// CloseAsync method stop new messages from being processed while allowing in-flight messages to be processed.
if (_singleDispatch)
Expand Down Expand Up @@ -273,7 +275,7 @@ public async Task StopAsync(CancellationToken cancellationToken)
public void Cancel()
{
ThrowIfDisposed();
_cancellationTokenSource.Cancel();
_stoppingCancellationTokenSource.Cancel();
}

[SuppressMessage("Microsoft.Usage", "CA2213:DisposableFieldsShouldBeDisposed", MessageId = "_cancellationTokenSource")]
Expand All @@ -290,7 +292,8 @@ public void Dispose()
// Mark it canceled but don't dispose of the source while the callers are running.
// Otherwise, callers would receive ObjectDisposedException when calling token.Register.
// For now, rely on finalization to clean up _cancellationTokenSource's wait handle (if allocated).
_cancellationTokenSource.Cancel();
_stoppingCancellationTokenSource.Cancel();
_disposingCancellationTokenSource.Cancel();

if (_batchReceiver.IsValueCreated)
{
Expand All @@ -299,14 +302,17 @@ public void Dispose()
#pragma warning restore AZC0102 // Do not use GetAwaiter().GetResult().
}

if (_messageProcessor.IsValueCreated)
// Only close the processor if it is not processing. If it is still processing, then calling CloseAsync would wait
// for the processing to complete, which is not what we want to do here since this is the final step of shutdown.
// The connection will still be closed when the client is disposed.
if (_messageProcessor.IsValueCreated && !_messageProcessor.Value.Processor.IsProcessing)
{
#pragma warning disable AZC0102 // Do not use GetAwaiter().GetResult().
_messageProcessor.Value.Processor.CloseAsync(CancellationToken.None).GetAwaiter().GetResult();
#pragma warning restore AZC0102 // Do not use GetAwaiter().GetResult().
}

if (_sessionMessageProcessor.IsValueCreated)
if (_sessionMessageProcessor.IsValueCreated && !_sessionMessageProcessor.Value.Processor.IsProcessing)
{
#pragma warning disable AZC0102 // Do not use GetAwaiter().GetResult().
_sessionMessageProcessor.Value.Processor.CloseAsync(CancellationToken.None).GetAwaiter().GetResult();
Expand All @@ -321,7 +327,7 @@ public void Dispose()
}

_stopAsyncSemaphore.Dispose();
_cancellationTokenSource.Dispose();
_stoppingCancellationTokenSource.Dispose();
_batchReceiveRegistration.Dispose();
_concurrencyUpdateManager?.Dispose();

Expand All @@ -336,7 +342,7 @@ internal async Task ProcessMessageAsync(ProcessMessageEventArgs args)

_concurrencyUpdateManager?.MessageProcessed();

using (CancellationTokenSource linkedCts = CancellationTokenSource.CreateLinkedTokenSource(args.CancellationToken, _cancellationTokenSource.Token))
using (CancellationTokenSource linkedCts = CancellationTokenSource.CreateLinkedTokenSource(args.CancellationToken, _stoppingCancellationTokenSource.Token))
{
var actions = new ServiceBusMessageActions(args);
if (!await _messageProcessor.Value.BeginProcessingMessageAsync(actions, args.Message, linkedCts.Token).ConfigureAwait(false))
Expand All @@ -349,7 +355,7 @@ internal async Task ProcessMessageAsync(ProcessMessageEventArgs args)

TriggeredFunctionData data = input.GetTriggerFunctionData();

FunctionResult result = await _triggerExecutor.TryExecuteAsync(data, linkedCts.Token).ConfigureAwait(false);
FunctionResult result = await _triggerExecutor.TryExecuteAsync(data, _disposingCancellationTokenSource.Token).ConfigureAwait(false);
try
{
await _messageProcessor.Value.CompleteProcessingMessageAsync(actions, args.Message, result, linkedCts.Token)
Expand All @@ -369,7 +375,7 @@ internal async Task ProcessSessionMessageAsync(ProcessSessionMessageEventArgs ar
_concurrencyUpdateManager?.MessageProcessed();

using (CancellationTokenSource linkedCts =
CancellationTokenSource.CreateLinkedTokenSource(args.CancellationToken, _cancellationTokenSource.Token))
CancellationTokenSource.CreateLinkedTokenSource(args.CancellationToken, _stoppingCancellationTokenSource.Token))
{
var actions = new ServiceBusSessionMessageActions(args);
if (!await _sessionMessageProcessor.Value.BeginProcessingMessageAsync(actions, args.Message, linkedCts.Token)
Expand All @@ -382,7 +388,7 @@ internal async Task ProcessSessionMessageAsync(ProcessSessionMessageEventArgs ar
ServiceBusTriggerInput input = ServiceBusTriggerInput.CreateSingle(args.Message, actions, receiveActions, _client.Value);

TriggeredFunctionData data = input.GetTriggerFunctionData();
FunctionResult result = await _triggerExecutor.TryExecuteAsync(data, linkedCts.Token).ConfigureAwait(false);
FunctionResult result = await _triggerExecutor.TryExecuteAsync(data, _disposingCancellationTokenSource.Token).ConfigureAwait(false);

if (actions.ShouldReleaseSession)
{
Expand Down Expand Up @@ -643,7 +649,7 @@ private async Task TriggerAndCompleteMessagesInternal(ServiceBusReceivedMessage[
scope.SetMessageData(messagesArray);

scope.Start();
FunctionResult result = await _triggerExecutor.TryExecuteAsync(input.GetTriggerFunctionData(), cancellationToken).ConfigureAwait(false);
FunctionResult result = await _triggerExecutor.TryExecuteAsync(input.GetTriggerFunctionData(), _disposingCancellationTokenSource.Token).ConfigureAwait(false);
if (result.Exception != null)
{
scope.Failed(result.Exception);
Expand Down
64 changes: 56 additions & 8 deletions ...servicebus/Microsoft.Azure.WebJobs.Extensions.ServiceBus/tests/ServiceBusEndToEndTests.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -242,6 +242,30 @@ public async Task TestSingle_AutoCompleteEnabledOnTrigger_CompleteInFunction()
}
}

[Test]
public async Task TestSingle_Dispose()
{
await WriteQueueMessage("{'Name': 'Test1', 'Value': 'Value'}");
var host = BuildHost<TestSingleDispose>();

bool result = _waitHandle1.WaitOne(SBTimeoutMills);
Assert.True(result);
host.Dispose();
_waitHandle2.Set();
}

[Test]
public async Task TestBatch_Dispose()
{
await WriteQueueMessage("{'Name': 'Test1', 'Value': 'Value'}");
var host = BuildHost<TestBatchDispose>();

bool result = _waitHandle1.WaitOne(SBTimeoutMills);
Assert.True(result);
host.Dispose();
_waitHandle2.Set();
}

[Test]
public async Task TestSingle_AutoCompleteDisabledOnTrigger_AbandonsWhenException()
{
Expand Down Expand Up @@ -1588,6 +1612,34 @@ public static async Task RunAsync(
}
}

public class TestSingleDispose
{
public static void RunAsync(
[ServiceBusTrigger(FirstQueueNameKey)]
ServiceBusReceivedMessage message,
CancellationToken cancellationToken)
{
_waitHandle1.Set();
bool result = _waitHandle2.WaitOne(SBTimeoutMills);
Assert.IsTrue(result);
Assert.IsTrue(cancellationToken.IsCancellationRequested);
}
}

public class TestBatchDispose
{
public static void RunAsync(
[ServiceBusTrigger(FirstQueueNameKey)]
ServiceBusReceivedMessage[] message,
CancellationToken cancellationToken)
{
_waitHandle1.Set();
bool result = _waitHandle2.WaitOne(SBTimeoutMills);
Assert.IsTrue(result);
Assert.IsTrue(cancellationToken.IsCancellationRequested);
}
}

public class TestCrossEntityTransaction
{
public static async Task RunAsync(
Expand Down Expand Up @@ -1750,8 +1802,7 @@ public static async Task RunAsync(
{
logger.LogInformation($"DrainModeValidationFunctions.QueueNoSessions: message data {msg.Body}");
_drainValidationPreDelay.Set();
await DrainModeHelper.WaitForCancellationAsync(cancellationToken);
Assert.True(cancellationToken.IsCancellationRequested);
Assert.False(cancellationToken.IsCancellationRequested);
await messageActions.CompleteMessageAsync(msg);
_drainValidationPostDelay.Set();
}
Expand All @@ -1768,8 +1819,7 @@ public static async Task RunAsync(
{
logger.LogInformation($"DrainModeValidationFunctions.NoSessions: message data {msg.Body}");
_drainValidationPreDelay.Set();
await DrainModeHelper.WaitForCancellationAsync(cancellationToken);
Assert.True(cancellationToken.IsCancellationRequested);
Assert.False(cancellationToken.IsCancellationRequested);
await messageActions.CompleteMessageAsync(msg);
_drainValidationPostDelay.Set();
}
Expand All @@ -1787,8 +1837,7 @@ public static async Task RunAsync(
Assert.True(array.Length > 0);
logger.LogInformation($"DrainModeTestJobBatch.QueueNoSessionsBatch: received {array.Length} messages");
_drainValidationPreDelay.Set();
await DrainModeHelper.WaitForCancellationAsync(cancellationToken);
Assert.True(cancellationToken.IsCancellationRequested);
Assert.False(cancellationToken.IsCancellationRequested);
foreach (ServiceBusReceivedMessage msg in array)
{
await messageActions.CompleteMessageAsync(msg);
Expand All @@ -1808,8 +1857,7 @@ public static async Task RunAsync(
Assert.True(array.Length > 0);
logger.LogInformation($"DrainModeTestJobBatch.TopicNoSessionsBatch: received {array.Length} messages");
_drainValidationPreDelay.Set();
await DrainModeHelper.WaitForCancellationAsync(cancellationToken);
Assert.True(cancellationToken.IsCancellationRequested);
Assert.False(cancellationToken.IsCancellationRequested);
foreach (ServiceBusReceivedMessage msg in array)
{
// validate that manual lock renewal works
Expand Down
64 changes: 56 additions & 8 deletions ...us/Microsoft.Azure.WebJobs.Extensions.ServiceBus/tests/ServiceBusSessionsEndToEndTests.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -677,6 +677,30 @@ public async Task TestBatch_ProcessMessagesSpan_FailedScope()
Assert.IsTrue(scope.IsFailed);
}

[Test]
public async Task TestSingle_Dispose()
{
await WriteQueueMessage("{'Name': 'Test1', 'Value': 'Value'}", "sessionId1");
var host = BuildHost<TestSingleDispose>();

bool result = _waitHandle1.WaitOne(SBTimeoutMills);
Assert.True(result);
host.Dispose();
_waitHandle2.Set();
}

[Test]
public async Task TestBatch_Dispose()
{
await WriteQueueMessage("{'Name': 'Test1', 'Value': 'Value'}", "sessionId1");
var host = BuildHost<TestBatchDispose>();

bool result = _waitHandle1.WaitOne(SBTimeoutMills);
Assert.True(result);
host.Dispose();
_waitHandle2.Set();
}

private async Task TestMultiple<T>(bool isXml = false)
{
if (isXml)
Expand Down Expand Up @@ -830,8 +854,7 @@ public static async Task QueueWithSessions(
Assert.AreEqual(msg.PartitionKey, partitionKey);
Assert.AreEqual(msg.TransactionPartitionKey, transactionPartitionKey);
_drainValidationPreDelay.Set();
await DrainModeHelper.WaitForCancellationAsync(cancellationToken);
Assert.True(cancellationToken.IsCancellationRequested);
Assert.False(cancellationToken.IsCancellationRequested);
try
{
await messageActions.CompleteMessageAsync(msg);
Expand All @@ -856,8 +879,7 @@ public static async Task TopicWithSessions(
$"DrainModeValidationFunctions.TopicWithSessions: message data {msg.Body} with session id {msg.SessionId}");
Assert.AreEqual(_drainModeSessionId, msg.SessionId);
_drainValidationPreDelay.Set();
await DrainModeHelper.WaitForCancellationAsync(cancellationToken);
Assert.True(cancellationToken.IsCancellationRequested);
Assert.False(cancellationToken.IsCancellationRequested);
try
{
await messageSession.CompleteMessageAsync(msg);
Expand Down Expand Up @@ -885,8 +907,7 @@ public static async Task QueueWithSessionsBatch(
$"DrainModeTestJobBatch.QueueWithSessionsBatch: received {array.Length} messages with session id {array[0].SessionId}");
Assert.AreEqual(_drainModeSessionId, array[0].SessionId);
_drainValidationPreDelay.Set();
await DrainModeHelper.WaitForCancellationAsync(cancellationToken);
Assert.True(cancellationToken.IsCancellationRequested);
Assert.False(cancellationToken.IsCancellationRequested);
for (int i = 0; i < array.Length; i++)
{
var message = array[i];
Expand Down Expand Up @@ -918,8 +939,7 @@ public static async Task TopicWithSessionsBatch(
$"DrainModeTestJobBatch.TopicWithSessionsBatch: received {array.Length} messages with session id {array[0].SessionId}");
Assert.AreEqual(_drainModeSessionId, array[0].SessionId);
_drainValidationPreDelay.Set();
await DrainModeHelper.WaitForCancellationAsync(cancellationToken);
Assert.True(cancellationToken.IsCancellationRequested);
Assert.False(cancellationToken.IsCancellationRequested);
foreach (ServiceBusReceivedMessage msg in array)
{
await messageSession.CompleteMessageAsync(msg);
Expand Down Expand Up @@ -1160,6 +1180,34 @@ public static void RunAsync(
}
}

public class TestSingleDispose
{
public static void RunAsync(
[ServiceBusTrigger(FirstQueueNameKey, IsSessionsEnabled = true)]
ServiceBusReceivedMessage message,
CancellationToken cancellationToken)
{
_waitHandle1.Set();
bool result = _waitHandle2.WaitOne(SBTimeoutMills);
Assert.IsTrue(result);
Assert.IsTrue(cancellationToken.IsCancellationRequested);
}
}

public class TestBatchDispose
{
public static void RunAsync(
[ServiceBusTrigger(FirstQueueNameKey, IsSessionsEnabled = true)]
ServiceBusReceivedMessage[] message,
CancellationToken cancellationToken)
{
_waitHandle1.Set();
bool result = _waitHandle2.WaitOne(SBTimeoutMills);
Assert.IsTrue(result);
Assert.IsTrue(cancellationToken.IsCancellationRequested);
}
}

public class CustomMessagingProvider : MessagingProvider
{
public const string CustomMessagingCategory = "CustomMessagingProvider";
Expand Down