diff --git a/sdk/eventhubs/azure_messaging_eventhubs/tests/eventhubs_consumer_groups.rs b/sdk/eventhubs/azure_messaging_eventhubs/tests/eventhubs_consumer_groups.rs new file mode 100644 index 00000000000..052e121a0e9 --- /dev/null +++ b/sdk/eventhubs/azure_messaging_eventhubs/tests/eventhubs_consumer_groups.rs @@ -0,0 +1,363 @@ +// Copyright (c) Microsoft Corporation. All Rights reserved +// Licensed under the MIT license. + +//! Live tests for Event Hubs consumer groups. +//! +//! Every event carries a marker for the run in its application properties. A +//! reader keeps only the events that have the marker of its own run, so +//! foreign traffic on the shared namespace does not change a result. +//! +//! `open_receiver_on_partition` does no network I/O. The broker attaches the +//! link on the first poll of `stream_events()`. An attach error therefore +//! arrives as the first item of the stream. + +use azure_core_amqp::{error::AmqpErrorKind, AmqpDescribedError, AmqpError, AmqpErrorCondition}; +use azure_core_test::{recorded, TestContext}; +use azure_messaging_eventhubs::{ + error::ErrorKind, + models::{AmqpSimpleValue, EventData, ReceivedEventData}, + ConsumerClient, EventDataBatchOptions, EventHubsError, OpenReceiverOptions, ProducerClient, + Result, StartLocation, StartPosition, +}; +use futures::{Stream, StreamExt}; +use std::{sync::LazyLock, time::Duration}; +use tokio::time::timeout; + +const DEFAULT_GROUP: &str = "$Default"; +const SECOND_GROUP: &str = "defaultGroup"; +const MISSING_GROUP: &str = "ThisIsFake"; +const MISSING_PARTITION: &str = "-1"; +const PARTITION: &str = "0"; +const EVENT_COUNT: usize = 5; +const RUN_MARKER_KEY: &str = "consumer-group-run"; +const EVENT_LABEL_KEY: &str = "consumer-group-label"; + +/// The time each consumer group gets to deliver every event of the run. +const READ_DEADLINE: Duration = Duration::from_secs(60); + +/// The time the broker gets to reject one attach. The attach is not retried. +const ATTACH_ERROR_DEADLINE: Duration = Duration::from_secs(30); + +// The three tests share partition 0, and cargo runs them on parallel threads. +static PARTITION_LOCK: LazyLock> = + LazyLock::new(|| async_lock::Mutex::new(())); + +/// Finds the AMQP described error that the broker sent, either in this error +/// or in its source chain. +fn described_error(amqp: &AmqpError) -> Option<&AmqpDescribedError> { + if let AmqpErrorKind::AmqpDescribedError(described) = amqp.kind() { + return Some(described); + } + let mut source = std::error::Error::source(amqp); + for _ in 0..16 { + let current = source?; + let inner = current + .downcast_ref::() + .or_else(|| current.downcast_ref::>().map(|e| e.as_ref())); + if let Some(inner) = inner { + if let AmqpErrorKind::AmqpDescribedError(described) = inner.kind() { + return Some(described); + } + } + source = current.source(); + } + None +} + +/// Takes exactly one item from the stream and returns the error that it holds. +async fn first_stream_error(stream: &mut S, limit: Duration, context: &str) -> EventHubsError +where + S: Stream> + Unpin, +{ + match timeout(limit, stream.next()).await { + Err(_) => panic!( + "{context}: no stream item arrived within {limit:?}. The attach must fail, not hang." + ), + Ok(None) => panic!("{context}: the stream ended with no item."), + Ok(Some(Ok(event))) => panic!( + "{context}: expected an attach error, got an event at sequence number {:?}", + event.sequence_number() + ), + Ok(Some(Err(err))) => err, + } +} + +/// Makes sure that the attach failed with one of the expected conditions. +/// +/// The expected sets come from the .NET mapping in +/// `Azure.Messaging.EventHubs/src/Amqp/AmqpError.cs`. They are not verified +/// against this broker. `authorizer.authorize_path` runs at CBS put-token time +/// before the attach, and that path carries the consumer group segment, so a +/// rejection there can carry a different condition. Step 4 names the observed +/// condition, so the first live run can correct the set. +fn assert_attach_condition(err: &EventHubsError, expected: &[AmqpErrorCondition], context: &str) { + assert!( + !matches!(err.kind, ErrorKind::ConsumerDisconnected(_)), + "{context}: the error says that the broker stole the link. Expected one of {expected:?}." + ); + let ErrorKind::AmqpError(amqp) = &err.kind else { + panic!( + "{context}: expected ErrorKind::AmqpError, got {:?}", + err.kind + ); + }; + let Some(described) = described_error(amqp) else { + panic!("{context}: no AMQP described error in the source chain: {amqp:?}"); + }; + assert!( + expected.contains(&described.condition), + "{context}: observed condition {:?} with description {:?}. Expected one of {expected:?}.", + described.condition, + described.description + ); +} + +/// Sends one batch to `partition`, with one event for each label. +async fn send_labeled_batch( + producer: &ProducerClient, + partition: &str, + run_marker: &str, + labels: &[String], +) -> Result<()> { + let batch = producer + .create_batch(Some(EventDataBatchOptions { + partition_id: Some(partition.to_string()), + ..Default::default() + })) + .await?; + for label in labels { + let event = EventData::builder() + .with_body(label.clone()) + .add_property(RUN_MARKER_KEY.to_string(), run_marker.to_string()) + .add_property(EVENT_LABEL_KEY.to_string(), label.clone()) + .build(); + assert!( + batch.try_add_event_data(event, None)?, + "run {run_marker}: the batch had no room for {label}" + ); + } + producer.send_batch(batch, None).await +} + +/// Returns the label of an event that carries the marker of this run. +fn tagged_label(event: &ReceivedEventData, run_marker: &str) -> Option { + let properties = event.event_data().properties()?; + match properties.get(RUN_MARKER_KEY) { + Some(AmqpSimpleValue::String(marker)) if marker == run_marker => {} + _ => return None, + } + match properties.get(EVENT_LABEL_KEY) { + Some(AmqpSimpleValue::String(label)) => Some(label.clone()), + _ => None, + } +} + +/// Reads the events of this run, and drops the events of every other run. +async fn read_run_events( + stream: &mut S, + run_marker: &str, + want: usize, + deadline: Duration, +) -> Vec<(String, i64)> +where + S: Stream> + Unpin, +{ + let mut found: Vec<(String, i64)> = Vec::with_capacity(want); + let _ = timeout(deadline, async { + while let Some(item) = stream.next().await { + let event = match item { + Ok(event) => event, + Err(err) => panic!("run {run_marker}: the stream failed while reading: {err:?}"), + }; + let Some(label) = tagged_label(&event, run_marker) else { + continue; + }; + let sequence = event.sequence_number().unwrap_or_else(|| { + panic!("run {run_marker}: the event for {label} has no sequence number") + }); + found.push((label, sequence)); + if found.len() == want { + break; + } + } + }) + .await; + found +} + +#[recorded::test(live)] +async fn two_consumer_groups_read_the_same_events(ctx: TestContext) -> Result<()> { + let _guard = PARTITION_LOCK.lock().await; + + let recording = ctx.recording(); + let host = recording.var("EVENTHUBS_HOST", None); + let eventhub = recording.var("EVENTHUB_NAME", None); + let credential = recording.credential(); + + let producer = ProducerClient::builder() + .open(host.as_str(), eventhub.as_str(), credential.clone()) + .await?; + + let run_marker = format!("consumer-groups-{}", azure_core::Uuid::new_v4()); + let start_sequence = producer + .get_partition_properties(PARTITION) + .await? + .last_enqueued_sequence_number; + + let expected: Vec = (0..EVENT_COUNT).map(|i| format!("event-{i}")).collect(); + send_labeled_batch(&producer, PARTITION, &run_marker, &expected).await?; + + let default_consumer = ConsumerClient::builder() + .with_consumer_group(DEFAULT_GROUP.to_string()) + .open(host.as_str(), eventhub.clone(), credential.clone()) + .await?; + let second_consumer = ConsumerClient::builder() + .with_consumer_group(SECOND_GROUP.to_string()) + .open(host.as_str(), eventhub.clone(), credential.clone()) + .await?; + + // Do not set owner_level. An epoch consumer displaces the other reader. + let start_position = StartPosition { + location: StartLocation::SequenceNumber(start_sequence), + inclusive: false, + }; + let default_receiver = default_consumer + .open_receiver_on_partition( + PARTITION.to_string(), + Some(OpenReceiverOptions { + start_position: Some(start_position.clone()), + ..Default::default() + }), + ) + .await?; + let second_receiver = second_consumer + .open_receiver_on_partition( + PARTITION.to_string(), + Some(OpenReceiverOptions { + start_position: Some(start_position), + ..Default::default() + }), + ) + .await?; + + let default_read = { + let mut stream = default_receiver.stream_events(); + read_run_events(&mut stream, &run_marker, EVENT_COUNT, READ_DEADLINE).await + }; + let second_read = { + let mut stream = second_receiver.stream_events(); + read_run_events(&mut stream, &run_marker, EVENT_COUNT, READ_DEADLINE).await + }; + + default_receiver.close().await?; + second_receiver.close().await?; + default_consumer.close().await?; + second_consumer.close().await?; + producer.close().await?; + + let (labels_default, sequences_default): (Vec, Vec) = + default_read.into_iter().unzip(); + let (labels_second, sequences_second): (Vec, Vec) = + second_read.into_iter().unzip(); + + assert_eq!( + labels_default, expected, + "run {run_marker}: consumer group {DEFAULT_GROUP} on partition {PARTITION} did not read every event of this run within {READ_DEADLINE:?}" + ); + assert_eq!( + labels_second, expected, + "run {run_marker}: consumer group {SECOND_GROUP} on partition {PARTITION} did not read every event of this run within {READ_DEADLINE:?}" + ); + assert_eq!( + sequences_default, sequences_second, + "run {run_marker}: a sequence number is a property of the partition, so the two consumer groups must report the same sequence numbers" + ); + + Ok(()) +} + +#[recorded::test(live)] +async fn receiver_on_unknown_consumer_group_fails(ctx: TestContext) -> Result<()> { + const CONTEXT: &str = "receiver on an unknown consumer group"; + let _guard = PARTITION_LOCK.lock().await; + + let recording = ctx.recording(); + let host = recording.var("EVENTHUBS_HOST", None); + let eventhub = recording.var("EVENTHUB_NAME", None); + + let consumer = ConsumerClient::builder() + .with_consumer_group(MISSING_GROUP.to_string()) + .open(host.as_str(), eventhub, recording.credential()) + .await?; + // Earliest makes the test fail fast. If the consumer group segment stopped + // reaching the wire, the receiver attaches on $Default and gets an event. + let receiver = consumer + .open_receiver_on_partition( + PARTITION.to_string(), + Some(OpenReceiverOptions { + start_position: Some(StartPosition { + location: StartLocation::Earliest, + ..Default::default() + }), + ..Default::default() + }), + ) + .await?; + + let observed = { + let mut stream = receiver.stream_events(); + first_stream_error(&mut stream, ATTACH_ERROR_DEADLINE, CONTEXT).await + }; + + let _ = receiver.close().await; + let _ = consumer.close().await; + + assert_attach_condition(&observed, &[AmqpErrorCondition::NotFound], CONTEXT); + + Ok(()) +} + +#[recorded::test(live)] +async fn receiver_on_unknown_partition_fails(ctx: TestContext) -> Result<()> { + const CONTEXT: &str = "receiver on an unknown partition"; + let _guard = PARTITION_LOCK.lock().await; + + let recording = ctx.recording(); + let host = recording.var("EVENTHUBS_HOST", None); + let eventhub = recording.var("EVENTHUB_NAME", None); + + let consumer = ConsumerClient::builder() + .open(host.as_str(), eventhub, recording.credential()) + .await?; + let receiver = consumer + .open_receiver_on_partition( + MISSING_PARTITION.to_string(), + Some(OpenReceiverOptions { + start_position: Some(StartPosition { + location: StartLocation::Earliest, + ..Default::default() + }), + ..Default::default() + }), + ) + .await?; + + let observed = { + let mut stream = receiver.stream_events(); + first_stream_error(&mut stream, ATTACH_ERROR_DEADLINE, CONTEXT).await + }; + + let _ = receiver.close().await; + let _ = consumer.close().await; + + assert_attach_condition( + &observed, + &[ + AmqpErrorCondition::ArgumentOutOfRangeError, + AmqpErrorCondition::NotFound, + ], + CONTEXT, + ); + + Ok(()) +}