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