[Event Hubs] Set GeoDR on management link (Azure#48204)

The focus of these changes is to set the desired capability for GeoDR on the
management link, so that requests for Event Hub and partition metadata can take
it into account for offset formatting.  Also included are tweaks to the tests for
`EventPosition.Latest` to reduce flakininess due to timing issues during nightly
runs.

Author: jsquire

sdk/eventhub/Azure.Messaging.EventHubs/src/Amqp/AmqpConnectionScope.cs

@@ -603,6 +603,9 @@ protected virtual async Task<RequestResponseAmqpLink> CreateManagementLinkAsync(
                 var linkSettings = new AmqpLinkSettings { OperationTimeout = operationTimeout };
                 linkSettings.AddProperty(AmqpProperty.Timeout, (uint)linkTimeout.TotalMilliseconds);
 
+                linkSettings.DesiredCapabilities ??= new Multiple<AmqpSymbol>();
+                linkSettings.DesiredCapabilities.Add(AmqpProperty.GeoReplication);
+
                 link = new RequestResponseAmqpLink(AmqpManagement.LinkType, session, AmqpManagement.Address, linkSettings.Properties);
 
                 // Track the link before returning it, so that it can be managed with the scope.

sdk/eventhub/Azure.Messaging.EventHubs/tests/Consumer/EventHubConsumerClientLiveTests.cs

@@ -821,16 +821,44 @@ public async Task ConsumerCanReadFromLatest()
 
                     // Give the receiver a moment to ensure that it is established and then send events for it to read.
 
-                    await Task.Delay(250);
-                    await SendEventsAsync(scope.EventHubName, EventGenerator.CreateEvents(50), new CreateBatchOptions { PartitionId = partition }, cancellationSource.Token);
+                    var sendTask = Task.Run(async () =>
+                    {
+                        try
+                        {
+                            while (!cancellationSource.IsCancellationRequested)
+                            {
+                                await Task.Delay(150);
+                                await SendEventsAsync(scope.EventHubName, EventGenerator.CreateEvents(5), new CreateBatchOptions { PartitionId = partition }, cancellationSource.Token);
+                            }
+                        }
+                        catch (TaskCanceledException)
+                        {
+                            // Expected
+                        }
+                    });
+
+                    while ((eventsRead < 1) && (!cancellationSource.IsCancellationRequested))
+                    {
+                        try
+                        {
+                            await Task.Delay(250, cancellationSource.Token);
+                        }
+                        catch (TaskCanceledException)
+                        {
+                            // Expected
+                        }
+                    }
 
                     // Await reading of the events and validate that we were able to read at least one event.
 
                     Assert.That(cancellationSource.IsCancellationRequested, Is.False, "The cancellation token should not have been signaled.");
                     Assert.That(eventsRead, Is.GreaterThanOrEqualTo(1), "At least one event should have been read.");
-                }
 
-                cancellationSource.Cancel();
+                    cancellationSource.Cancel();
+
+                    await readTask;
+                    await sendTask;
+                }
             }
         }
 
@@ -2402,44 +2430,83 @@ public async Task ConsumerCanReadFromAllPartitionsStartingWithLatest()
                 using var cancellationSource = new CancellationTokenSource();
                 cancellationSource.CancelAfter(EventHubsTestEnvironment.Instance.TestExecutionTimeLimit);
 
-                var sourceEvents = EventGenerator.CreateEvents(50).ToList();
-
-                await using (var consumer = new EventHubConsumerClient(
+                await using var consumer = new EventHubConsumerClient(
                     EventHubConsumerClient.DefaultConsumerGroupName,
                     EventHubsTestEnvironment.Instance.FullyQualifiedNamespace,
                     scope.EventHubName,
-                    EventHubsTestEnvironment.Instance.Credential))
-                {
-                    // Send a set of seed events, which should not be read.
+                    EventHubsTestEnvironment.Instance.Credential);
 
-                    var partitions = (await consumer.GetPartitionIdsAsync(cancellationSource.Token)).ToArray();
-                    await SendEventsToAllPartitionsAsync(scope.EventHubName, EventGenerator.CreateEvents(50), partitions, cancellationSource.Token);
+                // Send a set of seed events to the partition, which should not be read.
 
-                    // Begin reading though no events have been published.  This is necessary to open the connection and
-                    // ensure that the consumer is watching the partition.
+                var partitions = await consumer.GetPartitionIdsAsync(cancellationSource.Token);
+                await SendEventsToAllPartitionsAsync(scope.EventHubName, EventGenerator.CreateEvents(50), partitions, cancellationSource.Token);
 
-                    var readTask = ReadEventsFromAllPartitionsAsync(consumer, sourceEvents.Select(evt => evt.MessageId), cancellationSource.Token, startFromEarliest: false);
+                // Begin reading in the background, though no events should be read until the next publish.  This is necessary to open the connection and
+                // ensure that the receiver is watching the partition.
 
-                    // Give the consumer a moment to ensure that it is established and then send events for it to read.
+                var eventsRead = 0;
 
-                    await Task.Delay(1500);
-                    await SendEventsToAllPartitionsAsync(scope.EventHubName, sourceEvents, partitions, cancellationSource.Token);
+                var readTask = Task.Run(async () =>
+                {
+                    await Task.Yield();
 
-                    // Read the events and validate the resulting state.
+                    try
+                    {
+                       await foreach (var item in consumer.ReadEventsAsync(startReadingAtEarliestEvent: false, cancellationToken: cancellationSource.Token))
+                        {
+                            // If more than one event was read, no need to keep going.
 
-                    var readState = await readTask;
-                    Assert.That(cancellationSource.IsCancellationRequested, Is.False, "The cancellation token should not have been signaled.");
-                    Assert.That(readState.Events.Count, Is.EqualTo(sourceEvents.Count), "Only the source events should have been read.");
+                            if (++eventsRead > 1)
+                            {
+                                break;
+                            }
+                        }
+                    }
+                    catch (TaskCanceledException)
+                    {
+                        // Expected
+                    }
+                });
 
-                    foreach (var sourceEvent in sourceEvents)
+                // Give the receiver a moment to ensure that it is established and then send events for it to read.
+
+                var sendTask = Task.Run(async () =>
+                {
+                    try
                     {
-                        var sourceId = sourceEvent.MessageId;
-                        Assert.That(readState.Events.TryGetValue(sourceId, out var readEvent), Is.True, $"The event with custom identifier [{ sourceId }] was not processed.");
-                        Assert.That(sourceEvent.IsEquivalentTo(readEvent.Data), $"The event with custom identifier [{ sourceId }] did not match the corresponding processed event.");
+                        while (!cancellationSource.IsCancellationRequested)
+                        {
+                            await Task.Delay(150);
+                            await SendEventsToAllPartitionsAsync(scope.EventHubName, EventGenerator.CreateEvents(5), partitions, cancellationSource.Token);
+                        }
+                    }
+                    catch (TaskCanceledException)
+                    {
+                        // Expected
+                    }
+                });
+
+                while ((eventsRead < 1) && (!cancellationSource.IsCancellationRequested))
+                {
+                    try
+                    {
+                        await Task.Delay(250, cancellationSource.Token);
+                    }
+                    catch (TaskCanceledException)
+                    {
+                        // Expected
                     }
                 }
 
+                // Await reading of the events and validate that we were able to read at least one event.
+
+                Assert.That(cancellationSource.IsCancellationRequested, Is.False, "The cancellation token should not have been signaled.");
+                Assert.That(eventsRead, Is.GreaterThanOrEqualTo(1), "At least one event should have been read.");
+
                 cancellationSource.Cancel();
+
+                await readTask;
+                await sendTask;
             }
         }
 

sdk/eventhub/Azure.Messaging.EventHubs/tests/Primitives/PartitionReceiverLiveTests.cs

@@ -622,16 +622,44 @@ public async Task ReceiverCanReadFromLatest()
 
                     // Give the receiver a moment to ensure that it is established and then send events for it to read.
 
-                    await Task.Delay(250);
-                    await SendEventsAsync(scope.EventHubName, EventGenerator.CreateEvents(50), new CreateBatchOptions { PartitionId = partition }, cancellationSource.Token);
+                    var sendTask = Task.Run(async () =>
+                    {
+                        try
+                        {
+                            while (!cancellationSource.IsCancellationRequested)
+                            {
+                                await Task.Delay(150);
+                                await SendEventsAsync(scope.EventHubName, EventGenerator.CreateEvents(5), new CreateBatchOptions { PartitionId = partition }, cancellationSource.Token);
+                            }
+                        }
+                        catch (TaskCanceledException)
+                        {
+                            // Expected
+                        }
+                    });
+
+                    while ((eventsRead < 1) && (!cancellationSource.IsCancellationRequested))
+                    {
+                        try
+                        {
+                            await Task.Delay(250, cancellationSource.Token);
+                        }
+                        catch (TaskCanceledException)
+                        {
+                            // Expected
+                        }
+                    }
 
                     // Await reading of the events and validate that we were able to read at least one event.
 
                     Assert.That(cancellationSource.IsCancellationRequested, Is.False, "The cancellation token should not have been signaled.");
                     Assert.That(eventsRead, Is.GreaterThanOrEqualTo(1), "At least one event should have been read.");
-                }
 
-                cancellationSource.Cancel();
+                    cancellationSource.Cancel();
+
+                    await readTask;
+                    await sendTask;
+                }
             }
         }