Fast slow store directions

deployment-examples/docker-compose/worker.json5

@@ -41,6 +41,7 @@
             },
           },
         },
+        fast_direction: "get",
         slow: {
           ref_store: {
             name: "GRPC_LOCAL_STORE",

kubernetes/components/worker/worker.json5

@@ -40,6 +40,7 @@
             },
           },
         },
+        fast_direction: "get",
         slow: {
           ref_store: {
             name: "GRPC_LOCAL_STORE",

nativelink-config/src/stores.rs

@@ -639,16 +639,46 @@ pub struct OntapS3ExistenceCacheSpec {
     pub backend: Box<ExperimentalOntapS3Spec>,
 }
 
+#[derive(Serialize, Deserialize, Default, Debug, Clone, Copy, PartialEq)]
+#[serde(rename_all = "snake_case")]
+pub enum StoreDirection {
+    /// The store operates normally and all get and put operations are
+    /// handled by it.
+    #[default]
+    Both,
+    /// Update operations will cause persistence to this store, but Get
+    /// operations will be ignored.
+    /// This only makes sense on the fast store as the slow store will
+    /// never get written to on Get anyway.
+    Update,
+    /// Get operations will cause persistence to this store, but Update
+    /// operations will be ignored.
+    Get,
+    /// Operate as a read only store, only really makes sense if there's
+    /// another way to write to it.
+    ReadOnly,
+}
+
 #[derive(Serialize, Deserialize, Debug, Clone)]
 #[serde(deny_unknown_fields)]
 pub struct FastSlowSpec {
     /// Fast store that will be attempted to be contacted before reaching
     /// out to the `slow` store.
     pub fast: StoreSpec,
 
+    /// How to handle the fast store.  This can be useful to set to Get for
+    /// worker nodes such that results are persisted to the slow store only.
+    #[serde(default)]
+    pub fast_direction: StoreDirection,
+
     /// If the object does not exist in the `fast` store it will try to
     /// get it from this store.
     pub slow: StoreSpec,
+
+    /// How to handle the slow store.  This can be useful if creating a diode
+    /// and you wish to have an upstream read only store.
+    #[serde(default)]
+    pub slow_direction: StoreDirection,
 }
 
 #[derive(Serialize, Deserialize, Debug, Default, Clone, Copy)]

nativelink-store/src/fast_slow_store.rs

@@ -23,7 +23,7 @@ use std::sync::{Arc, Weak};
 
 use async_trait::async_trait;
 use futures::{FutureExt, join};
-use nativelink_config::stores::FastSlowSpec;
+use nativelink_config::stores::{FastSlowSpec, StoreDirection};
 use nativelink_error::{Code, Error, ResultExt, make_err};
 use nativelink_metric::MetricsComponent;
 use nativelink_util::buf_channel::{
@@ -51,8 +51,10 @@ type Loader = Arc<OnceCell<()>>;
 pub struct FastSlowStore {
     #[metric(group = "fast_store")]
     fast_store: Store,
+    fast_direction: StoreDirection,
     #[metric(group = "slow_store")]
     slow_store: Store,
+    slow_direction: StoreDirection,
     weak_self: Weak<Self>,
     #[metric]
     metrics: FastSlowStoreMetrics,
@@ -113,10 +115,12 @@ impl Drop for LoaderGuard<'_> {
 }
 
 impl FastSlowStore {
-    pub fn new(_spec: &FastSlowSpec, fast_store: Store, slow_store: Store) -> Arc<Self> {
+    pub fn new(spec: &FastSlowSpec, fast_store: Store, slow_store: Store) -> Arc<Self> {
         Arc::new_cyclic(|weak_self| Self {
             fast_store,
+            fast_direction: spec.fast_direction,
             slow_store,
+            slow_direction: spec.slow_direction,
             weak_self: weak_self.clone(),
             metrics: FastSlowStoreMetrics::default(),
             populating_digests: Mutex::new(HashMap::new()),
@@ -182,6 +186,29 @@ impl FastSlowStore {
             .slow_store_hit_count
             .fetch_add(1, Ordering::Acquire);
 
+        // If the fast store is noop or read only or update only then bypass it.
+        if self
+            .fast_store
+            .inner_store(Some(key.borrow()))
+            .optimized_for(StoreOptimizations::NoopUpdates)
+            || self.fast_direction == StoreDirection::ReadOnly
+            || self.fast_direction == StoreDirection::Update
+        {
+            let Some(writer) = maybe_writer else {
+                return Err(make_err!(
+                    Code::Internal,
+                    "Attempt to populate fast store that is read only or noop"
+                ));
+            };
+            self.slow_store
+                .get_part(key, writer.borrow_mut(), offset, length)
+                .await?;
+            self.metrics
+                .slow_store_downloaded_bytes
+                .fetch_add(writer.get_bytes_written(), Ordering::Acquire);
+            return Ok(());
+        }
+
         let send_range = offset..length.map_or(u64::MAX, |length| length + offset);
         let mut bytes_received: u64 = 0;
 
@@ -327,12 +354,31 @@ impl StoreDriver for FastSlowStore {
     ) -> Result<(), Error> {
         // If either one of our stores is a noop store, bypass the multiplexing
         // and just use the store that is not a noop store.
-        let slow_store = self.slow_store.inner_store(Some(key.borrow()));
-        if slow_store.optimized_for(StoreOptimizations::NoopUpdates) {
+        let ignore_slow = self
+            .slow_store
+            .inner_store(Some(key.borrow()))
+            .optimized_for(StoreOptimizations::NoopUpdates)
+            || self.slow_direction == StoreDirection::ReadOnly
+            || self.slow_direction == StoreDirection::Get;
+        let ignore_fast = self
+            .fast_store
+            .inner_store(Some(key.borrow()))
+            .optimized_for(StoreOptimizations::NoopUpdates)
+            || self.fast_direction == StoreDirection::ReadOnly
+            || self.fast_direction == StoreDirection::Get;
+        if ignore_slow && ignore_fast {
+            // We need to drain the reader to avoid the writer complaining that we dropped
+            // the connection prematurely.
+            reader
+                .drain()
+                .await
+                .err_tip(|| "In FastFlowStore::update")?;
+            return Ok(());
+        }
+        if ignore_slow {
             return self.fast_store.update(key, reader, size_info).await;
         }
-        let fast_store = self.fast_store.inner_store(Some(key.borrow()));
-        if fast_store.optimized_for(StoreOptimizations::NoopUpdates) {
+        if ignore_fast {
             return self.slow_store.update(key, reader, size_info).await;
         }
 
@@ -406,7 +452,10 @@ impl StoreDriver for FastSlowStore {
         {
             if !self
                 .slow_store
+                .inner_store(Some(key.borrow()))
                 .optimized_for(StoreOptimizations::NoopUpdates)
+                && self.slow_direction != StoreDirection::ReadOnly
+                && self.slow_direction != StoreDirection::Get
             {
                 slow_update_store_with_file(
                     self.slow_store.as_store_driver_pin(),
@@ -417,6 +466,11 @@ impl StoreDriver for FastSlowStore {
                 .await
                 .err_tip(|| "In FastSlowStore::update_with_whole_file slow_store")?;
             }
+            if self.fast_direction == StoreDirection::ReadOnly
+                || self.fast_direction == StoreDirection::Get
+            {
+                return Ok(Some(file));
+            }
             return self
                 .fast_store
                 .update_with_whole_file(key, path, file, upload_size)
@@ -427,10 +481,13 @@ impl StoreDriver for FastSlowStore {
             .slow_store
             .optimized_for(StoreOptimizations::FileUpdates)
         {
-            if !self
+            let ignore_fast = self
                 .fast_store
+                .inner_store(Some(key.borrow()))
                 .optimized_for(StoreOptimizations::NoopUpdates)
-            {
+                || self.fast_direction == StoreDirection::ReadOnly
+                || self.fast_direction == StoreDirection::Get;
+            if !ignore_fast {
                 slow_update_store_with_file(
                     self.fast_store.as_store_driver_pin(),
                     key.borrow(),
@@ -440,6 +497,11 @@ impl StoreDriver for FastSlowStore {
                 .await
                 .err_tip(|| "In FastSlowStore::update_with_whole_file fast_store")?;
             }
+            let ignore_slow = self.slow_direction == StoreDirection::ReadOnly
+                || self.slow_direction == StoreDirection::Get;
+            if ignore_slow {
+                return Ok(Some(file));
+            }
             return self
                 .slow_store
                 .update_with_whole_file(key, path, file, upload_size)

nativelink-store/tests/fast_slow_store_test.rs

@@ -18,7 +18,7 @@ use std::sync::{Arc, Mutex};
 
 use async_trait::async_trait;
 use bytes::Bytes;
-use nativelink_config::stores::{FastSlowSpec, MemorySpec, NoopSpec, StoreSpec};
+use nativelink_config::stores::{FastSlowSpec, MemorySpec, NoopSpec, StoreDirection, StoreSpec};
 use nativelink_error::{Code, Error, ResultExt, make_err};
 use nativelink_macro::nativelink_test;
 use nativelink_metric::MetricsComponent;
@@ -35,20 +35,29 @@ use rand::{Rng, SeedableRng};
 
 const MEGABYTE_SZ: usize = 1024 * 1024;
 
-fn make_stores() -> (Store, Store, Store) {
+fn make_stores_direction(
+    fast_direction: StoreDirection,
+    slow_direction: StoreDirection,
+) -> (Store, Store, Store) {
     let fast_store = Store::new(MemoryStore::new(&MemorySpec::default()));
     let slow_store = Store::new(MemoryStore::new(&MemorySpec::default()));
     let fast_slow_store = Store::new(FastSlowStore::new(
         &FastSlowSpec {
             fast: StoreSpec::Memory(MemorySpec::default()),
             slow: StoreSpec::Memory(MemorySpec::default()),
+            fast_direction,
+            slow_direction,
         },
         fast_store.clone(),
         slow_store.clone(),
     ));
     (fast_slow_store, fast_store, slow_store)
 }
 
+fn make_stores() -> (Store, Store, Store) {
+    make_stores_direction(StoreDirection::default(), StoreDirection::default())
+}
+
 fn make_random_data(sz: usize) -> Vec<u8> {
     let mut value = vec![0u8; sz];
     let mut rng = SmallRng::seed_from_u64(1);
@@ -339,6 +348,8 @@ async fn drop_on_eof_completes_store_futures() -> Result<(), Error> {
         &FastSlowSpec {
             fast: StoreSpec::Memory(MemorySpec::default()),
             slow: StoreSpec::Memory(MemorySpec::default()),
+            fast_direction: StoreDirection::default(),
+            slow_direction: StoreDirection::default(),
         },
         fast_store,
         slow_store,
@@ -380,6 +391,8 @@ async fn ignore_value_in_fast_store() -> Result<(), Error> {
         &FastSlowSpec {
             fast: StoreSpec::Memory(MemorySpec::default()),
             slow: StoreSpec::Memory(MemorySpec::default()),
+            fast_direction: StoreDirection::default(),
+            slow_direction: StoreDirection::default(),
         },
         fast_store.clone(),
         slow_store,
@@ -403,6 +416,8 @@ async fn has_checks_fast_store_when_noop() -> Result<(), Error> {
     let fast_slow_store_config = FastSlowSpec {
         fast: StoreSpec::Memory(MemorySpec::default()),
         slow: StoreSpec::Noop(NoopSpec::default()),
+        fast_direction: StoreDirection::default(),
+        slow_direction: StoreDirection::default(),
     };
     let fast_slow_store = Arc::new(FastSlowStore::new(
         &fast_slow_store_config,
@@ -437,3 +452,95 @@ async fn has_checks_fast_store_when_noop() -> Result<(), Error> {
     );
     Ok(())
 }
+
+#[nativelink_test]
+async fn fast_get_only_not_updated() -> Result<(), Error> {
+    let (fast_slow_store, fast_store, slow_store) =
+        make_stores_direction(StoreDirection::Get, StoreDirection::Both);
+    let digest = DigestInfo::try_new(VALID_HASH, 100).unwrap();
+    fast_slow_store
+        .update_oneshot(digest, make_random_data(100).into())
+        .await?;
+    assert!(
+        fast_store.has(digest).await?.is_none(),
+        "Expected data to not be in the fast store"
+    );
+    assert!(
+        slow_store.has(digest).await?.is_some(),
+        "Expected data in the slow store"
+    );
+    Ok(())
+}
+
+#[nativelink_test]
+async fn fast_readonly_only_not_updated() -> Result<(), Error> {
+    let (fast_slow_store, fast_store, slow_store) =
+        make_stores_direction(StoreDirection::ReadOnly, StoreDirection::Both);
+    let digest = DigestInfo::try_new(VALID_HASH, 100).unwrap();
+    fast_slow_store
+        .update_oneshot(digest, make_random_data(100).into())
+        .await?;
+    assert!(
+        fast_store.has(digest).await?.is_none(),
+        "Expected data to not be in the fast store"
+    );
+    assert!(
+        slow_store.has(digest).await?.is_some(),
+        "Expected data in the slow store"
+    );
+    Ok(())
+}
+
+#[nativelink_test]
+async fn slow_readonly_only_not_updated() -> Result<(), Error> {
+    let (fast_slow_store, fast_store, slow_store) =
+        make_stores_direction(StoreDirection::Both, StoreDirection::ReadOnly);
+    let digest = DigestInfo::try_new(VALID_HASH, 100).unwrap();
+    fast_slow_store
+        .update_oneshot(digest, make_random_data(100).into())
+        .await?;
+    assert!(
+        fast_store.has(digest).await?.is_some(),
+        "Expected data to be in the fast store"
+    );
+    assert!(
+        slow_store.has(digest).await?.is_none(),
+        "Expected data to not be in the slow store"
+    );
+    Ok(())
+}
+
+#[nativelink_test]
+async fn slow_get_only_not_updated() -> Result<(), Error> {
+    let (fast_slow_store, fast_store, slow_store) =
+        make_stores_direction(StoreDirection::Both, StoreDirection::Get);
+    let digest = DigestInfo::try_new(VALID_HASH, 100).unwrap();
+    fast_slow_store
+        .update_oneshot(digest, make_random_data(100).into())
+        .await?;
+    assert!(
+        fast_store.has(digest).await?.is_some(),
+        "Expected data to be in the fast store"
+    );
+    assert!(
+        slow_store.has(digest).await?.is_none(),
+        "Expected data to not be in the slow store"
+    );
+    Ok(())
+}
+
+#[nativelink_test]
+async fn fast_put_only_not_updated() -> Result<(), Error> {
+    let (fast_slow_store, fast_store, slow_store) =
+        make_stores_direction(StoreDirection::Update, StoreDirection::Both);
+    let digest = DigestInfo::try_new(VALID_HASH, 100).unwrap();
+    slow_store
+        .update_oneshot(digest, make_random_data(100).into())
+        .await?;
+    fast_slow_store.get_part_unchunked(digest, 0, None).await?;
+    assert!(
+        fast_store.has(digest).await?.is_none(),
+        "Expected data to not be in the fast store"
+    );
+    Ok(())
+}