mobc 0.9.0

A generic connection pool with async/await support
Documentation
use std::{
    marker::PhantomData,
    sync::{
        atomic::{AtomicU64, Ordering},
        Arc,
    },
    time::{Duration, Instant},
};

use metrics::counter;
use tokio::sync::OwnedSemaphorePermit;

use crate::metrics_utils::{
    GaugeGuard, ACTIVE_CONNECTIONS, CLOSED_TOTAL, IDLE_CONNECTIONS, OPENED_TOTAL, OPEN_CONNECTIONS,
};

pub(crate) struct ActiveConn<C> {
    inner: C,
    state: ConnState,
    _permit: OwnedSemaphorePermit,
    _active_connections_gauge: GaugeGuard,
}

impl<C> ActiveConn<C> {
    pub(crate) fn new(inner: C, permit: OwnedSemaphorePermit, state: ConnState) -> ActiveConn<C> {
        Self {
            inner,
            state,
            _permit: permit,
            _active_connections_gauge: GaugeGuard::increment(ACTIVE_CONNECTIONS),
        }
    }

    pub(crate) fn into_idle(self) -> IdleConn<C> {
        IdleConn {
            inner: self.inner,
            state: self.state,
            _idle_connections_gauge: GaugeGuard::increment(IDLE_CONNECTIONS),
        }
    }

    pub(crate) fn is_brand_new(&self) -> bool {
        self.state.brand_new
    }

    pub(crate) fn set_brand_new(&mut self, brand_new: bool) {
        self.state.brand_new = brand_new;
    }

    pub(crate) fn into_raw(self) -> C {
        self.inner
    }

    pub(crate) fn as_raw_ref(&self) -> &C {
        &self.inner
    }

    pub(crate) fn as_raw_mut(&mut self) -> &mut C {
        &mut self.inner
    }
}

pub(crate) struct IdleConn<C> {
    inner: C,
    state: ConnState,
    _idle_connections_gauge: GaugeGuard,
}

impl<C> IdleConn<C> {
    pub(crate) fn is_brand_new(&self) -> bool {
        self.state.brand_new
    }

    pub(crate) fn into_active(self, permit: OwnedSemaphorePermit) -> ActiveConn<C> {
        ActiveConn::new(self.inner, permit, self.state)
    }

    pub(crate) fn created_at(&self) -> Instant {
        self.state.created_at
    }

    pub(crate) fn expired(&self, timeout: Option<Duration>) -> bool {
        timeout
            .and_then(|check_interval| {
                Instant::now()
                    .checked_duration_since(self.state.created_at)
                    .map(|dur_since| dur_since >= check_interval)
            })
            .unwrap_or(false)
    }

    pub(crate) fn idle_expired(&self, timeout: Option<Duration>) -> bool {
        timeout
            .and_then(|check_interval| {
                Instant::now()
                    .checked_duration_since(self.state.last_used_at)
                    .map(|dur_since| dur_since >= check_interval)
            })
            .unwrap_or(false)
    }

    pub(crate) fn needs_health_check(&self, timeout: Option<Duration>) -> bool {
        timeout
            .and_then(|check_interval| {
                Instant::now()
                    .checked_duration_since(self.state.last_checked_at)
                    .map(|dur_since| dur_since >= check_interval)
            })
            .unwrap_or(true)
    }

    pub(crate) fn mark_checked(&mut self) {
        self.state.last_checked_at = Instant::now()
    }

    pub(crate) fn split_raw(self) -> (C, ConnSplit<C>) {
        (
            self.inner,
            ConnSplit::new(self.state, self._idle_connections_gauge),
        )
    }
}

pub(crate) struct ConnState {
    pub(crate) created_at: Instant,
    pub(crate) last_used_at: Instant,
    pub(crate) last_checked_at: Instant,
    pub(crate) brand_new: bool,
    total_connections_open: Arc<AtomicU64>,
    total_connections_closed: Arc<AtomicU64>,
    _open_connections_gauge: GaugeGuard,
}

impl ConnState {
    pub(crate) fn new(
        total_connections_open: Arc<AtomicU64>,
        total_connections_closed: Arc<AtomicU64>,
    ) -> Self {
        counter!(OPENED_TOTAL).increment(1);
        Self {
            created_at: Instant::now(),
            last_used_at: Instant::now(),
            last_checked_at: Instant::now(),
            brand_new: true,
            total_connections_open,
            total_connections_closed,
            _open_connections_gauge: GaugeGuard::increment(OPEN_CONNECTIONS),
        }
    }
}

impl Drop for ConnState {
    fn drop(&mut self) {
        self.total_connections_open.fetch_sub(1, Ordering::Relaxed);
        self.total_connections_closed
            .fetch_add(1, Ordering::Relaxed);
        counter!(CLOSED_TOTAL).increment(1);
    }
}

pub(crate) struct ConnSplit<C> {
    state: ConnState,
    gauge: GaugeGuard,
    _phantom: PhantomData<C>,
}

impl<C> ConnSplit<C> {
    fn new(state: ConnState, gauge: GaugeGuard) -> Self {
        Self {
            state,
            gauge,
            _phantom: PhantomData,
        }
    }

    pub(crate) fn restore(self, raw: C) -> IdleConn<C> {
        IdleConn {
            inner: raw,
            state: self.state,
            _idle_connections_gauge: self.gauge,
        }
    }
}