From 9965e611080e343abbc6773346642d06c25b5ec3 Mon Sep 17 00:00:00 2001
From: Vincent Breitmoser <look@my.amazin.horse>
Date: Tue, 20 Jul 2021 12:10:57 +0200
Subject: [PATCH] web: introduce a cache for loaded epochs

---
 src/web/mod.rs         | 10 ++++++--
 src/web/vks_updates.rs | 53 ++++++++++++++++++++++++++++++++----------
 2 files changed, 49 insertions(+), 14 deletions(-)

diff --git a/src/web/mod.rs b/src/web/mod.rs
index 381fc07..5018e3a 100644
--- a/src/web/mod.rs
+++ b/src/web/mod.rs
@@ -1,4 +1,3 @@
-use rocket;
 use rocket::http::{Header, Status};
 use rocket::request;
 use rocket::outcome::Outcome;
@@ -30,6 +29,8 @@ use crate::database::{Database, KeyDatabase, Query};
 use crate::database::types::Fingerprint;
 use crate::Result;
 
+use crate::web::vks_updates::UpdateEpochCache;
+
 use std::convert::TryInto;
 
 mod hkp;
@@ -38,6 +39,7 @@ mod maintenance;
 mod vks;
 mod vks_web;
 mod vks_api;
+mod vks_updates;
 mod debug_web;
 
 use crate::web::maintenance::MaintenanceMode;
@@ -417,6 +419,8 @@ fn rocket_factory(mut rocket: rocket::Rocket) -> Result<rocket::Rocket> {
         vks_web::verify_confirm_form,
         vks_web::quick_upload,
         vks_web::quick_upload_proceed,
+        // Update Manifests
+        vks_updates::get_update_manifest,
         // Debug
         debug_web::debug_info,
         // HKP
@@ -443,6 +447,7 @@ fn rocket_factory(mut rocket: rocket::Rocket) -> Result<rocket::Rocket> {
     let rate_limiter = configure_rate_limiter(rocket.config())?;
     let maintenance_mode = configure_maintenance_mode(rocket.config())?;
     let localized_template_list = configure_localized_template_list(rocket.config())?;
+    let update_cache = UpdateEpochCache::new();
     println!("{:?}", localized_template_list);
 
     let prometheus = configure_prometheus(rocket.config());
@@ -462,6 +467,7 @@ fn rocket_factory(mut rocket: rocket::Rocket) -> Result<rocket::Rocket> {
        .manage(db_service)
        .manage(rate_limiter)
        .manage(localized_template_list)
+       .manage(update_cache)
        .mount("/", routes);
 
     if let Some(prometheus) = prometheus {
@@ -568,7 +574,7 @@ pub mod tests {
     use std::io::Write;
     use std::path::Path;
     use tempfile::{tempdir, TempDir};
-    use super::rocket;
+    use rocket;
     use rocket::local::{Client, LocalResponse};
     use rocket::http::Status;
     use rocket::http::ContentType;
diff --git a/src/web/vks_updates.rs b/src/web/vks_updates.rs
index 29de9bd..fc32518 100644
--- a/src/web/vks_updates.rs
+++ b/src/web/vks_updates.rs
@@ -1,4 +1,4 @@
-use std::io::{BufRead, BufReader};
+use std::{collections::HashMap, sync::RwLock};
 
 use rocket::http::hyper::header::{Expires, HttpDate};
 
@@ -8,6 +8,14 @@ use crate::Result;
 
 const EPOCH_SERVE_LIMIT: u32 = 120;
 
+pub struct UpdateEpochCache(RwLock<HashMap<Epoch, Vec<u32>>>);
+
+impl UpdateEpochCache {
+    pub fn new() -> Self {
+        UpdateEpochCache(RwLock::new(HashMap::new()))
+    }
+}
+
 #[derive(Responder)]
 pub enum ManifestUpdateResponse {
     #[response(status = 200, content_type = "application/pgp-keystore-update-manifest")]
@@ -21,6 +29,7 @@ pub enum ManifestUpdateResponse {
 #[get("/vks/v1/updates/<epoch>")]
 pub fn get_update_manifest(
     db: rocket::State<KeyDatabase>,
+    cache: rocket::State<UpdateEpochCache>,
     epoch: u32,
 ) -> ManifestUpdateResponse {
     let epoch_now = Epoch::current().unwrap();
@@ -33,19 +42,15 @@ pub fn get_update_manifest(
             format!("Updafe manifest data only available for the last {} epochs", EPOCH_SERVE_LIMIT));
     }
 
+    if let Err(e) = provision_cache_since(&db, &cache, &epoch_since) {
+        return ManifestUpdateResponse::NotFound(e.to_string());
+    }
+
+    let cache_lock = cache.0.read().expect("lock can't be poisoned");
     let mut epoch_data: Vec<&[u32]> = Vec::with_capacity((epoch_now - epoch_since) as usize);
     for e in epoch_since.until(epoch_now).expect("epoch_since is before epoch_now") {
-        match db.read_log_epoch(e) {
-            Err(e) => {
-                eprintln!("Error reading epoch: {:?}", e);
-                return ManifestUpdateResponse::NotFound(
-                    "No data found for requested update epoch".to_owned());
-            }
-            Ok(None) => {
-                return ManifestUpdateResponse::NotFound(
-                    "No data found for requested update epoch".to_owned());
-            }
-            Ok(Some(v)) => epoch_data.push(&v),
+        if let Some(v) = cache_lock.get(&e) {
+            epoch_data.push(&v);
         }
     }
 
@@ -54,3 +59,27 @@ pub fn get_update_manifest(
     let expires = Expires(HttpDate(epoch_now.succ().expect("We're not at the end of time").into()));
     ManifestUpdateResponse::Binary(manifest.to_vec(), expires)
 }
+
+fn provision_cache_since(db: &KeyDatabase, cache: &UpdateEpochCache, epoch_since: &Epoch) -> Result<()> {
+    let epoch_now = Epoch::current().expect("not the end of time");
+    let mut cache_lock = cache.0.write().expect("lock can't be poisoned");
+
+    for epoch in epoch_since.until(epoch_now).expect("epoch_since is before epoch_now") {
+        if cache_lock.contains_key(&epoch) {
+            continue;
+        }
+        match db.read_log_epoch(epoch) {
+            Err(e) => {
+                eprintln!("{:?}", e);
+                Err(anyhow!("No update manifest data available for requested epoch"))?
+            },
+            Ok(None) => Err(anyhow!("No update manifest data available for requested epoch"))?,
+            Ok(Some(prefixes)) => cache_lock.insert(epoch, prefixes),
+        };
+    }
+
+    let ref epoch_earliest = epoch_now - EPOCH_SERVE_LIMIT;
+    cache_lock.retain(|k, _| k > epoch_earliest);
+
+    Ok(())
+}