fix: Ensure connection pool metrics stay consistent (#99)

* fix: Ensure connection pool metrics stay consistent

This started as a bugfix to several prisma issues regarding incorrect
metrics:

prisma/prisma#25177
prisma/prisma#23525

And a couple of more discovered the testing, but not reported in the
issues.
There were several causes for this:
1. Following pattern appears quite a lot in mobc code:

```rust
gauge!("something").increment(1.0);
do_a_thing_that_could_fail()?;
gauge!("something").decrement(1.0);
```

So, in case `do_a_thing_that_could_fail` actually fails, gauge will get
incremented, but never will get decremented.

2. Couple of metrics were relying on `Conn::close` being manually called
and that was not the case every once in a while.

To prevent both of those problems, I rewrote the internals of library to
rely on RAII rather than manual counters and resources management.

`Conn` struct is now split into two:
- `ActiveConn` - represents currently checked out connection that have
  been actively used by the client. Holds onto semaphore permit and can
  be converted into `IdleConn`. Doing so will free the permit.
- `IdleConn` - represents idle connection, currently checked into the
  pool. Can be converted to `ActiveConn` by providing a valid permit.

`ConnState` represents the shared state of the connection that is
retained between different activity states.

Both `IdleConn` and `ActiveConn` manage their corresponding gauges -
increment them on creation and decrement them during drop. `ConnState`
manages `CONNECTIONS_OPEN` gauge and `CONNECTIONS_TOTAL` and
`CLOSED_TOTAL` counters in the same way.

This system ensures that metrics stay consistent: since metrics are
automatically incremented and decremented on state conversions, we can
always be sure that:
- Connection is always either idle or active, there is no in between
state.
- Idle connections and active connections gauges will always add up to
the currently open connections gauge
- Total connections open counter, minus total connections closed gauge
will always be equal to number of currently open connections.

Since resources are now managed by `Drop::drop` implementations, that
removes the need for manual `close` method and simiplifies the code
quite in a few places, also ensuring  it is safer against future
changes.

* Update src/conn.rs

Co-authored-by: Alexey Orlenko <[email protected]>

---------

Co-authored-by: Alexey Orlenko <[email protected]>

Author: SevInf

Cargo.toml

@@ -27,7 +27,7 @@ async-trait = "0.1"
 futures-timer = "3.0.2"
 log = "0.4"
 thiserror = "1.0"
-metrics = "0.22.1"
+metrics = "0.23.0"
 tracing = { version = "0.1", features = ["attributes"] }
 tracing-subscriber = "0.3.11"
 

src/conn.rs

@@ -0,0 +1,183 @@
+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,
+        }
+    }
+}