diff --git a/src/server/test.rs b/src/server/test.rs index cc7c04b8cc..832b820dd9 100644 --- a/src/server/test.rs +++ b/src/server/test.rs @@ -76,7 +76,10 @@ async fn get_test_state(settings: &Settings) -> ServerState { metrics: Box::new(metrics), port: settings.port, quota_enabled: settings.enable_quota, - deadman: Arc::new(RwLock::new(Deadman::default())), + deadman: Arc::new(RwLock::new(Deadman { + max_size: settings.database_pool_max_size, + ..Default::default() + })), } } @@ -695,3 +698,60 @@ async fn overquota() { let resp = app.call(req).await.unwrap(); assert!(resp.response().status().is_success()); } + +#[actix_rt::test] +async fn lbheartbeat_check() { + use actix_web::web::Buf; + + let mut settings = get_test_settings(); + settings.database_pool_max_size = Some(10); + + let mut app = init_app!(settings).await; + + // Test all is well. + let lb_req = create_request(http::Method::GET, "/__lbheartbeat__", None, None).to_request(); + let sresp = app.call(lb_req).await.unwrap(); + let status = sresp.status(); + // dbg!(status, test::read_body(sresp).await); + assert!(status.is_success()); + + // Exhaust the connections. + let mut headers: HashMap<&str, String> = HashMap::new(); + headers.insert("TEST_CONNECTIONS", "10".to_owned()); + headers.insert("TEST_IDLES", "0".to_owned()); + let req = create_request( + http::Method::GET, + "/__lbheartbeat__", + Some(headers.clone()), + None, + ) + .to_request(); + let sresp = app.call(req).await.unwrap(); + let status = sresp.status(); + // dbg!(status, test::read_body(sresp).await); + assert!(status == StatusCode::INTERNAL_SERVER_ERROR); + + // check duration for exhausted connections + std::thread::sleep(std::time::Duration::from_secs(1)); + let req = + create_request(http::Method::GET, "/__lbheartbeat__", Some(headers), None).to_request(); + let sresp = app.call(req).await.unwrap(); + let status = sresp.status(); + let body = test::read_body(sresp).await; + let resp: HashMap = + serde_json::de::from_str(std::str::from_utf8(body.bytes()).unwrap()).unwrap(); + // dbg!(status, body, &resp); + assert!(status == StatusCode::INTERNAL_SERVER_ERROR); + assert!(resp.get("duration_ms").unwrap().as_u64().unwrap() > 1000); + + // check recovery + let mut headers: HashMap<&str, String> = HashMap::new(); + headers.insert("TEST_CONNECTIONS", "5".to_owned()); + headers.insert("TEST_IDLES", "5".to_owned()); + let req = + create_request(http::Method::GET, "/__lbheartbeat__", Some(headers), None).to_request(); + let sresp = app.call(req).await.unwrap(); + let status = sresp.status(); + // dbg!(status, test::read_body(sresp).await); + assert!(status == StatusCode::OK); +} diff --git a/src/web/handlers.rs b/src/web/handlers.rs index e73a8ba48d..345da2ae96 100644 --- a/src/web/handlers.rs +++ b/src/web/handlers.rs @@ -1,7 +1,9 @@ //! API Handlers use std::collections::HashMap; -use actix_web::{http::StatusCode, web::Data, Error, HttpRequest, HttpResponse}; +use actix_web::{ + dev::HttpResponseBuilder, http::StatusCode, web::Data, Error, HttpRequest, HttpResponse, +}; use serde::Serialize; use serde_json::{json, Value}; @@ -553,15 +555,44 @@ pub async fn lbheartbeat(req: HttpRequest) -> Result { let deadarc = state.deadman.clone(); let mut deadman = *deadarc.read().await; - let db_state = state.db_pool.clone().state(); + let db_state = if cfg!(test) { + use crate::db::results::PoolState; + use actix_web::http::header::HeaderValue; + use std::str::FromStr; + + let test_pool = PoolState { + connections: u32::from_str( + req.headers() + .get("TEST_CONNECTIONS") + .unwrap_or(&HeaderValue::from_static("0")) + .to_str() + .unwrap_or("0"), + ) + .unwrap_or_default(), + idle_connections: u32::from_str( + req.headers() + .get("TEST_IDLES") + .unwrap_or(&HeaderValue::from_static("0")) + .to_str() + .unwrap_or("0"), + ) + .unwrap_or_default(), + }; + // dbg!(&test_pool, deadman.max_size); + test_pool + } else { + state.db_pool.clone().state() + }; let active = db_state.connections - db_state.idle_connections; + let mut status_code = StatusCode::OK; if let Some(max_size) = deadman.max_size { if active >= max_size && db_state.idle_connections == 0 { - if deadman.previous_count > 0 { + if deadman.clock_start.is_none() { deadman.clock_start = Some(time::Instant::now()); } + status_code = StatusCode::INTERNAL_SERVER_ERROR; } else if deadman.clock_start.is_some() { deadman.clock_start = None } @@ -583,7 +614,7 @@ pub async fn lbheartbeat(req: HttpRequest) -> Result { }; } - Ok(HttpResponse::Ok().json(json!(resp))) + Ok(HttpResponseBuilder::new(status_code).json(json!(resp))) } // try returning an API error