diff --git a/src/server/mod.rs b/src/server/mod.rs index b5146e6..08f1a53 100644 --- a/src/server/mod.rs +++ b/src/server/mod.rs @@ -1,6 +1,8 @@ mod host; +use std::collections::HashMap; use std::sync::Arc; +use std::time::{Duration, Instant}; use axum::async_trait; use axum::extract::FromRequestParts; @@ -11,7 +13,9 @@ use axum::http::StatusCode; use axum::response::{IntoResponse, Response}; use axum::routing::{get, post}; use axum::{RequestPartsExt, Router, TypedHeader}; +use tokio::sync::RwLock; use tracing::debug; +use uuid::Uuid; use crate::auth::{self, Claims}; use crate::config::Configuration; @@ -19,6 +23,7 @@ use crate::machine_id; struct Context { config: Arc, + cache: RwLock>, } type State = Arc; @@ -57,12 +62,10 @@ impl FromRequestParts> for Claims { AuthError })?; let machine_id = - machine_id::get_machine_id(&hostname, ctx.config.clone()) - .await - .ok_or_else(|| { - debug!("No machine ID found for host {}", hostname); - AuthError - })?; + get_machine_id(&hostname, ctx).await.ok_or_else(|| { + debug!("No machine ID found for host {}", hostname); + AuthError + })?; let claims = auth::validate_token( bearer.token(), &hostname, @@ -81,6 +84,7 @@ impl FromRequestParts> for Claims { pub fn make_app(config: Configuration) -> Router { let ctx = Arc::new(Context { config: config.into(), + cache: RwLock::new(Default::default()), }); Router::new() .route("/", get(|| async { "UP" })) @@ -88,6 +92,25 @@ pub fn make_app(config: Configuration) -> Router { .with_state(ctx) } +async fn get_machine_id(hostname: &str, ctx: &State) -> Option { + let cache = ctx.cache.read().await; + if let Some((ts, m)) = cache.get(hostname) { + if ts.elapsed() < Duration::from_secs(60) { + debug!("Found cached machine ID for {}", hostname); + return Some(*m); + } else { + debug!("Cached machine ID for {} has expired", hostname); + } + } + drop(cache); + let machine_id = + machine_id::get_machine_id(hostname, ctx.config.clone()).await?; + let mut cache = ctx.cache.write().await; + debug!("Caching machine ID for {}", hostname); + cache.insert(hostname.into(), (Instant::now(), machine_id)); + Some(machine_id) +} + #[cfg(test)] mod test { use axum::body::Body;