Improve cancellation error handling (#28713)

* Improve cancellation error handling

* Increase threshold for flaky tests

* More flaky tests

Author: JoshLove-msft

sdk/servicebus/Azure.Messaging.ServiceBus/src/Amqp/AmqpReceiver.cs

@@ -241,8 +241,8 @@ private async Task<ReceivingAmqpLink> OpenReceiverLinkAsync(
                 // to log here.
                 throw;
             }
-            catch (TaskCanceledException)
-                when (_isProcessor)
+            catch (OperationCanceledException)
+                when (_isProcessor && cancellationToken.IsCancellationRequested)
             {
                 // do not log this as the processor is shutting down
                 throw;

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

@@ -302,8 +302,8 @@ internal async Task<IReadOnlyList<ServiceBusReceivedMessage>> ReceiveMessagesAsy
                     maxWaitTime,
                     cancellationToken).ConfigureAwait(false);
             }
-            catch (TaskCanceledException ex)
-                when (isProcessor)
+            catch (OperationCanceledException ex)
+                when (isProcessor && cancellationToken.IsCancellationRequested)
             {
                 Logger.ProcessorStoppingReceiveCanceled(Identifier, ex.ToString());
                 scope.Failed(ex);

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

@@ -157,7 +157,7 @@ public async Task ReceiveMessagesWhenQueueEmpty()
                 Assert.ThrowsAsync<TaskCanceledException>(async () => await receiver.ReceiveMessagesAsync(1, cancellationToken: cancellationTokenSource.Token));
                 var stop = DateTime.UtcNow;
 
-                Assert.That(stop - start, Is.EqualTo(TimeSpan.FromSeconds(3)).Within(TimeSpan.FromSeconds(3)));
+                Assert.That(stop - start, Is.EqualTo(TimeSpan.FromSeconds(3)).Within(TimeSpan.FromSeconds(5)));
             }
         }
 
@@ -230,7 +230,7 @@ public async Task CancellingDoesNotBlockSubsequentReceives(bool prefetch)
                 Assert.AreEqual(1, msg.DeliveryCount);
                 var end = DateTime.UtcNow;
                 Assert.NotNull(msg);
-                Assert.Less(end - start, TimeSpan.FromSeconds(5));
+                Assert.Less(end - start, TimeSpan.FromSeconds(10));
             }
         }
 

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

@@ -244,7 +244,7 @@ public async Task ReceiveMessagesWhenQueueEmpty()
                 Assert.ThrowsAsync<TaskCanceledException>(async () => await receiver.ReceiveMessagesAsync(1, cancellationToken: cancellationTokenSource.Token));
                 var stop = DateTime.UtcNow;
 
-                Assert.That(stop - start,  Is.EqualTo(TimeSpan.FromSeconds(3)).Within(TimeSpan.FromSeconds(3)));
+                Assert.That(stop - start,  Is.EqualTo(TimeSpan.FromSeconds(3)).Within(TimeSpan.FromSeconds(5)));
             }
         }
 
@@ -1014,7 +1014,7 @@ public async Task CancellingDoesNotBlockSubsequentReceives(bool prefetch)
                 Assert.AreEqual(1, msg.DeliveryCount);
                 var end = DateTime.UtcNow;
                 Assert.NotNull(msg);
-                Assert.Less(end - start, TimeSpan.FromSeconds(5));
+                Assert.Less(end - start, TimeSpan.FromSeconds(10));
             }
         }
     }