Store thumbnails in rocksdb.

src/library.rs

@@ -3,22 +3,33 @@ use std::fs::File;
 use std::io;
 use std::path::Path;
 use std::path::PathBuf;
+use std::sync::Arc;
 
 use google_photoslibrary1 as photos;
+use log::error;
 use log::info;
+use log::warn;
 use photos::schemas::Album;
 use photos::schemas::MediaItem;
+use rocksdb::DB;
+
+// Used to ensure DB is invalidated after schema changes.
+const LIBRARY_GENERATION: &'static str = "1";
 
 #[derive(Clone)]
 pub struct Library {
     root: PathBuf,
     originals_dir: PathBuf,
+    cache_db: Arc<DB>,
 }
 
 impl Library {
-    pub fn new(root: PathBuf) -> io::Result<Library> {
+    pub fn new(root: PathBuf) -> Result<Library, Box<dyn std::error::Error>> {
+        let db = DB::open_default(root.join("cache"))?;
+        let cache_db = Arc::new(db);
         let lib = Library {
             originals_dir: root.join("images").join("originals"),
+            cache_db,
             root,
         };
         if !lib.originals_dir.exists() {
@@ -30,6 +41,20 @@ impl Library {
         }
         Ok(lib)
     }
+    // Removes all data in the database from older schema.
+    pub fn clean_db(&self) -> Result<(), rocksdb::Error> {
+        Library::gc(LIBRARY_GENERATION, &self.cache_db)
+    }
+    fn gc(generation: &str, db: &DB) -> Result<(), rocksdb::Error> {
+        let gen = format!("{}/", generation);
+        for (k, _v) in db.iterator(rocksdb::IteratorMode::Start) {
+            if !k.starts_with(gen.as_bytes()) {
+                info!("deleting stale key: {}", String::from_utf8_lossy(&k));
+                db.delete(k)?;
+            }
+        }
+        Ok(())
+    }
     pub fn create_album_index(&self, albums: &Vec<Album>) -> io::Result<()> {
         // Serialize it to a JSON string.
         let j = serde_json::to_string(albums)?;
@@ -102,4 +127,113 @@ impl Library {
             None
         }
     }
+    // TODO(wathiede): make this a macro like format! to skip the second string create and copy.
+    fn generational_key(generation: &str, key: &str) -> String {
+        format!("{}/{}", generation, key)
+    }
+    pub fn thumbnail(&self, media_items_id: &str, dimensions: (u32, u32)) -> Option<Vec<u8>> {
+        fn generate_thumbnail(
+            lib: &Library,
+            media_items_id: &str,
+            dimensions: (u32, u32),
+        ) -> Result<Vec<u8>, io::Error> {
+            let (w, h) = dimensions;
+            match lib.original(&media_items_id) {
+                None => {
+                    warn!("Couldn't find original {}", &media_items_id);
+                    Err(io::Error::new(
+                        io::ErrorKind::NotFound,
+                        format!("{}", media_items_id),
+                    ))
+                }
+                Some(path) => {
+                    let orig_img = image::io::Reader::open(&path)?
+                        .with_guessed_format()?
+                        .decode()
+                        .map_err(|e| io::Error::new(io::ErrorKind::Other, e))?;
+                    // TODO (wathiede) resize..fill
+                    let img = orig_img.resize(w, h, image::imageops::FilterType::CatmullRom);
+                    let mut buf = Vec::new();
+                    img.write_to(&mut buf, image::ImageFormat::Jpeg)
+                        .map_err(|e| io::Error::new(io::ErrorKind::Other, e))?;
+                    Ok(buf)
+                }
+            }
+        }
+        fn cache_key(media_items_id: &str, dimensions: (u32, u32)) -> String {
+            let (w, h) = dimensions;
+            Library::generational_key(
+                LIBRARY_GENERATION,
+                &format!("{}-w={}-h={}", media_items_id, w, h),
+            )
+        }
+        let key = cache_key(media_items_id, dimensions);
+        let db = self.cache_db.clone();
+        match db.get(key.as_bytes()) {
+            // Cache hit, return bytes as-is.
+            Ok(Some(bytes)) => {
+                info!("cache HIT {}", key);
+                Some(bytes)
+            }
+            // Cache miss, fill cache and return.
+            Ok(None) => {
+                info!("cache MISS {}", key);
+                let bytes = match generate_thumbnail(self, media_items_id, dimensions) {
+                    Ok(bytes) => bytes,
+                    Err(e) => {
+                        error!("Failed to generate thumbnail for {}: {}", media_items_id, e);
+                        return None;
+                    }
+                };
+                match db.put(key.as_bytes(), &bytes) {
+                    Ok(_) => Some(bytes),
+                    Err(e) => {
+                        error!("Failed to put bytes to {}: {}", key, e);
+                        None
+                    }
+                }
+            }
+            // RocksDB error.
+            Err(e) => {
+                error!("Failed to search DB for {}: {}", key, e);
+                None
+            }
+        }
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use super::*;
+    use tempdir::TempDir;
+
+    #[test]
+    fn clean_db() {
+        let td = TempDir::new("photosync_test").expect("failed to create temporary directory");
+        eprintln!("creating database in {}", td.path().to_string_lossy());
+        let db = DB::open_default(td.path()).expect("failed to open DB");
+        let keys = vec!["one", "two", "three"];
+
+        fn get_keys(db: &DB) -> Vec<String> {
+            db.iterator(rocksdb::IteratorMode::Start)
+                .map(|(k, _v)| String::from_utf8(k.to_vec()).expect("key not utf-8"))
+                .collect()
+        }
+        for k in &keys {
+            for g in vec!["1", "2", "3"] {
+                db.put(Library::generational_key(g, k), k)
+                    .expect("failed to put");
+            }
+        }
+
+        assert_eq!(
+            get_keys(&db),
+            vec![
+                "1/one", "1/three", "1/two", "2/one", "2/three", "2/two", "3/one", "3/three",
+                "3/two"
+            ]
+        );
+        Library::gc("2", &db).expect("failed to GC DB");
+        assert_eq!(get_keys(&db), vec!["2/one", "2/three", "2/two",]);
+    }
 }

src/web.rs

@@ -2,7 +2,6 @@ use std::error::Error;
 use std::net::SocketAddr;
 use std::path::PathBuf;
 
-use log::info;
 use log::warn;
 use prometheus::Encoder;
 use rust_embed::RustEmbed;
@@ -32,6 +31,7 @@ fn metrics() -> impl Filter<Extract = (impl warp::reply::Reply,), Error = Reject
         .with(warp::reply::with::headers(text_headers))
 }
 
+// TODO(wathiede): add caching for hashed files.  Add at least etag for everything.
 fn index(path: warp::path::FullPath) -> Result<impl warp::Reply, warp::Rejection> {
     let path = path.as_str();
     let path = if path.ends_with("/") {
@@ -47,7 +47,6 @@ fn index(path: warp::path::FullPath) -> Result<impl warp::Reply, warp::Rejection
 
             Ok(warp::http::Response::builder()
                 .header("Content-Type", mime.essence_str())
-                .header("Content-Length", bytes.len())
                 .body(bytes.into_owned()))
         }
         None => Err(warp::reject::not_found()),
@@ -74,7 +73,6 @@ fn album(lib: Library, id: String) -> Result<impl warp::Reply, warp::Rejection>
 struct ImageParams {
     w: Option<u32>,
     h: Option<u32>,
-    c: Option<bool>,
 }
 
 fn image(
@@ -82,42 +80,17 @@ fn image(
     media_items_id: String,
     params: ImageParams,
 ) -> Result<impl warp::Reply, warp::Rejection> {
-    info!("image_id {} params {:?}", media_items_id, params);
-
-    match lib.original(&media_items_id) {
+    match lib.thumbnail(
+        &media_items_id,
+        (params.w.unwrap_or(0), params.h.unwrap_or(0)),
+    ) {
         None => {
             warn!("Couldn't find original {}", &media_items_id);
             Err(warp::reject::not_found())
         }
-        Some(path) => {
-            let orig_img = image::io::Reader::open(&path)
-                .map_err(|e| {
-                    warn!("Couldn't open {}: {}", path.to_string_lossy(), e);
-                    warp::reject::not_found()
-                })?
-                .with_guessed_format()
-                .map_err(|e| {
-                    warn!("Couldn't guess format {}: {}", path.to_string_lossy(), e);
-                    warp::reject::not_found()
-                })?
-                .decode()
-                .map_err(|e| {
-                    warn!("Couldn't decode {}: {}", path.to_string_lossy(), e);
-                    warp::reject::not_found()
-                })?;
-            let img = orig_img.resize(
-                params.w.unwrap_or(0),
-                params.h.unwrap_or(0),
-                image::imageops::FilterType::CatmullRom,
-            );
-            let mut buf = Vec::new();
-            img.write_to(&mut buf, image::ImageFormat::Jpeg)
-                .map_err(|e| {
-                    warn!("Couldn't write_to {}: {}", path.to_string_lossy(), e);
-                    warp::reject::not_found()
-                })?;
-            Ok(buf)
-        }
+        Some(bytes) => Ok(warp::http::Response::builder()
+            .header("Content-Type", "image/jpeg")
+            .body(bytes)),
     }
 }