rust/neon
1
0
Fork 0
mirror of https://github.com/neondatabase/neon.git synced 2026-01-15 17:32:56 +00:00
Code Issues Packages Projects Releases Wiki Activity

Implement a cache for get_auth_info method

Browse Source
This commit is contained in:
Dmitry Ivanov
2023-05-19 18:43:00 +03:00
committed by Vadim Kharitonov
parent f7e9ec49be
commit ce416c8160
5 changed files with 52 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files

21
proxy/src/bin/proxy.rs
View File

@@ -124,14 +124,25 @@ fn build_config(args: &clap::ArgMatches) -> anyhow::Result<&'static ProxyConfig>
let auth_backend = match args.get_one::<String>("auth-backend").unwrap().as_str() {
"console" => {
let config::CacheOptions { size, ttl } = args
.get_one::<String>("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::<String>("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::<String>("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")

5
proxy/src/config.rs
View File

@@ -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`].

2
proxy/src/console.rs
View File

@@ -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.

6
proxy/src/console/provider.rs
View File

@@ -166,7 +166,7 @@ pub struct NodeInfo {
pub type NodeInfoCache = TimedLru<Box<str>, NodeInfo>;
pub type CachedNodeInfo = Cached<NodeInfo>;
pub type AuthInfoCache = TimedLru<Box<str>, AuthInfo>;
pub type AuthInfoCache = TimedLru<(Box<str>, Box<str>), AuthInfo>;
pub type CachedAuthInfo = Cached<AuthInfo>;
/// 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,
}

33
proxy/src/console/provider/neon.rs
View File

@@ -113,9 +113,25 @@ impl super::Api for Api {
extra: &ConsoleReqExtra<'_>,
creds: &ClientCredentials<'_>,
) -> Result<Option<CachedAuthInfo>, 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<str>, Box<str>) = (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<CachedNodeInfo, WakeComputeError> {
let key = creds.project().expect("impossible");
let key: Box<str> = 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::<str>::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)
}