Update add_named_blob to create directory under aliases atomically (#36)

* Update `add_named_blob` so blob alias dirs are created atomically

* Add some comments to `add_named_blob` function

* Update `add_named_blob` to copy+hash the contents in one pass

Author: kylewlacy

crates/brioche-autopack/src/lib.rs

@@ -1043,16 +1043,17 @@ fn add_named_blob_from(
         Path::new(filename)
     };
 
-    let mut file = std::fs::File::open(path)?;
+    let file = std::fs::File::open(path)?;
     let metadata = file.metadata()?;
 
     let permissions = metadata.permissions();
     let mode = permissions.mode();
     let is_executable = mode & 0o111 != 0;
 
+    let file_reader = std::io::BufReader::new(file);
     let resource_path = brioche_resources::add_named_blob(
         &ctx.config.resource_dir,
-        &mut file,
+        file_reader,
         is_executable,
         alias_name,
     )?;

crates/brioche-packer/src/main.rs

@@ -240,9 +240,10 @@ fn run_update_source(args: &UpdateSourceArgs) -> eyre::Result<()> {
             let new_source_permissions = new_source.metadata()?.permissions();
             let is_executable = is_executable(&new_source_permissions);
 
+            let new_source_reader = std::io::BufReader::new(new_source);
             let new_source_resource = brioche_resources::add_named_blob(
                 &output_resource_dir,
-                &new_source,
+                new_source_reader,
                 is_executable,
                 new_name,
             )?;

crates/brioche-resources/src/lib.rs

@@ -106,50 +106,114 @@ fn find_resource_dirs_from_program(
 
 pub fn add_named_blob(
     resource_dir: &Path,
-    mut contents: impl std::io::Seek + std::io::Read,
+    mut contents: impl std::io::BufRead,
     executable: bool,
     name: &Path,
 ) -> Result<PathBuf, AddBlobError> {
-    let mut hasher = blake3::Hasher::new();
-    std::io::copy(&mut contents, &mut hasher)?;
-    let hash = hasher.finalize();
-
-    let blob_suffix = if executable { ".x" } else { "" };
-    let blob_name = format!("{hash}{blob_suffix}");
-
-    contents.seek(std::io::SeekFrom::Start(0))?;
-
+    // Create the 'blobs' directory
     let blob_dir = resource_dir.join("blobs");
-    let blob_path = blob_dir.join(&blob_name);
-    let blob_temp_id = ulid::Ulid::new();
-    let blob_temp_path = blob_dir.join(format!("{blob_name}-{blob_temp_id}"));
     std::fs::create_dir_all(&blob_dir)?;
 
+    // Open a temporary file to copy the contents to
+    let blob_temp_id = ulid::Ulid::new();
+    let blob_temp_path = blob_dir.join(blob_temp_id.to_string());
     let mut blob_file_options = std::fs::OpenOptions::new();
     blob_file_options.create_new(true).write(true);
     if executable {
         blob_file_options.mode(0o777);
     }
     let mut blob_file = blob_file_options.open(&blob_temp_path)?;
-    std::io::copy(&mut contents, &mut blob_file)?;
+
+    // Read the contents, both copying it to the temporary file and hashing
+    // as we go
+    let mut hasher = blake3::Hasher::new();
+    loop {
+        let buf = contents.fill_buf()?;
+        if buf.is_empty() {
+            break;
+        }
+
+        hasher.update(buf);
+        blob_file.write_all(buf)?;
+
+        let consumed = buf.len();
+        contents.consume(consumed);
+    }
+
+    // Get the hash of the contents, which we'll use as the blob's name
+    let hash = hasher.finalize();
+
+    // Get the final blob's filename. We use a suffix to distinguish identical
+    // blobs with different permissions
+    let blob_suffix = if executable { ".x" } else { "" };
+    let blob_name = format!("{hash}{blob_suffix}");
+    let blob_path = blob_dir.join(&blob_name);
+
+    // Rename the blob to its final path
     drop(blob_file);
     std::fs::rename(&blob_temp_path, &blob_path)?;
 
-    let alias_dir = resource_dir.join("aliases").join(name).join(&blob_name);
-    std::fs::create_dir_all(&alias_dir)?;
+    // Create a temporary directory for the alias dir
+    let alias_temp_dir = resource_dir.join(format!("{blob_name}-{blob_temp_id}-alias"));
+    let alias_temp_path = alias_temp_dir.join(name);
+    std::fs::create_dir(&alias_temp_dir)?;
 
-    let temp_alias_path = alias_dir.join(format!("{}-{blob_temp_id}", name.display()));
-    let alias_path = alias_dir.join(name);
+    // Create the symlink within the temporary dir
+    let alias_parent_dir = resource_dir.join("aliases").join(name);
+    let alias_dir = alias_parent_dir.join(&blob_name);
     let blob_pack_relative_path = pathdiff::diff_paths(&blob_path, &alias_dir)
         .expect("blob path is not a prefix of alias path");
-    std::os::unix::fs::symlink(blob_pack_relative_path, &temp_alias_path)?;
-    std::fs::rename(&temp_alias_path, &alias_path)?;
+    std::os::unix::fs::symlink(&blob_pack_relative_path, &alias_temp_path)?;
+
+    // Create directory for the alias dir
+    std::fs::create_dir_all(&alias_parent_dir)?;
 
+    // Rename the temp dir to the final alias path. This ensures that the alias
+    // dir itself is atomic, and never appears empty
+    let alias_path = alias_dir.join(name);
+    let result = std::fs::rename(&alias_temp_dir, alias_dir);
+    match result {
+        Ok(()) => {
+            // Alias dir created successfully
+        }
+        Err(err)
+            if err.kind() == std::io::ErrorKind::AlreadyExists
+                || err.kind() == std::io::ErrorKind::DirectoryNotEmpty =>
+        {
+            // Could not rename temp alias dir to final path. On Unix, this
+            // means that the alias dir already exists and is non-empty
+
+            // Clean up the temporary dir first
+            std::fs::remove_dir_all(&alias_temp_dir)?;
+
+            // Try to create the symlink again-- this time in its final path
+            let result = std::os::unix::fs::symlink(&blob_pack_relative_path, &alias_path);
+            match result {
+                Ok(()) => {
+                    // Symlink created successfully. This means the alias
+                    // dir already existed and was not empty, but contained
+                    // something else? This probably shouldn't happen...
+                }
+                Err(err) if err.kind() == std::io::ErrorKind::AlreadyExists => {
+                    // Path already exists, nothing to do
+                }
+                Err(err) => {
+                    return Err(err.into());
+                }
+            }
+        }
+        Err(err) => {
+            return Err(err.into());
+        }
+    }
+
+    // Return the symlink alias path relative to the resource dir
     let alias_path = alias_path
         .strip_prefix(resource_dir)
         .expect("alias path is not in resource dir");
     Ok(alias_path.to_owned())
 }
+
 pub fn add_named_resource_directory(
     resource_dir: &Path,
     source: &Path,

crates/brioche-strip/src/main.rs

@@ -438,9 +438,10 @@ fn finish_remapped_file(remapped_file: RemapFile) -> eyre::Result<()> {
                     // Add the temp file as a new resource. We re-use the
                     // original program's name and permissions
                     temp_file.rewind()?;
+                    let temp_file_reader = std::io::BufReader::new(temp_file);
                     let new_source_resource = brioche_resources::add_named_blob(
                         &output_resource_dir,
-                        &mut temp_file,
+                        temp_file_reader,
                         is_executable,
                         program_name,
                     )?;