From 4be3d85c162b300e00f005433ab6a40fcc36419f Mon Sep 17 00:00:00 2001
From: Marcelo Altmann <marcelo@readyset.io>
Date: Mon, 4 Nov 2024 17:35:20 -0300
Subject: [PATCH] Add connection / query timeout

Add a timeout to the connection phase and query phase. In case a
connection get stuck, this will move the server to shunned state.

Fixes: #24
---
 src/hosts.rs | 29 ++++++++++++++++++++++-------
 1 file changed, 22 insertions(+), 7 deletions(-)

diff --git a/src/hosts.rs b/src/hosts.rs
index 1110fae..e616f48 100644
--- a/src/hosts.rs
+++ b/src/hosts.rs
@@ -1,5 +1,6 @@
 use crate::{config::Config, queries::Query};
 use core::fmt;
+use std::time::Duration;
 use mysql::{prelude::Queryable, Conn, OptsBuilder};
 
 #[allow(dead_code)]
@@ -68,7 +69,10 @@ impl Host {
                 .tcp_port(port)
                 .user(Some(config.readyset_user.clone()))
                 .pass(Some(config.readyset_password.clone()))
-                .prefer_socket(false),
+                .prefer_socket(false)
+                .read_timeout(Some(Duration::from_secs(5)))
+                .write_timeout(Some(Duration::from_secs(5)))
+                .tcp_connect_timeout(Some(Duration::from_secs(5))),
         ) {
             Ok(conn) => conn,
             Err(err) => {
@@ -144,14 +148,25 @@ impl Host {
     pub fn check_readyset_is_ready(&mut self) -> Result<bool, mysql::Error> {
         match &mut self.conn {
             Some(conn) => {
-                let rows: Vec<(String, String)> =
-                    conn.query("SHOW READYSET STATUS").unwrap_or(vec![]);
-                for (field, value) in rows {
-                    if field == "Snapshot Status" {
-                        return Ok(value == "Completed");
+                let result  =
+                    conn.query("SHOW READYSET STATUS");
+                match result {
+                    Ok(rows) => {
+                        let rows: Vec<(String, String)> = rows;
+                        for (field, value) in rows {
+                            if field == "Snapshot Status" {
+                                return Ok(value == "Completed");
+                            }
+                        }
+                        Ok(false)
+                    },
+                    Err(err) => {
+                        Err(mysql::Error::IoError(std::io::Error::new(
+                            std::io::ErrorKind::Other,
+                            format!("Failed to execute query: {}", err),
+                        )))
                     }
                 }
-                Ok(false)
             }
             None => Err(mysql::Error::IoError(std::io::Error::new(
                 std::io::ErrorKind::Other,