healthchecks!

src/errors.rs

@@ -5,4 +5,5 @@ pub enum Error {
     ClientBadStartup,
     ProtocolSyncError,
     ServerError,
+    ServerTimeout,
 }

src/pool.rs

@@ -43,8 +43,20 @@ impl ManageConnection for ServerPool {
     }
 
     /// Determines if the connection is still connected to the database.
-    async fn is_valid(&self, _conn: &mut PooledConnection<'_, Self>) -> Result<(), Self::Error> {
+    async fn is_valid(&self, conn: &mut PooledConnection<'_, Self>) -> Result<(), Self::Error> {
-        Ok(())
+        let server = &mut *conn;
+
+        // If this fails, the connection will be closed and another will be grabbed from the pool quietly :-).
+        // Failover, step 1, complete.
+        match tokio::time::timeout(
+            tokio::time::Duration::from_millis(1000),
+            server.query("SELECT 1"),
+        )
+        .await
+        {
+            Ok(_) => Ok(()),
+            Err(_err) => Err(Error::ServerTimeout),
+        }
     }
 
     /// Synchronously determine if the connection is no longer usable, if possible.

src/server.rs

@@ -209,8 +209,8 @@ impl Server {
         self.bad = true;
     }
 
-    pub async fn set_name(&mut self, name: &str) -> Result<(), Error> {
+    pub async fn query(&mut self, query: &str) -> Result<(), Error> {
-        let mut query = BytesMut::from(&format!("SET application_name = {}", name).as_bytes()[..]);
+        let mut query = BytesMut::from(&query.as_bytes()[..]);
         query.put_u8(0);
 
         let len = query.len() as i32 + 4;
@@ -226,4 +226,10 @@ impl Server {
 
         Ok(())
     }
+
+    pub async fn set_name(&mut self, name: &str) -> Result<(), Error> {
+        Ok(self
+            .query(&format!("SET application_name = '{}'", name))
+            .await?)
+    }
 }