LCOV - code coverage report
Current view: top level - proxy/src/binary - proxy.rs (source / functions) Coverage Total Hit
Test: 07bee600374ccd486c69370d0972d9035964fe68.info Lines: 8.3 % 504 42
Test Date: 2025-02-20 13:11:02 Functions: 18.4 % 152 28

            Line data    Source code
       1              : use std::net::SocketAddr;
       2              : use std::pin::pin;
       3              : use std::sync::Arc;
       4              : use std::time::Duration;
       5              : 
       6              : use anyhow::bail;
       7              : use futures::future::Either;
       8              : use remote_storage::RemoteStorageConfig;
       9              : use tokio::net::TcpListener;
      10              : use tokio::task::JoinSet;
      11              : use tokio_util::sync::CancellationToken;
      12              : use tracing::{info, warn, Instrument};
      13              : use utils::sentry_init::init_sentry;
      14              : use utils::{project_build_tag, project_git_version};
      15              : 
      16              : use crate::auth::backend::jwt::JwkCache;
      17              : use crate::auth::backend::{AuthRateLimiter, ConsoleRedirectBackend, MaybeOwned};
      18              : use crate::cancellation::{handle_cancel_messages, CancellationHandler};
      19              : use crate::config::{
      20              :     self, remote_storage_from_toml, AuthenticationConfig, CacheOptions, ComputeConfig, HttpConfig,
      21              :     ProjectInfoCacheOptions, ProxyConfig, ProxyProtocolV2,
      22              : };
      23              : use crate::context::parquet::ParquetUploadArgs;
      24              : use crate::http::health_server::AppMetrics;
      25              : use crate::metrics::Metrics;
      26              : use crate::rate_limiter::{
      27              :     EndpointRateLimiter, LeakyBucketConfig, RateBucketInfo, WakeComputeRateLimiter,
      28              : };
      29              : use crate::redis::connection_with_credentials_provider::ConnectionWithCredentialsProvider;
      30              : use crate::redis::kv_ops::RedisKVClient;
      31              : use crate::redis::{elasticache, notifications};
      32              : use crate::scram::threadpool::ThreadPool;
      33              : use crate::serverless::cancel_set::CancelSet;
      34              : use crate::serverless::GlobalConnPoolOptions;
      35              : use crate::tls::client_config::compute_client_config_with_root_certs;
      36              : use crate::{auth, control_plane, http, serverless, usage_metrics};
      37              : 
      38              : project_git_version!(GIT_VERSION);
      39              : project_build_tag!(BUILD_TAG);
      40              : 
      41              : use clap::{Parser, ValueEnum};
      42              : 
      43              : #[derive(Clone, Debug, ValueEnum)]
      44              : enum AuthBackendType {
      45              :     #[value(name("cplane-v1"), alias("control-plane"))]
      46              :     ControlPlaneV1,
      47              : 
      48              :     #[value(name("link"), alias("control-redirect"))]
      49              :     ConsoleRedirect,
      50              : 
      51              :     #[cfg(any(test, feature = "testing"))]
      52              :     Postgres,
      53              : }
      54              : 
      55              : /// Neon proxy/router
      56              : #[derive(Parser)]
      57              : #[command(version = GIT_VERSION, about)]
      58              : struct ProxyCliArgs {
      59              :     /// Name of the region this proxy is deployed in
      60            1 :     #[clap(long, default_value_t = String::new())]
      61            0 :     region: String,
      62              :     /// listen for incoming client connections on ip:port
      63              :     #[clap(short, long, default_value = "127.0.0.1:4432")]
      64            0 :     proxy: String,
      65            1 :     #[clap(value_enum, long, default_value_t = AuthBackendType::ConsoleRedirect)]
      66            0 :     auth_backend: AuthBackendType,
      67              :     /// listen for management callback connection on ip:port
      68              :     #[clap(short, long, default_value = "127.0.0.1:7000")]
      69            0 :     mgmt: String,
      70              :     /// listen for incoming http connections (metrics, etc) on ip:port
      71              :     #[clap(long, default_value = "127.0.0.1:7001")]
      72            0 :     http: String,
      73              :     /// listen for incoming wss connections on ip:port
      74              :     #[clap(long)]
      75              :     wss: Option<String>,
      76              :     /// redirect unauthenticated users to the given uri in case of console redirect auth
      77              :     #[clap(short, long, default_value = "http://localhost:3000/psql_session/")]
      78            0 :     uri: String,
      79              :     /// cloud API endpoint for authenticating users
      80              :     #[clap(
      81              :         short,
      82              :         long,
      83              :         default_value = "http://localhost:3000/authenticate_proxy_request/"
      84              :     )]
      85            0 :     auth_endpoint: String,
      86              :     /// JWT used to connect to control plane.
      87              :     #[clap(
      88              :         long,
      89              :         value_name = "JWT",
      90              :         default_value = "",
      91              :         env = "NEON_PROXY_TO_CONTROLPLANE_TOKEN"
      92              :     )]
      93            0 :     control_plane_token: Arc<str>,
      94              :     /// if this is not local proxy, this toggles whether we accept jwt or passwords for http
      95            1 :     #[clap(long, default_value_t = false, value_parser = clap::builder::BoolishValueParser::new(), action = clap::ArgAction::Set)]
      96            0 :     is_auth_broker: bool,
      97              :     /// path to TLS key for client postgres connections
      98              :     ///
      99              :     /// tls-key and tls-cert are for backwards compatibility, we can put all certs in one dir
     100              :     #[clap(short = 'k', long, alias = "ssl-key")]
     101              :     tls_key: Option<String>,
     102              :     /// path to TLS cert for client postgres connections
     103              :     ///
     104              :     /// tls-key and tls-cert are for backwards compatibility, we can put all certs in one dir
     105              :     #[clap(short = 'c', long, alias = "ssl-cert")]
     106              :     tls_cert: Option<String>,
     107              :     /// Allow writing TLS session keys to the given file pointed to by the environment variable `SSLKEYLOGFILE`.
     108              :     #[clap(long, alias = "allow-ssl-keylogfile")]
     109            0 :     allow_tls_keylogfile: bool,
     110              :     /// path to directory with TLS certificates for client postgres connections
     111              :     #[clap(long)]
     112              :     certs_dir: Option<String>,
     113              :     /// timeout for the TLS handshake
     114              :     #[clap(long, default_value = "15s", value_parser = humantime::parse_duration)]
     115            0 :     handshake_timeout: tokio::time::Duration,
     116              :     /// http endpoint to receive periodic metric updates
     117              :     #[clap(long)]
     118              :     metric_collection_endpoint: Option<String>,
     119              :     /// how often metrics should be sent to a collection endpoint
     120              :     #[clap(long)]
     121              :     metric_collection_interval: Option<String>,
     122              :     /// cache for `wake_compute` api method (use `size=0` to disable)
     123              :     #[clap(long, default_value = config::CacheOptions::CACHE_DEFAULT_OPTIONS)]
     124            0 :     wake_compute_cache: String,
     125              :     /// lock for `wake_compute` api method. example: "shards=32,permits=4,epoch=10m,timeout=1s". (use `permits=0` to disable).
     126              :     #[clap(long, default_value = config::ConcurrencyLockOptions::DEFAULT_OPTIONS_WAKE_COMPUTE_LOCK)]
     127            0 :     wake_compute_lock: String,
     128              :     /// lock for `connect_compute` api method. example: "shards=32,permits=4,epoch=10m,timeout=1s". (use `permits=0` to disable).
     129              :     #[clap(long, default_value = config::ConcurrencyLockOptions::DEFAULT_OPTIONS_CONNECT_COMPUTE_LOCK)]
     130            0 :     connect_compute_lock: String,
     131              :     #[clap(flatten)]
     132              :     sql_over_http: SqlOverHttpArgs,
     133              :     /// timeout for scram authentication protocol
     134              :     #[clap(long, default_value = "15s", value_parser = humantime::parse_duration)]
     135            0 :     scram_protocol_timeout: tokio::time::Duration,
     136              :     /// size of the threadpool for password hashing
     137            1 :     #[clap(long, default_value_t = 4)]
     138            0 :     scram_thread_pool_size: u8,
     139              :     /// Endpoint rate limiter max number of requests per second.
     140              :     ///
     141              :     /// Provided in the form `<Requests Per Second>@<Bucket Duration Size>`.
     142              :     /// Can be given multiple times for different bucket sizes.
     143            4 :     #[clap(long, default_values_t = RateBucketInfo::DEFAULT_ENDPOINT_SET)]
     144            1 :     endpoint_rps_limit: Vec<RateBucketInfo>,
     145              :     /// Wake compute rate limiter max number of requests per second.
     146            4 :     #[clap(long, default_values_t = RateBucketInfo::DEFAULT_SET)]
     147            1 :     wake_compute_limit: Vec<RateBucketInfo>,
     148              :     /// Whether the auth rate limiter actually takes effect (for testing)
     149            1 :     #[clap(long, default_value_t = false, value_parser = clap::builder::BoolishValueParser::new(), action = clap::ArgAction::Set)]
     150            0 :     auth_rate_limit_enabled: bool,
     151              :     /// Authentication rate limiter max number of hashes per second.
     152            4 :     #[clap(long, default_values_t = RateBucketInfo::DEFAULT_AUTH_SET)]
     153            1 :     auth_rate_limit: Vec<RateBucketInfo>,
     154              :     /// The IP subnet to use when considering whether two IP addresses are considered the same.
     155            1 :     #[clap(long, default_value_t = 64)]
     156            0 :     auth_rate_limit_ip_subnet: u8,
     157              :     /// Redis rate limiter max number of requests per second.
     158            3 :     #[clap(long, default_values_t = RateBucketInfo::DEFAULT_REDIS_SET)]
     159            1 :     redis_rps_limit: Vec<RateBucketInfo>,
     160              :     /// Cancellation channel size (max queue size for redis kv client)
     161              :     #[clap(long, default_value = "1024")]
     162            0 :     cancellation_ch_size: usize,
     163              :     /// cache for `allowed_ips` (use `size=0` to disable)
     164              :     #[clap(long, default_value = config::CacheOptions::CACHE_DEFAULT_OPTIONS)]
     165            0 :     allowed_ips_cache: String,
     166              :     /// cache for `role_secret` (use `size=0` to disable)
     167              :     #[clap(long, default_value = config::CacheOptions::CACHE_DEFAULT_OPTIONS)]
     168            0 :     role_secret_cache: String,
     169              :     /// redis url for notifications (if empty, redis_host:port will be used for both notifications and streaming connections)
     170              :     #[clap(long)]
     171              :     redis_notifications: Option<String>,
     172              :     /// what from the available authentications type to use for the regional redis we have. Supported are "irsa" and "plain".
     173              :     #[clap(long, default_value = "irsa")]
     174            0 :     redis_auth_type: String,
     175              :     /// redis host for streaming connections (might be different from the notifications host)
     176              :     #[clap(long)]
     177              :     redis_host: Option<String>,
     178              :     /// redis port for streaming connections (might be different from the notifications host)
     179              :     #[clap(long)]
     180              :     redis_port: Option<u16>,
     181              :     /// redis cluster name, used in aws elasticache
     182              :     #[clap(long)]
     183              :     redis_cluster_name: Option<String>,
     184              :     /// redis user_id, used in aws elasticache
     185              :     #[clap(long)]
     186              :     redis_user_id: Option<String>,
     187              :     /// aws region to retrieve credentials
     188            1 :     #[clap(long, default_value_t = String::new())]
     189            0 :     aws_region: String,
     190              :     /// cache for `project_info` (use `size=0` to disable)
     191              :     #[clap(long, default_value = config::ProjectInfoCacheOptions::CACHE_DEFAULT_OPTIONS)]
     192            0 :     project_info_cache: String,
     193              :     /// cache for all valid endpoints
     194              :     #[clap(long, default_value = config::EndpointCacheConfig::CACHE_DEFAULT_OPTIONS)]
     195            0 :     endpoint_cache_config: String,
     196              :     #[clap(flatten)]
     197              :     parquet_upload: ParquetUploadArgs,
     198              : 
     199              :     /// interval for backup metric collection
     200              :     #[clap(long, default_value = "10m", value_parser = humantime::parse_duration)]
     201            0 :     metric_backup_collection_interval: std::time::Duration,
     202              :     /// remote storage configuration for backup metric collection
     203              :     /// Encoded as toml (same format as pageservers), eg
     204              :     /// `{bucket_name='the-bucket',bucket_region='us-east-1',prefix_in_bucket='proxy',endpoint='http://minio:9000'}`
     205              :     #[clap(long, value_parser = remote_storage_from_toml)]
     206              :     metric_backup_collection_remote_storage: Option<RemoteStorageConfig>,
     207              :     /// chunk size for backup metric collection
     208              :     /// Size of each event is no more than 400 bytes, so 2**22 is about 200MB before the compression.
     209              :     #[clap(long, default_value = "4194304")]
     210            0 :     metric_backup_collection_chunk_size: usize,
     211              :     /// Whether to retry the connection to the compute node
     212              :     #[clap(long, default_value = config::RetryConfig::CONNECT_TO_COMPUTE_DEFAULT_VALUES)]
     213            0 :     connect_to_compute_retry: String,
     214              :     /// Whether to retry the wake_compute request
     215              :     #[clap(long, default_value = config::RetryConfig::WAKE_COMPUTE_DEFAULT_VALUES)]
     216            0 :     wake_compute_retry: String,
     217              : 
     218              :     /// Configure if this is a private access proxy for the POC: In that case the proxy will ignore the IP allowlist
     219            1 :     #[clap(long, default_value_t = false, value_parser = clap::builder::BoolishValueParser::new(), action = clap::ArgAction::Set)]
     220            0 :     is_private_access_proxy: bool,
     221              : 
     222              :     /// Configure whether all incoming requests have a Proxy Protocol V2 packet.
     223              :     // TODO(conradludgate): switch default to rejected or required once we've updated all deployments
     224            1 :     #[clap(value_enum, long, default_value_t = ProxyProtocolV2::Supported)]
     225            0 :     proxy_protocol_v2: ProxyProtocolV2,
     226              : 
     227              :     /// Time the proxy waits for the webauth session to be confirmed by the control plane.
     228              :     // TODO: rename to `console_redirect_confirmation_timeout`.
     229              :     #[clap(long, default_value = "2m", value_parser = humantime::parse_duration)]
     230            0 :     webauth_confirmation_timeout: std::time::Duration,
     231              : }
     232              : 
     233              : #[derive(clap::Args, Clone, Copy, Debug)]
     234              : struct SqlOverHttpArgs {
     235              :     /// timeout for http connection requests
     236              :     #[clap(long, default_value = "15s", value_parser = humantime::parse_duration)]
     237            0 :     sql_over_http_timeout: tokio::time::Duration,
     238              : 
     239              :     /// Whether the SQL over http pool is opt-in
     240            1 :     #[clap(long, default_value_t = true, value_parser = clap::builder::BoolishValueParser::new(), action = clap::ArgAction::Set)]
     241            0 :     sql_over_http_pool_opt_in: bool,
     242              : 
     243              :     /// How many connections to pool for each endpoint. Excess connections are discarded
     244            1 :     #[clap(long, default_value_t = 20)]
     245            0 :     sql_over_http_pool_max_conns_per_endpoint: usize,
     246              : 
     247              :     /// How many connections to pool for each endpoint. Excess connections are discarded
     248            1 :     #[clap(long, default_value_t = 20000)]
     249            0 :     sql_over_http_pool_max_total_conns: usize,
     250              : 
     251              :     /// How long pooled connections should remain idle for before closing
     252              :     #[clap(long, default_value = "5m", value_parser = humantime::parse_duration)]
     253            0 :     sql_over_http_idle_timeout: tokio::time::Duration,
     254              : 
     255              :     /// Duration each shard will wait on average before a GC sweep.
     256              :     /// A longer time will causes sweeps to take longer but will interfere less frequently.
     257              :     #[clap(long, default_value = "10m", value_parser = humantime::parse_duration)]
     258            0 :     sql_over_http_pool_gc_epoch: tokio::time::Duration,
     259              : 
     260              :     /// How many shards should the global pool have. Must be a power of two.
     261              :     /// More shards will introduce less contention for pool operations, but can
     262              :     /// increase memory used by the pool
     263            1 :     #[clap(long, default_value_t = 128)]
     264            0 :     sql_over_http_pool_shards: usize,
     265              : 
     266            1 :     #[clap(long, default_value_t = 10000)]
     267            0 :     sql_over_http_client_conn_threshold: u64,
     268              : 
     269            1 :     #[clap(long, default_value_t = 64)]
     270            0 :     sql_over_http_cancel_set_shards: usize,
     271              : 
     272            1 :     #[clap(long, default_value_t = 10 * 1024 * 1024)] // 10 MiB
     273            0 :     sql_over_http_max_request_size_bytes: usize,
     274              : 
     275            1 :     #[clap(long, default_value_t = 10 * 1024 * 1024)] // 10 MiB
     276            0 :     sql_over_http_max_response_size_bytes: usize,
     277              : }
     278              : 
     279            0 : pub async fn run() -> anyhow::Result<()> {
     280            0 :     let _logging_guard = crate::logging::init().await?;
     281            0 :     let _panic_hook_guard = utils::logging::replace_panic_hook_with_tracing_panic_hook();
     282            0 :     let _sentry_guard = init_sentry(Some(GIT_VERSION.into()), &[]);
     283            0 : 
     284            0 :     // TODO: refactor these to use labels
     285            0 :     info!("Version: {GIT_VERSION}");
     286            0 :     info!("Build_tag: {BUILD_TAG}");
     287            0 :     let neon_metrics = ::metrics::NeonMetrics::new(::metrics::BuildInfo {
     288            0 :         revision: GIT_VERSION,
     289            0 :         build_tag: BUILD_TAG,
     290            0 :     });
     291              : 
     292            0 :     let jemalloc = match crate::jemalloc::MetricRecorder::new() {
     293            0 :         Ok(t) => Some(t),
     294            0 :         Err(e) => {
     295            0 :             tracing::error!(error = ?e, "could not start jemalloc metrics loop");
     296            0 :             None
     297              :         }
     298              :     };
     299              : 
     300            0 :     let args = ProxyCliArgs::parse();
     301            0 :     let config = build_config(&args)?;
     302            0 :     let auth_backend = build_auth_backend(&args)?;
     303              : 
     304            0 :     match auth_backend {
     305            0 :         Either::Left(auth_backend) => info!("Authentication backend: {auth_backend}"),
     306            0 :         Either::Right(auth_backend) => info!("Authentication backend: {auth_backend:?}"),
     307              :     }
     308            0 :     info!("Using region: {}", args.aws_region);
     309              : 
     310              :     // TODO: untangle the config args
     311            0 :     let regional_redis_client = match (args.redis_auth_type.as_str(), &args.redis_notifications) {
     312            0 :         ("plain", redis_url) => match redis_url {
     313              :             None => {
     314            0 :                 bail!("plain auth requires redis_notifications to be set");
     315              :             }
     316            0 :             Some(url) => Some(
     317            0 :                 ConnectionWithCredentialsProvider::new_with_static_credentials(url.to_string()),
     318            0 :             ),
     319              :         },
     320            0 :         ("irsa", _) => match (&args.redis_host, args.redis_port) {
     321            0 :             (Some(host), Some(port)) => Some(
     322            0 :                 ConnectionWithCredentialsProvider::new_with_credentials_provider(
     323            0 :                     host.to_string(),
     324            0 :                     port,
     325            0 :                     elasticache::CredentialsProvider::new(
     326            0 :                         args.aws_region,
     327            0 :                         args.redis_cluster_name,
     328            0 :                         args.redis_user_id,
     329            0 :                     )
     330            0 :                     .await,
     331              :                 ),
     332              :             ),
     333              :             (None, None) => {
     334            0 :                 warn!("irsa auth requires redis-host and redis-port to be set, continuing without regional_redis_client");
     335            0 :                 None
     336              :             }
     337              :             _ => {
     338            0 :                 bail!("redis-host and redis-port must be specified together");
     339              :             }
     340              :         },
     341              :         _ => {
     342            0 :             bail!("unknown auth type given");
     343              :         }
     344              :     };
     345              : 
     346            0 :     let redis_notifications_client = if let Some(url) = args.redis_notifications {
     347            0 :         Some(ConnectionWithCredentialsProvider::new_with_static_credentials(url))
     348              :     } else {
     349            0 :         regional_redis_client.clone()
     350              :     };
     351              : 
     352              :     // Check that we can bind to address before further initialization
     353            0 :     let http_address: SocketAddr = args.http.parse()?;
     354            0 :     info!("Starting http on {http_address}");
     355            0 :     let http_listener = TcpListener::bind(http_address).await?.into_std()?;
     356              : 
     357            0 :     let mgmt_address: SocketAddr = args.mgmt.parse()?;
     358            0 :     info!("Starting mgmt on {mgmt_address}");
     359            0 :     let mgmt_listener = TcpListener::bind(mgmt_address).await?;
     360              : 
     361            0 :     let proxy_listener = if args.is_auth_broker {
     362            0 :         None
     363              :     } else {
     364            0 :         let proxy_address: SocketAddr = args.proxy.parse()?;
     365            0 :         info!("Starting proxy on {proxy_address}");
     366              : 
     367            0 :         Some(TcpListener::bind(proxy_address).await?)
     368              :     };
     369              : 
     370              :     // TODO: rename the argument to something like serverless.
     371              :     // It now covers more than just websockets, it also covers SQL over HTTP.
     372            0 :     let serverless_listener = if let Some(serverless_address) = args.wss {
     373            0 :         let serverless_address: SocketAddr = serverless_address.parse()?;
     374            0 :         info!("Starting wss on {serverless_address}");
     375            0 :         Some(TcpListener::bind(serverless_address).await?)
     376            0 :     } else if args.is_auth_broker {
     377            0 :         bail!("wss arg must be present for auth-broker")
     378              :     } else {
     379            0 :         None
     380              :     };
     381              : 
     382            0 :     let cancellation_token = CancellationToken::new();
     383            0 : 
     384            0 :     let redis_rps_limit = Vec::leak(args.redis_rps_limit.clone());
     385            0 :     RateBucketInfo::validate(redis_rps_limit)?;
     386              : 
     387            0 :     let redis_kv_client = regional_redis_client
     388            0 :         .as_ref()
     389            0 :         .map(|redis_publisher| RedisKVClient::new(redis_publisher.clone(), redis_rps_limit));
     390            0 : 
     391            0 :     // channel size should be higher than redis client limit to avoid blocking
     392            0 :     let cancel_ch_size = args.cancellation_ch_size;
     393            0 :     let (tx_cancel, rx_cancel) = tokio::sync::mpsc::channel(cancel_ch_size);
     394            0 :     let cancellation_handler = Arc::new(CancellationHandler::new(
     395            0 :         &config.connect_to_compute,
     396            0 :         Some(tx_cancel),
     397            0 :     ));
     398            0 : 
     399            0 :     // bit of a hack - find the min rps and max rps supported and turn it into
     400            0 :     // leaky bucket config instead
     401            0 :     let max = args
     402            0 :         .endpoint_rps_limit
     403            0 :         .iter()
     404            0 :         .map(|x| x.rps())
     405            0 :         .max_by(f64::total_cmp)
     406            0 :         .unwrap_or(EndpointRateLimiter::DEFAULT.max);
     407            0 :     let rps = args
     408            0 :         .endpoint_rps_limit
     409            0 :         .iter()
     410            0 :         .map(|x| x.rps())
     411            0 :         .min_by(f64::total_cmp)
     412            0 :         .unwrap_or(EndpointRateLimiter::DEFAULT.rps);
     413            0 :     let endpoint_rate_limiter = Arc::new(EndpointRateLimiter::new_with_shards(
     414            0 :         LeakyBucketConfig { rps, max },
     415            0 :         64,
     416            0 :     ));
     417            0 : 
     418            0 :     // client facing tasks. these will exit on error or on cancellation
     419            0 :     // cancellation returns Ok(())
     420            0 :     let mut client_tasks = JoinSet::new();
     421            0 :     match auth_backend {
     422            0 :         Either::Left(auth_backend) => {
     423            0 :             if let Some(proxy_listener) = proxy_listener {
     424            0 :                 client_tasks.spawn(crate::proxy::task_main(
     425            0 :                     config,
     426            0 :                     auth_backend,
     427            0 :                     proxy_listener,
     428            0 :                     cancellation_token.clone(),
     429            0 :                     cancellation_handler.clone(),
     430            0 :                     endpoint_rate_limiter.clone(),
     431            0 :                 ));
     432            0 :             }
     433              : 
     434            0 :             if let Some(serverless_listener) = serverless_listener {
     435            0 :                 client_tasks.spawn(serverless::task_main(
     436            0 :                     config,
     437            0 :                     auth_backend,
     438            0 :                     serverless_listener,
     439            0 :                     cancellation_token.clone(),
     440            0 :                     cancellation_handler.clone(),
     441            0 :                     endpoint_rate_limiter.clone(),
     442            0 :                 ));
     443            0 :             }
     444              :         }
     445            0 :         Either::Right(auth_backend) => {
     446            0 :             if let Some(proxy_listener) = proxy_listener {
     447            0 :                 client_tasks.spawn(crate::console_redirect_proxy::task_main(
     448            0 :                     config,
     449            0 :                     auth_backend,
     450            0 :                     proxy_listener,
     451            0 :                     cancellation_token.clone(),
     452            0 :                     cancellation_handler.clone(),
     453            0 :                 ));
     454            0 :             }
     455              :         }
     456              :     }
     457              : 
     458            0 :     client_tasks.spawn(crate::context::parquet::worker(
     459            0 :         cancellation_token.clone(),
     460            0 :         args.parquet_upload,
     461            0 :     ));
     462            0 : 
     463            0 :     // maintenance tasks. these never return unless there's an error
     464            0 :     let mut maintenance_tasks = JoinSet::new();
     465            0 :     maintenance_tasks.spawn(crate::signals::handle(cancellation_token.clone(), || {}));
     466            0 :     maintenance_tasks.spawn(http::health_server::task_main(
     467            0 :         http_listener,
     468            0 :         AppMetrics {
     469            0 :             jemalloc,
     470            0 :             neon_metrics,
     471            0 :             proxy: crate::metrics::Metrics::get(),
     472            0 :         },
     473            0 :     ));
     474            0 :     maintenance_tasks.spawn(control_plane::mgmt::task_main(mgmt_listener));
     475              : 
     476            0 :     if let Some(metrics_config) = &config.metric_collection {
     477            0 :         // TODO: Add gc regardles of the metric collection being enabled.
     478            0 :         maintenance_tasks.spawn(usage_metrics::task_main(metrics_config));
     479            0 :     }
     480              : 
     481              :     #[cfg_attr(not(any(test, feature = "testing")), expect(irrefutable_let_patterns))]
     482            0 :     if let Either::Left(auth::Backend::ControlPlane(api, ())) = &auth_backend {
     483            0 :         if let crate::control_plane::client::ControlPlaneClient::ProxyV1(api) = &**api {
     484            0 :             match (redis_notifications_client, regional_redis_client.clone()) {
     485            0 :                 (None, None) => {}
     486            0 :                 (client1, client2) => {
     487            0 :                     let cache = api.caches.project_info.clone();
     488            0 :                     if let Some(client) = client1 {
     489            0 :                         maintenance_tasks.spawn(notifications::task_main(
     490            0 :                             client,
     491            0 :                             cache.clone(),
     492            0 :                             args.region.clone(),
     493            0 :                         ));
     494            0 :                     }
     495            0 :                     if let Some(client) = client2 {
     496            0 :                         maintenance_tasks.spawn(notifications::task_main(
     497            0 :                             client,
     498            0 :                             cache.clone(),
     499            0 :                             args.region.clone(),
     500            0 :                         ));
     501            0 :                     }
     502            0 :                     maintenance_tasks.spawn(async move { cache.clone().gc_worker().await });
     503            0 :                 }
     504              :             }
     505              : 
     506            0 :             if let Some(mut redis_kv_client) = redis_kv_client {
     507            0 :                 maintenance_tasks.spawn(async move {
     508            0 :                     redis_kv_client.try_connect().await?;
     509            0 :                     handle_cancel_messages(&mut redis_kv_client, rx_cancel).await
     510            0 :                 });
     511            0 :             }
     512              : 
     513            0 :             if let Some(regional_redis_client) = regional_redis_client {
     514            0 :                 let cache = api.caches.endpoints_cache.clone();
     515            0 :                 let con = regional_redis_client;
     516            0 :                 let span = tracing::info_span!("endpoints_cache");
     517            0 :                 maintenance_tasks.spawn(
     518            0 :                     async move { cache.do_read(con, cancellation_token.clone()).await }
     519            0 :                         .instrument(span),
     520            0 :                 );
     521            0 :             }
     522            0 :         }
     523            0 :     }
     524              : 
     525              :     let maintenance = loop {
     526              :         // get one complete task
     527            0 :         match futures::future::select(
     528            0 :             pin!(maintenance_tasks.join_next()),
     529            0 :             pin!(client_tasks.join_next()),
     530            0 :         )
     531            0 :         .await
     532              :         {
     533              :             // exit immediately on maintenance task completion
     534            0 :             Either::Left((Some(res), _)) => break crate::error::flatten_err(res)?,
     535              :             // exit with error immediately if all maintenance tasks have ceased (should be caught by branch above)
     536            0 :             Either::Left((None, _)) => bail!("no maintenance tasks running. invalid state"),
     537              :             // exit immediately on client task error
     538            0 :             Either::Right((Some(res), _)) => crate::error::flatten_err(res)?,
     539              :             // exit if all our client tasks have shutdown gracefully
     540            0 :             Either::Right((None, _)) => return Ok(()),
     541              :         }
     542              :     };
     543              : 
     544              :     // maintenance tasks return Infallible success values, this is an impossible value
     545              :     // so this match statically ensures that there are no possibilities for that value
     546              :     match maintenance {}
     547            0 : }
     548              : 
     549              : /// ProxyConfig is created at proxy startup, and lives forever.
     550            0 : fn build_config(args: &ProxyCliArgs) -> anyhow::Result<&'static ProxyConfig> {
     551            0 :     let thread_pool = ThreadPool::new(args.scram_thread_pool_size);
     552            0 :     Metrics::install(thread_pool.metrics.clone());
     553              : 
     554            0 :     let tls_config = match (&args.tls_key, &args.tls_cert) {
     555            0 :         (Some(key_path), Some(cert_path)) => Some(config::configure_tls(
     556            0 :             key_path,
     557            0 :             cert_path,
     558            0 :             args.certs_dir.as_ref(),
     559            0 :             args.allow_tls_keylogfile,
     560            0 :         )?),
     561            0 :         (None, None) => None,
     562            0 :         _ => bail!("either both or neither tls-key and tls-cert must be specified"),
     563              :     };
     564              : 
     565            0 :     let backup_metric_collection_config = config::MetricBackupCollectionConfig {
     566            0 :         remote_storage_config: args.metric_backup_collection_remote_storage.clone(),
     567            0 :         chunk_size: args.metric_backup_collection_chunk_size,
     568            0 :     };
     569              : 
     570            0 :     let metric_collection = match (
     571            0 :         &args.metric_collection_endpoint,
     572            0 :         &args.metric_collection_interval,
     573              :     ) {
     574            0 :         (Some(endpoint), Some(interval)) => Some(config::MetricCollectionConfig {
     575            0 :             endpoint: endpoint.parse()?,
     576            0 :             interval: humantime::parse_duration(interval)?,
     577            0 :             backup_metric_collection_config,
     578              :         }),
     579            0 :         (None, None) => None,
     580            0 :         _ => bail!(
     581            0 :             "either both or neither metric-collection-endpoint \
     582            0 :              and metric-collection-interval must be specified"
     583            0 :         ),
     584              :     };
     585              : 
     586              :     let config::ConcurrencyLockOptions {
     587            0 :         shards,
     588            0 :         limiter,
     589            0 :         epoch,
     590            0 :         timeout,
     591            0 :     } = args.connect_compute_lock.parse()?;
     592            0 :     info!(
     593              :         ?limiter,
     594              :         shards,
     595              :         ?epoch,
     596            0 :         "Using NodeLocks (connect_compute)"
     597              :     );
     598            0 :     let connect_compute_locks = control_plane::locks::ApiLocks::new(
     599            0 :         "connect_compute_lock",
     600            0 :         limiter,
     601            0 :         shards,
     602            0 :         timeout,
     603            0 :         epoch,
     604            0 :         &Metrics::get().proxy.connect_compute_lock,
     605            0 :     );
     606            0 : 
     607            0 :     let http_config = HttpConfig {
     608            0 :         accept_websockets: !args.is_auth_broker,
     609            0 :         pool_options: GlobalConnPoolOptions {
     610            0 :             max_conns_per_endpoint: args.sql_over_http.sql_over_http_pool_max_conns_per_endpoint,
     611            0 :             gc_epoch: args.sql_over_http.sql_over_http_pool_gc_epoch,
     612            0 :             pool_shards: args.sql_over_http.sql_over_http_pool_shards,
     613            0 :             idle_timeout: args.sql_over_http.sql_over_http_idle_timeout,
     614            0 :             opt_in: args.sql_over_http.sql_over_http_pool_opt_in,
     615            0 :             max_total_conns: args.sql_over_http.sql_over_http_pool_max_total_conns,
     616            0 :         },
     617            0 :         cancel_set: CancelSet::new(args.sql_over_http.sql_over_http_cancel_set_shards),
     618            0 :         client_conn_threshold: args.sql_over_http.sql_over_http_client_conn_threshold,
     619            0 :         max_request_size_bytes: args.sql_over_http.sql_over_http_max_request_size_bytes,
     620            0 :         max_response_size_bytes: args.sql_over_http.sql_over_http_max_response_size_bytes,
     621            0 :     };
     622            0 :     let authentication_config = AuthenticationConfig {
     623            0 :         jwks_cache: JwkCache::default(),
     624            0 :         thread_pool,
     625            0 :         scram_protocol_timeout: args.scram_protocol_timeout,
     626            0 :         rate_limiter_enabled: args.auth_rate_limit_enabled,
     627            0 :         rate_limiter: AuthRateLimiter::new(args.auth_rate_limit.clone()),
     628            0 :         rate_limit_ip_subnet: args.auth_rate_limit_ip_subnet,
     629            0 :         ip_allowlist_check_enabled: !args.is_private_access_proxy,
     630            0 :         is_vpc_acccess_proxy: args.is_private_access_proxy,
     631            0 :         is_auth_broker: args.is_auth_broker,
     632            0 :         accept_jwts: args.is_auth_broker,
     633            0 :         console_redirect_confirmation_timeout: args.webauth_confirmation_timeout,
     634            0 :     };
     635              : 
     636            0 :     let compute_config = ComputeConfig {
     637            0 :         retry: config::RetryConfig::parse(&args.connect_to_compute_retry)?,
     638            0 :         tls: Arc::new(compute_client_config_with_root_certs()?),
     639            0 :         timeout: Duration::from_secs(2),
     640              :     };
     641              : 
     642            0 :     let config = ProxyConfig {
     643            0 :         tls_config,
     644            0 :         metric_collection,
     645            0 :         http_config,
     646            0 :         authentication_config,
     647            0 :         proxy_protocol_v2: args.proxy_protocol_v2,
     648            0 :         handshake_timeout: args.handshake_timeout,
     649            0 :         region: args.region.clone(),
     650            0 :         wake_compute_retry_config: config::RetryConfig::parse(&args.wake_compute_retry)?,
     651            0 :         connect_compute_locks,
     652            0 :         connect_to_compute: compute_config,
     653            0 :     };
     654            0 : 
     655            0 :     let config = Box::leak(Box::new(config));
     656            0 : 
     657            0 :     tokio::spawn(config.connect_compute_locks.garbage_collect_worker());
     658            0 : 
     659            0 :     Ok(config)
     660            0 : }
     661              : 
     662              : /// auth::Backend is created at proxy startup, and lives forever.
     663            0 : fn build_auth_backend(
     664            0 :     args: &ProxyCliArgs,
     665            0 : ) -> anyhow::Result<Either<&'static auth::Backend<'static, ()>, &'static ConsoleRedirectBackend>> {
     666            0 :     match &args.auth_backend {
     667              :         AuthBackendType::ControlPlaneV1 => {
     668            0 :             let wake_compute_cache_config: CacheOptions = args.wake_compute_cache.parse()?;
     669            0 :             let project_info_cache_config: ProjectInfoCacheOptions =
     670            0 :                 args.project_info_cache.parse()?;
     671            0 :             let endpoint_cache_config: config::EndpointCacheConfig =
     672            0 :                 args.endpoint_cache_config.parse()?;
     673              : 
     674            0 :             info!("Using NodeInfoCache (wake_compute) with options={wake_compute_cache_config:?}");
     675            0 :             info!(
     676            0 :                 "Using AllowedIpsCache (wake_compute) with options={project_info_cache_config:?}"
     677              :             );
     678            0 :             info!("Using EndpointCacheConfig with options={endpoint_cache_config:?}");
     679            0 :             let caches = Box::leak(Box::new(control_plane::caches::ApiCaches::new(
     680            0 :                 wake_compute_cache_config,
     681            0 :                 project_info_cache_config,
     682            0 :                 endpoint_cache_config,
     683            0 :             )));
     684              : 
     685              :             let config::ConcurrencyLockOptions {
     686            0 :                 shards,
     687            0 :                 limiter,
     688            0 :                 epoch,
     689            0 :                 timeout,
     690            0 :             } = args.wake_compute_lock.parse()?;
     691            0 :             info!(?limiter, shards, ?epoch, "Using NodeLocks (wake_compute)");
     692            0 :             let locks = Box::leak(Box::new(control_plane::locks::ApiLocks::new(
     693            0 :                 "wake_compute_lock",
     694            0 :                 limiter,
     695            0 :                 shards,
     696            0 :                 timeout,
     697            0 :                 epoch,
     698            0 :                 &Metrics::get().wake_compute_lock,
     699            0 :             )));
     700            0 :             tokio::spawn(locks.garbage_collect_worker());
     701              : 
     702            0 :             let url: crate::url::ApiUrl = args.auth_endpoint.parse()?;
     703              : 
     704            0 :             let endpoint = http::Endpoint::new(url, http::new_client());
     705            0 : 
     706            0 :             let mut wake_compute_rps_limit = args.wake_compute_limit.clone();
     707            0 :             RateBucketInfo::validate(&mut wake_compute_rps_limit)?;
     708            0 :             let wake_compute_endpoint_rate_limiter =
     709            0 :                 Arc::new(WakeComputeRateLimiter::new(wake_compute_rps_limit));
     710            0 : 
     711            0 :             let api = control_plane::client::cplane_proxy_v1::NeonControlPlaneClient::new(
     712            0 :                 endpoint,
     713            0 :                 args.control_plane_token.clone(),
     714            0 :                 caches,
     715            0 :                 locks,
     716            0 :                 wake_compute_endpoint_rate_limiter,
     717            0 :             );
     718            0 : 
     719            0 :             let api = control_plane::client::ControlPlaneClient::ProxyV1(api);
     720            0 :             let auth_backend = auth::Backend::ControlPlane(MaybeOwned::Owned(api), ());
     721            0 :             let config = Box::leak(Box::new(auth_backend));
     722            0 : 
     723            0 :             Ok(Either::Left(config))
     724              :         }
     725              : 
     726              :         #[cfg(any(test, feature = "testing"))]
     727              :         AuthBackendType::Postgres => {
     728            0 :             let url = args.auth_endpoint.parse()?;
     729            0 :             let api = control_plane::client::mock::MockControlPlane::new(
     730            0 :                 url,
     731            0 :                 !args.is_private_access_proxy,
     732            0 :             );
     733            0 :             let api = control_plane::client::ControlPlaneClient::PostgresMock(api);
     734            0 : 
     735            0 :             let auth_backend = auth::Backend::ControlPlane(MaybeOwned::Owned(api), ());
     736            0 : 
     737            0 :             let config = Box::leak(Box::new(auth_backend));
     738            0 : 
     739            0 :             Ok(Either::Left(config))
     740              :         }
     741              : 
     742              :         AuthBackendType::ConsoleRedirect => {
     743            0 :             let wake_compute_cache_config: CacheOptions = args.wake_compute_cache.parse()?;
     744            0 :             let project_info_cache_config: ProjectInfoCacheOptions =
     745            0 :                 args.project_info_cache.parse()?;
     746            0 :             let endpoint_cache_config: config::EndpointCacheConfig =
     747            0 :                 args.endpoint_cache_config.parse()?;
     748              : 
     749            0 :             info!("Using NodeInfoCache (wake_compute) with options={wake_compute_cache_config:?}");
     750            0 :             info!(
     751            0 :                 "Using AllowedIpsCache (wake_compute) with options={project_info_cache_config:?}"
     752              :             );
     753            0 :             info!("Using EndpointCacheConfig with options={endpoint_cache_config:?}");
     754            0 :             let caches = Box::leak(Box::new(control_plane::caches::ApiCaches::new(
     755            0 :                 wake_compute_cache_config,
     756            0 :                 project_info_cache_config,
     757            0 :                 endpoint_cache_config,
     758            0 :             )));
     759              : 
     760              :             let config::ConcurrencyLockOptions {
     761            0 :                 shards,
     762            0 :                 limiter,
     763            0 :                 epoch,
     764            0 :                 timeout,
     765            0 :             } = args.wake_compute_lock.parse()?;
     766            0 :             info!(?limiter, shards, ?epoch, "Using NodeLocks (wake_compute)");
     767            0 :             let locks = Box::leak(Box::new(control_plane::locks::ApiLocks::new(
     768            0 :                 "wake_compute_lock",
     769            0 :                 limiter,
     770            0 :                 shards,
     771            0 :                 timeout,
     772            0 :                 epoch,
     773            0 :                 &Metrics::get().wake_compute_lock,
     774            0 :             )));
     775              : 
     776            0 :             let url = args.uri.clone().parse()?;
     777            0 :             let ep_url: crate::url::ApiUrl = args.auth_endpoint.parse()?;
     778            0 :             let endpoint = http::Endpoint::new(ep_url, http::new_client());
     779            0 :             let mut wake_compute_rps_limit = args.wake_compute_limit.clone();
     780            0 :             RateBucketInfo::validate(&mut wake_compute_rps_limit)?;
     781            0 :             let wake_compute_endpoint_rate_limiter =
     782            0 :                 Arc::new(WakeComputeRateLimiter::new(wake_compute_rps_limit));
     783            0 : 
     784            0 :             // Since we use only get_allowed_ips_and_secret() wake_compute_endpoint_rate_limiter
     785            0 :             // and locks are not used in ConsoleRedirectBackend,
     786            0 :             // but they are required by the NeonControlPlaneClient
     787            0 :             let api = control_plane::client::cplane_proxy_v1::NeonControlPlaneClient::new(
     788            0 :                 endpoint,
     789            0 :                 args.control_plane_token.clone(),
     790            0 :                 caches,
     791            0 :                 locks,
     792            0 :                 wake_compute_endpoint_rate_limiter,
     793            0 :             );
     794            0 : 
     795            0 :             let backend = ConsoleRedirectBackend::new(url, api);
     796            0 :             let config = Box::leak(Box::new(backend));
     797            0 : 
     798            0 :             Ok(Either::Right(config))
     799              :         }
     800              :     }
     801            0 : }
     802              : 
     803              : #[cfg(test)]
     804              : mod tests {
     805              :     use std::time::Duration;
     806              : 
     807              :     use clap::Parser;
     808              : 
     809              :     use crate::rate_limiter::RateBucketInfo;
     810              : 
     811              :     #[test]
     812            1 :     fn parse_endpoint_rps_limit() {
     813            1 :         let config = super::ProxyCliArgs::parse_from([
     814            1 :             "proxy",
     815            1 :             "--endpoint-rps-limit",
     816            1 :             "100@1s",
     817            1 :             "--endpoint-rps-limit",
     818            1 :             "20@30s",
     819            1 :         ]);
     820            1 : 
     821            1 :         assert_eq!(
     822            1 :             config.endpoint_rps_limit,
     823            1 :             vec![
     824            1 :                 RateBucketInfo::new(100, Duration::from_secs(1)),
     825            1 :                 RateBucketInfo::new(20, Duration::from_secs(30)),
     826            1 :             ]
     827            1 :         );
     828            1 :     }
     829              : }
        

Generated by: LCOV version 2.1-beta