diff --git a/sdk/eventhubs/azure_messaging_eventhubs/CHANGELOG.md b/sdk/eventhubs/azure_messaging_eventhubs/CHANGELOG.md index 8575377c254..1347078012d 100644 --- a/sdk/eventhubs/azure_messaging_eventhubs/CHANGELOG.md +++ b/sdk/eventhubs/azure_messaging_eventhubs/CHANGELOG.md @@ -4,6 +4,8 @@ ### Features Added +- Added `BufferedProducerClient`, a producer that accepts single events and publishes them in the background. One worker owns each partition, so events keep their enqueued order and a slow partition does not stop the others. A successful enqueue means only that the local buffer accepted the event; the client reports each terminal outcome through the handlers that `with_on_send_succeeded` and `with_on_send_failed` register, and a handler for failed batches is required. The client sends a batch when the next event does not fit, when the batch reaches the configured event count, when the maximum wait time expires, on `flush`, and on `close`. The defaults are a 1 second maximum wait time and 1500 buffered events for each partition, matching the .NET, JavaScript, and Python clients. `flush` sets a barrier over the events accepted before the call, `close` sends the buffered events, and `abort` abandons them. New public types: `BufferedProducerClient`, `BufferedProducerClientBuilder`, `EnqueueEventOptions`, `SendBatchSucceededContext`, and `SendBatchFailedContext`. ([#4873](https://github.com/Azure/azure-sdk-for-rust/pull/4873)) +- Added the `ErrorKind::SendNotAccepted` error variant. The buffered producer reports an AMQP `Modified` or `Released` outcome as a delivery failure, because neither outcome means that the service durably stored the events. `ProducerClient::send_batch` keeps its historical behavior and treats both as success with a warning. - Added connection-string authentication. `ProducerClientBuilder` and `ConsumerClientBuilder` now have an `open_with_connection_string` method that authenticates with a Shared Access Signature parsed from an Event Hubs connection string (`Endpoint=sb://...;SharedAccessKeyName=...;SharedAccessKey=...`, optionally with `EntityPath`, or a pre-formed `SharedAccessSignature`). The connection-string parser is exposed publicly as `ConnectionString`. This reaches parity with the other Azure SDKs for development and test scenarios; Microsoft Entra ID via `open` with a `TokenCredential` remains the recommended path for production. The parser rejects empty required values and empty Event Hub names up front, and a pre-formed `SharedAccessSignature` reports its own `se` as the token expiry (rather than a rolling client-side window); because such a token cannot be renewed, the connection's token refresher detects the non-advancing expiry and leaves the broker to enforce it. ([#3459](https://github.com/Azure/azure-sdk-for-rust/issues/3459)) - The `EventProcessor` now opens every partition receiver with AMQP epoch (owner level) `0` and surfaces broker-initiated displacement as the new `EventHubsError::ConsumerDisconnected` error kind. When a second `EventProcessor` instance claims a partition this instance is currently holding, the broker disconnects this instance's receiver and the consumer's `stream_events()` resolves with `ConsumerDisconnected`. This matches the behavior of `EventProcessorClient` in the .NET and Java Azure SDKs. Consumers should pattern-match on `ErrorKind::ConsumerDisconnected` to detect a stolen partition and re-acquire a client via `next_partition_client()`. - Added `EventHubsError::ConsumerDisconnected(Option)` error variant. diff --git a/sdk/eventhubs/azure_messaging_eventhubs/README.md b/sdk/eventhubs/azure_messaging_eventhubs/README.md index 4b27e1ca7f5..fda00696e5e 100644 --- a/sdk/eventhubs/azure_messaging_eventhubs/README.md +++ b/sdk/eventhubs/azure_messaging_eventhubs/README.md @@ -120,6 +120,10 @@ Additional examples for various scenarios can be found on in the examples direct - [Send events](#send-events) - [Send events directly to the Event Hub](#send-events-directly-to-the-event-hub) - [Send events using a batch operation](#send-events-using-a-batch-operation) +- [Send events with the buffered producer](#send-events-with-the-buffered-producer) + - [Route events to a partition](#route-events-to-a-partition) + - [Flush and shut down](#flush-and-shut-down) + - [Trade-offs of buffered publishing](#trade-offs-of-buffered-publishing) - [Open an Event Hubs message consumer on an Event Hubs instance](#open-an-event-hubs-message-consumer-on-an-event-hub-instance) - [Receive events](#receive-events) @@ -178,6 +182,125 @@ async fn send_events(producer: &ProducerClient) -> Result<(), Box Result<(), Box> { + let host = ""; + let eventhub = ""; + let credential = DeveloperToolsCredential::new(None)?; + + let producer = BufferedProducerClient::builder() + .with_on_send_succeeded(|context| async move { + println!( + "The service accepted {} events on partition {}.", + context.events.len(), + context.partition_id + ); + }) + .with_on_send_failed(|context| async move { + eprintln!( + "{} events failed on partition {}: {}", + context.events.len(), + context.partition_id, + context.error + ); + }) + .open(host, eventhub, credential.clone()) + .await?; + + for index in 0..1000 { + producer.enqueue_event(format!("event {index}"), None).await?; + } + + producer.close().await?; + Ok(()) +} +``` + +#### Route events to a partition + +Give a partition ID to send an event to one partition. Give a partition key to send every event +with that key to the same partition. Set at most one of the two; the client rejects a request that +sets both. When you set neither, the client assigns the partitions in round-robin order. + +```rust no_run +use azure_messaging_eventhubs::{BufferedProducerClient, EnqueueEventOptions}; + +async fn route_events( + producer: &BufferedProducerClient, +) -> Result<(), Box> { + producer + .enqueue_event( + "to partition 0", + Some(EnqueueEventOptions { + partition_id: Some("0".to_string()), + ..Default::default() + }), + ) + .await?; + + producer + .enqueue_event( + "grouped by key", + Some(EnqueueEventOptions { + partition_key: Some("customer-17".to_string()), + ..Default::default() + }), + ) + .await?; + + Ok(()) +} +``` + +#### Flush and shut down + +`flush` sets a barrier. It completes once every event that the client accepted before the call +reaches a terminal outcome. An event that arrives after the barrier does not delay the call. + +`close` sends the buffered events and then shuts the client down. `abort` shuts the client down at +once and abandons the buffered events. + +```rust no_run +use azure_messaging_eventhubs::BufferedProducerClient; + +async fn flush_and_close( + producer: &BufferedProducerClient, +) -> Result<(), Box> { + producer.enqueue_event("an event", None).await?; + + // Wait for the events that the client already accepted. + producer.flush().await?; + println!("{} events are still buffered.", producer.total_buffered_event_count()); + + // Send what is left, then shut down. + producer.close().await?; + Ok(()) +} +``` + +#### Trade-offs of buffered publishing + +- A successful enqueue means only that the local buffer accepted the event. It does not mean that + Event Hubs accepted the event. +- The process loses the buffered events if it stops before a flush or a close. Call `flush` or + `close` when the delivery of the buffered events matters. +- A send failure arrives after the enqueue call already returned, through the failure handler. +- Buffering gives a higher throughput, but the latency of one event is less predictable. +- Use `ProducerClient` when the application needs the result of each send. + ### Open an Event Hubs message consumer on an Event Hub instance ```rust no_run diff --git a/sdk/eventhubs/azure_messaging_eventhubs/examples/eventhubs_buffered_produce_events.rs b/sdk/eventhubs/azure_messaging_eventhubs/examples/eventhubs_buffered_produce_events.rs new file mode 100644 index 00000000000..670092ce895 --- /dev/null +++ b/sdk/eventhubs/azure_messaging_eventhubs/examples/eventhubs_buffered_produce_events.rs @@ -0,0 +1,124 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT license. + +// cspell: ignore retryable + +//! This sample shows how to publish events with the buffered producer client. +//! +//! The caller enqueues single events. The client buffers them, groups them into +//! batches for each partition, and publishes them in the background. +//! +//! A successful enqueue means only that the local buffer accepted the event. It +//! does not mean that Event Hubs accepted the event. The client reports the real +//! outcome through the two handlers below. + +use azure_core::time::Duration; +use azure_identity::DeveloperToolsCredential; +use azure_messaging_eventhubs::{BufferedProducerClient, EnqueueEventOptions}; +use std::sync::{ + atomic::{AtomicUsize, Ordering}, + Arc, +}; + +#[tokio::main] +async fn main() -> Result<(), Box> { + let eventhub_namespace = std::env::var("EVENTHUBS_HOST")?; + let eventhub_name = std::env::var("EVENTHUB_NAME")?; + let credential = DeveloperToolsCredential::new(None)?; + + // The handlers report the outcome of each batch. Count the events so the + // sample can print a summary at the end. + let accepted = Arc::new(AtomicUsize::new(0)); + let rejected = Arc::new(AtomicUsize::new(0)); + + let for_success = accepted.clone(); + let for_failure = rejected.clone(); + + let producer = BufferedProducerClient::builder() + // Send a batch that is not full after this time. + .with_max_wait_time(Duration::seconds(1)) + // Wait for space once a partition buffer holds this many events. + .with_max_buffered_event_count_per_partition(1500) + .with_on_send_succeeded(move |context| { + let accepted = for_success.clone(); + async move { + accepted.fetch_add(context.events.len(), Ordering::AcqRel); + println!( + "The service accepted {} events on partition {}.", + context.events.len(), + context.partition_id + ); + } + }) + // A handler for failed batches is required. The client calls it only + // after the retry policy is exhausted, or when the error is not + // retryable. The client does not enqueue the events again, so the + // application decides what to do with them. + .with_on_send_failed(move |context| { + let rejected = for_failure.clone(); + async move { + rejected.fetch_add(context.events.len(), Ordering::AcqRel); + eprintln!( + "{} events failed on partition {}: {}", + context.events.len(), + context.partition_id, + context.error + ); + } + }) + .open( + eventhub_namespace.as_str(), + eventhub_name.as_str(), + credential.clone(), + ) + .await?; + + // The client assigns these events to the partitions in round-robin order. + for index in 0..100 { + producer + .enqueue_event(format!("automatic event {index}"), None) + .await?; + } + + // These events all go to partition 0. + producer + .enqueue_events( + vec!["first", "second", "third"], + Some(EnqueueEventOptions { + partition_id: Some("0".to_string()), + ..Default::default() + }), + ) + .await?; + + // Every event with the same key goes to the same partition. + producer + .enqueue_event( + "an event for one customer", + Some(EnqueueEventOptions { + partition_key: Some("customer-17".to_string()), + ..Default::default() + }), + ) + .await?; + + println!( + "{} events are waiting in the buffer.", + producer.total_buffered_event_count() + ); + + // The flush completes once every event that the client accepted before this + // call reaches a terminal outcome. + producer.flush().await?; + + // A graceful close sends what is left, then it releases the connection. Use + // `abort` instead to shut down at once and abandon the buffered events. + producer.close().await?; + + println!( + "Done. The service accepted {} events and rejected {} events.", + accepted.load(Ordering::Acquire), + rejected.load(Ordering::Acquire) + ); + Ok(()) +} diff --git a/sdk/eventhubs/azure_messaging_eventhubs/src/error.rs b/sdk/eventhubs/azure_messaging_eventhubs/src/error.rs index 38c7f126592..1011b14f5cc 100644 --- a/sdk/eventhubs/azure_messaging_eventhubs/src/error.rs +++ b/sdk/eventhubs/azure_messaging_eventhubs/src/error.rs @@ -45,6 +45,17 @@ pub enum ErrorKind { /// AmqpError(AmqpError), + /// The service settled the transfer, but it did not durably accept it. + /// + /// The broker returned an AMQP `Modified` or `Released` outcome. Neither + /// outcome means the service stored the events. The buffered producer + /// reports this as a delivery failure. + /// + /// A `Released` or `Modified` outcome does not prove that the service + /// discarded the events either. If the caller sends the same events again, + /// the service can store them two times. + SendNotAccepted(Cow<'static, str>), + /// Receiver was disconnected by the broker because another receiver /// attached with the same or higher epoch (owner level). The inner /// `AmqpDescribedError` is for logging; match on the variant: @@ -95,6 +106,9 @@ impl std::fmt::Display for EventHubsError { ErrorKind::SendRejected(e) => write!(f, "Send rejected: {:?}", e), ErrorKind::InvalidManagementResponse => f.write_str("Invalid management response"), ErrorKind::AmqpError(source) => write!(f, "AMQP Error: {:?}", source), + ErrorKind::SendNotAccepted(msg) => { + write!(f, "Send was not durably accepted: {}", msg) + } ErrorKind::ConsumerDisconnected(e) => { write!( f, diff --git a/sdk/eventhubs/azure_messaging_eventhubs/src/lib.rs b/sdk/eventhubs/azure_messaging_eventhubs/src/lib.rs index 6949ec13b31..abee6aab8d6 100644 --- a/sdk/eventhubs/azure_messaging_eventhubs/src/lib.rs +++ b/sdk/eventhubs/azure_messaging_eventhubs/src/lib.rs @@ -23,6 +23,10 @@ pub use consumer::{ }; pub use producer::{ batch::{EventDataBatch, EventDataBatchOptions}, + buffered::{ + BufferedProducerClient, EnqueueEventOptions, SendBatchFailedContext, + SendBatchSucceededContext, + }, ProducerClient, SendBatchOptions, SendEventOptions, SendMessageOptions, }; @@ -36,6 +40,7 @@ pub use event_processor::{processor::EventProcessor, CheckpointStore, ProcessorS pub mod builders { pub use crate::consumer::builders::ConsumerClientBuilder; pub use crate::event_processor::processor::builders::EventProcessorBuilder; + pub use crate::producer::buffered::builders::BufferedProducerClientBuilder; pub use crate::producer::builders::ProducerClientBuilder; } pub use common::connection_string::ConnectionString; diff --git a/sdk/eventhubs/azure_messaging_eventhubs/src/producer/batch.rs b/sdk/eventhubs/azure_messaging_eventhubs/src/producer/batch.rs index a664749abfe..317cec635c9 100644 --- a/sdk/eventhubs/azure_messaging_eventhubs/src/producer/batch.rs +++ b/sdk/eventhubs/azure_messaging_eventhubs/src/producer/batch.rs @@ -11,10 +11,166 @@ use tracing::debug; /// Represents the options that can be set when adding event data to an [`EventDataBatch`]. pub struct AddEventDataOptions {} -struct EventDataBatchState { +/// The owned serialization core of a batch. +/// +/// This holds every piece of state needed to accumulate messages and to produce +/// the AMQP batch envelope. It borrows nothing, so a background task can own one +/// for the whole life of the task. [`EventDataBatch`] wraps one of these behind a +/// mutex and keeps the borrow of the [`ProducerClient`] for itself. +pub(crate) struct EventDataBatchInner { serialized_messages: Vec>, size_in_bytes: u64, batch_envelope: Option, + max_size_in_bytes: u64, + partition_key: Option, +} + +impl EventDataBatchInner { + pub(crate) fn new(max_size_in_bytes: u64, partition_key: Option) -> Self { + Self { + serialized_messages: Vec::new(), + size_in_bytes: 0, + batch_envelope: None, + max_size_in_bytes, + partition_key, + } + } + + pub(crate) fn size(&self) -> u64 { + self.size_in_bytes + } + + pub(crate) fn len(&self) -> usize { + self.serialized_messages.len() + } + + pub(crate) fn is_empty(&self) -> bool { + self.serialized_messages.is_empty() + } + + fn arithmetic_error() -> EventHubsError { + EventHubsError::with_message("Arithmetic error calculating Batch size.") + } + + fn calculate_actual_size_for_payload(length: usize) -> Result { + const MESSAGE_HEADER_SIZE_32: usize = 8; + const MESSAGE_HEADER_SIZE_8: usize = 5; + if length < 256 { + Ok(length + .checked_add(MESSAGE_HEADER_SIZE_8) + .ok_or_else(Self::arithmetic_error)? as u64) + } else { + Ok(length + .checked_add(MESSAGE_HEADER_SIZE_32) + .ok_or_else(Self::arithmetic_error)? as u64) + } + } + + /// Tries to add an AMQP message to the batch. + /// + /// Returns `true` when the message was added. Returns `false` when the + /// message does not fit; in that case the batch is left unchanged. + pub(crate) fn try_add(&mut self, message: impl Into) -> Result { + let mut message = message.into(); + if message.properties.is_none() || message.properties.as_ref().unwrap().message_id.is_none() + { + message.set_message_id(Uuid::new_v4()); + } + if let Some(partition_key) = self.partition_key.as_ref() { + message.add_message_annotation( + AmqpSymbol::from("x-opt-partition-key"), + partition_key.clone(), + ); + } + + let message_len = AmqpMessage::serialize(&message)?.len(); + if self.serialized_messages.is_empty() { + // The first message serialized is the batch envelope - we capture the parameters from the first message to use for the batch + self.size_in_bytes = self + .size_in_bytes + .checked_add(message_len as u64) + .ok_or_else(Self::arithmetic_error)?; + self.batch_envelope = Some(Self::create_batch_envelope(&message)); + } + let serialized_message = AmqpMessage::serialize(&message)?; + let actual_message_size = + Self::calculate_actual_size_for_payload(serialized_message.len())?; + if self + .size_in_bytes + .checked_add(actual_message_size) + .ok_or_else(Self::arithmetic_error)? + > self.max_size_in_bytes + { + debug!("Batch is full. Cannot add more messages."); + debug!("Message size: {actual_message_size}"); + debug!("Current batch size: {:?}", self.size_in_bytes); + debug!("Max batch size: {:?}", self.max_size_in_bytes); + if self.serialized_messages.is_empty() { + self.batch_envelope = None; + self.size_in_bytes = 0; + } + return Ok(false); + } + self.size_in_bytes += actual_message_size; + self.serialized_messages.push(serialized_message); + + Ok(true) + } + + /// Takes the accumulated messages as a single AMQP batch envelope and resets + /// the batch so that it can accumulate again. + /// + /// # Panics + /// + /// Panics when the batch is empty. Callers must not send an empty batch. + pub(crate) fn take_envelope(&mut self) -> AmqpMessage { + let mut batch_envelope = self.batch_envelope.clone().expect( + "Batch envelope is missing when getting messages; \ + send_batch was called on an empty batch (add at least one event before sending).", + ); + + // Move the messages out of the batch state into a local variable so we + // can subsequently move it to the message body. + let mut serialized_messages = Vec::>::new(); + serialized_messages.append(&mut self.serialized_messages); + + batch_envelope.set_message_body(serialized_messages); + + // Reset the batch state for the next batch + self.batch_envelope = None; + self.size_in_bytes = 0; + self.serialized_messages.clear(); + + batch_envelope + } + + fn create_batch_envelope(message: &AmqpMessage) -> AmqpMessage { + // Transfer all the message options from the original message to the batch envelope + // Do NOT transfer the body, that will be handled later. + let mut batch_builder = AmqpMessage::builder(); + + if let Some(message_header) = message.header.as_ref() { + batch_builder = batch_builder.with_header(message_header.clone()); + } + if let Some(message_properties) = message.properties.as_ref() { + batch_builder = batch_builder.with_properties(message_properties.clone()); + } + if let Some(application_properties) = message.application_properties.as_ref() { + batch_builder = + batch_builder.with_application_properties(application_properties.clone()); + } + if let Some(delivery_annotations) = message.delivery_annotations.as_ref() { + batch_builder = batch_builder.with_delivery_annotations(delivery_annotations.clone()); + } + if let Some(message_annotations) = message.message_annotations.as_ref() { + batch_builder = batch_builder.with_message_annotations(message_annotations.clone()); + } + if let Some(footer) = message.footer.as_ref() { + batch_builder = batch_builder.with_footer(footer.clone()); + } + + batch_builder.build() + } } /// Represents a collections of event data that can be sent to an Event Hubs instance in one operation. @@ -47,12 +203,7 @@ struct EventDataBatchState { /// ``` pub struct EventDataBatch<'a> { producer: &'a ProducerClient, - batch_state: Mutex, - /// The size that [`EventDataBatch::resolve_max_size_in_bytes`] decided. - /// [`EventDataBatch::try_add_amqp_message`] refuses a message that does not - /// fit under it. - max_size_in_bytes: u64, - partition_key: Option, + inner: Mutex, partition_id: Option, } @@ -68,15 +219,10 @@ impl<'a> EventDataBatch<'a> { options: Option, max_size_in_bytes: u64, ) -> Self { + let partition_key = options.as_ref().and_then(|o| o.partition_key.clone()); Self { producer, - batch_state: Mutex::new(EventDataBatchState { - serialized_messages: Vec::new(), - size_in_bytes: 0, - batch_envelope: None, - }), - max_size_in_bytes, - partition_key: options.as_ref().and_then(|o| o.partition_key.clone()), + inner: Mutex::new(EventDataBatchInner::new(max_size_in_bytes, partition_key)), partition_id: options.and_then(|o| o.partition_id), } } @@ -134,7 +280,7 @@ impl<'a> EventDataBatch<'a> { /// pub fn size(&self) -> u64 { // Note that lock() returns an infallible result. - self.batch_state.lock().unwrap().size_in_bytes + self.inner.lock().unwrap().size() } /// Gets the number of messages in the batch. @@ -144,7 +290,7 @@ impl<'a> EventDataBatch<'a> { /// The number of messages in the batch. /// pub fn len(&self) -> usize { - self.batch_state.lock().unwrap().serialized_messages.len() + self.inner.lock().unwrap().len() } /// Determines whether the batch is empty. @@ -153,25 +299,7 @@ impl<'a> EventDataBatch<'a> { /// `true` if the batch is empty; otherwise, `false`. /// pub fn is_empty(&self) -> bool { - self.len() == 0 - } - - fn arithmetic_error() -> EventHubsError { - EventHubsError::with_message("Arithmetic error calculating Batch size.") - } - - fn calculate_actual_size_for_payload(length: usize) -> Result { - const MESSAGE_HEADER_SIZE_32: usize = 8; - const MESSAGE_HEADER_SIZE_8: usize = 5; - if length < 256 { - Ok(length - .checked_add(MESSAGE_HEADER_SIZE_8) - .ok_or_else(Self::arithmetic_error)? as u64) - } else { - Ok(length - .checked_add(MESSAGE_HEADER_SIZE_32) - .ok_or_else(Self::arithmetic_error)? as u64) - } + self.inner.lock().unwrap().is_empty() } /// Tries to add an event data to the batch. @@ -258,107 +386,16 @@ impl<'a> EventDataBatch<'a> { message: impl Into, #[allow(unused_variables)] options: Option, ) -> Result { - let mut message = message.into(); - if message.properties.is_none() || message.properties.as_ref().unwrap().message_id.is_none() - { - message.set_message_id(Uuid::new_v4()); - } - if let Some(partition_key) = self.partition_key.as_ref() { - message.add_message_annotation( - AmqpSymbol::from("x-opt-partition-key"), - partition_key.clone(), - ); - } - - let mut batch_state = self.batch_state.lock().unwrap(); - let message_len = AmqpMessage::serialize(&message)?.len(); - if batch_state.serialized_messages.is_empty() { - // The first message serialized is the batch envelope - we capture the parameters from the first message to use for the batch - batch_state.size_in_bytes = batch_state - .size_in_bytes - .checked_add(message_len as u64) - .ok_or_else(Self::arithmetic_error)?; - batch_state.batch_envelope = Some(self.create_batch_envelope(&message)); - } - let serialized_message = AmqpMessage::serialize(&message)?; - let actual_message_size = - Self::calculate_actual_size_for_payload(serialized_message.len())?; - if batch_state - .size_in_bytes - .checked_add(actual_message_size) - .ok_or_else(Self::arithmetic_error)? - > self.max_size_in_bytes - { - debug!("Batch is full. Cannot add more messages."); - debug!("Message size: {actual_message_size}"); - debug!("Current batch size: {:?}", batch_state.size_in_bytes); - debug!("Max batch size: {:?}", self.max_size_in_bytes); - if batch_state.serialized_messages.is_empty() { - batch_state.batch_envelope = None; - batch_state.size_in_bytes = 0; - } - return Ok(false); - } - batch_state.size_in_bytes += actual_message_size; - batch_state.serialized_messages.push(serialized_message); - - Ok(true) + self.inner.lock().unwrap().try_add(message) } pub(crate) fn get_messages(&self) -> AmqpMessage { - let mut batch_state = self.batch_state.lock().unwrap(); - - let mut batch_envelope = batch_state.batch_envelope.clone().expect( - "Batch envelope is missing when getting messages; \ - send_batch was called on an empty batch (add at least one event before sending).", - ); - - // Move the messages out of the batch state into a local variable so we - // can subsequently move it to the message body. - let mut serialized_messages = Vec::>::new(); - serialized_messages.append(&mut batch_state.serialized_messages); - - batch_envelope.set_message_body(serialized_messages); - - // Reset the batch state for the next batch - batch_state.batch_envelope = None; - batch_state.size_in_bytes = 0; - batch_state.serialized_messages.clear(); - - batch_envelope + self.inner.lock().unwrap().take_envelope() } pub(crate) fn get_batch_path(&self) -> Result { Self::batch_path(self.producer.base_url(), self.partition_id.as_deref()) } - - fn create_batch_envelope(&self, message: &AmqpMessage) -> AmqpMessage { - // Transfer all the message options from the original message to the batch envelope - // Do NOT transfer the body, that will be handled later. - let mut batch_builder = AmqpMessage::builder(); - - if let Some(message_header) = message.header.as_ref() { - batch_builder = batch_builder.with_header(message_header.clone()); - } - if let Some(message_properties) = message.properties.as_ref() { - batch_builder = batch_builder.with_properties(message_properties.clone()); - } - if let Some(application_properties) = message.application_properties.as_ref() { - batch_builder = - batch_builder.with_application_properties(application_properties.clone()); - } - if let Some(delivery_annotations) = message.delivery_annotations.as_ref() { - batch_builder = batch_builder.with_delivery_annotations(delivery_annotations.clone()); - } - if let Some(message_annotations) = message.message_annotations.as_ref() { - batch_builder = batch_builder.with_message_annotations(message_annotations.clone()); - } - if let Some(footer) = message.footer.as_ref() { - batch_builder = batch_builder.with_footer(footer.clone()); - } - - batch_builder.build() - } } /// Represents the options that can be set when creating an [`EventDataBatch`]. @@ -417,6 +454,11 @@ mod tests { } } + // The cap the batch actually enforces, read out of the serialization core. + fn effective_max_size(batch: &EventDataBatch<'_>) -> u64 { + batch.inner.lock().unwrap().max_size_in_bytes + } + // A client that never opens a connection. `try_add_event_data` only // serializes and measures, so a batch can be driven without a broker. fn offline_producer() -> ProducerClient { @@ -533,7 +575,7 @@ mod tests { EventDataBatch::resolve_max_size_in_bytes(Some(&options), LINK_MAX_SIZE) .expect("1024 bytes is below the link maximum"); let batch = EventDataBatch::new(&producer, Some(options), max_size_in_bytes); - assert_eq!(batch.max_size_in_bytes, MAX_SIZE); + assert_eq!(effective_max_size(&batch), MAX_SIZE); let body = "x".repeat(128); let mut accepted = 0; @@ -580,7 +622,7 @@ mod tests { .expect("the link maximum is always allowed"); let batch = EventDataBatch::new(&producer, None, max_size_in_bytes); - assert_eq!(batch.max_size_in_bytes, LINK_MAX_SIZE); + assert_eq!(effective_max_size(&batch), LINK_MAX_SIZE); assert!(batch .try_add_event_data("x".repeat(128), None) .expect("adding an event of a known size cannot fail")); diff --git a/sdk/eventhubs/azure_messaging_eventhubs/src/producer/buffered/mod.rs b/sdk/eventhubs/azure_messaging_eventhubs/src/producer/buffered/mod.rs new file mode 100644 index 00000000000..e185a09347b --- /dev/null +++ b/sdk/eventhubs/azure_messaging_eventhubs/src/producer/buffered/mod.rs @@ -0,0 +1,1066 @@ +// Copyright (c) Microsoft Corporation. All Rights reserved +// Licensed under the MIT license. + +// cspell: ignore retryable + +//! A producer client that buffers events and publishes them in the background. + +pub(crate) mod partition_resolver; +pub(crate) mod send_client; +pub(crate) mod worker; + +use crate::{ + error::Result, + models::EventData, + producer::{ + buffered::{ + partition_resolver::PartitionResolver, + send_client::BufferedSendClient, + worker::{Command, PartitionWorker}, + }, + ProducerClient, + }, + EventHubsError, +}; +use async_lock::Semaphore; +use azure_core::{ + async_runtime::{get_async_runtime, SpawnedTask}, + fmt::SafeDebug, + time::Duration, + Uuid, +}; +use azure_core_amqp::{AmqpMessage, AmqpSymbol}; +use futures::{ + channel::{mpsc, oneshot}, + future::{BoxFuture, Shared}, + FutureExt, +}; +use std::{ + collections::HashMap, + sync::{ + atomic::{AtomicBool, AtomicUsize, Ordering}, + Arc, Mutex, + }, +}; +use tracing::{debug, trace, warn}; + +/// The default maximum time that the client waits before it sends a batch that +/// is not full. +/// +/// This matches the .NET, JavaScript, and Python clients. +pub(crate) const DEFAULT_MAX_WAIT_TIME_SECONDS: i64 = 1; + +/// The default maximum number of events that the client buffers for one partition. +/// +/// This matches the .NET, Java, JavaScript, and Python clients. +pub(crate) const DEFAULT_MAX_BUFFERED_EVENT_COUNT_PER_PARTITION: usize = 1500; + +/// Options for [`BufferedProducerClient::enqueue_event`] and +/// [`BufferedProducerClient::enqueue_events`]. +/// +/// Set at most one of `partition_id` and `partition_key`. The client rejects a +/// request that sets both. When neither is set, the client assigns a partition +/// in round-robin order. +/// +/// [`BufferedProducerClient::enqueue_event`]: crate::BufferedProducerClient::enqueue_event +/// [`BufferedProducerClient::enqueue_events`]: crate::BufferedProducerClient::enqueue_events +/// +/// # Examples +/// +/// ``` +/// use azure_messaging_eventhubs::EnqueueEventOptions; +/// +/// let to_partition = EnqueueEventOptions { +/// partition_id: Some("0".to_string()), +/// ..Default::default() +/// }; +/// +/// let by_key = EnqueueEventOptions { +/// partition_key: Some("customer-17".to_string()), +/// ..Default::default() +/// }; +/// ``` +// This type stays constructible with a struct expression, like the other option +// types in this crate (`SendEventOptions`, `EventDataBatchOptions`). Callers use +// `..Default::default()`, so a new field later does not break them. +#[derive(Default, Clone, SafeDebug)] +pub struct EnqueueEventOptions { + /// The ID of the partition that receives the event. + /// + /// The partition ID must be one of the partitions that the client read when + /// it opened. + pub partition_id: Option, + + /// The partition key that selects the partition for the event. + /// + /// Events with the same partition key go to the same partition. The client + /// keeps the key on the event, so the service also sees it. + pub partition_key: Option, +} + +/// Reports that the service accepted a batch of events. +/// +/// The client passes this to the handler that +/// [`BufferedProducerClientBuilder::with_on_send_succeeded`] registers. +/// +/// [`BufferedProducerClientBuilder::with_on_send_succeeded`]: crate::builders::BufferedProducerClientBuilder::with_on_send_succeeded +#[derive(SafeDebug)] +#[non_exhaustive] +pub struct SendBatchSucceededContext { + /// The ID of the partition that received the events. + pub partition_id: String, + + /// The events in the batch, in the order that the caller enqueued them. + pub events: Vec, +} + +/// Reports that the service did not durably accept a batch of events. +/// +/// The client passes this to the handler that +/// [`BufferedProducerClientBuilder::with_on_send_failed`] registers. The client +/// reports a failure only after the retry policy is exhausted, or when the +/// error is not retryable. +/// +/// A failure does not always mean that the batch never reached the service. An +/// AMQP `Modified` or `Released` outcome settles the transfer without a durable +/// accept, and neither outcome proves whether the service stored the events. +/// +/// The client does not enqueue the events again. Re-enqueueing can change the +/// order of events, and it can store an event two times. The events are in this +/// context, so the application decides what to do with them. +/// +/// [`BufferedProducerClientBuilder::with_on_send_failed`]: crate::builders::BufferedProducerClientBuilder::with_on_send_failed +#[derive(SafeDebug)] +#[non_exhaustive] +pub struct SendBatchFailedContext { + /// The ID of the partition that the client tried to send to. + pub partition_id: String, + + /// The events in the batch, in the order that the caller enqueued them. + pub events: Vec, + + /// The error that stopped the batch. + pub error: EventHubsError, +} + +/// A handler that the client calls after the service accepts a batch. +pub(crate) type SucceededHandler = + Arc BoxFuture<'static, ()> + Send + Sync>; + +/// A handler that the client calls after a batch fails for the last time. +pub(crate) type FailedHandler = + Arc BoxFuture<'static, ()> + Send + Sync>; + +/// The handlers that a partition worker reports outcomes to. +#[derive(Clone)] +pub(crate) struct DeliveryHandlers { + pub(crate) succeeded: Option, + pub(crate) failed: FailedHandler, +} + +/// The per-partition state that the client owns. +struct PartitionState { + /// The queue that feeds the worker. + /// + /// A close takes the sender out. That ends the queue, which tells the worker + /// to drain and stop. An enqueue clones the sender and then releases this + /// lock, so a call that waits for space never holds it. + sender: Mutex>>, + + /// The capacity of the partition buffer. + /// + /// One permit stands for one event that the client accepted but has not + /// finished. The permit returns once the event reaches a terminal outcome. + /// An enqueue that finds no permit waits, so the buffer never grows without + /// a bound. + capacity: Arc, + + /// The number of events for this partition that have no terminal outcome yet. + buffered: Arc, + + /// The worker task. + task: Mutex>, + + /// Resolves once the worker stopped. + /// + /// The worker completes this channel when it leaves its loop, and a runtime + /// that drops the task cancels it. A close waits for it, so the client never + /// reports that it closed while a worker still sends, still calls a delivery + /// handler, or still holds the connection. `AbortableTask::abort` cannot + /// carry that promise: on the standard thread runtime it detaches the thread + /// and lets the await return at once. + stopped: Mutex>>, +} + +/// A producer client that buffers events and publishes them in the background. +/// +/// The caller enqueues single events. The client groups them into batches for +/// each partition, and it publishes each batch from a background worker. This +/// gives a higher throughput than [`ProducerClient`], because the caller does +/// not wait for each send. +/// +/// # Enqueue does not mean delivery +/// +/// A successful enqueue means only that the client accepted the event into the +/// local buffer. It does not mean that Event Hubs accepted the event. The +/// client reports the real outcome later, through the handlers that the builder +/// registers. +/// +/// The application must handle these trade-offs: +/// +/// * The process loses buffered events if it stops before a flush or a close. +/// * A send failure arrives after the enqueue call already returned. +/// * Buffering gives a higher throughput, but the latency of one event is less +/// predictable. +/// * Use [`ProducerClient`] when the application needs the result of each send. +/// +/// # Examples +/// +/// ```no_run +/// use azure_messaging_eventhubs::BufferedProducerClient; +/// use azure_identity::DeveloperToolsCredential; +/// use std::error::Error; +/// +/// #[tokio::main] +/// async fn main() -> Result<(), Box> { +/// let namespace = std::env::var("EVENT_HUB_NAMESPACE")?; +/// let eventhub = std::env::var("EVENT_HUB_NAME")?; +/// let credential = DeveloperToolsCredential::new(None)?; +/// +/// let producer = BufferedProducerClient::builder() +/// .with_on_send_failed(|context| async move { +/// eprintln!( +/// "{} events failed on partition {}: {}", +/// context.events.len(), +/// context.partition_id, +/// context.error +/// ); +/// }) +/// .open(&namespace, &eventhub, credential.clone()) +/// .await?; +/// +/// producer.enqueue_event("Hello, Event Hubs!", None).await?; +/// producer.flush().await?; +/// producer.close().await?; +/// Ok(()) +/// } +/// ``` +pub struct BufferedProducerClient { + /// The underlying client. Tests build the buffered client over a mock + /// instead, so this is `None` in those tests. + producer: Mutex>>, + + resolver: PartitionResolver, + partitions: HashMap, + total_buffered: Arc, + closed: AtomicBool, + abandon: Arc, + close_signal: Mutex>>, + closing: Shared>, + next_flush_id: AtomicUsize, +} + +impl BufferedProducerClient { + /// Returns a builder that creates a [`BufferedProducerClient`]. + pub fn builder() -> builders::BufferedProducerClientBuilder { + builders::BufferedProducerClientBuilder::new() + } + + /// Builds the client and starts one worker for each partition. + async fn start( + producer: Option>, + send_client: Arc, + max_wait_time: Duration, + max_buffered_event_count_per_partition: usize, + handlers: DeliveryHandlers, + ) -> Result { + let partition_ids = send_client.partition_ids().await?; + if partition_ids.is_empty() { + return Err(EventHubsError::with_message( + "The Event Hub reported no partitions.", + )); + } + + let total_buffered = Arc::new(AtomicUsize::new(0)); + let abandon = Arc::new(AtomicBool::new(false)); + let (close_signal, closing) = oneshot::channel(); + + let mut partitions = HashMap::with_capacity(partition_ids.len()); + for partition_id in &partition_ids { + let (sender, receiver) = mpsc::unbounded(); + let buffered = Arc::new(AtomicUsize::new(0)); + let capacity = Arc::new(Semaphore::new(max_buffered_event_count_per_partition)); + + let (stopped_sender, stopped) = oneshot::channel(); + let worker = PartitionWorker::new( + partition_id.clone(), + receiver, + send_client.clone(), + max_wait_time, + max_buffered_event_count_per_partition, + handlers.clone(), + buffered.clone(), + total_buffered.clone(), + abandon.clone(), + stopped_sender, + ); + + let task = get_async_runtime().spawn(Box::pin(worker.run())); + + partitions.insert( + partition_id.clone(), + PartitionState { + sender: Mutex::new(Some(sender)), + capacity, + buffered, + task: Mutex::new(Some(task)), + stopped: Mutex::new(Some(stopped)), + }, + ); + } + + debug!( + partition_count = partition_ids.len(), + buffered_event_count = max_buffered_event_count_per_partition, + "Buffered producer client started." + ); + + Ok(Self { + producer: Mutex::new(producer), + resolver: PartitionResolver::new(partition_ids), + partitions, + total_buffered, + closed: AtomicBool::new(false), + abandon, + close_signal: Mutex::new(Some(close_signal)), + closing: closing.shared(), + next_flush_id: AtomicUsize::new(0), + }) + } + + /// Adds one event to the buffer. + /// + /// The call returns once the client accepts the event into the local + /// buffer. It does not wait for Event Hubs to accept the event. The client + /// reports the delivery outcome through the registered handlers. + /// + /// When the buffer for the target partition is full, the call waits for + /// space. A close of the client makes a waiting call return an error. + /// + /// # Arguments + /// + /// * `event` - The event to add to the buffer. + /// * `options` - The routing options for the event. + /// + /// # Returns + /// + /// A `Result` that shows whether the client accepted the event. + /// + /// # Examples + /// + /// ```no_run + /// # use azure_messaging_eventhubs::{BufferedProducerClient, EnqueueEventOptions}; + /// # async fn example(producer: BufferedProducerClient) -> Result<(), Box> { + /// producer.enqueue_event("Hello, Event Hubs!", None).await?; + /// + /// producer + /// .enqueue_event( + /// "For one partition", + /// Some(EnqueueEventOptions { + /// partition_id: Some("0".to_string()), + /// ..Default::default() + /// }), + /// ) + /// .await?; + /// # Ok(()) + /// # } + /// ``` + pub async fn enqueue_event( + &self, + event: impl Into, + options: Option, + ) -> Result<()> { + let options = options.unwrap_or_default(); + if options.partition_id.is_some() && options.partition_key.is_some() { + return Err(EventHubsError::with_message( + "Set a partition ID or a partition key, not both.", + )); + } + self.enqueue_one(event.into(), &options).await + } + + /// Adds several events to the buffer. + /// + /// The client adds the events one at a time, in order. The call stops at the + /// first event that it cannot accept, and it returns that error. The client + /// keeps the events that it already accepted. + /// + /// # Arguments + /// + /// * `events` - The events to add to the buffer. + /// * `options` - The routing options for every event in the call. + /// + /// # Examples + /// + /// ```no_run + /// # use azure_messaging_eventhubs::BufferedProducerClient; + /// # async fn example(producer: BufferedProducerClient) -> Result<(), Box> { + /// producer + /// .enqueue_events(vec!["first", "second", "third"], None) + /// .await?; + /// # Ok(()) + /// # } + /// ``` + pub async fn enqueue_events( + &self, + events: impl IntoIterator, + options: Option, + ) -> Result<()> + where + E: Into, + { + let options = options.unwrap_or_default(); + if options.partition_id.is_some() && options.partition_key.is_some() { + return Err(EventHubsError::with_message( + "Set a partition ID or a partition key, not both.", + )); + } + for event in events { + self.enqueue_one(event.into(), &options).await?; + } + Ok(()) + } + + async fn enqueue_one(&self, event: EventData, options: &EnqueueEventOptions) -> Result<()> { + if self.closed.load(Ordering::Acquire) { + return Err(Self::closed_error()); + } + + let partition_id = match (&options.partition_id, &options.partition_key) { + (Some(partition_id), _) => { + if !self.resolver.contains(partition_id) { + return Err(EventHubsError::with_message(format!( + "The Event Hub has no partition with the ID {partition_id}." + ))); + } + partition_id.clone() + } + (None, Some(partition_key)) => self.resolver.assign_for_key(partition_key).to_string(), + (None, None) => self.resolver.assign_round_robin().to_string(), + }; + + let mut message = AmqpMessage::from(event.clone()); + if message.properties.is_none() || message.properties.as_ref().unwrap().message_id.is_none() + { + message.set_message_id(Uuid::new_v4()); + } + if let Some(partition_key) = options.partition_key.as_ref() { + // Keep the key on the message so the service sees it too. + message.add_message_annotation( + AmqpSymbol::from("x-opt-partition-key"), + partition_key.clone(), + ); + } + + let state = self + .partitions + .get(&partition_id) + .expect("the resolver only returns known partitions"); + + // Take one unit of buffer capacity. The permit travels with the event and + // returns once the event reaches a terminal outcome. + let permit = match state.capacity.try_acquire_arc() { + Some(permit) => permit, + None => { + debug!( + partition_id = %partition_id, + buffered_event_count = state.buffered.load(Ordering::Acquire), + "The partition buffer is full; the enqueue is waiting for space." + ); + let mut acquire = Box::pin(state.capacity.acquire_arc()).fuse(); + let mut closing = self.closing.clone().fuse(); + futures::select! { + permit = acquire => permit, + _ = closing => return Err(Self::closed_error()), + } + } + }; + + let command = Command::Event { + event, + message: Box::new(message), + permit, + }; + + // Take a clone of the sender and release the lock at once, so a slow + // enqueue never blocks a close. + let sender = { + let guard = state.sender.lock().unwrap(); + guard.as_ref().ok_or_else(Self::closed_error)?.clone() + }; + // Count the event before the worker can see it. The worker decrements + // the counts as soon as the event reaches a terminal outcome, and a + // fast terminal path (an oversized event with a handler that returns at + // once) can run before this call returns. Counting afterwards lets that + // decrement reach zero first and wrap the counts to `usize::MAX`. + state.buffered.fetch_add(1, Ordering::AcqRel); + self.total_buffered.fetch_add(1, Ordering::AcqRel); + + if sender.unbounded_send(command).is_err() { + // The worker never saw the event, so it never decrements for it. + state.buffered.fetch_sub(1, Ordering::AcqRel); + self.total_buffered.fetch_sub(1, Ordering::AcqRel); + return Err(Self::closed_error()); + } + + trace!( + partition_id = %partition_id, + "The client accepted an event into the buffer." + ); + Ok(()) + } + + /// Sends every event that the client accepted before this call. + /// + /// The call sets a barrier. It completes once every event in front of the + /// barrier reaches a terminal outcome, either a success or a failure. An + /// event that the caller enqueues after the barrier does not delay the call. + /// + /// # Examples + /// + /// ```no_run + /// # use azure_messaging_eventhubs::BufferedProducerClient; + /// # async fn example(producer: BufferedProducerClient) -> Result<(), Box> { + /// producer.enqueue_event("Hello, Event Hubs!", None).await?; + /// producer.flush().await?; + /// # Ok(()) + /// # } + /// ``` + pub async fn flush(&self) -> Result<()> { + let flush_id = self.next_flush_id.fetch_add(1, Ordering::Relaxed); + debug!(flush_id, "Flush started."); + + let mut waiters = Vec::with_capacity(self.partitions.len()); + for (partition_id, state) in &self.partitions { + let (completed, waiter) = oneshot::channel(); + let sender = state.sender.lock().unwrap().as_ref().cloned(); + let sent = match sender { + Some(sender) => sender.unbounded_send(Command::Flush(completed)).is_ok(), + None => false, + }; + if sent { + waiters.push(waiter); + } else { + trace!( + flush_id, + partition_id = %partition_id, + "The partition worker already stopped; the flush skips it." + ); + } + } + + let mut cancelled = false; + for waiter in waiters { + if waiter.await.is_err() { + cancelled = true; + } + } + + if cancelled { + debug!(flush_id, "Flush stopped because the client is closing."); + return Err(Self::closed_error()); + } + + debug!(flush_id, "Flush completed."); + Ok(()) + } + + /// Returns the number of events that have no terminal outcome yet. + /// + /// The count covers every partition. It includes the events in the queues, + /// the events in an active batch, and the events of a send that is in + /// flight. + pub fn total_buffered_event_count(&self) -> usize { + self.total_buffered.load(Ordering::Acquire) + } + + /// Returns the number of events for one partition that have no terminal + /// outcome yet. + /// + /// The method returns 0 for a partition ID that the Event Hub does not have. + pub fn buffered_event_count(&self, partition_id: &str) -> usize { + self.partitions + .get(partition_id) + .map(|state| state.buffered.load(Ordering::Acquire)) + .unwrap_or(0) + } + + /// Sends the buffered events, then closes the client. + /// + /// The client stops accepting new events, sends every event that it already + /// accepted, stops the workers, and releases the connection. + /// + /// Use [`abort`](Self::abort) to close without sending the buffered events. + /// + /// The method takes `&self`, so an application can hold the client in an + /// `Arc`, enqueue from many tasks, and close it from one of them. A second + /// call does nothing and returns `Ok`. + pub async fn close(&self) -> Result<()> { + self.shutdown(false).await + } + + /// Closes the client at once and abandons the buffered events. + /// + /// The client drops every event that it did not send yet. It reports the + /// number of abandoned events in a warning, and it removes them from the + /// buffered counts. + /// + /// The call waits for every worker to stop, so it does not return while a + /// worker still publishes, still calls a delivery handler, or still holds + /// the connection. A worker that is inside a send when the call starts + /// finishes that send first on a runtime that cannot cancel a task, so the + /// call can take as long as one send. The retry policy bounds that send. No + /// batch that the client has not started to send goes to the service. + /// + /// A second call does nothing and returns `Ok`. + pub async fn abort(&self) -> Result<()> { + self.shutdown(true).await + } + + async fn shutdown(&self, abandon: bool) -> Result<()> { + if self.closed.swap(true, Ordering::AcqRel) { + return Ok(()); + } + if abandon { + self.abandon.store(true, Ordering::Release); + } + + debug!( + abandon, + buffered_event_count = self.total_buffered.load(Ordering::Acquire), + "Closing the buffered producer client." + ); + + // Fail every enqueue that is waiting for space. + self.signal_closing(); + + // Taking the sender ends each queue. A worker then sends what it still + // holds and stops, unless the client abandons the events. The abandon + // flag is already set above, so a worker that takes this path drops its + // events instead of sending them. + let mut tasks = Vec::with_capacity(self.partitions.len()); + let mut acknowledgements = Vec::with_capacity(self.partitions.len()); + for state in self.partitions.values() { + drop(state.sender.lock().unwrap().take()); + if let Some(stopped) = state.stopped.lock().unwrap().take() { + acknowledgements.push(stopped); + } + if let Some(task) = state.task.lock().unwrap().take() { + if abandon { + // On a runtime with cancellation this ends an in-flight send + // at once. On the standard thread runtime it only detaches + // the thread, so the acknowledgement below, not this call, is + // what makes the close wait for the worker. + task.abort(); + } + tasks.push(task); + } + } + + // Wait for every worker to stop. A cancelled task drops its end of the + // channel, which resolves the receiver with an error, so this waits for + // the worker to finish or to be dropped, and never for both. + for acknowledgement in acknowledgements { + let _ = acknowledgement.await; + } + + for task in tasks { + if let Err(error) = task.await { + debug!("A partition worker stopped with an error: {error}"); + } + } + + if abandon { + // A worker that the runtime stopped inside a send cannot clear its + // own counters, so the client clears them here. The counts and the + // abandoned events then agree. + let abandoned = self.total_buffered.swap(0, Ordering::AcqRel); + for state in self.partitions.values() { + state.buffered.store(0, Ordering::Release); + } + if abandoned > 0 { + warn!( + event_count = abandoned, + "Abandoned buffered events during an immediate close." + ); + } + } + + let producer = self.producer.lock().unwrap().take(); + if let Some(producer) = producer { + match Arc::try_unwrap(producer) { + Ok(producer) => producer.close().await?, + Err(_) => { + // A worker thread that the runtime could not stop still + // holds a reference. Dropping this one closes the + // connection once that thread finishes. + warn!( + "Could not close the connection now, because another reference exists; \ + it closes when the last reference drops." + ); + } + } + } + + debug!("Buffered producer client closed."); + Ok(()) + } + + fn signal_closing(&self) { + if let Some(signal) = self.close_signal.lock().unwrap().take() { + let _ = signal.send(()); + } + } + + fn closed_error() -> EventHubsError { + EventHubsError::with_message("The buffered producer client is closed.") + } + + /// Returns the underlying producer, so a test can force an error on the + /// connection. + #[cfg(test)] + pub(crate) fn inner_producer(&self) -> Option> { + self.producer.lock().unwrap().clone() + } +} + +impl std::fmt::Debug for BufferedProducerClient { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("BufferedProducerClient") + .field("partition_count", &self.partitions.len()) + .field( + "total_buffered_event_count", + &self.total_buffered.load(Ordering::Acquire), + ) + .field("closed", &self.closed.load(Ordering::Acquire)) + .finish() + } +} + +impl Drop for BufferedProducerClient { + fn drop(&mut self) { + if self.closed.load(Ordering::Acquire) { + return; + } + + warn!( + buffered_event_count = self.total_buffered.load(Ordering::Acquire), + "The buffered producer client was dropped without a close; \ + buffered events are abandoned. Call close or flush when delivery matters." + ); + + self.closed.store(true, Ordering::Release); + self.abandon.store(true, Ordering::Release); + self.signal_closing(); + + // Stop the workers. Dropping the partition map also drops every sender, + // which ends each queue. The abandon flag is set above, so a worker that + // reaches the end of its queue drops its events instead of sending them, + // whatever the runtime does with the abort below. A drop cannot wait for + // the workers, so call `close` or `abort` when that matters. + for state in self.partitions.values_mut() { + if let Some(task) = state.task.lock().unwrap().take() { + task.abort(); + } + } + } +} + +/// Builders for the buffered producer client. +pub mod builders { + use super::{ + BufferedProducerClient, DeliveryHandlers, FailedHandler, SendBatchFailedContext, + SendBatchSucceededContext, SucceededHandler, + DEFAULT_MAX_BUFFERED_EVENT_COUNT_PER_PARTITION, DEFAULT_MAX_WAIT_TIME_SECONDS, + }; + use crate::{ + error::Result, + producer::{buffered::send_client::ProducerSendClient, ProducerClient}, + EventHubsError, RetryOptions, + }; + use azure_core::time::Duration; + use futures::FutureExt; + use std::{future::Future, sync::Arc}; + + /// A builder that creates a [`BufferedProducerClient`]. + /// + /// The builder needs a handler for failed batches. A buffered send reports + /// its failure later, so a client with no failure handler would lose events + /// without a report. Use + /// [`with_on_send_failed`](Self::with_on_send_failed) before `open`. + /// + /// # Examples + /// + /// ```no_run + /// use azure_messaging_eventhubs::BufferedProducerClient; + /// use azure_identity::DeveloperToolsCredential; + /// use azure_core::time::Duration; + /// + /// #[tokio::main] + /// async fn main() -> Result<(), Box> { + /// let credential = DeveloperToolsCredential::new(None)?; + /// let producer = BufferedProducerClient::builder() + /// .with_max_wait_time(Duration::seconds(1)) + /// .with_max_buffered_event_count_per_partition(1500) + /// .with_on_send_succeeded(|context| async move { + /// println!("sent {} events", context.events.len()); + /// }) + /// .with_on_send_failed(|context| async move { + /// eprintln!("failed {} events: {}", context.events.len(), context.error); + /// }) + /// .open("my_namespace", "my_eventhub", credential) + /// .await?; + /// producer.close().await?; + /// Ok(()) + /// } + /// ``` + pub struct BufferedProducerClientBuilder { + application_id: Option, + retry_options: Option, + custom_endpoint: Option, + max_wait_time: Duration, + max_buffered_event_count_per_partition: usize, + on_send_succeeded: Option, + on_send_failed: Option, + } + + impl Default for BufferedProducerClientBuilder { + fn default() -> Self { + Self::new() + } + } + + impl BufferedProducerClientBuilder { + pub(super) fn new() -> Self { + Self { + application_id: None, + retry_options: None, + custom_endpoint: None, + max_wait_time: Duration::seconds(DEFAULT_MAX_WAIT_TIME_SECONDS), + max_buffered_event_count_per_partition: + DEFAULT_MAX_BUFFERED_EVENT_COUNT_PER_PARTITION, + on_send_succeeded: None, + on_send_failed: None, + } + } + + /// Sets the application ID that identifies the client. + pub fn with_application_id(mut self, application_id: String) -> Self { + self.application_id = Some(application_id); + self + } + + /// Sets the options that configure retry operations. + pub fn with_retry_options(mut self, retry_options: RetryOptions) -> Self { + self.retry_options = Some(retry_options); + self + } + + /// Sets a custom endpoint for the Event Hub. + pub fn with_custom_endpoint(mut self, endpoint: String) -> Self { + self.custom_endpoint = Some(endpoint); + self + } + + /// Sets how long the client waits before it sends a batch that is not full. + /// + /// A short time lowers the latency of one event. A long time makes the + /// batches larger, which raises the throughput. The default is 1 second. + pub fn with_max_wait_time(mut self, max_wait_time: Duration) -> Self { + self.max_wait_time = max_wait_time; + self + } + + /// Sets how many events the client buffers for one partition. + /// + /// An enqueue waits for space once a partition buffer holds this many + /// events. The default is 1500. + pub fn with_max_buffered_event_count_per_partition(mut self, count: usize) -> Self { + self.max_buffered_event_count_per_partition = count; + self + } + + /// Registers the handler that runs after the service accepts a batch. + /// + /// The handler is optional. The partition worker waits for the handler, + /// so a slow handler slows only its own partition. The handler must not + /// call `flush`, `close`, or `abort` on the same client, because that + /// would wait for the worker that is waiting for the handler. + pub fn with_on_send_succeeded(mut self, handler: F) -> Self + where + F: Fn(SendBatchSucceededContext) -> Fut + Send + Sync + 'static, + Fut: Future + Send + 'static, + { + self.on_send_succeeded = Some(Arc::new(move |context| handler(context).boxed())); + self + } + + /// Registers the handler that runs after a batch fails for the last time. + /// + /// The handler is required. The client calls it only once the retry + /// policy is exhausted, or when the error is not retryable. The context + /// holds the events, so the application decides what to do with them. + /// The client does not enqueue them again. + /// + /// The partition worker waits for the handler, so a slow handler slows + /// only its own partition. The handler must not call `flush`, `close`, + /// or `abort` on the same client, because that would wait for the worker + /// that is waiting for the handler. + pub fn with_on_send_failed(mut self, handler: F) -> Self + where + F: Fn(SendBatchFailedContext) -> Fut + Send + Sync + 'static, + Fut: Future + Send + 'static, + { + self.on_send_failed = Some(Arc::new(move |context| handler(context).boxed())); + self + } + + fn build_handlers(&self) -> Result { + let failed = self.on_send_failed.clone().ok_or_else(|| { + EventHubsError::with_message( + "A buffered producer client needs a handler for failed batches. \ + Call with_on_send_failed before open.", + ) + })?; + Ok(DeliveryHandlers { + succeeded: self.on_send_succeeded.clone(), + failed, + }) + } + + fn validate(&self) -> Result<()> { + if self.max_buffered_event_count_per_partition == 0 { + return Err(EventHubsError::with_message( + "The maximum buffered event count for one partition must be at least 1.", + )); + } + if self.max_wait_time <= Duration::ZERO { + return Err(EventHubsError::with_message( + "The maximum wait time must be longer than zero.", + )); + } + Ok(()) + } + + /// Opens a connection to the Event Hub and starts the background workers. + /// + /// # Arguments + /// + /// * `fully_qualified_namespace` - The fully qualified namespace of the Event Hubs instance. + /// * `eventhub` - The name of the Event Hub. + /// * `credential` - The token credential used for authorization. + pub async fn open( + self, + fully_qualified_namespace: &str, + eventhub: &str, + credential: Arc, + ) -> Result { + self.validate()?; + let handlers = self.build_handlers()?; + + let mut builder = ProducerClient::builder(); + if let Some(application_id) = self.application_id { + builder = builder.with_application_id(application_id); + } + if let Some(retry_options) = self.retry_options { + builder = builder.with_retry_options(retry_options); + } + if let Some(custom_endpoint) = self.custom_endpoint { + builder = builder.with_custom_endpoint(custom_endpoint); + } + + let producer = Arc::new( + builder + .open(fully_qualified_namespace, eventhub, credential) + .await?, + ); + let send_client = Arc::new(ProducerSendClient::new(producer.clone())); + + BufferedProducerClient::start( + Some(producer), + send_client, + self.max_wait_time, + self.max_buffered_event_count_per_partition, + handlers, + ) + .await + } + + /// Opens a connection to the Event Hub with a connection string, and + /// starts the background workers. + /// + /// Prefer [`open`](Self::open) with a `TokenCredential` for production. + /// + /// # Arguments + /// + /// * `connection_string` - An Event Hubs connection string. + /// * `eventhub` - The Event Hub name. This is required unless the + /// connection string includes an `EntityPath`. + pub async fn open_with_connection_string( + self, + connection_string: &str, + eventhub: Option<&str>, + ) -> Result { + self.validate()?; + let handlers = self.build_handlers()?; + + let mut builder = ProducerClient::builder(); + if let Some(application_id) = self.application_id { + builder = builder.with_application_id(application_id); + } + if let Some(retry_options) = self.retry_options { + builder = builder.with_retry_options(retry_options); + } + if let Some(custom_endpoint) = self.custom_endpoint { + builder = builder.with_custom_endpoint(custom_endpoint); + } + + let producer = Arc::new( + builder + .open_with_connection_string(connection_string, eventhub) + .await?, + ); + let send_client = Arc::new(ProducerSendClient::new(producer.clone())); + + BufferedProducerClient::start( + Some(producer), + send_client, + self.max_wait_time, + self.max_buffered_event_count_per_partition, + handlers, + ) + .await + } + + /// Starts a client over a supplied send client, with no network. + #[cfg(test)] + pub(crate) async fn open_with_send_client( + self, + send_client: Arc, + ) -> Result { + self.validate()?; + let handlers = self.build_handlers()?; + BufferedProducerClient::start( + None, + send_client, + self.max_wait_time, + self.max_buffered_event_count_per_partition, + handlers, + ) + .await + } + } +} + +#[cfg(test)] +mod tests; diff --git a/sdk/eventhubs/azure_messaging_eventhubs/src/producer/buffered/partition_resolver.rs b/sdk/eventhubs/azure_messaging_eventhubs/src/producer/buffered/partition_resolver.rs new file mode 100644 index 00000000000..a0ed5a10cc8 --- /dev/null +++ b/sdk/eventhubs/azure_messaging_eventhubs/src/producer/buffered/partition_resolver.rs @@ -0,0 +1,297 @@ +// Copyright (c) Microsoft Corporation. All Rights reserved +// Licensed under the MIT license. + +use std::sync::atomic::{AtomicUsize, Ordering}; + +/// Assigns events to partitions. +/// +/// The resolver holds the partition IDs that the client read when it opened. +/// It assigns a partition in one of two ways: +/// +/// * Round robin, when the caller gives no partition ID and no partition key. +/// * A hash of the partition key, when the caller gives a partition key. +/// +/// The hash matches the other Azure SDKs and the Event Hubs gateway. Two +/// clients in two languages send the same partition key to the same partition. +pub(crate) struct PartitionResolver { + partitions: Vec, + next: AtomicUsize, +} + +impl PartitionResolver { + /// Creates a resolver over the given partition IDs. + /// + /// The caller must supply at least one partition ID. + pub(crate) fn new(partitions: Vec) -> Self { + debug_assert!( + !partitions.is_empty(), + "a partition resolver needs at least one partition" + ); + Self { + partitions, + next: AtomicUsize::new(0), + } + } + + /// Returns `true` when the given partition ID is one of the known partitions. + pub(crate) fn contains(&self, partition_id: &str) -> bool { + self.partitions.iter().any(|p| p == partition_id) + } + + /// Assigns the next partition in round-robin order. + pub(crate) fn assign_round_robin(&self) -> &str { + // The counter wraps. A wrap changes which partition follows which, but + // every index stays inside the partition range. + let index = self.next.fetch_add(1, Ordering::Relaxed); + &self.partitions[index % self.partitions.len()] + } + + /// Assigns the partition that the partition key hashes to. + pub(crate) fn assign_for_key(&self, partition_key: &str) -> &str { + let hash = Self::generate_hash_code(partition_key); + let index = ((hash as i32) % (self.partitions.len() as i32)).unsigned_abs() as usize; + &self.partitions[index] + } + + /// Generates the hash code for a partition key with the Jenkins lookup3 algorithm. + /// + /// This is a port of the .NET implementation, which is itself a port of the + /// Event Hubs service code. The value must match the gateway, so do not + /// change it without careful thought. + /// + /// Source: + fn generate_hash_code(partition_key: &str) -> i16 { + let (hash1, hash2) = Self::compute_hash(partition_key.as_bytes(), 0, 0); + (hash1 ^ hash2) as u16 as i16 + } + + /// Computes the two lookup3 hash values for the given bytes. + /// + /// Source: + fn compute_hash(data: &[u8], seed1: u32, seed2: u32) -> (u32, u32) { + fn le32(bytes: &[u8]) -> u32 { + u32::from_le_bytes([bytes[0], bytes[1], bytes[2], bytes[3]]) + } + + let len = data.len() as u32; + let mut a = 0xDEAD_BEEF_u32.wrapping_add(len).wrapping_add(seed1); + let mut b = a; + let mut c = a.wrapping_add(seed2); + + let chunks = if data.len() > 12 { + (data.len() - 1) / 12 + } else { + 0 + }; + + let mut offset = 0usize; + for _ in 0..chunks { + a = a.wrapping_add(le32(&data[offset..offset + 4])); + b = b.wrapping_add(le32(&data[offset + 4..offset + 8])); + c = c.wrapping_add(le32(&data[offset + 8..offset + 12])); + offset += 12; + + a = a.wrapping_sub(c); + a ^= c.rotate_left(4); + c = c.wrapping_add(b); + + b = b.wrapping_sub(a); + b ^= a.rotate_left(6); + a = a.wrapping_add(c); + + c = c.wrapping_sub(b); + c ^= b.rotate_left(8); + b = b.wrapping_add(a); + + a = a.wrapping_sub(c); + a ^= c.rotate_left(16); + c = c.wrapping_add(b); + + b = b.wrapping_sub(a); + b ^= a.rotate_left(19); + a = a.wrapping_add(c); + + c = c.wrapping_sub(b); + c ^= b.rotate_left(4); + b = b.wrapping_add(a); + } + + let tail = &data[offset..]; + match tail.len() { + 12 => { + a = a.wrapping_add(le32(&tail[0..4])); + b = b.wrapping_add(le32(&tail[4..8])); + c = c.wrapping_add(le32(&tail[8..12])); + } + left @ (9..=11) => { + if left == 11 { + c = c.wrapping_add((tail[10] as u32) << 16); + } + if left >= 10 { + c = c.wrapping_add((tail[9] as u32) << 8); + } + c = c.wrapping_add(tail[8] as u32); + b = b.wrapping_add(le32(&tail[4..8])); + a = a.wrapping_add(le32(&tail[0..4])); + } + 8 => { + b = b.wrapping_add(le32(&tail[4..8])); + a = a.wrapping_add(le32(&tail[0..4])); + } + left @ (5..=7) => { + if left == 7 { + b = b.wrapping_add((tail[6] as u32) << 16); + } + if left >= 6 { + b = b.wrapping_add((tail[5] as u32) << 8); + } + b = b.wrapping_add(tail[4] as u32); + a = a.wrapping_add(le32(&tail[0..4])); + } + 4 => { + a = a.wrapping_add(le32(&tail[0..4])); + } + left @ (1..=3) => { + if left == 3 { + a = a.wrapping_add((tail[2] as u32) << 16); + } + if left >= 2 { + a = a.wrapping_add((tail[1] as u32) << 8); + } + a = a.wrapping_add(tail[0] as u32); + } + _ => return (c, b), + } + + c ^= b; + c = c.wrapping_sub(b.rotate_left(14)); + + a ^= c; + a = a.wrapping_sub(c.rotate_left(11)); + + b ^= a; + b = b.wrapping_sub(a.rotate_left(25)); + + c ^= b; + c = c.wrapping_sub(b.rotate_left(16)); + + a ^= c; + a = a.wrapping_sub(c.rotate_left(4)); + + b ^= a; + b = b.wrapping_sub(a.rotate_left(14)); + + c ^= b; + c = c.wrapping_sub(b.rotate_left(24)); + + (c, b) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn partitions(count: usize) -> Vec { + (0..count).map(|i| i.to_string()).collect() + } + + /// The expected values come from the .NET test suite. The Rust hash must + /// agree with .NET, with the other SDKs, and with the Event Hubs gateway. + /// + /// Source: + #[test] + fn hash_matches_dotnet_vectors() { + // cspell:disable + let cases: &[(&str, i16)] = &[ + ("7", -15263), + ("131", 30562), + ("7149583486996073602", 12977), + ("FWfAT", -22341), + ("sOdeEAsyQoEuEFPGerWO", -6503), + ( + "FAyAIctPeCgmiwLKbJcyswoHglHVjQdvtBowLACDNORsYvOcLddNJYDmhAVkbyLOrHTKLneMNcbgWVlasVywOByANjs", + 5226, + ), + ( + "1XYM6!(7(lF5wq4k4m*e$Nc!1ezLJv*1YK1Y-C^*&B$O)lq^iUkG(TNzXG;Zi#z2Og*Qq0#^*k):vXh$3,C7We7%W0meJ;b3,rQCg^J;^twXgs5E$$hWKxqp", + 23950, + ), + ( + "E(x;RRIaQcJs*P;D&jTPau-4K04oqr:lF6Z):ERpo&;9040qyV@G1_c9mgOs-8_8/10Fwa-7b7-yP!T-!IH&968)FWuI;(^g$2fN;)HJ^^yTn:", + -29304, + ), + ("!c*_!I@1^c", 15372), + ("p4*!jioeO/z-!-;w:dh", -3104), + ("$0cb", 26269), + ("-4189260826195535198", 453), + ]; + // cspell:enable + + for (key, expected) in cases { + assert_eq!( + PartitionResolver::generate_hash_code(key), + *expected, + "the hash for key {key} was incorrect" + ); + } + } + + #[test] + fn round_robin_walks_every_partition_then_wraps() { + let resolver = PartitionResolver::new(partitions(4)); + + let first: Vec = (0..4) + .map(|_| resolver.assign_round_robin().to_string()) + .collect(); + assert_eq!(first, vec!["0", "1", "2", "3"]); + + // The next pass repeats the same order. + let second: Vec = (0..4) + .map(|_| resolver.assign_round_robin().to_string()) + .collect(); + assert_eq!(second, first); + } + + #[test] + fn round_robin_with_one_partition_always_returns_it() { + let resolver = PartitionResolver::new(partitions(1)); + for _ in 0..5 { + assert_eq!(resolver.assign_round_robin(), "0"); + } + } + + #[test] + fn key_assignment_is_stable() { + let resolver = PartitionResolver::new(partitions(8)); + let first = resolver.assign_for_key("some-key").to_string(); + for _ in 0..5 { + assert_eq!(resolver.assign_for_key("some-key"), first); + } + } + + #[test] + fn key_assignment_stays_in_range() { + // A key whose hash is negative must still map into the partition range. + for count in 1..=32 { + let resolver = PartitionResolver::new(partitions(count)); + // cspell:disable-next-line + for key in ["7", "FWfAT", "p4*!jioeO/z-!-;w:dh", "", "$0cb"] { + let assigned = resolver.assign_for_key(key); + assert!( + resolver.contains(assigned), + "key {key} mapped outside the partition range for {count} partitions" + ); + } + } + } + + #[test] + fn contains_reports_known_partitions() { + let resolver = PartitionResolver::new(partitions(3)); + assert!(resolver.contains("0")); + assert!(resolver.contains("2")); + assert!(!resolver.contains("3")); + assert!(!resolver.contains("not-a-partition")); + } +} diff --git a/sdk/eventhubs/azure_messaging_eventhubs/src/producer/buffered/send_client.rs b/sdk/eventhubs/azure_messaging_eventhubs/src/producer/buffered/send_client.rs new file mode 100644 index 00000000000..b587748f0e2 --- /dev/null +++ b/sdk/eventhubs/azure_messaging_eventhubs/src/producer/buffered/send_client.rs @@ -0,0 +1,264 @@ +// Copyright (c) Microsoft Corporation. All Rights reserved +// Licensed under the MIT license. + +// cspell: ignore retryable + +use crate::{error::Result, producer::ProducerClient}; +use azure_core::http::Url; +use azure_core_amqp::{AmqpMessage, AmqpSendOutcome, AmqpSenderApis}; +use std::sync::Arc; + +/// The operations that a partition worker needs from the AMQP layer. +/// +/// The buffered producer talks to the service only through this trait. The +/// production implementation forwards to [`ProducerClient`], which applies the +/// retry policy and the connection recovery. A test implementation supplies +/// scripted outcomes, so the worker tests do not need a network. +#[async_trait::async_trait] +pub(crate) trait BufferedSendClient: Send + Sync + 'static { + /// Returns the partition IDs of the Event Hub. + async fn partition_ids(&self) -> Result>; + + /// Returns the largest message that the link to the partition accepts. + async fn max_message_size(&self, partition_id: &str) -> Result; + + /// Sends one batch envelope to a partition and returns the AMQP outcome. + /// + /// An `Err` means the retry policy is exhausted, or the error is not + /// retryable. The caller must treat an `Err` as a terminal failure. + async fn send_envelope( + &self, + partition_id: &str, + envelope: AmqpMessage, + ) -> Result; +} + +/// The production [`BufferedSendClient`], backed by a [`ProducerClient`]. +pub(crate) struct ProducerSendClient { + producer: Arc, +} + +impl ProducerSendClient { + pub(crate) fn new(producer: Arc) -> Self { + Self { producer } + } + + fn partition_path(&self, partition_id: &str) -> Result { + let path = format!("{}/Partitions/{}", self.producer.base_url(), partition_id); + Url::parse(&path).map_err(|e| azure_core::Error::from(e).into()) + } +} + +#[async_trait::async_trait] +impl BufferedSendClient for ProducerSendClient { + async fn partition_ids(&self) -> Result> { + Ok(self.producer.get_eventhub_properties().await?.partition_ids) + } + + async fn max_message_size(&self, partition_id: &str) -> Result { + let path = self.partition_path(partition_id)?; + let sender = self.producer.ensure_sender(path).await?; + sender.max_message_size().await?.ok_or_else(|| { + crate::EventHubsError::with_message( + "No maximum message size available from the sender link.", + ) + }) + } + + async fn send_envelope( + &self, + partition_id: &str, + envelope: AmqpMessage, + ) -> Result { + let path = self.partition_path(partition_id)?; + self.producer.send_batch_envelope(path, envelope).await + } +} + +#[cfg(test)] +pub(crate) mod mock { + use super::*; + use crate::EventHubsError; + use azure_core_amqp::message::AmqpMessageBody; + use futures::channel::{mpsc, oneshot}; + use std::{ + collections::{HashMap, VecDeque}, + sync::Mutex, + }; + + /// The outcome that the mock returns for one send. + #[derive(Clone, Debug)] + pub(crate) enum SendScript { + Accepted, + Modified, + Released, + /// The recoverable sender turns a rejected outcome into an error, so a + /// rejected send reaches the worker as an error. + Rejected, + /// The retry policy is exhausted, or the error is not retryable. + Error(&'static str), + } + + /// One send that the mock observed. + #[derive(Clone, Debug, PartialEq, Eq)] + pub(crate) struct RecordedSend { + pub(crate) partition_id: String, + pub(crate) event_count: usize, + } + + #[derive(Default)] + struct MockState { + script: HashMap>, + gates: HashMap>>, + sends: Vec, + } + + /// A [`BufferedSendClient`] for tests. + /// + /// The mock records every send, returns scripted outcomes, and can hold a + /// send open until the test releases it. Tests use the gates instead of + /// sleeps, so no test depends on timing to prove a race. + pub(crate) struct MockSendClient { + partitions: Vec, + max_message_size: u64, + state: Mutex, + started_tx: mpsc::UnboundedSender, + } + + impl MockSendClient { + /// Creates a mock over the given partition IDs. + /// + /// The returned receiver reports the partition ID of every send as it + /// starts, before the mock waits on any gate for that send. + pub(crate) fn new(partitions: &[&str]) -> (Arc, mpsc::UnboundedReceiver) { + let (started_tx, started_rx) = mpsc::unbounded(); + let client = Arc::new(Self { + partitions: partitions.iter().map(|p| p.to_string()).collect(), + max_message_size: 1024 * 1024, + state: Mutex::new(MockState::default()), + started_tx, + }); + (client, started_rx) + } + + /// Creates a mock whose link reports the given maximum message size. + pub(crate) fn with_max_message_size( + partitions: &[&str], + max_message_size: u64, + ) -> (Arc, mpsc::UnboundedReceiver) { + let (started_tx, started_rx) = mpsc::unbounded(); + let client = Arc::new(Self { + partitions: partitions.iter().map(|p| p.to_string()).collect(), + max_message_size, + state: Mutex::new(MockState::default()), + started_tx, + }); + (client, started_rx) + } + + /// Queues one scripted outcome for a partition. + /// + /// The mock returns `Accepted` when a partition has no queued outcome. + pub(crate) fn push_outcome(&self, partition_id: &str, outcome: SendScript) { + self.state + .lock() + .unwrap() + .script + .entry(partition_id.to_string()) + .or_default() + .push_back(outcome); + } + + /// Holds the next un-gated send on a partition until the test drops or + /// completes the returned sender. + pub(crate) fn gate(&self, partition_id: &str) -> oneshot::Sender<()> { + let (tx, rx) = oneshot::channel(); + self.state + .lock() + .unwrap() + .gates + .entry(partition_id.to_string()) + .or_default() + .push_back(rx); + tx + } + + /// Returns every send that the mock observed, in order. + pub(crate) fn sends(&self) -> Vec { + self.state.lock().unwrap().sends.clone() + } + + /// Returns the total number of events across every observed send. + pub(crate) fn total_events(&self) -> usize { + self.sends().iter().map(|s| s.event_count).sum() + } + + fn count_events(envelope: &AmqpMessage) -> usize { + match &envelope.body { + AmqpMessageBody::Binary(items) => items.len(), + _ => 0, + } + } + } + + #[async_trait::async_trait] + impl BufferedSendClient for MockSendClient { + async fn partition_ids(&self) -> Result> { + Ok(self.partitions.clone()) + } + + async fn max_message_size(&self, _partition_id: &str) -> Result { + Ok(self.max_message_size) + } + + async fn send_envelope( + &self, + partition_id: &str, + envelope: AmqpMessage, + ) -> Result { + let event_count = Self::count_events(&envelope); + + // Report the start before waiting on a gate, so a test can observe + // that a send is in flight while it is still held. + let _ = self.started_tx.unbounded_send(partition_id.to_string()); + + let gate = self + .state + .lock() + .unwrap() + .gates + .get_mut(partition_id) + .and_then(|g| g.pop_front()); + if let Some(gate) = gate { + // A dropped sender resolves the receiver with an error. Either + // way the send proceeds once the test releases the gate. + let _ = gate.await; + } + + let outcome = { + let mut state = self.state.lock().unwrap(); + state.sends.push(RecordedSend { + partition_id: partition_id.to_string(), + event_count, + }); + state + .script + .get_mut(partition_id) + .and_then(|s| s.pop_front()) + .unwrap_or(SendScript::Accepted) + }; + + match outcome { + SendScript::Accepted => Ok(AmqpSendOutcome::Accepted), + // `SendModification` is not re-exported from `azure_core_amqp`, + // so build it through `Default` instead of naming the type. + SendScript::Modified => Ok(AmqpSendOutcome::Modified(Default::default())), + SendScript::Released => Ok(AmqpSendOutcome::Released), + SendScript::Rejected => Err(EventHubsError::with_message( + "Batch was rejected by the Event Hub.", + )), + SendScript::Error(message) => Err(EventHubsError::with_message(message)), + } + } + } +} diff --git a/sdk/eventhubs/azure_messaging_eventhubs/src/producer/buffered/tests.rs b/sdk/eventhubs/azure_messaging_eventhubs/src/producer/buffered/tests.rs new file mode 100644 index 00000000000..6573347af30 --- /dev/null +++ b/sdk/eventhubs/azure_messaging_eventhubs/src/producer/buffered/tests.rs @@ -0,0 +1,1388 @@ +// Copyright (c) Microsoft Corporation. All Rights reserved +// Licensed under the MIT license. + +// cspell: ignore backpressure retryable + +use super::{ + send_client::mock::{MockSendClient, SendScript}, + *, +}; +use crate::models::EventData; +use azure_core::time::Duration; +use futures::{pin_mut, poll, StreamExt}; +use std::collections::HashSet; + +/// One delivery report that a test collected from a handler. +#[derive(Debug, Clone, PartialEq, Eq)] +enum Report { + Succeeded { + partition_id: String, + bodies: Vec, + }, + Failed { + partition_id: String, + bodies: Vec, + error: String, + }, +} + +impl Report { + fn partition_id(&self) -> &str { + match self { + Report::Succeeded { partition_id, .. } | Report::Failed { partition_id, .. } => { + partition_id + } + } + } + + fn bodies(&self) -> &[String] { + match self { + Report::Succeeded { bodies, .. } | Report::Failed { bodies, .. } => bodies, + } + } + + fn is_success(&self) -> bool { + matches!(self, Report::Succeeded { .. }) + } +} + +fn bodies_of(events: &[EventData]) -> Vec { + events + .iter() + .map(|event| String::from_utf8_lossy(event.body().unwrap_or_default()).into_owned()) + .collect() +} + +/// The settings that a test needs from the client. +struct Config { + max_wait_time: Duration, + max_buffered: usize, + max_message_size: u64, + with_success_handler: bool, +} + +impl Default for Config { + fn default() -> Self { + Self { + // Long enough that only an explicit trigger sends a batch. A test + // that wants the timer sets a short time itself. + max_wait_time: Duration::seconds(30), + max_buffered: 64, + max_message_size: 1024 * 1024, + with_success_handler: true, + } + } +} + +struct Harness { + client: Arc, + mock: Arc, + reports: mpsc::UnboundedReceiver, + started: mpsc::UnboundedReceiver, +} + +impl Harness { + /// Waits for the next delivery report. + async fn next_report(&mut self) -> Report { + self.reports + .next() + .await + .expect("a delivery report was expected") + } + + /// Waits for the next `count` delivery reports. + async fn next_reports(&mut self, count: usize) -> Vec { + let mut reports = Vec::with_capacity(count); + for _ in 0..count { + reports.push(self.next_report().await); + } + reports + } + + /// Waits until the mock starts a send on any partition. + async fn next_started(&mut self) -> String { + self.started.next().await.expect("a send was expected") + } +} + +async fn harness(partitions: &[&str], config: Config) -> Harness { + let (mock, started) = + MockSendClient::with_max_message_size(partitions, config.max_message_size); + let (report_tx, reports) = mpsc::unbounded(); + + let mut builder = BufferedProducerClient::builder() + .with_max_wait_time(config.max_wait_time) + .with_max_buffered_event_count_per_partition(config.max_buffered); + + if config.with_success_handler { + let tx = report_tx.clone(); + builder = builder.with_on_send_succeeded(move |context| { + let tx = tx.clone(); + async move { + let _ = tx.unbounded_send(Report::Succeeded { + partition_id: context.partition_id, + bodies: bodies_of(&context.events), + }); + } + }); + } + + let tx = report_tx; + let client = builder + .with_on_send_failed(move |context| { + let tx = tx.clone(); + async move { + let _ = tx.unbounded_send(Report::Failed { + partition_id: context.partition_id, + bodies: bodies_of(&context.events), + error: context.error.to_string(), + }); + } + }) + .open_with_send_client(mock.clone() as Arc) + .await + .expect("the client opened"); + + Harness { + client: Arc::new(client), + mock, + reports, + started, + } +} + +fn to_partition(partition_id: &str) -> Option { + Some(EnqueueEventOptions { + partition_id: Some(partition_id.to_string()), + ..Default::default() + }) +} + +// 1. A full batch sends immediately. +#[tokio::test] +async fn full_batch_sends_without_waiting_for_the_timer() { + // The wait time is 30 seconds, so only a full batch can trigger this send. + let mut h = harness( + &["0"], + Config { + max_buffered: 2, + ..Default::default() + }, + ) + .await; + + h.client + .enqueue_events(vec!["a", "b"], to_partition("0")) + .await + .unwrap(); + + let report = h.next_report().await; + assert!(report.is_success()); + assert_eq!(report.bodies(), ["a", "b"]); + h.client.close().await.unwrap(); +} + +// 1b. A batch also sends when the next event does not fit the maximum message size. +#[tokio::test] +async fn batch_sends_when_the_next_event_does_not_fit() { + // One 1000 byte event fits. Two do not. + let mut h = harness( + &["0"], + Config { + max_message_size: 2500, + ..Default::default() + }, + ) + .await; + + let first = "a".repeat(1000); + let second = "b".repeat(1000); + h.client + .enqueue_event(first.clone(), to_partition("0")) + .await + .unwrap(); + h.client + .enqueue_event(second.clone(), to_partition("0")) + .await + .unwrap(); + + // The second event does not fit, so the client sends the first on its own. + let report = h.next_report().await; + assert!(report.is_success()); + assert_eq!(report.bodies(), [first]); + + h.client.flush().await.unwrap(); + let report = h.next_report().await; + assert_eq!(report.bodies(), [second]); + + assert_eq!(h.mock.sends().len(), 2); + h.client.close().await.unwrap(); +} + +// 2. A partial batch sends after the maximum wait time. +#[tokio::test] +async fn partial_batch_sends_after_the_maximum_wait_time() { + let mut h = harness( + &["0"], + Config { + max_wait_time: Duration::milliseconds(50), + ..Default::default() + }, + ) + .await; + + h.client + .enqueue_event("only", to_partition("0")) + .await + .unwrap(); + + // No flush and no close. Only the timer can send this batch. + let report = h.next_report().await; + assert!(report.is_success()); + assert_eq!(report.bodies(), ["only"]); + h.client.close().await.unwrap(); +} + +// 3. An explicit partition ID routes the event to that partition. +#[tokio::test] +async fn explicit_partition_id_routes_the_event() { + let mut h = harness(&["0", "1", "2"], Config::default()).await; + + h.client + .enqueue_event("routed", to_partition("2")) + .await + .unwrap(); + h.client.flush().await.unwrap(); + + let report = h.next_report().await; + assert_eq!(report.partition_id(), "2"); + assert_eq!(h.mock.sends()[0].partition_id, "2"); + h.client.close().await.unwrap(); +} + +// 3b. An unknown partition ID is an error. +#[tokio::test] +async fn unknown_partition_id_is_rejected() { + let h = harness(&["0", "1"], Config::default()).await; + + let error = h + .client + .enqueue_event("nowhere", to_partition("9")) + .await + .unwrap_err(); + assert!(error.to_string().contains("no partition")); + h.client.close().await.unwrap(); +} + +// 4. A partition key routes every event with that key to one partition. +#[tokio::test] +async fn partition_key_routes_every_event_to_one_partition() { + let mut h = harness(&["0", "1", "2", "3"], Config::default()).await; + + let options = Some(EnqueueEventOptions { + partition_key: Some("customer-17".to_string()), + ..Default::default() + }); + h.client + .enqueue_events(vec!["a", "b", "c"], options) + .await + .unwrap(); + h.client.flush().await.unwrap(); + + let report = h.next_report().await; + assert_eq!(report.bodies(), ["a", "b", "c"]); + + // The resolver decides the partition, and every event went to that one. + let expected = h.client.resolver.assign_for_key("customer-17"); + assert_eq!(report.partition_id(), expected); + h.client.close().await.unwrap(); +} + +// 5. Automatic assignment spreads events over the partitions in round-robin order. +#[tokio::test] +async fn automatic_assignment_uses_round_robin() { + let h = harness( + &["0", "1", "2", "3"], + Config { + // One event for each batch, so each event is its own send. + max_buffered: 1, + ..Default::default() + }, + ) + .await; + + for index in 0..8 { + h.client + .enqueue_event(format!("e{index}"), None) + .await + .unwrap(); + } + h.client.close().await.unwrap(); + + let sends = h.mock.sends(); + assert_eq!(sends.len(), 8); + for partition_id in ["0", "1", "2", "3"] { + let count = sends + .iter() + .filter(|s| s.partition_id == partition_id) + .count(); + assert_eq!(count, 2, "partition {partition_id} did not get two events"); + } +} + +// 6. A request that sets a partition ID and a partition key is rejected. +#[tokio::test] +async fn conflicting_routing_options_are_rejected() { + let h = harness(&["0", "1"], Config::default()).await; + + let options = Some(EnqueueEventOptions { + partition_id: Some("0".to_string()), + partition_key: Some("a-key".to_string()), + }); + + let error = h + .client + .enqueue_event("conflict", options.clone()) + .await + .unwrap_err(); + assert!(error.to_string().contains("not both")); + + let error = h + .client + .enqueue_events(vec!["conflict"], options) + .await + .unwrap_err(); + assert!(error.to_string().contains("not both")); + + assert_eq!(h.client.total_buffered_event_count(), 0); + h.client.close().await.unwrap(); +} + +// 7. Events keep the enqueue order inside one partition. +#[tokio::test] +async fn events_keep_their_order_inside_a_partition() { + let mut h = harness(&["0"], Config::default()).await; + + for index in 0..10 { + h.client + .enqueue_event(format!("e{index}"), to_partition("0")) + .await + .unwrap(); + } + h.client.flush().await.unwrap(); + + let report = h.next_report().await; + let expected: Vec = (0..10).map(|index| format!("e{index}")).collect(); + assert_eq!(report.bodies(), expected.as_slice()); + h.client.close().await.unwrap(); +} + +// 8. Two partitions publish at the same time. +#[tokio::test] +async fn different_partitions_publish_concurrently() { + let mut h = harness( + &["0", "1"], + Config { + max_buffered: 1, + ..Default::default() + }, + ) + .await; + + // Hold the first send on each partition. + let release_zero = h.mock.gate("0"); + let release_one = h.mock.gate("1"); + + h.client + .enqueue_event("a", to_partition("0")) + .await + .unwrap(); + h.client + .enqueue_event("b", to_partition("1")) + .await + .unwrap(); + + // Both sends start while both are held, so one partition does not block the + // other. + let first = h.next_started().await; + let second = h.next_started().await; + let started: HashSet = [first, second].into_iter().collect(); + assert_eq!( + started.len(), + 2, + "both partitions should have a send in flight" + ); + + let _ = release_zero.send(()); + let _ = release_one.send(()); + + let reports = h.next_reports(2).await; + assert!(reports.iter().all(|r| r.is_success())); + h.client.close().await.unwrap(); +} + +// 9. A full buffer makes an enqueue wait. +#[tokio::test] +async fn a_full_buffer_applies_backpressure() { + let mut h = harness( + &["0"], + Config { + max_buffered: 1, + ..Default::default() + }, + ) + .await; + + // Hold the worker inside its first send, so it stops draining the queue. + let _release = h.mock.gate("0"); + h.client + .enqueue_event("first", to_partition("0")) + .await + .unwrap(); + let _ = h.next_started().await; + + // The queue now fills, and the enqueue cannot finish. + let pending = h + .client + .enqueue_events(vec!["a", "b", "c", "d", "e", "f"], to_partition("0")); + pin_mut!(pending); + assert!( + poll!(pending.as_mut()).is_pending(), + "the enqueue should wait for space in a full buffer" + ); + + // A close must not wait for the enqueue that is still parked. + h.client.abort().await.unwrap(); +} + +// 10. A waiting enqueue continues once the buffer has space. +#[tokio::test] +async fn a_waiting_enqueue_resumes_when_space_appears() { + let mut h = harness( + &["0"], + Config { + max_buffered: 1, + ..Default::default() + }, + ) + .await; + + let release = h.mock.gate("0"); + h.client + .enqueue_event("first", to_partition("0")) + .await + .unwrap(); + let _ = h.next_started().await; + + let pending = h + .client + .enqueue_events(vec!["a", "b", "c", "d", "e", "f"], to_partition("0")); + pin_mut!(pending); + assert!(poll!(pending.as_mut()).is_pending()); + + // Letting the send finish drains the queue, so the enqueue continues. + let _ = release.send(()); + pending.await.unwrap(); + + h.client.close().await.unwrap(); + assert_eq!(h.mock.total_events(), 7); + assert_eq!(h.client.total_buffered_event_count(), 0); +} + +// 11. A waiting enqueue fails once the client closes. +#[tokio::test] +async fn a_waiting_enqueue_fails_when_the_client_closes() { + let mut h = harness( + &["0"], + Config { + max_buffered: 1, + ..Default::default() + }, + ) + .await; + + let _release = h.mock.gate("0"); + h.client + .enqueue_event("first", to_partition("0")) + .await + .unwrap(); + let _ = h.next_started().await; + + let client = h.client.clone(); + let pending = client.enqueue_events(vec!["a", "b", "c", "d", "e", "f"], to_partition("0")); + pin_mut!(pending); + assert!(poll!(pending.as_mut()).is_pending()); + + // The close signal must wake the waiting enqueue with an error. + let closing = h.client.abort(); + pin_mut!(closing); + let _ = poll!(closing.as_mut()); + + let error = pending.await.unwrap_err(); + assert!(error.to_string().contains("closed")); + + closing.await.unwrap(); +} + +// 12. One event that is larger than an empty batch fails on its own. +#[tokio::test] +async fn an_oversized_event_produces_one_failure() { + let mut h = harness( + &["0"], + Config { + max_message_size: 200, + ..Default::default() + }, + ) + .await; + + let big = "x".repeat(4096); + h.client + .enqueue_event(big, to_partition("0")) + .await + .unwrap(); + + let report = h.next_report().await; + assert!(!report.is_success()); + assert_eq!(report.bodies().len(), 1); + assert!(report.bodies()[0].starts_with("xxxx")); + + // The worker stays healthy, so a later event still goes out. + h.client + .enqueue_event("small", to_partition("0")) + .await + .unwrap(); + h.client.flush().await.unwrap(); + let report = h.next_report().await; + assert!(report.is_success()); + assert_eq!(report.bodies(), ["small"]); + + assert_eq!(h.client.total_buffered_event_count(), 0); + h.client.close().await.unwrap(); +} + +// 13. A retryable failure reaches the client only after the retry policy runs. +// +// The retry policy sits below the send seam, inside `RecoverableSender`. The +// tests for the policy itself are in `common/retry.rs`, and the live +// forced-error test covers recovery from end to end. This test states the +// contract that the worker depends on: one error from the seam is already +// terminal, so the worker reports one failure and does not try again. +#[tokio::test] +async fn a_terminal_error_is_not_retried_by_the_worker() { + let mut h = harness(&["0"], Config::default()).await; + + h.mock.push_outcome("0", SendScript::Error("server busy")); + + h.client + .enqueue_event("a", to_partition("0")) + .await + .unwrap(); + h.client.flush().await.unwrap(); + + let report = h.next_report().await; + assert!(!report.is_success()); + + // Exactly one send attempt reached the seam. + assert_eq!(h.mock.sends().len(), 1); + h.client.close().await.unwrap(); +} + +// 14. An exhausted retry produces exactly one failure result. +#[tokio::test] +async fn retry_exhaustion_produces_one_failure_result() { + let mut h = harness(&["0"], Config::default()).await; + + h.mock + .push_outcome("0", SendScript::Error("retries exhausted")); + + h.client + .enqueue_events(vec!["a", "b", "c"], to_partition("0")) + .await + .unwrap(); + h.client.flush().await.unwrap(); + + let report = h.next_report().await; + match report { + Report::Failed { bodies, error, .. } => { + assert_eq!(bodies, ["a", "b", "c"]); + assert!(error.contains("retries exhausted")); + } + other => panic!("expected a failure report, got {other:?}"), + } + + // The client reports the batch one time only. + h.client.close().await.unwrap(); + assert!(h.reports.next().await.is_none()); +} + +// 15. A batch that the service accepts produces one success result. +#[tokio::test] +async fn a_successful_batch_produces_one_success_result() { + let mut h = harness(&["0"], Config::default()).await; + + h.client + .enqueue_events(vec!["a", "b"], to_partition("0")) + .await + .unwrap(); + h.client.flush().await.unwrap(); + + let report = h.next_report().await; + assert!(report.is_success()); + assert_eq!(report.bodies(), ["a", "b"]); + + h.client.close().await.unwrap(); + assert!(h.reports.next().await.is_none()); +} + +// 16. Modified, Released, and Rejected are never a success. +#[tokio::test] +async fn modified_released_and_rejected_are_not_reported_as_success() { + for outcome in [ + SendScript::Modified, + SendScript::Released, + SendScript::Rejected, + ] { + let mut h = harness(&["0"], Config::default()).await; + h.mock.push_outcome("0", outcome.clone()); + + h.client + .enqueue_event("a", to_partition("0")) + .await + .unwrap(); + h.client.flush().await.unwrap(); + + let report = h.next_report().await; + assert!( + !report.is_success(), + "outcome {outcome:?} must not be a success" + ); + assert_eq!(report.bodies(), ["a"]); + + h.client.close().await.unwrap(); + } +} + +// 16b. A Modified or Released outcome carries the SendNotAccepted error kind. +#[tokio::test] +async fn a_not_accepted_outcome_uses_the_send_not_accepted_error() { + let mut h = harness(&["0"], Config::default()).await; + h.mock.push_outcome("0", SendScript::Modified); + + h.client + .enqueue_event("a", to_partition("0")) + .await + .unwrap(); + h.client.flush().await.unwrap(); + + match h.next_report().await { + Report::Failed { error, .. } => { + assert!(error.contains("not durably accepted"), "got {error}"); + } + other => panic!("expected a failure report, got {other:?}"), + } + h.client.close().await.unwrap(); +} + +// 17. A flush waits for the events that the client accepted before the barrier. +#[tokio::test] +async fn flush_waits_for_events_accepted_before_the_barrier() { + let mut h = harness(&["0"], Config::default()).await; + + let release = h.mock.gate("0"); + h.client + .enqueue_event("before", to_partition("0")) + .await + .unwrap(); + + let client = h.client.clone(); + let flush = client.flush(); + pin_mut!(flush); + assert!( + poll!(flush.as_mut()).is_pending(), + "the flush must wait for the held send" + ); + + let _ = release.send(()); + flush.await.unwrap(); + + let report = h.next_report().await; + assert!(report.is_success()); + assert_eq!(report.bodies(), ["before"]); + h.client.close().await.unwrap(); +} + +// 18. An event that arrives after the barrier does not delay that flush. +#[tokio::test] +async fn events_after_the_barrier_do_not_delay_the_flush() { + let mut h = harness(&["0"], Config::default()).await; + + h.client + .enqueue_event("before", to_partition("0")) + .await + .unwrap(); + + let client = h.client.clone(); + let flush = client.flush(); + pin_mut!(flush); + // The first poll puts the barrier into the queue. + let _ = poll!(flush.as_mut()); + + // This event sits behind the barrier. + h.client + .enqueue_event("after", to_partition("0")) + .await + .unwrap(); + + flush.await.unwrap(); + + // The flush covered the first event only. The wait time is 30 seconds, so + // the second event is still in the buffer. + let report = h.next_report().await; + assert_eq!(report.bodies(), ["before"]); + assert_eq!(h.client.total_buffered_event_count(), 1); + + h.client.close().await.unwrap(); +} + +// 19. Two flush calls at the same time both complete. +#[tokio::test] +async fn concurrent_flush_calls_both_complete() { + let mut h = harness(&["0", "1"], Config::default()).await; + + h.client + .enqueue_event("a", to_partition("0")) + .await + .unwrap(); + h.client + .enqueue_event("b", to_partition("1")) + .await + .unwrap(); + + let (first, second) = futures::join!(h.client.flush(), h.client.flush()); + first.unwrap(); + second.unwrap(); + + assert_eq!(h.client.total_buffered_event_count(), 0); + let reports = h.next_reports(2).await; + assert!(reports.iter().all(|r| r.is_success())); + h.client.close().await.unwrap(); +} + +// 20. A graceful close sends the buffered events. +#[tokio::test] +async fn graceful_close_sends_buffered_events() { + let mut h = harness(&["0", "1"], Config::default()).await; + + h.client + .enqueue_events(vec!["a", "b", "c"], to_partition("0")) + .await + .unwrap(); + h.client + .enqueue_event("d", to_partition("1")) + .await + .unwrap(); + + // No flush. The close must send them. + h.client.close().await.unwrap(); + + assert_eq!(h.mock.total_events(), 4); + assert_eq!(h.client.total_buffered_event_count(), 0); + + let mut seen = Vec::new(); + while let Some(report) = h.reports.next().await { + assert!(report.is_success()); + seen.extend(report.bodies().to_vec()); + } + seen.sort(); + assert_eq!(seen, ["a", "b", "c", "d"]); +} + +// 21. An immediate close abandons the buffered events and clears the counts. +#[tokio::test] +async fn immediate_close_abandons_buffered_events() { + let h = harness(&["0"], Config::default()).await; + + h.client + .enqueue_events(vec!["a", "b", "c"], to_partition("0")) + .await + .unwrap(); + assert_eq!(h.client.total_buffered_event_count(), 3); + + h.client.abort().await.unwrap(); + + // The events never reached the service, and the counts agree. + assert_eq!(h.mock.total_events(), 0); + assert_eq!(h.client.total_buffered_event_count(), 0); + assert_eq!(h.client.buffered_event_count("0"), 0); +} + +// 22. A close on an idle client completes. +#[tokio::test] +async fn closing_an_idle_client_completes() { + let h = harness(&["0", "1", "2"], Config::default()).await; + + h.client.close().await.unwrap(); + + // A second close does nothing and still succeeds. + h.client.close().await.unwrap(); + assert_eq!(h.client.total_buffered_event_count(), 0); +} + +// 22b. An enqueue after a close is rejected. +#[tokio::test] +async fn enqueue_after_close_is_rejected() { + let h = harness(&["0"], Config::default()).await; + h.client.close().await.unwrap(); + + let error = h.client.enqueue_event("late", None).await.unwrap_err(); + assert!(error.to_string().contains("closed")); +} + +// 23. A worker shutdown releases every reference that the worker held. +#[tokio::test] +async fn worker_shutdown_releases_references() { + let h = harness(&["0", "1", "2"], Config::default()).await; + + // The workers each hold a reference to the send client. + assert!(Arc::strong_count(&h.mock) > 1); + + h.client.close().await.unwrap(); + + assert_eq!( + Arc::strong_count(&h.mock), + 1, + "the workers still hold a reference to the send client" + ); +} + +// 24. A terminal failure does not stop the worker, and no event is lost or sent +// two times. +// +// A real connection recovery happens below the send seam, inside +// `RecoverableSender`. The live forced-error test covers that path. This test +// covers the part that the buffered producer owns: the worker survives a +// terminal failure and keeps serving its queue. +#[tokio::test] +async fn a_worker_keeps_serving_after_a_terminal_failure() { + let mut h = harness( + &["0"], + Config { + max_buffered: 1, + ..Default::default() + }, + ) + .await; + + h.mock.push_outcome("0", SendScript::Error("link detached")); + + h.client + .enqueue_event("lost", to_partition("0")) + .await + .unwrap(); + let first = h.next_report().await; + assert!(!first.is_success()); + assert_eq!(first.bodies(), ["lost"]); + + for index in 0..3 { + h.client + .enqueue_event(format!("after{index}"), to_partition("0")) + .await + .unwrap(); + } + h.client.close().await.unwrap(); + + let mut delivered = Vec::new(); + while let Some(report) = h.reports.next().await { + if report.is_success() { + delivered.extend(report.bodies().to_vec()); + } + } + delivered.sort(); + assert_eq!(delivered, ["after0", "after1", "after2"]); + + // Four sends in total: the one that failed, and the three that followed. + assert_eq!(h.mock.sends().len(), 4); +} + +// 25. Many enqueues at the same time neither lose nor repeat an event. +#[tokio::test] +async fn concurrent_enqueues_do_not_lose_or_repeat_events() { + let mut h = harness(&["0", "1", "2", "3"], Config::default()).await; + + let mut tasks = Vec::new(); + for task_index in 0..4 { + let client = h.client.clone(); + tasks.push(tokio::spawn(async move { + for event_index in 0..25 { + client + .enqueue_event(format!("t{task_index}-e{event_index}"), None) + .await + .unwrap(); + } + })); + } + for task in tasks { + task.await.unwrap(); + } + + h.client.close().await.unwrap(); + + assert_eq!(h.mock.total_events(), 100); + + let mut delivered = Vec::new(); + while let Some(report) = h.reports.next().await { + assert!(report.is_success()); + delivered.extend(report.bodies().to_vec()); + } + + let unique: HashSet<&String> = delivered.iter().collect(); + assert_eq!(delivered.len(), 100, "an event was lost or repeated"); + assert_eq!(unique.len(), 100, "an event was repeated"); + assert_eq!(h.client.total_buffered_event_count(), 0); +} + +// 26. Dropping the client stops the background work. +#[tokio::test] +async fn dropping_the_client_cancels_background_work() { + let mut h = harness(&["0", "1"], Config::default()).await; + + h.client + .enqueue_event("a", to_partition("0")) + .await + .unwrap(); + + drop(h.client); + + // Every worker ends, so it releases the send client and the handlers. The + // report stream ends once the last handler reference drops. The test holds + // the only other reference to the send client. + while h.reports.next().await.is_some() {} + assert_eq!( + Arc::strong_count(&h.mock), + 1, + "a worker outlived the client that was dropped" + ); +} + +// A client with no failure handler cannot open. +#[tokio::test] +async fn a_missing_failure_handler_is_an_error() { + let (mock, _started) = MockSendClient::new(&["0"]); + let error = BufferedProducerClient::builder() + .open_with_send_client(mock as Arc) + .await + .unwrap_err(); + assert!(error.to_string().contains("with_on_send_failed")); +} + +// The builder rejects settings that cannot work. +#[tokio::test] +async fn the_builder_rejects_invalid_settings() { + let (mock, _started) = MockSendClient::new(&["0"]); + let error = BufferedProducerClient::builder() + .with_max_buffered_event_count_per_partition(0) + .with_on_send_failed(|_| async {}) + .open_with_send_client(mock.clone() as Arc) + .await + .unwrap_err(); + assert!(error.to_string().contains("at least 1")); + + let error = BufferedProducerClient::builder() + .with_max_wait_time(Duration::ZERO) + .with_on_send_failed(|_| async {}) + .open_with_send_client(mock as Arc) + .await + .unwrap_err(); + assert!(error.to_string().contains("longer than zero")); +} + +// The per-partition count follows the total count. +#[tokio::test] +async fn buffered_counts_track_each_partition() { + let h = harness(&["0", "1"], Config::default()).await; + + h.client + .enqueue_events(vec!["a", "b"], to_partition("0")) + .await + .unwrap(); + h.client + .enqueue_event("c", to_partition("1")) + .await + .unwrap(); + + assert_eq!(h.client.total_buffered_event_count(), 3); + assert_eq!(h.client.buffered_event_count("0"), 2); + assert_eq!(h.client.buffered_event_count("1"), 1); + assert_eq!(h.client.buffered_event_count("unknown"), 0); + + h.client.close().await.unwrap(); + assert_eq!(h.client.total_buffered_event_count(), 0); +} + +// 27. A worker that nothing cancels still abandons its active batch. +// +// `AbortableTask::abort` cancels a task on a runtime that supports it. On the +// standard thread runtime it only detaches the thread, so the worker keeps +// running and reaches the end of its queue, which is the same path as a +// graceful close. This test drives the worker with no cancellation at all, so +// only the abandon flag can stop the batch from going to the service. +#[tokio::test] +async fn an_abandoning_worker_that_nothing_cancels_does_not_publish() { + let (mock, _started) = MockSendClient::with_max_message_size(&["0"], 1024 * 1024); + + let (sender, receiver) = mpsc::unbounded(); + let buffered = Arc::new(AtomicUsize::new(0)); + let total_buffered = Arc::new(AtomicUsize::new(0)); + let abandon = Arc::new(AtomicBool::new(false)); + let (stopped_tx, stopped_rx) = oneshot::channel(); + + let reported = Arc::new(AtomicUsize::new(0)); + let counter = reported.clone(); + let handlers = DeliveryHandlers { + succeeded: None, + failed: Arc::new(move |_context| { + let counter = counter.clone(); + Box::pin(async move { + counter.fetch_add(1, Ordering::AcqRel); + }) + }), + }; + + let worker = PartitionWorker::new( + "0".to_string(), + receiver, + mock.clone() as Arc, + // Long enough that the timer cannot send the batch during the test. + Duration::seconds(30), + 64, + handlers, + buffered.clone(), + total_buffered.clone(), + abandon.clone(), + stopped_tx, + ); + + let capacity = Arc::new(Semaphore::new(4)); + let permit = capacity + .try_acquire_arc() + .expect("a new semaphore has capacity"); + let event = EventData::from("abandoned"); + buffered.fetch_add(1, Ordering::AcqRel); + total_buffered.fetch_add(1, Ordering::AcqRel); + sender + .unbounded_send(Command::Event { + message: Box::new(AmqpMessage::from(event.clone())), + event, + permit, + }) + .expect("the worker holds the receiver"); + + let run = worker.run(); + pin_mut!(run); + + // One poll takes the event into the active batch. The batch is not full and + // the wait time is long, so the worker then waits for the next command. + assert!(poll!(&mut run).is_pending()); + + // Abandon the events and end the queue, exactly as an immediate close does. + abandon.store(true, Ordering::Release); + drop(sender); + + // Nothing cancels this future, so the worker runs its close path in full. + run.await; + + assert_eq!( + mock.total_events(), + 0, + "an immediate close promised to drop these events, so none may reach the service" + ); + assert_eq!( + reported.load(Ordering::Acquire), + 0, + "an abandoned event has no delivery outcome to report" + ); + assert!( + stopped_rx.await.is_ok(), + "the worker must tell the client that it stopped" + ); +} + +// 28. The buffered counts stay sane when every event reaches a terminal outcome +// as soon as the worker sees it. +// +// An oversized event fails inside the worker without a send. The client counts +// the event before it publishes the command, so that failure can never +// decrement a count that is still zero and wrap it to `usize::MAX`. The +// interleaving that this guards against is a race, so this test does not +// reproduce it on demand; it states the invariant and exercises the path with +// many events on more than one thread. +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn a_fast_terminal_outcome_does_not_wrap_the_buffered_counts() { + const EVENT_COUNT: usize = 200; + + let mut h = harness( + &["0"], + Config { + max_message_size: 200, + with_success_handler: false, + ..Default::default() + }, + ) + .await; + + let big = "x".repeat(4096); + for _ in 0..EVENT_COUNT { + h.client + .enqueue_event(big.clone(), to_partition("0")) + .await + .unwrap(); + // A wrapped count is astronomically large, and the real count can never + // pass the number of events that the test enqueued. + assert!( + h.client.total_buffered_event_count() <= EVENT_COUNT, + "the buffered count wrapped: {}", + h.client.total_buffered_event_count() + ); + } + + for _ in 0..EVENT_COUNT { + assert!(!h.next_report().await.is_success()); + } + + h.client.close().await.unwrap(); + assert_eq!(h.client.total_buffered_event_count(), 0); + assert_eq!(h.client.buffered_event_count("0"), 0); +} + +// 29. A delivery handler can enqueue again without deadlocking the worker. +// +// The handler runs on the worker task, and the worker is the only thing that +// returns capacity permits. While the worker held the permits of the batch it +// was reporting, a handler that enqueued to the same partition waited for a +// permit that only the worker could return, and the worker waited for the +// handler. The event is already at a terminal outcome when the handler runs, so +// the permit goes back first. +#[tokio::test] +async fn a_failure_handler_can_enqueue_again() { + use std::sync::{OnceLock, Weak}; + + // One permit for the partition, so the retry can only proceed if the + // failing event already gave its permit back. + const BUFFER: usize = 1; + + let (mock, _started) = MockSendClient::with_max_message_size(&["0"], 200); + + let slot: Arc>> = Arc::new(OnceLock::new()); + let retried = Arc::new(AtomicUsize::new(0)); + + let for_handler = slot.clone(); + let counter = retried.clone(); + let client = BufferedProducerClient::builder() + .with_max_wait_time(Duration::seconds(30)) + .with_max_buffered_event_count_per_partition(BUFFER) + .with_on_send_failed(move |_context| { + let slot = for_handler.clone(); + let counter = counter.clone(); + async move { + // Retry once. A second retry would recurse without an end. + if counter.fetch_add(1, Ordering::AcqRel) > 0 { + return; + } + let client = slot + .get() + .expect("the test sets the client before it enqueues") + .upgrade() + .expect("the client is alive while the handler runs"); + client + .enqueue_event("small", to_partition("0")) + .await + .expect("the retry must not be rejected"); + } + }) + .open_with_send_client(mock.clone() as Arc) + .await + .expect("the client opened"); + + let client = Arc::new(client); + slot.set(Arc::downgrade(&client)).expect("set once"); + + // Too large for the link, so the worker fails it without a send. That is + // the path that calls the handler while it still holds the permit. + let oversized = "x".repeat(4096); + client + .enqueue_event(oversized, to_partition("0")) + .await + .unwrap(); + + // The deadlock shows up as a hang, so bound it. + tokio::time::timeout(std::time::Duration::from_secs(10), async { + while retried.load(Ordering::Acquire) == 0 { + tokio::task::yield_now().await; + } + client.flush().await + }) + .await + .expect("the failure handler deadlocked with the worker") + .expect("the flush completed"); + + assert_eq!( + mock.total_events(), + 1, + "the event that the handler enqueued must reach the service" + ); + + client.close().await.unwrap(); +} + +/// Live tests for the buffered producer. +/// +/// These tests need a real Event Hub. They live in the crate because +/// `ProducerClient::force_error` is only available to crate tests. +#[cfg(test)] +mod live { + use crate::{ + common::tests::force_errors, BufferedProducerClient, EnqueueEventOptions, + SendBatchFailedContext, SendBatchSucceededContext, + }; + use azure_core::time::Duration; + use azure_core_amqp::{error::AmqpErrorKind, AmqpError}; + use azure_core_test::{recorded, TestContext}; + use std::{ + collections::HashSet, + sync::{Arc, Mutex}, + }; + + /// Records which event bodies reached a terminal outcome. + #[derive(Default)] + struct Outcomes { + succeeded: Mutex>, + failed: Mutex>, + } + + fn bodies(events: &[crate::models::EventData]) -> Vec { + events + .iter() + .map(|event| String::from_utf8_lossy(event.body().unwrap_or_default()).into_owned()) + .collect() + } + + /// A connection recovery must not lose an event, and it must not report an + /// event two times. + #[recorded::test(live)] + async fn buffered_recovery_keeps_every_event(ctx: TestContext) -> crate::Result<()> { + const TEST_NAME: &str = "buffered_recovery_keeps_every_event"; + const PARTITION: &str = "1"; + + let recording = ctx.recording(); + let host = recording.var("EVENTHUBS_HOST", None); + let eventhub = recording.var("EVENTHUB_NAME", None); + let credential = recording.credential(); + + let outcomes = Arc::new(Outcomes::default()); + let for_success = outcomes.clone(); + let for_failure = outcomes.clone(); + + let producer = Arc::new( + BufferedProducerClient::builder() + .with_application_id(TEST_NAME.to_string()) + .with_max_wait_time(Duration::milliseconds(200)) + .with_on_send_succeeded(move |context: SendBatchSucceededContext| { + let outcomes = for_success.clone(); + async move { + outcomes + .succeeded + .lock() + .unwrap() + .extend(bodies(&context.events)); + } + }) + .with_on_send_failed(move |context: SendBatchFailedContext| { + let outcomes = for_failure.clone(); + async move { + outcomes + .failed + .lock() + .unwrap() + .extend(bodies(&context.events)); + } + }) + .open(host.as_str(), eventhub.as_str(), credential.clone()) + .await?, + ); + + let enqueued = Arc::new(Mutex::new(Vec::::new())); + let for_test = enqueued.clone(); + + force_errors( + producer.clone(), + move |producer: Arc| { + let enqueued = for_test.clone(); + async move { + let mut index = 0usize; + loop { + let body = format!("recovery-{index}"); + // An enqueue can fail only when the client closes. + if producer + .enqueue_event( + body.clone(), + Some(EnqueueEventOptions { + partition_id: Some(PARTITION.to_string()), + ..Default::default() + }), + ) + .await + .is_err() + { + break; + } + enqueued.lock().unwrap().push(body); + index += 1; + } + } + }, + |producer: Arc| { + // Break the link under the worker. The recoverable sender must + // rebuild it without the buffered producer losing an event. + producer + .inner_producer() + .expect("a live client has a producer") + .force_error(AmqpError::from(AmqpErrorKind::LinkClosedByRemote( + Box::new(azure_core::error::Error::new( + azure_core::error::ErrorKind::Other, + "Forced error", + )), + ))) + .unwrap(); + }, + Duration::seconds(5), + Duration::seconds(20), + ) + .await?; + + producer.close().await?; + + let enqueued = enqueued.lock().unwrap().clone(); + let succeeded = outcomes.succeeded.lock().unwrap().clone(); + let failed = outcomes.failed.lock().unwrap().clone(); + + // Every event that the client accepted reached exactly one terminal + // outcome. + let mut reported: Vec = succeeded.iter().chain(failed.iter()).cloned().collect(); + reported.sort(); + let unique: HashSet<&String> = reported.iter().collect(); + assert_eq!( + reported.len(), + unique.len(), + "an event reached a terminal outcome two times" + ); + assert_eq!( + reported.len(), + enqueued.len(), + "the client did not report every accepted event" + ); + assert_eq!(producer.total_buffered_event_count(), 0); + + Ok(()) + } +} diff --git a/sdk/eventhubs/azure_messaging_eventhubs/src/producer/buffered/worker.rs b/sdk/eventhubs/azure_messaging_eventhubs/src/producer/buffered/worker.rs new file mode 100644 index 00000000000..4e5d7b8d0ef --- /dev/null +++ b/sdk/eventhubs/azure_messaging_eventhubs/src/producer/buffered/worker.rs @@ -0,0 +1,489 @@ +// Copyright (c) Microsoft Corporation. All Rights reserved +// Licensed under the MIT license. + +use super::{ + send_client::BufferedSendClient, DeliveryHandlers, SendBatchFailedContext, + SendBatchSucceededContext, +}; +use crate::{ + error::ErrorKind, models::EventData, producer::batch::EventDataBatchInner, EventHubsError, +}; +use async_lock::SemaphoreGuardArc; +use azure_core::{sleep, time::Duration}; +use azure_core_amqp::{AmqpMessage, AmqpSendOutcome}; +use futures::{ + channel::{mpsc, oneshot}, + FutureExt, StreamExt, +}; +use std::{ + future::Future, + pin::Pin, + sync::{ + atomic::{AtomicBool, AtomicUsize, Ordering}, + Arc, + }, +}; +use tracing::{debug, trace, warn}; + +/// A message that the client sends to a partition worker. +pub(crate) enum Command { + /// One event to add to the active batch. + /// + /// The command carries the capacity permit for the event. The permit returns + /// to the partition semaphore once the event reaches a terminal outcome, so + /// the buffer bound covers every event that the client accepted. + Event { + event: EventData, + message: Box, + permit: SemaphoreGuardArc, + }, + + /// A flush barrier. + /// + /// The worker sends the active batch, then it completes the sender. The + /// position of this command in the queue is the barrier: the worker + /// processes every command in front of it first, and no command behind it + /// can delay it. + Flush(oneshot::Sender<()>), +} + +/// One event that the worker holds, with the capacity permit for the event. +type PendingEvent = (EventData, SemaphoreGuardArc); + +/// Publishes the events of one partition. +/// +/// One worker owns one partition. The worker is the only reader of the +/// partition queue, so the events keep the order that the caller enqueued them, +/// and only one send is active for the partition at a time. +pub(crate) struct PartitionWorker { + partition_id: String, + receiver: mpsc::UnboundedReceiver, + send_client: Arc, + max_wait_time: Duration, + max_events_per_batch: usize, + handlers: DeliveryHandlers, + buffered: Arc, + total_buffered: Arc, + abandon: Arc, + + /// Tells the client that this worker stopped. + /// + /// The worker sends on this channel when it leaves [`Self::run`]. A runtime + /// that drops the task instead cancels the channel. Either way the client + /// learns that the worker holds nothing more, without depending on what + /// [`AbortableTask::abort`] does on the runtime in use. + /// + /// [`AbortableTask::abort`]: azure_core::async_runtime::AbortableTask + stopped: Option>, +} + +impl PartitionWorker { + #[allow(clippy::too_many_arguments)] + pub(crate) fn new( + partition_id: String, + receiver: mpsc::UnboundedReceiver, + send_client: Arc, + max_wait_time: Duration, + max_events_per_batch: usize, + handlers: DeliveryHandlers, + buffered: Arc, + total_buffered: Arc, + abandon: Arc, + stopped: oneshot::Sender<()>, + ) -> Self { + Self { + partition_id, + receiver, + send_client, + max_wait_time, + max_events_per_batch, + handlers, + buffered, + total_buffered, + abandon, + stopped: Some(stopped), + } + } + + /// Runs the worker until the queue closes or the client abandons the events. + pub(crate) async fn run(mut self) { + debug!( + partition_id = %self.partition_id, + "Buffered producer partition worker started." + ); + + let mut batch: Option = None; + let mut pending: Vec = Vec::new(); + let mut timer: Option + Send>>> = None; + + loop { + if self.abandon.load(Ordering::Acquire) { + break; + } + + let command = match timer.as_mut() { + // The batch is empty. Wait for a command with no timer running. + None => self.receiver.next().await, + // The batch holds events. Send it when the wait time expires. + Some(deadline) => { + futures::select! { + command = self.receiver.next().fuse() => command, + _ = deadline.as_mut().fuse() => { + debug!( + partition_id = %self.partition_id, + "Maximum wait time expired; sending the partial batch." + ); + self.send_batch(&mut batch, &mut pending).await; + timer = None; + continue; + } + } + } + }; + + let Some(command) = command else { + // The client dropped every sender. Send what is left and stop. + // A client that abandons its events takes the same path, so + // `send_batch` must see the flag; it returns without a send. + debug!( + partition_id = %self.partition_id, + "Partition queue closed; draining the active batch." + ); + self.send_batch(&mut batch, &mut pending).await; + break; + }; + + match command { + Command::Event { + event, + message, + permit, + } => { + self.add_event( + &mut batch, + &mut pending, + &mut timer, + event, + *message, + permit, + ) + .await; + + if pending.len() >= self.max_events_per_batch { + debug!( + partition_id = %self.partition_id, + event_count = pending.len(), + "Batch reached the configured event count; sending it." + ); + self.send_batch(&mut batch, &mut pending).await; + timer = None; + } + } + Command::Flush(completed) => { + self.send_batch(&mut batch, &mut pending).await; + timer = None; + if self.abandon.load(Ordering::Acquire) { + // The client abandoned the events during this flush, so + // the barrier cannot report success. Dropping the sender + // cancels the waiter. + drop(completed); + break; + } + // A dropped receiver means the caller stopped waiting. + let _ = completed.send(()); + } + } + } + + if self.abandon.load(Ordering::Acquire) { + self.discard_remaining(pending).await; + } + + debug!( + partition_id = %self.partition_id, + "Buffered producer partition worker stopped." + ); + + // Tell the client that this worker holds nothing more. The client waits + // for this, so a close does not return while a worker still runs. + if let Some(stopped) = self.stopped.take() { + let _ = stopped.send(()); + } + } + + /// Adds one event to the active batch, and sends the batch when the event + /// does not fit. + #[allow(clippy::too_many_arguments)] + async fn add_event( + &self, + batch: &mut Option, + pending: &mut Vec, + timer: &mut Option + Send>>>, + event: EventData, + message: AmqpMessage, + permit: SemaphoreGuardArc, + ) { + // The batch needs the maximum message size of the link, so it is created + // on the first event. + if batch.is_none() { + match self.send_client.max_message_size(&self.partition_id).await { + Ok(max_size) => { + *batch = Some(EventDataBatchInner::new(max_size, None)); + } + Err(error) => { + warn!( + partition_id = %self.partition_id, + "Could not read the maximum message size; failing the event." + ); + self.fail_one(event, permit, error).await; + return; + } + } + } + let active = batch.as_mut().expect("the batch was just created"); + + match active.try_add(message.clone()) { + Ok(true) => { + trace!( + partition_id = %self.partition_id, + "Added an event to the active batch." + ); + pending.push((event, permit)); + if timer.is_none() { + *timer = Some(Box::pin(sleep(self.max_wait_time))); + } + } + Ok(false) if active.is_empty() => { + // One event on its own is larger than the whole batch. Fail that + // event only. The worker stays healthy. + self.fail_oversized(event, permit).await; + } + Ok(false) => { + // The batch is full. Send it, then start a new batch with this event. + debug!( + partition_id = %self.partition_id, + "The next event does not fit; sending the full batch." + ); + self.send_batch(batch, pending).await; + *timer = None; + + let active = batch.as_mut().expect("the batch is reused after a send"); + match active.try_add(message) { + Ok(true) => { + pending.push((event, permit)); + *timer = Some(Box::pin(sleep(self.max_wait_time))); + } + Ok(false) => self.fail_oversized(event, permit).await, + Err(error) => self.fail_one(event, permit, error).await, + } + } + Err(error) => { + warn!( + partition_id = %self.partition_id, + "Could not add an event to the batch; failing that event." + ); + self.fail_one(event, permit, error).await; + } + } + } + + async fn fail_oversized(&self, event: EventData, permit: SemaphoreGuardArc) { + warn!( + partition_id = %self.partition_id, + "An event is too large for an empty batch; failing that event." + ); + self.fail_one( + event, + permit, + EventHubsError::with_message( + "The event is too large for the maximum message size of the link.", + ), + ) + .await; + } + + /// Reports one event as failed and returns its capacity permit. + async fn fail_one(&self, event: EventData, permit: SemaphoreGuardArc, error: EventHubsError) { + // Give the capacity back before the handler runs. The outcome is + // already terminal, and the handler runs on this task, so a handler + // that enqueues to this partition would otherwise wait for a permit + // that only this task can return. + self.release(1); + drop(permit); + self.report_failure(vec![event], error).await; + } + + /// Sends the active batch and reports exactly one outcome for it. + /// + /// The method does nothing when the batch holds no events, so the worker + /// never sends an empty batch. It also does nothing once the client + /// abandons its events: every path that ends the worker calls this method, + /// and an immediate close must not publish what it promised to drop. The + /// events stay in `pending`, and [`Self::discard_remaining`] drops them. + async fn send_batch( + &self, + batch: &mut Option, + pending: &mut Vec, + ) { + let Some(active) = batch.as_mut() else { + return; + }; + if active.is_empty() { + return; + } + if self.abandon.load(Ordering::Acquire) { + debug!( + partition_id = %self.partition_id, + event_count = pending.len(), + "The client abandoned its events; not sending the active batch." + ); + return; + } + + let batch_size_in_bytes = active.size(); + let envelope = active.take_envelope(); + let (events, permits): (Vec, Vec) = + std::mem::take(pending).into_iter().unzip(); + let event_count = events.len(); + + debug!( + partition_id = %self.partition_id, + event_count, + batch_size_in_bytes, + "Sending a batch of events." + ); + + let outcome = self + .send_client + .send_envelope(&self.partition_id, envelope) + .await; + + // The send settled, so every event in the batch is at a terminal + // outcome whatever that outcome is. Give the capacity back before the + // handlers run. The handlers run on this task, so a handler that + // enqueues to this partition would otherwise wait for a permit that + // only this task can return. + self.release(event_count); + drop(permits); + + match outcome { + Ok(AmqpSendOutcome::Accepted) => { + debug!( + partition_id = %self.partition_id, + event_count, + "The service accepted the batch." + ); + self.report_success(events).await; + } + Ok(AmqpSendOutcome::Modified(reason)) => { + // Modified does not mean that the service stored the events, so + // this is a delivery failure, not a success. + warn!( + partition_id = %self.partition_id, + event_count, + modification = ?reason, + "The service modified the batch; it did not durably accept it." + ); + self.report_failure( + events, + EventHubsError::from(ErrorKind::SendNotAccepted( + "the service returned a Modified outcome".into(), + )), + ) + .await; + } + Ok(AmqpSendOutcome::Released) => { + warn!( + partition_id = %self.partition_id, + event_count, + "The service released the batch; it did not durably accept it." + ); + self.report_failure( + events, + EventHubsError::from(ErrorKind::SendNotAccepted( + "the service returned a Released outcome".into(), + )), + ) + .await; + } + Ok(AmqpSendOutcome::Rejected(reason)) => { + warn!( + partition_id = %self.partition_id, + event_count, + "The service rejected the batch." + ); + self.report_failure( + events, + EventHubsError::from(ErrorKind::SendRejected(reason)), + ) + .await; + } + Err(error) => { + // The recoverable sender already applied the retry policy, so + // this error is terminal. + warn!( + partition_id = %self.partition_id, + event_count, + "The batch failed after the retry policy was exhausted." + ); + self.report_failure(events, error).await; + } + } + } + + async fn report_success(&self, events: Vec) { + if let Some(handler) = self.handlers.succeeded.as_ref() { + handler(SendBatchSucceededContext { + partition_id: self.partition_id.clone(), + events, + }) + .await; + } + } + + async fn report_failure(&self, events: Vec, error: EventHubsError) { + (self.handlers.failed)(SendBatchFailedContext { + partition_id: self.partition_id.clone(), + events, + error, + }) + .await; + } + + /// Removes events from the buffered counts once they reach a terminal outcome. + fn release(&self, count: usize) { + if count == 0 { + return; + } + self.buffered.fetch_sub(count, Ordering::AcqRel); + self.total_buffered.fetch_sub(count, Ordering::AcqRel); + } + + /// Drops every event that the worker still holds after an immediate close. + /// + /// Dropping a command also returns its capacity permit, so a later caller is + /// not blocked by an abandoned event. + async fn discard_remaining(&mut self, pending: Vec) { + let mut abandoned = pending.len(); + drop(pending); + + // Take everything that is still in the queue. The client already took + // the sender, so this loop ends. + self.receiver.close(); + while let Some(command) = self.receiver.next().await { + match command { + Command::Event { .. } => abandoned += 1, + // A waiting flush must not hang. Dropping the sender cancels it. + Command::Flush(completed) => drop(completed), + } + } + + if abandoned > 0 { + warn!( + partition_id = %self.partition_id, + event_count = abandoned, + "Abandoned buffered events during an immediate close." + ); + } + } +} diff --git a/sdk/eventhubs/azure_messaging_eventhubs/src/producer/mod.rs b/sdk/eventhubs/azure_messaging_eventhubs/src/producer/mod.rs index 7e8e14c8b9f..70584faf216 100644 --- a/sdk/eventhubs/azure_messaging_eventhubs/src/producer/mod.rs +++ b/sdk/eventhubs/azure_messaging_eventhubs/src/producer/mod.rs @@ -1,6 +1,8 @@ // Copyright (c) Microsoft Corporation. All Rights reserved // Licensed under the MIT license. +// cspell: ignore retryable + use crate::{ common::{ recoverable::{RecoverableConnection, RecoverableSender}, @@ -25,6 +27,9 @@ use tracing::{trace, warn}; /// Types used to collect messages into a "batch" before submitting them to an Event Hub. pub(crate) mod batch; +/// A producer client that buffers events and publishes them in the background. +pub(crate) mod buffered; + pub(crate) const DEFAULT_EVENTHUBS_APPLICATION: &str = "DefaultApplicationName"; #[derive(Default, Debug, Clone)] @@ -403,18 +408,8 @@ impl ProducerClient { #[allow(unused_variables)] options: Option, ) -> Result<()> { let path = batch.get_batch_path()?; - let sender = self.connection.get_sender(path.clone()).await?; - let messages = batch.get_messages(); - let outcome = sender - .send( - messages, - Some(AmqpSendOptions { - message_format: Some(Self::BATCH_MESSAGE_FORMAT), - ..Default::default() - }), - ) - .await?; + let outcome = self.send_batch_envelope(path.clone(), messages).await?; match outcome { AmqpSendOutcome::Accepted => Ok(()), AmqpSendOutcome::Rejected(reason) => { @@ -459,6 +454,35 @@ impl ProducerClient { } } + /// Sends a batch envelope to a path and returns the raw AMQP outcome. + /// + /// The caller decides what each outcome means. [`ProducerClient::send_batch`] + /// treats `Modified` and `Released` as success with a warning, for backward + /// compatibility. The buffered producer treats them as a delivery failure, + /// because neither outcome means the service stored the events. + /// + /// An `Err` from this method means the retry policy is exhausted, or the + /// error is not retryable. [`RecoverableSender`] applies the retry policy and + /// the connection recovery, and it converts a `Rejected` outcome into an + /// error inside the retry loop. + pub(crate) async fn send_batch_envelope( + &self, + path: Url, + envelope: AmqpMessage, + ) -> Result { + let sender = self.connection.get_sender(path).await?; + let outcome = sender + .send( + envelope, + Some(AmqpSendOptions { + message_format: Some(Self::BATCH_MESSAGE_FORMAT), + ..Default::default() + }), + ) + .await?; + Ok(outcome) + } + /// Gets the properties of the Event Hub. /// # Returns /// A `Result` containing the properties of the Event Hub. diff --git a/sdk/eventhubs/azure_messaging_eventhubs/tests/eventhubs_buffered_producer.rs b/sdk/eventhubs/azure_messaging_eventhubs/tests/eventhubs_buffered_producer.rs new file mode 100644 index 00000000000..13c2a4a316b --- /dev/null +++ b/sdk/eventhubs/azure_messaging_eventhubs/tests/eventhubs_buffered_producer.rs @@ -0,0 +1,499 @@ +// Copyright (c) Microsoft Corporation. All Rights reserved +// Licensed under the MIT license. + +//! Live tests for the buffered producer client. +//! +//! The tests need a real Event Hub. They read the namespace from +//! `EVENTHUBS_HOST` and the Event Hub name from `EVENTHUB_NAME`. The +//! connection string test also reads `EVENTHUBS_CONNECTION_STRING`. + +use azure_core::time::Duration; +use azure_core_test::{recorded, TestContext}; +use azure_messaging_eventhubs::{ + models::EventData, BufferedProducerClient, ConsumerClient, EnqueueEventOptions, + OpenReceiverOptions, SendBatchFailedContext, SendBatchSucceededContext, StartLocation, + StartPosition, +}; +use futures::stream::StreamExt; +use std::{ + env, + error::Error, + sync::{ + atomic::{AtomicUsize, Ordering}, + Arc, Mutex, + }, +}; +use tracing::info; + +/// Collects the delivery reports of a test. +#[derive(Default)] +struct Reports { + succeeded: AtomicUsize, + failed: AtomicUsize, + failures: Mutex>, +} + +impl Reports { + fn on_success(&self, context: &SendBatchSucceededContext) { + self.succeeded + .fetch_add(context.events.len(), Ordering::AcqRel); + } + + fn on_failure(&self, context: &SendBatchFailedContext) { + self.failed + .fetch_add(context.events.len(), Ordering::AcqRel); + self.failures + .lock() + .unwrap() + .push(context.error.to_string()); + } + + fn succeeded(&self) -> usize { + self.succeeded.load(Ordering::Acquire) + } + + fn failed(&self) -> usize { + self.failed.load(Ordering::Acquire) + } + + fn failures(&self) -> Vec { + self.failures.lock().unwrap().clone() + } +} + +/// Builds a buffered producer that records every delivery outcome. +async fn open_producer( + test_name: &str, + host: &str, + eventhub: &str, + credential: Arc, + max_wait_time: Duration, +) -> Result<(BufferedProducerClient, Arc), Box> { + let reports = Arc::new(Reports::default()); + + let for_success = reports.clone(); + let for_failure = reports.clone(); + + let producer = BufferedProducerClient::builder() + .with_application_id(test_name.to_string()) + .with_max_wait_time(max_wait_time) + .with_on_send_succeeded(move |context| { + let reports = for_success.clone(); + async move { + reports.on_success(&context); + } + }) + .with_on_send_failed(move |context| { + let reports = for_failure.clone(); + async move { + reports.on_failure(&context); + } + }) + .open(host, eventhub, credential) + .await?; + + Ok((producer, reports)) +} + +/// Reads events from one partition and returns the bodies that carry the prefix. +/// +/// Each test tags its events with its own prefix, so a test that shares a +/// partition with another test does not count the events of that other test. +async fn receive_bodies( + consumer: &ConsumerClient, + partition_id: &str, + start_sequence: i64, + count: usize, + prefix: &str, +) -> Result, Box> { + let receiver = consumer + .open_receiver_on_partition( + partition_id.to_string(), + Some(OpenReceiverOptions { + start_position: Some(StartPosition { + location: StartLocation::SequenceNumber(start_sequence), + inclusive: false, + }), + ..Default::default() + }), + ) + .await?; + + let mut bodies = Vec::with_capacity(count); + let mut stream = receiver.stream_events(); + while let Some(event) = stream.next().await { + let event = event?; + if let Some(body) = event.event_data().body() { + let body = String::from_utf8_lossy(body).into_owned(); + if body.starts_with(prefix) { + bodies.push(body); + } + } + if bodies.len() >= count { + break; + } + } + Ok(bodies) +} + +/// Enqueue, batch, flush, and then read the events back. +#[recorded::test(live)] +async fn buffered_round_trip(ctx: TestContext) -> Result<(), Box> { + const TEST_NAME: &str = "buffered_round_trip"; + const PARTITION: &str = "0"; + const EVENT_COUNT: usize = 20; + + let recording = ctx.recording(); + let host = env::var("EVENTHUBS_HOST")?; + let eventhub = env::var("EVENTHUB_NAME")?; + let credential = recording.credential(); + + let consumer = ConsumerClient::builder() + .with_application_id(TEST_NAME.to_string()) + .open(host.as_str(), eventhub.clone(), credential.clone()) + .await?; + let start_sequence = consumer + .get_partition_properties(PARTITION) + .await? + .last_enqueued_sequence_number; + + let (producer, reports) = open_producer( + TEST_NAME, + host.as_str(), + eventhub.as_str(), + credential.clone(), + Duration::seconds(1), + ) + .await?; + + for index in 0..EVENT_COUNT { + producer + .enqueue_event( + EventData::builder() + .with_body(format!("buffered-{index}").into_bytes()) + .build(), + Some(EnqueueEventOptions { + partition_id: Some(PARTITION.to_string()), + ..Default::default() + }), + ) + .await?; + } + + // A successful enqueue only means the local buffer accepted the event, so + // the test flushes before it reads the events back. + producer.flush().await?; + + assert_eq!(reports.succeeded(), EVENT_COUNT); + assert_eq!(reports.failed(), 0, "failures: {:?}", reports.failures()); + assert_eq!(producer.total_buffered_event_count(), 0); + + // The client sent more than one event in each batch. + info!("Reading the events back from partition {PARTITION}."); + let bodies = receive_bodies( + &consumer, + PARTITION, + start_sequence, + EVENT_COUNT, + "buffered-", + ) + .await?; + assert_eq!(bodies.len(), EVENT_COUNT); + for index in 0..EVENT_COUNT { + assert!( + bodies.contains(&format!("buffered-{index}")), + "the service did not return event {index}" + ); + } + + producer.close().await?; + consumer.close().await?; + Ok(()) +} + +/// An explicit partition ID sends the events to that partition only. +#[recorded::test(live)] +async fn buffered_explicit_partition_routing(ctx: TestContext) -> Result<(), Box> { + const TEST_NAME: &str = "buffered_explicit_partition_routing"; + const PARTITION: &str = "1"; + const EVENT_COUNT: usize = 5; + + let recording = ctx.recording(); + let host = env::var("EVENTHUBS_HOST")?; + let eventhub = env::var("EVENTHUB_NAME")?; + let credential = recording.credential(); + + let consumer = ConsumerClient::builder() + .with_application_id(TEST_NAME.to_string()) + .open(host.as_str(), eventhub.clone(), credential.clone()) + .await?; + let start_sequence = consumer + .get_partition_properties(PARTITION) + .await? + .last_enqueued_sequence_number; + + let (producer, reports) = open_producer( + TEST_NAME, + host.as_str(), + eventhub.as_str(), + credential.clone(), + Duration::seconds(1), + ) + .await?; + + for index in 0..EVENT_COUNT { + producer + .enqueue_event( + format!("routed-{index}"), + Some(EnqueueEventOptions { + partition_id: Some(PARTITION.to_string()), + ..Default::default() + }), + ) + .await?; + } + producer.flush().await?; + + assert_eq!(reports.succeeded(), EVENT_COUNT); + assert_eq!(reports.failed(), 0, "failures: {:?}", reports.failures()); + + let bodies = + receive_bodies(&consumer, PARTITION, start_sequence, EVENT_COUNT, "routed-").await?; + for index in 0..EVENT_COUNT { + assert!(bodies.contains(&format!("routed-{index}"))); + } + + producer.close().await?; + consumer.close().await?; + Ok(()) +} + +/// A batch that is not full still goes out after the maximum wait time. +#[recorded::test(live)] +async fn buffered_partial_batch_timeout(ctx: TestContext) -> Result<(), Box> { + const TEST_NAME: &str = "buffered_partial_batch_timeout"; + const PARTITION: &str = "2"; + + let recording = ctx.recording(); + let host = env::var("EVENTHUBS_HOST")?; + let eventhub = env::var("EVENTHUB_NAME")?; + let credential = recording.credential(); + + let consumer = ConsumerClient::builder() + .with_application_id(TEST_NAME.to_string()) + .open(host.as_str(), eventhub.clone(), credential.clone()) + .await?; + let start_sequence = consumer + .get_partition_properties(PARTITION) + .await? + .last_enqueued_sequence_number; + + let (producer, reports) = open_producer( + TEST_NAME, + host.as_str(), + eventhub.as_str(), + credential.clone(), + Duration::milliseconds(500), + ) + .await?; + + producer + .enqueue_event( + "partial-batch", + Some(EnqueueEventOptions { + partition_id: Some(PARTITION.to_string()), + ..Default::default() + }), + ) + .await?; + + // No flush and no close. The wait time alone must send this event, so the + // read below returns once the timer fires. + let bodies = receive_bodies(&consumer, PARTITION, start_sequence, 1, "partial-batch").await?; + assert_eq!(bodies, vec!["partial-batch".to_string()]); + assert_eq!(reports.failed(), 0, "failures: {:?}", reports.failures()); + + producer.close().await?; + consumer.close().await?; + Ok(()) +} + +/// A graceful close sends the events that the client still holds. +#[recorded::test(live)] +async fn buffered_graceful_close_sends_events(ctx: TestContext) -> Result<(), Box> { + const TEST_NAME: &str = "buffered_graceful_close_sends_events"; + const PARTITION: &str = "3"; + const EVENT_COUNT: usize = 10; + + let recording = ctx.recording(); + let host = env::var("EVENTHUBS_HOST")?; + let eventhub = env::var("EVENTHUB_NAME")?; + let credential = recording.credential(); + + let consumer = ConsumerClient::builder() + .with_application_id(TEST_NAME.to_string()) + .open(host.as_str(), eventhub.clone(), credential.clone()) + .await?; + let start_sequence = consumer + .get_partition_properties(PARTITION) + .await? + .last_enqueued_sequence_number; + + // A long wait time makes sure that the close, and not the timer, sends the + // events. + let (producer, reports) = open_producer( + TEST_NAME, + host.as_str(), + eventhub.as_str(), + credential.clone(), + Duration::seconds(120), + ) + .await?; + + for index in 0..EVENT_COUNT { + producer + .enqueue_event( + format!("closed-{index}"), + Some(EnqueueEventOptions { + partition_id: Some(PARTITION.to_string()), + ..Default::default() + }), + ) + .await?; + } + + producer.close().await?; + + assert_eq!(reports.succeeded(), EVENT_COUNT); + assert_eq!(reports.failed(), 0, "failures: {:?}", reports.failures()); + assert_eq!(producer.total_buffered_event_count(), 0); + + let bodies = + receive_bodies(&consumer, PARTITION, start_sequence, EVENT_COUNT, "closed-").await?; + for index in 0..EVENT_COUNT { + assert!(bodies.contains(&format!("closed-{index}"))); + } + + consumer.close().await?; + Ok(()) +} + +/// A connection string opens the client, and the delivery path still works. +#[recorded::test(live)] +async fn buffered_round_trip_with_connection_string( + _ctx: TestContext, +) -> Result<(), Box> { + const TEST_NAME: &str = "buffered_round_trip_with_connection_string"; + const EVENT_COUNT: usize = 5; + + // SAS credentials come from the connection string, not the recording. + let connection_string = env::var("EVENTHUBS_CONNECTION_STRING")?; + let eventhub = env::var("EVENTHUB_NAME").ok(); + + let consumer = ConsumerClient::builder() + .with_application_id(TEST_NAME.to_string()) + .open_with_connection_string(&connection_string, eventhub.as_deref()) + .await?; + + // The Event Hub decides which partitions exist, so the test asks for one + // instead of naming it. The standard test resource has four partitions, and + // there are more tests than partitions, so this one shares a partition. The + // marker below keeps the events apart. + let partition = consumer + .get_eventhub_properties() + .await? + .partition_ids + .last() + .ok_or("the Event Hub reported no partitions")? + .clone(); + info!("Using partition {partition} for the connection-string test."); + + let start_sequence = consumer + .get_partition_properties(&partition) + .await? + .last_enqueued_sequence_number; + + let reports = Arc::new(Reports::default()); + let for_success = reports.clone(); + let for_failure = reports.clone(); + + let producer = BufferedProducerClient::builder() + .with_application_id(TEST_NAME.to_string()) + .with_max_wait_time(Duration::seconds(1)) + .with_on_send_succeeded(move |context| { + let reports = for_success.clone(); + async move { + reports.on_success(&context); + } + }) + .with_on_send_failed(move |context| { + let reports = for_failure.clone(); + async move { + reports.on_failure(&context); + } + }) + .open_with_connection_string(&connection_string, eventhub.as_deref()) + .await?; + + // Tag the events, so the test finds them among the other events of the partition. + let marker = format!("sas-buffered-{start_sequence}"); + for index in 0..EVENT_COUNT { + producer + .enqueue_event( + format!("{marker}-{index}"), + Some(EnqueueEventOptions { + partition_id: Some(partition.clone()), + ..Default::default() + }), + ) + .await?; + } + producer.flush().await?; + + assert_eq!(reports.succeeded(), EVENT_COUNT); + assert_eq!(reports.failed(), 0, "failures: {:?}", reports.failures()); + + let bodies = + receive_bodies(&consumer, &partition, start_sequence, EVENT_COUNT, &marker).await?; + for index in 0..EVENT_COUNT { + assert!(bodies.contains(&format!("{marker}-{index}"))); + } + + producer.close().await?; + consumer.close().await?; + Ok(()) +} + +/// Automatic assignment spreads the events over the partitions. +#[recorded::test(live)] +async fn buffered_automatic_partition_assignment(ctx: TestContext) -> Result<(), Box> { + const TEST_NAME: &str = "buffered_automatic_partition_assignment"; + const EVENT_COUNT: usize = 32; + + let recording = ctx.recording(); + let host = env::var("EVENTHUBS_HOST")?; + let eventhub = env::var("EVENTHUB_NAME")?; + let credential = recording.credential(); + + let (producer, reports) = open_producer( + TEST_NAME, + host.as_str(), + eventhub.as_str(), + credential.clone(), + Duration::seconds(1), + ) + .await?; + + for index in 0..EVENT_COUNT { + producer + .enqueue_event(format!("auto-{index}"), None) + .await?; + } + producer.flush().await?; + + assert_eq!(reports.succeeded(), EVENT_COUNT); + assert_eq!(reports.failed(), 0, "failures: {:?}", reports.failures()); + + producer.close().await?; + Ok(()) +}