store: Move ErrorHandler functionality to ConnectionManager

Author: lutter

store/postgres/src/pool/manager.rs

@@ -4,7 +4,7 @@
 //! tracking availability of the underlying database
 
 use deadpool::managed::{Hook, RecycleError, RecycleResult};
-use diesel::{r2d2, IntoSql};
+use diesel::IntoSql;
 
 use diesel_async::pooled_connection::{PoolError as DieselPoolError, PoolableConnection};
 use diesel_async::{AsyncConnection, RunQueryDsl};
@@ -18,7 +18,6 @@ use graph::slog::info;
 use graph::slog::Logger;
 
 use std::collections::HashMap;
-use std::fmt;
 use std::sync::atomic::AtomicBool;
 use std::sync::atomic::Ordering;
 use std::sync::Arc;
@@ -30,13 +29,63 @@ use crate::pool::AsyncPool;
 /// Our own connection manager. It is pretty much the same as
 /// `AsyncDieselConnectionManager` but makes it easier to instrument and
 /// track connection errors
+#[derive(Clone)]
 pub struct ConnectionManager {
+    logger: Logger,
     connection_url: String,
+    state_tracker: StateTracker,
+    error_counter: Counter,
 }
 
 impl ConnectionManager {
-    pub(super) fn new(connection_url: String) -> Self {
-        Self { connection_url }
+    pub(super) fn new(
+        logger: Logger,
+        connection_url: String,
+        state_tracker: StateTracker,
+        registry: &MetricsRegistry,
+        const_labels: HashMap<String, String>,
+    ) -> Self {
+        let error_counter = registry
+            .global_counter(
+                "store_connection_error_count",
+                "The number of Postgres connections errors",
+                const_labels,
+            )
+            .expect("failed to create `store_connection_error_count` counter");
+
+        Self {
+            logger,
+            connection_url,
+            state_tracker,
+            error_counter,
+        }
+    }
+
+    fn handle_error(&self, error: &dyn std::error::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.error_counter.inc();
+        if self.state_tracker.is_available() {
+            error!(self.logger, "Connection checkout"; "error" => msg);
+        }
+        self.state_tracker.mark_unavailable(Duration::from_secs(0));
     }
 }
 
@@ -46,9 +95,11 @@ impl deadpool::managed::Manager for ConnectionManager {
     type Error = DieselPoolError;
 
     async fn create(&self) -> Result<Self::Type, Self::Error> {
-        diesel_async::AsyncPgConnection::establish(&self.connection_url)
-            .await
-            .map_err(DieselPoolError::ConnectionError)
+        let res = diesel_async::AsyncPgConnection::establish(&self.connection_url).await;
+        if let Err(ref e) = res {
+            self.handle_error(e);
+        }
+        res.map_err(DieselPoolError::ConnectionError)
     }
 
     async fn recycle(
@@ -59,11 +110,14 @@ impl deadpool::managed::Manager for ConnectionManager {
         if std::thread::panicking() || obj.is_broken() {
             return Err(RecycleError::Message("Broken connection".into()));
         }
-        diesel::select(67_i32.into_sql::<diesel::sql_types::Integer>())
+        let res = diesel::select(67_i32.into_sql::<diesel::sql_types::Integer>())
             .execute(obj)
             .await
-            .map(|_| ())
-            .map_err(DieselPoolError::QueryError)?;
+            .map(|_| ());
+        if let Err(ref e) = res {
+            self.handle_error(e);
+        }
+        res.map_err(DieselPoolError::QueryError)?;
         Ok(())
     }
 }
@@ -139,57 +193,6 @@ impl StateTracker {
     }
 }
 
-#[derive(Clone)]
-pub(super) struct ErrorHandler {
-    logger: Logger,
-    counter: Counter,
-    state_tracker: StateTracker,
-}
-
-impl ErrorHandler {
-    pub(super) fn new(logger: Logger, counter: Counter, state_tracker: StateTracker) -> 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(Duration::from_secs(0));
-    }
-}
-
 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

store/postgres/src/pool/mod.rs

@@ -36,7 +36,7 @@ pub use diesel_async::scoped_futures::ScopedFutureExt;
 
 pub use coordinator::PoolCoordinator;
 pub use foreign_server::ForeignServer;
-use manager::{ErrorHandler, StateTracker};
+use manager::StateTracker;
 
 type AsyncPool = deadpool::managed::Pool<ConnectionManager>;
 /// A database connection for asynchronous diesel operations
@@ -429,29 +429,14 @@ impl PoolInner {
 
         let state_tracker = StateTracker::new(logger_pool.cheap_clone());
 
-        // Note: deadpool provides built-in metrics via pool.status()
-        // The r2d2-style ErrorHandler and EventHandler are not needed with deadpool.
-        // Metrics can be obtained from pool.status() and custom hooks can be added
-        // to the pool builder if needed.
-        let error_counter = registry
-            .global_counter(
-                "store_connection_error_count",
-                "The number of Postgres connections errors",
-                const_labels.clone(),
-            )
-            .expect("failed to create `store_connection_error_count` counter");
-        crit!(
-            logger_pool,
-            "Unfinished: not to replicate ErrorHandler with mobc"
-        );
-        let _error_handler = Box::new(ErrorHandler::new(
+        // Connect to Postgres
+        let conn_manager = ConnectionManager::new(
             logger_pool.clone(),
-            error_counter,
+            postgres_url.clone(),
             state_tracker.clone(),
-        ));
-
-        // Connect to Postgres
-        let conn_manager = ConnectionManager::new(postgres_url.clone());
+            &registry,
+            const_labels.clone(),
+        );
 
         // Note: deadpool does not support min_idle configuration
         if let Some(min_idle) = ENV_VARS.store.connection_min_idle {
@@ -471,7 +456,7 @@ impl PoolInner {
         // The post_create and post_recycle hooks are only called when
         // create and recycle succeed; we can therefore mark the pool
         // available
-        let pool = AsyncPool::builder(conn_manager)
+        let pool = AsyncPool::builder(conn_manager.clone())
             .max_size(pool_size as usize)
             .timeouts(timeouts)
             .runtime(Runtime::Tokio1)
@@ -485,7 +470,6 @@ impl PoolInner {
         let wait_meter = WaitMeter::new(&registry, const_labels.clone());
 
         let fdw_pool = fdw_pool_size.map(|pool_size| {
-            let conn_manager = ConnectionManager::new(postgres_url.clone());
             let fdw_timeouts = Timeouts {
                 wait: Some(ENV_VARS.store.connection_timeout),
                 create: None,