pool.rs
raw
use std::{
sync::{Condvar, Mutex},
time::Duration,
};
use repository::Error;
/// How long a lonely client backs off between attempts when the server
/// refuses new connections; roughly twenty seconds of patience in all.
const BUSY_BACKOFF: &[Duration] = &[
Duration::from_millis(500),
Duration::from_millis(1000),
Duration::from_millis(2000),
Duration::from_millis(4000),
Duration::from_millis(6000),
Duration::from_millis(8000),
];
/// A capped pool of protocol connections shared by concurrent pipeline
/// threads.
///
/// FTP and SFTP sessions serve one command at a time, so each thread
/// needs a connection to itself while it works — but servers limit how
/// many logins they accept, so the pool never holds more than `max`
/// connections; threads beyond that wait for one to be returned.
///
/// Connections are created on demand, reused once returned healthy, and
/// discarded on any failure. An operation that fails on a *pooled*
/// connection is retried on another (ultimately a freshly dialed one),
/// so a connection gone stale in the pool never surfaces as a backup
/// failure — every backend operation is idempotent, which makes the
/// blind retry safe.
///
/// A dial refused because the server is full (per the backend's
/// `server_is_full` classifier) is handled adaptively: while this pool
/// holds connections, the thread simply waits for one of them — the
/// server's limit acts as the effective cap; when it holds none (other
/// clients own every slot), the dial is retried on a backoff schedule
/// before the refusal becomes an error.
pub(crate) struct Pool<C> {
state: Mutex<PoolState<C>>,
returned: Condvar,
max: usize,
backoff: &'static [Duration],
}
struct PoolState<C> {
idle: Vec<C>,
/// Connections in existence: idle plus checked out.
total: usize,
}
impl<C> Pool<C> {
pub(crate) fn new(max: usize) -> Pool<C> {
Pool {
state: Mutex::new(PoolState {
idle: Vec::new(),
total: 0,
}),
returned: Condvar::new(),
max: max.max(1),
backoff: BUSY_BACKOFF,
}
}
#[cfg(test)]
pub(crate) fn with_backoff(max: usize, backoff: &'static [Duration]) -> Pool<C> {
Pool {
backoff,
..Pool::new(max)
}
}
/// Runs `op` with a connection, dialing via `connect` as needed.
pub(crate) fn with<T>(
&self,
connect: impl Fn() -> Result<C, Error>,
server_is_full: impl Fn(&Error) -> bool,
mut op: impl FnMut(&mut C) -> Result<T, Error>,
) -> Result<T, Error> {
loop {
let (mut connection, was_pooled) = self.checkout(&connect, &server_is_full)?;
match op(&mut connection) {
Ok(value) => {
let mut state = self.state.lock().unwrap();
state.idle.push(connection);
drop(state);
self.returned.notify_one();
return Ok(value);
}
Err(err) => {
// The connection is suspect either way; discard it
drop(connection);
let mut state = self.state.lock().unwrap();
state.total -= 1;
drop(state);
self.returned.notify_one();
if !was_pooled {
// Even a fresh connection fails: the error is
// real
return Err(err);
}
// A pooled connection may just have gone stale
// (idle timeout, say); try again on another
}
}
}
}
/// Obtains a connection: an idle pooled one, or a fresh dial while
/// under the cap, waiting otherwise. Returns whether the connection
/// came from the pool.
fn checkout(
&self,
connect: &impl Fn() -> Result<C, Error>,
server_is_full: &impl Fn(&Error) -> bool,
) -> Result<(C, bool), Error> {
let mut backoff = self.backoff.iter();
let mut state = self.state.lock().unwrap();
loop {
if let Some(connection) = state.idle.pop() {
return Ok((connection, true));
}
if state.total < self.max {
state.total += 1;
drop(state);
match connect() {
Ok(connection) => return Ok((connection, false)),
Err(err) => {
state = self.state.lock().unwrap();
state.total -= 1;
if !server_is_full(&err) {
return Err(err);
}
if state.total > 0 {
// The server is at its limit but some of
// its slots are ours; wait to reuse one —
// the server's limit is the effective cap
} else {
// Every slot belongs to someone else; back
// off and try again, for a while
let Some(delay) = backoff.next() else {
return Err(err);
};
drop(state);
std::thread::sleep(*delay);
state = self.state.lock().unwrap();
continue;
}
}
}
}
// At the cap (ours or the server's): wait for a connection
// to come back. The timeout only guards against a missed
// wakeup; the loop re-decides from current state.
state = self
.returned
.wait_timeout(state, Duration::from_millis(500))
.unwrap()
.0;
}
}
}
/// A process-unique suffix for temporary upload names, so concurrent
/// writers (even from different processes) never collide.
pub(crate) fn temp_suffix() -> String {
use std::sync::atomic::{AtomicU64, Ordering};
static COUNTER: AtomicU64 = AtomicU64::new(0);
format!(
".part-{}-{}",
std::process::id(),
COUNTER.fetch_add(1, Ordering::Relaxed)
)
}
#[cfg(test)]
mod tests {
use std::sync::{
Arc, Barrier,
atomic::{AtomicUsize, Ordering},
};
use super::*;
fn busy_error() -> Error {
Error::Backend("server is full".to_string())
}
fn is_busy(err: &Error) -> bool {
matches!(err, Error::Backend(text) if text == "server is full")
}
/// No matter how many threads pile in, connections never exceed the
/// cap.
#[test]
fn the_cap_is_respected() {
let pool: Arc<Pool<()>> = Arc::new(Pool::new(3));
let dialed = Arc::new(AtomicUsize::new(0));
let in_use = Arc::new(AtomicUsize::new(0));
let high_water = Arc::new(AtomicUsize::new(0));
let workers: Vec<_> = (0..16)
.map(|_| {
let pool = pool.clone();
let dialed = dialed.clone();
let in_use = in_use.clone();
let high_water = high_water.clone();
std::thread::spawn(move || {
pool.with(
|| {
dialed.fetch_add(1, Ordering::AcqRel);
Ok(())
},
|_| false,
|_| {
let now = in_use.fetch_add(1, Ordering::AcqRel) + 1;
high_water.fetch_max(now, Ordering::AcqRel);
std::thread::sleep(Duration::from_millis(5));
in_use.fetch_sub(1, Ordering::AcqRel);
Ok(())
},
)
.unwrap();
})
})
.collect();
for worker in workers {
worker.join().unwrap();
}
assert!(high_water.load(Ordering::Acquire) <= 3);
assert!(dialed.load(Ordering::Acquire) <= 3);
}
/// Operations failing on stale pooled connections retry until a
/// fresh connection settles the question.
#[test]
fn stale_pooled_connections_are_retried_through() {
let pool: Arc<Pool<usize>> = Arc::new(Pool::new(2));
let dialed = Arc::new(AtomicUsize::new(0));
let connect = {
let dialed = dialed.clone();
move || Ok(dialed.fetch_add(1, Ordering::AcqRel))
};
// Seed the pool with two connections (ids 0 and 1) by holding
// both simultaneously
let barrier = Arc::new(Barrier::new(2));
let holders: Vec<_> = (0..2)
.map(|_| {
let pool = pool.clone();
let connect = connect.clone();
let barrier = barrier.clone();
std::thread::spawn(move || {
pool.with(
connect,
|_| false,
|_| {
barrier.wait();
Ok(())
},
)
.unwrap();
})
})
.collect();
for holder in holders {
holder.join().unwrap();
}
assert_eq!(dialed.load(Ordering::Acquire), 2);
// Both pooled connections (ids 0, 1) now act stale; only a
// fresh one (id >= 2) works
let value = pool
.with(
connect,
|_| false,
|id| {
if *id < 2 {
Err(Error::Backend("stale".to_string()))
} else {
Ok(*id)
}
},
)
.unwrap();
assert_eq!(value, 2);
}
/// When the server refuses a dial but we already hold a connection,
/// the thread waits for it instead of failing.
#[test]
fn full_server_waits_for_our_own_connections() {
let pool: Arc<Pool<()>> = Arc::new(Pool::new(4));
let dialed = Arc::new(AtomicUsize::new(0));
// The "server" accepts exactly one connection, ever
let connect = {
let dialed = dialed.clone();
move || {
if dialed.fetch_add(1, Ordering::AcqRel) == 0 {
Ok(())
} else {
Err(busy_error())
}
}
};
let workers: Vec<_> = (0..4)
.map(|_| {
let pool = pool.clone();
let connect = connect.clone();
std::thread::spawn(move || {
pool.with(connect, is_busy, |_| {
std::thread::sleep(Duration::from_millis(10));
Ok(())
})
.unwrap();
})
})
.collect();
for worker in workers {
worker.join().unwrap();
}
}
/// When the server is full and we hold nothing, the dial is retried
/// on the backoff schedule and then fails honestly.
#[test]
fn lonely_full_server_backs_off_then_errors() {
static QUICK: &[Duration] = &[Duration::from_millis(1), Duration::from_millis(1)];
let pool: Pool<()> = Pool::with_backoff(2, QUICK);
let dialed = AtomicUsize::new(0);
let result = pool.with(
|| {
dialed.fetch_add(1, Ordering::AcqRel);
Err(busy_error())
},
is_busy,
|_: &mut ()| Ok(()),
);
assert!(result.is_err());
assert_eq!(
dialed.load(Ordering::Acquire),
3,
"one initial dial plus one per backoff step"
);
}
}