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

cas/grpc_service/ac_server.rs

@@ -65,19 +65,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?,
         ))
@@ -95,20 +96,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")?;

cas/grpc_service/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) = buf_channel::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"));

cas/grpc_service/cas_server.rs

@@ -102,7 +102,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;
@@ -155,7 +157,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;
@@ -212,12 +216,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"))

cas/scheduler/cache_lookup_scheduler.rs

@@ -61,7 +61,7 @@ async fn get_action_from_store(
     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 {

cas/store/compression_store.rs

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

cas/store/dedup_store.rs

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

cas/store/existence_store.rs

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

cas/store/fast_slow_store.rs

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

cas/store/filesystem_store.rs

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

cas/store/grpc_store.rs

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

cas/store/memory_store.rs

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

cas/store/noop_store.rs

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

cas/store/ref_store.rs

@@ -20,7 +20,7 @@ use async_trait::async_trait;
 use error::ResultExt;
 
 use buf_channel::{DropCloserReadHalf, DropCloserWriteHalf};
-use common::DigestInfo;
+use common::{log, DigestInfo};
 use error::{make_err, make_input_err, Code, Error};
 use store::StoreManager;
 use traits::{StoreTrait, UploadSizeInfo};
@@ -123,6 +123,16 @@ impl StoreTrait for RefStore {
             .await
     }
 
+    fn inner_store(self: Arc<Self>, digest: Option<&DigestInfo>) -> Arc<dyn StoreTrait> {
+        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)
     }

cas/store/s3_store.rs

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

cas/store/shard_store.rs

@@ -170,6 +170,14 @@ impl StoreTrait for ShardStore {
         Box::new(self)
     }
 
+    fn inner_store(self: Arc<Self>, digest: Option<&DigestInfo>) -> Arc<dyn StoreTrait> {
+        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}"));

cas/store/size_partitioning_store.rs

@@ -115,6 +115,16 @@ impl StoreTrait for SizePartitioningStore {
             .await
     }
 
+    fn inner_store(self: Arc<Self>, digest: Option<&DigestInfo>) -> Arc<dyn StoreTrait> {
+        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)
     }

cas/store/store_trait.rs

@@ -168,6 +168,13 @@ pub trait StoreTrait: Sync + Send + Unpin {
             .merge(data_res.err_tip(|| "Failed to read stream to completion in get_part_unchunked"))
     }
 
+    /// Gets the underlying store for the given digest. This can be used to find out
+    /// what any underlying store is for a given digest will be and hand it to the caller.
+    /// A caller might want to use this to obtain a reference to the "real" underlying store
+    /// (if applicable) and check if it implements some special traits that allow optimizations.
+    /// Note: If the store performs complex operations on the data, it should return itself.
+    fn inner_store(self: Arc<Self>, _digest: Option<&DigestInfo>) -> Arc<dyn StoreTrait>;
+
     /// Expect the returned Any to be `Arc<Self>`.
     fn as_any(self: Arc<Self>) -> Box<dyn std::any::Any + Send>;
 

cas/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(&config::stores::MemoryStore::default()));
+        store_manager.add_store("mem_store", memory_store.clone());
+
+        let ref_store_inner = Arc::new(RefStore::new(
+            &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(
+            &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(())
+    }
 }

cas/store/verify_store.rs

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

cas/worker/local_worker.rs

@@ -322,6 +322,7 @@ pub async fn new_local_worker(
     historical_store: Arc<dyn Store>,
 ) -> Result<LocalWorker<WorkerApiClientWrapper, RunningActionsManagerImpl>, Error> {
     let fast_slow_store = cas_store
+        .inner_store(None)
         .as_any()
         .downcast_ref::<Arc<FastSlowStore>>()
         .err_tip(|| "Expected store for LocalWorker's store to be a FastSlowStore")?