Skip to content

Commit 08fecec

Browse files
itowlsonitowlson
committed
Pool Postgres connections
Signed-off-by: itowlson <[email protected]>
1 parent 7417676 commit 08fecec

File tree

6 files changed

+154
-70
lines changed

6 files changed

+154
-70
lines changed

‎Cargo.lock‎

Lines changed: 35 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎crates/factor-outbound-pg/Cargo.toml‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ edition = { workspace = true }
77
[dependencies]
88
anyhow = { workspace = true }
99
chrono = { workspace = true }
10+
deadpool-postgres = { version = "0.14", features = ["rt_tokio_1"] }
1011
native-tls = "0.2"
1112
postgres-native-tls = "0.5"
1213
spin-core = { path = "../core" }

‎crates/factor-outbound-pg/src/client.rs‎

Lines changed: 76 additions & 41 deletions
Original file line numberDiff line numberDiff line change
@@ -1,20 +1,85 @@
1-
use anyhow::{anyhow, Result};
1+
use anyhow::{anyhow, Context, Result};
22
use native_tls::TlsConnector;
33
use postgres_native_tls::MakeTlsConnector;
44
use spin_world::async_trait;
55
use spin_world::spin::postgres::postgres::{
66
self as v3, Column, DbDataType, DbValue, ParameterValue, RowSet,
77
};
88
use tokio_postgres::types::Type;
9-
use tokio_postgres::{config::SslMode, types::ToSql, Row};
10-
use tokio_postgres::{Client as TokioClient, NoTls, Socket};
9+
use tokio_postgres::{config::SslMode, types::ToSql, NoTls, Row};
10+
11+
const CONNECTION_POOL_SIZE: usize = 64;
1112

1213
#[async_trait]
13-
pub trait Client {
14-
async fn build_client(address: &str) -> Result<Self>
15-
where
16-
Self: Sized;
14+
pub trait ClientFactory: Send + Sync {
15+
type Client: Client + Send + Sync + 'static;
16+
fn new() -> Self;
17+
async fn build_client(&mut self, address: &str) -> Result<Self::Client>;
18+
}
19+
20+
pub struct PooledTokioClientFactory {
21+
pools: std::collections::HashMap<String, deadpool_postgres::Pool>,
22+
}
23+
24+
#[async_trait]
25+
impl ClientFactory for PooledTokioClientFactory {
26+
type Client = deadpool_postgres::Object;
27+
28+
fn new() -> Self {
29+
Self {
30+
pools: Default::default(),
31+
}
32+
}
33+
34+
async fn build_client(&mut self, address: &str) -> Result<Self::Client> {
35+
let pool_entry = self.pools.entry(address.to_owned());
36+
let pool = match pool_entry {
37+
std::collections::hash_map::Entry::Occupied(entry) => entry.into_mut(),
38+
std::collections::hash_map::Entry::Vacant(entry) => {
39+
let pool = create_connection_pool(address)
40+
.context("establishing PostgreSQL connection pool")?;
41+
entry.insert(pool)
42+
}
43+
};
44+
45+
Ok(pool.get().await?)
46+
}
47+
}
48+
49+
fn create_connection_pool(address: &str) -> Result<deadpool_postgres::Pool> {
50+
let config = address
51+
.parse::<tokio_postgres::Config>()
52+
.context("parsing Postgres connection string")?;
53+
54+
tracing::debug!("Build new connection: {}", address);
1755

56+
// TODO: This is slower but safer. Is it the right tradeoff?
57+
// https://docs.rs/deadpool-postgres/latest/deadpool_postgres/enum.RecyclingMethod.html
58+
let mgr_config = deadpool_postgres::ManagerConfig {
59+
recycling_method: deadpool_postgres::RecyclingMethod::Clean,
60+
};
61+
62+
let mgr = if config.get_ssl_mode() == SslMode::Disable {
63+
deadpool_postgres::Manager::from_config(config, NoTls, mgr_config)
64+
} else {
65+
let builder = TlsConnector::builder();
66+
let connector = MakeTlsConnector::new(builder.build()?);
67+
deadpool_postgres::Manager::from_config(config, connector, mgr_config)
68+
};
69+
70+
// TODO: what is our max size heuristic? Should this be passed in soe that different
71+
// hosts can manage it according to their needs? Will a plain number suffice for
72+
// sophisticated hosts anyway?
73+
let pool = deadpool_postgres::Pool::builder(mgr)
74+
.max_size(CONNECTION_POOL_SIZE)
75+
.build()
76+
.context("building Postgres connection pool")?;
77+
78+
Ok(pool)
79+
}
80+
81+
#[async_trait]
82+
pub trait Client {
1883
async fn execute(
1984
&self,
2085
statement: String,
@@ -29,28 +94,7 @@ pub trait Client {
2994
}
3095

3196
#[async_trait]
32-
impl Client for TokioClient {
33-
async fn build_client(address: &str) -> Result<Self>
34-
where
35-
Self: Sized,
36-
{
37-
let config = address.parse::<tokio_postgres::Config>()?;
38-
39-
tracing::debug!("Build new connection: {}", address);
40-
41-
if config.get_ssl_mode() == SslMode::Disable {
42-
let (client, connection) = config.connect(NoTls).await?;
43-
spawn_connection(connection);
44-
Ok(client)
45-
} else {
46-
let builder = TlsConnector::builder();
47-
let connector = MakeTlsConnector::new(builder.build()?);
48-
let (client, connection) = config.connect(connector).await?;
49-
spawn_connection(connection);
50-
Ok(client)
51-
}
52-
}
53-
97+
impl Client for deadpool_postgres::Object {
5498
async fn execute(
5599
&self,
56100
statement: String,
@@ -67,7 +111,8 @@ impl Client for TokioClient {
67111
.map(|b| b.as_ref() as &(dyn ToSql + Sync))
68112
.collect();
69113

70-
self.execute(&statement, params_refs.as_slice())
114+
self.as_ref()
115+
.execute(&statement, params_refs.as_slice())
71116
.await
72117
.map_err(|e| v3::Error::QueryFailed(format!("{e:?}")))
73118
}
@@ -89,6 +134,7 @@ impl Client for TokioClient {
89134
.collect();
90135

91136
let results = self
137+
.as_ref()
92138
.query(&statement, params_refs.as_slice())
93139
.await
94140
.map_err(|e| v3::Error::QueryFailed(format!("{e:?}")))?;
@@ -111,17 +157,6 @@ impl Client for TokioClient {
111157
}
112158
}
113159

114-
fn spawn_connection<T>(connection: tokio_postgres::Connection<Socket, T>)
115-
where
116-
T: tokio_postgres::tls::TlsStream + std::marker::Unpin + std::marker::Send + 'static,
117-
{
118-
tokio::spawn(async move {
119-
if let Err(e) = connection.await {
120-
tracing::error!("Postgres connection error: {}", e);
121-
}
122-
});
123-
}
124-
125160
fn to_sql_parameter(value: &ParameterValue) -> Result<Box<dyn ToSql + Send + Sync>> {
126161
match value {
127162
ParameterValue::Boolean(v) => Ok(Box::new(*v)),

‎crates/factor-outbound-pg/src/host.rs‎

Lines changed: 14 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -9,17 +9,20 @@ use tracing::field::Empty;
99
use tracing::instrument;
1010
use tracing::Level;
1111

12-
use crate::client::Client;
12+
use crate::client::{Client, ClientFactory};
1313
use crate::InstanceState;
1414

15-
impl<C: Client> InstanceState<C> {
15+
impl<CF: ClientFactory> InstanceState<CF> {
1616
async fn open_connection<Conn: 'static>(
1717
&mut self,
1818
address: &str,
1919
) -> Result<Resource<Conn>, v3::Error> {
2020
self.connections
2121
.push(
22-
C::build_client(address)
22+
self.client_factory
23+
.write()
24+
.await
25+
.build_client(address)
2326
.await
2427
.map_err(|e| v3::Error::ConnectionFailed(format!("{e:?}")))?,
2528
)
@@ -30,7 +33,7 @@ impl<C: Client> InstanceState<C> {
3033
async fn get_client<Conn: 'static>(
3134
&mut self,
3235
connection: Resource<Conn>,
33-
) -> Result<&C, v3::Error> {
36+
) -> Result<&CF::Client, v3::Error> {
3437
self.connections
3538
.get(connection.rep())
3639
.ok_or_else(|| v3::Error::ConnectionFailed("no connection found".into()))
@@ -71,8 +74,8 @@ fn v2_params_to_v3(
7174
params.into_iter().map(|p| p.try_into()).collect()
7275
}
7376

74-
impl<C: Send + Sync + Client> spin_world::spin::postgres::postgres::HostConnection
75-
for InstanceState<C>
77+
impl<CF: ClientFactory + Send + Sync> spin_world::spin::postgres::postgres::HostConnection
78+
for InstanceState<CF>
7679
{
7780
#[instrument(name = "spin_outbound_pg.open", skip(self, address), err(level = Level::INFO), fields(otel.kind = "client", db.system = "postgresql", db.address = Empty, server.port = Empty, db.namespace = Empty))]
7881
async fn open(&mut self, address: String) -> Result<Resource<v3::Connection>, v3::Error> {
@@ -122,13 +125,13 @@ impl<C: Send + Sync + Client> spin_world::spin::postgres::postgres::HostConnecti
122125
}
123126
}
124127

125-
impl<C: Send> v2_types::Host for InstanceState<C> {
128+
impl<CF: ClientFactory + Send> v2_types::Host for InstanceState<CF> {
126129
fn convert_error(&mut self, error: v2::Error) -> Result<v2::Error> {
127130
Ok(error)
128131
}
129132
}
130133

131-
impl<C: Send + Sync + Client> v3::Host for InstanceState<C> {
134+
impl<CF: Send + Sync + ClientFactory> v3::Host for InstanceState<CF> {
132135
fn convert_error(&mut self, error: v3::Error) -> Result<v3::Error> {
133136
Ok(error)
134137
}
@@ -152,9 +155,9 @@ macro_rules! delegate {
152155
}};
153156
}
154157

155-
impl<C: Send + Sync + Client> v2::Host for InstanceState<C> {}
158+
impl<CF: Send + Sync + ClientFactory> v2::Host for InstanceState<CF> {}
156159

157-
impl<C: Send + Sync + Client> v2::HostConnection for InstanceState<C> {
160+
impl<CF: Send + Sync + ClientFactory> v2::HostConnection for InstanceState<CF> {
158161
#[instrument(name = "spin_outbound_pg.open", skip(self, address), err(level = Level::INFO), fields(otel.kind = "client", db.system = "postgresql", db.address = Empty, server.port = Empty, db.namespace = Empty))]
159162
async fn open(&mut self, address: String) -> Result<Resource<v2::Connection>, v2::Error> {
160163
spin_factor_outbound_networking::record_address_fields(&address);
@@ -206,7 +209,7 @@ impl<C: Send + Sync + Client> v2::HostConnection for InstanceState<C> {
206209
}
207210
}
208211

209-
impl<C: Send + Sync + Client> v1::Host for InstanceState<C> {
212+
impl<CF: Send + Sync + ClientFactory> v1::Host for InstanceState<CF> {
210213
async fn execute(
211214
&mut self,
212215
address: String,

‎crates/factor-outbound-pg/src/lib.rs‎

Lines changed: 15 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -1,24 +1,26 @@
11
pub mod client;
22
mod host;
33

4-
use client::Client;
4+
use std::sync::Arc;
5+
6+
use client::ClientFactory;
57
use spin_factor_outbound_networking::{
68
config::allowed_hosts::OutboundAllowedHosts, OutboundNetworkingFactor,
79
};
810
use spin_factors::{
911
anyhow, ConfigureAppContext, Factor, FactorData, PrepareContext, RuntimeFactors,
1012
SelfInstanceBuilder,
1113
};
12-
use tokio_postgres::Client as PgClient;
14+
use tokio::sync::RwLock;
1315

14-
pub struct OutboundPgFactor<C = PgClient> {
15-
_phantom: std::marker::PhantomData<C>,
16+
pub struct OutboundPgFactor<CF = crate::client::PooledTokioClientFactory> {
17+
_phantom: std::marker::PhantomData<CF>,
1618
}
1719

18-
impl<C: Send + Sync + Client + 'static> Factor for OutboundPgFactor<C> {
20+
impl<CF: ClientFactory + Send + Sync + 'static> Factor for OutboundPgFactor<CF> {
1921
type RuntimeConfig = ();
20-
type AppState = ();
21-
type InstanceBuilder = InstanceState<C>;
22+
type AppState = Arc<RwLock<CF>>;
23+
type InstanceBuilder = InstanceState<CF>;
2224

2325
fn init(&mut self, ctx: &mut impl spin_factors::InitContext<Self>) -> anyhow::Result<()> {
2426
ctx.link_bindings(spin_world::v1::postgres::add_to_linker::<_, FactorData<Self>>)?;
@@ -33,7 +35,7 @@ impl<C: Send + Sync + Client + 'static> Factor for OutboundPgFactor<C> {
3335
&self,
3436
_ctx: ConfigureAppContext<T, Self>,
3537
) -> anyhow::Result<Self::AppState> {
36-
Ok(())
38+
Ok(Arc::new(RwLock::new(CF::new())))
3739
}
3840

3941
fn prepare<T: RuntimeFactors>(
@@ -45,6 +47,7 @@ impl<C: Send + Sync + Client + 'static> Factor for OutboundPgFactor<C> {
4547
.allowed_hosts();
4648
Ok(InstanceState {
4749
allowed_hosts,
50+
client_factory: ctx.app_state().clone(),
4851
connections: Default::default(),
4952
})
5053
}
@@ -64,9 +67,10 @@ impl<C> OutboundPgFactor<C> {
6467
}
6568
}
6669

67-
pub struct InstanceState<C> {
70+
pub struct InstanceState<CF: ClientFactory> {
6871
allowed_hosts: OutboundAllowedHosts,
69-
connections: spin_resource_table::Table<C>,
72+
client_factory: Arc<RwLock<CF>>,
73+
connections: spin_resource_table::Table<CF::Client>,
7074
}
7175

72-
impl<C: Send + 'static> SelfInstanceBuilder for InstanceState<C> {}
76+
impl<CF: ClientFactory + Send + 'static> SelfInstanceBuilder for InstanceState<CF> {}

0 commit comments

Comments
 (0)