picodata-driver 3.0.0

use std::{
    fmt,
    net::SocketAddr,
    sync::atomic::{AtomicU8, AtomicUsize, Ordering},
};

use uuid::Uuid;

use crate::error::Error;

/// State of instance.
/// As far as we need only online and offline states - other states are ignored.
#[derive(Clone, Debug, Hash, PartialEq, Eq)]
pub(crate) enum State {
    /// Instance is up and running.
    Online,
    /// Instance if down
    Offline,
}

impl From<&str> for State {
    fn from(value: &str) -> Self {
        match value {
            "Online" | "online" => Self::Online,
            _ => Self::Offline,
        }
    }
}

impl From<State> for u8 {
    fn from(value: State) -> Self {
        match value {
            State::Online => 1,
            State::Offline => 0,
        }
    }
}

impl From<State> for AtomicU8 {
    fn from(value: State) -> Self {
        match value {
            State::Online => AtomicU8::new(1),
            State::Offline => AtomicU8::new(0),
        }
    }
}

impl From<u8> for State {
    fn from(value: u8) -> Self {
        if value == 1 {
            Self::Online
        } else {
            Self::Offline
        }
    }
}

/// Picodata instance info.
#[derive(Debug)]
pub(crate) struct Instance {
    /// Index into Pool slots. Assigned by Topology via atomic store.
    idx: AtomicUsize,
    /// Instance name — unique stable identifier in Picodata.
    pub(crate) name: String,
    /// Address of instance, as advertised by the cluster.
    pub(crate) addr: Addr,
    /// Actual state represented as u8. May be online of offline.
    cond: AtomicU8,
    /// Tier name this instance belongs to.
    pub(crate) tier: String,
    /// Bucket count for this instance's tier.
    pub(crate) bucket_count: u64,
    /// Replicaset UUID.
    pub(crate) replicaset_uuid: Uuid,
    /// Whether this instance is the current master of its replicaset.
    pub(crate) is_master: bool,
    /// Raft id of the instance.
    pub(crate) raft_id: u64,
    /// Instance UUID.
    pub(crate) uuid: Uuid,
    /// Name of the replicaset this instance belongs to.
    pub(crate) replicaset_name: String,
}

impl Instance {
    /// Get current state of intance.
    ///
    /// # Returns
    ///
    /// A `State` of instance.
    pub(crate) fn state(&self) -> State {
        State::from(self.cond.load(Ordering::Acquire))
    }

    /// Sets instance state atomically.
    pub(crate) fn set_state(&self, state: State) {
        self.cond.store(state.into(), Ordering::Release);
    }

    /// Returns `true` if the instance is online.
    pub(crate) fn is_online(&self) -> bool {
        self.state() == State::Online
    }

    /// Returns the pool slot index assigned by topology.
    pub(crate) fn idx(&self) -> usize {
        self.idx.load(Ordering::Acquire)
    }

    /// Sets the pool slot index atomically.
    pub(crate) fn set_idx(&self, idx: usize) {
        self.idx.store(idx, Ordering::Release);
    }
}

/// Address of an instance: a literal socket address or a host name.
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
pub(crate) enum Addr {
    /// A literal address.
    Ip(SocketAddr),
    /// A host name, resolved by the client on every connection.
    Name { host: Box<str>, port: u16 },
}

impl fmt::Display for Addr {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Addr::Ip(addr) => write!(f, "{addr}"),
            Addr::Name { host, port } => write!(f, "{host}:{port}"),
        }
    }
}

impl std::str::FromStr for Addr {
    type Err = Error;

    /// Parses an advertised `HOST:PORT`.
    fn from_str(address: &str) -> Result<Self, Self::Err> {
        if let Ok(addr) = address.parse::<SocketAddr>() {
            return Ok(Addr::Ip(addr));
        }
        let (host, port) = address
            .rsplit_once(':')
            .ok_or_else(|| Error::InvalidAddress(format!("{address}: expected HOST:PORT")))?;
        if host.is_empty() {
            return Err(Error::InvalidAddress(format!("{address}: empty host")));
        }
        let port = port
            .parse::<u16>()
            .map_err(|e| Error::InvalidAddress(format!("{address}: {e}")))?;
        Ok(Addr::Name {
            host: host.into(),
            port,
        })
    }
}

impl TryFrom<pg::Row> for Instance {
    type Error = Error;

    fn try_from(value: pg::Row) -> Result<Self, Self::Error> {
        let address = value.get::<_, &str>("address");
        let current_state = value.get::<_, serde_json::Value>("current_state");
        let state = current_state
            .as_array()
            .and_then(|x| x.first())
            .and_then(|x| x.as_str())
            .unwrap_or_default();
        let addr = address.parse::<Addr>()?;
        let tier = value.get::<_, &str>("tier").to_string();
        let bucket_count = u64::try_from(value.get::<_, i64>("bucket_count"))
            .expect("bucket_count is non-negative");
        let replicaset_uuid = Uuid::parse_str(value.try_get::<_, &str>("replicaset_uuid")?)?;
        let name = value.get::<_, &str>("instance_name").to_string();
        let current_master_name = value
            .try_get::<_, &str>("current_master_name")
            .unwrap_or_default();
        let is_master = name == current_master_name;

        Ok(Self {
            idx: AtomicUsize::new(0),
            name,
            addr,
            cond: State::from(state).into(),
            tier,
            bucket_count,
            replicaset_uuid,
            is_master,
            raft_id: value.try_get::<_, i64>("raft_id")?.cast_unsigned(),
            uuid: Uuid::parse_str(value.try_get::<_, &str>("instance_uuid")?)?,
            replicaset_name: value.try_get::<_, &str>("replicaset_name")?.to_string(),
        })
    }
}

impl PartialEq for Instance {
    fn eq(&self, other: &Self) -> bool {
        self.addr == other.addr
    }
}

impl std::hash::Hash for Instance {
    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
        self.addr.hash(state);
    }
}

impl Eq for Instance {}

#[cfg(test)]
mod tests {
    use super::*;

    /// Regression: the address was parsed with `SocketAddr::from_str` alone,
    /// which takes a literal address only, so a cluster advertising service
    /// names lost every instance and left the driver with an empty pool.
    #[test]
    fn addr_tells_a_name_from_a_literal_address() {
        assert_eq!(
            "picodata-1:5432".parse::<Addr>().expect("a name is valid"),
            Addr::Name {
                host: "picodata-1".into(),
                port: 5432
            }
        );
        assert_eq!(
            "127.0.0.1:5433".parse::<Addr>().expect("an IP still works"),
            Addr::Ip("127.0.0.1:5433".parse().expect("literal"))
        );
        // picodata cannot advertise IPv6 yet; a bracketed one must still parse as literal.
        assert_eq!(
            "[::1]:5433"
                .parse::<Addr>()
                .expect("IPv6 is a literal address"),
            Addr::Ip("[::1]:5433".parse().expect("literal"))
        );
    }

    /// Only what picodata itself refuses to store is rejected here.
    #[test]
    fn addr_rejects_a_malformed_address() {
        for bad in ["localhost", "localhost:", "localhost:abc", ":5433", ""] {
            assert!(
                matches!(bad.parse::<Addr>(), Err(Error::InvalidAddress(_))),
                "{bad:?} must not parse"
            );
        }
    }

    impl Instance {
        pub(crate) fn new_for_test(
            name: &str,
            addr: &str,
            tier: &str,
            rs_uuid: Uuid,
            is_master: bool,
            bucket_count: u64,
        ) -> Self {
            Self {
                idx: AtomicUsize::new(0),
                name: name.to_string(),
                addr: addr.parse().expect("valid HOST:PORT"),
                cond: State::Online.into(),
                tier: tier.to_string(),
                bucket_count,
                replicaset_uuid: rs_uuid,
                is_master,
                raft_id: 0,
                uuid: Uuid::nil(),
                replicaset_name: String::new(),
            }
        }
    }
}