diff --git a/Cargo.toml b/Cargo.toml index 27a3e2a..4e7c487 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,7 +1,7 @@ [package] name = "tmp-postgrust" -version = "0.11.0" -authors = ["John Children "] +version = "0.11.1" +authors = ["John Children "] license = "MIT" edition = "2018" description = "Temporary postgresql instances for testing" @@ -23,6 +23,7 @@ tokio = { version = "1.8", features = ["parking_lot", "rt", "sync", "io-util", " tracing = "0.1" which = "4.0" futures-util = { version = "0.3", optional = true } +num_cpus = "1.17.0" [dev-dependencies] test-log = { version = "0.2", default-features = false, features = ["trace"] } diff --git a/src/asynchronous.rs b/src/asynchronous.rs index b93c813..f52913e 100644 --- a/src/asynchronous.rs +++ b/src/asynchronous.rs @@ -1,14 +1,15 @@ -use nix::sys::signal::{self, Signal}; -use nix::unistd::{Pid, Uid, User}; use std::convert::TryInto; use std::path::Path; use std::process::Stdio; -use std::sync::Arc; +use std::sync::{Arc, OnceLock}; +use nix::sys::signal::{self, Signal}; +use nix::unistd::{Pid, Uid, User}; use tempfile::TempDir; use tokio::fs::create_dir_all; use tokio::io::Lines; use tokio::process::{ChildStderr, ChildStdout}; +use tokio::sync::{Semaphore, SemaphorePermit}; use tokio::{ io::BufReader, process::{Child, Command}, @@ -19,6 +20,9 @@ use crate::errors::{ProcessCapture, TmpPostgrustError, TmpPostgrustResult}; use crate::search::{all_dir_entries, build_copy_dst_path}; use crate::POSTGRES_UID_GID; +/// Limit the total processes that can be running at any one time. +pub(crate) static MAX_CONCURRENT_PROCESSES: OnceLock = OnceLock::new(); + #[instrument(skip(command, fail))] async fn exec_process( command: &mut Command, @@ -196,6 +200,8 @@ pub struct ProcessGuard { /// Socket directory for connection to the running process. pub(crate) socket_dir: Arc, pub(crate) postgres_process: Child, + // Limit the total concurrent processes. + pub(crate) _process_permit: SemaphorePermit<'static>, } impl Drop for ProcessGuard { diff --git a/src/errors.rs b/src/errors.rs index 9d61e80..fed75cd 100644 --- a/src/errors.rs +++ b/src/errors.rs @@ -75,6 +75,10 @@ pub enum TmpPostgrustError { /// The explicit directory that was searched, if one was provided. searched_dir: Option, }, + /// Error when a process permit cannot be acquired for a new postgres subprocess. + #[cfg(feature = "tokio-process")] + #[error("acquiring process permit failed")] + AsyncProcessPermitAcquireError(#[source] tokio::sync::AcquireError), } /// Result type for `TmpPostgrustError`, used by functions in this crate. diff --git a/src/lib.rs b/src/lib.rs index 5e2ef6c..5d0cb8b 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -516,6 +516,13 @@ impl TmpPostgrustFactory { .fetch_add(1, std::sync::atomic::Ordering::SeqCst); asynchronous::chown_to_non_root(dir.path()).await?; + + let permit = crate::asynchronous::MAX_CONCURRENT_PROCESSES + .get_or_init(|| tokio::sync::Semaphore::new(num_cpus::get())) + .acquire() + .await + .map_err(TmpPostgrustError::AsyncProcessPermitAcquireError)?; + let mut postgres_process_handle = asynchronous::start_postgres_subprocess(&self.binaries.postgres, dir.path(), port)?; let stdout = postgres_process_handle.stdout.take().unwrap(); @@ -542,6 +549,7 @@ impl TmpPostgrustFactory { _cache_directory: Arc::clone(&self.cache_dir), socket_dir: Arc::clone(&self.socket_dir), postgres_process: postgres_process_handle, + _process_permit: permit, }) } }