Skip to content

Commit

Permalink
feat: Add support for multiple digest functions with backward compat…
Browse files Browse the repository at this point in the history
…ibility

- Updated `ResourceInfo` to default `digest_function` to `sha256` if not provided.
- Ensured backward compatibility for existing functionality.
- Updated tests to reflect default `sha256` behavior when `digest_function` is not set
  • Loading branch information
asr2003 authored Dec 17, 2024
1 parent f672af7 commit 6b47cd7
Show file tree
Hide file tree
Showing 2 changed files with 34 additions and 2 deletions.
32 changes: 30 additions & 2 deletions nativelink-store/src/grpc_store.rs
Original file line number Diff line number Diff line change
Expand Up @@ -45,7 +45,7 @@ use nativelink_util::origin_context::ActiveOriginContext;
use nativelink_util::proto_stream_utils::{
FirstStream, WriteRequestStreamWrapper, WriteState, WriteStateWrapper,
};
use nativelink_util::resource_info::ResourceInfo;
use nativelink_util::resource_info::{is_supported_digest_function, ResourceInfo};
use nativelink_util::retry::{Retrier, RetryResult};
use nativelink_util::store_trait::{StoreDriver, StoreKey, UploadSizeInfo};
use nativelink_util::{default_health_status_indicator, tls_utils};
Expand Down Expand Up @@ -276,12 +276,28 @@ impl GrpcStore {
&self,
grpc_request: impl IntoRequest<ReadRequest>,
) -> Result<impl Stream<Item = Result<ReadResponse, Status>>, Error> {
const IS_UPLOAD_FALSE: bool = false;

let request = self.get_read_request(grpc_request.into_request().into_inner())?;
let resource_name = &request.resource_name;
let resource_info = ResourceInfo::new(resource_name, IS_UPLOAD_FALSE)
.err_tip(|| "Failed to parse resource_name in GrpcStore::read")?;

let digest_function = resource_info.digest_function.as_deref().unwrap_or("sha256");

if !is_supported_digest_function(digest_function) {
return Err(make_input_err!(
"Unsupported digest_function: {} in resource_name '{}'",
digest_function,
resource_name
));
}

error_if!(
matches!(self.store_type, nativelink_config::stores::StoreType::ac),
"CAS operation on AC store"
);

let request = self.get_read_request(grpc_request.into_request().into_inner())?;
self.perform_request(request, |request| async move {
self.read_internal(request).await
})
Expand Down Expand Up @@ -514,6 +530,18 @@ impl StoreDriver for GrpcStore {
keys: &[StoreKey<'_>],
results: &mut [Option<u64>],
) -> Result<(), Error> {
let digest_function = ActiveOriginContext::get_value(&ACTIVE_HASHER_FUNC)
.err_tip(|| "In GrpcStore::has_with_results")?
.map_or_else(default_digest_hasher_func, |v| *v)
.to_string();

if !is_supported_digest_function(&digest_function) {
return Err(make_input_err!(
"Unsupported digest_function: {}",
digest_function
));
}

if matches!(self.store_type, nativelink_config::stores::StoreType::ac) {
keys.iter()
.zip(results.iter_mut())
Expand Down
4 changes: 4 additions & 0 deletions nativelink-util/src/resource_info.rs
Original file line number Diff line number Diff line change
Expand Up @@ -206,6 +206,10 @@ enum State {
OptionalMetadata,
}

pub fn is_supported_digest_function(digest_function: &str) -> bool {
DIGEST_FUNCTIONS.contains(&digest_function.to_lowercase().as_str())
}

// Iterate backwards looking for "(compressed-)blobs", once found, move forward
// populating the output struct. This recursive function utilises the stack to
// temporarily hold the reference to the previous item reducing the need for
Expand Down

0 comments on commit 6b47cd7

Please sign in to comment.