Add function to Store API to get the inner store when possible

native-link-scheduler/src/cache_lookup_scheduler.rs

@@ -57,11 +57,11 @@ pub struct CacheLookupScheduler {
 
 async fn get_action_from_store(
     ac_store: &Arc<dyn Store>,
-    action_digest: &DigestInfo,
+    action_digest: DigestInfo,
     instance_name: String,
 ) -> Option<ProtoActionResult> {
     // If we are a GrpcStore we shortcut here, as this is a special store.
-    let any_store = ac_store.clone().as_any();
+    let any_store = ac_store.clone().inner_store(Some(action_digest)).as_any();
     let maybe_grpc_store = any_store.downcast_ref::<Arc<GrpcStore>>();
     if let Some(grpc_store) = maybe_grpc_store {
         let action_result_request = GetActionResultRequest {
@@ -78,7 +78,7 @@ async fn get_action_from_store(
             .map(|response| response.into_inner())
             .ok()
     } else {
-        get_and_decode_digest::<ProtoActionResult>(Pin::new(ac_store.as_ref()), action_digest)
+        get_and_decode_digest::<ProtoActionResult>(Pin::new(ac_store.as_ref()), &action_digest)
             .await
             .ok()
     }
@@ -185,7 +185,7 @@ impl ActionScheduler for CacheLookupScheduler {
             // Perform cache check.
             let action_digest = current_state.action_digest();
             let instance_name = action_info.instance_name().clone();
-            if let Some(action_result) = get_action_from_store(&ac_store, action_digest, instance_name).await {
+            if let Some(action_result) = get_action_from_store(&ac_store, *action_digest, instance_name).await {
                 if validate_outputs_exist(&cas_store, &action_result).await {
                     // Found in the cache, return the result immediately.
                     Arc::make_mut(&mut current_state).stage = ActionStage::CompletedFromCache(action_result);

native-link-service/src/ac_server.rs

@@ -63,19 +63,20 @@ impl AcServer {
             .get(instance_name)
             .err_tip(|| format!("'instance_name' not configured for '{}'", instance_name))?;
 
-        // If we are a GrpcStore we shortcut here, as this is a special store.
-        let any_store = store.clone().as_any();
-        let maybe_grpc_store = any_store.downcast_ref::<Arc<GrpcStore>>();
-        if let Some(grpc_store) = maybe_grpc_store {
-            return grpc_store.get_action_result(Request::new(get_action_request)).await;
-        }
-
         // TODO(blaise.bruer) We should write a test for these errors.
         let digest: DigestInfo = get_action_request
             .action_digest
+            .clone()
             .err_tip(|| "Action digest was not set in message")?
             .try_into()?;
 
+        // If we are a GrpcStore we shortcut here, as this is a special store.
+        let any_store = store.clone().inner_store(Some(digest)).as_any();
+        let maybe_grpc_store = any_store.downcast_ref::<Arc<GrpcStore>>();
+        if let Some(grpc_store) = maybe_grpc_store {
+            return grpc_store.get_action_result(Request::new(get_action_request)).await;
+        }
+
         Ok(Response::new(
             get_and_decode_digest::<ActionResult>(Pin::new(store.as_ref()), &digest).await?,
         ))
@@ -93,20 +94,21 @@ impl AcServer {
             .get(instance_name)
             .err_tip(|| format!("'instance_name' not configured for '{}'", instance_name))?;
 
+        let digest: DigestInfo = update_action_request
+            .action_digest
+            .clone()
+            .err_tip(|| "Action digest was not set in message")?
+            .try_into()?;
+
         // If we are a GrpcStore we shortcut here, as this is a special store.
-        let any_store = store.clone().as_any();
+        let any_store = store.clone().inner_store(Some(digest)).as_any();
         let maybe_grpc_store = any_store.downcast_ref::<Arc<GrpcStore>>();
         if let Some(grpc_store) = maybe_grpc_store {
             return grpc_store
                 .update_action_result(Request::new(update_action_request))
                 .await;
         }
 
-        let digest: DigestInfo = update_action_request
-            .action_digest
-            .err_tip(|| "Action digest was not set in message")?
-            .try_into()?;
-
         let action_result = update_action_request
             .action_result
             .err_tip(|| "Action result was not set in message")?;

native-link-service/src/bytestream_server.rs

@@ -237,16 +237,16 @@ impl ByteStreamServer {
             .err_tip(|| format!("'instance_name' not configured for '{}'", instance_name))?
             .clone();
 
+        let digest = DigestInfo::try_new(resource_info.hash, resource_info.expected_size)?;
+
         // If we are a GrpcStore we shortcut here, as this is a special store.
-        let any_store = store.clone().as_any();
+        let any_store = store.clone().inner_store(Some(digest)).as_any();
         let maybe_grpc_store = any_store.downcast_ref::<Arc<GrpcStore>>();
         if let Some(grpc_store) = maybe_grpc_store {
             let stream = grpc_store.read(Request::new(read_request)).await?.into_inner();
             return Ok(Response::new(Box::pin(stream)));
         }
 
-        let digest = DigestInfo::try_new(resource_info.hash, resource_info.expected_size)?;
-
         let (tx, rx) = make_buf_channel_pair();
 
         struct ReaderState {
@@ -345,8 +345,11 @@ impl ByteStreamServer {
             .err_tip(|| format!("'instance_name' not configured for '{}'", instance_name))?
             .clone();
 
+        let digest = DigestInfo::try_new(&stream.hash, stream.expected_size)
+            .err_tip(|| "Invalid digest input in ByteStream::write")?;
+
         // If we are a GrpcStore we shortcut here, as this is a special store.
-        let any_store = store.clone().as_any();
+        let any_store = store.clone().inner_store(Some(digest)).as_any();
         let maybe_grpc_store = any_store.downcast_ref::<Arc<GrpcStore>>();
         if let Some(grpc_store) = maybe_grpc_store {
             return grpc_store.write(stream).await;
@@ -356,8 +359,6 @@ impl ByteStreamServer {
             .uuid
             .take()
             .ok_or_else(|| make_input_err!("UUID must be set if writing data"))?;
-        let digest = DigestInfo::try_new(&stream.hash, stream.expected_size)
-            .err_tip(|| "Invalid digest input in ByteStream::write")?;
         let mut active_stream_guard = self.create_or_join_upload_stream(uuid, store, digest)?;
         let expected_size = stream.expected_size as u64;
 
@@ -444,8 +445,10 @@ impl ByteStreamServer {
             .err_tip(|| format!("'instance_name' not configured for '{}'", &resource_info.instance_name))?
             .clone();
 
+        let digest = DigestInfo::try_new(resource_info.hash, resource_info.expected_size)?;
+
         // If we are a GrpcStore we shortcut here, as this is a special store.
-        let any_store = store_clone.clone().as_any();
+        let any_store = store_clone.clone().inner_store(Some(digest)).as_any();
         let maybe_grpc_store = any_store.downcast_ref::<Arc<GrpcStore>>();
         if let Some(grpc_store) = maybe_grpc_store {
             return grpc_store.query_write_status(Request::new(query_request.clone())).await;
@@ -468,8 +471,6 @@ impl ByteStreamServer {
             }
         }
 
-        let digest = DigestInfo::try_new(resource_info.hash, resource_info.expected_size)?;
-
         let has_fut = Pin::new(store_clone.as_ref()).has(digest);
         let Some(item_size) = has_fut.await.err_tip(|| "Failed to call .has() on store")? else {
             return Err(make_err!(Code::NotFound, "{}", "not found"));

native-link-service/src/cas_server.rs

@@ -104,7 +104,9 @@ impl CasServer {
             .clone();
 
         // If we are a GrpcStore we shortcut here, as this is a special store.
-        let any_store = store.clone().as_any();
+        // Note: We don't know the digests here, so we try perform a very shallow
+        // check to see if it's a grpc store.
+        let any_store = store.clone().inner_store(None).as_any();
         let maybe_grpc_store = any_store.downcast_ref::<Arc<GrpcStore>>();
         if let Some(grpc_store) = maybe_grpc_store {
             return grpc_store.batch_update_blobs(Request::new(inner_request)).await;
@@ -157,7 +159,9 @@ impl CasServer {
             .clone();
 
         // If we are a GrpcStore we shortcut here, as this is a special store.
-        let any_store = store.clone().as_any();
+        // Note: We don't know the digests here, so we try perform a very shallow
+        // check to see if it's a grpc store.
+        let any_store = store.clone().inner_store(None).as_any();
         let maybe_grpc_store = any_store.downcast_ref::<Arc<GrpcStore>>();
         if let Some(grpc_store) = maybe_grpc_store {
             return grpc_store.batch_read_blobs(Request::new(inner_request)).await;
@@ -214,12 +218,13 @@ impl CasServer {
             .clone();
 
         // If we are a GrpcStore we shortcut here, as this is a special store.
-        let any_store = store.clone().as_any();
+        // Note: We don't know the digests here, so we try perform a very shallow
+        // check to see if it's a grpc store.
+        let any_store = store.clone().inner_store(None).as_any();
         let maybe_grpc_store = any_store.downcast_ref::<Arc<GrpcStore>>();
 
         if let Some(grpc_store) = maybe_grpc_store {
             let stream = grpc_store.get_tree(Request::new(inner_request)).await?.into_inner();
-            // let stream = grpc_store.read(Request::new(read_request)).await?.into_inner();
             return Ok(Response::new(Box::pin(stream)));
         }
         Err(make_err!(Code::Unimplemented, "get_tree is not implemented"))

native-link-store/src/compression_store.rs

@@ -566,6 +566,10 @@ impl Store for CompressionStore {
         Ok(())
     }
 
+    fn inner_store(self: Arc<Self>, _digest: Option<DigestInfo>) -> Arc<dyn Store> {
+        self
+    }
+
     fn as_any(self: Arc<Self>) -> Box<dyn std::any::Any + Send> {
         Box::new(self)
     }

native-link-store/src/dedup_store.rs

@@ -340,6 +340,10 @@ impl Store for DedupStore {
         Ok(())
     }
 
+    fn inner_store(self: Arc<Self>, _digest: Option<DigestInfo>) -> Arc<dyn Store> {
+        self
+    }
+
     fn as_any(self: Arc<Self>) -> Box<dyn std::any::Any + Send> {
         Box::new(self)
     }

native-link-store/src/existence_cache_store.rs

@@ -178,6 +178,10 @@ impl Store for ExistenceCacheStore {
         result
     }
 
+    fn inner_store(self: Arc<Self>, _digest: Option<DigestInfo>) -> Arc<dyn Store> {
+        self
+    }
+
     fn as_any(self: Arc<Self>) -> Box<dyn std::any::Any + Send> {
         Box::new(self)
     }

native-link-store/src/fast_slow_store.rs

@@ -260,6 +260,10 @@ impl Store for FastSlowStore {
         Ok(())
     }
 
+    fn inner_store(self: Arc<Self>, _digest: Option<DigestInfo>) -> Arc<dyn Store> {
+        self
+    }
+
     fn as_any(self: Arc<Self>) -> Box<dyn std::any::Any + Send> {
         Box::new(self)
     }

native-link-store/src/filesystem_store.rs

@@ -678,6 +678,10 @@ impl<Fe: FileEntry> Store for FilesystemStore<Fe> {
         Ok(())
     }
 
+    fn inner_store(self: Arc<Self>, _digest: Option<DigestInfo>) -> Arc<dyn Store> {
+        self
+    }
+
     fn as_any(self: Arc<Self>) -> Box<dyn std::any::Any + Send> {
         Box::new(self)
     }

native-link-store/src/grpc_store.rs

@@ -672,6 +672,10 @@ impl Store for GrpcStore {
         Ok(())
     }
 
+    fn inner_store(self: Arc<Self>, _digest: Option<DigestInfo>) -> Arc<dyn Store> {
+        self
+    }
+
     fn as_any(self: Arc<Self>) -> Box<dyn std::any::Any + Send> {
         Box::new(self)
     }

native-link-store/src/memory_store.rs

@@ -130,6 +130,10 @@ impl Store for MemoryStore {
         Ok(())
     }
 
+    fn inner_store(self: Arc<Self>, _digest: Option<DigestInfo>) -> Arc<dyn Store> {
+        self
+    }
+
     fn as_any(self: Arc<Self>) -> Box<dyn std::any::Any + Send> {
         Box::new(self)
     }

native-link-store/src/noop_store.rs

@@ -63,6 +63,10 @@ impl Store for NoopStore {
         Err(make_err!(Code::NotFound, "Not found in noop store"))
     }
 
+    fn inner_store(self: Arc<Self>, _digest: Option<DigestInfo>) -> Arc<dyn Store> {
+        self
+    }
+
     fn as_any(self: Arc<Self>) -> Box<dyn std::any::Any + Send> {
         Box::new(self)
     }

native-link-store/src/ref_store.rs

@@ -19,7 +19,7 @@ use std::sync::{Arc, Mutex, Weak};
 use async_trait::async_trait;
 use error::{make_err, make_input_err, Code, Error, ResultExt};
 use native_link_util::buf_channel::{DropCloserReadHalf, DropCloserWriteHalf};
-use native_link_util::common::DigestInfo;
+use native_link_util::common::{log, DigestInfo};
 use native_link_util::store_trait::{Store, UploadSizeInfo};
 
 use crate::store_manager::StoreManager;
@@ -122,6 +122,16 @@ impl Store for RefStore {
             .await
     }
 
+    fn inner_store(self: Arc<Self>, digest: Option<DigestInfo>) -> Arc<dyn Store> {
+        match self.get_store() {
+            Ok(store) => store.clone().inner_store(digest),
+            Err(e) => {
+                log::error!("Failed to get store for digest: {e:?}");
+                self
+            }
+        }
+    }
+
     fn as_any(self: Arc<Self>) -> Box<dyn std::any::Any + Send> {
         Box::new(self)
     }

native-link-store/src/s3_store.rs

@@ -530,6 +530,10 @@ impl Store for S3Store {
             .await
     }
 
+    fn inner_store(self: Arc<Self>, _digest: Option<DigestInfo>) -> Arc<dyn Store> {
+        self
+    }
+
     fn as_any(self: Arc<Self>) -> Box<dyn std::any::Any + Send> {
         Box::new(self)
     }

native-link-store/src/shard_store.rs

@@ -169,6 +169,14 @@ impl Store for ShardStore {
         Box::new(self)
     }
 
+    fn inner_store(self: Arc<Self>, digest: Option<DigestInfo>) -> Arc<dyn Store> {
+        let Some(digest) = digest else {
+            return self;
+        };
+        let index = self.get_store_index(&digest);
+        self.weights_and_stores[index].1.clone().inner_store(Some(digest))
+    }
+
     fn register_metrics(self: Arc<Self>, registry: &mut Registry) {
         for (i, (_, store)) in self.weights_and_stores.iter().enumerate() {
             let store_registry = registry.sub_registry_with_prefix(format!("store_{i}"));

native-link-store/src/size_partitioning_store.rs

@@ -114,6 +114,16 @@ impl Store for SizePartitioningStore {
             .await
     }
 
+    fn inner_store(self: Arc<Self>, digest: Option<DigestInfo>) -> Arc<dyn Store> {
+        let Some(digest) = digest else {
+            return self;
+        };
+        if digest.size_bytes < self.size {
+            return self.lower_store.clone().inner_store(Some(digest));
+        }
+        self.upper_store.clone().inner_store(Some(digest))
+    }
+
     fn as_any(self: Arc<Self>) -> Box<dyn std::any::Any + Send> {
         Box::new(self)
     }

native-link-store/src/verify_store.rs

@@ -160,6 +160,10 @@ impl Store for VerifyStore {
         self.pin_inner().get_part_ref(digest, writer, offset, length).await
     }
 
+    fn inner_store(self: Arc<Self>, _digest: Option<DigestInfo>) -> Arc<dyn Store> {
+        self
+    }
+
     fn as_any(self: Arc<Self>) -> Box<dyn std::any::Any + Send> {
         Box::new(self)
     }

native-link-store/tests/ref_store_test.rs

@@ -124,4 +124,36 @@ mod ref_store_tests {
         }
         Ok(())
     }
+
+    #[tokio::test]
+    async fn inner_store_test() -> Result<(), Error> {
+        let store_manager = Arc::new(StoreManager::new());
+
+        let memory_store = Arc::new(MemoryStore::new(&native_link_config::stores::MemoryStore::default()));
+        store_manager.add_store("mem_store", memory_store.clone());
+
+        let ref_store_inner = Arc::new(RefStore::new(
+            &native_link_config::stores::RefStore {
+                name: "mem_store".to_string(),
+            },
+            Arc::downgrade(&store_manager),
+        ));
+        store_manager.add_store("ref_store_inner", ref_store_inner.clone());
+
+        let ref_store_outer = Arc::new(RefStore::new(
+            &native_link_config::stores::RefStore {
+                name: "ref_store_inner".to_string(),
+            },
+            Arc::downgrade(&store_manager),
+        ));
+        store_manager.add_store("ref_store_outer", ref_store_outer.clone());
+
+        // Ensure the result of inner_store() points to exact same memory store.
+        assert_eq!(
+            Arc::as_ptr(&ref_store_outer.inner_store(None)) as *const (),
+            Arc::as_ptr(&memory_store) as *const (),
+            "Expected inner store to be memory store"
+        );
+        Ok(())
+    }
 }