LCOV - code coverage report
Current view: top level - proxy/src/serverless - conn_pool_lib.rs (source / functions) Coverage Total Hit
Test: 07bee600374ccd486c69370d0972d9035964fe68.info Lines: 54.6 % 456 249
Test Date: 2025-02-20 13:11:02 Functions: 24.7 % 93 23

            Line data    Source code
       1              : use std::collections::HashMap;
       2              : use std::marker::PhantomData;
       3              : use std::ops::Deref;
       4              : use std::sync::atomic::{self, AtomicUsize};
       5              : use std::sync::{Arc, Weak};
       6              : use std::time::Duration;
       7              : 
       8              : use clashmap::ClashMap;
       9              : use parking_lot::RwLock;
      10              : use postgres_client::ReadyForQueryStatus;
      11              : use rand::Rng;
      12              : use tracing::{debug, info, Span};
      13              : 
      14              : use super::backend::HttpConnError;
      15              : use super::conn_pool::ClientDataRemote;
      16              : use super::http_conn_pool::ClientDataHttp;
      17              : use super::local_conn_pool::ClientDataLocal;
      18              : use crate::auth::backend::ComputeUserInfo;
      19              : use crate::context::RequestContext;
      20              : use crate::control_plane::messages::{ColdStartInfo, MetricsAuxInfo};
      21              : use crate::metrics::{HttpEndpointPoolsGuard, Metrics};
      22              : use crate::types::{DbName, EndpointCacheKey, RoleName};
      23              : use crate::usage_metrics::{Ids, MetricCounter, USAGE_METRICS};
      24              : 
      25              : #[derive(Debug, Clone)]
      26              : pub(crate) struct ConnInfo {
      27              :     pub(crate) user_info: ComputeUserInfo,
      28              :     pub(crate) dbname: DbName,
      29              : }
      30              : 
      31              : impl ConnInfo {
      32              :     // hm, change to hasher to avoid cloning?
      33            3 :     pub(crate) fn db_and_user(&self) -> (DbName, RoleName) {
      34            3 :         (self.dbname.clone(), self.user_info.user.clone())
      35            3 :     }
      36              : 
      37            2 :     pub(crate) fn endpoint_cache_key(&self) -> Option<EndpointCacheKey> {
      38            2 :         // We don't want to cache http connections for ephemeral endpoints.
      39            2 :         if self.user_info.options.is_ephemeral() {
      40            0 :             None
      41              :         } else {
      42            2 :             Some(self.user_info.endpoint_cache_key())
      43              :         }
      44            2 :     }
      45              : }
      46              : 
      47              : #[derive(Clone)]
      48              : pub(crate) enum ClientDataEnum {
      49              :     Remote(ClientDataRemote),
      50              :     Local(ClientDataLocal),
      51              :     Http(ClientDataHttp),
      52              : }
      53              : 
      54              : #[derive(Clone)]
      55              : pub(crate) struct ClientInnerCommon<C: ClientInnerExt> {
      56              :     pub(crate) inner: C,
      57              :     pub(crate) aux: MetricsAuxInfo,
      58              :     pub(crate) conn_id: uuid::Uuid,
      59              :     pub(crate) data: ClientDataEnum, // custom client data like session, key, jti
      60              : }
      61              : 
      62              : impl<C: ClientInnerExt> Drop for ClientInnerCommon<C> {
      63            7 :     fn drop(&mut self) {
      64            7 :         match &mut self.data {
      65            7 :             ClientDataEnum::Remote(remote_data) => {
      66            7 :                 remote_data.cancel();
      67            7 :             }
      68            0 :             ClientDataEnum::Local(local_data) => {
      69            0 :                 local_data.cancel();
      70            0 :             }
      71            0 :             ClientDataEnum::Http(_http_data) => (),
      72              :         }
      73            7 :     }
      74              : }
      75              : 
      76              : impl<C: ClientInnerExt> ClientInnerCommon<C> {
      77            6 :     pub(crate) fn get_conn_id(&self) -> uuid::Uuid {
      78            6 :         self.conn_id
      79            6 :     }
      80              : 
      81            0 :     pub(crate) fn get_data(&mut self) -> &mut ClientDataEnum {
      82            0 :         &mut self.data
      83            0 :     }
      84              : }
      85              : 
      86              : pub(crate) struct ConnPoolEntry<C: ClientInnerExt> {
      87              :     pub(crate) conn: ClientInnerCommon<C>,
      88              :     pub(crate) _last_access: std::time::Instant,
      89              : }
      90              : 
      91              : // Per-endpoint connection pool, (dbname, username) -> DbUserConnPool
      92              : // Number of open connections is limited by the `max_conns_per_endpoint`.
      93              : pub(crate) struct EndpointConnPool<C: ClientInnerExt> {
      94              :     pools: HashMap<(DbName, RoleName), DbUserConnPool<C>>,
      95              :     total_conns: usize,
      96              :     /// max # connections per endpoint
      97              :     max_conns: usize,
      98              :     _guard: HttpEndpointPoolsGuard<'static>,
      99              :     global_connections_count: Arc<AtomicUsize>,
     100              :     global_pool_size_max_conns: usize,
     101              :     pool_name: String,
     102              : }
     103              : 
     104              : impl<C: ClientInnerExt> EndpointConnPool<C> {
     105            0 :     pub(crate) fn new(
     106            0 :         hmap: HashMap<(DbName, RoleName), DbUserConnPool<C>>,
     107            0 :         tconns: usize,
     108            0 :         max_conns_per_endpoint: usize,
     109            0 :         global_connections_count: Arc<AtomicUsize>,
     110            0 :         max_total_conns: usize,
     111            0 :         pname: String,
     112            0 :     ) -> Self {
     113            0 :         Self {
     114            0 :             pools: hmap,
     115            0 :             total_conns: tconns,
     116            0 :             max_conns: max_conns_per_endpoint,
     117            0 :             _guard: Metrics::get().proxy.http_endpoint_pools.guard(),
     118            0 :             global_connections_count,
     119            0 :             global_pool_size_max_conns: max_total_conns,
     120            0 :             pool_name: pname,
     121            0 :         }
     122            0 :     }
     123              : 
     124            0 :     pub(crate) fn get_conn_entry(
     125            0 :         &mut self,
     126            0 :         db_user: (DbName, RoleName),
     127            0 :     ) -> Option<ConnPoolEntry<C>> {
     128            0 :         let Self {
     129            0 :             pools,
     130            0 :             total_conns,
     131            0 :             global_connections_count,
     132            0 :             ..
     133            0 :         } = self;
     134            0 :         pools.get_mut(&db_user).and_then(|pool_entries| {
     135            0 :             let (entry, removed) = pool_entries.get_conn_entry(total_conns);
     136            0 :             global_connections_count.fetch_sub(removed, atomic::Ordering::Relaxed);
     137            0 :             entry
     138            0 :         })
     139            0 :     }
     140              : 
     141            0 :     pub(crate) fn remove_client(
     142            0 :         &mut self,
     143            0 :         db_user: (DbName, RoleName),
     144            0 :         conn_id: uuid::Uuid,
     145            0 :     ) -> bool {
     146            0 :         let Self {
     147            0 :             pools,
     148            0 :             total_conns,
     149            0 :             global_connections_count,
     150            0 :             ..
     151            0 :         } = self;
     152            0 :         if let Some(pool) = pools.get_mut(&db_user) {
     153            0 :             let old_len = pool.get_conns().len();
     154            0 :             pool.get_conns()
     155            0 :                 .retain(|conn| conn.conn.get_conn_id() != conn_id);
     156            0 :             let new_len = pool.get_conns().len();
     157            0 :             let removed = old_len - new_len;
     158            0 :             if removed > 0 {
     159            0 :                 global_connections_count.fetch_sub(removed, atomic::Ordering::Relaxed);
     160            0 :                 Metrics::get()
     161            0 :                     .proxy
     162            0 :                     .http_pool_opened_connections
     163            0 :                     .get_metric()
     164            0 :                     .dec_by(removed as i64);
     165            0 :             }
     166            0 :             *total_conns -= removed;
     167            0 :             removed > 0
     168              :         } else {
     169            0 :             false
     170              :         }
     171            0 :     }
     172              : 
     173            6 :     pub(crate) fn get_name(&self) -> &str {
     174            6 :         &self.pool_name
     175            6 :     }
     176              : 
     177            0 :     pub(crate) fn get_pool(&self, db_user: (DbName, RoleName)) -> Option<&DbUserConnPool<C>> {
     178            0 :         self.pools.get(&db_user)
     179            0 :     }
     180              : 
     181            0 :     pub(crate) fn get_pool_mut(
     182            0 :         &mut self,
     183            0 :         db_user: (DbName, RoleName),
     184            0 :     ) -> Option<&mut DbUserConnPool<C>> {
     185            0 :         self.pools.get_mut(&db_user)
     186            0 :     }
     187              : 
     188            6 :     pub(crate) fn put(pool: &RwLock<Self>, conn_info: &ConnInfo, client: ClientInnerCommon<C>) {
     189            6 :         let conn_id = client.get_conn_id();
     190            6 :         let (max_conn, conn_count, pool_name) = {
     191            6 :             let pool = pool.read();
     192            6 :             (
     193            6 :                 pool.global_pool_size_max_conns,
     194            6 :                 pool.global_connections_count
     195            6 :                     .load(atomic::Ordering::Relaxed),
     196            6 :                 pool.get_name().to_string(),
     197            6 :             )
     198            6 :         };
     199            6 : 
     200            6 :         if client.inner.is_closed() {
     201            1 :             info!(%conn_id, "{}: throwing away connection '{conn_info}' because connection is closed", pool_name);
     202            1 :             return;
     203            5 :         }
     204            5 : 
     205            5 :         if conn_count >= max_conn {
     206            1 :             info!(%conn_id, "{}: throwing away connection '{conn_info}' because pool is full", pool_name);
     207            1 :             return;
     208            4 :         }
     209            4 : 
     210            4 :         // return connection to the pool
     211            4 :         let mut returned = false;
     212            4 :         let mut per_db_size = 0;
     213            4 :         let total_conns = {
     214            4 :             let mut pool = pool.write();
     215            4 : 
     216            4 :             if pool.total_conns < pool.max_conns {
     217            3 :                 let pool_entries = pool.pools.entry(conn_info.db_and_user()).or_default();
     218            3 :                 pool_entries.get_conns().push(ConnPoolEntry {
     219            3 :                     conn: client,
     220            3 :                     _last_access: std::time::Instant::now(),
     221            3 :                 });
     222            3 : 
     223            3 :                 returned = true;
     224            3 :                 per_db_size = pool_entries.get_conns().len();
     225            3 : 
     226            3 :                 pool.total_conns += 1;
     227            3 :                 pool.global_connections_count
     228            3 :                     .fetch_add(1, atomic::Ordering::Relaxed);
     229            3 :                 Metrics::get()
     230            3 :                     .proxy
     231            3 :                     .http_pool_opened_connections
     232            3 :                     .get_metric()
     233            3 :                     .inc();
     234            3 :             }
     235              : 
     236            4 :             pool.total_conns
     237            4 :         };
     238            4 : 
     239            4 :         // do logging outside of the mutex
     240            4 :         if returned {
     241            3 :             debug!(%conn_id, "{pool_name}: returning connection '{conn_info}' back to the pool, total_conns={total_conns}, for this (db, user)={per_db_size}");
     242              :         } else {
     243            1 :             info!(%conn_id, "{pool_name}: throwing away connection '{conn_info}' because pool is full, total_conns={total_conns}");
     244              :         }
     245            6 :     }
     246              : }
     247              : 
     248              : impl<C: ClientInnerExt> Drop for EndpointConnPool<C> {
     249            2 :     fn drop(&mut self) {
     250            2 :         if self.total_conns > 0 {
     251            2 :             self.global_connections_count
     252            2 :                 .fetch_sub(self.total_conns, atomic::Ordering::Relaxed);
     253            2 :             Metrics::get()
     254            2 :                 .proxy
     255            2 :                 .http_pool_opened_connections
     256            2 :                 .get_metric()
     257            2 :                 .dec_by(self.total_conns as i64);
     258            2 :         }
     259            2 :     }
     260              : }
     261              : 
     262              : pub(crate) struct DbUserConnPool<C: ClientInnerExt> {
     263              :     pub(crate) conns: Vec<ConnPoolEntry<C>>,
     264              :     pub(crate) initialized: Option<bool>, // a bit ugly, exists only for local pools
     265              : }
     266              : 
     267              : impl<C: ClientInnerExt> Default for DbUserConnPool<C> {
     268            2 :     fn default() -> Self {
     269            2 :         Self {
     270            2 :             conns: Vec::new(),
     271            2 :             initialized: None,
     272            2 :         }
     273            2 :     }
     274              : }
     275              : 
     276              : pub(crate) trait DbUserConn<C: ClientInnerExt>: Default {
     277              :     fn set_initialized(&mut self);
     278              :     fn is_initialized(&self) -> bool;
     279              :     fn clear_closed_clients(&mut self, conns: &mut usize) -> usize;
     280              :     fn get_conn_entry(&mut self, conns: &mut usize) -> (Option<ConnPoolEntry<C>>, usize);
     281              :     fn get_conns(&mut self) -> &mut Vec<ConnPoolEntry<C>>;
     282              : }
     283              : 
     284              : impl<C: ClientInnerExt> DbUserConn<C> for DbUserConnPool<C> {
     285            0 :     fn set_initialized(&mut self) {
     286            0 :         self.initialized = Some(true);
     287            0 :     }
     288              : 
     289            0 :     fn is_initialized(&self) -> bool {
     290            0 :         self.initialized.unwrap_or(false)
     291            0 :     }
     292              : 
     293            1 :     fn clear_closed_clients(&mut self, conns: &mut usize) -> usize {
     294            1 :         let old_len = self.conns.len();
     295            1 : 
     296            2 :         self.conns.retain(|conn| !conn.conn.inner.is_closed());
     297            1 : 
     298            1 :         let new_len = self.conns.len();
     299            1 :         let removed = old_len - new_len;
     300            1 :         *conns -= removed;
     301            1 :         removed
     302            1 :     }
     303              : 
     304            0 :     fn get_conn_entry(&mut self, conns: &mut usize) -> (Option<ConnPoolEntry<C>>, usize) {
     305            0 :         let mut removed = self.clear_closed_clients(conns);
     306            0 :         let conn = self.conns.pop();
     307            0 :         if conn.is_some() {
     308            0 :             *conns -= 1;
     309            0 :             removed += 1;
     310            0 :         }
     311              : 
     312            0 :         Metrics::get()
     313            0 :             .proxy
     314            0 :             .http_pool_opened_connections
     315            0 :             .get_metric()
     316            0 :             .dec_by(removed as i64);
     317            0 : 
     318            0 :         (conn, removed)
     319            0 :     }
     320              : 
     321            6 :     fn get_conns(&mut self) -> &mut Vec<ConnPoolEntry<C>> {
     322            6 :         &mut self.conns
     323            6 :     }
     324              : }
     325              : 
     326              : pub(crate) trait EndpointConnPoolExt<C: ClientInnerExt> {
     327              :     fn clear_closed(&mut self) -> usize;
     328              :     fn total_conns(&self) -> usize;
     329              : }
     330              : 
     331              : impl<C: ClientInnerExt> EndpointConnPoolExt<C> for EndpointConnPool<C> {
     332            1 :     fn clear_closed(&mut self) -> usize {
     333            1 :         let mut clients_removed: usize = 0;
     334            1 :         for db_pool in self.pools.values_mut() {
     335            1 :             clients_removed += db_pool.clear_closed_clients(&mut self.total_conns);
     336            1 :         }
     337            1 :         clients_removed
     338            1 :     }
     339              : 
     340            1 :     fn total_conns(&self) -> usize {
     341            1 :         self.total_conns
     342            1 :     }
     343              : }
     344              : 
     345              : pub(crate) struct GlobalConnPool<C, P>
     346              : where
     347              :     C: ClientInnerExt,
     348              :     P: EndpointConnPoolExt<C>,
     349              : {
     350              :     // endpoint -> per-endpoint connection pool
     351              :     //
     352              :     // That should be a fairly conteded map, so return reference to the per-endpoint
     353              :     // pool as early as possible and release the lock.
     354              :     pub(crate) global_pool: ClashMap<EndpointCacheKey, Arc<RwLock<P>>>,
     355              : 
     356              :     /// Number of endpoint-connection pools
     357              :     ///
     358              :     /// [`ClashMap::len`] iterates over all inner pools and acquires a read lock on each.
     359              :     /// That seems like far too much effort, so we're using a relaxed increment counter instead.
     360              :     /// It's only used for diagnostics.
     361              :     pub(crate) global_pool_size: AtomicUsize,
     362              : 
     363              :     /// Total number of connections in the pool
     364              :     pub(crate) global_connections_count: Arc<AtomicUsize>,
     365              : 
     366              :     pub(crate) config: &'static crate::config::HttpConfig,
     367              : 
     368              :     _marker: PhantomData<C>,
     369              : }
     370              : 
     371              : #[derive(Debug, Clone, Copy)]
     372              : pub struct GlobalConnPoolOptions {
     373              :     // Maximum number of connections per one endpoint.
     374              :     // Can mix different (dbname, username) connections.
     375              :     // When running out of free slots for a particular endpoint,
     376              :     // falls back to opening a new connection for each request.
     377              :     pub max_conns_per_endpoint: usize,
     378              : 
     379              :     pub gc_epoch: Duration,
     380              : 
     381              :     pub pool_shards: usize,
     382              : 
     383              :     pub idle_timeout: Duration,
     384              : 
     385              :     pub opt_in: bool,
     386              : 
     387              :     // Total number of connections in the pool.
     388              :     pub max_total_conns: usize,
     389              : }
     390              : 
     391              : impl<C, P> GlobalConnPool<C, P>
     392              : where
     393              :     C: ClientInnerExt,
     394              :     P: EndpointConnPoolExt<C>,
     395              : {
     396            1 :     pub(crate) fn new(config: &'static crate::config::HttpConfig) -> Arc<Self> {
     397            1 :         let shards = config.pool_options.pool_shards;
     398            1 :         Arc::new(Self {
     399            1 :             global_pool: ClashMap::with_shard_amount(shards),
     400            1 :             global_pool_size: AtomicUsize::new(0),
     401            1 :             config,
     402            1 :             global_connections_count: Arc::new(AtomicUsize::new(0)),
     403            1 :             _marker: PhantomData,
     404            1 :         })
     405            1 :     }
     406              : 
     407              :     #[cfg(test)]
     408            9 :     pub(crate) fn get_global_connections_count(&self) -> usize {
     409            9 :         self.global_connections_count
     410            9 :             .load(atomic::Ordering::Relaxed)
     411            9 :     }
     412              : 
     413            0 :     pub(crate) fn get_idle_timeout(&self) -> Duration {
     414            0 :         self.config.pool_options.idle_timeout
     415            0 :     }
     416              : 
     417            0 :     pub(crate) fn shutdown(&self) {
     418            0 :         // drops all strong references to endpoint-pools
     419            0 :         self.global_pool.clear();
     420            0 :     }
     421              : 
     422            0 :     pub(crate) async fn gc_worker(&self, mut rng: impl Rng) {
     423            0 :         let epoch = self.config.pool_options.gc_epoch;
     424            0 :         let mut interval = tokio::time::interval(epoch / (self.global_pool.shards().len()) as u32);
     425              :         loop {
     426            0 :             interval.tick().await;
     427              : 
     428            0 :             let shard = rng.gen_range(0..self.global_pool.shards().len());
     429            0 :             self.gc(shard);
     430              :         }
     431              :     }
     432              : 
     433            2 :     pub(crate) fn gc(&self, shard: usize) {
     434            2 :         debug!(shard, "pool: performing epoch reclamation");
     435              : 
     436              :         // acquire a random shard lock
     437            2 :         let mut shard = self.global_pool.shards()[shard].write();
     438            2 : 
     439            2 :         let timer = Metrics::get()
     440            2 :             .proxy
     441            2 :             .http_pool_reclaimation_lag_seconds
     442            2 :             .start_timer();
     443            2 :         let current_len = shard.len();
     444            2 :         let mut clients_removed = 0;
     445            2 :         shard.retain(|(endpoint, x)| {
     446              :             // if the current endpoint pool is unique (no other strong or weak references)
     447              :             // then it is currently not in use by any connections.
     448            2 :             if let Some(pool) = Arc::get_mut(x) {
     449            1 :                 let endpoints = pool.get_mut();
     450            1 :                 clients_removed = endpoints.clear_closed();
     451            1 : 
     452            1 :                 if endpoints.total_conns() == 0 {
     453            0 :                     info!("pool: discarding pool for endpoint {endpoint}");
     454            0 :                     return false;
     455            1 :                 }
     456            1 :             }
     457              : 
     458            2 :             true
     459            2 :         });
     460            2 : 
     461            2 :         let new_len = shard.len();
     462            2 :         drop(shard);
     463            2 :         timer.observe();
     464            2 : 
     465            2 :         // Do logging outside of the lock.
     466            2 :         if clients_removed > 0 {
     467            1 :             let size = self
     468            1 :                 .global_connections_count
     469            1 :                 .fetch_sub(clients_removed, atomic::Ordering::Relaxed)
     470            1 :                 - clients_removed;
     471            1 :             Metrics::get()
     472            1 :                 .proxy
     473            1 :                 .http_pool_opened_connections
     474            1 :                 .get_metric()
     475            1 :                 .dec_by(clients_removed as i64);
     476            1 :             info!("pool: performed global pool gc. removed {clients_removed} clients, total number of clients in pool is {size}");
     477            1 :         }
     478            2 :         let removed = current_len - new_len;
     479            2 : 
     480            2 :         if removed > 0 {
     481            0 :             let global_pool_size = self
     482            0 :                 .global_pool_size
     483            0 :                 .fetch_sub(removed, atomic::Ordering::Relaxed)
     484            0 :                 - removed;
     485            0 :             info!("pool: performed global pool gc. size now {global_pool_size}");
     486            2 :         }
     487            2 :     }
     488              : }
     489              : 
     490              : impl<C: ClientInnerExt> GlobalConnPool<C, EndpointConnPool<C>> {
     491            0 :     pub(crate) fn get(
     492            0 :         self: &Arc<Self>,
     493            0 :         ctx: &RequestContext,
     494            0 :         conn_info: &ConnInfo,
     495            0 :     ) -> Result<Option<Client<C>>, HttpConnError> {
     496            0 :         let mut client: Option<ClientInnerCommon<C>> = None;
     497            0 :         let Some(endpoint) = conn_info.endpoint_cache_key() else {
     498            0 :             return Ok(None);
     499              :         };
     500              : 
     501            0 :         let endpoint_pool = self.get_or_create_endpoint_pool(&endpoint);
     502            0 :         if let Some(entry) = endpoint_pool
     503            0 :             .write()
     504            0 :             .get_conn_entry(conn_info.db_and_user())
     505            0 :         {
     506            0 :             client = Some(entry.conn);
     507            0 :         }
     508            0 :         let endpoint_pool = Arc::downgrade(&endpoint_pool);
     509              : 
     510              :         // ok return cached connection if found and establish a new one otherwise
     511            0 :         if let Some(mut client) = client {
     512            0 :             if client.inner.is_closed() {
     513            0 :                 info!("pool: cached connection '{conn_info}' is closed, opening a new one");
     514            0 :                 return Ok(None);
     515            0 :             }
     516            0 :             tracing::Span::current()
     517            0 :                 .record("conn_id", tracing::field::display(client.get_conn_id()));
     518            0 :             tracing::Span::current().record(
     519            0 :                 "pid",
     520            0 :                 tracing::field::display(client.inner.get_process_id()),
     521            0 :             );
     522            0 :             debug!(
     523            0 :                 cold_start_info = ColdStartInfo::HttpPoolHit.as_str(),
     524            0 :                 "pool: reusing connection '{conn_info}'"
     525              :             );
     526              : 
     527            0 :             match client.get_data() {
     528            0 :                 ClientDataEnum::Local(data) => {
     529            0 :                     data.session().send(ctx.session_id())?;
     530              :                 }
     531              : 
     532            0 :                 ClientDataEnum::Remote(data) => {
     533            0 :                     data.session().send(ctx.session_id())?;
     534              :                 }
     535            0 :                 ClientDataEnum::Http(_) => (),
     536              :             }
     537              : 
     538            0 :             ctx.set_cold_start_info(ColdStartInfo::HttpPoolHit);
     539            0 :             ctx.success();
     540            0 :             return Ok(Some(Client::new(client, conn_info.clone(), endpoint_pool)));
     541            0 :         }
     542            0 :         Ok(None)
     543            0 :     }
     544              : 
     545            2 :     pub(crate) fn get_or_create_endpoint_pool(
     546            2 :         self: &Arc<Self>,
     547            2 :         endpoint: &EndpointCacheKey,
     548            2 :     ) -> Arc<RwLock<EndpointConnPool<C>>> {
     549              :         // fast path
     550            2 :         if let Some(pool) = self.global_pool.get(endpoint) {
     551            0 :             return pool.clone();
     552            2 :         }
     553            2 : 
     554            2 :         // slow path
     555            2 :         let new_pool = Arc::new(RwLock::new(EndpointConnPool {
     556            2 :             pools: HashMap::new(),
     557            2 :             total_conns: 0,
     558            2 :             max_conns: self.config.pool_options.max_conns_per_endpoint,
     559            2 :             _guard: Metrics::get().proxy.http_endpoint_pools.guard(),
     560            2 :             global_connections_count: self.global_connections_count.clone(),
     561            2 :             global_pool_size_max_conns: self.config.pool_options.max_total_conns,
     562            2 :             pool_name: String::from("remote"),
     563            2 :         }));
     564            2 : 
     565            2 :         // find or create a pool for this endpoint
     566            2 :         let mut created = false;
     567            2 :         let pool = self
     568            2 :             .global_pool
     569            2 :             .entry(endpoint.clone())
     570            2 :             .or_insert_with(|| {
     571            2 :                 created = true;
     572            2 :                 new_pool
     573            2 :             })
     574            2 :             .clone();
     575            2 : 
     576            2 :         // log new global pool size
     577            2 :         if created {
     578            2 :             let global_pool_size = self
     579            2 :                 .global_pool_size
     580            2 :                 .fetch_add(1, atomic::Ordering::Relaxed)
     581            2 :                 + 1;
     582            2 :             info!(
     583            0 :                 "pool: created new pool for '{endpoint}', global pool size now {global_pool_size}"
     584              :             );
     585            0 :         }
     586              : 
     587            2 :         pool
     588            2 :     }
     589              : }
     590              : pub(crate) struct Client<C: ClientInnerExt> {
     591              :     span: Span,
     592              :     inner: Option<ClientInnerCommon<C>>,
     593              :     conn_info: ConnInfo,
     594              :     pool: Weak<RwLock<EndpointConnPool<C>>>,
     595              : }
     596              : 
     597              : pub(crate) struct Discard<'a, C: ClientInnerExt> {
     598              :     conn_info: &'a ConnInfo,
     599              :     pool: &'a mut Weak<RwLock<EndpointConnPool<C>>>,
     600              : }
     601              : 
     602              : impl<C: ClientInnerExt> Client<C> {
     603            7 :     pub(crate) fn new(
     604            7 :         inner: ClientInnerCommon<C>,
     605            7 :         conn_info: ConnInfo,
     606            7 :         pool: Weak<RwLock<EndpointConnPool<C>>>,
     607            7 :     ) -> Self {
     608            7 :         Self {
     609            7 :             inner: Some(inner),
     610            7 :             span: Span::current(),
     611            7 :             conn_info,
     612            7 :             pool,
     613            7 :         }
     614            7 :     }
     615              : 
     616            0 :     pub(crate) fn client_inner(&mut self) -> (&mut ClientInnerCommon<C>, Discard<'_, C>) {
     617            0 :         let Self {
     618            0 :             inner,
     619            0 :             pool,
     620            0 :             conn_info,
     621            0 :             span: _,
     622            0 :         } = self;
     623            0 :         let inner_m = inner.as_mut().expect("client inner should not be removed");
     624            0 :         (inner_m, Discard { conn_info, pool })
     625            0 :     }
     626              : 
     627            1 :     pub(crate) fn inner(&mut self) -> (&mut C, Discard<'_, C>) {
     628            1 :         let Self {
     629            1 :             inner,
     630            1 :             pool,
     631            1 :             conn_info,
     632            1 :             span: _,
     633            1 :         } = self;
     634            1 :         let inner = inner.as_mut().expect("client inner should not be removed");
     635            1 :         (&mut inner.inner, Discard { conn_info, pool })
     636            1 :     }
     637              : 
     638            0 :     pub(crate) fn metrics(&self) -> Arc<MetricCounter> {
     639            0 :         let aux = &self
     640            0 :             .inner
     641            0 :             .as_ref()
     642            0 :             .expect("client inner should not be removed")
     643            0 :             .aux;
     644            0 :         USAGE_METRICS.register(Ids {
     645            0 :             endpoint_id: aux.endpoint_id,
     646            0 :             branch_id: aux.branch_id,
     647            0 :         })
     648            0 :     }
     649              : }
     650              : 
     651              : impl<C: ClientInnerExt> Drop for Client<C> {
     652            7 :     fn drop(&mut self) {
     653            7 :         let conn_info = self.conn_info.clone();
     654            7 :         let client = self
     655            7 :             .inner
     656            7 :             .take()
     657            7 :             .expect("client inner should not be removed");
     658            7 :         if let Some(conn_pool) = std::mem::take(&mut self.pool).upgrade() {
     659            6 :             let _current_span = self.span.enter();
     660            6 :             // return connection to the pool
     661            6 :             EndpointConnPool::put(&conn_pool, &conn_info, client);
     662            6 :         }
     663            7 :     }
     664              : }
     665              : 
     666              : impl<C: ClientInnerExt> Deref for Client<C> {
     667              :     type Target = C;
     668              : 
     669            0 :     fn deref(&self) -> &Self::Target {
     670            0 :         &self
     671            0 :             .inner
     672            0 :             .as_ref()
     673            0 :             .expect("client inner should not be removed")
     674            0 :             .inner
     675            0 :     }
     676              : }
     677              : 
     678              : pub(crate) trait ClientInnerExt: Sync + Send + 'static {
     679              :     fn is_closed(&self) -> bool;
     680              :     fn get_process_id(&self) -> i32;
     681              : }
     682              : 
     683              : impl ClientInnerExt for postgres_client::Client {
     684            0 :     fn is_closed(&self) -> bool {
     685            0 :         self.is_closed()
     686            0 :     }
     687              : 
     688            0 :     fn get_process_id(&self) -> i32 {
     689            0 :         self.get_process_id()
     690            0 :     }
     691              : }
     692              : 
     693              : impl<C: ClientInnerExt> Discard<'_, C> {
     694            0 :     pub(crate) fn check_idle(&mut self, status: ReadyForQueryStatus) {
     695            0 :         let conn_info = &self.conn_info;
     696            0 :         if status != ReadyForQueryStatus::Idle && std::mem::take(self.pool).strong_count() > 0 {
     697            0 :             info!("pool: throwing away connection '{conn_info}' because connection is not idle");
     698            0 :         }
     699            0 :     }
     700            1 :     pub(crate) fn discard(&mut self) {
     701            1 :         let conn_info = &self.conn_info;
     702            1 :         if std::mem::take(self.pool).strong_count() > 0 {
     703            1 :             info!("pool: throwing away connection '{conn_info}' because connection is potentially in a broken state");
     704            0 :         }
     705            1 :     }
     706              : }
        

Generated by: LCOV version 2.1-beta