LCOV - code coverage report
Current view: top level - proxy/src/context - mod.rs (source / functions) Coverage Total Hit
Test: 07bee600374ccd486c69370d0972d9035964fe68.info Lines: 34.3 % 312 107
Test Date: 2025-02-20 13:11:02 Functions: 27.5 % 40 11

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

Generated by: LCOV version 2.1-beta