azure_messaging_eventhubs 0.16.0

// Copyright (c) Microsoft Corp. All Rights Reserved.
// Licensed under the MIT License.

//use async_channel::{bounded, Receiver, Sender};
use super::{
    load_balancer::LoadBalancer,
    models::{Checkpoint, StartPositions},
    partition_client::PartitionClient,
    CheckpointStore, ProcessorStrategy,
};
use crate::{
    error::Result, models::ConsumerClientDetails, ConsumerClient, EventHubsError,
    OpenReceiverOptions, StartLocation, StartPosition,
};
//use async_io::Timer;
use async_lock::Mutex as AsyncMutex;
use azure_core::{error::ErrorKind as AzureErrorKind, time::Duration, Error};
use futures::{
    channel::mpsc::{channel, Receiver, Sender},
    SinkExt, StreamExt,
};
use std::{
    sync::{
        Arc,
        Mutex as SyncMutex, // Mutex for blocking operations
    },
    {collections::HashMap, sync::Weak},
};
use tracing::{debug, error, info, warn};

// AMQP epoch (owner level) used for every partition receiver opened by
// `EventProcessor`. Matches `EventProcessorClient` in the .NET and Java
// SDKs. `0` is itself an exclusive epoch (a new receiver-at-0 displaces a
// prior receiver-at-0), which is how the processor detects steals.
const PROCESSOR_OWNER_LEVEL: i64 = 0;

/// Represents the event processor responsible for processing events
/// from Event Hub partitions.
///
/// This struct manages the load balancing strategy, checkpoint store,
/// and consumer client for processing events.
/// It provides methods for starting the event processor, dispatching
/// events, and managing partition clients.
///
/// The event processor uses a load balancer to distribute the load
/// across partitions and a checkpoint store to manage checkpoints.
///
/// Each per-partition receiver opens with AMQP epoch `0`, matching the .NET
/// and Java `EventProcessorClient`. When another `EventProcessor` attaches
/// to the same partition, the broker disconnects this receiver and
/// `stream_events()` resolves with `EventHubsError::ConsumerDisconnected`.
/// To use a different epoch, open receivers directly via
/// `ConsumerClient::open_receiver_on_partition`.
///
/// For more information on Event Processors and scenarios in which you would
/// use an Event Processor, see the [Event Processor documentation](https://learn.microsoft.com/azure/event-hubs/event-processor-balance-partition-load).
///
pub struct EventProcessor {
    checkpoint_store: Arc<dyn CheckpointStore + Send + Sync>,
    load_balancer: Arc<AsyncMutex<LoadBalancer>>,
    consumer_client: ConsumerClient,
    next_partition_clients: AsyncMutex<Receiver<Arc<PartitionClient>>>,
    next_partition_client_sender: Sender<Arc<PartitionClient>>,
    client_details: ConsumerClientDetails,
    prefetch: u32,
    update_interval: Duration,
    start_positions: StartPositions,
    is_running: std::sync::Mutex<bool>,
    partition_ids: Vec<String>,
}

struct EventProcessorOptions {
    strategy: ProcessorStrategy,
    partition_expiration_duration: Duration,
    update_interval: Duration,
    start_positions: StartPositions,
    prefetch: u32,
    partition_ids: Vec<String>,
}

pub(crate) struct ProcessorConsumersMap {
    consumers: SyncMutex<HashMap<String, Weak<PartitionClient>>>,
}

impl ProcessorConsumersMap {
    fn new() -> Self {
        ProcessorConsumersMap {
            consumers: SyncMutex::new(HashMap::new()),
        }
    }

    /// Adds a partition client to the consumers map.
    /// If a partition client already exists for the given partition ID,
    /// it will not be added again.
    /// Returns `true` if the partition client was added successfully,
    /// or `false` if it already exists.
    ///
    /// # Arguments
    /// * `partition_id` - The ID of the partition for which the client is being added.
    /// * `partition_client` - The partition client to be added.
    ///
    /// # Returns
    /// A `Result` indicating the success or failure of the operation.
    /// If successful, returns `true` if the partition client was added,
    /// or `false` if it already exists.
    ///
    pub async fn add_partition_client(
        &self,
        partition_id: &str,
        partition_client: Arc<PartitionClient>,
    ) -> Result<bool> {
        debug!(partition_id = %partition_id, "Adding partition client for partition.");
        let mut consumers = self
            .consumers
            .lock()
            .map_err(|_| EventHubsError::with_message("Could not lock consumers mutex."))?;
        if consumers.contains_key(partition_id) {
            debug!(
                partition_id = %partition_id,
                "Partition client already exists for partition."
            );
            return Ok(false);
        }
        consumers.insert(partition_id.to_string(), Arc::downgrade(&partition_client));
        debug!(partitions = ?consumers.keys(), "Consumers for partition.");
        Ok(true)
    }

    pub fn remove_partition_client(&self, partition_id: &str) -> Result<()> {
        debug!(partition_id = %partition_id, "Removing partition client for partition.");
        let mut consumers = self
            .consumers
            .lock()
            .map_err(|_| EventHubsError::with_message("Could not lock consumers mutex."))?;
        consumers.remove(partition_id);
        debug!(partitions = ?consumers.keys(), "Consumers for partition now.");
        Ok(())
    }

    /// Returns the set of partition IDs that have active partition clients.
    fn get_active_partition_ids(&self) -> Result<Vec<String>> {
        let consumers = self
            .consumers
            .lock()
            .map_err(|_| EventHubsError::with_message("Could not lock consumers mutex."))?;
        Ok(consumers.keys().cloned().collect())
    }

    /// Removes partitions reassigned away from this processor and closes
    /// their receivers so consumer streams terminate. Backstop for the
    /// broker's epoch-based disconnect.
    async fn revoke_partition_clients(&self, partition_ids: &[String]) -> Result<()> {
        // Collect under the sync lock, then release before awaiting:
        // SyncMutex guards cannot be held across `.await`.
        let to_close: Vec<Arc<PartitionClient>> = {
            let mut consumers = self
                .consumers
                .lock()
                .map_err(|_| EventHubsError::with_message("Could not lock consumers mutex."))?;
            partition_ids
                .iter()
                .filter_map(|id| consumers.remove(id).and_then(|w| w.upgrade()))
                .collect()
        };
        for client in to_close {
            client.request_close_receiver().await;
        }
        Ok(())
    }
}

//pub(crate) type ConsumersType = std::sync::Mutex<HashMap<String, Arc<PartitionClient>>>;

unsafe impl Send for EventProcessor {}
unsafe impl Sync for EventProcessor {}

impl EventProcessor {
    /// Creates a new `EventProcessorBuilder` instance.
    /// This builder allows you to configure various options for the event processor,
    /// such as load balancing strategy, update interval, start positions, and more.
    ///
    /// # Returns a new [`builders::EventProcessorBuilder`] instance.
    pub fn builder() -> builders::EventProcessorBuilder {
        builders::EventProcessorBuilder::new()
    }

    fn new(
        consumer_client: ConsumerClient,
        checkpoint_store: Arc<dyn CheckpointStore + Send + Sync>,
        options: EventProcessorOptions,
    ) -> Result<Arc<Self>> {
        let (sender, receiver) = channel(options.partition_ids.len());

        let client_details = consumer_client.get_details()?;

        Ok(Arc::new(EventProcessor {
            checkpoint_store: checkpoint_store.clone(),
            consumer_client,

            // Default to Balanced strategy if not provided
            load_balancer: Arc::new(AsyncMutex::new(LoadBalancer::new(
                checkpoint_store.clone(),
                client_details.clone(),
                options.strategy,
                options.partition_expiration_duration,
                None,
            ))),
            client_details,
            prefetch: options.prefetch,
            update_interval: options.update_interval,
            start_positions: options.start_positions,
            next_partition_client_sender: sender,
            next_partition_clients: AsyncMutex::new(receiver),
            is_running: std::sync::Mutex::new(false),
            partition_ids: options.partition_ids,
        }))
    }

    /// Starts the event processor.
    /// This method initiates the event processing loop and begins
    /// processing events from the Event Hub partitions.
    /// It uses the specified checkpoint store and load balancing strategy
    /// to manage the ownership of partitions and distribute the load
    /// among consumers.
    /// The event processor will run until it is stopped or interrupted.
    /// # Errors
    /// Returns an error if the event processor fails to start.
    /// # Examples
    /// ```
    /// use azure_messaging_eventhubs::EventProcessor;
    /// use azure_messaging_eventhubs::ConsumerClient;
    /// use std::sync::Arc;
    /// use azure_core::time::Duration;
    /// use azure_messaging_eventhubs::ProcessorStrategy;
    /// use azure_messaging_eventhubs::CheckpointStore;
    ///
    /// async fn run_processor(consumer_client: ConsumerClient, checkpoint_store: impl CheckpointStore+Send+Sync+'static) -> Result<(), Box<dyn std::error::Error>> {
    ///   // Create an instance of the EventProcessor
    ///   let event_processor = EventProcessor::builder()
    ///       .with_load_balancing_strategy(ProcessorStrategy::Balanced)
    ///       .with_update_interval(Duration::seconds(30))
    ///       .with_partition_expiration_duration(Duration::seconds(10))
    ///       .with_prefetch(300)
    ///       .build(
    ///          consumer_client,
    ///          Arc::new(checkpoint_store)).await?;
    ///
    ///   // Start the event processor
    ///   {
    ///     tokio::select!{
    ///          result = event_processor.run() => {
    ///              if let Err(e) = result {
    ///                  println!("Event processor failed: {:?}", e);
    ///              } else {
    ///                  println!("Event processor finished successfully");
    ///              }
    ///          }
    ///          _ = tokio::time::sleep(std::time::Duration::from_secs(60)) => {}
    ///     }
    ///   }
    ///   Ok(())
    /// }
    /// ```
    ///
    pub async fn run(&self) -> Result<()> {
        {
            let mut is_running = self.is_running.lock().map_err(|_| {
                Error::new(AzureErrorKind::Io, "Could not lock is_running on startup")
            })?;
            *is_running = true;
        }

        let consumers = Arc::new(ProcessorConsumersMap::new());
        let partition_ids = &self
            .partition_ids
            .iter()
            .map(String::as_str)
            .collect::<Vec<&str>>();

        loop {
            let result = self.dispatch(partition_ids, &consumers).await;
            match result {
                Ok(_) => {
                    debug!("Event processor dispatched successfully.");
                }
                Err(e) => {
                    error!(err = ?e, "Error dispatching event processor.");
                    return Err(e);
                }
            }
            debug!("Event processor sleeping for {:?}", self.update_interval);
            azure_core::sleep::sleep(self.update_interval).await;
            debug!("Event processor woke up from sleep.");
            if self.is_shutdown()? {
                info!("Event processor shutting down.");
                break Ok(());
            }
        }
    }

    /// Shuts down the event processor.
    pub async fn shutdown(&self) -> Result<()> {
        // Implement shutdown logic if needed

        let mut is_running = self.is_running.lock().map_err(|_| {
            EventHubsError::with_message("Failed to acquire lock on is_running for shutdown")
        })?;

        *is_running = false;
        Ok(())
    }

    fn is_shutdown(&self) -> Result<bool> {
        // Implement shutdown logic if needed
        let is_running = self
            .is_running
            .lock()
            .map_err(|_| EventHubsError::with_message("Failed to acquire lock on is_running"))?;
        if *is_running {
            Ok(false)
        } else {
            Ok(true)
        }
    }

    #[tracing::instrument(
        level = "debug",
        skip_all,
        fields(
            eventhub = %self.client_details.eventhub_name,
            consumer_group = %self.client_details.consumer_group,
            fully_qualified_namespace = %self.client_details.fully_qualified_namespace,
            owner_id = %self.client_details.client_id,
        ),
        err,
    )]
    async fn dispatch(
        &self,
        partition_ids: &[&str],
        consumers: &Arc<ProcessorConsumersMap>,
    ) -> Result<()> {
        debug!("Dispatch partition clients to consumers.");
        let load_balancer = self.load_balancer.lock().await;

        let ownerships = load_balancer.load_balance(partition_ids).await;
        let ownerships = ownerships.map_err(|e| {
            error!(err = ?e, "Error in load balancing.");
            e
        })?;

        // Revoke clients for any partitions no longer in the ownership set.
        let owned_ids: std::collections::HashSet<&str> =
            ownerships.iter().map(|o| o.partition_id.as_str()).collect();
        let active_ids = consumers.get_active_partition_ids()?;
        let stolen: Vec<String> = active_ids
            .into_iter()
            .filter(|id| !owned_ids.contains(id.as_str()))
            .collect();
        if !stolen.is_empty() {
            info!(
                partitions = %stolen.join(", "),
                "Partitions no longer owned, revoking."
            );
            consumers.revoke_partition_clients(&stolen).await?;
        }

        let checkpoints = self.get_checkpoint_map().await;
        let checkpoints = checkpoints.map_err(|e| {
            error!(err = ?e, "Error in getting checkpoint map.");
            e
        })?;

        debug!(
            "Adding partition clients for {} ownerships ",
            ownerships.len()
        );
        for ownership in ownerships {
            let err = self
                .add_partition_client(
                    ownership.partition_id,
                    &checkpoints,
                    Arc::downgrade(consumers),
                )
                .await;
            if let Err(e) = err {
                error!(err = ?e, "Error adding partition client.");
                return Err(e);
            }
        }

        Ok(())
    }

    #[tracing::instrument(
        level = "debug",
        skip_all,
        fields(partition_id = %partition_id),
        err,
    )]
    async fn add_partition_client(
        &self,
        partition_id: String,
        checkpoints: &HashMap<String, Checkpoint>,
        consumers: Weak<ProcessorConsumersMap>,
    ) -> Result<()> {
        debug!(partition_id = %partition_id, "Add partition client for partition.");

        let partition_client = Arc::new(PartitionClient::new(
            partition_id.clone(),
            self.checkpoint_store.clone(),
            self.client_details.clone(),
            consumers.clone(),
        ));

        if let Some(strong_consumers) = consumers.upgrade() {
            if !strong_consumers
                .add_partition_client(&partition_id, partition_client.clone())
                .await?
            {
                debug!(
                    partition_id = %partition_id,
                    "Partition client already exists for partition, ignoring."
                );
                return Ok(());
            }
        } else {
            error!("Consumers map is no longer valid.");
            return Err(EventHubsError::with_message(
                "Consumers map is no longer valid.",
            ));
        }

        // Since we can only have a single EventReceiver on a partition, we don't actually attempt to create the receiver until
        let start_position = self.get_start_position(&partition_id, checkpoints);
        debug!(
            partition_id = %partition_id,
            start_position = ?start_position,
            "Start position for partition."
        );
        let receiver = self
            .consumer_client
            .open_receiver_on_partition(
                partition_id.clone(),
                Some(OpenReceiverOptions {
                    start_position: Some(start_position),
                    prefetch: Some(self.prefetch),
                    owner_level: Some(PROCESSOR_OWNER_LEVEL),
                    ..Default::default()
                }),
            )
            .await;
        // Roll back the consumers-map entry on failure; otherwise the
        // partition is stuck (the map's `contains_key` check would
        // short-circuit every retry) until steal-revocation or restart.
        let receiver = match receiver {
            Ok(r) => r,
            Err(e) => {
                error!(
                    partition_id = %partition_id,
                    err = ?e,
                    "Error opening receiver for partition client."
                );
                if let Some(strong_consumers) = consumers.upgrade() {
                    let _ = strong_consumers.remove_partition_client(&partition_id);
                }
                return Err(e);
            }
        };
        info!(partition_id = %partition_id, "Receiver opened for partition client.");
        if let Err(e) = partition_client.set_event_receiver(receiver) {
            error!(
                partition_id = %partition_id,
                err = ?e,
                "Error setting event receiver for partition."
            );
            if let Some(strong_consumers) = consumers.upgrade() {
                let _ = strong_consumers.remove_partition_client(&partition_id);
            }
            return Err(e);
        }

        debug!(partition_id = %partition_id, "Adding partition client to queue.");

        // Send the partition client to the next partition client receiver
        {
            let mut sender = self.next_partition_client_sender.clone();
            sender.send(partition_client).await.map_err(|e| {
                EventHubsError::from(azure_core::Error::with_message(
                    AzureErrorKind::Other,
                    format!("Failed to send partition client: {:?}", e),
                ))
            })?;
        }
        debug!(
            partition_id = %partition_id,
            "add_partition_client: Partition client added for partition."
        );

        Ok(())
    }

    /// Retrieves the next partition client for processing events.
    ///
    /// This method returns the next available partition client.
    pub async fn next_partition_client(&self) -> Result<Arc<PartitionClient>> {
        // Implement the function or remove it if not needed
        debug!("next_partition_client: Waiting to receive the next partition client.");

        {
            // Wait for the next partition client to be available
            let mut clients = self.next_partition_clients.lock().await;
            let next_client = clients.next().await.ok_or_else(|| {
                EventHubsError::with_message("No next partition client available: ")
            })?;

            debug!(
                partition_id = %next_client.get_partition_id(),
                "next_partition_client: Returning partition client for partition."
            );
            Ok(next_client)
        }
    }

    /// Closes the event processor.
    pub async fn close(self) -> Result<()> {
        // Close all partition clients.
        info!("Closing all partition clients.");
        let mut clients = self.next_partition_clients.lock().await;
        while let Ok(client) = clients.try_recv() {
            info!(
                partition_id = %client.get_partition_id(),
                "Closing partition client for partition."
            );
            // A partition client that the application still holds cannot be
            // taken out of its `Arc`. Report it and continue, so that one such
            // client does not stop the processor from closing the clients that
            // follow.
            //
            // The connection solved the same problem by taking `&self`, and a
            // partition client could do the same through
            // `EventReceiver::request_close`, which detaches the receiver the
            // way `EventReceiver::close` does. That needs a new signature for
            // the public `PartitionClient::close`, which is a breaking change.
            // This fix therefore leaves such a client to the application that
            // holds it.
            let Ok(client) = Arc::try_unwrap(client) else {
                warn!(
                    "Could not close a partition client, because the application still holds it."
                );
                continue;
            };
            let res = client.close().await;
            if let Err(e) = res {
                error!(err = ?e, "Failed to close partition client.");
            } else {
                info!("Partition client closed successfully.");
            }
        }

        // Close the event processor and release resources.
        info!("Closing consumer client.");
        let res = self.consumer_client.close().await;
        if let Err(e) = res {
            error!(err = ?e, "Failed to close consumer client.");
        } else {
            info!("Consumer client closed successfully.");
        }
        Ok(())
    }

    /// Retrieves the checkpoint map for the Event Hub.
    ///
    /// This method fetches the checkpoints for all partitions in the Event Hub
    /// and returns them as a `HashMap` where the keys are partition IDs
    ///
    /// # Returns
    /// A `Result` containing a `HashMap` of partition IDs and their corresponding `Checkpoint` objects.
    ///
    ///
    async fn get_checkpoint_map(&self) -> Result<HashMap<String, Checkpoint>> {
        let checkpoints = self.checkpoint_store.list_checkpoints(
            &self.client_details.fully_qualified_namespace,
            &self.client_details.eventhub_name,
            &self.client_details.consumer_group,
        );
        let mut checkpoint_map = HashMap::new();
        for checkpoint in checkpoints.await? {
            checkpoint_map.insert(checkpoint.partition_id.clone(), checkpoint);
        }
        Ok(checkpoint_map)
    }

    /// Retrieve the start position for the specified ownership.
    ///
    /// This method determines the starting position for event processing
    /// based on the ownership information and the provided checkpoints.
    /// It checks if the ownership has a corresponding checkpoint and
    /// returns the appropriate start position.
    ///
    /// If no checkpoint is found for the partition in the ownership, a start
    /// position is chosen from the configured default start positions.
    ///
    /// # Arguments
    /// * partition_id - The partition for which to determine the start position.
    /// * `checkpoints` - A map of checkpoints for all partitions.
    ///
    fn get_start_position(
        &self,
        partition_id: &str,
        checkpoints: &HashMap<String, Checkpoint>,
    ) -> StartPosition {
        let mut start_position = self.start_positions.default.clone();
        if checkpoints.contains_key(partition_id) {
            let checkpoint = checkpoints.get(partition_id).unwrap();
            if let Some(offset) = &checkpoint.offset {
                start_position.location = StartLocation::Offset(offset.clone());
            } else if let Some(sequence_number) = checkpoint.sequence_number {
                start_position.location = StartLocation::SequenceNumber(sequence_number);
            }
        } else if self
            .start_positions
            .per_partition
            .contains_key(partition_id)
        {
            start_position = self
                .start_positions
                .per_partition
                .get(partition_id)
                .unwrap()
                .clone();
        } else {
            start_position = self.start_positions.default.clone();
        }
        start_position
    }
}

pub mod builders {
    use super::{CheckpointStore, EventProcessor};
    use crate::{error::Result, event_processor::models::StartPositions, ConsumerClient};
    use azure_core::time::Duration;
    use std::sync::Arc;

    const DEFAULT_PREFETCH: u32 = 300;
    const DEFAULT_UPDATE_INTERVAL: Duration = Duration::seconds(30);
    const DEFAULT_PARTITION_EXPIRATION_DURATION: Duration = Duration::seconds(60);

    /// Builder for creating an `EventProcessor`.
    /// This builder allows you to configure various options for the event processor,
    /// such as load balancing strategy, update interval, start positions, and more.
    /// It provides a fluent interface for setting these options and building the event processor.
    /// # Examples
    /// ``` no_run
    /// use azure_messaging_eventhubs::{EventProcessor,CheckpointStore ,ConsumerClient};
    /// use std::sync::Arc;
    ///
    /// async fn create_processor(checkpoint_store: Arc<dyn CheckpointStore>) -> Result<(), Box<dyn std::error::Error>> {
    /// use azure_core::Result;
    /// use azure_identity::DeveloperToolsCredential;
    ///
    /// let eventhub_namespace = std::env::var("EVENTHUBS_HOST")?;
    /// let eventhub_name = std::env::var("EVENTHUB_NAME")?;
    /// let consumer = ConsumerClient::builder()
    ///         .open(
    ///             &eventhub_namespace,
    ///             eventhub_name,
    ///             DeveloperToolsCredential::new(None)?.clone(),
    ///         )
    ///         .await?;
    /// println!("Opened consumer client");
    /// let processor = EventProcessor::builder()
    ///     .build(consumer, checkpoint_store.clone())
    ///     .await?;
    /// Ok(())
    /// }
    /// ```
    #[derive(Default)]
    pub struct EventProcessorBuilder {
        update_interval: Option<Duration>,
        start_positions: Option<StartPositions>,
        max_partition_count: Option<usize>,
        prefetch: Option<u32>,
        load_balancing_strategy: Option<super::ProcessorStrategy>,
        partition_expiration_duration: Option<Duration>,
    }
    /// Returns an error if `partition_expiration_duration` is not strictly
    /// greater than `update_interval`. When expiration is shorter than the
    /// load-balancing cycle, every consumer's ownership record expires before
    /// the next cycle observes it, so the load balancer perpetually re-claims
    /// every partition. This is extracted from `build()` to make the rule
    /// directly unit-testable without needing a live `ConsumerClient`.
    pub(crate) fn validate_expiration_vs_update_interval(
        partition_expiration_duration: Duration,
        update_interval: Duration,
    ) -> Result<()> {
        if partition_expiration_duration <= update_interval {
            return Err(crate::EventHubsError::with_message(format!(
                "partition_expiration_duration ({partition_expiration_duration:?}) must be \
                 greater than update_interval ({update_interval:?}); otherwise ownership \
                 records expire between load-balancing cycles and the processor will \
                 perpetually re-claim partitions, causing duplicate processing. A ratio of \
                 at least 2x is recommended."
            )));
        }
        Ok(())
    }

    impl EventProcessorBuilder {
        pub(super) fn new() -> Self {
            EventProcessorBuilder {
                ..Default::default()
            }
        }

        /// Sets the load balancing strategy for the event processor.
        /// The default strategy is `Greedy`.
        pub fn with_load_balancing_strategy(
            mut self,
            load_balancing_strategy: super::ProcessorStrategy,
        ) -> Self {
            self.load_balancing_strategy = Some(load_balancing_strategy);
            self
        }

        /// Sets the processor update interval for the event processor.
        ///
        /// The processor will sleep for the update interval between each iteration.
        /// The default update interval is 30 seconds.
        pub fn with_update_interval(mut self, update_interval: Duration) -> Self {
            self.update_interval = Some(update_interval);
            self
        }

        /// Sets the start positions for each partition and the default start position.
        pub fn with_start_positions(mut self, start_positions: StartPositions) -> Self {
            self.start_positions = Some(start_positions);
            self
        }

        /// Sets the maximum number of partitions to process.
        pub fn with_max_partition_count(mut self, max_partition_count: usize) -> Self {
            self.max_partition_count = Some(max_partition_count);
            self
        }

        /// Sets the prefetch count for the event processor.
        pub fn with_prefetch(mut self, prefetch: u32) -> Self {
            self.prefetch = Some(prefetch);
            self
        }

        /// Sets the partition expiration duration for the event processor.
        pub fn with_partition_expiration_duration(
            mut self,
            partition_expiration_duration: Duration,
        ) -> Self {
            self.partition_expiration_duration = Some(partition_expiration_duration);
            self
        }

        /// Builds the event processor with the specified consumer client and checkpoint store.
        /// Returns a `Result` containing the constructed `EventProcessor`.
        ///
        /// # Connection options (including transport)
        ///
        /// The event processor does not open its own connection. It processes
        /// partitions using the [`ConsumerClient`] passed here, and every
        /// per-partition receiver reuses that client's connection. Connection-level
        /// options, such as the transport, a custom endpoint, retry options, and the
        /// application id, are therefore configured on the [`ConsumerClient`] before
        /// it is passed to `build`.
        ///
        /// To run the processor over AMQP-over-WebSockets (port 443, useful when the
        /// native AMQP ports are blocked), select the transport on the consumer
        /// client with
        /// [`ConsumerClientBuilder::with_transport`](crate::builders::ConsumerClientBuilder::with_transport):
        ///
        /// ```no_run
        /// use azure_messaging_eventhubs::{EventProcessor, CheckpointStore, ConsumerClient};
        /// use azure_messaging_eventhubs::models::AmqpTransport;
        /// use std::sync::Arc;
        ///
        /// async fn create_processor(checkpoint_store: Arc<dyn CheckpointStore>) -> Result<(), Box<dyn std::error::Error>> {
        /// use azure_identity::DeveloperToolsCredential;
        ///
        /// let eventhub_namespace = std::env::var("EVENTHUBS_HOST")?;
        /// let eventhub_name = std::env::var("EVENTHUB_NAME")?;
        /// let consumer = ConsumerClient::builder()
        ///     .with_transport(AmqpTransport::WebSocket)
        ///     .open(
        ///         &eventhub_namespace,
        ///         eventhub_name,
        ///         DeveloperToolsCredential::new(None)?.clone(),
        ///     )
        ///     .await?;
        /// let processor = EventProcessor::builder()
        ///     .build(consumer, checkpoint_store.clone())
        ///     .await?;
        /// Ok(())
        /// }
        /// ```
        pub async fn build(
            self,
            consumer_client: ConsumerClient,
            checkpoint_store: Arc<dyn CheckpointStore + Send + Sync>,
        ) -> Result<Arc<EventProcessor>> {
            let update_interval = self.update_interval.unwrap_or(DEFAULT_UPDATE_INTERVAL);
            let partition_expiration_duration = self
                .partition_expiration_duration
                .unwrap_or(DEFAULT_PARTITION_EXPIRATION_DURATION);

            validate_expiration_vs_update_interval(partition_expiration_duration, update_interval)?;

            // Retrieve the set of partitions from the consumer client
            // and limit the number of partitions to the specified max_partition_count.
            let mut eh_properties = consumer_client.get_eventhub_properties().await?;
            if let Some(max_partition_count) = self.max_partition_count {
                eh_properties.partition_ids.truncate(max_partition_count);
            }

            EventProcessor::new(
                consumer_client,
                checkpoint_store,
                super::EventProcessorOptions {
                    strategy: self
                        .load_balancing_strategy
                        .unwrap_or(super::ProcessorStrategy::Greedy),
                    partition_expiration_duration,
                    update_interval,
                    start_positions: self.start_positions.unwrap_or_default(),
                    prefetch: self.prefetch.unwrap_or(DEFAULT_PREFETCH),
                    partition_ids: eh_properties.partition_ids,
                },
            )
        }
    }
}

#[cfg(test)]
mod tests {
    use super::builders::validate_expiration_vs_update_interval;
    use super::{
        EventProcessor, EventProcessorOptions, PartitionClient, ProcessorConsumersMap,
        ProcessorStrategy, StartPositions,
    };
    use crate::{ConsumerClient, InMemoryCheckpointStore};
    use azure_core::time::Duration;
    use azure_core_test::credentials::MockCredential;
    use futures::SinkExt;
    use std::sync::Arc;

    /// Builds a processor that holds `partition_ids` queued partition clients,
    /// with no connection to the service. The returned map is the one that a
    /// `PartitionClient::close` removes itself from, so a test reads it to
    /// find out which clients closed.
    async fn processor_with_queued_clients(
        partition_ids: &[&str],
    ) -> (Arc<EventProcessor>, Arc<ProcessorConsumersMap>) {
        let consumer_client = ConsumerClient::new_unconnected(
            "example.servicebus.windows.net",
            "test-eventhub",
            Arc::new(MockCredential),
        )
        .expect("the client must build");
        let client_details = consumer_client.get_details().expect("details must parse");
        let checkpoint_store = Arc::new(InMemoryCheckpointStore::new());

        let processor = EventProcessor::new(
            consumer_client,
            checkpoint_store.clone(),
            EventProcessorOptions {
                strategy: ProcessorStrategy::Greedy,
                partition_expiration_duration: Duration::seconds(60),
                update_interval: Duration::seconds(30),
                start_positions: StartPositions::default(),
                prefetch: 300,
                partition_ids: partition_ids.iter().map(|id| id.to_string()).collect(),
            },
        )
        .expect("the processor must build");

        let consumers = Arc::new(ProcessorConsumersMap::new());
        let mut sender = processor.next_partition_client_sender.clone();
        for partition_id in partition_ids {
            let client = Arc::new(PartitionClient::new(
                partition_id.to_string(),
                checkpoint_store.clone(),
                client_details.clone(),
                Arc::downgrade(&consumers),
            ));
            consumers
                .add_partition_client(partition_id, client.clone())
                .await
                .expect("the map must accept the client");
            sender.send(client).await.expect("the queue must accept it");
        }

        (processor, consumers)
    }

    /// `close` must not stop at a partition client that the application still
    /// holds. It used to take each client out of its `Arc` and return an error
    /// when that failed, which left the clients behind it open and skipped the
    /// close of the consumer connection.
    #[tokio::test]
    async fn close_continues_past_a_retained_partition_client() {
        let (processor, consumers) = processor_with_queued_clients(&["0", "1"]).await;

        // Stand in for an application that holds the first client it took.
        let retained = consumers
            .consumers
            .lock()
            .expect("the map must lock")
            .get("0")
            .expect("partition 0 must be in the map")
            .upgrade()
            .expect("partition 0 must still be alive");

        let connection = {
            let Ok(processor) = Arc::try_unwrap(processor) else {
                panic!("the test must be the only holder of the processor");
            };
            let connection = processor.consumer_client.recoverable_connection();
            processor.close().await.expect("close must succeed");
            connection
        };

        let active = consumers
            .get_active_partition_ids()
            .expect("the map must lock");
        assert!(
            active.contains(&"0".to_string()),
            "the retained client must stay in the map, because it did not close"
        );
        assert!(
            !active.contains(&"1".to_string()),
            "the client behind the retained one must close, got: {active:?}"
        );
        assert!(
            connection.is_closed(),
            "the consumer connection must close after the partition clients"
        );

        drop(retained);
    }

    /// The validation must reject the historical default (expiration=10s,
    /// update_interval=30s). This combination is the root cause of issue
    /// #3851: the ownership record expires 20s before the next load-balancing
    /// cycle observes it, so every consumer reports `current=0` and
    /// re-claims every partition every cycle.
    #[test]
    fn historical_default_combination_is_rejected() {
        let result =
            validate_expiration_vs_update_interval(Duration::seconds(10), Duration::seconds(30));
        assert!(
            result.is_err(),
            "10s expiration with 30s interval must be rejected"
        );
    }

    /// The new default (expiration=60s, update_interval=30s) must be accepted.
    #[test]
    fn new_default_combination_is_accepted() {
        validate_expiration_vs_update_interval(Duration::seconds(60), Duration::seconds(30))
            .expect("60s expiration with 30s interval must be valid");
    }

    /// Equal values are rejected: if the record expires exactly at the
    /// instant the next cycle reads it, the read is racing the expiry.
    #[test]
    fn equal_values_are_rejected() {
        let result =
            validate_expiration_vs_update_interval(Duration::seconds(30), Duration::seconds(30));
        assert!(
            result.is_err(),
            "equal expiration and interval must be rejected"
        );
    }

    /// Custom configurations with adequate headroom are accepted.
    #[test]
    fn larger_expiration_is_accepted() {
        validate_expiration_vs_update_interval(Duration::seconds(120), Duration::seconds(60))
            .expect("2x ratio should be accepted");
    }
}