Support Timeout.Zero in ReceiveMessagesAsync (Azure#22934)

* Support Timeout.Zero in ReceiveMessagesAsync

* upgrading AMQP in Data.props

* updating AMQP version for Track 2 in Package.Data.props

* Update Packages.Data.props

reverting version change of AMQP for track 1

Co-authored-by: Olha Tkachenko <[email protected]>

Author: OlhaTkachenko

eng/Packages.Data.props

@@ -90,7 +90,7 @@
     <PackageReference Update="Azure.Storage.Blobs" Version="12.8.0" />
 
     <!-- Other approved packages -->
-    <PackageReference Update="Microsoft.Azure.Amqp" Version="2.4.13" />
+    <PackageReference Update="Microsoft.Azure.Amqp" Version="2.5.1" />
     <PackageReference Update="Microsoft.Identity.Client" Version="4.30.1" />
     <PackageReference Update="Microsoft.Identity.Client.Extensions.Msal" Version="2.18.4" />
 

sdk/servicebus/Azure.Messaging.ServiceBus/src/Receiver/ServiceBusReceiver.cs

@@ -260,7 +260,15 @@ internal async Task<IReadOnlyList<ServiceBusReceivedMessage>> ReceiveMessagesAsy
 
             if (maxWaitTime.HasValue)
             {
-                Argument.AssertPositive(maxWaitTime.Value, nameof(maxWaitTime));
+                // maxWaitTime could be zero only when prefetch enabled
+                if (PrefetchCount > 0)
+                {
+                    Argument.AssertNotNegative(maxWaitTime.Value, nameof(maxWaitTime));
+                }
+                else
+                {
+                    Argument.AssertPositive(maxWaitTime.Value, nameof(maxWaitTime));
+                }
             }
             if (PrefetchCount > 0 && maxMessages > PrefetchCount)
             {

sdk/servicebus/Azure.Messaging.ServiceBus/tests/Receiver/ReceiverLiveTests.cs

@@ -3,6 +3,7 @@
 
 using System;
 using System.Collections.Generic;
+using System.Diagnostics;
 using System.Linq;
 using System.Text;
 using System.Threading;
@@ -46,6 +47,28 @@ public async Task PeekUsingConnectionStringWithSharedKey()
             }
         }
 
+        [Test]
+        public async Task PeekWithZeroTimeout()
+        {
+            await using (var scope =
+                await ServiceBusScope.CreateWithQueue(enablePartitioning: false, enableSession: false))
+            {
+                await using var client = new ServiceBusClient(TestEnvironment.ServiceBusConnectionString);
+                await using var receiverWithPrefetch = client.CreateReceiver(scope.QueueName,
+                    options: new ServiceBusReceiverOptions() {PrefetchCount = 10});
+
+                var stopwatch = new Stopwatch();
+                stopwatch.Start();
+                await receiverWithPrefetch.ReceiveMessagesAsync(10, TimeSpan.Zero).ConfigureAwait(false);
+                stopwatch.Stop();
+                var durationWithPrefetchModeInSecs = stopwatch.Elapsed.TotalSeconds;
+
+                // If prefetch is enabled, timeout 0 secs will not be replaced with default timeout.
+                // In such case, only prefetched messages will be returned and no call to server will be made and call will be very fast.
+                Assert.IsTrue(durationWithPrefetchModeInSecs < 1);
+            }
+        }
+
         [Test]
         public async Task PeekUsingConnectionStringWithSas()
         {

sdk/servicebus/Azure.Messaging.ServiceBus/tests/Receiver/ReceiverTests.cs

@@ -2,11 +2,11 @@
 // Licensed under the MIT License.
 
 using System;
+using System.Collections.Generic;
 using System.Text;
 using System.Threading;
 using System.Threading.Tasks;
 using Azure.Core;
-using Azure.Messaging.ServiceBus.Amqp;
 using Azure.Messaging.ServiceBus.Core;
 using Moq;
 using NUnit.Framework;
@@ -121,6 +121,48 @@ public void ReceiveValidatesMaxWaitTime()
                 Throws.InstanceOf<ArgumentOutOfRangeException>());
         }
 
+        [Test]
+        public async Task ReceiveValidatesMaxWaitTimePrefetchMode()
+        {
+            var mockTransportReceiver = new Mock<TransportReceiver>();
+            var mockConnection = CreateMockConnection();
+            mockConnection.Setup(
+                    connection => connection.CreateTransportReceiver(
+                        It.IsAny<string>(),
+                        It.IsAny<ServiceBusRetryPolicy>(),
+                        It.IsAny<ServiceBusReceiveMode>(),
+                        It.IsAny<uint>(),
+                        It.IsAny<string>(),
+                        It.IsAny<string>(),
+                        It.IsAny<bool>(),
+                        It.IsAny<bool>(),
+                        It.IsAny<CancellationToken>()))
+                .Returns(mockTransportReceiver.Object);
+            IReadOnlyList<ServiceBusReceivedMessage> receivedMessages = new[]
+            {
+                ServiceBusModelFactory.ServiceBusReceivedMessage(messageId: Guid.NewGuid().ToString())
+            };
+            mockTransportReceiver.Setup(transportReceiver =>
+                    transportReceiver.ReceiveMessagesAsync(It.IsAny<int>(), It.IsAny<TimeSpan?>(),
+                        It.IsAny<CancellationToken>()))
+                .Returns(Task.FromResult(receivedMessages));
+
+            var receiver = new ServiceBusReceiver(mockConnection.Object, "queue", default,
+                new ServiceBusReceiverOptions {PrefetchCount = 10});
+
+            Assert.That(
+                async () => await receiver.ReceiveMessagesAsync(10, TimeSpan.FromSeconds(-1)),
+                Throws.InstanceOf<ArgumentOutOfRangeException>());
+            Assert.That(
+                async () => await receiver.ReceiveMessageAsync(TimeSpan.FromSeconds(-1)),
+                Throws.InstanceOf<ArgumentOutOfRangeException>());
+
+            var actuallyReceivedMessages =
+                await receiver.ReceiveMessagesAsync(10, TimeSpan.FromSeconds(0)).ConfigureAwait(false);
+            Assert.AreEqual(receivedMessages.Count, actuallyReceivedMessages.Count);
+            Assert.AreEqual(receivedMessages[0].MessageId, actuallyReceivedMessages[0].MessageId);
+        }
+
         [Test]
         public async Task ReceiveMessageValidatesClientIsNotDisposed()
         {