LCOV - code coverage report
Current view: top level - proxy/src - lib.rs (source / functions) Coverage Total Hit
Test: 2b0730d767f560e20b6748f57465922aa8bb805e.info Lines: 26.8 % 71 19
Test Date: 2024-09-25 14:04:07 Functions: 20.5 % 88 18

            Line data    Source code
       1              : // rustc lints/lint groups
       2              : // https://doc.rust-lang.org/rustc/lints/groups.html
       3              : #![deny(
       4              :     deprecated,
       5              :     future_incompatible,
       6              :     let_underscore,
       7              :     nonstandard_style,
       8              :     rust_2024_compatibility
       9              : )]
      10              : #![warn(clippy::all, clippy::pedantic, clippy::cargo)]
      11              : // List of denied lints from the clippy::restriction group.
      12              : // https://rust-lang.github.io/rust-clippy/master/index.html#?groups=restriction
      13              : #![warn(
      14              :     clippy::undocumented_unsafe_blocks,
      15              :     // TODO: Enable once all individual checks are enabled.
      16              :     //clippy::as_conversions,
      17              :     clippy::dbg_macro,
      18              :     clippy::empty_enum_variants_with_brackets,
      19              :     clippy::exit,
      20              :     clippy::float_cmp_const,
      21              :     clippy::lossy_float_literal,
      22              :     clippy::macro_use_imports,
      23              :     clippy::manual_ok_or,
      24              :     // TODO: consider clippy::map_err_ignore
      25              :     // TODO: consider clippy::mem_forget
      26              :     clippy::rc_mutex,
      27              :     clippy::rest_pat_in_fully_bound_structs,
      28              :     clippy::string_add,
      29              :     clippy::string_to_string,
      30              :     clippy::todo,
      31              :     // TODO: consider clippy::unimplemented
      32              :     // TODO: consider clippy::unwrap_used
      33              : )]
      34              : // List of permanently allowed lints.
      35              : #![allow(
      36              :     // It's ok to cast bool to u8, etc.
      37              :     clippy::cast_lossless,
      38              :     // Seems unavoidable.
      39              :     clippy::multiple_crate_versions,
      40              :     // While #[must_use] is a great feature this check is too noisy.
      41              :     clippy::must_use_candidate,
      42              :     // Inline consts, structs, fns, imports, etc. are ok if they're used by
      43              :     // the following statement(s).
      44              :     clippy::items_after_statements,
      45              : )]
      46              : // List of temporarily allowed lints.
      47              : // TODO: fix code and reduce list or move to permanent list above.
      48              : #![expect(
      49              :     clippy::cargo_common_metadata,
      50              :     clippy::cast_possible_truncation,
      51              :     clippy::cast_possible_wrap,
      52              :     clippy::cast_precision_loss,
      53              :     clippy::cast_sign_loss,
      54              :     clippy::doc_markdown,
      55              :     clippy::inline_always,
      56              :     clippy::match_same_arms,
      57              :     clippy::match_wild_err_arm,
      58              :     clippy::missing_errors_doc,
      59              :     clippy::missing_panics_doc,
      60              :     clippy::module_name_repetitions,
      61              :     clippy::needless_pass_by_value,
      62              :     clippy::redundant_closure_for_method_calls,
      63              :     clippy::similar_names,
      64              :     clippy::single_match_else,
      65              :     clippy::struct_excessive_bools,
      66              :     clippy::struct_field_names,
      67              :     clippy::too_many_lines,
      68              :     clippy::unused_self
      69              : )]
      70              : #![cfg_attr(
      71              :     any(test, feature = "testing"),
      72              :     allow(
      73              :         clippy::needless_raw_string_hashes,
      74              :         clippy::unreadable_literal,
      75              :         clippy::unused_async,
      76              :     )
      77              : )]
      78              : // List of temporarily allowed lints to unblock beta/nightly.
      79              : #![allow(
      80              :     unknown_lints,
      81              :     // TODO: 1.82: Add `use<T>` where necessary and remove from this list.
      82              :     impl_trait_overcaptures,
      83              : )]
      84              : 
      85              : use std::{convert::Infallible, future::Future};
      86              : 
      87              : use anyhow::{bail, Context};
      88              : use intern::{EndpointIdInt, EndpointIdTag, InternId};
      89              : use tokio::task::JoinError;
      90              : use tokio_util::sync::CancellationToken;
      91              : use tracing::warn;
      92              : 
      93              : pub mod auth;
      94              : pub mod cache;
      95              : pub mod cancellation;
      96              : pub mod compute;
      97              : pub mod config;
      98              : pub mod console;
      99              : pub mod context;
     100              : pub mod error;
     101              : pub mod http;
     102              : pub mod intern;
     103              : pub mod jemalloc;
     104              : pub mod logging;
     105              : pub mod metrics;
     106              : pub mod parse;
     107              : pub mod protocol2;
     108              : pub mod proxy;
     109              : pub mod rate_limiter;
     110              : pub mod redis;
     111              : pub mod sasl;
     112              : pub mod scram;
     113              : pub mod serverless;
     114              : pub mod stream;
     115              : pub mod url;
     116              : pub mod usage_metrics;
     117              : pub mod waiters;
     118              : 
     119              : /// Handle unix signals appropriately.
     120            0 : pub async fn handle_signals<F, Fut>(
     121            0 :     token: CancellationToken,
     122            0 :     mut refresh_config: F,
     123            0 : ) -> anyhow::Result<Infallible>
     124            0 : where
     125            0 :     F: FnMut() -> Fut,
     126            0 :     Fut: Future<Output = anyhow::Result<()>>,
     127            0 : {
     128              :     use tokio::signal::unix::{signal, SignalKind};
     129              : 
     130            0 :     let mut hangup = signal(SignalKind::hangup())?;
     131            0 :     let mut interrupt = signal(SignalKind::interrupt())?;
     132            0 :     let mut terminate = signal(SignalKind::terminate())?;
     133              : 
     134              :     loop {
     135            0 :         tokio::select! {
     136              :             // Hangup is commonly used for config reload.
     137            0 :             _ = hangup.recv() => {
     138            0 :                 warn!("received SIGHUP");
     139            0 :                 refresh_config().await?;
     140              :             }
     141              :             // Shut down the whole application.
     142            0 :             _ = interrupt.recv() => {
     143            0 :                 warn!("received SIGINT, exiting immediately");
     144            0 :                 bail!("interrupted");
     145              :             }
     146            0 :             _ = terminate.recv() => {
     147            0 :                 warn!("received SIGTERM, shutting down once all existing connections have closed");
     148            0 :                 token.cancel();
     149              :             }
     150              :         }
     151              :     }
     152            0 : }
     153              : 
     154              : /// Flattens `Result<Result<T>>` into `Result<T>`.
     155            0 : pub fn flatten_err<T>(r: Result<anyhow::Result<T>, JoinError>) -> anyhow::Result<T> {
     156            0 :     r.context("join error").and_then(|x| x)
     157            0 : }
     158              : 
     159              : macro_rules! smol_str_wrapper {
     160              :     ($name:ident) => {
     161              :         #[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Default)]
     162              :         pub struct $name(smol_str::SmolStr);
     163              : 
     164              :         impl $name {
     165              :             #[allow(unused)]
     166            0 :             pub(crate) fn as_str(&self) -> &str {
     167            0 :                 self.0.as_str()
     168            0 :             }
     169              :         }
     170              : 
     171              :         impl std::fmt::Display for $name {
     172            3 :             fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
     173            3 :                 self.0.fmt(f)
     174            3 :             }
     175              :         }
     176              : 
     177              :         impl<T> std::cmp::PartialEq<T> for $name
     178              :         where
     179              :             smol_str::SmolStr: std::cmp::PartialEq<T>,
     180              :         {
     181           19 :             fn eq(&self, other: &T) -> bool {
     182           19 :                 self.0.eq(other)
     183           19 :             }
     184              :         }
     185              : 
     186              :         impl<T> From<T> for $name
     187              :         where
     188              :             smol_str::SmolStr: From<T>,
     189              :         {
     190          164 :             fn from(x: T) -> Self {
     191          164 :                 Self(x.into())
     192          164 :             }
     193              :         }
     194              : 
     195              :         impl AsRef<str> for $name {
     196           10 :             fn as_ref(&self) -> &str {
     197           10 :                 self.0.as_ref()
     198           10 :             }
     199              :         }
     200              : 
     201              :         impl std::ops::Deref for $name {
     202              :             type Target = str;
     203          141 :             fn deref(&self) -> &str {
     204          141 :                 &*self.0
     205          141 :             }
     206              :         }
     207              : 
     208              :         impl<'de> serde::de::Deserialize<'de> for $name {
     209            0 :             fn deserialize<D: serde::de::Deserializer<'de>>(d: D) -> Result<Self, D::Error> {
     210            0 :                 <smol_str::SmolStr as serde::de::Deserialize<'de>>::deserialize(d).map(Self)
     211            0 :             }
     212              :         }
     213              : 
     214              :         impl serde::Serialize for $name {
     215            0 :             fn serialize<S: serde::Serializer>(&self, s: S) -> Result<S::Ok, S::Error> {
     216            0 :                 self.0.serialize(s)
     217            0 :             }
     218              :         }
     219              :     };
     220              : }
     221              : 
     222              : const POOLER_SUFFIX: &str = "-pooler";
     223              : 
     224              : impl EndpointId {
     225            3 :     fn normalize(&self) -> Self {
     226            3 :         if let Some(stripped) = self.as_ref().strip_suffix(POOLER_SUFFIX) {
     227            0 :             stripped.into()
     228              :         } else {
     229            3 :             self.clone()
     230              :         }
     231            3 :     }
     232              : 
     233            0 :     fn normalize_intern(&self) -> EndpointIdInt {
     234            0 :         if let Some(stripped) = self.as_ref().strip_suffix(POOLER_SUFFIX) {
     235            0 :             EndpointIdTag::get_interner().get_or_intern(stripped)
     236              :         } else {
     237            0 :             self.into()
     238              :         }
     239            0 :     }
     240              : }
     241              : 
     242              : // 90% of role name strings are 20 characters or less.
     243              : smol_str_wrapper!(RoleName);
     244              : // 50% of endpoint strings are 23 characters or less.
     245              : smol_str_wrapper!(EndpointId);
     246              : // 50% of branch strings are 23 characters or less.
     247              : smol_str_wrapper!(BranchId);
     248              : // 90% of project strings are 23 characters or less.
     249              : smol_str_wrapper!(ProjectId);
     250              : 
     251              : // will usually equal endpoint ID
     252              : smol_str_wrapper!(EndpointCacheKey);
     253              : 
     254              : smol_str_wrapper!(DbName);
     255              : 
     256              : // postgres hostname, will likely be a port:ip addr
     257              : smol_str_wrapper!(Host);
     258              : 
     259              : // Endpoints are a bit tricky. Rare they might be branches or projects.
     260              : impl EndpointId {
     261            0 :     pub(crate) fn is_endpoint(&self) -> bool {
     262            0 :         self.0.starts_with("ep-")
     263            0 :     }
     264            0 :     pub(crate) fn is_branch(&self) -> bool {
     265            0 :         self.0.starts_with("br-")
     266            0 :     }
     267              :     // pub(crate) fn is_project(&self) -> bool {
     268              :     //     !self.is_endpoint() && !self.is_branch()
     269              :     // }
     270            0 :     pub(crate) fn as_branch(&self) -> BranchId {
     271            0 :         BranchId(self.0.clone())
     272            0 :     }
     273            0 :     pub(crate) fn as_project(&self) -> ProjectId {
     274            0 :         ProjectId(self.0.clone())
     275            0 :     }
     276              : }
        

Generated by: LCOV version 2.1-beta