344 lines
14 KiB
Rust
344 lines
14 KiB
Rust
//! Process-wide connection pool keyed by `(url, principal)`.
|
|
//!
|
|
//! Two [`crate::ToolServer`] builds with the same `(url, credential)`
|
|
//! observe the same `Arc<HubConnection>`; distinct credentials open
|
|
//! distinct sockets. The pool is the canonical entry point — direct
|
|
//! [`crate::HubConnection::connect`] calls are reserved for tests and
|
|
//! one-shot programs that explicitly want unpooled behaviour.
|
|
|
|
use std::sync::Arc;
|
|
use std::time::{Duration, Instant};
|
|
|
|
use dashmap::DashMap;
|
|
use tokio::sync::OnceCell;
|
|
use tokio::task::JoinHandle;
|
|
use url::Url;
|
|
use xai_tool_protocol::ConnectionKind;
|
|
|
|
use crate::auth::AuthProvider;
|
|
use crate::connection::{
|
|
ConnKey, ConnectCallback, ConnectionConfig, ConnectionTuning, DisconnectCallback,
|
|
HubConnection, ReconnectCallback,
|
|
};
|
|
use crate::error::ClientError;
|
|
|
|
/// Idle window for the reaper: a pooled connection is evictable once it is
|
|
/// unused (`Arc::strong_count == 1`, i.e. only the pool holds it) **and**
|
|
/// `now - last_handout >= DEFAULT_POOL_IDLE_TTL`.
|
|
///
|
|
/// Note the clock is `last_handout` (the last time the pool returned the
|
|
/// connection), not the moment the last consumer `Arc` was dropped: a
|
|
/// connection held longer than the TTL and then released is eligible on the
|
|
/// very next sweep, with no extra post-drop grace period. The only hard
|
|
/// guarantee is that an in-use connection (`strong_count > 1`) is never
|
|
/// reaped. Tuned well above the server's own 90s dead-peer idle timeout so a
|
|
/// short borrow between turns of an active conversation isn't churned.
|
|
pub const DEFAULT_POOL_IDLE_TTL: Duration = Duration::from_secs(300);
|
|
|
|
/// How often the shared pool's idle reaper scans for evictable entries.
|
|
pub const DEFAULT_POOL_SWEEP_INTERVAL: Duration = Duration::from_secs(60);
|
|
|
|
/// A pooled connection plus the last time it was handed out to a caller.
|
|
///
|
|
/// `last_handout` is refreshed on every [`HubConnectionPool::get_or_connect`]
|
|
/// hit (and on the initial insert), so a connection that is repeatedly
|
|
/// re-fetched never looks idle even if its [`Arc`] strong count briefly
|
|
/// returns to 1 between fetches. Eviction additionally requires
|
|
/// `Arc::strong_count == 1` (only the pool holds it), so a connection a
|
|
/// consumer still holds is never reaped regardless of `last_handout`.
|
|
struct Pooled {
|
|
conn: Arc<HubConnection>,
|
|
last_handout: Instant,
|
|
}
|
|
|
|
/// The process-global pool used by [`HubConnectionPool::shared`].
|
|
///
|
|
/// `tokio::sync::OnceCell` is preferred over `std::sync::OnceLock` /
|
|
/// `LazyLock` here because the pool is only ever observed from
|
|
/// async contexts (the connection actor lives on a tokio runtime
|
|
/// already), so the async-aware `get_or_init` semantics avoid the
|
|
/// blocking-init footgun of the sync alternatives without taking a
|
|
/// hard dependency on additional sync primitives.
|
|
///
|
|
/// Tests MUST use [`HubConnectionPool::new`] to avoid cross-test
|
|
/// pollution: cargo runs all integration tests in the same binary
|
|
/// unless otherwise configured, so any test that touches
|
|
/// `HubConnectionPool::shared()` leaves the pool populated for
|
|
/// subsequent tests.
|
|
static SHARED: OnceCell<Arc<HubConnectionPool>> = OnceCell::const_new();
|
|
|
|
/// Pool of live server connections.
|
|
pub struct HubConnectionPool {
|
|
connections: DashMap<ConnKey, Pooled>,
|
|
}
|
|
|
|
impl std::fmt::Debug for HubConnectionPool {
|
|
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
|
f.debug_struct("HubConnectionPool")
|
|
.field("connection_count", &self.connections.len())
|
|
.finish()
|
|
}
|
|
}
|
|
|
|
impl HubConnectionPool {
|
|
/// Build a fresh, unshared pool. Tests typically use this so each
|
|
/// test sees an isolated registry.
|
|
pub fn new() -> Arc<Self> {
|
|
Arc::new(Self {
|
|
connections: DashMap::new(),
|
|
})
|
|
}
|
|
|
|
/// Return the process-wide shared pool, lazily initialising it on
|
|
/// the first call. Subsequent callers in the same process observe
|
|
/// the same `Arc`.
|
|
///
|
|
/// The shared pool spawns an idle reaper (see [`Self::spawn_idle_reaper`])
|
|
/// exactly once, so a connection that is unused (`strong_count == 1`) and
|
|
/// has not been handed out for [`DEFAULT_POOL_IDLE_TTL`] is closed instead
|
|
/// of living for the whole process lifetime. (Unpooled / test pools built
|
|
/// via [`Self::new`] do not
|
|
/// get a reaper; they can call [`Self::sweep_idle`] directly.)
|
|
pub async fn shared() -> Arc<Self> {
|
|
SHARED
|
|
.get_or_init(|| async {
|
|
let pool = Self::new();
|
|
pool.spawn_idle_reaper(DEFAULT_POOL_IDLE_TTL, DEFAULT_POOL_SWEEP_INTERVAL);
|
|
pool
|
|
})
|
|
.await
|
|
.clone()
|
|
}
|
|
|
|
/// Look up an existing pooled connection for `(url, credential)`,
|
|
/// or open a fresh one if no pooled entry exists.
|
|
///
|
|
/// `kind` is the connection role announced in the hello frame. The
|
|
/// pool is keyed by `(url, principal)` only; mixing
|
|
/// [`ConnectionKind`] values for the same `(url, principal)` is a
|
|
/// caller error and surfaces as a [`ClientError::InvalidConfig`].
|
|
///
|
|
/// The optional extra access key is not part of the pool key, so the first
|
|
/// caller's key is the one carried on a shared connection's handshake (in
|
|
/// practice it is a per-deployment constant). The plaintext-scheme guard is
|
|
/// re-checked on every call below so it can't be bypassed by a cached
|
|
/// insecure entry.
|
|
pub async fn get_or_connect(
|
|
self: &Arc<Self>,
|
|
url: Url,
|
|
credential: Arc<dyn AuthProvider>,
|
|
kind: ConnectionKind,
|
|
on_reconnect: Option<Arc<ReconnectCallback>>,
|
|
on_disconnect: Option<Arc<DisconnectCallback>>,
|
|
server_id: Option<xai_tool_protocol::ServerId>,
|
|
alpha_test_key: Option<String>,
|
|
allow_insecure_ws: bool,
|
|
) -> Result<Arc<HubConnection>, ClientError> {
|
|
self.get_or_connect_tuned(
|
|
url,
|
|
credential,
|
|
kind,
|
|
on_reconnect,
|
|
on_disconnect,
|
|
None, // on_connect (unused by the simple wrapper)
|
|
server_id,
|
|
None,
|
|
None,
|
|
alpha_test_key,
|
|
allow_insecure_ws,
|
|
ConnectionTuning::default(),
|
|
)
|
|
.await
|
|
}
|
|
|
|
/// Like [`Self::get_or_connect`] but carries optional connection-tuning
|
|
/// knobs ([`ConnectionTuning`]) onto a freshly-opened connection. A
|
|
/// `Default` tuning is behaviourally identical to `get_or_connect`, so
|
|
/// existing callers are unaffected.
|
|
///
|
|
/// Tuning binds to the socket at open time: it takes effect only when
|
|
/// THIS call opens the connection. Because the pool dedups by
|
|
/// `(url, principal)`, a hit on an existing entry returns that
|
|
/// connection as-is and the `tuning` argument is ignored — the first
|
|
/// opener's ping/backoff settings win for the lifetime of the pooled
|
|
/// connection. Callers that need distinct tuning must use a distinct
|
|
/// `(url, principal)` or an unpooled [`HubConnection::connect`].
|
|
pub(crate) async fn get_or_connect_tuned(
|
|
self: &Arc<Self>,
|
|
url: Url,
|
|
credential: Arc<dyn AuthProvider>,
|
|
kind: ConnectionKind,
|
|
on_reconnect: Option<Arc<ReconnectCallback>>,
|
|
on_disconnect: Option<Arc<DisconnectCallback>>,
|
|
on_connect: Option<Arc<ConnectCallback>>,
|
|
server_id: Option<xai_tool_protocol::ServerId>,
|
|
server_description: Option<String>,
|
|
server_metadata: Option<serde_json::Value>,
|
|
alpha_test_key: Option<String>,
|
|
allow_insecure_ws: bool,
|
|
tuning: ConnectionTuning,
|
|
) -> Result<Arc<HubConnection>, ClientError> {
|
|
if url.scheme() != "wss" && !crate::connection::host_is_loopback(&url) && !allow_insecure_ws
|
|
{
|
|
return Err(ClientError::InsecureScheme { url });
|
|
}
|
|
let key = ConnKey {
|
|
url: url.as_str().to_owned(),
|
|
principal: credential.principal_key(),
|
|
};
|
|
if let Some(mut existing) = self.connections.get_mut(&key) {
|
|
existing.last_handout = Instant::now();
|
|
let conn = existing.conn.clone();
|
|
drop(existing);
|
|
if conn.kind() != kind {
|
|
return Err(ClientError::InvalidConfig(format!(
|
|
"pool entry for {} bound to {:?}; rebuild requested {:?}",
|
|
key.url,
|
|
conn.kind(),
|
|
kind
|
|
)));
|
|
}
|
|
return Ok(conn);
|
|
}
|
|
let config = ConnectionConfig {
|
|
url,
|
|
credential,
|
|
kind,
|
|
on_reconnect,
|
|
on_disconnect,
|
|
on_connect,
|
|
server_id,
|
|
server_description,
|
|
server_metadata,
|
|
outbound_buffer: None,
|
|
tuning,
|
|
alpha_test_key,
|
|
allow_insecure_ws,
|
|
on_fatal: Some(Arc::downgrade(self)),
|
|
};
|
|
let conn = HubConnection::connect(config).await?;
|
|
// Race window: another caller may have inserted between our
|
|
// `get` and `connect`. Resolve via `entry().or_insert_with`
|
|
// semantics — if we lose the race we drop our fresh
|
|
// connection and adopt the winning one.
|
|
match self.connections.entry(key.clone()) {
|
|
dashmap::Entry::Occupied(mut existing) => {
|
|
existing.get_mut().last_handout = Instant::now();
|
|
let winner = existing.get().conn.clone();
|
|
drop(conn);
|
|
if winner.kind() != kind {
|
|
return Err(ClientError::InvalidConfig(format!(
|
|
"pool entry for {} bound to {:?}; rebuild requested {:?}",
|
|
key.url,
|
|
winner.kind(),
|
|
kind
|
|
)));
|
|
}
|
|
Ok(winner)
|
|
}
|
|
dashmap::Entry::Vacant(slot) => {
|
|
crate::metrics::pool_connections_inc();
|
|
slot.insert(Pooled {
|
|
conn: conn.clone(),
|
|
last_handout: Instant::now(),
|
|
});
|
|
Ok(conn)
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Number of pooled connections. Intended for tests and metrics.
|
|
pub fn len(&self) -> usize {
|
|
self.connections.len()
|
|
}
|
|
|
|
/// `true` when no connection is pooled.
|
|
pub fn is_empty(&self) -> bool {
|
|
self.connections.is_empty()
|
|
}
|
|
|
|
/// Forget the pooled connection for `key`. The actual underlying
|
|
/// `Arc<HubConnection>` is dropped only when no other holder
|
|
/// keeps a reference; the next [`Self::get_or_connect`] for the
|
|
/// same key opens a fresh socket.
|
|
pub fn forget(&self, key: &ConnKey) {
|
|
if self.connections.remove(key).is_some() {
|
|
crate::metrics::pool_connections_dec();
|
|
}
|
|
}
|
|
|
|
/// Close and remove every pooled connection that is BOTH unused (no live
|
|
/// consumer holds an `Arc` — only the pool does, so `strong_count == 1`)
|
|
/// AND idle longer than `idle_ttl` (no hand-out within the window).
|
|
/// Removing the entry drops the pool's last `Arc<HubConnection>`, whose
|
|
/// `Drop` closes the socket.
|
|
///
|
|
/// The strong-count check runs inside the map's per-shard lock (via
|
|
/// [`DashMap::retain`]), serialised against `get_or_connect`, so a
|
|
/// connection handed out concurrently is never evicted out from under a
|
|
/// caller. Returns the number of connections evicted.
|
|
pub fn sweep_idle(&self, idle_ttl: Duration) -> usize {
|
|
let now = Instant::now();
|
|
let mut evicted = 0usize;
|
|
self.connections.retain(|_key, pooled| {
|
|
let idle_for = now.saturating_duration_since(pooled.last_handout);
|
|
// `strong_count == 1` ⇒ only this pool entry references the
|
|
// connection, so no consumer can still be using it.
|
|
let unused = Arc::strong_count(&pooled.conn) == 1;
|
|
let evict = unused && idle_for >= idle_ttl;
|
|
if evict {
|
|
evicted += 1;
|
|
}
|
|
!evict
|
|
});
|
|
for _ in 0..evicted {
|
|
crate::metrics::pool_connections_dec();
|
|
crate::metrics::pool_evictions_inc();
|
|
}
|
|
evicted
|
|
}
|
|
|
|
/// Spawn a background task that calls [`Self::sweep_idle`] every
|
|
/// `sweep_interval`, closing connections idle longer than `idle_ttl`.
|
|
///
|
|
/// The task holds a [`std::sync::Weak`] to the pool, so it exits on its
|
|
/// own once the last strong `Arc<HubConnectionPool>` is dropped (it never
|
|
/// keeps the pool alive). The first interval tick is skipped so a
|
|
/// freshly-handed-out connection is never swept on the immediate tick.
|
|
pub fn spawn_idle_reaper(
|
|
self: &Arc<Self>,
|
|
idle_ttl: Duration,
|
|
sweep_interval: Duration,
|
|
) -> JoinHandle<()> {
|
|
let weak = Arc::downgrade(self);
|
|
tokio::spawn(async move {
|
|
let mut ticker = tokio::time::interval(sweep_interval);
|
|
ticker.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
|
|
// `interval`'s first tick resolves immediately; skip it.
|
|
ticker.tick().await;
|
|
loop {
|
|
ticker.tick().await;
|
|
let Some(pool) = weak.upgrade() else { break };
|
|
pool.sweep_idle(idle_ttl);
|
|
}
|
|
})
|
|
}
|
|
|
|
/// Like [`Self::forget`] but identity-checked: only removes the slot
|
|
/// when `predicate` accepts the currently-stored connection. The
|
|
/// self-evicting actor passes an `Arc::ptr_eq` check so a race-loser
|
|
/// can never drop the winner's fresh entry (ABA-safe).
|
|
pub(crate) fn forget_if(
|
|
&self,
|
|
key: &ConnKey,
|
|
predicate: impl FnOnce(&Arc<HubConnection>) -> bool,
|
|
) {
|
|
if self
|
|
.connections
|
|
.remove_if(key, |_, pooled| predicate(&pooled.conn))
|
|
.is_some()
|
|
{
|
|
crate::metrics::pool_connections_dec();
|
|
}
|
|
}
|
|
}
|