feat(castore/fs): optional refscanner for ingest
Change-Id: Ieca06de4c2e2680d89fe05a380079fafa5454837 Reviewed-on: https://cl.tvl.fyi/c/depot/+/12529 Autosubmit: yuka <yuka@yuka.dev> Tested-by: BuildkiteCI Reviewed-by: flokli <flokli@flokli.de>
This commit is contained in:
		
							parent
							
								
									d277bd9fbf
								
							
						
					
					
						commit
						b82cacb449
					
				
					 5 changed files with 60 additions and 22 deletions
				
			
		|  | @ -16,6 +16,7 @@ use walkdir::WalkDir; | ||||||
| 
 | 
 | ||||||
| use crate::blobservice::BlobService; | use crate::blobservice::BlobService; | ||||||
| use crate::directoryservice::DirectoryService; | use crate::directoryservice::DirectoryService; | ||||||
|  | use crate::refscan::{ReferenceReader, ReferenceScanner}; | ||||||
| use crate::{B3Digest, Node}; | use crate::{B3Digest, Node}; | ||||||
| 
 | 
 | ||||||
| use super::ingest_entries; | use super::ingest_entries; | ||||||
|  | @ -29,16 +30,18 @@ use super::IngestionError; | ||||||
| ///
 | ///
 | ||||||
| /// This function will walk the filesystem using `walkdir` and will consume
 | /// This function will walk the filesystem using `walkdir` and will consume
 | ||||||
| /// `O(#number of entries)` space.
 | /// `O(#number of entries)` space.
 | ||||||
| #[instrument(skip(blob_service, directory_service), fields(path, indicatif.pb_show=1), err)] | #[instrument(skip(blob_service, directory_service, reference_scanner), fields(path, indicatif.pb_show=1), err)] | ||||||
| pub async fn ingest_path<BS, DS, P>( | pub async fn ingest_path<BS, DS, P, P2>( | ||||||
|     blob_service: BS, |     blob_service: BS, | ||||||
|     directory_service: DS, |     directory_service: DS, | ||||||
|     path: P, |     path: P, | ||||||
|  |     reference_scanner: Option<&ReferenceScanner<P2>>, | ||||||
| ) -> Result<Node, IngestionError<Error>> | ) -> Result<Node, IngestionError<Error>> | ||||||
| where | where | ||||||
|     P: AsRef<std::path::Path> + std::fmt::Debug, |     P: AsRef<std::path::Path> + std::fmt::Debug, | ||||||
|     BS: BlobService + Clone, |     BS: BlobService + Clone, | ||||||
|     DS: DirectoryService, |     DS: DirectoryService, | ||||||
|  |     P2: AsRef<[u8]> + Send + Sync, | ||||||
| { | { | ||||||
|     let span = Span::current(); |     let span = Span::current(); | ||||||
|     span.pb_set_message(&format!("Ingesting {:?}", path)); |     span.pb_set_message(&format!("Ingesting {:?}", path)); | ||||||
|  | @ -50,7 +53,8 @@ where | ||||||
|         .contents_first(true) |         .contents_first(true) | ||||||
|         .into_iter(); |         .into_iter(); | ||||||
| 
 | 
 | ||||||
|     let entries = dir_entries_to_ingestion_stream(blob_service, iter, path.as_ref()); |     let entries = | ||||||
|  |         dir_entries_to_ingestion_stream(blob_service, iter, path.as_ref(), reference_scanner); | ||||||
|     ingest_entries( |     ingest_entries( | ||||||
|         directory_service, |         directory_service, | ||||||
|         entries.inspect({ |         entries.inspect({ | ||||||
|  | @ -71,14 +75,16 @@ where | ||||||
| /// The produced stream is buffered, so uploads can happen concurrently.
 | /// The produced stream is buffered, so uploads can happen concurrently.
 | ||||||
| ///
 | ///
 | ||||||
| /// The root is the [Path] in the filesystem that is being ingested into the castore.
 | /// The root is the [Path] in the filesystem that is being ingested into the castore.
 | ||||||
| pub fn dir_entries_to_ingestion_stream<'a, BS, I>( | pub fn dir_entries_to_ingestion_stream<'a, BS, I, P>( | ||||||
|     blob_service: BS, |     blob_service: BS, | ||||||
|     iter: I, |     iter: I, | ||||||
|     root: &'a std::path::Path, |     root: &'a std::path::Path, | ||||||
|  |     reference_scanner: Option<&'a ReferenceScanner<P>>, | ||||||
| ) -> BoxStream<'a, Result<IngestionEntry, Error>> | ) -> BoxStream<'a, Result<IngestionEntry, Error>> | ||||||
| where | where | ||||||
|     BS: BlobService + Clone + 'a, |     BS: BlobService + Clone + 'a, | ||||||
|     I: Iterator<Item = Result<DirEntry, walkdir::Error>> + Send + 'a, |     I: Iterator<Item = Result<DirEntry, walkdir::Error>> + Send + 'a, | ||||||
|  |     P: AsRef<[u8]> + Send + Sync, | ||||||
| { | { | ||||||
|     let prefix = root.parent().unwrap_or_else(|| std::path::Path::new("")); |     let prefix = root.parent().unwrap_or_else(|| std::path::Path::new("")); | ||||||
| 
 | 
 | ||||||
|  | @ -89,7 +95,13 @@ where | ||||||
|                 async move { |                 async move { | ||||||
|                     match x { |                     match x { | ||||||
|                         Ok(dir_entry) => { |                         Ok(dir_entry) => { | ||||||
|                             dir_entry_to_ingestion_entry(blob_service, &dir_entry, prefix).await |                             dir_entry_to_ingestion_entry( | ||||||
|  |                                 blob_service, | ||||||
|  |                                 &dir_entry, | ||||||
|  |                                 prefix, | ||||||
|  |                                 reference_scanner, | ||||||
|  |                             ) | ||||||
|  |                             .await | ||||||
|                         } |                         } | ||||||
|                         Err(e) => Err(Error::Stat( |                         Err(e) => Err(Error::Stat( | ||||||
|                             prefix.to_path_buf(), |                             prefix.to_path_buf(), | ||||||
|  | @ -107,13 +119,15 @@ where | ||||||
| ///
 | ///
 | ||||||
| /// The prefix path is stripped from the path of each entry. This is usually the parent path
 | /// The prefix path is stripped from the path of each entry. This is usually the parent path
 | ||||||
| /// of the path being ingested so that the last element of the stream only has one component.
 | /// of the path being ingested so that the last element of the stream only has one component.
 | ||||||
| pub async fn dir_entry_to_ingestion_entry<BS>( | pub async fn dir_entry_to_ingestion_entry<BS, P>( | ||||||
|     blob_service: BS, |     blob_service: BS, | ||||||
|     entry: &DirEntry, |     entry: &DirEntry, | ||||||
|     prefix: &std::path::Path, |     prefix: &std::path::Path, | ||||||
|  |     reference_scanner: Option<&ReferenceScanner<P>>, | ||||||
| ) -> Result<IngestionEntry, Error> | ) -> Result<IngestionEntry, Error> | ||||||
| where | where | ||||||
|     BS: BlobService, |     BS: BlobService, | ||||||
|  |     P: AsRef<[u8]>, | ||||||
| { | { | ||||||
|     let file_type = entry.file_type(); |     let file_type = entry.file_type(); | ||||||
| 
 | 
 | ||||||
|  | @ -134,13 +148,18 @@ where | ||||||
|             .into_os_string() |             .into_os_string() | ||||||
|             .into_vec(); |             .into_vec(); | ||||||
| 
 | 
 | ||||||
|  |         if let Some(reference_scanner) = &reference_scanner { | ||||||
|  |             reference_scanner.scan(&target); | ||||||
|  |         } | ||||||
|  | 
 | ||||||
|         Ok(IngestionEntry::Symlink { path, target }) |         Ok(IngestionEntry::Symlink { path, target }) | ||||||
|     } else if file_type.is_file() { |     } else if file_type.is_file() { | ||||||
|         let metadata = entry |         let metadata = entry | ||||||
|             .metadata() |             .metadata() | ||||||
|             .map_err(|e| Error::Stat(entry.path().to_path_buf(), e.into()))?; |             .map_err(|e| Error::Stat(entry.path().to_path_buf(), e.into()))?; | ||||||
| 
 | 
 | ||||||
|         let digest = upload_blob(blob_service, entry.path().to_path_buf()).await?; |         let digest = | ||||||
|  |             upload_blob(blob_service, entry.path().to_path_buf(), reference_scanner).await?; | ||||||
| 
 | 
 | ||||||
|         Ok(IngestionEntry::Regular { |         Ok(IngestionEntry::Regular { | ||||||
|             path, |             path, | ||||||
|  | @ -156,13 +175,15 @@ where | ||||||
| } | } | ||||||
| 
 | 
 | ||||||
| /// Uploads the file at the provided [Path] the the [BlobService].
 | /// Uploads the file at the provided [Path] the the [BlobService].
 | ||||||
| #[instrument(skip(blob_service), fields(path, indicatif.pb_show=1), err)] | #[instrument(skip(blob_service, reference_scanner), fields(path, indicatif.pb_show=1), err)] | ||||||
| async fn upload_blob<BS>( | async fn upload_blob<BS, P>( | ||||||
|     blob_service: BS, |     blob_service: BS, | ||||||
|     path: impl AsRef<std::path::Path>, |     path: impl AsRef<std::path::Path>, | ||||||
|  |     reference_scanner: Option<&ReferenceScanner<P>>, | ||||||
| ) -> Result<B3Digest, Error> | ) -> Result<B3Digest, Error> | ||||||
| where | where | ||||||
|     BS: BlobService, |     BS: BlobService, | ||||||
|  |     P: AsRef<[u8]>, | ||||||
| { | { | ||||||
|     let span = Span::current(); |     let span = Span::current(); | ||||||
|     span.pb_set_style(&tvix_tracing::PB_TRANSFER_STYLE); |     span.pb_set_style(&tvix_tracing::PB_TRANSFER_STYLE); | ||||||
|  | @ -184,9 +205,16 @@ where | ||||||
|     }); |     }); | ||||||
| 
 | 
 | ||||||
|     let mut writer = blob_service.open_write().await; |     let mut writer = blob_service.open_write().await; | ||||||
|     tokio::io::copy(&mut BufReader::new(reader), &mut writer) |     if let Some(reference_scanner) = reference_scanner { | ||||||
|         .await |         let mut reader = ReferenceReader::new(reference_scanner, BufReader::new(reader)); | ||||||
|         .map_err(|e| Error::BlobRead(path.as_ref().to_path_buf(), e))?; |         tokio::io::copy(&mut reader, &mut writer) | ||||||
|  |             .await | ||||||
|  |             .map_err(|e| Error::BlobRead(path.as_ref().to_path_buf(), e))?; | ||||||
|  |     } else { | ||||||
|  |         tokio::io::copy(&mut BufReader::new(reader), &mut writer) | ||||||
|  |             .await | ||||||
|  |             .map_err(|e| Error::BlobRead(path.as_ref().to_path_buf(), e))?; | ||||||
|  |     } | ||||||
| 
 | 
 | ||||||
|     let digest = writer |     let digest = writer | ||||||
|         .close() |         .close() | ||||||
|  |  | ||||||
|  | @ -21,10 +21,11 @@ async fn symlink() { | ||||||
|     ) |     ) | ||||||
|     .unwrap(); |     .unwrap(); | ||||||
| 
 | 
 | ||||||
|     let root_node = ingest_path( |     let root_node = ingest_path::<_, _, _, &[u8]>( | ||||||
|         blob_service, |         blob_service, | ||||||
|         directory_service, |         directory_service, | ||||||
|         tmpdir.path().join("doesntmatter"), |         tmpdir.path().join("doesntmatter"), | ||||||
|  |         None, | ||||||
|     ) |     ) | ||||||
|     .await |     .await | ||||||
|     .expect("must succeed"); |     .expect("must succeed"); | ||||||
|  | @ -46,10 +47,11 @@ async fn single_file() { | ||||||
| 
 | 
 | ||||||
|     std::fs::write(tmpdir.path().join("root"), HELLOWORLD_BLOB_CONTENTS).unwrap(); |     std::fs::write(tmpdir.path().join("root"), HELLOWORLD_BLOB_CONTENTS).unwrap(); | ||||||
| 
 | 
 | ||||||
|     let root_node = ingest_path( |     let root_node = ingest_path::<_, _, _, &[u8]>( | ||||||
|         blob_service.clone(), |         blob_service.clone(), | ||||||
|         directory_service, |         directory_service, | ||||||
|         tmpdir.path().join("root"), |         tmpdir.path().join("root"), | ||||||
|  |         None, | ||||||
|     ) |     ) | ||||||
|     .await |     .await | ||||||
|     .expect("must succeed"); |     .expect("must succeed"); | ||||||
|  | @ -84,9 +86,14 @@ async fn complicated() { | ||||||
|     // File ``keep/.keep`
 |     // File ``keep/.keep`
 | ||||||
|     std::fs::write(tmpdir.path().join("keep").join(".keep"), vec![]).unwrap(); |     std::fs::write(tmpdir.path().join("keep").join(".keep"), vec![]).unwrap(); | ||||||
| 
 | 
 | ||||||
|     let root_node = ingest_path(blob_service.clone(), &directory_service, tmpdir.path()) |     let root_node = ingest_path::<_, _, _, &[u8]>( | ||||||
|         .await |         blob_service.clone(), | ||||||
|         .expect("must succeed"); |         &directory_service, | ||||||
|  |         tmpdir.path(), | ||||||
|  |         None, | ||||||
|  |     ) | ||||||
|  |     .await | ||||||
|  |     .expect("must succeed"); | ||||||
| 
 | 
 | ||||||
|     // ensure root_node matched expectations
 |     // ensure root_node matched expectations
 | ||||||
|     assert_eq!( |     assert_eq!( | ||||||
|  |  | ||||||
|  | @ -89,10 +89,11 @@ async fn filtered_ingest( | ||||||
|     let dir_entries = entries.into_iter().rev().map(Ok); |     let dir_entries = entries.into_iter().rev().map(Ok); | ||||||
| 
 | 
 | ||||||
|     state.tokio_handle.block_on(async { |     state.tokio_handle.block_on(async { | ||||||
|         let entries = tvix_castore::import::fs::dir_entries_to_ingestion_stream( |         let entries = tvix_castore::import::fs::dir_entries_to_ingestion_stream::<'_, _, _, &[u8]>( | ||||||
|             &state.blob_service, |             &state.blob_service, | ||||||
|             dir_entries, |             dir_entries, | ||||||
|             path, |             path, | ||||||
|  |             None, // TODO re-scan
 | ||||||
|         ); |         ); | ||||||
|         ingest_entries(&state.directory_service, entries) |         ingest_entries(&state.directory_service, entries) | ||||||
|             .await |             .await | ||||||
|  |  | ||||||
|  | @ -339,10 +339,11 @@ async fn run_cli(cli: Cli) -> Result<(), Box<dyn std::error::Error + Send + Sync | ||||||
|                     async move { |                     async move { | ||||||
|                         // Ingest the given path.
 |                         // Ingest the given path.
 | ||||||
| 
 | 
 | ||||||
|                         ingest_path( |                         ingest_path::<_, _, _, &[u8]>( | ||||||
|                             blob_service, |                             blob_service, | ||||||
|                             directory_service, |                             directory_service, | ||||||
|                             PathBuf::from(elem.path.to_absolute_path()), |                             PathBuf::from(elem.path.to_absolute_path()), | ||||||
|  |                             None, | ||||||
|                         ) |                         ) | ||||||
|                         .await |                         .await | ||||||
|                         .map(|root_node| (elem, root_node)) |                         .map(|root_node| (elem, root_node)) | ||||||
|  |  | ||||||
|  | @ -123,9 +123,10 @@ where | ||||||
|     PS: AsRef<dyn PathInfoService>, |     PS: AsRef<dyn PathInfoService>, | ||||||
|     NS: NarCalculationService, |     NS: NarCalculationService, | ||||||
| { | { | ||||||
|     let root_node = ingest_path(blob_service, directory_service, path.as_ref()) |     let root_node = | ||||||
|         .await |         ingest_path::<_, _, _, &[u8]>(blob_service, directory_service, path.as_ref(), None) | ||||||
|         .map_err(|e| std::io::Error::new(std::io::ErrorKind::Other, e))?; |             .await | ||||||
|  |             .map_err(|e| std::io::Error::new(std::io::ErrorKind::Other, e))?; | ||||||
| 
 | 
 | ||||||
|     // Ask for the NAR size and sha256
 |     // Ask for the NAR size and sha256
 | ||||||
|     let (nar_size, nar_sha256) = nar_calculation_service.calculate_nar(&root_node).await?; |     let (nar_size, nar_sha256) = nar_calculation_service.calculate_nar(&root_node).await?; | ||||||
|  |  | ||||||
		Loading…
	
	Add table
		Add a link
		
	
		Reference in a new issue