Line data Source code
1 : use crate::{
2 : auth::{self, backend::AuthRateLimiter},
3 : console::locks::ApiLocks,
4 : rate_limiter::{RateBucketInfo, RateLimitAlgorithm, RateLimiterConfig},
5 : scram::threadpool::ThreadPool,
6 : serverless::{cancel_set::CancelSet, GlobalConnPoolOptions},
7 : Host,
8 : };
9 : use anyhow::{bail, ensure, Context, Ok};
10 : use itertools::Itertools;
11 : use remote_storage::RemoteStorageConfig;
12 : use rustls::{
13 : crypto::ring::sign,
14 : pki_types::{CertificateDer, PrivateKeyDer},
15 : };
16 : use sha2::{Digest, Sha256};
17 : use std::{
18 : collections::{HashMap, HashSet},
19 : str::FromStr,
20 : sync::Arc,
21 : time::Duration,
22 : };
23 : use tracing::{error, info};
24 : use x509_parser::oid_registry;
25 :
26 : pub struct ProxyConfig {
27 : pub tls_config: Option<TlsConfig>,
28 : pub auth_backend: auth::BackendType<'static, (), ()>,
29 : pub metric_collection: Option<MetricCollectionConfig>,
30 : pub allow_self_signed_compute: bool,
31 : pub http_config: HttpConfig,
32 : pub authentication_config: AuthenticationConfig,
33 : pub require_client_ip: bool,
34 : pub region: String,
35 : pub handshake_timeout: Duration,
36 : pub wake_compute_retry_config: RetryConfig,
37 : pub connect_compute_locks: ApiLocks<Host>,
38 : pub connect_to_compute_retry_config: RetryConfig,
39 : }
40 :
41 : #[derive(Debug)]
42 : pub struct MetricCollectionConfig {
43 : pub endpoint: reqwest::Url,
44 : pub interval: Duration,
45 : pub backup_metric_collection_config: MetricBackupCollectionConfig,
46 : }
47 :
48 : pub struct TlsConfig {
49 : pub config: Arc<rustls::ServerConfig>,
50 : pub common_names: HashSet<String>,
51 : pub cert_resolver: Arc<CertResolver>,
52 : }
53 :
54 : pub struct HttpConfig {
55 : pub accept_websockets: bool,
56 : pub pool_options: GlobalConnPoolOptions,
57 : pub cancel_set: CancelSet,
58 : pub client_conn_threshold: u64,
59 : }
60 :
61 : pub struct AuthenticationConfig {
62 : pub thread_pool: Arc<ThreadPool>,
63 : pub scram_protocol_timeout: tokio::time::Duration,
64 : pub rate_limiter_enabled: bool,
65 : pub rate_limiter: AuthRateLimiter,
66 : pub rate_limit_ip_subnet: u8,
67 : }
68 :
69 : impl TlsConfig {
70 40 : pub fn to_server_config(&self) -> Arc<rustls::ServerConfig> {
71 40 : self.config.clone()
72 40 : }
73 : }
74 :
75 : /// <https://github.com/postgres/postgres/blob/ca481d3c9ab7bf69ff0c8d71ad3951d407f6a33c/src/include/libpq/pqcomm.h#L159>
76 : pub const PG_ALPN_PROTOCOL: &[u8] = b"postgresql";
77 :
78 : /// Configure TLS for the main endpoint.
79 0 : pub fn configure_tls(
80 0 : key_path: &str,
81 0 : cert_path: &str,
82 0 : certs_dir: Option<&String>,
83 0 : ) -> anyhow::Result<TlsConfig> {
84 0 : let mut cert_resolver = CertResolver::new();
85 0 :
86 0 : // add default certificate
87 0 : cert_resolver.add_cert_path(key_path, cert_path, true)?;
88 :
89 : // add extra certificates
90 0 : if let Some(certs_dir) = certs_dir {
91 0 : for entry in std::fs::read_dir(certs_dir)? {
92 0 : let entry = entry?;
93 0 : let path = entry.path();
94 0 : if path.is_dir() {
95 : // file names aligned with default cert-manager names
96 0 : let key_path = path.join("tls.key");
97 0 : let cert_path = path.join("tls.crt");
98 0 : if key_path.exists() && cert_path.exists() {
99 0 : cert_resolver.add_cert_path(
100 0 : &key_path.to_string_lossy(),
101 0 : &cert_path.to_string_lossy(),
102 0 : false,
103 0 : )?;
104 0 : }
105 0 : }
106 : }
107 0 : }
108 :
109 0 : let common_names = cert_resolver.get_common_names();
110 0 :
111 0 : let cert_resolver = Arc::new(cert_resolver);
112 0 :
113 0 : // allow TLS 1.2 to be compatible with older client libraries
114 0 : let mut config = rustls::ServerConfig::builder_with_protocol_versions(&[
115 0 : &rustls::version::TLS13,
116 0 : &rustls::version::TLS12,
117 0 : ])
118 0 : .with_no_client_auth()
119 0 : .with_cert_resolver(cert_resolver.clone());
120 0 :
121 0 : config.alpn_protocols = vec![PG_ALPN_PROTOCOL.to_vec()];
122 0 :
123 0 : Ok(TlsConfig {
124 0 : config: Arc::new(config),
125 0 : common_names,
126 0 : cert_resolver,
127 0 : })
128 0 : }
129 :
130 : /// Channel binding parameter
131 : ///
132 : /// <https://www.rfc-editor.org/rfc/rfc5929#section-4>
133 : /// Description: The hash of the TLS server's certificate as it
134 : /// appears, octet for octet, in the server's Certificate message. Note
135 : /// that the Certificate message contains a certificate_list, in which
136 : /// the first element is the server's certificate.
137 : ///
138 : /// The hash function is to be selected as follows:
139 : ///
140 : /// * if the certificate's signatureAlgorithm uses a single hash
141 : /// function, and that hash function is either MD5 or SHA-1, then use SHA-256;
142 : ///
143 : /// * if the certificate's signatureAlgorithm uses a single hash
144 : /// function and that hash function neither MD5 nor SHA-1, then use
145 : /// the hash function associated with the certificate's
146 : /// signatureAlgorithm;
147 : ///
148 : /// * if the certificate's signatureAlgorithm uses no hash functions or
149 : /// uses multiple hash functions, then this channel binding type's
150 : /// channel bindings are undefined at this time (updates to is channel
151 : /// binding type may occur to address this issue if it ever arises).
152 : #[derive(Debug, Clone, Copy)]
153 : pub enum TlsServerEndPoint {
154 : Sha256([u8; 32]),
155 : Undefined,
156 : }
157 :
158 : impl TlsServerEndPoint {
159 42 : pub fn new(cert: &CertificateDer<'_>) -> anyhow::Result<Self> {
160 42 : let sha256_oids = [
161 42 : // I'm explicitly not adding MD5 or SHA1 here... They're bad.
162 42 : oid_registry::OID_SIG_ECDSA_WITH_SHA256,
163 42 : oid_registry::OID_PKCS1_SHA256WITHRSA,
164 42 : ];
165 :
166 42 : let pem = x509_parser::parse_x509_certificate(cert)
167 42 : .context("Failed to parse PEM object from cerficiate")?
168 : .1;
169 :
170 42 : info!(subject = %pem.subject, "parsing TLS certificate");
171 :
172 42 : let reg = oid_registry::OidRegistry::default().with_all_crypto();
173 42 : let oid = pem.signature_algorithm.oid();
174 42 : let alg = reg.get(oid);
175 42 : if sha256_oids.contains(oid) {
176 42 : let tls_server_end_point: [u8; 32] = Sha256::new().chain_update(cert).finalize().into();
177 42 : info!(subject = %pem.subject, signature_algorithm = alg.map(|a| a.description()), tls_server_end_point = %base64::encode(tls_server_end_point), "determined channel binding");
178 42 : Ok(Self::Sha256(tls_server_end_point))
179 : } else {
180 0 : error!(subject = %pem.subject, signature_algorithm = alg.map(|a| a.description()), "unknown channel binding");
181 0 : Ok(Self::Undefined)
182 : }
183 42 : }
184 :
185 32 : pub fn supported(&self) -> bool {
186 32 : !matches!(self, TlsServerEndPoint::Undefined)
187 32 : }
188 : }
189 :
190 : #[derive(Default, Debug)]
191 : pub struct CertResolver {
192 : certs: HashMap<String, (Arc<rustls::sign::CertifiedKey>, TlsServerEndPoint)>,
193 : default: Option<(Arc<rustls::sign::CertifiedKey>, TlsServerEndPoint)>,
194 : }
195 :
196 : impl CertResolver {
197 42 : pub fn new() -> Self {
198 42 : Self::default()
199 42 : }
200 :
201 0 : fn add_cert_path(
202 0 : &mut self,
203 0 : key_path: &str,
204 0 : cert_path: &str,
205 0 : is_default: bool,
206 0 : ) -> anyhow::Result<()> {
207 0 : let priv_key = {
208 0 : let key_bytes = std::fs::read(key_path)
209 0 : .context(format!("Failed to read TLS keys at '{key_path}'"))?;
210 0 : let mut keys = rustls_pemfile::pkcs8_private_keys(&mut &key_bytes[..]).collect_vec();
211 0 :
212 0 : ensure!(keys.len() == 1, "keys.len() = {} (should be 1)", keys.len());
213 : PrivateKeyDer::Pkcs8(
214 0 : keys.pop()
215 0 : .unwrap()
216 0 : .context(format!("Failed to parse TLS keys at '{key_path}'"))?,
217 : )
218 : };
219 :
220 0 : let cert_chain_bytes = std::fs::read(cert_path)
221 0 : .context(format!("Failed to read TLS cert file at '{cert_path}.'"))?;
222 :
223 0 : let cert_chain = {
224 0 : rustls_pemfile::certs(&mut &cert_chain_bytes[..])
225 0 : .try_collect()
226 0 : .with_context(|| {
227 0 : format!("Failed to read TLS certificate chain from bytes from file at '{cert_path}'.")
228 0 : })?
229 : };
230 :
231 0 : self.add_cert(priv_key, cert_chain, is_default)
232 0 : }
233 :
234 42 : pub fn add_cert(
235 42 : &mut self,
236 42 : priv_key: PrivateKeyDer<'static>,
237 42 : cert_chain: Vec<CertificateDer<'static>>,
238 42 : is_default: bool,
239 42 : ) -> anyhow::Result<()> {
240 42 : let key = sign::any_supported_type(&priv_key).context("invalid private key")?;
241 :
242 42 : let first_cert = &cert_chain[0];
243 42 : let tls_server_end_point = TlsServerEndPoint::new(first_cert)?;
244 42 : let pem = x509_parser::parse_x509_certificate(first_cert)
245 42 : .context("Failed to parse PEM object from cerficiate")?
246 : .1;
247 :
248 42 : let common_name = pem.subject().to_string();
249 :
250 : // We only use non-wildcard certificates in link proxy so it seems okay to treat them the same as
251 : // wildcard ones as we don't use SNI there. That treatment only affects certificate selection, so
252 : // verify-full will still check wildcard match. Old coding here just ignored non-wildcard common names
253 : // and passed None instead, which blows up number of cases downstream code should handle. Proper coding
254 : // here should better avoid Option for common_names, and do wildcard-based certificate selection instead
255 : // of cutting off '*.' parts.
256 42 : let common_name = if common_name.starts_with("CN=*.") {
257 0 : common_name.strip_prefix("CN=*.").map(|s| s.to_string())
258 : } else {
259 42 : common_name.strip_prefix("CN=").map(|s| s.to_string())
260 : }
261 42 : .context("Failed to parse common name from certificate")?;
262 :
263 42 : let cert = Arc::new(rustls::sign::CertifiedKey::new(cert_chain, key));
264 42 :
265 42 : if is_default {
266 42 : self.default = Some((cert.clone(), tls_server_end_point));
267 42 : }
268 :
269 42 : self.certs.insert(common_name, (cert, tls_server_end_point));
270 42 :
271 42 : Ok(())
272 42 : }
273 :
274 42 : pub fn get_common_names(&self) -> HashSet<String> {
275 42 : self.certs.keys().map(|s| s.to_string()).collect()
276 42 : }
277 : }
278 :
279 : impl rustls::server::ResolvesServerCert for CertResolver {
280 0 : fn resolve(
281 0 : &self,
282 0 : client_hello: rustls::server::ClientHello<'_>,
283 0 : ) -> Option<Arc<rustls::sign::CertifiedKey>> {
284 0 : self.resolve(client_hello.server_name()).map(|x| x.0)
285 0 : }
286 : }
287 :
288 : impl CertResolver {
289 40 : pub fn resolve(
290 40 : &self,
291 40 : server_name: Option<&str>,
292 40 : ) -> Option<(Arc<rustls::sign::CertifiedKey>, TlsServerEndPoint)> {
293 : // loop here and cut off more and more subdomains until we find
294 : // a match to get a proper wildcard support. OTOH, we now do not
295 : // use nested domains, so keep this simple for now.
296 : //
297 : // With the current coding foo.com will match *.foo.com and that
298 : // repeats behavior of the old code.
299 40 : if let Some(mut sni_name) = server_name {
300 : loop {
301 80 : if let Some(cert) = self.certs.get(sni_name) {
302 40 : return Some(cert.clone());
303 40 : }
304 40 : if let Some((_, rest)) = sni_name.split_once('.') {
305 40 : sni_name = rest;
306 40 : } else {
307 0 : return None;
308 : }
309 : }
310 : } else {
311 : // No SNI, use the default certificate, otherwise we can't get to
312 : // options parameter which can be used to set endpoint name too.
313 : // That means that non-SNI flow will not work for CNAME domains in
314 : // verify-full mode.
315 : //
316 : // If that will be a problem we can:
317 : //
318 : // a) Instead of multi-cert approach use single cert with extra
319 : // domains listed in Subject Alternative Name (SAN).
320 : // b) Deploy separate proxy instances for extra domains.
321 0 : self.default.as_ref().cloned()
322 : }
323 40 : }
324 : }
325 :
326 : #[derive(Debug)]
327 : pub struct EndpointCacheConfig {
328 : /// Batch size to receive all endpoints on the startup.
329 : pub initial_batch_size: usize,
330 : /// Batch size to receive endpoints.
331 : pub default_batch_size: usize,
332 : /// Timeouts for the stream read operation.
333 : pub xread_timeout: Duration,
334 : /// Stream name to read from.
335 : pub stream_name: String,
336 : /// Limiter info (to distinguish when to enable cache).
337 : pub limiter_info: Vec<RateBucketInfo>,
338 : /// Disable cache.
339 : /// If true, cache is ignored, but reports all statistics.
340 : pub disable_cache: bool,
341 : /// Retry interval for the stream read operation.
342 : pub retry_interval: Duration,
343 : }
344 :
345 : impl EndpointCacheConfig {
346 : /// Default options for [`crate::console::provider::NodeInfoCache`].
347 : /// Notice that by default the limiter is empty, which means that cache is disabled.
348 : pub const CACHE_DEFAULT_OPTIONS: &'static str =
349 : "initial_batch_size=1000,default_batch_size=10,xread_timeout=5m,stream_name=controlPlane,disable_cache=true,limiter_info=1000@1s,retry_interval=1s";
350 :
351 : /// Parse cache options passed via cmdline.
352 : /// Example: [`Self::CACHE_DEFAULT_OPTIONS`].
353 0 : fn parse(options: &str) -> anyhow::Result<Self> {
354 0 : let mut initial_batch_size = None;
355 0 : let mut default_batch_size = None;
356 0 : let mut xread_timeout = None;
357 0 : let mut stream_name = None;
358 0 : let mut limiter_info = vec![];
359 0 : let mut disable_cache = false;
360 0 : let mut retry_interval = None;
361 :
362 0 : for option in options.split(',') {
363 0 : let (key, value) = option
364 0 : .split_once('=')
365 0 : .with_context(|| format!("bad key-value pair: {option}"))?;
366 :
367 0 : match key {
368 0 : "initial_batch_size" => initial_batch_size = Some(value.parse()?),
369 0 : "default_batch_size" => default_batch_size = Some(value.parse()?),
370 0 : "xread_timeout" => xread_timeout = Some(humantime::parse_duration(value)?),
371 0 : "stream_name" => stream_name = Some(value.to_string()),
372 0 : "limiter_info" => limiter_info.push(RateBucketInfo::from_str(value)?),
373 0 : "disable_cache" => disable_cache = value.parse()?,
374 0 : "retry_interval" => retry_interval = Some(humantime::parse_duration(value)?),
375 0 : unknown => bail!("unknown key: {unknown}"),
376 : }
377 : }
378 0 : RateBucketInfo::validate(&mut limiter_info)?;
379 :
380 : Ok(Self {
381 0 : initial_batch_size: initial_batch_size.context("missing `initial_batch_size`")?,
382 0 : default_batch_size: default_batch_size.context("missing `default_batch_size`")?,
383 0 : xread_timeout: xread_timeout.context("missing `xread_timeout`")?,
384 0 : stream_name: stream_name.context("missing `stream_name`")?,
385 0 : disable_cache,
386 0 : limiter_info,
387 0 : retry_interval: retry_interval.context("missing `retry_interval`")?,
388 : })
389 0 : }
390 : }
391 :
392 : impl FromStr for EndpointCacheConfig {
393 : type Err = anyhow::Error;
394 :
395 0 : fn from_str(options: &str) -> Result<Self, Self::Err> {
396 0 : let error = || format!("failed to parse endpoint cache options '{options}'");
397 0 : Self::parse(options).with_context(error)
398 0 : }
399 : }
400 : #[derive(Debug)]
401 : pub struct MetricBackupCollectionConfig {
402 : pub interval: Duration,
403 : pub remote_storage_config: Option<RemoteStorageConfig>,
404 : pub chunk_size: usize,
405 : }
406 :
407 2 : pub fn remote_storage_from_toml(s: &str) -> anyhow::Result<RemoteStorageConfig> {
408 2 : RemoteStorageConfig::from_toml(&s.parse()?)
409 2 : }
410 :
411 : /// Helper for cmdline cache options parsing.
412 : #[derive(Debug)]
413 : pub struct CacheOptions {
414 : /// Max number of entries.
415 : pub size: usize,
416 : /// Entry's time-to-live.
417 : pub ttl: Duration,
418 : }
419 :
420 : impl CacheOptions {
421 : /// Default options for [`crate::console::provider::NodeInfoCache`].
422 : pub const CACHE_DEFAULT_OPTIONS: &'static str = "size=4000,ttl=4m";
423 :
424 : /// Parse cache options passed via cmdline.
425 : /// Example: [`Self::CACHE_DEFAULT_OPTIONS`].
426 8 : fn parse(options: &str) -> anyhow::Result<Self> {
427 8 : let mut size = None;
428 8 : let mut ttl = None;
429 :
430 14 : for option in options.split(',') {
431 14 : let (key, value) = option
432 14 : .split_once('=')
433 14 : .with_context(|| format!("bad key-value pair: {option}"))?;
434 :
435 14 : match key {
436 14 : "size" => size = Some(value.parse()?),
437 6 : "ttl" => ttl = Some(humantime::parse_duration(value)?),
438 0 : unknown => bail!("unknown key: {unknown}"),
439 : }
440 : }
441 :
442 : // TTL doesn't matter if cache is always empty.
443 8 : if let Some(0) = size {
444 4 : ttl.get_or_insert(Duration::default());
445 4 : }
446 :
447 : Ok(Self {
448 8 : size: size.context("missing `size`")?,
449 8 : ttl: ttl.context("missing `ttl`")?,
450 : })
451 8 : }
452 : }
453 :
454 : impl FromStr for CacheOptions {
455 : type Err = anyhow::Error;
456 :
457 8 : fn from_str(options: &str) -> Result<Self, Self::Err> {
458 8 : let error = || format!("failed to parse cache options '{options}'");
459 8 : Self::parse(options).with_context(error)
460 8 : }
461 : }
462 :
463 : /// Helper for cmdline cache options parsing.
464 : #[derive(Debug)]
465 : pub struct ProjectInfoCacheOptions {
466 : /// Max number of entries.
467 : pub size: usize,
468 : /// Entry's time-to-live.
469 : pub ttl: Duration,
470 : /// Max number of roles per endpoint.
471 : pub max_roles: usize,
472 : /// Gc interval.
473 : pub gc_interval: Duration,
474 : }
475 :
476 : impl ProjectInfoCacheOptions {
477 : /// Default options for [`crate::console::provider::NodeInfoCache`].
478 : pub const CACHE_DEFAULT_OPTIONS: &'static str =
479 : "size=10000,ttl=4m,max_roles=10,gc_interval=60m";
480 :
481 : /// Parse cache options passed via cmdline.
482 : /// Example: [`Self::CACHE_DEFAULT_OPTIONS`].
483 0 : fn parse(options: &str) -> anyhow::Result<Self> {
484 0 : let mut size = None;
485 0 : let mut ttl = None;
486 0 : let mut max_roles = None;
487 0 : let mut gc_interval = None;
488 :
489 0 : for option in options.split(',') {
490 0 : let (key, value) = option
491 0 : .split_once('=')
492 0 : .with_context(|| format!("bad key-value pair: {option}"))?;
493 :
494 0 : match key {
495 0 : "size" => size = Some(value.parse()?),
496 0 : "ttl" => ttl = Some(humantime::parse_duration(value)?),
497 0 : "max_roles" => max_roles = Some(value.parse()?),
498 0 : "gc_interval" => gc_interval = Some(humantime::parse_duration(value)?),
499 0 : unknown => bail!("unknown key: {unknown}"),
500 : }
501 : }
502 :
503 : // TTL doesn't matter if cache is always empty.
504 0 : if let Some(0) = size {
505 0 : ttl.get_or_insert(Duration::default());
506 0 : }
507 :
508 : Ok(Self {
509 0 : size: size.context("missing `size`")?,
510 0 : ttl: ttl.context("missing `ttl`")?,
511 0 : max_roles: max_roles.context("missing `max_roles`")?,
512 0 : gc_interval: gc_interval.context("missing `gc_interval`")?,
513 : })
514 0 : }
515 : }
516 :
517 : impl FromStr for ProjectInfoCacheOptions {
518 : type Err = anyhow::Error;
519 :
520 0 : fn from_str(options: &str) -> Result<Self, Self::Err> {
521 0 : let error = || format!("failed to parse cache options '{options}'");
522 0 : Self::parse(options).with_context(error)
523 0 : }
524 : }
525 :
526 : /// This is a config for connect to compute and wake compute.
527 : #[derive(Clone, Copy, Debug)]
528 : pub struct RetryConfig {
529 : /// Number of times we should retry.
530 : pub max_retries: u32,
531 : /// Retry duration is base_delay * backoff_factor ^ n, where n starts at 0
532 : pub base_delay: tokio::time::Duration,
533 : /// Exponential base for retry wait duration
534 : pub backoff_factor: f64,
535 : }
536 :
537 : impl RetryConfig {
538 : /// Default options for RetryConfig.
539 :
540 : /// Total delay for 5 retries with 200ms base delay and 2 backoff factor is about 6s.
541 : pub const CONNECT_TO_COMPUTE_DEFAULT_VALUES: &'static str =
542 : "num_retries=5,base_retry_wait_duration=200ms,retry_wait_exponent_base=2";
543 : /// Total delay for 8 retries with 100ms base delay and 1.6 backoff factor is about 7s.
544 : /// Cplane has timeout of 60s on each request. 8m7s in total.
545 : pub const WAKE_COMPUTE_DEFAULT_VALUES: &'static str =
546 : "num_retries=8,base_retry_wait_duration=100ms,retry_wait_exponent_base=1.6";
547 :
548 : /// Parse retry options passed via cmdline.
549 : /// Example: [`Self::CONNECT_TO_COMPUTE_DEFAULT_VALUES`].
550 0 : pub fn parse(options: &str) -> anyhow::Result<Self> {
551 0 : let mut num_retries = None;
552 0 : let mut base_retry_wait_duration = None;
553 0 : let mut retry_wait_exponent_base = None;
554 :
555 0 : for option in options.split(',') {
556 0 : let (key, value) = option
557 0 : .split_once('=')
558 0 : .with_context(|| format!("bad key-value pair: {option}"))?;
559 :
560 0 : match key {
561 0 : "num_retries" => num_retries = Some(value.parse()?),
562 0 : "base_retry_wait_duration" => {
563 0 : base_retry_wait_duration = Some(humantime::parse_duration(value)?);
564 : }
565 0 : "retry_wait_exponent_base" => retry_wait_exponent_base = Some(value.parse()?),
566 0 : unknown => bail!("unknown key: {unknown}"),
567 : }
568 : }
569 :
570 : Ok(Self {
571 0 : max_retries: num_retries.context("missing `num_retries`")?,
572 0 : base_delay: base_retry_wait_duration.context("missing `base_retry_wait_duration`")?,
573 0 : backoff_factor: retry_wait_exponent_base
574 0 : .context("missing `retry_wait_exponent_base`")?,
575 : })
576 0 : }
577 : }
578 :
579 : /// Helper for cmdline cache options parsing.
580 16 : #[derive(serde::Deserialize)]
581 : pub struct ConcurrencyLockOptions {
582 : /// The number of shards the lock map should have
583 : pub shards: usize,
584 : /// The number of allowed concurrent requests for each endpoitn
585 : #[serde(flatten)]
586 : pub limiter: RateLimiterConfig,
587 : /// Garbage collection epoch
588 : #[serde(deserialize_with = "humantime_serde::deserialize")]
589 : pub epoch: Duration,
590 : /// Lock timeout
591 : #[serde(deserialize_with = "humantime_serde::deserialize")]
592 : pub timeout: Duration,
593 : }
594 :
595 : impl ConcurrencyLockOptions {
596 : /// Default options for [`crate::console::provider::ApiLocks`].
597 : pub const DEFAULT_OPTIONS_WAKE_COMPUTE_LOCK: &'static str = "permits=0";
598 : /// Default options for [`crate::console::provider::ApiLocks`].
599 : pub const DEFAULT_OPTIONS_CONNECT_COMPUTE_LOCK: &'static str =
600 : "shards=64,permits=100,epoch=10m,timeout=10ms";
601 :
602 : // pub const DEFAULT_OPTIONS_WAKE_COMPUTE_LOCK: &'static str = "shards=32,permits=4,epoch=10m,timeout=1s";
603 :
604 : /// Parse lock options passed via cmdline.
605 : /// Example: [`Self::DEFAULT_OPTIONS_WAKE_COMPUTE_LOCK`].
606 8 : fn parse(options: &str) -> anyhow::Result<Self> {
607 8 : let options = options.trim();
608 8 : if options.starts_with('{') && options.ends_with('}') {
609 2 : return Ok(serde_json::from_str(options)?);
610 6 : }
611 6 :
612 6 : let mut shards = None;
613 6 : let mut permits = None;
614 6 : let mut epoch = None;
615 6 : let mut timeout = None;
616 :
617 18 : for option in options.split(',') {
618 18 : let (key, value) = option
619 18 : .split_once('=')
620 18 : .with_context(|| format!("bad key-value pair: {option}"))?;
621 :
622 18 : match key {
623 18 : "shards" => shards = Some(value.parse()?),
624 14 : "permits" => permits = Some(value.parse()?),
625 8 : "epoch" => epoch = Some(humantime::parse_duration(value)?),
626 4 : "timeout" => timeout = Some(humantime::parse_duration(value)?),
627 0 : unknown => bail!("unknown key: {unknown}"),
628 : }
629 : }
630 :
631 : // these dont matter if lock is disabled
632 6 : if let Some(0) = permits {
633 2 : timeout = Some(Duration::default());
634 2 : epoch = Some(Duration::default());
635 2 : shards = Some(2);
636 4 : }
637 :
638 6 : let permits = permits.context("missing `permits`")?;
639 6 : let out = Self {
640 6 : shards: shards.context("missing `shards`")?,
641 6 : limiter: RateLimiterConfig {
642 6 : algorithm: RateLimitAlgorithm::Fixed,
643 6 : initial_limit: permits,
644 6 : },
645 6 : epoch: epoch.context("missing `epoch`")?,
646 6 : timeout: timeout.context("missing `timeout`")?,
647 : };
648 :
649 6 : ensure!(out.shards > 1, "shard count must be > 1");
650 6 : ensure!(
651 6 : out.shards.is_power_of_two(),
652 0 : "shard count must be a power of two"
653 : );
654 :
655 6 : Ok(out)
656 8 : }
657 : }
658 :
659 : impl FromStr for ConcurrencyLockOptions {
660 : type Err = anyhow::Error;
661 :
662 8 : fn from_str(options: &str) -> Result<Self, Self::Err> {
663 8 : let error = || format!("failed to parse cache lock options '{options}'");
664 8 : Self::parse(options).with_context(error)
665 8 : }
666 : }
667 :
668 : #[cfg(test)]
669 : mod tests {
670 : use crate::rate_limiter::Aimd;
671 :
672 : use super::*;
673 :
674 : #[test]
675 2 : fn test_parse_cache_options() -> anyhow::Result<()> {
676 2 : let CacheOptions { size, ttl } = "size=4096,ttl=5min".parse()?;
677 2 : assert_eq!(size, 4096);
678 2 : assert_eq!(ttl, Duration::from_secs(5 * 60));
679 :
680 2 : let CacheOptions { size, ttl } = "ttl=4m,size=2".parse()?;
681 2 : assert_eq!(size, 2);
682 2 : assert_eq!(ttl, Duration::from_secs(4 * 60));
683 :
684 2 : let CacheOptions { size, ttl } = "size=0,ttl=1s".parse()?;
685 2 : assert_eq!(size, 0);
686 2 : assert_eq!(ttl, Duration::from_secs(1));
687 :
688 2 : let CacheOptions { size, ttl } = "size=0".parse()?;
689 2 : assert_eq!(size, 0);
690 2 : assert_eq!(ttl, Duration::default());
691 :
692 2 : Ok(())
693 2 : }
694 :
695 : #[test]
696 2 : fn test_parse_lock_options() -> anyhow::Result<()> {
697 : let ConcurrencyLockOptions {
698 2 : epoch,
699 2 : limiter,
700 2 : shards,
701 2 : timeout,
702 2 : } = "shards=32,permits=4,epoch=10m,timeout=1s".parse()?;
703 2 : assert_eq!(epoch, Duration::from_secs(10 * 60));
704 2 : assert_eq!(timeout, Duration::from_secs(1));
705 2 : assert_eq!(shards, 32);
706 2 : assert_eq!(limiter.initial_limit, 4);
707 2 : assert_eq!(limiter.algorithm, RateLimitAlgorithm::Fixed);
708 :
709 : let ConcurrencyLockOptions {
710 2 : epoch,
711 2 : limiter,
712 2 : shards,
713 2 : timeout,
714 2 : } = "epoch=60s,shards=16,timeout=100ms,permits=8".parse()?;
715 2 : assert_eq!(epoch, Duration::from_secs(60));
716 2 : assert_eq!(timeout, Duration::from_millis(100));
717 2 : assert_eq!(shards, 16);
718 2 : assert_eq!(limiter.initial_limit, 8);
719 2 : assert_eq!(limiter.algorithm, RateLimitAlgorithm::Fixed);
720 :
721 : let ConcurrencyLockOptions {
722 2 : epoch,
723 2 : limiter,
724 2 : shards,
725 2 : timeout,
726 2 : } = "permits=0".parse()?;
727 2 : assert_eq!(epoch, Duration::ZERO);
728 2 : assert_eq!(timeout, Duration::ZERO);
729 2 : assert_eq!(shards, 2);
730 2 : assert_eq!(limiter.initial_limit, 0);
731 2 : assert_eq!(limiter.algorithm, RateLimitAlgorithm::Fixed);
732 :
733 2 : Ok(())
734 2 : }
735 :
736 : #[test]
737 2 : fn test_parse_json_lock_options() -> anyhow::Result<()> {
738 : let ConcurrencyLockOptions {
739 2 : epoch,
740 2 : limiter,
741 2 : shards,
742 2 : timeout,
743 2 : } = r#"{"shards":32,"initial_limit":44,"aimd":{"min":5,"max":500,"inc":10,"dec":0.9,"utilisation":0.8},"epoch":"10m","timeout":"1s"}"#
744 2 : .parse()?;
745 2 : assert_eq!(epoch, Duration::from_secs(10 * 60));
746 2 : assert_eq!(timeout, Duration::from_secs(1));
747 2 : assert_eq!(shards, 32);
748 2 : assert_eq!(limiter.initial_limit, 44);
749 2 : assert_eq!(
750 2 : limiter.algorithm,
751 2 : RateLimitAlgorithm::Aimd {
752 2 : conf: Aimd {
753 2 : min: 5,
754 2 : max: 500,
755 2 : dec: 0.9,
756 2 : inc: 10,
757 2 : utilisation: 0.8
758 2 : }
759 2 : },
760 2 : );
761 :
762 2 : Ok(())
763 2 : }
764 : }
|
