nusadb 0.1.0

//! A small, thread-safe connection pool. [`Pool::get`] hands out a [`PooledConnection`] that returns
//! the connection to the pool on drop; checkout blocks (with a [`std::sync::Condvar`]) when all
//! connections are in use, up to `max_size`.

use std::ops::{Deref, DerefMut};
use std::sync::{Arc, Condvar, Mutex};

use crate::connection::{Config, Connection};
use crate::error::{Error, Result};

struct Inner {
    idle: Vec<Connection>,
    open: usize,
}

/// A pool of NusaDB connections sharing one [`Config`].
#[derive(Clone)]
pub struct Pool {
    config: Arc<Config>,
    max_size: usize,
    state: Arc<(Mutex<Inner>, Condvar)>,
}

impl Pool {
    /// Create a pool of up to `max_size` connections (created lazily on demand). `max_size` is
    /// clamped to at least 1.
    pub fn new(config: Config, max_size: usize) -> Result<Self> {
        let max_size = max_size.max(1);
        Ok(Self {
            config: Arc::new(config),
            max_size,
            state: Arc::new((
                Mutex::new(Inner {
                    idle: Vec::with_capacity(max_size),
                    open: 0,
                }),
                Condvar::new(),
            )),
        })
    }

    /// Check out a connection, blocking until one is free if the pool is at `max_size`. A fresh
    /// connection is opened on demand if the pool has not yet reached its cap.
    pub fn get(&self) -> Result<PooledConnection> {
        let (lock, cvar) = &*self.state;
        let mut inner = lock
            .lock()
            .map_err(|_| Error::Pool("pool mutex poisoned".to_owned()))?;
        loop {
            if let Some(conn) = inner.idle.pop() {
                return Ok(PooledConnection {
                    conn: Some(conn),
                    pool: self.clone(),
                });
            }
            if inner.open < self.max_size {
                inner.open += 1;
                // Release the lock while connecting (network I/O) so other threads progress.
                drop(inner);
                let conn = match Connection::connect_config(&self.config) {
                    Ok(c) => c,
                    Err(e) => {
                        // Roll back the reservation so the slot is reusable.
                        if let Ok(mut g) = lock.lock() {
                            g.open -= 1;
                            cvar.notify_one();
                        }
                        return Err(e);
                    }
                };
                return Ok(PooledConnection {
                    conn: Some(conn),
                    pool: self.clone(),
                });
            }
            // At capacity: wait for a returned connection.
            inner = cvar
                .wait(inner)
                .map_err(|_| Error::Pool("pool mutex poisoned".to_owned()))?;
        }
    }

    /// The maximum number of connections this pool will open.
    #[must_use]
    pub fn max_size(&self) -> usize {
        self.max_size
    }

    fn put_back(&self, conn: Connection) {
        let (lock, cvar) = &*self.state;
        if let Ok(mut inner) = lock.lock() {
            inner.idle.push(conn);
            cvar.notify_one();
        }
    }

    fn discard(&self) {
        let (lock, cvar) = &*self.state;
        if let Ok(mut inner) = lock.lock() {
            inner.open = inner.open.saturating_sub(1);
            cvar.notify_one();
        }
    }
}

/// A connection checked out of a [`Pool`]. Derefs to [`Connection`]; returns to the pool on drop
/// unless [`PooledConnection::discard`] was called.
pub struct PooledConnection {
    conn: Option<Connection>,
    pool: Pool,
}

impl PooledConnection {
    /// Drop this connection instead of returning it to the pool (e.g. after a fatal I/O error left it
    /// in an unknown state). The pool may open a fresh one to replace it.
    pub fn discard(mut self) {
        self.conn.take();
        self.pool.discard();
    }
}

impl Deref for PooledConnection {
    type Target = Connection;
    fn deref(&self) -> &Connection {
        self.conn.as_ref().expect("connection present until drop")
    }
}

impl DerefMut for PooledConnection {
    fn deref_mut(&mut self) -> &mut Connection {
        self.conn.as_mut().expect("connection present until drop")
    }
}

impl Drop for PooledConnection {
    fn drop(&mut self) {
        if let Some(conn) = self.conn.take() {
            self.pool.put_back(conn);
        }
    }
}