diff --git a/Cargo.lock b/Cargo.lock index 6faf4b72f0..810945729c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4314,6 +4314,7 @@ dependencies = [ "http 1.1.0", "http-body-util", "humantime", + "humantime-serde", "hyper 0.14.26", "hyper 1.2.0", "hyper-tungstenite", @@ -4355,6 +4356,7 @@ dependencies = [ "scopeguard", "serde", "serde_json", + "serde_with", "sha2", "smallvec", "smol_str", diff --git a/proxy/Cargo.toml b/proxy/Cargo.toml index 6b8f2ecbf4..c774786edc 100644 --- a/proxy/Cargo.toml +++ b/proxy/Cargo.toml @@ -35,6 +35,7 @@ hmac.workspace = true hostname.workspace = true http.workspace = true humantime.workspace = true +humantime-serde.workspace = true hyper-tungstenite.workspace = true hyper.workspace = true hyper1 = { package = "hyper", version = "1.2", features = ["server"] } @@ -70,6 +71,7 @@ rustls.workspace = true scopeguard.workspace = true serde.workspace = true serde_json.workspace = true +serde_with.workspace = true sha2 = { workspace = true, features = ["asm"] } smol_str.workspace = true smallvec.workspace = true diff --git a/proxy/src/config.rs b/proxy/src/config.rs index f9519c7645..8b67e5a7ec 100644 --- a/proxy/src/config.rs +++ b/proxy/src/config.rs @@ -3,13 +3,18 @@ use crate::{ rate_limiter::RateBucketInfo, serverless::GlobalConnPoolOptions, }; -use anyhow::{bail, ensure, Context, Ok}; +use anyhow::{bail, ensure, Context}; use itertools::Itertools; use remote_storage::RemoteStorageConfig; use rustls::{ crypto::ring::sign, pki_types::{CertificateDer, PrivateKeyDer}, }; +use serde::{ + de::{value::BorrowedStrDeserializer, MapAccess}, + forward_to_deserialize_any, Deserialize, Deserializer, +}; +use serde_with::serde_as; use sha2::{Digest, Sha256}; use std::{ collections::{HashMap, HashSet}, @@ -518,14 +523,23 @@ impl FromStr for ProjectInfoCacheOptions { } /// Helper for cmdline cache options parsing. +#[serde_as] +#[derive(Deserialize)] pub struct WakeComputeLockOptions { /// The number of shards the lock map should have + #[serde_as(as = "serde_with::DisplayFromStr")] + #[serde(default)] pub shards: usize, /// The number of allowed concurrent requests for each endpoitn + #[serde_as(as = "serde_with::DisplayFromStr")] pub permits: usize, /// Garbage collection epoch + #[serde(with = "humantime_serde")] + #[serde(default)] pub epoch: Duration, /// Lock timeout + #[serde(with = "humantime_serde")] + #[serde(default)] pub timeout: Duration, } @@ -538,45 +552,23 @@ impl WakeComputeLockOptions { /// Parse lock options passed via cmdline. /// Example: [`Self::DEFAULT_OPTIONS_WAKE_COMPUTE_LOCK`]. fn parse(options: &str) -> anyhow::Result { - let mut shards = None; - let mut permits = None; - let mut epoch = None; - let mut timeout = None; - - for option in options.split(',') { - let (key, value) = option - .split_once('=') - .with_context(|| format!("bad key-value pair: {option}"))?; - - match key { - "shards" => shards = Some(value.parse()?), - "permits" => permits = Some(value.parse()?), - "epoch" => epoch = Some(humantime::parse_duration(value)?), - "timeout" => timeout = Some(humantime::parse_duration(value)?), - unknown => bail!("unknown key: {unknown}"), - } + let out = Self::deserialize(SimpleKVConfig(options))?; + if out.permits != 0 { + ensure!( + out.timeout > Duration::ZERO, + "wake compute lock timeout should be non-zero" + ); + ensure!( + out.epoch > Duration::ZERO, + "wake compute lock gc epoch should be non-zero" + ); + ensure!(out.shards > 1, "shard count must be > 1"); + ensure!( + out.shards.is_power_of_two(), + "shard count must be a power of two" + ); } - // these dont matter if lock is disabled - if let Some(0) = permits { - timeout = Some(Duration::default()); - epoch = Some(Duration::default()); - shards = Some(2); - } - - let out = Self { - shards: shards.context("missing `shards`")?, - permits: permits.context("missing `permits`")?, - epoch: epoch.context("missing `epoch`")?, - timeout: timeout.context("missing `timeout`")?, - }; - - ensure!(out.shards > 1, "shard count must be > 1"); - ensure!( - out.shards.is_power_of_two(), - "shard count must be a power of two" - ); - Ok(out) } } @@ -590,6 +582,100 @@ impl FromStr for WakeComputeLockOptions { } } +struct SimpleKVConfig<'a>(&'a str); +struct SimpleKVConfigMapAccess<'a> { + kv: std::str::Split<'a, char>, + val: Option<&'a str>, +} + +#[derive(Debug)] +struct SimpleKVConfigErr(String); + +impl std::fmt::Display for SimpleKVConfigErr { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str(&self.0) + } +} + +impl std::error::Error for SimpleKVConfigErr {} + +impl serde::de::Error for SimpleKVConfigErr { + fn custom(msg: T) -> Self + where + T: std::fmt::Display, + { + Self(msg.to_string()) + } +} + +impl<'de> MapAccess<'de> for SimpleKVConfigMapAccess<'de> { + type Error = SimpleKVConfigErr; + + fn next_key_seed(&mut self, seed: K) -> Result, Self::Error> + where + K: serde::de::DeserializeSeed<'de>, + { + let Some(kv) = self.kv.next() else { + return Ok(None); + }; + let (key, value) = kv + .split_once('=') + .ok_or_else(|| SimpleKVConfigErr("invalid kv pair".to_string()))?; + self.val = Some(value); + + seed.deserialize(BorrowedStrDeserializer::new(key)) + .map(Some) + } + + fn next_value_seed(&mut self, seed: V) -> Result + where + V: serde::de::DeserializeSeed<'de>, + { + seed.deserialize(BorrowedStrDeserializer::new(self.val.take().unwrap())) + } + + fn next_entry_seed( + &mut self, + kseed: K, + vseed: V, + ) -> Result, Self::Error> + where + K: serde::de::DeserializeSeed<'de>, + V: serde::de::DeserializeSeed<'de>, + { + let Some(kv) = self.kv.next() else { + return Ok(None); + }; + let (key, value) = kv + .split_once('=') + .ok_or_else(|| SimpleKVConfigErr("invalid kv pair".to_string()))?; + + let key = kseed.deserialize(BorrowedStrDeserializer::new(key))?; + let value = vseed.deserialize(BorrowedStrDeserializer::new(value))?; + Ok(Some((key, value))) + } +} + +impl<'de> Deserializer<'de> for SimpleKVConfig<'de> { + type Error = SimpleKVConfigErr; + + fn deserialize_any(self, visitor: V) -> Result + where + V: serde::de::Visitor<'de>, + { + visitor.visit_map(SimpleKVConfigMapAccess { + kv: self.0.split(','), + val: None, + }) + } + + forward_to_deserialize_any! { + bool i8 i16 i32 i64 i128 u8 u16 u32 u64 u128 f32 f64 char str string + bytes byte_buf option unit struct unit_struct newtype_struct seq tuple + tuple_struct map enum identifier ignored_any + } +} + #[cfg(test)] mod tests { use super::*; @@ -647,7 +733,7 @@ mod tests { } = "permits=0".parse()?; assert_eq!(epoch, Duration::ZERO); assert_eq!(timeout, Duration::ZERO); - assert_eq!(shards, 2); + assert_eq!(shards, 0); assert_eq!(permits, 0); Ok(()) diff --git a/proxy/src/console/provider.rs b/proxy/src/console/provider.rs index aa1800a9da..278a9c1979 100644 --- a/proxy/src/console/provider.rs +++ b/proxy/src/console/provider.rs @@ -17,7 +17,7 @@ use crate::{ scram, EndpointCacheKey, }; use dashmap::DashMap; -use std::{sync::Arc, time::Duration}; +use std::{num::NonZeroUsize, sync::Arc, time::Duration}; use tokio::sync::{OwnedSemaphorePermit, Semaphore}; use tokio::time::Instant; use tracing::info; @@ -449,13 +449,17 @@ impl ApiCaches { /// Various caches for [`console`](super). pub struct ApiLocks { name: &'static str, - node_locks: DashMap>, - permits: usize, + inner: Option, timeout: Duration, epoch: std::time::Duration, metrics: &'static ApiLockMetrics, } +struct ApiLocksInner { + permits: NonZeroUsize, + node_locks: DashMap>, +} + impl ApiLocks { pub fn new( name: &'static str, @@ -465,10 +469,14 @@ impl ApiLocks { epoch: std::time::Duration, metrics: &'static ApiLockMetrics, ) -> prometheus::Result { + let inner = NonZeroUsize::new(permits).map(|permits| ApiLocksInner { + permits, + node_locks: DashMap::with_shard_amount(shards), + }); + Ok(Self { name, - node_locks: DashMap::with_shard_amount(shards), - permits, + inner, timeout, epoch, metrics, @@ -479,20 +487,21 @@ impl ApiLocks { &self, key: &EndpointCacheKey, ) -> Result { - if self.permits == 0 { + let Some(inner) = &self.inner else { return Ok(WakeComputePermit { permit: None }); - } + }; let now = Instant::now(); let semaphore = { // get fast path - if let Some(semaphore) = self.node_locks.get(key) { + if let Some(semaphore) = inner.node_locks.get(key) { semaphore.clone() } else { - self.node_locks + inner + .node_locks .entry(key.clone()) .or_insert_with(|| { self.metrics.semaphores_registered.inc(); - Arc::new(Semaphore::new(self.permits)) + Arc::new(Semaphore::new(inner.permits.get())) }) .clone() } @@ -509,13 +518,13 @@ impl ApiLocks { } pub async fn garbage_collect_worker(&self) { - if self.permits == 0 { + let Some(inner) = &self.inner else { return; - } + }; let mut interval = - tokio::time::interval(self.epoch / (self.node_locks.shards().len()) as u32); + tokio::time::interval(self.epoch / (inner.node_locks.shards().len()) as u32); loop { - for (i, shard) in self.node_locks.shards().iter().enumerate() { + for (i, shard) in inner.node_locks.shards().iter().enumerate() { interval.tick().await; // temporary lock a single shard and then clear any semaphores that aren't currently checked out // race conditions: if strong_count == 1, there's no way that it can increase while the shard is locked