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