diff --git a/src/DurableTask.ServiceBus/Common/Abstraction/ServiceBusAbstraction.cs b/src/DurableTask.ServiceBus/Common/Abstraction/ServiceBusAbstraction.cs
index 26cecadc..0a8f72e2 100644
--- a/src/DurableTask.ServiceBus/Common/Abstraction/ServiceBusAbstraction.cs
+++ b/src/DurableTask.ServiceBus/Common/Abstraction/ServiceBusAbstraction.cs
@@ -4,6 +4,7 @@ namespace DurableTask.ServiceBus.Common.Abstraction
using System;
using System.Collections.Generic;
using System.Linq;
+ using System.Threading;
using System.Threading.Tasks;
///
@@ -497,19 +498,31 @@ public SessionClient(ServiceBusConnection serviceBusConnection, string entityPat
this.serviceBusClient = serviceBusConnection.GetSendViaEntityScopedClient(entityPath);
}
+ ///
+ /// Accepts the next available session, or returns null if none becomes available within .
+ /// A non-positive timeout falls back to the client's retry policy timeout.
+ ///
public async Task AcceptMessageSessionAsync(TimeSpan operationTimeout)
{
- try
+ var options = new Azure.Messaging.ServiceBus.ServiceBusSessionReceiverOptions
{
- var options = new Azure.Messaging.ServiceBus.ServiceBusSessionReceiverOptions
- {
- ReceiveMode = receiveMode,
- };
- return new IMessageSession(await this.serviceBusClient.AcceptNextSessionAsync(entityPath, options));
- }
- catch (Azure.Messaging.ServiceBus.ServiceBusException e) when (e.Reason.Equals(Azure.Messaging.ServiceBus.ServiceBusFailureReason.ServiceTimeout))
+ ReceiveMode = receiveMode,
+ };
+
+ using (var timeoutCts = operationTimeout > TimeSpan.Zero ? new CancellationTokenSource(operationTimeout) : new CancellationTokenSource())
{
- return null;
+ try
+ {
+ return new IMessageSession(await this.serviceBusClient.AcceptNextSessionAsync(entityPath, options, timeoutCts.Token));
+ }
+ catch (Azure.Messaging.ServiceBus.ServiceBusException e) when (e.Reason.Equals(Azure.Messaging.ServiceBus.ServiceBusFailureReason.ServiceTimeout))
+ {
+ return null;
+ }
+ catch (OperationCanceledException) when (timeoutCts.IsCancellationRequested)
+ {
+ return null;
+ }
}
}
diff --git a/test/DurableTask.ServiceBus.Tests/SessionClientTests.cs b/test/DurableTask.ServiceBus.Tests/SessionClientTests.cs
new file mode 100644
index 00000000..0e880f27
--- /dev/null
+++ b/test/DurableTask.ServiceBus.Tests/SessionClientTests.cs
@@ -0,0 +1,157 @@
+// ----------------------------------------------------------------------------------
+// Copyright Microsoft Corporation
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+// http://www.apache.org/licenses/LICENSE-2.0
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+// ----------------------------------------------------------------------------------
+
+namespace DurableTask.ServiceBus.Tests
+{
+ using System;
+ using System.Diagnostics;
+ using System.Reflection;
+ using System.Runtime.CompilerServices;
+ using System.Threading;
+ using System.Threading.Tasks;
+ using Azure.Messaging.ServiceBus;
+ using DurableTask.ServiceBus.Common.Abstraction;
+ using Microsoft.VisualStudio.TestTools.UnitTesting;
+
+ [TestClass]
+ public class SessionClientTests
+ {
+ [TestMethod]
+ public async Task AcceptMessageSessionAsyncReturnsSessionWhenAccepted()
+ {
+ var client = new FakeServiceBusClient((ct) => Task.FromResult(new FakeSessionReceiver("session-1")));
+
+ IMessageSession session = await CreateSessionClient(client).AcceptMessageSessionAsync(TimeSpan.FromSeconds(30));
+
+ Assert.IsNotNull(session);
+ Assert.AreEqual("session-1", session.SessionId);
+ }
+
+ [TestMethod]
+ public async Task AcceptMessageSessionAsyncReturnsNullWhenOperationTimeoutElapses()
+ {
+ var client = new FakeServiceBusClient(async (ct) =>
+ {
+ await Task.Delay(TimeSpan.FromSeconds(10), ct);
+ return null;
+ });
+
+ Stopwatch timer = Stopwatch.StartNew();
+ IMessageSession session = await CreateSessionClient(client).AcceptMessageSessionAsync(TimeSpan.FromMilliseconds(200));
+ timer.Stop();
+
+ Assert.IsNull(session);
+ Assert.IsTrue(client.LastToken.IsCancellationRequested, "The operation timeout should cancel the accept call.");
+ Assert.IsTrue(timer.Elapsed < TimeSpan.FromSeconds(5), $"Accept should return shortly after the operation timeout, but took {timer.Elapsed}.");
+ }
+
+ [TestMethod]
+ public async Task AcceptMessageSessionAsyncReturnsNullOnServiceTimeout()
+ {
+ var client = new FakeServiceBusClient((ct) =>
+ throw new ServiceBusException("No session available.", ServiceBusFailureReason.ServiceTimeout));
+
+ IMessageSession session = await CreateSessionClient(client).AcceptMessageSessionAsync(TimeSpan.FromSeconds(30));
+
+ Assert.IsNull(session);
+ }
+
+ [TestMethod]
+ public async Task AcceptMessageSessionAsyncWithNonPositiveTimeoutDoesNotCancel()
+ {
+ var client = new FakeServiceBusClient(async (ct) =>
+ {
+ await Task.Delay(200);
+ Assert.IsFalse(ct.IsCancellationRequested, "A non-positive timeout should not cancel the accept call.");
+ throw new ServiceBusException("No session available.", ServiceBusFailureReason.ServiceTimeout);
+ });
+
+ IMessageSession session = await CreateSessionClient(client).AcceptMessageSessionAsync(TimeSpan.Zero);
+
+ Assert.IsNull(session);
+ }
+
+ [TestMethod]
+ public async Task AcceptMessageSessionAsyncPropagatesCancellationNotCausedByTimeout()
+ {
+ var client = new FakeServiceBusClient((ct) => throw new OperationCanceledException());
+
+ await Assert.ThrowsExceptionAsync(
+ () => CreateSessionClient(client).AcceptMessageSessionAsync(TimeSpan.FromSeconds(30)));
+ }
+
+ [TestMethod]
+ public async Task AcceptMessageSessionAsyncPropagatesOtherServiceBusErrors()
+ {
+ var client = new FakeServiceBusClient((ct) =>
+ throw new ServiceBusException("Entity not found.", ServiceBusFailureReason.MessagingEntityNotFound));
+
+ ServiceBusException exception = await Assert.ThrowsExceptionAsync(
+ () => CreateSessionClient(client).AcceptMessageSessionAsync(TimeSpan.FromSeconds(30)));
+ Assert.AreEqual(ServiceBusFailureReason.MessagingEntityNotFound, exception.Reason);
+ }
+
+ static SessionClient CreateSessionClient(ServiceBusClient serviceBusClient)
+ {
+#if NETFRAMEWORK
+ var sessionClient = (SessionClient)System.Runtime.Serialization.FormatterServices.GetUninitializedObject(typeof(SessionClient));
+#else
+ var sessionClient = (SessionClient)RuntimeHelpers.GetUninitializedObject(typeof(SessionClient));
+#endif
+ SetField(sessionClient, "serviceBusClient", serviceBusClient);
+ SetField(sessionClient, "entityPath", "test-entity");
+ SetField(sessionClient, "receiveMode", ServiceBusReceiveMode.PeekLock);
+ return sessionClient;
+ }
+
+ static void SetField(object target, string name, object value)
+ {
+ FieldInfo field = typeof(SessionClient).GetField(name, BindingFlags.Instance | BindingFlags.NonPublic);
+ Assert.IsNotNull(field, $"Field '{name}' was not found on {nameof(SessionClient)}.");
+ field.SetValue(target, value);
+ }
+
+ sealed class FakeServiceBusClient : ServiceBusClient
+ {
+ readonly Func> acceptNextSession;
+
+ public FakeServiceBusClient(Func> acceptNextSession)
+ {
+ this.acceptNextSession = acceptNextSession;
+ }
+
+ public CancellationToken LastToken { get; private set; }
+
+ public override Task AcceptNextSessionAsync(
+ string queueName,
+ ServiceBusSessionReceiverOptions options = default,
+ CancellationToken cancellationToken = default)
+ {
+ this.LastToken = cancellationToken;
+ return this.acceptNextSession(cancellationToken);
+ }
+ }
+
+ sealed class FakeSessionReceiver : ServiceBusSessionReceiver
+ {
+ readonly string sessionId;
+
+ public FakeSessionReceiver(string sessionId)
+ {
+ this.sessionId = sessionId;
+ }
+
+ public override string SessionId => this.sessionId;
+ }
+ }
+}