Repository: MystenLabs/sui
File: crates/sui-pg-db/src/lib.rs (Rust)

Code before:
    fn poll(self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<Self::Output> {
        let completed = self.completed.clone();
        let pool_stats = self.pool_stats.clone();
        match self.project().inner.poll(cx) {
            Poll::Ready(Ok(result)) => {
                inc_if_false(&completed, &pool_stats.acquired);
                Poll::Ready(Ok(result))
            }
            Poll::Ready(Err(e)) => {
                inc_if_false(&completed, &pool_stats.unacquired_error);
                Poll::Ready(Err(e))
            }
            Poll::Pending => Poll::Pending,
        }
    }
}

#[pinned_drop]
impl<F> PinnedDrop for PoolFuture<F> {
    fn drop(self: Pin<&mut Self>) {
        inc_if_false(&self.completed, &self.pool_stats.unacquired_canceled);
    }
}

// Ensure the counter is only incremented once per bool instance.
fn inc_if_false(bool: &AtomicBool, counter: &AtomicU64) {
    if !bool.swap(true, Ordering::Relaxed) {
        counter.fetch_add(1, Ordering::Relaxed);
    }
}

/// Wrapper struct over the remote `PooledConnection` type for dealing with the `Store` trait.
pub struct Connection<'a>(PooledConnection<'a, AsyncPgConnectionWithId>);

impl DbArgs {
    pub fn connection_timeout(&self) -> Duration {
        Duration::from_millis(self.db_connection_timeout_ms)
    }

    pub fn statement_timeout(&self) -> Option<Duration> {
        self.db_statement_timeout_ms.map(Duration::from_millis)
    }
}

impl Db {
    /// Construct a new DB connection pool talking to the database at `database_url` that supports
    /// write and reads. Instances of [Db] can be cloned to share access to the same pool.
    pub async fn for_write(database_url: Url, config: DbArgs) -> anyhow::Result<Self> {
        Self::from_pool(database_url, config, false).await
    }


Review comment (suggestion):
it seems this largely makes the second increment in `fn connect()` a no-op? `PoolFuture::poll()` increments stats

Code after:
    /// pool.
    pub async fn for_read(database_url: Url, config: DbArgs) -> anyhow::Result<Self> {
        Self::new(database_url, config, true).await
    }

    async fn new(database_url: Url, db_args: DbArgs, read_only: bool) -> anyhow::Result<Self> {
        Ok(Db {
            pool: pool(database_url, db_args, read_only).await?,
            pool_metrics: None,
        })
    }

    pub fn register_metrics(
        mut self,
        prefix: Option<&str>,
        registry: &Registry,
    ) -> anyhow::Result<Self> {
        let pool_metrics = PoolMetrics::new(prefix, registry)?;
        self.pool_metrics = Some(pool_metrics);
        Ok(self)
    }

    /// Retrieves a connection from the pool. Can fail with a timeout if a connection cannot be
    /// established before the [DbArgs::connection_timeout] has elapsed.
    pub async fn connect(&self) -> anyhow::Result<Connection<'_>> {
        if let Some(pool_metrics) = &self.pool_metrics {
            let guard = CancelGuard::request(pool_metrics);
            match self.pool.get().await {
                Ok(c) => {
                    guard.acquired();
                    Ok(Connection(c))
                }
                Err(e) => {
                    guard.unacquired_error();
                    Err(e.into())
                }
            }
        } else {
            Ok(Connection(self.pool.get().await?))
        }
    }

    /// Statistics about the connection pool
    pub fn state(&self) -> bb8::State {
        self.pool.state()
    }

    pub fn pool_metrics(&self) -> Option<&PoolMetrics> {
        self.pool_metrics.as_ref().map(|p| p.as_ref())
    }