Line data Source code
1 : use std::fmt;
2 : use std::pin::pin;
3 : use std::sync::{Arc, Weak};
4 : use std::task::{Poll, ready};
5 :
6 : use futures::future::poll_fn;
7 : use futures::{Future, FutureExt};
8 : use postgres_client::tls::MakeTlsConnect;
9 : use smallvec::SmallVec;
10 : use tokio::net::TcpStream;
11 : use tokio::time::Instant;
12 : use tokio_util::sync::CancellationToken;
13 : use tracing::{error, info, info_span};
14 : #[cfg(test)]
15 : use {
16 : super::conn_pool_lib::GlobalConnPoolOptions,
17 : crate::auth::backend::ComputeUserInfo,
18 : std::{sync::atomic, time::Duration},
19 : };
20 :
21 : use super::conn_pool_lib::{
22 : Client, ClientDataEnum, ClientInnerCommon, ClientInnerExt, ConnInfo, EndpointConnPool,
23 : GlobalConnPool,
24 : };
25 : use crate::config::ComputeConfig;
26 : use crate::context::RequestContext;
27 : use crate::control_plane::messages::MetricsAuxInfo;
28 : use crate::metrics::Metrics;
29 :
30 : type TlsStream = <ComputeConfig as MakeTlsConnect<TcpStream>>::Stream;
31 :
32 : #[derive(Debug, Clone)]
33 : pub(crate) struct ConnInfoWithAuth {
34 : pub(crate) conn_info: ConnInfo,
35 : pub(crate) auth: AuthData,
36 : }
37 :
38 : #[derive(Debug, Clone)]
39 : pub(crate) enum AuthData {
40 : Password(SmallVec<[u8; 16]>),
41 : Jwt(String),
42 : }
43 :
44 : impl fmt::Display for ConnInfo {
45 : // use custom display to avoid logging password
46 0 : fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
47 0 : write!(
48 0 : f,
49 0 : "{}@{}/{}?{}",
50 : self.user_info.user,
51 : self.user_info.endpoint,
52 : self.dbname,
53 0 : self.user_info.options.get_cache_key("")
54 : )
55 0 : }
56 : }
57 :
58 0 : pub(crate) fn poll_client<C: ClientInnerExt>(
59 0 : global_pool: Arc<GlobalConnPool<C, EndpointConnPool<C>>>,
60 0 : ctx: &RequestContext,
61 0 : conn_info: ConnInfo,
62 0 : client: C,
63 0 : mut connection: postgres_client::Connection<TcpStream, TlsStream>,
64 0 : conn_id: uuid::Uuid,
65 0 : aux: MetricsAuxInfo,
66 0 : ) -> Client<C> {
67 0 : let conn_gauge = Metrics::get().proxy.db_connections.guard(ctx.protocol());
68 0 : let mut session_id = ctx.session_id();
69 0 : let (tx, mut rx) = tokio::sync::watch::channel(session_id);
70 :
71 0 : let span = info_span!(parent: None, "connection", %conn_id);
72 0 : let cold_start_info = ctx.cold_start_info();
73 0 : span.in_scope(|| {
74 0 : info!(cold_start_info = cold_start_info.as_str(), %conn_info, %session_id, "new connection");
75 0 : });
76 0 : let pool = match conn_info.endpoint_cache_key() {
77 0 : Some(endpoint) => Arc::downgrade(&global_pool.get_or_create_endpoint_pool(&endpoint)),
78 0 : None => Weak::new(),
79 : };
80 0 : let pool_clone = pool.clone();
81 :
82 0 : let db_user = conn_info.db_and_user();
83 0 : let idle = global_pool.get_idle_timeout();
84 0 : let cancel = CancellationToken::new();
85 0 : let cancelled = cancel.clone().cancelled_owned();
86 :
87 0 : tokio::spawn(async move {
88 0 : let _conn_gauge = conn_gauge;
89 0 : let mut idle_timeout = pin!(tokio::time::sleep(idle));
90 0 : let mut cancelled = pin!(cancelled);
91 :
92 0 : poll_fn(move |cx| {
93 0 : let _instrument = span.enter();
94 :
95 0 : if cancelled.as_mut().poll(cx).is_ready() {
96 0 : info!("connection dropped");
97 0 : return Poll::Ready(());
98 0 : }
99 :
100 0 : match rx.has_changed() {
101 : Ok(true) => {
102 0 : session_id = *rx.borrow_and_update();
103 0 : info!(%session_id, "changed session");
104 0 : idle_timeout.as_mut().reset(Instant::now() + idle);
105 : }
106 : Err(_) => {
107 0 : info!("connection dropped");
108 0 : return Poll::Ready(());
109 : }
110 0 : _ => {}
111 : }
112 :
113 : // 5 minute idle connection timeout
114 0 : if idle_timeout.as_mut().poll(cx).is_ready() {
115 0 : idle_timeout.as_mut().reset(Instant::now() + idle);
116 0 : info!("connection idle");
117 0 : if let Some(pool) = pool.clone().upgrade() {
118 : // remove client from pool - should close the connection if it's idle.
119 : // does nothing if the client is currently checked-out and in-use
120 0 : if pool.write().remove_client(db_user.clone(), conn_id) {
121 0 : info!("idle connection removed");
122 0 : }
123 0 : }
124 0 : }
125 :
126 0 : match ready!(connection.poll_unpin(cx)) {
127 0 : Err(e) => error!(%session_id, "connection error: {}", e),
128 0 : Ok(()) => info!("connection closed"),
129 : }
130 :
131 : // remove from connection pool
132 0 : if let Some(pool) = pool.clone().upgrade()
133 0 : && pool.write().remove_client(db_user.clone(), conn_id)
134 : {
135 0 : info!("closed connection removed");
136 0 : }
137 :
138 0 : Poll::Ready(())
139 0 : })
140 0 : .await;
141 0 : });
142 0 : let inner = ClientInnerCommon {
143 0 : inner: client,
144 0 : aux,
145 0 : conn_id,
146 0 : data: ClientDataEnum::Remote(ClientDataRemote {
147 0 : session: tx,
148 0 : cancel,
149 0 : }),
150 0 : };
151 :
152 0 : Client::new(inner, conn_info, pool_clone)
153 0 : }
154 :
155 : #[derive(Clone)]
156 : pub(crate) struct ClientDataRemote {
157 : session: tokio::sync::watch::Sender<uuid::Uuid>,
158 : cancel: CancellationToken,
159 : }
160 :
161 : impl ClientDataRemote {
162 0 : pub fn session(&mut self) -> &mut tokio::sync::watch::Sender<uuid::Uuid> {
163 0 : &mut self.session
164 0 : }
165 :
166 7 : pub fn cancel(&mut self) {
167 7 : self.cancel.cancel();
168 7 : }
169 : }
170 :
171 : #[cfg(test)]
172 : mod tests {
173 : use std::sync::atomic::AtomicBool;
174 :
175 : use super::*;
176 : use crate::proxy::NeonOptions;
177 : use crate::serverless::cancel_set::CancelSet;
178 : use crate::types::{BranchId, EndpointId, ProjectId};
179 :
180 : struct MockClient(Arc<AtomicBool>);
181 : impl MockClient {
182 6 : fn new(is_closed: bool) -> Self {
183 6 : MockClient(Arc::new(is_closed.into()))
184 6 : }
185 : }
186 : impl ClientInnerExt for MockClient {
187 8 : fn is_closed(&self) -> bool {
188 8 : self.0.load(atomic::Ordering::Relaxed)
189 8 : }
190 0 : fn get_process_id(&self) -> i32 {
191 0 : 0
192 0 : }
193 : }
194 :
195 5 : fn create_inner() -> ClientInnerCommon<MockClient> {
196 5 : create_inner_with(MockClient::new(false))
197 5 : }
198 :
199 7 : fn create_inner_with(client: MockClient) -> ClientInnerCommon<MockClient> {
200 7 : ClientInnerCommon {
201 7 : inner: client,
202 7 : aux: MetricsAuxInfo {
203 7 : endpoint_id: (&EndpointId::from("endpoint")).into(),
204 7 : project_id: (&ProjectId::from("project")).into(),
205 7 : branch_id: (&BranchId::from("branch")).into(),
206 7 : compute_id: "compute".into(),
207 7 : cold_start_info: crate::control_plane::messages::ColdStartInfo::Warm,
208 7 : },
209 7 : conn_id: uuid::Uuid::new_v4(),
210 7 : data: ClientDataEnum::Remote(ClientDataRemote {
211 7 : session: tokio::sync::watch::Sender::new(uuid::Uuid::new_v4()),
212 7 : cancel: CancellationToken::new(),
213 7 : }),
214 7 : }
215 7 : }
216 :
217 : #[tokio::test]
218 1 : async fn test_pool() {
219 1 : let _ = env_logger::try_init();
220 1 : let config = Box::leak(Box::new(crate::config::HttpConfig {
221 1 : accept_websockets: false,
222 1 : pool_options: GlobalConnPoolOptions {
223 1 : max_conns_per_endpoint: 2,
224 1 : gc_epoch: Duration::from_secs(1),
225 1 : pool_shards: 2,
226 1 : idle_timeout: Duration::from_secs(1),
227 1 : opt_in: false,
228 1 : max_total_conns: 3,
229 1 : },
230 1 : cancel_set: CancelSet::new(0),
231 1 : client_conn_threshold: u64::MAX,
232 1 : max_request_size_bytes: usize::MAX,
233 1 : max_response_size_bytes: usize::MAX,
234 1 : }));
235 1 : let pool = GlobalConnPool::new(config);
236 1 : let conn_info = ConnInfo {
237 1 : user_info: ComputeUserInfo {
238 1 : user: "user".into(),
239 1 : endpoint: "endpoint".into(),
240 1 : options: NeonOptions::default(),
241 1 : },
242 1 : dbname: "dbname".into(),
243 1 : };
244 1 : let ep_pool = Arc::downgrade(
245 1 : &pool.get_or_create_endpoint_pool(&conn_info.endpoint_cache_key().unwrap()),
246 : );
247 : {
248 1 : let mut client = Client::new(create_inner(), conn_info.clone(), ep_pool.clone());
249 1 : assert_eq!(0, pool.get_global_connections_count());
250 1 : client.inner().1.discard();
251 : // Discard should not add the connection from the pool.
252 1 : assert_eq!(0, pool.get_global_connections_count());
253 : }
254 : {
255 1 : let client = Client::new(create_inner(), conn_info.clone(), ep_pool.clone());
256 1 : drop(client);
257 1 : assert_eq!(1, pool.get_global_connections_count());
258 : }
259 : {
260 1 : let closed_client = Client::new(
261 1 : create_inner_with(MockClient::new(true)),
262 1 : conn_info.clone(),
263 1 : ep_pool.clone(),
264 : );
265 1 : drop(closed_client);
266 1 : assert_eq!(1, pool.get_global_connections_count());
267 : }
268 1 : let is_closed: Arc<AtomicBool> = Arc::new(false.into());
269 : {
270 1 : let client = Client::new(
271 1 : create_inner_with(MockClient(is_closed.clone())),
272 1 : conn_info.clone(),
273 1 : ep_pool.clone(),
274 : );
275 1 : drop(client);
276 : // The client should be added to the pool.
277 1 : assert_eq!(2, pool.get_global_connections_count());
278 : }
279 : {
280 1 : let client = Client::new(create_inner(), conn_info, ep_pool);
281 1 : drop(client);
282 :
283 : // The client shouldn't be added to the pool. Because the ep-pool is full.
284 1 : assert_eq!(2, pool.get_global_connections_count());
285 : }
286 :
287 1 : let conn_info = ConnInfo {
288 1 : user_info: ComputeUserInfo {
289 1 : user: "user".into(),
290 1 : endpoint: "endpoint-2".into(),
291 1 : options: NeonOptions::default(),
292 1 : },
293 1 : dbname: "dbname".into(),
294 1 : };
295 1 : let ep_pool = Arc::downgrade(
296 1 : &pool.get_or_create_endpoint_pool(&conn_info.endpoint_cache_key().unwrap()),
297 : );
298 : {
299 1 : let client = Client::new(create_inner(), conn_info.clone(), ep_pool.clone());
300 1 : drop(client);
301 1 : assert_eq!(3, pool.get_global_connections_count());
302 : }
303 : {
304 1 : let client = Client::new(create_inner(), conn_info.clone(), ep_pool.clone());
305 1 : drop(client);
306 :
307 : // The client shouldn't be added to the pool. Because the global pool is full.
308 1 : assert_eq!(3, pool.get_global_connections_count());
309 : }
310 :
311 1 : is_closed.store(true, atomic::Ordering::Relaxed);
312 : // Do gc for all shards.
313 1 : pool.gc(0);
314 1 : pool.gc(1);
315 : // Closed client should be removed from the pool.
316 1 : assert_eq!(2, pool.get_global_connections_count());
317 1 : }
318 : }
|
