use serde for deser config

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",

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

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<Self> {
-        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<T>(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<K>(&mut self, seed: K) -> Result<Option<K::Value>, 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<V>(&mut self, seed: V) -> Result<V::Value, Self::Error>
+    where
+        V: serde::de::DeserializeSeed<'de>,
+    {
+        seed.deserialize(BorrowedStrDeserializer::new(self.val.take().unwrap()))
+    }
+
+    fn next_entry_seed<K, V>(
+        &mut self,
+        kseed: K,
+        vseed: V,
+    ) -> Result<Option<(K::Value, V::Value)>, 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<V>(self, visitor: V) -> Result<V::Value, Self::Error>
+    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(())

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<EndpointCacheKey, Arc<Semaphore>>,
-    permits: usize,
+    inner: Option<ApiLocksInner>,
     timeout: Duration,
     epoch: std::time::Duration,
     metrics: &'static ApiLockMetrics,
 }
 
+struct ApiLocksInner {
+    permits: NonZeroUsize,
+    node_locks: DashMap<EndpointCacheKey, Arc<Semaphore>>,
+}
+
 impl ApiLocks {
     pub fn new(
         name: &'static str,
@@ -465,10 +469,14 @@ impl ApiLocks {
         epoch: std::time::Duration,
         metrics: &'static ApiLockMetrics,
     ) -> prometheus::Result<Self> {
+        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<WakeComputePermit, errors::WakeComputeError> {
-        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