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