picodata-driver 3.0.0

use std::sync::Arc;

use tokio::{
    sync::oneshot::{Receiver, channel},
    task::JoinHandle,
};
use tokio_util::sync::CancellationToken;

use crate::{Error, accessor::Accessor, config::Config, topology::Topology};

/// Receiver to recv topology init result.
pub(crate) type TopologyInitReceiver = Receiver<Result<(), Error>>;

/// Discovers cluster with provided interval to check if
/// there are new instances added or some instances removed.
pub(crate) struct Discovery {
    /// Accessor to a pool.
    accessor: Arc<Accessor>,
    /// Config to create new pools.
    config: Config,
    /// Token which is checked to stop underlying discovery task.
    token: CancellationToken,
}

impl Discovery {
    /// SQL query to fetch cluster topology from picodata system tables, one row per address.
    const TOPOLOGY_QUERY: &str = r"
    SELECT ppa.address, pi.current_state, pi.tier, pi.name AS instance_name,
           pi.replicaset_name, pt.bucket_count,
           pr.uuid AS replicaset_uuid, pr.current_master_name,
           pi.raft_id, pi.uuid AS instance_uuid, ppa.connection_type
	FROM _pico_peer_address AS ppa
    JOIN _pico_instance AS pi ON ppa.raft_id = pi.raft_id
    LEFT JOIN _pico_tier AS pt ON pi.tier = pt.name
    LEFT JOIN _pico_replicaset AS pr ON pr.name = pi.replicaset_name
	WHERE  ppa.connection_type IS NOT NULL
    ";

    /// Creates a new instance of `Manager`.
    ///
    /// # Parameters
    ///
    /// - `accessor`: An accessor to pool of connections.
    /// - `config`: A config to create new pools.
    /// - `token`: A token to cancel task.
    ///
    /// # Returns
    ///
    /// An instance of `Discovery`.
    pub(crate) fn new(accessor: Arc<Accessor>, config: Config, token: CancellationToken) -> Self {
        Self {
            accessor,
            config,
            token,
        }
    }

    /// Scans cluster instances.
    ///
    /// This function does check following conditions:
    ///
    /// - Instances were added. In this case new connection pools will be added.
    /// - Instances changed their state. If instance state is online - new pool will be added. If state is offline - instance pool will be removed.
    ///
    /// # Errors
    ///
    /// - Any [`Error`] returned by inner driver.
    ///
    /// # Returns
    ///
    /// A unit on success.
    /// Fetches topology and bucket ranges, merges into current, synchronizes pools.
    async fn refresh(&self, client: &pg::Client) -> Result<(), Error> {
        let excluded_tiers = self.config.excluded_tiers.clone().unwrap_or_default();
        let query = if excluded_tiers.is_empty() {
            Self::TOPOLOGY_QUERY.to_string()
        } else {
            let tiers = excluded_tiers
                .iter()
                .map(|x| format!("'{x}'"))
                .collect::<Vec<String>>()
                .join(",");
            format!("{} AND pi.tier NOT IN ({})", Self::TOPOLOGY_QUERY, tiers)
        };
        let topology_rows = client.query(&query, &[]).await.inspect_err(|e| {
            log::error!("failed to query topology: {e}");
        })?;

        let mut tiers = std::collections::HashSet::new();
        for row in &topology_rows {
            if let Ok(tier) = row.try_get::<_, &str>("tier")
                && !tier.is_empty()
                && !tiers.contains(tier)
            {
                tiers.insert(tier.to_string());
            }
        }

        let mut bucket_rows = Vec::new();
        for tier in &tiers {
            let query = format!("SELECT _pico_bucket('{tier}')");
            if let Ok(row) = client.query_one(&query, &[]).await {
                bucket_rows.push(row);
            }
        }

        let incoming = Topology::from((topology_rows, bucket_rows));
        let config = self.config.clone();
        let mut guard = self.accessor.write().await;
        guard.apply_topology(incoming);
        guard.synchronize(|inst| Arc::new(build_pool(&config.clone_for(inst))));
        Ok(())
    }

    /// Runs refresh using any available pooled connection.
    async fn once(&self) -> Result<(), Error> {
        let guard = self.accessor.read().await;
        let pools: Vec<_> = guard.active_pools().cloned().collect();
        drop(guard);

        for pool in &pools {
            match pool.get().await {
                Ok(client) => return self.refresh(&client).await,
                Err(e) => {
                    log::debug!("discovery: pool unavailable, trying next: {e}");
                }
            }
        }

        Err(Error::IO(std::io::Error::new(
            std::io::ErrorKind::ConnectionRefused,
            "discovery: all pools unavailable",
        )))
    }

    /// Starts new tokio task which discovers cluster each time interval is reached.
    ///
    /// # Returns
    ///
    /// A `JoinHandle` for running tokio task and topology init receiver.
    pub(crate) fn start(self) -> (JoinHandle<()>, TopologyInitReceiver) {
        let (sender, receiver) = channel();
        let handle = tokio::spawn(async move {
            // First refresh using a temporary deadpool pool (it spawns Connection internally).
            let init_result = async {
                let pool = build_pool(&self.config);
                let client = pool.get().await.map_err(Error::PostgresPool)?;
                self.refresh(&client).await
            }
            .await;

            sender
                .send(init_result)
                .expect("sending topology init result should not fail");
            let mut ticker = tokio::time::interval(self.config.discorery_interval);
            loop {
                tokio::select! {
                    _ = ticker.tick() => {
                        _ = self.once().await;
                    }
                    () = self.token.cancelled() => {
                        break
                    }
                }
            }
        });
        (handle, receiver)
    }
}

/// Creates a deadpool connection pool from config (handles TLS if enabled).
fn build_pool(cfg: &Config) -> pgpool::Pool {
    #[cfg(feature = "tls")]
    let mgr = match crate::tls::connector(cfg) {
        Ok(Some(tls)) => pgpool::Manager::new(cfg.postgres.clone(), tls),
        _ => pgpool::Manager::new(cfg.postgres.clone(), pg::NoTls),
    };
    #[cfg(not(feature = "tls"))]
    let mgr = pgpool::Manager::new(cfg.postgres.clone(), pg::NoTls);

    pgpool::Pool::builder(mgr)
        .max_size(cfg.pool_max_size)
        .build()
        .expect("building new deadpool should not fail")
}