From ff8b9b03848bf44303bb3dc7d04572823fef28f8 Mon Sep 17 00:00:00 2001 From: Ruben Fiszel Date: Mon, 13 Nov 2023 19:40:25 +0100 Subject: [PATCH] feat: cache postgres connection (#2621) * feat: cache pg connection * fix wmill dev * adjust timings * don't keep connection alive if not latest con * reduce sleep --- backend/windmill-worker/src/pg_executor.rs | 107 ++++++++++++++++-- frontend/src/lib/components/AppConnect.svelte | 4 +- 2 files changed, 102 insertions(+), 9 deletions(-) diff --git a/backend/windmill-worker/src/pg_executor.rs b/backend/windmill-worker/src/pg_executor.rs index 12b499355c..c8d5c2e79f 100644 --- a/backend/windmill-worker/src/pg_executor.rs +++ b/backend/windmill-worker/src/pg_executor.rs @@ -1,3 +1,7 @@ +use std::sync::atomic::{AtomicBool, AtomicU64}; +use std::sync::Arc; +use std::time::Duration; + use anyhow::Context; use chrono::Utc; use futures::TryStreamExt; @@ -8,6 +12,7 @@ use serde::Deserialize; use serde_json::value::RawValue; use serde_json::Map; use serde_json::{json, Value}; +use tokio::sync::Mutex; use tokio_postgres::types::IsNull; use tokio_postgres::{ types::{to_sql_checked, ToSql}, @@ -19,13 +24,14 @@ use tokio_postgres::{ }; use uuid::Uuid; use windmill_common::error::{self, Error}; -use windmill_common::worker::to_raw_value; +use windmill_common::worker::{to_raw_value, CLOUD_HOSTED}; use windmill_common::{error::to_anyhow, jobs::QueuedJob}; use windmill_parser_sql::parse_pgsql_sig; use crate::common::build_args_values; use crate::AuthedClientBackgroundTask; use bytes::BytesMut; +use lazy_static::lazy_static; use urlencoding::encode; #[derive(Deserialize)] @@ -39,6 +45,13 @@ struct PgDatabase { root_certificate_pem: Option, } +lazy_static! { + pub static ref CONNECTION_CACHE: Arc>> = + Arc::new(Mutex::new(None)); + pub static ref LAST_QUERY: AtomicU64 = AtomicU64::new(0); + pub static ref RUNNING: AtomicBool = AtomicBool::new(false); +} + pub async fn do_postgresql( job: &QueuedJob, client: &AuthedClientBackgroundTask, @@ -63,7 +76,28 @@ pub async fn do_postgresql( dbname = database.dbname, sslmode = sslmode ); - let (client, handle) = if sslmode == "require" { + let database_string_clone = database_string.clone(); + + RUNNING.store(true, std::sync::atomic::Ordering::Relaxed); + LAST_QUERY.store( + chrono::Utc::now().timestamp().try_into().unwrap_or(0), + std::sync::atomic::Ordering::Relaxed, + ); + let mtex; + if !*CLOUD_HOSTED { + mtex = Some(CONNECTION_CACHE.lock().await); + } else { + mtex = None; + } + + let has_cached_con = mtex + .as_ref() + .is_some_and(|x| x.as_ref().is_some_and(|y| y.0 == database_string)); + let new_client = if has_cached_con { + tracing::info!("Using cached connection"); + None + } else if sslmode == "require" { + tracing::info!("Creating new connection"); let mut connector = TlsConnector::builder(); if let Some(root_certificate_pem) = database.root_certificate_pem { if !root_certificate_pem.is_empty() { @@ -90,20 +124,25 @@ pub async fn do_postgresql( let handle = tokio::spawn(async move { if let Err(e) = connection.await { - eprintln!("connection error: {}", e); + let mut mtex = CONNECTION_CACHE.lock().await; + *mtex = None; + tracing::error!("connection error: {}", e); } }); - (client, handle) + Some((client, handle)) } else { + tracing::info!("Creating new connection"); let (client, connection) = tokio_postgres::connect(&database_string, NoTls) .await .map_err(to_anyhow)?; let handle = tokio::spawn(async move { if let Err(e) = connection.await { - eprintln!("connection error: {}", e); + let mut mtex = CONNECTION_CACHE.lock().await; + *mtex = None; + tracing::error!("connection error: {}", e); } }); - (client, handle) + Some((client, handle)) }; let mut statement_values: Vec = vec![]; @@ -133,6 +172,13 @@ pub async fn do_postgresql( convert_val(value, arg_t) }) .collect::>>()?; + + let (client, handle) = if let Some((client, handle)) = new_client.as_ref() { + (client, Some(handle)) + } else { + let (_, client) = mtex.as_ref().unwrap().as_ref().unwrap(); + (client, None) + }; // Now we can execute a simple statement that just returns its parameter. let rows = client .query_raw(query, query_params) @@ -146,8 +192,55 @@ pub async fn do_postgresql( .into_iter() .map(postgres_row_to_json_value) .collect::, _>>()?); + RUNNING.store(false, std::sync::atomic::Ordering::Relaxed); - handle.abort(); + if let Some(handle) = handle { + if let Some(mut mtex) = mtex { + let abort_handler = handle.abort_handle(); + + if let Some(new_client) = new_client { + *mtex = Some((database_string, new_client.0)); + } + drop(mtex); + LAST_QUERY.store( + chrono::Utc::now().timestamp().try_into().unwrap_or(0), + std::sync::atomic::Ordering::Relaxed, + ); + + tokio::spawn(async move { + loop { + tokio::time::sleep(Duration::from_secs(5)).await; + let last_query = LAST_QUERY.load(std::sync::atomic::Ordering::Relaxed); + let now = chrono::Utc::now().timestamp().try_into().unwrap_or(0); + + //we cache connection for 5 minutes at most + if last_query + 60 * 5 < now + && !RUNNING.load(std::sync::atomic::Ordering::Relaxed) + { + tracing::info!("Closing cache connection due to inactivity"); + break; + } + let mtex = CONNECTION_CACHE.lock().await; + if mtex.is_none() { + // connection is not in the mutex anymore + break; + } else if let Some(mtex) = mtex.as_ref() { + if mtex.0.as_str() != &database_string_clone { + // connection is not the latest one + break; + } + } + + tracing::debug!("Keeping cached connection alive due to activity") + } + let mut mtex = CONNECTION_CACHE.lock().await; + *mtex = None; + abort_handler.abort(); + }); + } else { + handle.abort(); + } + } // And then check that we got back the same string we sent over. return Ok(to_raw_value(&result)); } diff --git a/frontend/src/lib/components/AppConnect.svelte b/frontend/src/lib/components/AppConnect.svelte index dcc6d3e523..fd4f38cd4c 100644 --- a/frontend/src/lib/components/AppConnect.svelte +++ b/frontend/src/lib/components/AppConnect.svelte @@ -201,8 +201,8 @@ linkedSecret: undefined } ]) - const filteredNativeLanguages = filteredConnectsManual?.filter(([key, _]) => - nativeLanguagesCategory.includes(key) + const filteredNativeLanguages = filteredConnectsManual?.filter( + (o) => nativeLanguagesCategory?.includes(o[0]) ?? false ) filteredConnectsManual = [