store: Move PoolStateTracker and handlers to own module

Also, rename it to just StateTracker

Author: lutter

store/postgres/src/pool/mod.rs

@@ -1,7 +1,7 @@
 use diesel::r2d2::Builder;
 use diesel::{connection::SimpleConnection, pg::PgConnection};
 use diesel::{
-    r2d2::{self, event as e, ConnectionManager, HandleEvent, Pool, PooledConnection},
+    r2d2::{ConnectionManager, Pool, PooledConnection},
     Connection,
 };
 use diesel::{sql_query, RunQueryDsl};
@@ -14,15 +14,14 @@ use graph::internal_error;
 use graph::prelude::tokio::time::Instant;
 use graph::prelude::{
     anyhow::anyhow, crit, debug, error, info, o, tokio::sync::Semaphore, CancelGuard, CancelHandle,
-    CancelToken as _, CancelableError, Counter, Gauge, Logger, MovingStats, PoolWaitStats,
-    StoreError, ENV_VARS,
+    CancelToken as _, CancelableError, Gauge, Logger, MovingStats, PoolWaitStats, StoreError,
+    ENV_VARS,
 };
 use graph::prelude::{tokio, MetricsRegistry};
 use graph::slog::warn;
 use graph::util::timed_rw_lock::TimedMutex;
 
 use std::fmt::{self};
-use std::sync::atomic::{AtomicBool, Ordering};
 use std::sync::Arc;
 use std::time::Duration;
 use std::{collections::HashMap, sync::RwLock};
@@ -33,9 +32,11 @@ use crate::{Shard, PRIMARY_SHARD};
 
 mod coordinator;
 mod foreign_server;
+mod state_tracker;
 
 pub use coordinator::PoolCoordinator;
 pub use foreign_server::ForeignServer;
+use state_tracker::{ErrorHandler, EventHandler, StateTracker};
 
 /// The namespace under which the `PRIMARY_TABLES` are mapped into each
 /// shard
@@ -204,7 +205,7 @@ impl PoolState {
 pub struct ConnectionPool {
     inner: PoolState,
     pub shard: Shard,
-    state_tracker: PoolStateTracker,
+    state_tracker: StateTracker,
 }
 
 impl fmt::Debug for ConnectionPool {
@@ -240,47 +241,6 @@ impl PoolRole {
     }
 }
 
-#[derive(Clone)]
-struct PoolStateTracker {
-    available: Arc<AtomicBool>,
-    ignore_timeout: Arc<AtomicBool>,
-}
-
-impl PoolStateTracker {
-    fn new() -> Self {
-        Self {
-            available: Arc::new(AtomicBool::new(true)),
-            ignore_timeout: Arc::new(AtomicBool::new(false)),
-        }
-    }
-
-    fn mark_available(&self) {
-        self.available.store(true, Ordering::Relaxed);
-    }
-
-    fn mark_unavailable(&self) {
-        self.available.store(false, Ordering::Relaxed);
-    }
-
-    fn is_available(&self) -> bool {
-        self.available.load(Ordering::Relaxed)
-    }
-
-    fn timeout_is_ignored(&self) -> bool {
-        self.ignore_timeout.load(Ordering::Relaxed)
-    }
-
-    fn ignore_timeout<F, R>(&self, f: F) -> R
-    where
-        F: FnOnce() -> R,
-    {
-        self.ignore_timeout.store(true, Ordering::Relaxed);
-        let res = f();
-        self.ignore_timeout.store(false, Ordering::Relaxed);
-        res
-    }
-}
-
 impl ConnectionPool {
     fn create(
         shard_name: &str,
@@ -292,7 +252,7 @@ impl ConnectionPool {
         registry: Arc<MetricsRegistry>,
         coord: Arc<PoolCoordinator>,
     ) -> ConnectionPool {
-        let state_tracker = PoolStateTracker::new();
+        let state_tracker = StateTracker::new();
         let shard =
             Shard::new(shard_name.to_string()).expect("shard_name is a valid name for a shard");
         let inner = {
@@ -461,169 +421,6 @@ impl ConnectionPool {
     }
 }
 
-fn brief_error_msg(error: &dyn std::error::Error) -> String {
-    // For 'Connection refused' errors, Postgres includes the IP and
-    // port number in the error message. We want to suppress that and
-    // only use the first line from the error message. For more detailed
-    // analysis, 'Connection refused' manifests as a
-    // `ConnectionError(BadConnection("could not connect to server:
-    // Connection refused.."))`
-    error
-        .to_string()
-        .split('\n')
-        .next()
-        .unwrap_or("no error details provided")
-        .to_string()
-}
-
-#[derive(Clone)]
-struct ErrorHandler {
-    logger: Logger,
-    counter: Counter,
-    state_tracker: PoolStateTracker,
-}
-
-impl ErrorHandler {
-    fn new(logger: Logger, counter: Counter, state_tracker: PoolStateTracker) -> Self {
-        Self {
-            logger,
-            counter,
-            state_tracker,
-        }
-    }
-}
-impl std::fmt::Debug for ErrorHandler {
-    fn fmt(&self, _f: &mut fmt::Formatter) -> fmt::Result {
-        fmt::Result::Ok(())
-    }
-}
-
-impl r2d2::HandleError<r2d2::Error> for ErrorHandler {
-    fn handle_error(&self, error: r2d2::Error) {
-        let msg = brief_error_msg(&error);
-
-        // Don't count canceling statements for timeouts etc. as a
-        // connection error. Unfortunately, we only have the textual error
-        // and need to infer whether the error indicates that the database
-        // is down or if something else happened. When querying a replica,
-        // these messages indicate that a query was canceled because it
-        // conflicted with replication, but does not indicate that there is
-        // a problem with the database itself.
-        //
-        // This check will break if users run Postgres (or even graph-node)
-        // in a locale other than English. In that case, their database will
-        // be marked as unavailable even though it is perfectly fine.
-        if msg.contains("canceling statement")
-            || msg.contains("terminating connection due to conflict with recovery")
-        {
-            return;
-        }
-
-        self.counter.inc();
-        if self.state_tracker.is_available() {
-            error!(self.logger, "Postgres connection error"; "error" => msg);
-        }
-        self.state_tracker.mark_unavailable();
-    }
-}
-
-#[derive(Clone)]
-struct EventHandler {
-    logger: Logger,
-    count_gauge: Gauge,
-    wait_gauge: Gauge,
-    size_gauge: Gauge,
-    wait_stats: PoolWaitStats,
-    state_tracker: PoolStateTracker,
-}
-
-impl EventHandler {
-    fn new(
-        logger: Logger,
-        registry: Arc<MetricsRegistry>,
-        wait_stats: PoolWaitStats,
-        const_labels: HashMap<String, String>,
-        state_tracker: PoolStateTracker,
-    ) -> Self {
-        let count_gauge = registry
-            .global_gauge(
-                "store_connection_checkout_count",
-                "The number of Postgres connections currently checked out",
-                const_labels.clone(),
-            )
-            .expect("failed to create `store_connection_checkout_count` counter");
-        let wait_gauge = registry
-            .global_gauge(
-                "store_connection_wait_time_ms",
-                "Average connection wait time",
-                const_labels.clone(),
-            )
-            .expect("failed to create `store_connection_wait_time_ms` counter");
-        let size_gauge = registry
-            .global_gauge(
-                "store_connection_pool_size_count",
-                "Overall size of the connection pool",
-                const_labels,
-            )
-            .expect("failed to create `store_connection_pool_size_count` counter");
-        EventHandler {
-            logger,
-            count_gauge,
-            wait_gauge,
-            wait_stats,
-            size_gauge,
-            state_tracker,
-        }
-    }
-
-    fn add_conn_wait_time(&self, duration: Duration) {
-        self.wait_stats
-            .write()
-            .unwrap()
-            .add_and_register(duration, &self.wait_gauge);
-    }
-}
-
-impl std::fmt::Debug for EventHandler {
-    fn fmt(&self, _f: &mut fmt::Formatter) -> fmt::Result {
-        fmt::Result::Ok(())
-    }
-}
-
-impl HandleEvent for EventHandler {
-    fn handle_acquire(&self, _: e::AcquireEvent) {
-        self.size_gauge.inc();
-        self.state_tracker.mark_available();
-    }
-
-    fn handle_release(&self, _: e::ReleaseEvent) {
-        self.size_gauge.dec();
-    }
-
-    fn handle_checkout(&self, event: e::CheckoutEvent) {
-        self.count_gauge.inc();
-        self.add_conn_wait_time(event.duration());
-        self.state_tracker.mark_available();
-    }
-
-    fn handle_timeout(&self, event: e::TimeoutEvent) {
-        if self.state_tracker.timeout_is_ignored() {
-            return;
-        }
-        self.add_conn_wait_time(event.timeout());
-        if self.state_tracker.is_available() {
-            error!(self.logger, "Connection checkout timed out";
-               "wait_ms" => event.timeout().as_millis()
-            )
-        }
-        self.state_tracker.mark_unavailable();
-    }
-
-    fn handle_checkin(&self, _: e::CheckinEvent) {
-        self.count_gauge.dec();
-    }
-}
-
 #[derive(Clone)]
 pub struct PoolInner {
     logger: Logger,
@@ -662,7 +459,7 @@ impl PoolInner {
         fdw_pool_size: Option<u32>,
         logger: &Logger,
         registry: Arc<MetricsRegistry>,
-        state_tracker: PoolStateTracker,
+        state_tracker: StateTracker,
     ) -> PoolInner {
         check_mirrored_tables();