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

Adds `inner_store()` function to all stores that enables the
resolution of inner stores recursively to get an underlying store.
This is mostly for places that can perform optimizations when
specific code paths can be optimized with specific stores.

towards: #409

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(())
+    }
 }