Line data Source code
1 : use std::collections::VecDeque;
2 : use std::sync::atomic::{self, AtomicUsize};
3 : use std::sync::{Arc, Weak};
4 :
5 : use bytes::Bytes;
6 : use http_body_util::combinators::BoxBody;
7 : use hyper::client::conn::http2;
8 : use hyper_util::rt::{TokioExecutor, TokioIo};
9 : use parking_lot::RwLock;
10 : use smol_str::ToSmolStr;
11 : use tracing::{Instrument, debug, error, info, info_span};
12 :
13 : use super::AsyncRW;
14 : use super::backend::HttpConnError;
15 : use super::conn_pool_lib::{
16 : ClientDataEnum, ClientInnerCommon, ClientInnerExt, ConnInfo, ConnPoolEntry,
17 : EndpointConnPoolExt, GlobalConnPool,
18 : };
19 : use crate::context::RequestContext;
20 : use crate::control_plane::messages::{ColdStartInfo, MetricsAuxInfo};
21 : use crate::metrics::{HttpEndpointPoolsGuard, Metrics};
22 : use crate::protocol2::ConnectionInfoExtra;
23 : use crate::types::EndpointCacheKey;
24 : use crate::usage_metrics::{Ids, MetricCounter, USAGE_METRICS};
25 :
26 : pub(crate) type Send = http2::SendRequest<BoxBody<Bytes, hyper::Error>>;
27 : pub(crate) type Connect =
28 : http2::Connection<TokioIo<AsyncRW>, BoxBody<Bytes, hyper::Error>, TokioExecutor>;
29 :
30 : #[derive(Clone)]
31 : pub(crate) struct ClientDataHttp();
32 :
33 : // Per-endpoint connection pool
34 : // Number of open connections is limited by the `max_conns_per_endpoint`.
35 : pub(crate) struct HttpConnPool<C: ClientInnerExt + Clone> {
36 : // TODO(conrad):
37 : // either we should open more connections depending on stream count
38 : // (not exposed by hyper, need our own counter)
39 : // or we can change this to an Option rather than a VecDeque.
40 : //
41 : // Opening more connections to the same db because we run out of streams
42 : // seems somewhat redundant though.
43 : //
44 : // Probably we should run a semaphore and just the single conn. TBD.
45 : conns: VecDeque<ConnPoolEntry<C>>,
46 : _guard: HttpEndpointPoolsGuard<'static>,
47 : global_connections_count: Arc<AtomicUsize>,
48 : }
49 :
50 : impl<C: ClientInnerExt + Clone> HttpConnPool<C> {
51 0 : fn get_conn_entry(&mut self) -> Option<ConnPoolEntry<C>> {
52 0 : let Self { conns, .. } = self;
53 :
54 : loop {
55 0 : let conn = conns.pop_front()?;
56 0 : if !conn.conn.inner.is_closed() {
57 0 : let new_conn = ConnPoolEntry {
58 0 : conn: conn.conn.clone(),
59 0 : _last_access: std::time::Instant::now(),
60 0 : };
61 :
62 0 : conns.push_back(new_conn);
63 0 : return Some(conn);
64 0 : }
65 : }
66 0 : }
67 :
68 0 : fn remove_conn(&mut self, conn_id: uuid::Uuid) -> bool {
69 : let Self {
70 0 : conns,
71 0 : global_connections_count,
72 : ..
73 0 : } = self;
74 :
75 0 : let old_len = conns.len();
76 0 : conns.retain(|entry| entry.conn.conn_id != conn_id);
77 0 : let new_len = conns.len();
78 0 : let removed = old_len - new_len;
79 0 : if removed > 0 {
80 0 : global_connections_count.fetch_sub(removed, atomic::Ordering::Relaxed);
81 0 : Metrics::get()
82 0 : .proxy
83 0 : .http_pool_opened_connections
84 0 : .get_metric()
85 0 : .dec_by(removed as i64);
86 0 : }
87 0 : removed > 0
88 0 : }
89 : }
90 :
91 : impl<C: ClientInnerExt + Clone> EndpointConnPoolExt<C> for HttpConnPool<C> {
92 0 : fn clear_closed(&mut self) -> usize {
93 0 : let Self { conns, .. } = self;
94 0 : let old_len = conns.len();
95 0 : conns.retain(|entry| !entry.conn.inner.is_closed());
96 :
97 0 : let new_len = conns.len();
98 0 : old_len - new_len
99 0 : }
100 :
101 0 : fn total_conns(&self) -> usize {
102 0 : self.conns.len()
103 0 : }
104 : }
105 :
106 : impl<C: ClientInnerExt + Clone> Drop for HttpConnPool<C> {
107 0 : fn drop(&mut self) {
108 0 : if !self.conns.is_empty() {
109 0 : self.global_connections_count
110 0 : .fetch_sub(self.conns.len(), atomic::Ordering::Relaxed);
111 0 : Metrics::get()
112 0 : .proxy
113 0 : .http_pool_opened_connections
114 0 : .get_metric()
115 0 : .dec_by(self.conns.len() as i64);
116 0 : }
117 0 : }
118 : }
119 :
120 : impl<C: ClientInnerExt + Clone> GlobalConnPool<C, HttpConnPool<C>> {
121 : #[expect(unused_results)]
122 0 : pub(crate) fn get(
123 0 : self: &Arc<Self>,
124 0 : ctx: &RequestContext,
125 0 : conn_info: &ConnInfo,
126 0 : ) -> Result<Option<Client<C>>, HttpConnError> {
127 : let result: Result<Option<Client<C>>, HttpConnError>;
128 0 : let Some(endpoint) = conn_info.endpoint_cache_key() else {
129 0 : result = Ok(None);
130 0 : return result;
131 : };
132 0 : let endpoint_pool = self.get_or_create_endpoint_pool(&endpoint);
133 0 : let Some(client) = endpoint_pool.write().get_conn_entry() else {
134 0 : result = Ok(None);
135 0 : return result;
136 : };
137 :
138 0 : tracing::Span::current().record("conn_id", tracing::field::display(client.conn.conn_id));
139 0 : debug!(
140 0 : cold_start_info = ColdStartInfo::HttpPoolHit.as_str(),
141 0 : "pool: reusing connection '{conn_info}'"
142 : );
143 0 : ctx.set_cold_start_info(ColdStartInfo::HttpPoolHit);
144 0 : ctx.success();
145 :
146 0 : Ok(Some(Client::new(client.conn.clone())))
147 0 : }
148 :
149 0 : fn get_or_create_endpoint_pool(
150 0 : self: &Arc<Self>,
151 0 : endpoint: &EndpointCacheKey,
152 0 : ) -> Arc<RwLock<HttpConnPool<C>>> {
153 : // fast path
154 0 : if let Some(pool) = self.global_pool.get(endpoint) {
155 0 : return pool.clone();
156 0 : }
157 :
158 : // slow path
159 0 : let new_pool = Arc::new(RwLock::new(HttpConnPool {
160 0 : conns: VecDeque::new(),
161 0 : _guard: Metrics::get().proxy.http_endpoint_pools.guard(),
162 0 : global_connections_count: self.global_connections_count.clone(),
163 0 : }));
164 :
165 : // find or create a pool for this endpoint
166 0 : let mut created = false;
167 0 : let pool = self
168 0 : .global_pool
169 0 : .entry(endpoint.clone())
170 0 : .or_insert_with(|| {
171 0 : created = true;
172 0 : new_pool
173 0 : })
174 0 : .clone();
175 :
176 : // log new global pool size
177 0 : if created {
178 0 : let global_pool_size = self
179 0 : .global_pool_size
180 0 : .fetch_add(1, atomic::Ordering::Relaxed)
181 0 : + 1;
182 0 : info!(
183 0 : "pool: created new pool for '{endpoint}', global pool size now {global_pool_size}"
184 : );
185 0 : }
186 :
187 0 : pool
188 0 : }
189 : }
190 :
191 0 : pub(crate) fn poll_http2_client(
192 0 : global_pool: Arc<GlobalConnPool<Send, HttpConnPool<Send>>>,
193 0 : ctx: &RequestContext,
194 0 : conn_info: &ConnInfo,
195 0 : client: Send,
196 0 : connection: Connect,
197 0 : conn_id: uuid::Uuid,
198 0 : aux: MetricsAuxInfo,
199 0 : ) -> Client<Send> {
200 0 : let conn_gauge = Metrics::get().proxy.db_connections.guard(ctx.protocol());
201 0 : let session_id = ctx.session_id();
202 :
203 0 : let span = info_span!(parent: None, "connection", %conn_id);
204 0 : let cold_start_info = ctx.cold_start_info();
205 0 : span.in_scope(|| {
206 0 : info!(cold_start_info = cold_start_info.as_str(), %conn_info, %session_id, "new connection");
207 0 : });
208 :
209 0 : let pool = match conn_info.endpoint_cache_key() {
210 0 : Some(endpoint) => {
211 0 : let pool = global_pool.get_or_create_endpoint_pool(&endpoint);
212 0 : let client = ClientInnerCommon {
213 0 : inner: client.clone(),
214 0 : aux: aux.clone(),
215 0 : conn_id,
216 0 : data: ClientDataEnum::Http(ClientDataHttp()),
217 0 : };
218 0 : pool.write().conns.push_back(ConnPoolEntry {
219 0 : conn: client,
220 0 : _last_access: std::time::Instant::now(),
221 0 : });
222 0 : Metrics::get()
223 0 : .proxy
224 0 : .http_pool_opened_connections
225 0 : .get_metric()
226 0 : .inc();
227 :
228 0 : Arc::downgrade(&pool)
229 : }
230 0 : None => Weak::new(),
231 : };
232 :
233 0 : tokio::spawn(
234 0 : async move {
235 0 : let _conn_gauge = conn_gauge;
236 0 : let res = connection.await;
237 0 : match res {
238 0 : Ok(()) => info!("connection closed"),
239 0 : Err(e) => error!(%session_id, "connection error: {e:?}"),
240 : }
241 :
242 : // remove from connection pool
243 0 : if let Some(pool) = pool.clone().upgrade()
244 0 : && pool.write().remove_conn(conn_id)
245 : {
246 0 : info!("closed connection removed");
247 0 : }
248 0 : }
249 0 : .instrument(span),
250 : );
251 :
252 0 : let client = ClientInnerCommon {
253 0 : inner: client,
254 0 : aux,
255 0 : conn_id,
256 0 : data: ClientDataEnum::Http(ClientDataHttp()),
257 0 : };
258 :
259 0 : Client::new(client)
260 0 : }
261 :
262 : pub(crate) struct Client<C: ClientInnerExt + Clone> {
263 : pub(crate) inner: ClientInnerCommon<C>,
264 : }
265 :
266 : impl<C: ClientInnerExt + Clone> Client<C> {
267 0 : pub(self) fn new(inner: ClientInnerCommon<C>) -> Self {
268 0 : Self { inner }
269 0 : }
270 :
271 0 : pub(crate) fn metrics(&self, ctx: &RequestContext) -> Arc<MetricCounter> {
272 0 : let aux = &self.inner.aux;
273 :
274 0 : let private_link_id = match ctx.extra() {
275 0 : None => None,
276 0 : Some(ConnectionInfoExtra::Aws { vpce_id }) => Some(vpce_id.clone()),
277 0 : Some(ConnectionInfoExtra::Azure { link_id }) => Some(link_id.to_smolstr()),
278 : };
279 :
280 0 : USAGE_METRICS.register(Ids {
281 0 : endpoint_id: aux.endpoint_id,
282 0 : branch_id: aux.branch_id,
283 0 : private_link_id,
284 0 : })
285 0 : }
286 : }
287 :
288 : impl ClientInnerExt for Send {
289 0 : fn is_closed(&self) -> bool {
290 0 : self.is_closed()
291 0 : }
292 :
293 0 : fn get_process_id(&self) -> i32 {
294 : // ideally throw something meaningful
295 0 : -1
296 0 : }
297 : }
|
