blobrepo: add support for doing in-memory writes to blobstore

Summary:
This proves to be extremely useful for the upcoming bonsai
verification code.

The in-memory stuff is more complicated for the database backends, so punt on
that for now with some warnings.

Reviewed By: farnz

Differential Revision: D8909426

fbshipit-source-id: 1d66d877cfa48ef06bbe614f37c66cf6c2f0446c

blobrepo/src/repo.rs

@@ -25,8 +25,8 @@ use stats::Timeseries;
 use time_ext::DurationExt;
 use uuid::Uuid;
 
-use blobstore::{new_memcache_blobstore, Blobstore, EagerMemblob, MemoizedBlobstore,
-                PrefixBlobstore};
+use blobstore::{new_memcache_blobstore, Blobstore, EagerMemblob, MemWritesBlobstore,
+                MemoizedBlobstore, PrefixBlobstore};
 use bonsai_hg_mapping::{BonsaiHgMapping, CachingBonsaiHgMapping, MysqlBonsaiHgMapping,
                         SqliteBonsaiHgMapping};
 use bookmarks::{self, Bookmark, BookmarkPrefix, Bookmarks};
@@ -287,6 +287,40 @@ impl BlobRepo {
         ))
     }
 
+    /// Convert this BlobRepo instance into one that only does writes in memory.
+    ///
+    /// ------------
+    /// IMPORTANT!!!
+    /// ------------
+    /// Currently this applies to the blobstore *ONLY*. A future improvement would be to also
+    /// do database writes in-memory.
+    #[allow(non_snake_case)]
+    pub fn in_memory_writes_READ_DOC_COMMENT(self) -> BlobRepo {
+        let BlobRepo {
+            logger,
+            bookmarks,
+            blobstore,
+            filenodes,
+            changesets,
+            bonsai_hg_mapping,
+            repoid,
+        } = self;
+
+        // Drop the PrefixBlobstore (it will be wrapped up in one again by BlobRepo::new)
+        let blobstore = blobstore.into_inner();
+        let blobstore = Arc::new(MemWritesBlobstore::new(blobstore));
+
+        BlobRepo::new(
+            logger,
+            bookmarks,
+            blobstore,
+            filenodes,
+            changesets,
+            bonsai_hg_mapping,
+            repoid,
+        )
+    }
+
     fn fetch<K>(&self, key: &K) -> impl Future<Item = K::Value, Error = Error> + Send
     where
         K: MononokeId,

blobstore/src/lib.rs

@@ -58,6 +58,9 @@ pub use memblob::{EagerMemblob, LazyMemblob};
 mod memcache_cache_lease;
 pub use memcache_cache_lease::{new_memcache_blobstore, new_memcache_blobstore_no_lease};
 
+mod mem_writes;
+pub use mem_writes::MemWritesBlobstore;
+
 mod prefix;
 pub use prefix::PrefixBlobstore;
 

blobstore/src/mem_writes.rs

@@ -0,0 +1,190 @@
+// Copyright (c) 2018-present, Facebook, Inc.
+// All Rights Reserved.
+//
+// This software may be used and distributed according to the terms of the
+// GNU General Public License version 2 or any later version.
+
+use failure::Error;
+use futures::{Future, IntoFuture};
+use futures::future::Either;
+
+use futures_ext::{BoxFuture, FutureExt};
+
+use mononoke_types::BlobstoreBytes;
+
+use {Blobstore, EagerMemblob};
+
+/// A blobstore wrapper that reads from the underlying blobstore but writes to memory.
+#[derive(Clone, Debug)]
+pub struct MemWritesBlobstore<T: Blobstore + Clone> {
+    inner: T,
+    // TODO: replace with chashmap or evmap
+    memblob: EagerMemblob,
+}
+
+impl<T: Blobstore + Clone> MemWritesBlobstore<T> {
+    pub fn new(blobstore: T) -> Self {
+        Self {
+            inner: blobstore,
+            memblob: EagerMemblob::new(),
+        }
+    }
+}
+
+impl<T: Blobstore + Clone> Blobstore for MemWritesBlobstore<T> {
+    fn put(&self, key: String, value: BlobstoreBytes) -> BoxFuture<(), Error> {
+        // Don't write the key if it's already present.
+        self.is_present(key.clone())
+            .and_then({
+                let memblob = self.memblob.clone();
+                move |is_present| {
+                    if is_present {
+                        Either::A(Ok(()).into_future())
+                    } else {
+                        Either::B(memblob.put(key, value))
+                    }
+                }
+            })
+            .boxify()
+    }
+
+    fn get(&self, key: String) -> BoxFuture<Option<BlobstoreBytes>, Error> {
+        self.memblob
+            .get(key.clone())
+            .and_then({
+                let inner = self.inner.clone();
+                move |val| match val {
+                    Some(val) => Either::A(Ok(Some(val)).into_future()),
+                    None => Either::B(inner.get(key)),
+                }
+            })
+            .boxify()
+    }
+
+    fn is_present(&self, key: String) -> BoxFuture<bool, Error> {
+        self.memblob
+            .is_present(key.clone())
+            .and_then({
+                let inner = self.inner.clone();
+                move |is_present| {
+                    if is_present {
+                        Either::A(Ok(true).into_future())
+                    } else {
+                        Either::B(inner.is_present(key))
+                    }
+                }
+            })
+            .boxify()
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use super::*;
+
+    use bytes::Bytes;
+
+    #[test]
+    fn basic_read() {
+        let inner = EagerMemblob::new();
+        let foo_key = "foo".to_string();
+        inner
+            .put(foo_key.clone(), BlobstoreBytes::from_bytes("foobar"))
+            .wait()
+            .expect("initial put should work");
+        let outer = MemWritesBlobstore::new(inner.clone());
+
+        assert!(
+            outer
+                .is_present(foo_key.clone())
+                .wait()
+                .expect("is_present to inner should work")
+        );
+
+        assert_eq!(
+            outer
+                .get(foo_key.clone())
+                .wait()
+                .expect("get to inner should work")
+                .expect("value should be present")
+                .into_bytes(),
+            Bytes::from("foobar"),
+        );
+    }
+
+    #[test]
+    fn redirect_writes() {
+        let inner = EagerMemblob::new();
+        let foo_key = "foo".to_string();
+
+        let outer = MemWritesBlobstore::new(inner.clone());
+        outer
+            .put(foo_key.clone(), BlobstoreBytes::from_bytes("foobar"))
+            .wait()
+            .expect("put should work");
+
+        assert!(
+            !inner
+                .is_present(foo_key.clone())
+                .wait()
+                .expect("is_present on inner should work"),
+            "foo should not be present in inner",
+        );
+
+        assert!(
+            outer
+                .is_present(foo_key.clone())
+                .wait()
+                .expect("is_present on outer should work"),
+            "foo should be present in outer",
+        );
+
+        assert_eq!(
+            outer
+                .get(foo_key.clone())
+                .wait()
+                .expect("get to outer should work")
+                .expect("value should be present")
+                .into_bytes(),
+            Bytes::from("foobar"),
+        );
+    }
+
+    #[test]
+    fn present_in_inner() {
+        let inner = EagerMemblob::new();
+        let foo_key = "foo".to_string();
+        inner
+            .put(foo_key.clone(), BlobstoreBytes::from_bytes("foobar"))
+            .wait()
+            .expect("initial put should work");
+        let outer = MemWritesBlobstore::new(inner.clone());
+        outer
+            .put(foo_key.clone(), BlobstoreBytes::from_bytes("foobar"))
+            .wait()
+            .expect("put should work");
+
+        assert!(
+            outer
+                .is_present(foo_key.clone())
+                .wait()
+                .expect("is_present on outer should work"),
+            "foo should be present in outer",
+        );
+
+        // Change the value in inner.
+        inner
+            .put(foo_key.clone(), BlobstoreBytes::from_bytes("bazquux"))
+            .wait()
+            .expect("second put should work");
+        assert_eq!(
+            outer
+                .get(foo_key.clone())
+                .wait()
+                .expect("get to outer should work")
+                .expect("value should be present")
+                .into_bytes(),
+            Bytes::from("bazquux"),
+        );
+    }
+}

blobstore/src/prefix.rs

@@ -27,6 +27,11 @@ impl<T: Blobstore + Clone> PrefixBlobstore<T> {
         Self { prefix, blobstore }
     }
 
+    #[inline]
+    pub fn into_inner(self) -> T {
+        self.blobstore
+    }
+
     #[inline]
     fn prepend(&self, key: String) -> String {
         [&self.prefix, key.as_str()].concat()