diff --git a/nativelink-service/src/cas_server.rs b/nativelink-service/src/cas_server.rs index 372e1e2d2..65727f4cd 100644 --- a/nativelink-service/src/cas_server.rs +++ b/nativelink-service/src/cas_server.rs @@ -236,18 +236,23 @@ impl CasServer { let mut deque: VecDeque = VecDeque::new(); let mut directories: Vec = Vec::new(); // `page_token` will return the `{hash_str}-{size_bytes}` of the current request's first directory digest. - let mut page_token_parts = request.page_token.split('-'); - let page_token_digest = DigestInfo::try_new( - page_token_parts - .next() - .err_tip(|| "Failed to parse `hash_str` in `page_token`")?, - page_token_parts - .next() - .err_tip(|| "Failed to parse `size_bytes` in `page_token`")? - .parse::() - .err_tip(|| "Failed to parse `size_bytes` as i64")?, - ) - .err_tip(|| "Failed to parse `page_token` as `Digest` in `GetTreeRequest`")?; + let page_token_digest = match request.page_token.len() { + 0 => root_digest, + _ => { + let mut page_token_parts = request.page_token.split('-'); + DigestInfo::try_new( + page_token_parts + .next() + .err_tip(|| "Failed to parse `hash_str` in `page_token`")?, + page_token_parts + .next() + .err_tip(|| "Failed to parse `size_bytes` in `page_token`")? + .parse::() + .err_tip(|| "Failed to parse `size_bytes` as i64")?, + ) + .err_tip(|| "Failed to parse `page_token` as `Digest` in `GetTreeRequest`")? + } + }; let page_size = request.page_size; // If `page_size` is 0, paging is not necessary. let mut page_token_matched = page_size == 0; diff --git a/nativelink-service/tests/cas_server_test.rs b/nativelink-service/tests/cas_server_test.rs index a0726cfed..42ffab2b6 100644 --- a/nativelink-service/tests/cas_server_test.rs +++ b/nativelink-service/tests/cas_server_test.rs @@ -383,35 +383,71 @@ async fn get_tree_read_directories_without_paging() -> Result<(), Box>() - .await, - vec![GetTreeResponse { - directories: vec![ - root_directory.clone(), - sub_directories[0].clone(), - sub_directories[1].clone(), - sub_directories[2].clone(), - sub_directories[3].clone(), - sub_directories[4].clone() - ], - next_page_token: String::new() - }] - ); + + // First verify that using an empty page token is treated as if the client had sent the root + // digest. + { + let raw_response = cas_server + .get_tree(Request::new(GetTreeRequest { + instance_name: INSTANCE_NAME.to_string(), + page_size: 0, + page_token: String::new(), + root_digest: Some(root_directory_digest_info.into()), + digest_function: digest_function::Value::Sha256.into(), + })) + .await; + assert_eq!( + raw_response + .unwrap() + .into_inner() + .filter_map(|x| async move { Some(x.unwrap()) }) + .collect::>() + .await, + vec![GetTreeResponse { + directories: vec![ + root_directory.clone(), + sub_directories[0].clone(), + sub_directories[1].clone(), + sub_directories[2].clone(), + sub_directories[3].clone(), + sub_directories[4].clone() + ], + next_page_token: String::new() + }] + ); + } + + // Also verify that sending the root digest returns the entire tree as well. + { + let raw_response = cas_server + .get_tree(Request::new(GetTreeRequest { + instance_name: INSTANCE_NAME.to_string(), + page_size: 0, + page_token: format!("{root_directory_digest_info}"), + root_digest: Some(root_directory_digest_info.into()), + digest_function: digest_function::Value::Sha256.into(), + })) + .await; + assert_eq!( + raw_response + .unwrap() + .into_inner() + .filter_map(|x| async move { Some(x.unwrap()) }) + .collect::>() + .await, + vec![GetTreeResponse { + directories: vec![ + root_directory.clone(), + sub_directories[0].clone(), + sub_directories[1].clone(), + sub_directories[2].clone(), + sub_directories[3].clone(), + sub_directories[4].clone() + ], + next_page_token: String::new() + }] + ); + } Ok(()) } @@ -434,72 +470,105 @@ async fn get_tree_read_directories_with_paging() -> Result<(), Box>() - .await, - vec![GetTreeResponse { - directories: vec![root_directory.clone(), sub_directories[0].clone()], - next_page_token: format!("{}", sub_directory_digest_infos[1]), - }] - ); - let raw_response = cas_server - .get_tree(Request::new(GetTreeRequest { - instance_name: INSTANCE_NAME.to_string(), - page_size: 2, - page_token: format!("{}", sub_directory_digest_infos[1]), - root_digest: Some(root_directory_digest_info.into()), - digest_function: digest_function::Value::Sha256.into(), - })) - .await; - assert!(raw_response.is_ok()); - assert_eq!( - raw_response - .unwrap() - .into_inner() - .filter_map(|x| async move { Some(x.unwrap()) }) - .collect::>() - .await, - vec![GetTreeResponse { - directories: vec![sub_directories[1].clone(), sub_directories[2].clone()], - next_page_token: format!("{}", sub_directory_digest_infos[3]), - }] - ); - let raw_response = cas_server - .get_tree(Request::new(GetTreeRequest { - instance_name: INSTANCE_NAME.to_string(), - page_size: 2, - page_token: format!("{}", sub_directory_digest_infos[3]), - root_digest: Some(root_directory_digest_info.into()), - digest_function: digest_function::Value::Sha256.into(), - })) - .await; - assert!(raw_response.is_ok()); - assert_eq!( - raw_response - .unwrap() - .into_inner() - .filter_map(|x| async move { Some(x.unwrap()) }) - .collect::>() - .await, - vec![GetTreeResponse { - directories: vec![sub_directories[3].clone(), sub_directories[4].clone()], - next_page_token: String::new(), - }] - ); + + // First, verify that an empty initial page token is treated as if the client had sent the + // root digest and respects the page size. + { + let raw_response = cas_server + .get_tree(Request::new(GetTreeRequest { + instance_name: INSTANCE_NAME.to_string(), + page_size: 2, + page_token: String::new(), + root_digest: Some(root_directory_digest_info.into()), + digest_function: digest_function::Value::Sha256.into(), + })) + .await; + assert_eq!( + raw_response + .unwrap() + .into_inner() + .filter_map(|x| async move { Some(x.unwrap()) }) + .collect::>() + .await, + vec![GetTreeResponse { + directories: vec![root_directory.clone(), sub_directories[0].clone()], + next_page_token: format!("{}", sub_directory_digest_infos[1]), + }] + ); + } + + // Also verify that sending the root digest as the page token is treated as paging from the + // beginning and respects page size. + { + let raw_response = cas_server + .get_tree(Request::new(GetTreeRequest { + instance_name: INSTANCE_NAME.to_string(), + page_size: 2, + page_token: format!("{root_directory_digest_info}"), + root_digest: Some(root_directory_digest_info.into()), + digest_function: digest_function::Value::Sha256.into(), + })) + .await; + assert_eq!( + raw_response + .unwrap() + .into_inner() + .filter_map(|x| async move { Some(x.unwrap()) }) + .collect::>() + .await, + vec![GetTreeResponse { + directories: vec![root_directory.clone(), sub_directories[0].clone()], + next_page_token: format!("{}", sub_directory_digest_infos[1]), + }] + ); + } + + // Verify that paging from a non-initial page token will return the expected content. + { + let raw_response = cas_server + .get_tree(Request::new(GetTreeRequest { + instance_name: INSTANCE_NAME.to_string(), + page_size: 2, + page_token: format!("{}", sub_directory_digest_infos[1]), + root_digest: Some(root_directory_digest_info.into()), + digest_function: digest_function::Value::Sha256.into(), + })) + .await; + assert_eq!( + raw_response + .unwrap() + .into_inner() + .filter_map(|x| async move { Some(x.unwrap()) }) + .collect::>() + .await, + vec![GetTreeResponse { + directories: vec![sub_directories[1].clone(), sub_directories[2].clone()], + next_page_token: format!("{}", sub_directory_digest_infos[3]), + }] + ); + + let raw_response = cas_server + .get_tree(Request::new(GetTreeRequest { + instance_name: INSTANCE_NAME.to_string(), + page_size: 2, + page_token: format!("{}", sub_directory_digest_infos[3]), + root_digest: Some(root_directory_digest_info.into()), + digest_function: digest_function::Value::Sha256.into(), + })) + .await; + assert_eq!( + raw_response + .unwrap() + .into_inner() + .filter_map(|x| async move { Some(x.unwrap()) }) + .collect::>() + .await, + vec![GetTreeResponse { + directories: vec![sub_directories[3].clone(), sub_directories[4].clone()], + next_page_token: String::new(), + }] + ); + } Ok(()) }