diff --git a/proxy/src/bin/proxy.rs b/proxy/src/bin/proxy.rs index db51688612..ea729df0c2 100644 --- a/proxy/src/bin/proxy.rs +++ b/proxy/src/bin/proxy.rs @@ -124,14 +124,25 @@ fn build_config(args: &clap::ArgMatches) -> anyhow::Result<&'static ProxyConfig> let auth_backend = match args.get_one::("auth-backend").unwrap().as_str() { "console" => { + let config::CacheOptions { size, ttl } = args + .get_one::("get-auth-info-cache") + .unwrap() + .parse()?; + + info!("Using AuthInfoCache (get_auth_info) with size={size} ttl={ttl:?}"); + let auth_info = console::caches::AuthInfoCache::new("auth_info_cache", size, ttl); + let config::CacheOptions { size, ttl } = args .get_one::("wake-compute-cache") .unwrap() .parse()?; info!("Using NodeInfoCache (wake_compute) with size={size} ttl={ttl:?}"); + let node_info = console::caches::NodeInfoCache::new("node_info_cache", size, ttl); + let caches = Box::leak(Box::new(console::caches::ApiCaches { - node_info: console::caches::NodeInfoCache::new("node_info_cache", size, ttl), + auth_info, + node_info, })); let url = args.get_one::("auth-endpoint").unwrap().parse()?; @@ -240,11 +251,17 @@ fn cli() -> clap::Command { .long("metric-collection-interval") .help("how often metrics should be sent to a collection endpoint"), ) + .arg( + Arg::new("get-auth-info-cache") + .long("get-auth-info-cache") + .help("cache for `get_auth_info` api method (use `size=0` to disable)") + .default_value(config::CacheOptions::DEFAULT_AUTH_INFO), + ) .arg( Arg::new("wake-compute-cache") .long("wake-compute-cache") .help("cache for `wake_compute` api method (use `size=0` to disable)") - .default_value(config::CacheOptions::DEFAULT_OPTIONS_NODE_INFO), + .default_value(config::CacheOptions::DEFAULT_NODE_INFO), ) .arg( Arg::new("allow-self-signed-compute") diff --git a/proxy/src/config.rs b/proxy/src/config.rs index 497a8916e5..f8cf19b6c2 100644 --- a/proxy/src/config.rs +++ b/proxy/src/config.rs @@ -210,8 +210,11 @@ pub struct CacheOptions { } impl CacheOptions { + /// Default options for [`crate::auth::caches::AuthInfoCache`]. + pub const DEFAULT_AUTH_INFO: &str = "size=4000,ttl=5s"; + /// Default options for [`crate::auth::caches::NodeInfoCache`]. - pub const DEFAULT_OPTIONS_NODE_INFO: &str = "size=4000,ttl=5m"; + pub const DEFAULT_NODE_INFO: &str = "size=4000,ttl=5m"; /// Parse cache options passed via cmdline. /// Example: [`Self::DEFAULT_OPTIONS_NODE_INFO`]. diff --git a/proxy/src/console.rs b/proxy/src/console.rs index c67ff8a18c..38d41bec72 100644 --- a/proxy/src/console.rs +++ b/proxy/src/console.rs @@ -12,7 +12,7 @@ pub use provider::{CachedAuthInfo, CachedNodeInfo}; /// Various cache-related types. pub mod caches { - pub use super::provider::{ApiCaches, NodeInfoCache}; + pub use super::provider::{ApiCaches, AuthInfoCache, NodeInfoCache}; } /// Console's management API. diff --git a/proxy/src/console/provider.rs b/proxy/src/console/provider.rs index ebe8b2b3e2..aa41f2e3d9 100644 --- a/proxy/src/console/provider.rs +++ b/proxy/src/console/provider.rs @@ -166,7 +166,7 @@ pub struct NodeInfo { pub type NodeInfoCache = TimedLru, NodeInfo>; pub type CachedNodeInfo = Cached; -pub type AuthInfoCache = TimedLru, AuthInfo>; +pub type AuthInfoCache = TimedLru<(Box, Box), AuthInfo>; pub type CachedAuthInfo = Cached; /// This will allocate per each call, but the http requests alone @@ -190,6 +190,8 @@ pub trait Api { /// Various caches for [`console`]. pub struct ApiCaches { - /// Cache for the `wake_compute` API method. + /// Cache for the `get_auth_info` method. + pub auth_info: AuthInfoCache, + /// Cache for the `wake_compute` method. pub node_info: NodeInfoCache, } diff --git a/proxy/src/console/provider/neon.rs b/proxy/src/console/provider/neon.rs index 27631f546a..bd7316422a 100644 --- a/proxy/src/console/provider/neon.rs +++ b/proxy/src/console/provider/neon.rs @@ -113,9 +113,25 @@ impl super::Api for Api { extra: &ConsoleReqExtra<'_>, creds: &ClientCredentials<'_>, ) -> Result, GetAuthInfoError> { - // FIXME: add cache! - let res = self.do_get_auth_info(extra, creds).await?; - Ok(res.map(CachedAuthInfo::new_uncached)) + let project = creds.project().expect("impossible"); + let key: (Box, Box) = (project.into(), creds.user.into()); + + // Check if we already have a cached auth info for this project + user combo. + // Beware! We shouldn't flush this for unsuccessful auth attempts, otherwise + // the cache makes no sense whatsoever in the presence of unfaithful clients. + // Instead, we snoop an invalidation queue to keep the cache up-to-date. + if let Some(cached) = self.caches.auth_info.get(&key) { + info!(key = ?key, "found cached auth info"); + return Ok(Some(cached)); + } + + let info = self.do_get_auth_info(extra, creds).await?; + + Ok(info.map(|info| { + info!(key = ?key, "creating a cache entry for auth info"); + let (_, cached) = self.caches.auth_info.insert(key.into(), info.into()); + cached + })) } #[tracing::instrument(skip_all)] @@ -124,22 +140,21 @@ impl super::Api for Api { extra: &ConsoleReqExtra<'_>, creds: &ClientCredentials<'_>, ) -> Result { - let key = creds.project().expect("impossible"); + let key: Box = creds.project().expect("impossible").into(); // Every time we do a wakeup http request, the compute node will stay up // for some time (highly depends on the console's scale-to-zero policy); // The connection info remains the same during that period of time, // which means that we might cache it to reduce the load and latency. - if let Some(cached) = self.caches.node_info.get(key) { - info!(key, "found cached compute node info"); + if let Some(cached) = self.caches.node_info.get(&key) { + info!(key = ?key, "found cached compute node info"); return Ok(cached); } let info = self.do_wake_compute(extra, creds).await?; - let owned_key = Box::::from(key); - let (_, cached) = self.caches.node_info.insert(owned_key.into(), info.into()); - info!(key, "created a cache entry for compute node info"); + info!(key = ?key, "creating a cache entry for compute node info"); + let (_, cached) = self.caches.node_info.insert(key.into(), info.into()); Ok(cached) }