picodata-driver 3.0.0

use std::sync::{
    Arc,
    atomic::{AtomicBool, Ordering},
};

use tokio::task::JoinHandle;
use tokio_util::sync::CancellationToken;

use crate::{
    Strategy,
    accessor::Accessor,
    config::Config,
    discovery::Discovery,
    error::Error,
    metadata::{Metadata, MetadataHandle},
    topology::view::Topology,
};

/// Manager of each instance connection pool.
pub(crate) struct Manager {
    is_closed: Arc<AtomicBool>,
    accessor: Arc<Accessor>,
    token: CancellationToken,
    _discovery: JoinHandle<()>,
    metadata: MetadataHandle,
}

impl Manager {
    /// Creates a new manager: starts discovery and metadata services.
    pub(crate) async fn new<S: Strategy>(config: Config, strategy: S) -> Result<Self, Error> {
        let token = CancellationToken::new();
        let pool = crate::pool::Pool::new();
        let accessor = Arc::new(Accessor::new(pool, Box::new(strategy)));

        let (discovery, init) =
            Discovery::new(accessor.clone(), config.clone(), token.clone()).start();
        init.await
            .expect("waiting for discovery initializer should not fail")?;

        let (metadata, metadata_init) =
            Metadata::new(accessor.clone(), config, token.clone()).start();
        if let Err(e) = metadata_init
            .await
            .expect("waiting for metadata initializer should not fail")
        {
            log::warn!("metadata service init failed, shard routing disabled: {e}");
        }

        let is_closed = Arc::new(AtomicBool::new(false));

        Ok(Self {
            accessor,
            token,
            is_closed,
            _discovery: discovery,
            metadata,
        })
    }

    /// Returns the pool of the named instance; a pinned query never falls back to another one.
    pub(crate) async fn route_to(&self, name: &str) -> Result<Arc<pgpool::Pool>, Error> {
        self.check_closed()?;
        let guard = self.accessor.read().await;
        guard
            .topology()
            .instances
            .get(name)
            .and_then(|inst| guard.get(inst.idx()).cloned())
            .ok_or_else(|| Error::InstanceUnavailable(name.to_string()))
    }

    /// Routes a query to the appropriate pool using `dk_meta` and `bucket_id`.
    /// Falls back to round-robin when shard routing is not possible.
    pub(crate) async fn route(
        &self,
        sql: &str,
        params: &[&(dyn pg::types::ToSql + Sync)],
        on_master: bool,
        forced_bucket_id: Option<u64>,
    ) -> Result<Arc<pgpool::Pool>, Error> {
        self.check_closed()?;

        if let Some(bucket_id) = forced_bucket_id {
            log::debug!("route: forced bucket_id={bucket_id}, on_master={on_master}");
            let dk = self.metadata.resolve(sql).await;
            let tier = dk.as_ref().and_then(|dk| dk.tier.as_deref());
            return if let Some(tier) = tier {
                self.accessor.route(tier, bucket_id, on_master).await
            } else {
                self.accessor.next().await
            };
        }

        let dk = self.metadata.resolve(sql).await;

        let Some(dk) = dk else {
            log::debug!("route: no dk_meta, using round-robin");
            return self.accessor.next().await;
        };
        let Some(tier) = &dk.tier else {
            log::debug!("route: no tier in dk_meta, using round-robin");
            return self.accessor.next().await;
        };
        if dk.keys.is_empty() {
            log::debug!("route: empty dk_meta keys, using round-robin");
            return self.accessor.next().await;
        }

        let bucket_count = self.accessor.read().await.bucket_count(tier);
        let Some(bc) = bucket_count else {
            return self.accessor.next().await;
        };
        let Some(bucket_id) = crate::bucket::compute_bucket_id(params, &dk, bc) else {
            log::debug!("route: failed to compute bucket_id, using round-robin");
            return self.accessor.next().await;
        };

        log::debug!("route: bucket_id={bucket_id} in tier={tier}, on_master={on_master}");
        self.accessor.route(tier, bucket_id, on_master).await
    }

    /// Returns a snapshot of the topology.
    pub(crate) async fn topology(&self) -> Result<Topology, Error> {
        self.check_closed()?;
        Ok(self.accessor.read().await.topology().snapshot())
    }

    /// Handles query errors. Invalidates `dk_meta` cache on DDL invalidation (SQLSTATE 42999).
    pub(crate) async fn on_error(&self, sql: &str, error: &pg::Error) {
        if is_statement_invalidated(error) {
            self.metadata.invalidate(sql).await;
        }
    }

    /// Calls the callback for each active pool sequentially.
    pub(crate) async fn map<C, F>(&self, c: C) -> Result<(), Error>
    where
        F: Future<Output = Result<(), Error>>,
        C: Fn(Arc<pgpool::Pool>) -> F,
    {
        self.check_closed()?;
        self.accessor.map(c).await
    }

    /// Closes the manager and cancels background tasks.
    pub(crate) async fn close(&self) -> Result<(), Error> {
        self.check_closed()?;
        self.is_closed.store(true, Ordering::Relaxed);
        self.token.cancel();
        self.accessor.close().await;
        Ok(())
    }

    /// Returns `Err(Error::Closed)` if the manager has been shut down.
    fn check_closed(&self) -> Result<(), Error> {
        if self.is_closed.load(Ordering::Relaxed) {
            return Err(Error::Closed);
        }
        Ok(())
    }
}

/// Picodata returns SQLSTATE 42999 when a prepared statement has been
/// invalidated by a DDL change. The driver should evict cached `dk_meta`
/// so that the next call triggers a fresh Parse → Notice cycle.
fn is_statement_invalidated(error: &pg::Error) -> bool {
    error.code().is_some_and(|c| c.code() == "42999")
}