Do some connection pooling.

This prevents exhaustion of connections and also
has some performance benefits.

src/lib.rs

@@ -8,13 +8,14 @@ use std::path::{Path, PathBuf};
 use std::process::Command;
 use std::sync::Arc;
 use std::sync::Mutex;
+use std::sync::OnceLock;
 use std::sync::atomic::{AtomicU64, Ordering};
 
 use aws_config::BehaviorVersion;
 use aws_sdk_s3::Client;
 use aws_sdk_s3::config::{Credentials, Region};
 use aws_sdk_s3::primitives::ByteStream;
-use diesel::Connection as DieselConnection;
+use diesel::r2d2::{ConnectionManager, Pool, PooledConnection};
 use diesel::PgConnection;
 use futures_util::StreamExt;
 use lapin::options::{
@@ -49,6 +50,8 @@ type ProcessFn = dyn Fn(u64, Arc<ReadFileFn>, Arc<WriteFileFn>) -> Result<Vec<(u
 	+ 'static;
 
 static REQUEST_FILE_CONTEXT_COUNTER: AtomicU64 = AtomicU64::new(1);
+static DATABASE_URL_CACHE: OnceLock<String> = OnceLock::new();
+static DATABASE_POOL_CACHE: OnceLock<Pool<ConnectionManager<PgConnection>>> = OnceLock::new();
 
 fn init_tracing() {
 	let filter = EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("info"));
@@ -128,7 +131,7 @@ impl Microservice {
 		| -> Result<Vec<(u64, CaseKey)>, AnyError> {
 			let mut connection = establish_pg_connection().map_err(|e| {
 				format!(
-					"failed to establish PostgreSQL connection for request {}: {}",
+					"failed to checkout PostgreSQL connection from pool for request {}: {}",
 					request, e
 				)
 			})?;
@@ -601,10 +604,38 @@ fn resolve_password_from_pass(pass_key: &str) -> Result<String, AnyError> {
 	Ok(password)
 }
 
-fn establish_pg_connection() -> Result<PgConnection, AnyError> {
-	let database_url = fetch_database_url_from_sys_map()?;
-	PgConnection::establish(&database_url)
-		.map_err(|e| format!("failed to connect to PostgreSQL using sys-map DB config: {}", e).into())
+fn establish_pg_connection() -> Result<PooledConnection<ConnectionManager<PgConnection>>, AnyError> {
+	let pool = cached_pg_pool()?;
+	pool
+		.get()
+		.map_err(|e| format!("failed to get PostgreSQL connection from pool: {}", e).into())
+}
+
+fn cached_database_url() -> Result<String, AnyError> {
+	if let Some(url) = DATABASE_URL_CACHE.get() {
+		return Ok(url.clone());
+	}
+
+	let url = fetch_database_url_from_sys_map()?;
+	let _ = DATABASE_URL_CACHE.set(url.clone());
+	Ok(url)
+}
+
+fn cached_pg_pool() -> Result<&'static Pool<ConnectionManager<PgConnection>>, AnyError> {
+	if let Some(pool) = DATABASE_POOL_CACHE.get() {
+		return Ok(pool);
+	}
+
+	let database_url = cached_database_url()?;
+	let manager = ConnectionManager::<PgConnection>::new(database_url);
+	let pool = Pool::builder()
+		.build(manager)
+		.map_err(|e| format!("failed to build PostgreSQL connection pool: {}", e))?;
+
+	let _ = DATABASE_POOL_CACHE.set(pool);
+	DATABASE_POOL_CACHE
+		.get()
+		.ok_or_else(|| "failed to initialize PostgreSQL connection pool".into())
 }
 
 fn single_value<T: Clone>(values: &[T], field_name: &str) -> Result<T, AnyError> {
@@ -854,25 +885,18 @@ impl PyWriteFileFn {
 #[cfg(feature = "python")]
 fn run_python_process(
 	process: &Py<PyAny>,
+	engine: &Py<PyAny>,
 	request: u64,
 	read_file: Arc<ReadFileFn>,
 	write_file: Arc<WriteFileFn>,
 ) -> Result<Vec<(u64, CaseKey)>, AnyError> {
-	let database_url = fetch_database_url_from_sys_map()?;
-
 	Python::with_gil(|py| -> Result<Vec<(u64, CaseKey)>, AnyError> {
 		let py_read = Py::new(py, PyReadFileFn { inner: read_file })
 			.map_err(|e| format!("failed to build Python ReadFileFn wrapper: {}", e))?;
 		let py_write = Py::new(py, PyWriteFileFn { inner: write_file })
 			.map_err(|e| format!("failed to build Python WriteFileFn wrapper: {}", e))?;
-		let sqlalchemy = py
-			.import("sqlalchemy")
-			.map_err(|e| format!("failed to import sqlalchemy: {}", e))?;
-		let engine = sqlalchemy
-			.getattr("create_engine")
-			.and_then(|f| f.call1((database_url.as_str(),)))
-			.map_err(|e| format!("failed to create SQLAlchemy engine: {}", e))?;
 		let connection = engine
+			.bind(py)
 			.call_method0("connect")
 			.map_err(|e| format!("failed to open SQLAlchemy connection: {}", e))?;
 
@@ -909,6 +933,23 @@ fn run_python_process(
 	})
 }
 
+#[cfg(feature = "python")]
+fn create_python_sqlalchemy_engine() -> Result<Py<PyAny>, AnyError> {
+	let database_url = cached_database_url()?;
+
+	Python::with_gil(|py| {
+		let sqlalchemy = py
+			.import("sqlalchemy")
+			.map_err(|e| format!("failed to import sqlalchemy: {}", e))?;
+		let engine = sqlalchemy
+			.getattr("create_engine")
+			.and_then(|f| f.call1((database_url.as_str(),)))
+			.map_err(|e| format!("failed to create SQLAlchemy engine: {}", e))?;
+
+		Ok(engine.unbind())
+	})
+}
+
 #[cfg(feature = "python")]
 #[pyclass(name = "Microservice")]
 struct PyMicroservice {
@@ -937,10 +978,13 @@ impl PyMicroservice {
 
 	fn start(&self) -> PyResult<()> {
 		let process = Python::with_gil(|py| self.process.clone_ref(py));
+		let engine = create_python_sqlalchemy_engine().map_err(any_error_to_py)?;
 		let microservice = Microservice::new_case_key(
 			self.name.clone(),
 			self.config_host.clone(),
-			Arc::new(move |request, read_file, write_file| run_python_process(&process, request, read_file, write_file)),
+			Arc::new(move |request, read_file, write_file| {
+				run_python_process(&process, &engine, request, read_file, write_file)
+			}),
 		);
 
 		microservice.start().map_err(any_error_to_py)