picodata-driver 3.0.0

use std::sync::Arc;

use crate::topology::Topology;

/// Connection pool storage with stable slot indices.
///
/// Each instance has a permanent `idx` (sequence counter from Topology).
/// `slots[idx]` is `Some(pool)` when online, `None` when offline.
#[derive(Debug, Default)]
pub(crate) struct Pool {
    /// Connection pool slots. Index = `instance.idx()`.
    slots: Vec<Option<Arc<pgpool::Pool>>>,
    /// Cluster topology.
    topology: Topology,
}

impl Pool {
    /// Creates an empty pool with no slots.
    pub(crate) fn new() -> Self {
        Self::default()
    }

    /// Returns the pool at slot `idx`, if present.
    pub(crate) fn get(&self, idx: usize) -> Option<&Arc<pgpool::Pool>> {
        self.slots.get(idx)?.as_ref()
    }

    /// Iterates over all active (online) pools.
    pub(crate) fn active_pools(&self) -> impl Iterator<Item = &Arc<pgpool::Pool>> {
        self.slots.iter().filter_map(|s| s.as_ref())
    }

    /// Returns the number of active (online) pools.
    pub(crate) fn active_count(&self) -> usize {
        self.slots.iter().filter(|s| s.is_some()).count()
    }

    /// Picks a pool by strategy index (nth active pool).
    pub(crate) fn pick(&self, strategy_idx: usize) -> Option<Arc<pgpool::Pool>> {
        let count = self.active_count();
        if count == 0 {
            return None;
        }
        let target = strategy_idx % count;
        self.slots
            .iter()
            .filter_map(|s| s.as_ref())
            .nth(target)
            .cloned()
    }

    /// Resolves a pool by tier, `bucket_id`, and master preference via topology.
    pub(crate) fn resolve(
        &self,
        tier: &str,
        bucket_id: u64,
        on_master: bool,
    ) -> Option<Arc<pgpool::Pool>> {
        let idx = self.topology.resolve(tier, bucket_id, on_master)?;
        self.get(idx).cloned()
    }

    /// Returns a reference to the current topology.
    pub(crate) fn topology(&self) -> &Topology {
        &self.topology
    }

    /// Returns the bucket count for a tier.
    pub(crate) fn bucket_count(&self, tier: &str) -> Option<u64> {
        self.topology.bucket_count(tier)
    }

    /// Merges incoming topology into the current state.
    pub(crate) fn apply_topology(&mut self, incoming: Topology) {
        self.topology.merge(incoming);
    }

    /// Synchronizes pool slots with topology: creates pools for online instances,
    /// closes pools for offline or removed instances.
    pub(crate) fn synchronize(
        &mut self,
        create_pool: impl Fn(&crate::instance::Instance) -> Arc<pgpool::Pool>,
    ) {
        let mut alive: std::collections::HashSet<usize> = std::collections::HashSet::new();

        for inst in self.topology.instances.values() {
            let idx = inst.idx();
            if inst.is_online() {
                alive.insert(idx);
                if idx >= self.slots.len() {
                    self.slots.resize(idx + 1, None);
                    self.slots[idx] = Some(create_pool(inst));
                } else if self.slots[idx].is_none() {
                    self.slots[idx] = Some(create_pool(inst));
                }
            } else if idx < self.slots.len()
                && let Some(pool) = self.slots[idx].take()
            {
                pool.close();
            }
        }

        for (idx, slot) in self.slots.iter_mut().enumerate() {
            if !alive.contains(&idx)
                && let Some(pool) = slot.take()
            {
                pool.close();
            }
        }
    }

    /// Closes all pools and resets topology.
    pub(crate) fn close(&mut self) {
        for slot in &mut self.slots {
            if let Some(pool) = slot.take() {
                pool.close();
            }
        }
        self.topology = Topology::new();
    }
}