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

Code before:
    /// Time spent waiting for statements to complete, in milliseconds.
    pub db_statement_timeout_ms: Option<u64>,

    #[arg(long)]
    /// Enable server certificate verification. By default, this is set to false to match the
    /// default behavior of libpq.
    pub tls_verify_cert: bool,

    #[arg(long)]
    /// Path to a custom CA certificate to use for server certificate verification.
    pub tls_ca_cert_path: Option<PathBuf>,
}

#[derive(Clone)]
pub struct Db {
    pool: Pool<AsyncPgConnectionWithId>,
    pool_stats: Option<Arc<PoolStats>>,
}

// Ensures that unacquired_canceled is incremented if neither acquired nor unacquired_error are so
// that one of acquired, unacquired_error, unacquired_canceled is incremented for every requested
// connection. This is used to be able to calculate the number of pending connections:
// pending = requested - (acquired + unacquired_error + unacquired_canceled)
struct CancelGuard<'a> {
    stats: &'a PoolStats,
    disarmed: bool,
}

pub struct PoolStats {
    pub requested: IntCounter,
    pub acquired: IntCounter,
    pub unacquired_error: IntCounter,
    pub unacquired_canceled: IntCounter,
}

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

impl<'a> DbConfig<'a> {
    pub fn for_write(database_url: Url, db_args: DbArgs) -> Self {
        DbConfig {
            database_url,
            db_args,
            read_only: false,
            registry: None,
        }
    }

    pub fn for_read(database_url: Url, db_args: DbArgs) -> Self {
        DbConfig {
            database_url,

Review comment (refactor):
Here's a way to abstract this in Rust so the type system is doing more work for you to prove that it's impossible to break the invariants here (exactly one of `acquired`, `unacquired_error`, `unacquired_canceled` called by the time the guard is dropped).

```rust
struct CancelGuard<'m>(Option<&'m PoolMetrics>);

impl<'m> CancelGuard<'m> {
    fn request(metrics: &'m PoolMetrics) -> Self { 
        metrics.requested.inc();
        Self(Some(metrics))
    }

    fn acquired(mut self) {
        self.0.take().map(|m| m.acquired.inc());
    }

    fn errored(mut self) {
        self.0.take().map(|m| m.unacquired_error.inc());
    }
}

impl<'m> Drop for CancelGuard<'m> {
    fn drop(&mut self) {
        self.0.take().map(|m| m.unacquired_canceled.inc());
    }
}
```

Code after:
    /// Path to a custom CA certificate to use for server certificate verification.
    pub tls_ca_cert_path: Option<PathBuf>,
}

#[derive(Clone)]
pub struct Db {
    pool: Pool<AsyncPgConnectionWithId>,
    pool_metrics: Option<Arc<PoolMetrics>>,
}

// Ensures that unacquired_canceled is incremented if neither acquired nor unacquired_error are so
// that one of acquired, unacquired_error, unacquired_canceled is incremented for every requested
// connection. This is used to be able to calculate the number of pending connections:
// pending = requested - (acquired + unacquired_error + unacquired_canceled)
struct CancelGuard<'m>(Option<&'m PoolMetrics>);

/// 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::new(database_url, config, false).await
    }

    /// Construct a new DB connection pool talking to the database at `database_url` that defaults
    /// to read-only transactions. Instances of [Db] can be cloned to share access to the same
    /// 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(