feat: add lru memcache layer

Author: DevBlocky

Cargo.lock

@@ -22,6 +22,8 @@ serde = {version = "1.0.125", features = ["derive"]}
 tokio = {version = "1.5.0", features = ["fs", "io-util"]}
 bytes = "1.0.1"
 async-trait = "0.1.50"
+lru = "0.6.5"
+parking_lot = "0.11.1"
 
 [dev-dependencies]
 tokio = {version = "1.5.0", features = ["full"]}

Cargo.toml

@@ -47,6 +47,9 @@ impl CacheBuilder {
             dir_depth: 2,
             track_access: false,
 
+            // default to no in-mem lru
+            lru_size: 0,
+
             // default buffer sizes to 8kb
             rbuff_sz: 8192,
             wbuff_sz: 8192,
@@ -71,6 +74,17 @@ impl CacheBuilder {
         self
     }
 
+    /// Sets the maximum size (in bytes) for the in-memory Least-Recently-Used cache.
+    ///
+    /// **Default is `0`**
+    ///
+    /// Setting this above `0` will create an in-memory store for recently-used entries. This will
+    /// use up more RAM, but will also significantly increase speed on memcache `HIT`s.
+    pub fn memory_lru_max_size(mut self, size: usize) -> Self {
+        self.opts.lru_size = size;
+        self
+    }
+
     /// Changes the in-memory buffer sizes for reading and writing `fs` operations.
     ///
     /// **Default is `8kb` (`8196`)**

src/builder.rs

@@ -1,4 +1,4 @@
-use crate::{ForcepError, MetaDb, Metadata, Result};
+use crate::{mem_cache::MemCache, ForcepError, MetaDb, Metadata, Result};
 use bytes::Bytes;
 use std::io;
 use std::path;
@@ -25,6 +25,9 @@ pub(crate) struct Options {
     pub(crate) dir_depth: u8,
     pub(crate) track_access: bool,
 
+    // maximum size of the in-memory lru in bytes
+    pub(crate) lru_size: usize,
+
     // read and write buffer sizes
     pub(crate) rbuff_sz: usize,
     pub(crate) wbuff_sz: usize,
@@ -41,6 +44,12 @@ pub(crate) struct Options {
 /// This cache can evict items with a number of different eviction algorithms. To see more, see
 /// [`evict_with`] and the [`evictors`] module.
 ///
+/// # Memory Cache
+///
+/// An in-memory cache can be optionally enabled as a layer over the regular on-disk cache. The
+/// memcache provides fast `HIT`s for recently used entries, circumventing filesystem operations
+/// altogether. To enable, use the [`CacheBuilder`]`::memory_lru_max_size` method.
+///
 /// # Examples
 ///
 /// ```rust
@@ -57,9 +66,11 @@ pub(crate) struct Options {
 ///
 /// [`evict_with`]: #method.evict_with
 /// [`evictors`]: crate::evictors
+/// [`CacheBuilder`]: crate::CacheBuilder
 #[derive(Debug)]
 pub struct Cache {
     meta: MetaDb,
+    mem: MemCache,
     opts: Options,
 }
 
@@ -97,6 +108,7 @@ impl Cache {
         meta_path.push("index");
         Ok(Self {
             meta: MetaDb::new(&meta_path)?,
+            mem: MemCache::new(opts.lru_size),
             opts,
         })
     }
@@ -120,6 +132,15 @@ impl Cache {
         buf
     }
 
+    /// Tracks the access for a cache entry if the option is enabled
+    #[inline]
+    fn track_access_for(&self, k: &[u8]) -> Result<()> {
+        if self.opts.track_access {
+            self.meta.track_access_for(k)?;
+        }
+        Ok(())
+    }
+
     /// Reads an entry from the database, returning a vector of bytes that represent the entry.
     ///
     /// # Not Found
@@ -152,9 +173,15 @@ impl Cache {
     /// ```
     pub async fn read<K: AsRef<[u8]>>(&self, key: K) -> Result<Bytes> {
         use tokio::io::AsyncReadExt;
+        let k = key.as_ref();
+
+        // look in the memory cache to see if it's there and return if it is
+        if let Some(val) = self.mem.get(k) {
+            return self.track_access_for(k).map(|_| val);
+        }
 
         let file = {
-            let path = self.path_from_key(key.as_ref());
+            let path = self.path_from_key(k);
             afs::OpenOptions::new()
                 .read(true)
                 .open(&path)
@@ -175,12 +202,10 @@ impl Cache {
             .await
             .map_err(ForcepError::Io)?;
 
-        // track this access if the flag is set
-        if self.opts.track_access {
-            self.meta.track_access_for(key.as_ref())?;
-        }
-
-        Ok(Bytes::from(buf))
+        self.track_access_for(k)?;
+        let bytes = Bytes::from(buf);
+        self.mem.put(k, Bytes::clone(&bytes));
+        Ok(bytes)
     }
 
     /// Writes an entry with the specified key to the cache database. This will replace the
@@ -227,6 +252,9 @@ impl Cache {
             .await
             .map_err(ForcepError::Io)?;
 
+        if !self.mem.is_nil() {
+            self.mem.put(key, Bytes::from(Vec::from(value)));
+        }
         self.meta.insert_metadata_for(key, value)
     }
 

src/cache.rs

@@ -13,6 +13,7 @@
 //!
 //! - Asynchronous APIs
 //! - Fast and reliable reading/writing
+//! - Optional memory-cache layer
 //! - Tuned for large-file databases
 //! - Included cache eviction (LRU/FIFO)
 //! - Easily accessible value metadata
@@ -33,18 +34,20 @@
 //! # Examples
 //!
 //! ```rust,no_run
-//! #[tokio::main]
-//! async fn main() {
+//! use std::error::Error;
 //! use forceps::Cache;
 //!
-//! let cache = Cache::new("./cache")
-//!     .build()
-//!     .await
-//!     .unwrap();
+//! #[tokio::main]
+//! async fn main() -> Result<(), Box<dyn Error>> {
+//!     let cache = Cache::new("./cache")
+//!         .build()
+//!         .await?;
+//!
+//!     cache.write(b"MY_KEY", b"Hello World").await?;
+//!     let data = cache.read(b"MY_KEY").await?;
+//!     assert_eq!(data.as_ref(), b"Hello World");
 //!
-//! cache.write(b"MY_KEY", b"Hello World").await.unwrap();
-//! let data = cache.read(b"MY_KEY").await.unwrap();
-//! assert_eq!(data.as_ref(), b"Hello World");
+//!     Ok(())
 //! }
 //! ```
 
@@ -97,6 +100,7 @@ impl error::Error for ForcepError {
     }
 }
 
+mod mem_cache;
 mod tmp;
 
 mod builder;

src/lib.rs

@@ -0,0 +1,145 @@
+use crate::Md5Bytes;
+use bytes::Bytes;
+use lru::LruCache;
+use parking_lot::Mutex;
+use std::sync::atomic::{AtomicUsize, Ordering};
+
+type Lru = LruCache<Md5Bytes, Bytes>;
+
+#[inline]
+fn hash_key(k: &[u8]) -> Md5Bytes {
+    md5::compute(k).into()
+}
+
+#[derive(Debug)]
+struct MemCacheInner {
+    cache: Mutex<Lru>,
+
+    cap: usize,
+    current: AtomicUsize,
+}
+
+#[derive(Debug)]
+pub(crate) struct MemCache(Option<MemCacheInner>);
+
+impl MemCacheInner {
+    fn new(cap: usize) -> Self {
+        Self {
+            cache: Mutex::new(Lru::unbounded()),
+            cap,
+            current: AtomicUsize::new(0),
+        }
+    }
+
+    fn get(&self, k: &[u8]) -> Option<Bytes> {
+        let mut guard = self.cache.lock();
+        guard.get(&hash_key(k)).map(Bytes::clone)
+    }
+
+    fn put(&self, k: &[u8], v: Bytes) -> Option<Bytes> {
+        let mut guard = self.cache.lock();
+        let v_sz = v.len();
+        let other = guard.put(hash_key(k), v);
+
+        // sub the last value and add the new one to the total bytes counter
+        if let Some(ref other) = other {
+            self.current.fetch_sub(other.len(), Ordering::Relaxed);
+        }
+        let updated_sz = self.current.fetch_add(v_sz, Ordering::SeqCst) + v_sz;
+
+        // if the new size is greater than the cap, start evicting items
+        println!("{} / {}", updated_sz, self.cap);
+        if updated_sz > self.cap {
+            self.evict(&mut guard, updated_sz);
+        }
+
+        other
+    }
+
+    fn evict(&self, lru: &mut Lru, mut current: usize) -> usize {
+        // pop items until it meets size requirement
+        loop {
+            match lru.pop_lru() {
+                Some((_, b)) => current -= b.len(),
+                None => break,
+            }
+            if current <= self.cap {
+                break;
+            }
+        }
+        self.current.swap(current, Ordering::SeqCst);
+        current
+    }
+
+    #[cfg(test)]
+    fn peek(&self, k: &[u8]) -> Option<Bytes> {
+        let guard = self.cache.lock();
+        guard.peek(&hash_key(k)).map(Bytes::clone)
+    }
+}
+
+impl MemCache {
+    pub(crate) fn new(bytes_cap: usize) -> Self {
+        if bytes_cap == 0 {
+            Self(None)
+        } else {
+            Self(Some(MemCacheInner::new(bytes_cap)))
+        }
+    }
+
+    #[inline]
+    pub(crate) fn is_nil(&self) -> bool {
+        self.0.is_none()
+    }
+
+    #[inline]
+    pub(crate) fn get(&self, k: &[u8]) -> Option<Bytes> {
+        self.0.as_ref().and_then(|c| c.get(k))
+    }
+    #[inline]
+    pub(crate) fn put(&self, k: &[u8], v: Bytes) -> Option<Bytes> {
+        self.0.as_ref().and_then(|c| c.put(k, v))
+    }
+
+    // functions for tests
+    #[cfg(test)]
+    fn peek(&self, k: &[u8]) -> Option<Bytes> {
+        self.0.as_ref().and_then(|c| c.peek(k))
+    }
+    #[cfg(test)]
+    fn current_size(&self) -> Option<usize> {
+        self.0.as_ref().map(|x| x.current.load(Ordering::SeqCst))
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use super::*;
+    const D: &'static [u8] = &[0; 4096];
+
+    #[test]
+    fn verify_eviction() {
+        let cache = MemCache::new(D.len() * 2);
+        cache.put(b"ENT1", Bytes::from(D));
+        cache.put(b"ENT2", Bytes::from(D));
+        // verify that ENT1 still exists in cache (no eviction yet)
+        assert!(cache.peek(b"ENT1").is_some());
+
+        // with the new put, cache should evict ENT1
+        cache.put(b"ENT3", Bytes::from(D));
+        assert!(cache.peek(b"ENT1").is_none());
+    }
+
+    #[test]
+    fn size_updates() {
+        // make sure current size is being updated after insertion of items
+        let cache = MemCache::new(D.len() * 2);
+        cache.put(b"ENT1", Bytes::from(D));
+        cache.put(b"ENT2", Bytes::from(D));
+        assert_eq!(cache.current_size().unwrap(), D.len() * 2);
+
+        // this will verify that the last replaced entry has its byte count removed
+        cache.put(b"ENT2", Bytes::from(D));
+        assert_eq!(cache.current_size().unwrap(), D.len() * 2);
+    }
+}