feat: Add pool connection info to __lbheartbeat__ for ops

This will add several fields to `__lbheartbeat__` if
`database_pool_max_size` is specified. These include:

```json
{
  "active_connections": ... /* Number of active connections */,
  "idle_connections": ... /* number of idle connections */,
  "duration": ... /* how long no idle connections have been availble */,
}
```

Note that "duration" will only be present if `idle_connections` has been
zero since the last time a check was performed.

* this also adds `database_pool_max_size` as a config option.

Issue: #945

src/server/mod.rs

@@ -8,11 +8,12 @@ use actix_web::{
     HttpResponse, HttpServer,
 };
 use cadence::StatsdClient;
+use tokio::sync::RwLock;
 
 use crate::db::{pool_from_settings, spawn_pool_periodic_reporter, DbPool};
 use crate::error::ApiError;
 use crate::server::metrics::Metrics;
-use crate::settings::{Secrets, ServerLimits, Settings};
+use crate::settings::{Deadman, Secrets, ServerLimits, Settings};
 use crate::web::{handlers, middleware, tokenserver};
 
 pub const BSO_ID_REGEX: &str = r"[ -~]{1,64}";
@@ -51,6 +52,8 @@ pub struct ServerState {
     pub port: u16,
 
     pub quota_enabled: bool,
+
+    pub deadman: Arc<RwLock<Deadman>>,
 }
 
 pub fn cfg_path(path: &str) -> String {
@@ -144,14 +147,16 @@ macro_rules! build_app {
             // Remember to update .::web::middleware::DOCKER_FLOW_ENDPOINTS
             // when applying changes to endpoint names.
             .service(web::resource("/__heartbeat__").route(web::get().to(handlers::heartbeat)))
-            .service(
-                web::resource("/__lbheartbeat__").route(web::get().to(|_: HttpRequest| {
-                    // used by the load balancers, just return OK.
-                    HttpResponse::Ok()
-                        .content_type("application/json")
-                        .body("{}")
-                })),
-            )
+            .service(web::resource("/__lbheartbeat__").route(web::get().to(
+                handlers::lbheartbeat, /*
+                                           |_: HttpRequest| {
+                                           // used by the load balancers, just return OK.
+                                           HttpResponse::Ok()
+                                               .content_type("application/json")
+                                               .body("{}")
+                                       }
+                                       */
+            )))
             .service(
                 web::resource("/__version__").route(web::get().to(|_: HttpRequest| {
                     // return the contents of the version.json file created by circleci
@@ -182,6 +187,10 @@ impl Server {
         let fxa_metrics_hash_secret = Arc::new(settings.fxa_metrics_hash_secret.clone());
         let quota_enabled = settings.enable_quota;
         let actix_keep_alive = settings.actix_keep_alive;
+        let deadman = Arc::new(RwLock::new(Deadman {
+            max_size: settings.database_pool_max_size,
+            ..Default::default()
+        }));
 
         spawn_pool_periodic_reporter(Duration::from_secs(10), metrics.clone(), db_pool.clone())?;
 
@@ -199,6 +208,7 @@ impl Server {
                 metrics: Box::new(metrics.clone()),
                 port,
                 quota_enabled,
+                deadman: Arc::clone(&deadman),
             };
 
             build_app!(state, limits)

src/server/test.rs

@@ -76,6 +76,7 @@ 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())),
     }
 }
 

src/settings.rs

@@ -33,6 +33,13 @@ pub struct Quota {
     pub enforced: bool,
 }
 
+#[derive(Copy, Clone, Default, Debug)]
+pub struct Deadman {
+    pub max_size: Option<u32>,
+    pub previous_count: usize,
+    pub clock_start: Option<time::Instant>,
+}
+
 #[derive(Clone, Debug, Deserialize)]
 pub struct Settings {
     pub debug: bool,
@@ -114,6 +121,7 @@ impl Settings {
         s.set_default("database_pool_connection_timeout", Some(30))?;
         s.set_default("database_use_test_transactions", false)?;
         s.set_default("master_secret", "")?;
+        s.set_default::<Option<String>>("database_pool_max_size", Some("10".to_owned()))?;
         s.set_default::<Option<String>>("tokenserver_database_url", None)?;
         s.set_default::<Option<String>>("tokenserver_jwks_rsa_modulus", None)?;
         s.set_default::<Option<String>>("tokenserver_jwks_rsa_exponent", None)?;

src/web/extractors.rs

@@ -1673,13 +1673,14 @@ mod tests {
     use rand::{thread_rng, Rng};
     use serde_json::{self, json};
     use sha2::Sha256;
+    use tokio::sync::RwLock;
 
     use crate::db::{
         mock::{MockDb, MockDbPool},
         Db,
     };
     use crate::server::{metrics, ServerState};
-    use crate::settings::{Secrets, ServerLimits, Settings};
+    use crate::settings::{Deadman, Secrets, ServerLimits, Settings};
 
     use crate::web::auth::{hkdf_expand_32, HawkPayload};
 
@@ -1715,6 +1716,7 @@ mod tests {
             port: 8000,
             metrics: Box::new(metrics::metrics_from_opts(&settings).unwrap()),
             quota_enabled: settings.enable_quota,
+            deadman: Arc::new(RwLock::new(Deadman::default())),
         }
     }
 

src/web/handlers.rs

@@ -540,6 +540,53 @@ pub async fn heartbeat(hb: HeartbeatRequest) -> Result<HttpResponse, Error> {
     }
 }
 
+pub async fn lbheartbeat(req: HttpRequest) -> Result<HttpResponse, Error> {
+    let mut resp: HashMap<String, Value> = HashMap::new();
+
+    let state = match req.app_data::<Data<ServerState>>() {
+        Some(s) => s,
+        None => {
+            error!("⚠️ Could not load the app state");
+            return Ok(HttpResponse::InternalServerError().body(""));
+        }
+    };
+
+    let deadarc = state.deadman.clone();
+    let mut deadman = *deadarc.read().await;
+    let db_state = state.db_pool.clone().state();
+
+    if let Some(max_size) = deadman.max_size {
+        if db_state.connections >= max_size && db_state.idle_connections == 0 {
+            if deadman.previous_count > 0 {
+                deadman.clock_start = Some(time::Instant::now());
+            }
+        } else if deadman.clock_start.is_some() {
+                deadman.clock_start = None
+            }
+        deadman.previous_count = db_state.idle_connections as usize;
+        {
+            *deadarc.write().await = deadman;
+        }
+        resp.insert(
+            "active_connections".to_string(),
+            Value::from(db_state.connections),
+        );
+        resp.insert(
+            "idle_connections".to_string(),
+            Value::from(db_state.idle_connections),
+        );
+        if let Some(clock) = deadman.clock_start {
+            let duration: time::Duration = time::Instant::now() - clock;
+            resp.insert(
+                "duration_ms".to_string(),
+                Value::from(duration.whole_milliseconds()),
+            );
+        };
+    }
+
+    Ok(HttpResponse::Ok().json(json!(resp)))
+}
+
 // try returning an API error
 pub async fn test_error(
     _req: HttpRequest,