update

Cargo.toml

@@ -8,9 +8,9 @@ license = "MIT"
 [dependencies]
 filetime = "0.2"
 rand = "0.8"
+tempfile = "3"
 
 [dev-dependencies]
 proptest = "1"
 proptest-derive = "0.3"
-tempfile = "3"
 test_dir = "0.1"

src/lib.rs

@@ -13,6 +13,7 @@ pub use readonly::ReadOnlyCacheBuilder;
 pub use sharded::ShardedCache;
 pub use stack::Cache;
 pub use stack::CacheBuilder;
+pub use stack::CacheHitAction;
 
 /// Sharded cache keys consist of a filename and two hash values.  The
 /// two hashes should be computed by distinct functions of the key's

src/stack.rs

@@ -3,12 +3,14 @@
 //! between plain and sharded caches via late binding, and lets
 //! callers transparently handle misses by looking in a series of
 //! secondary cache directories.
+use std::borrow::Cow;
 use std::fs::File;
 use std::io::Error;
 use std::io::ErrorKind;
 use std::io::Result;
 use std::path::Path;
 use std::sync::Arc;
+use tempfile::NamedTempFile;
 
 use crate::Key;
 use crate::PlainCache;
@@ -27,6 +29,10 @@ trait FullCache:
     /// Implicitly "touches" the cached file if it exists.
     fn get(&self, key: Key) -> Result<Option<File>>;
 
+    /// Returns a temporary directory suitable for temporary files
+    /// that will be published as `key`.
+    fn temp_dir(&self, key: Key) -> Result<Cow<Path>>;
+
     /// Inserts or overwrites the file at `value` as `key` in the
     /// sharded cache directory.
     ///
@@ -53,6 +59,10 @@ impl FullCache for PlainCache {
         PlainCache::get(self, key.name)
     }
 
+    fn temp_dir(&self, _key: Key) -> Result<Cow<Path>> {
+        PlainCache::temp_dir(self)
+    }
+
     fn set(&self, key: Key, value: &Path) -> Result<()> {
         PlainCache::set(self, key.name, value)
     }
@@ -71,6 +81,10 @@ impl FullCache for ShardedCache {
         ShardedCache::get(self, key)
     }
 
+    fn temp_dir(&self, key: Key) -> Result<Cow<Path>> {
+        ShardedCache::temp_dir(self, Some(key))
+    }
+
     fn set(&self, key: Key, value: &Path) -> Result<()> {
         ShardedCache::set(self, key, value)
     }
@@ -104,6 +118,17 @@ pub struct Cache {
     read_side: ReadOnlyCache,
 }
 
+/// What to do with a cache hit in a `get_or_update` call?
+pub enum CacheHitAction {
+    /// Return the cache hit as is.
+    Accept,
+    /// Return the cache hit after promoting it to the current write
+    /// cache directory, if necessary.
+    Promote,
+    /// Replace with and return a new file.
+    Replace,
+}
+
 impl CacheBuilder {
     /// Returns a fresh empty builder.
     pub fn new() -> Self {
@@ -214,6 +239,90 @@ impl Cache {
         )
     }
 
+    /// Attempts to find a cache entry for `key`.  If there is none,
+    /// populates the cache with a file filled by `fill`.  Returns a
+    /// file in all cases (unless the call fails with an error).
+    pub fn ensure<'a>(
+        &self,
+        key: impl Into<Key<'a>>,
+        fill: impl FnOnce(&mut File) -> Result<()>,
+    ) -> Result<File> {
+        fn judge(_: &mut File) -> CacheHitAction {
+            CacheHitAction::Promote
+        }
+
+        self.get_or_update(key, judge, fill)
+    }
+
+    /// Attempts to find a cache entry for `key`.  If there is none,
+    /// populates the write cache (if possible) with a file, once
+    /// filled by `fill`; otherwise obeys the value returned by
+    /// `judge`.
+    pub fn get_or_update<'a>(
+        &self,
+        key: impl Into<Key<'a>>,
+        judge: impl FnOnce(&mut File) -> CacheHitAction,
+        fill: impl FnOnce(&mut File) -> Result<()>,
+    ) -> Result<File> {
+        // Attempts to return the `FullCache` for this `Cache`.
+        fn get_write_cache(this: &Cache) -> Result<&dyn FullCache> {
+            match this.write_side.as_ref() {
+                Some(cache) => Ok(cache.as_ref()),
+                None => Err(Error::new(
+                    ErrorKind::Unsupported,
+                    "no kismet write cache defined",
+                )),
+            }
+        }
+
+        // Promotes `file` to `cache`.
+        fn promote(cache: &dyn FullCache, key: Key, mut file: File) -> Result<File> {
+            use std::io::Seek;
+
+            let mut tmp = NamedTempFile::new_in(cache.temp_dir(key)?)?;
+            std::io::copy(&mut file, tmp.as_file_mut())?;
+            cache.put(key, tmp.path())?;
+
+            // We got a read-only file.  Rewind it before returning.
+            file.seek(std::io::SeekFrom::Start(0))?;
+            Ok(file)
+        }
+
+        let cache = get_write_cache(self)?;
+        let key: Key = key.into();
+        // Should we public with `set` to replace any old cached value?
+        let mut replace = false;
+        if let Some(mut file) = cache.get(key)? {
+            match judge(&mut file) {
+                // Promote is a no-op if the file is already in the write cache.
+                CacheHitAction::Accept | CacheHitAction::Promote => return Ok(file),
+                CacheHitAction::Replace => replace = true,
+            }
+        } else if let Some(mut file) = self.read_side.get(key)? {
+            match judge(&mut file) {
+                CacheHitAction::Accept => return Ok(file),
+                CacheHitAction::Promote => return promote(get_write_cache(self)?, key, file),
+                CacheHitAction::Replace => replace = true,
+            }
+        }
+
+        // We either have to replace or ensure there is a cache entry.
+        // Either way, start by populating a temporary file.
+        let mut tmp = NamedTempFile::new_in(cache.temp_dir(key)?)?;
+        fill(tmp.as_file_mut())?;
+
+        // Grab a read-only return value before publishing the file.
+        let path = tmp.path();
+        let ret = File::open(path)?;
+        if replace {
+            cache.set(key, path)?;
+        } else {
+            cache.put(key, path)?;
+        }
+
+        Ok(ret)
+    }
+
     /// Inserts or overwrites the file at `value` as `key` in the
     /// write cache directory.  This will always fail with
     /// `Unsupported` if no write cache was defined.
@@ -315,6 +424,114 @@ mod test {
         assert!(matches!(cache.touch(&TestKey::new("foo")), Ok(false)));
     }
 
+    // Fail to find a file, ensure it, then see that we can get it.
+    #[test]
+    fn test_ensure() {
+        use std::io::{Read, Write};
+        use test_dir::{DirBuilder, TestDir};
+
+        let temp = TestDir::temp();
+        let cache = CacheBuilder::new().writer(temp.path("."), 1, 10).build();
+        let key = TestKey::new("foo");
+
+        // The file doesn't exist initially.
+        assert!(matches!(cache.get(&key), Ok(None)));
+
+        {
+            let mut populated = cache
+                .ensure(&key, |file| file.write_all(b"test"))
+                .expect("ensure must succeed");
+
+            let mut dst = Vec::new();
+            populated.read_to_end(&mut dst).expect("read must succeed");
+            assert_eq!(&dst, b"test");
+        }
+
+        // And now get the file again.
+        {
+            let mut fetched = cache
+                .get(&key)
+                .expect("get must succeed")
+                .expect("file must be found");
+
+            let mut dst = Vec::new();
+            fetched.read_to_end(&mut dst).expect("read must succeed");
+            assert_eq!(&dst, b"test");
+        }
+    }
+
+    // Use a two-level cache, and make sure `ensure` promotes copies from
+    // the backup to the primary location.
+    #[test]
+    fn test_ensure_promote() {
+        use std::io::{Read, Write};
+        use tempfile::NamedTempFile;
+        use test_dir::{DirBuilder, FileType, TestDir};
+
+        let temp = TestDir::temp()
+            .create("cache", FileType::Dir)
+            .create("extra_plain", FileType::Dir);
+
+        // Populate the plain cache in `extra_plain` with one file.
+        {
+            let cache = PlainCache::new(temp.path("extra_plain"), 10);
+
+            let tmp = NamedTempFile::new_in(cache.temp_dir().expect("temp_dir must succeed"))
+                .expect("new temp file must succeed");
+            tmp.as_file()
+                .write_all(b"initial")
+                .expect("write must succeed");
+
+            cache.put("foo", tmp.path()).expect("put must succeed");
+        }
+
+        let cache = CacheBuilder::new()
+            .writer(temp.path("cache"), 1, 10)
+            .plain_reader(temp.path("extra_plain"))
+            .build();
+        let key = TestKey::new("foo");
+
+        // The file is found initially.
+        {
+            let mut fetched = cache
+                .get(&key)
+                .expect("get must succeed")
+                .expect("file must be found");
+
+            let mut dst = Vec::new();
+            fetched.read_to_end(&mut dst).expect("read must succeed");
+            assert_eq!(&dst, b"initial");
+        }
+
+        {
+            let mut populated = cache
+                .ensure(&key, |_| {
+                    unreachable!("should not be called for an extant file")
+                })
+                .expect("ensure must succeed");
+
+            let mut dst = Vec::new();
+            populated.read_to_end(&mut dst).expect("read must succeed");
+            assert_eq!(&dst, b"initial");
+        }
+
+        // And now get the file again, and make sure it doesn't come from the
+        // backup location.
+        {
+            let new_cache = CacheBuilder::new()
+                .writer(temp.path("cache"), 1, 10)
+                .build();
+            let mut fetched = new_cache
+                .get(&key)
+                .expect("get must succeed")
+                .expect("file must be found");
+
+            let mut dst = Vec::new();
+            fetched.read_to_end(&mut dst).expect("read must succeed");
+            assert_eq!(&dst, b"initial");
+        }
+    }
+
     // Smoke test a wrapped plain cache.
     #[test]
     fn smoke_test_plain() {