Line data Source code
1 : //! Connection request monitoring contexts
2 :
3 : use std::net::IpAddr;
4 :
5 : use chrono::Utc;
6 : use once_cell::sync::OnceCell;
7 : use smol_str::SmolStr;
8 : use tokio::sync::mpsc;
9 : use tracing::field::display;
10 : use tracing::{Span, debug, error, info_span};
11 : use try_lock::TryLock;
12 : use uuid::Uuid;
13 :
14 : use self::parquet::RequestData;
15 : use crate::control_plane::messages::{ColdStartInfo, MetricsAuxInfo};
16 : use crate::error::ErrorKind;
17 : use crate::intern::{BranchIdInt, ProjectIdInt};
18 : use crate::metrics::{
19 : ConnectOutcome, InvalidEndpointsGroup, LatencyAccumulated, LatencyTimer, Metrics, Protocol,
20 : Waiting,
21 : };
22 : use crate::pqproto::StartupMessageParams;
23 : use crate::protocol2::{ConnectionInfo, ConnectionInfoExtra};
24 : use crate::types::{DbName, EndpointId, RoleName};
25 :
26 : pub mod parquet;
27 :
28 : pub(crate) static LOG_CHAN: OnceCell<mpsc::WeakUnboundedSender<RequestData>> = OnceCell::new();
29 : pub(crate) static LOG_CHAN_DISCONNECT: OnceCell<mpsc::WeakUnboundedSender<RequestData>> =
30 : OnceCell::new();
31 :
32 : /// Context data for a single request to connect to a database.
33 : ///
34 : /// This data should **not** be used for connection logic, only for observability and limiting purposes.
35 : /// All connection logic should instead use strongly typed state machines, not a bunch of Options.
36 : pub struct RequestContext(
37 : /// To allow easier use of the ctx object, we have interior mutability.
38 : /// I would typically use a RefCell but that would break the `Send` requirements
39 : /// so we need something with thread-safety. `TryLock` is a cheap alternative
40 : /// that offers similar semantics to a `RefCell` but with synchronisation.
41 : TryLock<RequestContextInner>,
42 : );
43 :
44 : struct RequestContextInner {
45 : pub(crate) conn_info: ConnectionInfo,
46 : pub(crate) session_id: Uuid,
47 : pub(crate) protocol: Protocol,
48 : first_packet: chrono::DateTime<Utc>,
49 : region: &'static str,
50 : pub(crate) span: Span,
51 :
52 : // filled in as they are discovered
53 : project: Option<ProjectIdInt>,
54 : branch: Option<BranchIdInt>,
55 : endpoint_id: Option<EndpointId>,
56 : dbname: Option<DbName>,
57 : user: Option<RoleName>,
58 : application: Option<SmolStr>,
59 : user_agent: Option<SmolStr>,
60 : error_kind: Option<ErrorKind>,
61 : pub(crate) auth_method: Option<AuthMethod>,
62 : jwt_issuer: Option<String>,
63 : success: bool,
64 : pub(crate) cold_start_info: ColdStartInfo,
65 : pg_options: Option<StartupMessageParams>,
66 : testodrome_query_id: Option<SmolStr>,
67 :
68 : // extra
69 : // This sender is here to keep the request monitoring channel open while requests are taking place.
70 : sender: Option<mpsc::UnboundedSender<RequestData>>,
71 : // This sender is only used to log the length of session in case of success.
72 : disconnect_sender: Option<mpsc::UnboundedSender<RequestData>>,
73 : pub(crate) latency_timer: LatencyTimer,
74 : // Whether proxy decided that it's not a valid endpoint end rejected it before going to cplane.
75 : rejected: Option<bool>,
76 : disconnect_timestamp: Option<chrono::DateTime<Utc>>,
77 : }
78 :
79 : #[derive(Clone, Debug)]
80 : pub(crate) enum AuthMethod {
81 : // aka link
82 : ConsoleRedirect,
83 : ScramSha256,
84 : ScramSha256Plus,
85 : Cleartext,
86 : Jwt,
87 : }
88 :
89 : impl Clone for RequestContext {
90 0 : fn clone(&self) -> Self {
91 0 : let inner = self.0.try_lock().expect("should not deadlock");
92 0 : let new = RequestContextInner {
93 0 : conn_info: inner.conn_info.clone(),
94 0 : session_id: inner.session_id,
95 0 : protocol: inner.protocol,
96 0 : first_packet: inner.first_packet,
97 0 : region: inner.region,
98 0 : span: info_span!("background_task"),
99 :
100 0 : project: inner.project,
101 0 : branch: inner.branch,
102 0 : endpoint_id: inner.endpoint_id.clone(),
103 0 : dbname: inner.dbname.clone(),
104 0 : user: inner.user.clone(),
105 0 : application: inner.application.clone(),
106 0 : user_agent: inner.user_agent.clone(),
107 0 : error_kind: inner.error_kind,
108 0 : auth_method: inner.auth_method.clone(),
109 0 : jwt_issuer: inner.jwt_issuer.clone(),
110 0 : success: inner.success,
111 0 : rejected: inner.rejected,
112 0 : cold_start_info: inner.cold_start_info,
113 0 : pg_options: inner.pg_options.clone(),
114 0 : testodrome_query_id: inner.testodrome_query_id.clone(),
115 0 :
116 0 : sender: None,
117 0 : disconnect_sender: None,
118 0 : latency_timer: LatencyTimer::noop(inner.protocol),
119 0 : disconnect_timestamp: inner.disconnect_timestamp,
120 0 : };
121 0 :
122 0 : Self(TryLock::new(new))
123 0 : }
124 : }
125 :
126 : impl RequestContext {
127 75 : pub fn new(
128 75 : session_id: Uuid,
129 75 : conn_info: ConnectionInfo,
130 75 : protocol: Protocol,
131 75 : region: &'static str,
132 75 : ) -> Self {
133 : // TODO: be careful with long lived spans
134 75 : let span = info_span!(
135 75 : "connect_request",
136 75 : %protocol,
137 75 : ?session_id,
138 75 : %conn_info,
139 75 : ep = tracing::field::Empty,
140 75 : role = tracing::field::Empty,
141 75 : );
142 :
143 75 : let inner = RequestContextInner {
144 75 : conn_info,
145 75 : session_id,
146 75 : protocol,
147 75 : first_packet: Utc::now(),
148 75 : region,
149 75 : span,
150 75 :
151 75 : project: None,
152 75 : branch: None,
153 75 : endpoint_id: None,
154 75 : dbname: None,
155 75 : user: None,
156 75 : application: None,
157 75 : user_agent: None,
158 75 : error_kind: None,
159 75 : auth_method: None,
160 75 : jwt_issuer: None,
161 75 : success: false,
162 75 : rejected: None,
163 75 : cold_start_info: ColdStartInfo::Unknown,
164 75 : pg_options: None,
165 75 : testodrome_query_id: None,
166 75 :
167 75 : sender: LOG_CHAN.get().and_then(|tx| tx.upgrade()),
168 75 : disconnect_sender: LOG_CHAN_DISCONNECT.get().and_then(|tx| tx.upgrade()),
169 75 : latency_timer: LatencyTimer::new(protocol),
170 75 : disconnect_timestamp: None,
171 75 : };
172 75 :
173 75 : Self(TryLock::new(inner))
174 75 : }
175 :
176 : #[cfg(test)]
177 75 : pub(crate) fn test() -> Self {
178 : use std::net::SocketAddr;
179 75 : let ip = IpAddr::from([127, 0, 0, 1]);
180 75 : let addr = SocketAddr::new(ip, 5432);
181 75 : let conn_info = ConnectionInfo { addr, extra: None };
182 75 : RequestContext::new(Uuid::now_v7(), conn_info, Protocol::Tcp, "test")
183 75 : }
184 :
185 0 : pub(crate) fn console_application_name(&self) -> String {
186 0 : let this = self.0.try_lock().expect("should not deadlock");
187 0 : format!(
188 0 : "{}/{}",
189 0 : this.application.as_deref().unwrap_or_default(),
190 0 : this.protocol
191 0 : )
192 0 : }
193 :
194 0 : pub(crate) fn set_rejected(&self, rejected: bool) {
195 0 : let mut this = self.0.try_lock().expect("should not deadlock");
196 0 : this.rejected = Some(rejected);
197 0 : }
198 :
199 0 : pub(crate) fn set_cold_start_info(&self, info: ColdStartInfo) {
200 0 : self.0
201 0 : .try_lock()
202 0 : .expect("should not deadlock")
203 0 : .set_cold_start_info(info);
204 0 : }
205 :
206 0 : pub(crate) fn set_db_options(&self, options: StartupMessageParams) {
207 0 : let mut this = self.0.try_lock().expect("should not deadlock");
208 0 : this.set_application(options.get("application_name").map(SmolStr::from));
209 0 : if let Some(user) = options.get("user") {
210 0 : this.set_user(user.into());
211 0 : }
212 0 : if let Some(dbname) = options.get("database") {
213 0 : this.set_dbname(dbname.into());
214 0 : }
215 :
216 : // Try to get testodrome_query_id directly from parameters
217 0 : if let Some(options_str) = options.get("options") {
218 : // If not found directly, try to extract it from the options string
219 0 : for option in options_str.split_whitespace() {
220 0 : if option.starts_with("neon_query_id:") {
221 0 : if let Some(value) = option.strip_prefix("neon_query_id:") {
222 0 : this.set_testodrome_id(value.into());
223 0 : break;
224 0 : }
225 0 : }
226 : }
227 0 : }
228 :
229 0 : this.pg_options = Some(options);
230 0 : }
231 :
232 0 : pub(crate) fn set_project(&self, x: MetricsAuxInfo) {
233 0 : let mut this = self.0.try_lock().expect("should not deadlock");
234 0 : if this.endpoint_id.is_none() {
235 0 : this.set_endpoint_id(x.endpoint_id.as_str().into());
236 0 : }
237 0 : this.branch = Some(x.branch_id);
238 0 : this.project = Some(x.project_id);
239 0 : this.set_cold_start_info(x.cold_start_info);
240 0 : }
241 :
242 0 : pub(crate) fn set_project_id(&self, project_id: ProjectIdInt) {
243 0 : let mut this = self.0.try_lock().expect("should not deadlock");
244 0 : this.project = Some(project_id);
245 0 : }
246 :
247 29 : pub(crate) fn set_endpoint_id(&self, endpoint_id: EndpointId) {
248 29 : self.0
249 29 : .try_lock()
250 29 : .expect("should not deadlock")
251 29 : .set_endpoint_id(endpoint_id);
252 29 : }
253 :
254 0 : pub(crate) fn set_dbname(&self, dbname: DbName) {
255 0 : self.0
256 0 : .try_lock()
257 0 : .expect("should not deadlock")
258 0 : .set_dbname(dbname);
259 0 : }
260 :
261 0 : pub(crate) fn set_user(&self, user: RoleName) {
262 0 : self.0
263 0 : .try_lock()
264 0 : .expect("should not deadlock")
265 0 : .set_user(user);
266 0 : }
267 :
268 0 : pub(crate) fn set_user_agent(&self, user_agent: Option<SmolStr>) {
269 0 : self.0
270 0 : .try_lock()
271 0 : .expect("should not deadlock")
272 0 : .set_user_agent(user_agent);
273 0 : }
274 :
275 0 : pub(crate) fn set_testodrome_id(&self, query_id: SmolStr) {
276 0 : self.0
277 0 : .try_lock()
278 0 : .expect("should not deadlock")
279 0 : .set_testodrome_id(query_id);
280 0 : }
281 :
282 15 : pub(crate) fn set_auth_method(&self, auth_method: AuthMethod) {
283 15 : let mut this = self.0.try_lock().expect("should not deadlock");
284 15 : this.auth_method = Some(auth_method);
285 15 : }
286 :
287 12 : pub(crate) fn set_jwt_issuer(&self, jwt_issuer: String) {
288 12 : let mut this = self.0.try_lock().expect("should not deadlock");
289 12 : this.jwt_issuer = Some(jwt_issuer);
290 12 : }
291 :
292 0 : pub fn has_private_peer_addr(&self) -> bool {
293 0 : self.0
294 0 : .try_lock()
295 0 : .expect("should not deadlock")
296 0 : .has_private_peer_addr()
297 0 : }
298 :
299 0 : pub(crate) fn set_error_kind(&self, kind: ErrorKind) {
300 0 : let mut this = self.0.try_lock().expect("should not deadlock");
301 0 : // Do not record errors from the private address to metrics.
302 0 : if !this.has_private_peer_addr() {
303 0 : Metrics::get().proxy.errors_total.inc(kind);
304 0 : }
305 0 : if let Some(ep) = &this.endpoint_id {
306 0 : let metric = &Metrics::get().proxy.endpoints_affected_by_errors;
307 0 : let label = metric.with_labels(kind);
308 0 : metric.get_metric(label).measure(ep);
309 0 : }
310 0 : this.error_kind = Some(kind);
311 0 : }
312 :
313 0 : pub fn set_success(&self) {
314 0 : let mut this = self.0.try_lock().expect("should not deadlock");
315 0 : this.success = true;
316 0 : }
317 :
318 0 : pub fn log_connect(self) -> DisconnectLogger {
319 0 : let mut this = self.0.into_inner();
320 0 : this.log_connect();
321 0 :
322 0 : // close current span.
323 0 : this.span = Span::none();
324 0 :
325 0 : DisconnectLogger(this)
326 0 : }
327 :
328 16 : pub(crate) fn protocol(&self) -> Protocol {
329 16 : self.0.try_lock().expect("should not deadlock").protocol
330 16 : }
331 :
332 0 : pub(crate) fn span(&self) -> Span {
333 0 : self.0.try_lock().expect("should not deadlock").span.clone()
334 0 : }
335 :
336 0 : pub(crate) fn session_id(&self) -> Uuid {
337 0 : self.0.try_lock().expect("should not deadlock").session_id
338 0 : }
339 :
340 3 : pub(crate) fn peer_addr(&self) -> IpAddr {
341 3 : self.0
342 3 : .try_lock()
343 3 : .expect("should not deadlock")
344 3 : .conn_info
345 3 : .addr
346 3 : .ip()
347 3 : }
348 :
349 0 : pub(crate) fn extra(&self) -> Option<ConnectionInfoExtra> {
350 0 : self.0
351 0 : .try_lock()
352 0 : .expect("should not deadlock")
353 0 : .conn_info
354 0 : .extra
355 0 : .clone()
356 0 : }
357 :
358 0 : pub(crate) fn cold_start_info(&self) -> ColdStartInfo {
359 0 : self.0
360 0 : .try_lock()
361 0 : .expect("should not deadlock")
362 0 : .cold_start_info
363 0 : }
364 :
365 52 : pub(crate) fn latency_timer_pause(&self, waiting_for: Waiting) -> LatencyTimerPause<'_> {
366 52 : LatencyTimerPause {
367 52 : ctx: self,
368 52 : start: tokio::time::Instant::now(),
369 52 : waiting_for,
370 52 : }
371 52 : }
372 :
373 0 : pub(crate) fn latency_timer_pause_at(
374 0 : &self,
375 0 : at: tokio::time::Instant,
376 0 : waiting_for: Waiting,
377 0 : ) -> LatencyTimerPause<'_> {
378 0 : LatencyTimerPause {
379 0 : ctx: self,
380 0 : start: at,
381 0 : waiting_for,
382 0 : }
383 0 : }
384 :
385 0 : pub(crate) fn get_proxy_latency(&self) -> LatencyAccumulated {
386 0 : self.0
387 0 : .try_lock()
388 0 : .expect("should not deadlock")
389 0 : .latency_timer
390 0 : .accumulated()
391 0 : }
392 :
393 0 : pub(crate) fn get_testodrome_id(&self) -> Option<SmolStr> {
394 0 : self.0
395 0 : .try_lock()
396 0 : .expect("should not deadlock")
397 0 : .testodrome_query_id
398 0 : .clone()
399 0 : }
400 :
401 7 : pub(crate) fn success(&self) {
402 7 : self.0
403 7 : .try_lock()
404 7 : .expect("should not deadlock")
405 7 : .latency_timer
406 7 : .success();
407 7 : }
408 : }
409 :
410 : pub(crate) struct LatencyTimerPause<'a> {
411 : ctx: &'a RequestContext,
412 : start: tokio::time::Instant,
413 : waiting_for: Waiting,
414 : }
415 :
416 : impl Drop for LatencyTimerPause<'_> {
417 52 : fn drop(&mut self) {
418 52 : self.ctx
419 52 : .0
420 52 : .try_lock()
421 52 : .expect("should not deadlock")
422 52 : .latency_timer
423 52 : .unpause(self.start, self.waiting_for);
424 52 : }
425 : }
426 :
427 : impl RequestContextInner {
428 0 : fn set_cold_start_info(&mut self, info: ColdStartInfo) {
429 0 : self.cold_start_info = info;
430 0 : self.latency_timer.cold_start_info(info);
431 0 : }
432 :
433 29 : fn set_endpoint_id(&mut self, endpoint_id: EndpointId) {
434 29 : if self.endpoint_id.is_none() {
435 29 : self.span.record("ep", display(&endpoint_id));
436 29 : let metric = &Metrics::get().proxy.connecting_endpoints;
437 29 : let label = metric.with_labels(self.protocol);
438 29 : metric.get_metric(label).measure(&endpoint_id);
439 29 : self.endpoint_id = Some(endpoint_id);
440 29 : }
441 29 : }
442 :
443 0 : fn set_application(&mut self, app: Option<SmolStr>) {
444 0 : if let Some(app) = app {
445 0 : self.application = Some(app);
446 0 : }
447 0 : }
448 :
449 0 : fn set_user_agent(&mut self, user_agent: Option<SmolStr>) {
450 0 : self.user_agent = user_agent;
451 0 : }
452 :
453 0 : fn set_dbname(&mut self, dbname: DbName) {
454 0 : self.dbname = Some(dbname);
455 0 : }
456 :
457 0 : fn set_user(&mut self, user: RoleName) {
458 0 : self.span.record("role", display(&user));
459 0 : self.user = Some(user);
460 0 : }
461 :
462 0 : fn set_testodrome_id(&mut self, query_id: SmolStr) {
463 0 : self.testodrome_query_id = Some(query_id);
464 0 : }
465 :
466 0 : fn has_private_peer_addr(&self) -> bool {
467 0 : match self.conn_info.addr.ip() {
468 0 : IpAddr::V4(ip) => ip.is_private(),
469 0 : IpAddr::V6(_) => false,
470 : }
471 0 : }
472 :
473 0 : fn log_connect(&mut self) {
474 0 : let outcome = if self.success {
475 0 : ConnectOutcome::Success
476 : } else {
477 0 : ConnectOutcome::Failed
478 : };
479 :
480 : // TODO: get rid of entirely/refactor
481 : // check for false positives
482 : // AND false negatives
483 0 : if let Some(rejected) = self.rejected {
484 0 : let ep = self
485 0 : .endpoint_id
486 0 : .as_ref()
487 0 : .map(|x| x.as_str())
488 0 : .unwrap_or_default();
489 0 : // This makes sense only if cache is disabled
490 0 : debug!(
491 : ?outcome,
492 : ?rejected,
493 : ?ep,
494 0 : "check endpoint is valid with outcome"
495 : );
496 0 : Metrics::get()
497 0 : .proxy
498 0 : .invalid_endpoints_total
499 0 : .inc(InvalidEndpointsGroup {
500 0 : protocol: self.protocol,
501 0 : rejected: rejected.into(),
502 0 : outcome,
503 0 : });
504 0 : }
505 :
506 0 : if let Some(tx) = self.sender.take() {
507 : // If type changes, this error handling needs to be updated.
508 0 : let tx: mpsc::UnboundedSender<RequestData> = tx;
509 0 : if let Err(e) = tx.send(RequestData::from(&*self)) {
510 0 : error!("log_connect channel send failed: {e}");
511 0 : }
512 0 : }
513 0 : }
514 :
515 0 : fn log_disconnect(&mut self) {
516 0 : // If we are here, it's guaranteed that the user successfully connected to the endpoint.
517 0 : // Here we log the length of the session.
518 0 : self.disconnect_timestamp = Some(Utc::now());
519 0 : if let Some(tx) = self.disconnect_sender.take() {
520 : // If type changes, this error handling needs to be updated.
521 0 : let tx: mpsc::UnboundedSender<RequestData> = tx;
522 0 : if let Err(e) = tx.send(RequestData::from(&*self)) {
523 0 : error!("log_disconnect channel send failed: {e}");
524 0 : }
525 0 : }
526 0 : }
527 : }
528 :
529 : impl Drop for RequestContextInner {
530 75 : fn drop(&mut self) {
531 75 : if self.sender.is_some() {
532 0 : self.log_connect();
533 75 : }
534 75 : }
535 : }
536 :
537 : pub struct DisconnectLogger(RequestContextInner);
538 :
539 : impl Drop for DisconnectLogger {
540 0 : fn drop(&mut self) {
541 0 : self.0.log_disconnect();
542 0 : }
543 : }
|
