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