picodata-driver 3.0.0

use std::{collections::HashMap, sync::Arc};

use serde::Deserialize;
use tokio::{
    sync::{RwLock, mpsc},
    task::JoinHandle,
};
use tokio_util::sync::CancellationToken;

use crate::{accessor::Accessor, config::Config};

/// Distribution key metadata received from server via Notice after Parse.
#[derive(Clone, Debug)]
pub(crate) struct DkMeta {
    /// Tier name for bucket routing. None for global tables.
    pub tier: Option<String>,
    /// (`param_index`, `oid`) pairs forming the distribution key.
    pub keys: Vec<(u16, u32)>,
}

/// Parses a picodata Notice message into `(query, DkMeta)`.
/// Returns `None` if the message is not a `PreparedStatementMetadata` JSON.
pub(crate) fn parse_dk_notice(notice: &pg::error::DbError) -> Option<(String, DkMeta)> {
    let parsed: NoticePayload = serde_json::from_str(notice.message()).ok()?;
    Some((
        parsed.query,
        DkMeta {
            tier: parsed.tier,
            keys: parsed.dk_meta,
        },
    ))
}

/// JSON structure of the `PreparedStatementMetadata` Notice from picodata.
#[derive(Deserialize)]
struct NoticePayload {
    query: String,
    tier: Option<String>,
    dk_meta: Vec<(u16, u32)>,
}

/// Receiver for the metadata service initialization result.
pub(crate) type MetadataInitReceiver = tokio::sync::oneshot::Receiver<Result<(), crate::Error>>;

/// Shared cache: SQL query → optional `DkMeta` (None for queries without distribution key).
type DkCache = Arc<RwLock<HashMap<String, Option<Arc<DkMeta>>>>>;

/// Active connection: client for sending queries, notice channel for receiving
/// `dk_meta` from the spawned connection driver task.
struct ServiceConnection {
    client: pg::Client,
    notice_rx: mpsc::UnboundedReceiver<pg::error::DbError>,
    _conn_task: JoinHandle<()>,
}

/// Background service that caches distribution key metadata by sending Parse
/// requests to picodata and receiving Notice messages with `dk_meta`.
pub(crate) struct Metadata {
    accessor: Arc<Accessor>,
    config: Config,
    token: CancellationToken,
    cache: DkCache,
    request_tx: mpsc::UnboundedSender<String>,
    request_rx: mpsc::UnboundedReceiver<String>,
}

impl Metadata {
    /// Creates a new metadata service (not yet started).
    pub(crate) fn new(accessor: Arc<Accessor>, config: Config, token: CancellationToken) -> Self {
        let (request_tx, request_rx) = mpsc::unbounded_channel();
        Self {
            accessor,
            config,
            token,
            cache: Arc::new(RwLock::new(HashMap::new())),
            request_tx,
            request_rx,
        }
    }

    /// Spawns the background task and returns a handle + init receiver.
    pub(crate) fn start(self) -> (MetadataHandle, MetadataInitReceiver) {
        let (init_tx, init_rx) = tokio::sync::oneshot::channel();
        let handle = MetadataHandle {
            cache: self.cache.clone(),
            request_tx: self.request_tx.clone(),
            _task: tokio::spawn(self.run(init_tx)),
        };
        (handle, init_rx)
    }

    async fn handle_notice(&self, notice: pg::error::DbError) {
        if let Some((query, dk)) = parse_dk_notice(&notice) {
            log::debug!(
                "metadata: received dk_meta notice: query={query}, tier={:?}, keys={:?}",
                dk.tier,
                dk.keys
            );
            let meta = if dk.keys.is_empty() {
                None
            } else {
                Some(Arc::new(dk))
            };
            self.cache.write().await.insert(query, meta);
        } else {
            log::debug!("metadata: received non-dk notice: {}", notice.message());
        }
    }

    async fn handle_request(&self, client: &pg::Client, sql: &str) -> bool {
        if self.cache.read().await.contains_key(sql) {
            log::debug!("metadata: cache hit, skipping prepare for: {sql}");
            return true;
        }
        log::debug!("metadata: cache miss, preparing: {sql}");
        match client.prepare(sql).await {
            Ok(_) => true,
            Err(e) if e.is_closed() => false,
            Err(e) => {
                log::error!("metadata: prepare failed: {e}");
                true
            }
        }
    }

    async fn connect(accessor: &Accessor, config: &Config) -> Option<ServiceConnection> {
        let guard = accessor.read().await;
        let instances: Vec<_> = guard
            .topology()
            .instances
            .values()
            .filter(|i| i.is_online())
            .cloned()
            .collect();
        drop(guard);

        log::debug!(
            "metadata: found {} online instances to try",
            instances.len()
        );

        for inst in instances {
            let mut cfg = config.clone_for(&inst);
            cfg.postgres.options("pico_query_metadata=true");

            if let Some(svc) = open_service_connection(&cfg).await {
                log::debug!("metadata: connected to {}", inst.addr);
                return Some(svc);
            }
            log::error!("metadata: failed to connect to {}", inst.addr);
        }

        None
    }

    async fn run(mut self, init_tx: tokio::sync::oneshot::Sender<Result<(), crate::Error>>) {
        let Some(mut svc) = Self::connect(&self.accessor, &self.config).await else {
            init_tx
                .send(Err(crate::Error::IO(std::io::Error::new(
                    std::io::ErrorKind::ConnectionRefused,
                    "metadata: no online instances to connect",
                ))))
                .expect("metadata init receiver should not be dropped");
            return;
        };

        init_tx
            .send(Ok(()))
            .expect("metadata init receiver should not be dropped");

        loop {
            let reconnect = tokio::select! {
                notice = svc.notice_rx.recv() => match notice {
                    Some(n) => { self.handle_notice(n).await; false }
                    None => true,
                },
                Some(sql) = self.request_rx.recv() => {
                    !self.handle_request(&svc.client, &sql).await
                }
                () = self.token.cancelled() => return,
            };

            if reconnect {
                log::warn!("metadata: connection lost, reconnecting");
                if let Some(new_svc) = Self::connect(&self.accessor, &self.config).await {
                    svc = new_svc;
                }
            }
        }
    }
}

/// Handle used by `Manager` to interact with the running `MetadataService`.
pub(crate) struct MetadataHandle {
    cache: DkCache,
    request_tx: mpsc::UnboundedSender<String>,
    _task: JoinHandle<()>,
}

impl MetadataHandle {
    /// Returns cached `dk_meta` for the query, or sends a request to the service on cache miss.
    pub(crate) async fn resolve(&self, query: &str) -> Option<Arc<DkMeta>> {
        if let Some(meta) = self.cache.read().await.get(query) {
            log::debug!("metadata: resolve cache hit for: {query}");
            return meta.clone();
        }

        log::debug!("metadata: resolve cache miss, requesting: {query}");
        if self.request_tx.send(query.to_string()).is_err() {
            log::debug!("metadata: service not running, skipping request");
        }

        None
    }

    /// Removes cached `dk_meta` for the query (e.g. after DDL invalidation).
    pub(crate) async fn invalidate(&self, query: &str) {
        log::debug!("metadata: invalidating cache for: {query}");
        self.cache.write().await.remove(query);
    }
}

/// Timeout for establishing a metadata service connection to picodata.
const CONNECT_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(5);

/// Opens a raw pg connection with `pico_query_metadata=true` and spawns its driver task.
async fn open_service_connection(cfg: &Config) -> Option<ServiceConnection> {
    #[cfg(feature = "tls")]
    if let Ok(Some(tls)) = crate::tls::connector(cfg) {
        return open_tls_connection(cfg, tls).await;
    }

    open_plain_connection(cfg).await
}

/// Opens a TLS connection and spawns the driver task.
#[cfg(feature = "tls")]
async fn open_tls_connection(
    cfg: &Config,
    tls: tokio_postgres_rustls::MakeRustlsConnect,
) -> Option<ServiceConnection> {
    let result = tokio::time::timeout(CONNECT_TIMEOUT, cfg.postgres.connect(tls)).await;
    match result {
        Ok(Ok((client, mut conn))) => {
            let (notice_tx, notice_rx) = mpsc::unbounded_channel();
            let conn_task = tokio::spawn(async move {
                drive_connection(&mut conn, &notice_tx).await;
            });
            Some(ServiceConnection {
                client,
                notice_rx,
                _conn_task: conn_task,
            })
        }
        Ok(Err(e)) => {
            log::debug!("metadata: tls connect error: {e}");
            None
        }
        Err(_) => {
            log::debug!("metadata: tls connect timeout");
            None
        }
    }
}

/// Opens a plain (no TLS) connection and spawns the driver task.
async fn open_plain_connection(cfg: &Config) -> Option<ServiceConnection> {
    let result = tokio::time::timeout(CONNECT_TIMEOUT, cfg.postgres.connect(pg::NoTls)).await;
    match result {
        Ok(Ok((client, mut conn))) => {
            let (notice_tx, notice_rx) = mpsc::unbounded_channel();
            let conn_task = tokio::spawn(async move {
                drive_connection(&mut conn, &notice_tx).await;
            });
            Some(ServiceConnection {
                client,
                notice_rx,
                _conn_task: conn_task,
            })
        }
        Ok(Err(e)) => {
            log::debug!("metadata: connect error: {e}");
            None
        }
        Err(_) => {
            log::debug!("metadata: connect timeout");
            None
        }
    }
}

/// Polls the pg connection, forwarding Notice messages to the channel.
async fn drive_connection<S, T>(
    conn: &mut pg::Connection<S, T>,
    tx: &mpsc::UnboundedSender<pg::error::DbError>,
) where
    S: tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin,
    T: tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin,
{
    log::debug!("metadata: connection driver task started");
    loop {
        match std::future::poll_fn(|cx| conn.poll_message(cx)).await {
            Some(Ok(pg::AsyncMessage::Notice(n))) => {
                log::debug!("metadata: conn task got Notice: {}", n.message());
                if tx.send(n).is_err() {
                    break;
                }
            }
            Some(Ok(msg)) => {
                log::debug!("metadata: conn task got async msg: {msg:?}");
            }
            Some(Err(e)) => {
                log::error!("metadata: connection error: {e}");
                break;
            }
            None => {
                log::debug!("metadata: connection closed");
                break;
            }
        }
    }
}