ref: Move check_url into Checker struct (#34)

getsentry · Jun 10, 2024 · 60f91f2 · 60f91f2
1 parent f0a812f
commit 60f91f2
Show file tree

Hide file tree

Showing 2 changed files with 106 additions and 84 deletions.
diff --git a/src/checker.rs b/src/checker.rs
@@ -1,20 +1,35 @@
-use std::error::Error;
-
-use reqwest::{Response, StatusCode};
+use reqwest::{Client, ClientBuilder, Response, StatusCode};
+use std::{error::Error, time::Duration};
 use tokio::time::Instant;
 use uuid::Uuid;
 
 use crate::types::{
     CheckResult, CheckStatus, CheckStatusReason, CheckStatusReasonType, RequestInfo, RequestType,
 };
 
+pub struct CheckerConfig {
+    /// How long will we wait before we consider the request to have timed out.
+    pub timeout: Duration,
+}
+
+impl Default for CheckerConfig {
+    fn default() -> CheckerConfig {
+        CheckerConfig {
+            timeout: Duration::from_secs(5),
+        }
+    }
+}
+
+/// Responsible for making HTTP requests to check if a domain is up.
+#[derive(Clone, Debug)]
+pub struct Checker {
+    client: Client,
+}
+
 /// Fetches the response from a URL.
 ///
 /// First attempts to fetch just the head, and if not supported falls back to fetching the entire body.
-async fn do_request(
-    client: &reqwest::Client,
-    url: &str,
-) -> (RequestType, Result<Response, reqwest::Error>) {
+async fn do_request(client: &Client, url: &str) -> (RequestType, Result<Response, reqwest::Error>) {
     let head_response = match client.head(url).send().await {
         Ok(response) => response,
         Err(e) => return (RequestType::Head, Err(e)),
@@ -44,93 +59,104 @@ fn dns_error(err: &reqwest::Error) -> Option<String> {
     None
 }
 
-/// Makes a request to a url to determine whether it is up.
-/// Up is defined as returning a 2xx within a specific timeframe.
-pub async fn check_url(client: &reqwest::Client, url: &str) -> CheckResult {
-    let trace_id = Uuid::new_v4();
-
-    let start = Instant::now();
-    let (request_type, response) = do_request(client, url).await;
-    let duration_ms = Some(start.elapsed().as_millis());
+impl Checker {
+    pub fn new(config: CheckerConfig) -> Self {
+        let client = ClientBuilder::new()
+            .timeout(config.timeout)
+            .build()
+            .expect("Failed to build checker client");
 
-    let status = if response.as_ref().is_ok_and(|r| r.status().is_success()) {
-        CheckStatus::Success
-    } else {
-        CheckStatus::Failure
-    };
+        Self { client }
+    }
 
-    let http_status_code = match &response {
-        Ok(r) => Some(r.status().as_u16()),
-        Err(e) => e.status().map(|s| s.as_u16()),
-    };
+    /// Makes a request to a url to determine whether it is up.
+    /// Up is defined as returning a 2xx within a specific timeframe.
+    pub async fn check_url(&self, url: &str) -> CheckResult {
+        let trace_id = Uuid::new_v4();
+
+        let start = Instant::now();
+        let (request_type, response) = do_request(&self.client, url).await;
+        let duration_ms = Some(start.elapsed().as_millis());
+
+        let status = if response.as_ref().is_ok_and(|r| r.status().is_success()) {
+            CheckStatus::Success
+        } else {
+            CheckStatus::Failure
+        };
+
+        let http_status_code = match &response {
+            Ok(r) => Some(r.status().as_u16()),
+            Err(e) => e.status().map(|s| s.as_u16()),
+        };
+
+        let request_info = Some(RequestInfo {
+            http_status_code,
+            request_type,
+        });
 
-    let request_info = Some(RequestInfo {
-        http_status_code,
-        request_type,
-    });
-
-    let status_reason = match response {
-        Ok(r) if r.status().is_success() => None,
-        Ok(r) => Some(CheckStatusReason {
-            status_type: CheckStatusReasonType::Failure,
-            description: format!("Got non 2xx status: {}", r.status()),
-        }),
-        Err(e) => Some({
-            if e.is_timeout() {
-                CheckStatusReason {
-                    status_type: CheckStatusReasonType::Timeout,
-                    description: format!("{:?}", e),
-                }
-            } else if let Some(message) = dns_error(&e) {
-                CheckStatusReason {
-                    status_type: CheckStatusReasonType::DnsError,
-                    description: message,
+        let status_reason = match response {
+            Ok(r) if r.status().is_success() => None,
+            Ok(r) => Some(CheckStatusReason {
+                status_type: CheckStatusReasonType::Failure,
+                description: format!("Got non 2xx status: {}", r.status()),
+            }),
+            Err(e) => Some({
+                if e.is_timeout() {
+                    CheckStatusReason {
+                        status_type: CheckStatusReasonType::Timeout,
+                        description: format!("{:?}", e),
+                    }
+                } else if let Some(message) = dns_error(&e) {
+                    CheckStatusReason {
+                        status_type: CheckStatusReasonType::DnsError,
+                        description: message,
+                    }
+                } else {
+                    CheckStatusReason {
+                        status_type: CheckStatusReasonType::Failure,
+                        description: format!("{:?}", e),
+                    }
                 }
-            } else {
-                CheckStatusReason {
-                    status_type: CheckStatusReasonType::Failure,
-                    description: format!("{:?}", e),
-                }
-            }
-        }),
-    };
-
-    CheckResult {
-        guid: Uuid::new_v4(),
-        monitor_id: 0,
-        monitor_environment_id: 0,
-        status,
-        status_reason,
-        trace_id,
-        scheduled_check_time: 0,
-        actual_check_time: 0,
-        duration_ms,
-        request_info,
+            }),
+        };
+
+        CheckResult {
+            guid: Uuid::new_v4(),
+            monitor_id: 0,
+            monitor_environment_id: 0,
+            status,
+            status_reason,
+            trace_id,
+            scheduled_check_time: 0,
+            actual_check_time: 0,
+            duration_ms,
+            request_info,
+        }
     }
 }
 
 #[cfg(test)]
 mod tests {
     use crate::types::{CheckStatus, CheckStatusReasonType, RequestType};
 
-    use super::check_url;
+    use super::{Checker, CheckerConfig};
     use httpmock::prelude::*;
     use httpmock::Method;
-    use reqwest::{ClientBuilder, StatusCode};
+    use reqwest::StatusCode;
     use std::time::Duration;
     // use crate::checker::FailureReason;
 
     #[tokio::test]
     async fn test_simple_head() {
         let server = MockServer::start();
-        let client = ClientBuilder::new().build().unwrap();
+        let checker = Checker::new(CheckerConfig::default());
 
         let head_mock = server.mock(|when, then| {
             when.method(Method::HEAD).path("/head");
             then.delay(Duration::from_millis(50)).status(200);
         });
 
-        let result = check_url(&client, &server.url("/head")).await;
+        let result = checker.check_url(&server.url("/head")).await;
 
         assert_eq!(result.status, CheckStatus::Success);
         assert!(result.status_reason.is_none());
@@ -150,7 +176,7 @@ mod tests {
     #[tokio::test]
     async fn test_simple_get() {
         let server = MockServer::start();
-        let client = ClientBuilder::new().build().unwrap();
+        let checker = Checker::new(CheckerConfig::default());
 
         let head_disallowed_mock = server.mock(|when, then| {
             when.method(Method::HEAD).path("/no-head");
@@ -161,7 +187,7 @@ mod tests {
             then.status(200);
         });
 
-        let result = check_url(&client, &server.url("/no-head")).await;
+        let result = checker.check_url(&server.url("/no-head")).await;
 
         assert_eq!(result.status, CheckStatus::Success);
         assert_eq!(
@@ -180,14 +206,14 @@ mod tests {
         let server = MockServer::start();
 
         let timeout = Duration::from_millis(TIMEOUT);
-        let client = ClientBuilder::new().timeout(timeout).build().unwrap();
+        let checker = Checker::new(CheckerConfig { timeout });
 
         let timeout_mock = server.mock(|when, then| {
             when.method(Method::HEAD).path("/timeout");
             then.delay(Duration::from_millis(TIMEOUT + 100)).status(200);
         });
 
-        let result = check_url(&client, &server.url("/timeout")).await;
+        let result = checker.check_url(&server.url("/timeout")).await;
 
         assert_eq!(result.status, CheckStatus::Failure);
         assert!(result.duration_ms.unwrap_or(0) >= TIMEOUT as u128);
@@ -203,14 +229,14 @@ mod tests {
     #[tokio::test]
     async fn test_simple_400() {
         let server = MockServer::start();
-        let client = ClientBuilder::new().build().unwrap();
+        let checker = Checker::new(CheckerConfig::default());
 
         let head_mock = server.mock(|when, then| {
             when.method(Method::HEAD).path("/head");
             then.status(400);
         });
 
-        let result = check_url(&client, &server.url("/head")).await;
+        let result = checker.check_url(&server.url("/head")).await;
 
         assert_eq!(result.status, CheckStatus::Failure);
         assert_eq!(

diff --git a/src/scheduler.rs b/src/scheduler.rs
@@ -1,26 +1,22 @@
-use std::time::Duration;
+use std::sync::Arc;
 
 use chrono::Utc;
-use reqwest::ClientBuilder;
 use tokio_cron_scheduler::{Job, JobScheduler, JobSchedulerError};
 
-use crate::checker::check_url;
+use crate::checker::{Checker, CheckerConfig};
 
 pub async fn run_scheduler() -> Result<(), JobSchedulerError> {
     let scheduler = JobScheduler::new().await?;
 
-    let client = ClientBuilder::new()
-        .timeout(Duration::from_secs(5))
-        .build()
-        .expect("Failed to build checker client");
+    let checker = Arc::new(Checker::new(CheckerConfig::default()));
 
     let checker_job = Job::new_async("0 */5 * * * *", move |_uuid, mut _l| {
-        let job_client = client.clone();
+        let job_checker = checker.clone();
 
         Box::pin(async move {
             println!("Executing job at {:?}", Utc::now());
 
-            let check_result = check_url(&job_client, "https://sentry.io").await;
+            let check_result = job_checker.check_url("https://sentry.io").await;
 
             println!("checked sentry.io, got {:?}", check_result)
         })