picodata-driver 3.0.0

use std::sync::Arc;

use tokio::sync::{RwLock, RwLockReadGuard, RwLockWriteGuard};

use crate::{BoxStrategy, error::Error, pool::Pool};

/// Allows to access to instances connection pools.
pub(crate) struct Accessor {
    /// Actual pool of connections.
    pool: RwLock<Pool>,
    /// Strategy which picks which connection to choose.
    strategy: BoxStrategy,
}

impl Accessor {
    /// Creates a new `Accessor` with a provided pool and strategy.
    pub(crate) fn new(pool: Pool, strategy: BoxStrategy) -> Self {
        Self {
            pool: RwLock::new(pool),
            strategy,
        }
    }

    /// Acquire a read only lock for a pool.
    #[inline(always)]
    pub(crate) async fn read(&self) -> RwLockReadGuard<'_, Pool> {
        self.pool.read().await
    }

    /// Acquire a write lock for a pool.
    #[inline(always)]
    pub(crate) async fn write(&self) -> RwLockWriteGuard<'_, Pool> {
        self.pool.write().await
    }

    /// Route to a specific pool by tier, `bucket_id`, and master preference.
    /// Falls back to strategy if the bucket cannot be resolved.
    ///
    /// # Errors
    ///
    /// Returns [`Error::Unreachable`] if no active pools are available.
    pub(crate) async fn route(
        &self,
        tier: &str,
        bucket_id: u64,
        on_master: bool,
    ) -> Result<Arc<pgpool::Pool>, Error> {
        let guard = self.pool.read().await;
        if let Some(pool) = guard.resolve(tier, bucket_id, on_master) {
            return Ok(pool);
        }
        drop(guard);
        log::debug!("route: bucket_id={bucket_id} not resolved, fallback round-robin");
        self.next().await
    }

    /// Get next pool via strategy (round-robin / random).
    ///
    /// # Errors
    ///
    /// Returns [`Error::Unreachable`] if no active pools are available.
    pub(crate) async fn next(&self) -> Result<Arc<pgpool::Pool>, Error> {
        let guard = self.pool.read().await;
        let index = self.strategy.next(guard.active_count());
        guard.pick(index).ok_or(Error::Unreachable)
    }

    /// Calls a provided callback sequentially for each instance of a cluster.
    #[inline(always)]
    pub(crate) async fn map<C, F>(&self, c: C) -> Result<(), Error>
    where
        F: Future<Output = Result<(), Error>>,
        C: Fn(Arc<pgpool::Pool>) -> F,
    {
        let guard = self.pool.read().await;
        for pool in guard.active_pools() {
            c(pool.clone()).await?;
        }
        Ok(())
    }

    /// Closes pool.
    pub(crate) async fn close(&self) {
        self.pool.write().await.close();
    }
}