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 futures::future::Either;
9 : use proxy::auth;
10 : use proxy::auth::backend::MaybeOwned;
11 : use proxy::cancellation::CancelMap;
12 : use proxy::cancellation::CancellationHandler;
13 : use proxy::config::remote_storage_from_toml;
14 : use proxy::config::AuthenticationConfig;
15 : use proxy::config::CacheOptions;
16 : use proxy::config::HttpConfig;
17 : use proxy::config::ProjectInfoCacheOptions;
18 : use proxy::console;
19 : use proxy::context::parquet::ParquetUploadArgs;
20 : use proxy::http;
21 : use proxy::metrics::NUM_CANCELLATION_REQUESTS_SOURCE_FROM_CLIENT;
22 : use proxy::rate_limiter::AuthRateLimiter;
23 : use proxy::rate_limiter::EndpointRateLimiter;
24 : use proxy::rate_limiter::RateBucketInfo;
25 : use proxy::rate_limiter::RateLimiterConfig;
26 : use proxy::redis::cancellation_publisher::RedisPublisherClient;
27 : use proxy::redis::connection_with_credentials_provider::ConnectionWithCredentialsProvider;
28 : use proxy::redis::elasticache;
29 : use proxy::redis::notifications;
30 : use proxy::serverless::GlobalConnPoolOptions;
31 : use proxy::usage_metrics;
32 :
33 : use anyhow::bail;
34 : use proxy::config::{self, ProxyConfig};
35 : use proxy::serverless;
36 : use std::net::SocketAddr;
37 : use std::pin::pin;
38 : use std::sync::Arc;
39 : use tokio::net::TcpListener;
40 : use tokio::sync::Mutex;
41 : use tokio::task::JoinSet;
42 : use tokio_util::sync::CancellationToken;
43 : use tracing::info;
44 : use tracing::warn;
45 : use utils::{project_build_tag, project_git_version, sentry_init::init_sentry};
46 :
47 : project_git_version!(GIT_VERSION);
48 : project_build_tag!(BUILD_TAG);
49 :
50 : use clap::{Parser, ValueEnum};
51 :
52 : #[global_allocator]
53 5314 : static GLOBAL: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc;
54 :
55 20 : #[derive(Clone, Debug, ValueEnum)]
56 : enum AuthBackend {
57 : Console,
58 : #[cfg(feature = "testing")]
59 : Postgres,
60 : Link,
61 : }
62 :
63 : /// Neon proxy/router
64 4 : #[derive(Parser)]
65 : #[command(version = GIT_VERSION, about)]
66 : struct ProxyCliArgs {
67 : /// Name of the region this proxy is deployed in
68 2 : #[clap(long, default_value_t = String::new())]
69 0 : region: String,
70 : /// listen for incoming client connections on ip:port
71 : #[clap(short, long, default_value = "127.0.0.1:4432")]
72 0 : proxy: String,
73 2 : #[clap(value_enum, long, default_value_t = AuthBackend::Link)]
74 0 : auth_backend: AuthBackend,
75 : /// listen for management callback connection on ip:port
76 : #[clap(short, long, default_value = "127.0.0.1:7000")]
77 0 : mgmt: String,
78 : /// listen for incoming http connections (metrics, etc) on ip:port
79 : #[clap(long, default_value = "127.0.0.1:7001")]
80 0 : http: String,
81 : /// listen for incoming wss connections on ip:port
82 : #[clap(long)]
83 : wss: Option<String>,
84 : /// redirect unauthenticated users to the given uri in case of link auth
85 : #[clap(short, long, default_value = "http://localhost:3000/psql_session/")]
86 0 : uri: String,
87 : /// cloud API endpoint for authenticating users
88 : #[clap(
89 : short,
90 : long,
91 : default_value = "http://localhost:3000/authenticate_proxy_request/"
92 : )]
93 0 : auth_endpoint: String,
94 : /// path to TLS key for client postgres connections
95 : ///
96 : /// tls-key and tls-cert are for backwards compatibility, we can put all certs in one dir
97 : #[clap(short = 'k', long, alias = "ssl-key")]
98 : tls_key: Option<String>,
99 : /// path to TLS cert for client postgres connections
100 : ///
101 : /// tls-key and tls-cert are for backwards compatibility, we can put all certs in one dir
102 : #[clap(short = 'c', long, alias = "ssl-cert")]
103 : tls_cert: Option<String>,
104 : /// path to directory with TLS certificates for client postgres connections
105 : #[clap(long)]
106 : certs_dir: Option<String>,
107 : /// timeout for the TLS handshake
108 : #[clap(long, default_value = "15s", value_parser = humantime::parse_duration)]
109 0 : handshake_timeout: tokio::time::Duration,
110 : /// http endpoint to receive periodic metric updates
111 : #[clap(long)]
112 : metric_collection_endpoint: Option<String>,
113 : /// how often metrics should be sent to a collection endpoint
114 : #[clap(long)]
115 : metric_collection_interval: Option<String>,
116 : /// cache for `wake_compute` api method (use `size=0` to disable)
117 : #[clap(long, default_value = config::CacheOptions::CACHE_DEFAULT_OPTIONS)]
118 0 : wake_compute_cache: String,
119 : /// lock for `wake_compute` api method. example: "shards=32,permits=4,epoch=10m,timeout=1s". (use `permits=0` to disable).
120 : #[clap(long, default_value = config::WakeComputeLockOptions::DEFAULT_OPTIONS_WAKE_COMPUTE_LOCK)]
121 0 : wake_compute_lock: String,
122 : /// Allow self-signed certificates for compute nodes (for testing)
123 2 : #[clap(long, default_value_t = false, value_parser = clap::builder::BoolishValueParser::new(), action = clap::ArgAction::Set)]
124 0 : allow_self_signed_compute: bool,
125 : #[clap(flatten)]
126 : sql_over_http: SqlOverHttpArgs,
127 : /// timeout for scram authentication protocol
128 : #[clap(long, default_value = "15s", value_parser = humantime::parse_duration)]
129 0 : scram_protocol_timeout: tokio::time::Duration,
130 : /// Require that all incoming requests have a Proxy Protocol V2 packet **and** have an IP address associated.
131 2 : #[clap(long, default_value_t = false, value_parser = clap::builder::BoolishValueParser::new(), action = clap::ArgAction::Set)]
132 0 : require_client_ip: bool,
133 : /// Disable dynamic rate limiter and store the metrics to ensure its production behaviour.
134 2 : #[clap(long, default_value_t = false, value_parser = clap::builder::BoolishValueParser::new(), action = clap::ArgAction::Set)]
135 0 : disable_dynamic_rate_limiter: bool,
136 : /// Rate limit algorithm. Makes sense only if `disable_rate_limiter` is `false`.
137 2 : #[clap(value_enum, long, default_value_t = proxy::rate_limiter::RateLimitAlgorithm::Aimd)]
138 0 : rate_limit_algorithm: proxy::rate_limiter::RateLimitAlgorithm,
139 : /// Timeout for rate limiter. If it didn't manage to aquire a permit in this time, it will return an error.
140 : #[clap(long, default_value = "15s", value_parser = humantime::parse_duration)]
141 0 : rate_limiter_timeout: tokio::time::Duration,
142 : /// Endpoint rate limiter max number of requests per second.
143 : ///
144 : /// Provided in the form '<Requests Per Second>@<Bucket Duration Size>'.
145 : /// Can be given multiple times for different bucket sizes.
146 10 : #[clap(long, default_values_t = RateBucketInfo::DEFAULT_ENDPOINT_SET)]
147 2 : endpoint_rps_limit: Vec<RateBucketInfo>,
148 : /// Whether the auth rate limiter actually takes effect (for testing)
149 2 : #[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 10 : #[clap(long, default_values_t = RateBucketInfo::DEFAULT_AUTH_SET)]
153 2 : auth_rate_limit: Vec<RateBucketInfo>,
154 : /// Redis rate limiter max number of requests per second.
155 10 : #[clap(long, default_values_t = RateBucketInfo::DEFAULT_ENDPOINT_SET)]
156 2 : redis_rps_limit: Vec<RateBucketInfo>,
157 : /// Initial limit for dynamic rate limiter. Makes sense only if `rate_limit_algorithm` is *not* `None`.
158 2 : #[clap(long, default_value_t = 100)]
159 0 : initial_limit: usize,
160 : #[clap(flatten)]
161 : aimd_config: proxy::rate_limiter::AimdConfig,
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 : /// disable ip check for http requests. If it is too time consuming, it could be turned off.
169 2 : #[clap(long, default_value_t = false, value_parser = clap::builder::BoolishValueParser::new(), action = clap::ArgAction::Set)]
170 0 : disable_ip_check_for_http: bool,
171 : /// redis url for notifications (if empty, redis_host:port will be used for both notifications and streaming connections)
172 : #[clap(long)]
173 : redis_notifications: Option<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 2 : #[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 :
193 : #[clap(flatten)]
194 : parquet_upload: ParquetUploadArgs,
195 :
196 : /// interval for backup metric collection
197 : #[clap(long, default_value = "10m", value_parser = humantime::parse_duration)]
198 0 : metric_backup_collection_interval: std::time::Duration,
199 : /// remote storage configuration for backup metric collection
200 : /// Encoded as toml (same format as pageservers), eg
201 : /// `{bucket_name='the-bucket',bucket_region='us-east-1',prefix_in_bucket='proxy',endpoint='http://minio:9000'}`
202 : #[clap(long, default_value = "{}")]
203 0 : metric_backup_collection_remote_storage: String,
204 : /// chunk size for backup metric collection
205 : /// Size of each event is no more than 400 bytes, so 2**22 is about 200MB before the compression.
206 : #[clap(long, default_value = "4194304")]
207 0 : metric_backup_collection_chunk_size: usize,
208 : }
209 :
210 4 : #[derive(clap::Args, Clone, Copy, Debug)]
211 : struct SqlOverHttpArgs {
212 : /// timeout for http connection requests
213 : #[clap(long, default_value = "15s", value_parser = humantime::parse_duration)]
214 0 : sql_over_http_timeout: tokio::time::Duration,
215 :
216 : /// Whether the SQL over http pool is opt-in
217 2 : #[clap(long, default_value_t = true, value_parser = clap::builder::BoolishValueParser::new(), action = clap::ArgAction::Set)]
218 0 : sql_over_http_pool_opt_in: bool,
219 :
220 : /// How many connections to pool for each endpoint. Excess connections are discarded
221 2 : #[clap(long, default_value_t = 20)]
222 0 : sql_over_http_pool_max_conns_per_endpoint: usize,
223 :
224 : /// How many connections to pool for each endpoint. Excess connections are discarded
225 2 : #[clap(long, default_value_t = 20000)]
226 0 : sql_over_http_pool_max_total_conns: usize,
227 :
228 : /// How long pooled connections should remain idle for before closing
229 : #[clap(long, default_value = "5m", value_parser = humantime::parse_duration)]
230 0 : sql_over_http_idle_timeout: tokio::time::Duration,
231 :
232 : /// Duration each shard will wait on average before a GC sweep.
233 : /// A longer time will causes sweeps to take longer but will interfere less frequently.
234 : #[clap(long, default_value = "10m", value_parser = humantime::parse_duration)]
235 0 : sql_over_http_pool_gc_epoch: tokio::time::Duration,
236 :
237 : /// How many shards should the global pool have. Must be a power of two.
238 : /// More shards will introduce less contention for pool operations, but can
239 : /// increase memory used by the pool
240 2 : #[clap(long, default_value_t = 128)]
241 0 : sql_over_http_pool_shards: usize,
242 : }
243 :
244 : #[tokio::main]
245 0 : async fn main() -> anyhow::Result<()> {
246 0 : let _logging_guard = proxy::logging::init().await?;
247 0 : let _panic_hook_guard = utils::logging::replace_panic_hook_with_tracing_panic_hook();
248 0 : let _sentry_guard = init_sentry(Some(GIT_VERSION.into()), &[]);
249 0 :
250 0 : info!("Version: {GIT_VERSION}");
251 0 : info!("Build_tag: {BUILD_TAG}");
252 0 : ::metrics::set_build_info_metric(GIT_VERSION, BUILD_TAG);
253 0 :
254 0 : match proxy::jemalloc::MetricRecorder::new(prometheus::default_registry()) {
255 0 : Ok(t) => {
256 0 : t.start();
257 0 : }
258 0 : Err(e) => tracing::error!(error = ?e, "could not start jemalloc metrics loop"),
259 0 : }
260 0 :
261 0 : let args = ProxyCliArgs::parse();
262 0 : let config = build_config(&args)?;
263 0 :
264 0 : info!("Authentication backend: {}", config.auth_backend);
265 0 : info!("Using region: {}", config.aws_region);
266 0 :
267 0 : let region_provider = RegionProviderChain::default_provider().or_else(&*config.aws_region); // Replace with your Redis region if needed
268 0 : let provider_conf =
269 0 : ProviderConfig::without_region().with_region(region_provider.region().await);
270 0 : let aws_credentials_provider = {
271 0 : // uses "AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY"
272 0 : CredentialsProviderChain::first_try("env", EnvironmentVariableCredentialsProvider::new())
273 0 : // uses "AWS_PROFILE" / `aws sso login --profile <profile>`
274 0 : .or_else(
275 0 : "profile-sso",
276 0 : ProfileFileCredentialsProvider::builder()
277 0 : .configure(&provider_conf)
278 0 : .build(),
279 0 : )
280 0 : // uses "AWS_WEB_IDENTITY_TOKEN_FILE", "AWS_ROLE_ARN", "AWS_ROLE_SESSION_NAME"
281 0 : // needed to access remote extensions bucket
282 0 : .or_else(
283 0 : "token",
284 0 : WebIdentityTokenCredentialsProvider::builder()
285 0 : .configure(&provider_conf)
286 0 : .build(),
287 0 : )
288 0 : // uses imds v2
289 0 : .or_else("imds", ImdsCredentialsProvider::builder().build())
290 0 : };
291 0 : let elasticache_credentials_provider = Arc::new(elasticache::CredentialsProvider::new(
292 0 : elasticache::AWSIRSAConfig::new(
293 0 : config.aws_region.clone(),
294 0 : args.redis_cluster_name,
295 0 : args.redis_user_id,
296 0 : ),
297 0 : aws_credentials_provider,
298 0 : ));
299 0 : let redis_notifications_client =
300 0 : match (args.redis_notifications, (args.redis_host, args.redis_port)) {
301 0 : (Some(url), _) => {
302 0 : info!("Starting redis notifications listener ({url})");
303 0 : Some(ConnectionWithCredentialsProvider::new_with_static_credentials(url))
304 0 : }
305 0 : (None, (Some(host), Some(port))) => Some(
306 0 : ConnectionWithCredentialsProvider::new_with_credentials_provider(
307 0 : host,
308 0 : port,
309 0 : elasticache_credentials_provider.clone(),
310 0 : ),
311 0 : ),
312 0 : (None, (None, None)) => {
313 0 : warn!("Redis is disabled");
314 0 : None
315 0 : }
316 0 : _ => {
317 0 : bail!("redis-host and redis-port must be specified together");
318 0 : }
319 0 : };
320 0 :
321 0 : // Check that we can bind to address before further initialization
322 0 : let http_address: SocketAddr = args.http.parse()?;
323 0 : info!("Starting http on {http_address}");
324 0 : let http_listener = TcpListener::bind(http_address).await?.into_std()?;
325 0 :
326 0 : let mgmt_address: SocketAddr = args.mgmt.parse()?;
327 0 : info!("Starting mgmt on {mgmt_address}");
328 0 : let mgmt_listener = TcpListener::bind(mgmt_address).await?;
329 0 :
330 0 : let proxy_address: SocketAddr = args.proxy.parse()?;
331 0 : info!("Starting proxy on {proxy_address}");
332 0 : let proxy_listener = TcpListener::bind(proxy_address).await?;
333 0 : let cancellation_token = CancellationToken::new();
334 0 :
335 0 : let endpoint_rate_limiter = Arc::new(EndpointRateLimiter::new(&config.endpoint_rps_limit));
336 0 : let cancel_map = CancelMap::default();
337 0 :
338 0 : // let redis_notifications_client = redis_notifications_client.map(|x| Box::leak(Box::new(x)));
339 0 : let redis_publisher = match &redis_notifications_client {
340 0 : Some(redis_publisher) => Some(Arc::new(Mutex::new(RedisPublisherClient::new(
341 0 : redis_publisher.clone(),
342 0 : args.region.clone(),
343 0 : &config.redis_rps_limit,
344 0 : )?))),
345 0 : None => None,
346 0 : };
347 0 : let cancellation_handler = Arc::new(CancellationHandler::<
348 0 : Option<Arc<tokio::sync::Mutex<RedisPublisherClient>>>,
349 0 : >::new(
350 0 : cancel_map.clone(),
351 0 : redis_publisher,
352 0 : NUM_CANCELLATION_REQUESTS_SOURCE_FROM_CLIENT,
353 0 : ));
354 0 :
355 0 : // client facing tasks. these will exit on error or on cancellation
356 0 : // cancellation returns Ok(())
357 0 : let mut client_tasks = JoinSet::new();
358 0 : client_tasks.spawn(proxy::proxy::task_main(
359 0 : config,
360 0 : proxy_listener,
361 0 : cancellation_token.clone(),
362 0 : endpoint_rate_limiter.clone(),
363 0 : cancellation_handler.clone(),
364 0 : ));
365 0 :
366 0 : // TODO: rename the argument to something like serverless.
367 0 : // It now covers more than just websockets, it also covers SQL over HTTP.
368 0 : if let Some(serverless_address) = args.wss {
369 0 : let serverless_address: SocketAddr = serverless_address.parse()?;
370 0 : info!("Starting wss on {serverless_address}");
371 0 : let serverless_listener = TcpListener::bind(serverless_address).await?;
372 0 :
373 0 : client_tasks.spawn(serverless::task_main(
374 0 : config,
375 0 : serverless_listener,
376 0 : cancellation_token.clone(),
377 0 : endpoint_rate_limiter.clone(),
378 0 : cancellation_handler.clone(),
379 0 : ));
380 0 : }
381 0 :
382 0 : client_tasks.spawn(proxy::context::parquet::worker(
383 0 : cancellation_token.clone(),
384 0 : args.parquet_upload,
385 0 : ));
386 0 :
387 0 : // maintenance tasks. these never return unless there's an error
388 0 : let mut maintenance_tasks = JoinSet::new();
389 0 : maintenance_tasks.spawn(proxy::handle_signals(cancellation_token.clone()));
390 0 : maintenance_tasks.spawn(http::health_server::task_main(http_listener));
391 0 : maintenance_tasks.spawn(console::mgmt::task_main(mgmt_listener));
392 0 :
393 0 : if let Some(metrics_config) = &config.metric_collection {
394 0 : // TODO: Add gc regardles of the metric collection being enabled.
395 0 : maintenance_tasks.spawn(usage_metrics::task_main(metrics_config));
396 0 : client_tasks.spawn(usage_metrics::task_backup(
397 0 : &metrics_config.backup_metric_collection_config,
398 0 : cancellation_token,
399 0 : ));
400 0 : }
401 0 :
402 0 : if let auth::BackendType::Console(api, _) = &config.auth_backend {
403 0 : if let proxy::console::provider::ConsoleBackend::Console(api) = &**api {
404 0 : if let Some(redis_notifications_client) = redis_notifications_client {
405 0 : let cache = api.caches.project_info.clone();
406 0 : maintenance_tasks.spawn(notifications::task_main(
407 0 : redis_notifications_client.clone(),
408 0 : cache.clone(),
409 0 : cancel_map.clone(),
410 0 : args.region.clone(),
411 0 : ));
412 0 : maintenance_tasks.spawn(async move { cache.clone().gc_worker().await });
413 0 : }
414 0 : }
415 0 : }
416 0 :
417 0 : let maintenance = loop {
418 0 : // get one complete task
419 0 : match futures::future::select(
420 0 : pin!(maintenance_tasks.join_next()),
421 0 : pin!(client_tasks.join_next()),
422 0 : )
423 0 : .await
424 0 : {
425 0 : // exit immediately on maintenance task completion
426 0 : Either::Left((Some(res), _)) => break proxy::flatten_err(res)?,
427 0 : // exit with error immediately if all maintenance tasks have ceased (should be caught by branch above)
428 0 : Either::Left((None, _)) => bail!("no maintenance tasks running. invalid state"),
429 0 : // exit immediately on client task error
430 0 : Either::Right((Some(res), _)) => proxy::flatten_err(res)?,
431 0 : // exit if all our client tasks have shutdown gracefully
432 0 : Either::Right((None, _)) => return Ok(()),
433 0 : }
434 0 : };
435 0 :
436 0 : // maintenance tasks return Infallible success values, this is an impossible value
437 0 : // so this match statically ensures that there are no possibilities for that value
438 0 : match maintenance {}
439 0 : }
440 :
441 : /// ProxyConfig is created at proxy startup, and lives forever.
442 0 : fn build_config(args: &ProxyCliArgs) -> anyhow::Result<&'static ProxyConfig> {
443 0 : let tls_config = match (&args.tls_key, &args.tls_cert) {
444 0 : (Some(key_path), Some(cert_path)) => Some(config::configure_tls(
445 0 : key_path,
446 0 : cert_path,
447 0 : args.certs_dir.as_ref(),
448 0 : )?),
449 0 : (None, None) => None,
450 0 : _ => bail!("either both or neither tls-key and tls-cert must be specified"),
451 : };
452 :
453 0 : if args.allow_self_signed_compute {
454 0 : warn!("allowing self-signed compute certificates");
455 0 : }
456 0 : let backup_metric_collection_config = config::MetricBackupCollectionConfig {
457 0 : interval: args.metric_backup_collection_interval,
458 0 : remote_storage_config: remote_storage_from_toml(
459 0 : &args.metric_backup_collection_remote_storage,
460 0 : )?,
461 0 : chunk_size: args.metric_backup_collection_chunk_size,
462 : };
463 :
464 0 : let metric_collection = match (
465 0 : &args.metric_collection_endpoint,
466 0 : &args.metric_collection_interval,
467 : ) {
468 0 : (Some(endpoint), Some(interval)) => Some(config::MetricCollectionConfig {
469 0 : endpoint: endpoint.parse()?,
470 0 : interval: humantime::parse_duration(interval)?,
471 0 : backup_metric_collection_config,
472 : }),
473 0 : (None, None) => None,
474 0 : _ => bail!(
475 0 : "either both or neither metric-collection-endpoint \
476 0 : and metric-collection-interval must be specified"
477 0 : ),
478 : };
479 0 : let rate_limiter_config = RateLimiterConfig {
480 0 : disable: args.disable_dynamic_rate_limiter,
481 0 : algorithm: args.rate_limit_algorithm,
482 0 : timeout: args.rate_limiter_timeout,
483 0 : initial_limit: args.initial_limit,
484 0 : aimd_config: Some(args.aimd_config),
485 0 : };
486 :
487 0 : let auth_backend = match &args.auth_backend {
488 : AuthBackend::Console => {
489 0 : let wake_compute_cache_config: CacheOptions = args.wake_compute_cache.parse()?;
490 0 : let project_info_cache_config: ProjectInfoCacheOptions =
491 0 : args.project_info_cache.parse()?;
492 :
493 0 : info!("Using NodeInfoCache (wake_compute) with options={wake_compute_cache_config:?}");
494 0 : info!(
495 0 : "Using AllowedIpsCache (wake_compute) with options={project_info_cache_config:?}"
496 0 : );
497 0 : let caches = Box::leak(Box::new(console::caches::ApiCaches::new(
498 0 : wake_compute_cache_config,
499 0 : project_info_cache_config,
500 0 : )));
501 :
502 : let config::WakeComputeLockOptions {
503 0 : shards,
504 0 : permits,
505 0 : epoch,
506 0 : timeout,
507 0 : } = args.wake_compute_lock.parse()?;
508 0 : info!(permits, shards, ?epoch, "Using NodeLocks (wake_compute)");
509 0 : let locks = Box::leak(Box::new(
510 0 : console::locks::ApiLocks::new("wake_compute_lock", permits, shards, timeout)
511 0 : .unwrap(),
512 0 : ));
513 0 : tokio::spawn(locks.garbage_collect_worker(epoch));
514 :
515 0 : let url = args.auth_endpoint.parse()?;
516 0 : let endpoint = http::Endpoint::new(url, http::new_client(rate_limiter_config));
517 0 :
518 0 : let api = console::provider::neon::Api::new(endpoint, caches, locks);
519 0 : let api = console::provider::ConsoleBackend::Console(api);
520 0 : auth::BackendType::Console(MaybeOwned::Owned(api), ())
521 : }
522 : #[cfg(feature = "testing")]
523 : AuthBackend::Postgres => {
524 0 : let url = args.auth_endpoint.parse()?;
525 0 : let api = console::provider::mock::Api::new(url);
526 0 : let api = console::provider::ConsoleBackend::Postgres(api);
527 0 : auth::BackendType::Console(MaybeOwned::Owned(api), ())
528 : }
529 : AuthBackend::Link => {
530 0 : let url = args.uri.parse()?;
531 0 : auth::BackendType::Link(MaybeOwned::Owned(url), ())
532 : }
533 : };
534 0 : let http_config = HttpConfig {
535 0 : request_timeout: args.sql_over_http.sql_over_http_timeout,
536 0 : pool_options: GlobalConnPoolOptions {
537 0 : max_conns_per_endpoint: args.sql_over_http.sql_over_http_pool_max_conns_per_endpoint,
538 0 : gc_epoch: args.sql_over_http.sql_over_http_pool_gc_epoch,
539 0 : pool_shards: args.sql_over_http.sql_over_http_pool_shards,
540 0 : idle_timeout: args.sql_over_http.sql_over_http_idle_timeout,
541 0 : opt_in: args.sql_over_http.sql_over_http_pool_opt_in,
542 0 : max_total_conns: args.sql_over_http.sql_over_http_pool_max_total_conns,
543 0 : },
544 0 : };
545 0 : let authentication_config = AuthenticationConfig {
546 0 : scram_protocol_timeout: args.scram_protocol_timeout,
547 0 : rate_limiter_enabled: args.auth_rate_limit_enabled,
548 0 : rate_limiter: AuthRateLimiter::new(args.auth_rate_limit.clone()),
549 0 : };
550 0 :
551 0 : let mut endpoint_rps_limit = args.endpoint_rps_limit.clone();
552 0 : RateBucketInfo::validate(&mut endpoint_rps_limit)?;
553 0 : let mut redis_rps_limit = args.redis_rps_limit.clone();
554 0 : RateBucketInfo::validate(&mut redis_rps_limit)?;
555 :
556 0 : let config = Box::leak(Box::new(ProxyConfig {
557 0 : tls_config,
558 0 : auth_backend,
559 0 : metric_collection,
560 0 : allow_self_signed_compute: args.allow_self_signed_compute,
561 0 : http_config,
562 0 : authentication_config,
563 0 : require_client_ip: args.require_client_ip,
564 0 : disable_ip_check_for_http: args.disable_ip_check_for_http,
565 0 : endpoint_rps_limit,
566 0 : redis_rps_limit,
567 0 : handshake_timeout: args.handshake_timeout,
568 0 : region: args.region.clone(),
569 0 : aws_region: args.aws_region.clone(),
570 0 : }));
571 0 :
572 0 : Ok(config)
573 0 : }
574 :
575 : #[cfg(test)]
576 : mod tests {
577 : use std::time::Duration;
578 :
579 : use clap::Parser;
580 : use proxy::rate_limiter::RateBucketInfo;
581 :
582 : #[test]
583 2 : fn parse_endpoint_rps_limit() {
584 2 : let config = super::ProxyCliArgs::parse_from([
585 2 : "proxy",
586 2 : "--endpoint-rps-limit",
587 2 : "100@1s",
588 2 : "--endpoint-rps-limit",
589 2 : "20@30s",
590 2 : ]);
591 2 :
592 2 : assert_eq!(
593 2 : config.endpoint_rps_limit,
594 2 : vec![
595 2 : RateBucketInfo::new(100, Duration::from_secs(1)),
596 2 : RateBucketInfo::new(20, Duration::from_secs(30)),
597 2 : ]
598 2 : );
599 2 : }
600 : }
|
