Merge pull request #16 from QuentinMoss/qm-implement-origin-caching

feat: Cache origins

Cargo.toml

@@ -17,3 +17,4 @@ tower = { version = "0.4", features = ["util"] }
 tower-http = { version = "0.4.0", features = ["trace"] }
 clap = { version = "4.3.8", features = ["derive"] }
 toml = "0.7.5"
+parking_lot = "0.12.1"

src/cache.rs

@@ -0,0 +1,75 @@
+use crate::db::Origin;
+use parking_lot::RwLock;
+use std::collections::HashMap;
+use std::sync::Arc;
+
+use crate::error::AppError;
+
+#[derive(Debug)]
+pub struct OriginCache(pub(crate) Arc<OriginCacheInner>);
+
+impl OriginCache {
+    pub fn new() -> Self {
+        let inner = OriginCacheInner::new();
+        Self(Arc::new(inner))
+    }
+
+    pub fn refresh(&self, new_origins: Vec<Origin>) -> Result<(), AppError> {
+        self.0.refresh(new_origins)
+    }
+
+    pub fn get(&self, domain: &str) -> Option<Origin> {
+        self.0.get(domain)
+    }
+}
+
+impl Clone for OriginCache {
+    fn clone(&self) -> Self {
+        Self(Arc::clone(&self.0))
+    }
+}
+
+#[derive(Debug)]
+pub struct OriginCacheInner {
+    origins: Arc<RwLock<HashMap<String, Origin>>>,
+}
+
+impl OriginCacheInner {
+    pub fn new() -> Self {
+        Self {
+            origins: Arc::new(RwLock::new(HashMap::new())),
+        }
+    }
+
+    pub fn refresh(&self, new_origins: Vec<Origin>) -> Result<(), AppError> {
+        // Iterate over the fetched origins and insert them into the map
+        let map = new_origins
+            .into_iter()
+            .map(|origin| (origin.domain.clone(), origin))
+            .collect();
+
+        // Update the cache by acquiring a write lock and replacing the HashMap
+        *self.origins.write() = map;
+        Ok(())
+    }
+
+    pub fn get(&self, domain: &str) -> Option<Origin> {
+        tracing::info!("Got called on cache for domain: {}", domain);
+        // Look up domain in the cache and clone if found
+        let result = {
+            let origins = self.origins.read();
+
+            origins.get(domain).cloned()
+        };
+
+        // Mostly for development, but also useful if you want to see how often the cache is hit
+        if result.is_some() {
+            tracing::info!("Found origin in cache");
+        } else {
+            tracing::info!("Origin not found in cache");
+        }
+
+        // Return the result if found, otherwise None
+        result
+    }
+}

src/db.rs

@@ -5,7 +5,7 @@ use sqlx::Executor;
 
 use crate::ingest::HttpRequest;
 
-#[derive(Debug, Deserialize, Serialize, sqlx::FromRow)]
+#[derive(Debug, Deserialize, Serialize, sqlx::FromRow, Clone)]
 pub struct Origin {
     pub id: i64,
     pub domain: String,

src/error.rs

@@ -1,6 +1,8 @@
 use axum::http::StatusCode;
 use axum::response::{IntoResponse, Response};
+use std::fmt;
 
+#[derive(Debug)]
 pub struct AppError(anyhow::Error);
 
 impl IntoResponse for AppError {
@@ -14,6 +16,12 @@ impl IntoResponse for AppError {
     }
 }
 
+impl fmt::Display for AppError {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        fmt::Display::fmt(&self.0, f)
+    }
+}
+
 impl<E> From<E> for AppError
 where
     E: Into<anyhow::Error>,

src/lib.rs

@@ -1,3 +1,4 @@
+pub mod cache;
 pub mod db;
 pub mod error;
 pub mod ingest;
@@ -15,9 +16,11 @@ use axum::http::StatusCode;
 use axum::response::IntoResponse;
 use axum::{routing::post, Router};
 use hyper::HeaderMap;
+use queue::RetryQueue;
 use serde::Deserialize;
 use sqlx::sqlite::SqlitePool;
 
+use crate::cache::OriginCache;
 use crate::db::{ensure_schema, insert_request, mark_complete, mark_error};
 use crate::error::AppError;
 use crate::ingest::HttpRequest;
@@ -43,28 +46,32 @@ impl Default for Config {
     }
 }
 
-pub async fn app(config: &Config) -> Result<(Router, Router, SqlitePool)> {
+pub async fn app(config: &Config) -> Result<(Router, Router, RetryQueue)> {
     let pool = SqlitePool::connect(&config.database_url).await?;
-    let pool2 = pool.clone();
-
     ensure_schema(&pool).await?;
 
-    let mgmt_router = mgmt::router(pool.clone());
+    let origin_cache = OriginCache::new();
+
+    let mgmt_router = mgmt::router(pool.clone(), origin_cache.clone());
 
     let client = Client::new();
     let router = Router::new()
         .route("/", post(handler))
         .route("/*path", post(handler))
-        .layer(Extension(pool))
+        .layer(Extension(pool.clone()))
+        .layer(Extension(origin_cache.clone()))
         .with_state(client);
 
-    Ok((router, mgmt_router, pool2))
+    let retry_queue = RetryQueue::new(pool, origin_cache);
+
+    Ok((router, mgmt_router, retry_queue))
 }
 
 #[tracing::instrument(level = "trace", "ingest", skip_all)]
 async fn handler(
     State(client): State<Client>,
     Extension(pool): Extension<SqlitePool>,
+    Extension(origin_cache): Extension<OriginCache>,
     req: Request<Body>,
 ) -> StdResult<impl IntoResponse, AppError> {
     let method = req.method().to_string();
@@ -84,7 +91,7 @@ async fn handler(
     let queued_req = insert_request(&pool, r).await?;
     let req_id = queued_req.id;
 
-    let is_success = proxy(&pool, &client, queued_req).await?;
+    let is_success = proxy(&pool, &origin_cache, &client, queued_req).await?;
 
     if is_success {
         mark_complete(&pool, req_id).await?;

src/main.rs

@@ -1,8 +1,5 @@
-use std::time::Duration;
-
 use anyhow::Result;
 use clap::Parser;
-use tokio::time;
 use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt};
 
 #[derive(Parser, Debug)]
@@ -28,7 +25,7 @@ async fn main() -> Result<()> {
         None => soldr::Config::default(),
     };
 
-    let (ingest, mgmt, pool) = soldr::app(&config).await?;
+    let (ingest, mgmt, retry_queue) = soldr::app(&config).await?;
 
     let mgmt_listener = config.management_listener.parse()?;
     let ingest_listener = config.ingest_listener.parse()?;
@@ -45,14 +42,7 @@ async fn main() -> Result<()> {
 
     tokio::spawn(async move {
         tracing::info!("starting retry queue");
-        let mut interval = time::interval(Duration::from_secs(60));
-
-        loop {
-            interval.tick().await;
-            let pool2 = pool.clone();
-            tracing::trace!("retrying failed requests");
-            soldr::queue::tick(pool2).await;
-        }
+        retry_queue.start().await;
     });
 
     tracing::info!("ingest listening on {}", ingest_listener);

src/mgmt.rs

@@ -9,15 +9,17 @@ use serde::{Deserialize, Serialize};
 use sqlx::sqlite::SqlitePool;
 use tracing::Level;
 
+use crate::cache::OriginCache;
 use crate::db;
 use crate::error::AppError;
 
-pub fn router(pool: SqlitePool) -> Router {
+pub fn router(pool: SqlitePool, origin_cache: OriginCache) -> Router {
     Router::new()
         .route("/origins", post(create_origin))
         .route("/requests", get(list_requests))
         .route("/attempts", get(list_attempts))
         .layer(Extension(pool))
+        .layer(Extension(origin_cache))
 }
 
 async fn list_requests(
@@ -52,6 +54,7 @@ pub struct CreateOrigin {
 
 async fn create_origin(
     Extension(pool): Extension<SqlitePool>,
+    Extension(origin_cache): Extension<OriginCache>,
     Json(payload): Json<CreateOrigin>,
 ) -> StdResult<Json<db::Origin>, AppError> {
     let span = tracing::span!(Level::TRACE, "create_origin");
@@ -61,5 +64,8 @@ async fn create_origin(
     let origin = db::insert_origin(&pool, &payload.domain, &payload.origin_uri).await?;
     tracing::debug!("response = {:?}", &origin);
 
+    let origins = db::list_origins(&pool).await?;
+    origin_cache.refresh(origins).unwrap();
+
     Ok(Json(origin))
 }

src/proxy.rs

@@ -6,14 +6,19 @@ use hyper::Body;
 use hyper::Response;
 use sqlx::SqlitePool;
 
+use crate::cache::OriginCache;
 use crate::db::insert_attempt;
-use crate::db::list_origins;
 use crate::db::QueuedRequest;
 
 pub type Client = hyper::client::Client<HttpConnector, Body>;
 
-pub async fn proxy(pool: &SqlitePool, client: &Client, mut req: QueuedRequest) -> Result<bool> {
-    let uri = map_origin(pool, &req).await?;
+pub async fn proxy(
+    pool: &SqlitePool,
+    origin_cache: &OriginCache,
+    client: &Client,
+    mut req: QueuedRequest,
+) -> Result<bool> {
+    let uri = map_origin(origin_cache, &req).await?;
 
     if uri.is_none() {
         // no origin found, so mark as complete and move on
@@ -45,7 +50,7 @@ pub async fn proxy(pool: &SqlitePool, client: &Client, mut req: QueuedRequest) -
     Ok(is_success)
 }
 
-async fn map_origin(pool: &SqlitePool, req: &QueuedRequest) -> Result<Option<Uri>> {
+async fn map_origin(origin_cache: &OriginCache, req: &QueuedRequest) -> Result<Option<Uri>> {
     let uri = Uri::try_from(&req.uri)?;
     let parts = uri.into_parts();
 
@@ -72,14 +77,10 @@ async fn map_origin(pool: &SqlitePool, req: &QueuedRequest) -> Result<Option<Uri
     };
     tracing::debug!("authority = {}", &authority);
 
-    let origins = list_origins(pool).await?;
-    tracing::debug!("origins = {:?}", &origins);
-    let matching_origin = origins
-        .iter()
-        .find(|origin| origin.domain == authority.as_str());
+    let matching_origin = origin_cache.get(authority.as_str());
 
     let origin_uri = match matching_origin {
-        Some(origin) => &origin.origin_uri,
+        Some(origin) => origin.origin_uri,
         None => {
             tracing::trace!("no match found");
             return Ok(None);

src/queue.rs

@@ -1,25 +1,52 @@
+use std::time::Duration;
+
 use anyhow::Result;
 use sqlx::sqlite::SqlitePool;
+use tokio::time;
+
+use crate::cache::OriginCache;
 
 use crate::{
     db::{list_failed_requests, mark_complete, mark_error, QueuedRequest},
     proxy::{self, Client},
 };
 
-pub async fn tick(pool: SqlitePool) {
-    if let Err(err) = do_tick(pool).await {
-        // TODO flow through the request id
-        tracing::error!("tick error {:?}", err);
+pub struct RetryQueue {
+    pool: SqlitePool,
+    origin_cache: OriginCache,
+}
+
+impl RetryQueue {
+    pub fn new(pool: SqlitePool, origin_cache: OriginCache) -> Self {
+        Self { pool, origin_cache }
+    }
+
+    pub async fn start(&self) {
+        let mut interval = time::interval(Duration::from_secs(60));
+
+        loop {
+            interval.tick().await;
+            tracing::trace!("retrying failed requests");
+            self.tick().await;
+        }
+    }
+
+    pub async fn tick(&self) {
+        if let Err(err) = do_tick(&self.pool, &self.origin_cache).await {
+            // TODO flow through the request id
+            tracing::error!("tick error {:?}", err);
+        }
     }
 }
 
-async fn do_tick(pool: SqlitePool) -> Result<()> {
-    let requests = list_failed_requests(&pool).await?;
+async fn do_tick(pool: &SqlitePool, origin_cache: &OriginCache) -> Result<()> {
+    let requests = list_failed_requests(pool).await?;
 
     let mut tasks = Vec::with_capacity(requests.len());
     for request in requests {
         let pool2 = pool.clone();
-        tasks.push(tokio::spawn(retry_request(pool2, request)));
+        let origin_cache2 = origin_cache.clone();
+        tasks.push(tokio::spawn(retry_request(pool2, origin_cache2, request)));
     }
 
     for task in tasks {
@@ -32,12 +59,16 @@ async fn do_tick(pool: SqlitePool) -> Result<()> {
     Ok(())
 }
 
-async fn retry_request(pool: SqlitePool, request: QueuedRequest) -> Result<()> {
+async fn retry_request(
+    pool: SqlitePool,
+    origin_cache: OriginCache,
+    request: QueuedRequest,
+) -> Result<()> {
     tracing::trace!("retrying {:?}", &request);
 
     let req_id = request.id;
     let client = Client::new();
-    let is_success = proxy::proxy(&pool, &client, request).await?;
+    let is_success = proxy::proxy(&pool, &origin_cache, &client, request).await?;
 
     if is_success {
         mark_complete(&pool, req_id).await?;