commit 29bb633d12d6a100f478216a1c279474e171535d
parent e64fe9607d9834e24a5f1c14b6a2965d40f3db6e
Author: mauricefisher64 <92736594+mauricefisher64@users.noreply.github.com>
Date: Thu, 21 Jul 2022 12:18:48 -0400
Optimize performance of large assets (#84)
* Optimize file writing
* test load without verify
* more perf test
* Test hardware accelerated Sha2
* Test interleaved hashing
* Remove print statement
* chunked Blake3
* Add Blake3 support to hash_utils
* Fix build issue
* Remove unneeded const value
* optimize cleanup
* Remove asm until Windows builds are fixed.
* fix bad toml
* Document feature to disable interleaved IO
* ThreadSendError is not used
* Tweak wording
* Tweak wording
* Fix bad wording
* adjust logic on when to let save_to_asset perform the copy
Co-authored-by: mauricefisher64 <mfisher@Octopus>
Co-authored-by: Eric Scouten <scouten@adobe.com>
Diffstat:
9 files changed, 297 insertions(+), 30 deletions(-)
diff --git a/README.md b/README.md
@@ -58,6 +58,7 @@ The Rust SDK crate provides:
* `file_io` enables manifest generation, signing via OpenSSL, and embedding manifests in various file formats.
* `serialize_thumbnails` includes binary thumbnail data in the [Serde](https://serde.rs/) serialization output.
* `xmp_write` enables updating XMP on embed with the `dcterms:provenance` field. (Requires [xmp_toolkit](https://crates.io/crates/xmp_toolkit).)
+* `no_interleaved_io` the SDK uses threaded I/O for some operations to improve performance. Using this feature will force fully synchronous I/O.
## License
diff --git a/sdk/Cargo.toml b/sdk/Cargo.toml
@@ -24,6 +24,7 @@ bmff = [] # Work in progress support for BMFF-based containers
file_io = ["openssl"]
serialize_thumbnails = []
xmp_write = ["xmp_toolkit"]
+no_interleaved_io = ["file_io"]
# The diagnostics feature is unsupported and might be removed.
# It enables some low-overhead timing features used in our development cycle.
@@ -50,6 +51,7 @@ async-trait = { version = "0.1.48", optional = true }
atree = "0.5.2"
base64 = "0.13.0"
bcder = "0.6.0"
+blake3 = "1.0.0"
bytes = "1.1.0"
byteorder = "1.3.4"
chrono = { version = "0.4.19", features = ["wasmbind"] }
diff --git a/sdk/src/asset_handlers/bmff_io.rs b/sdk/src/asset_handlers/bmff_io.rs
@@ -13,7 +13,7 @@
use std::collections::HashMap;
use std::convert::{From, TryFrom};
-use std::fs::File;
+use std::fs::{File, OpenOptions};
use std::io::{Read, Seek, SeekFrom, Write};
use std::path::Path;
@@ -26,7 +26,7 @@ use atree::{Arena, Token};
use tempfile::{Builder, NamedTempFile};
use crate::assertions::ExclusionsMap;
-use crate::asset_io::{AssetIO, CAILoader, CAIRead, HashObjectPositions};
+use crate::asset_io::{AssetIO, AssetPatch, CAILoader, CAIRead, HashObjectPositions};
use crate::error::{Error, Result};
use crate::utils::hash_utils::{vec_compare, Exclusion};
@@ -904,6 +904,10 @@ impl CAILoader for BmffIO {
}
impl AssetIO for BmffIO {
+ fn asset_patch_ref(&self) -> Option<&dyn AssetPatch> {
+ Some(self)
+ }
+
fn read_cai_store(&self, asset_path: &Path) -> Result<Vec<u8>> {
let mut f = File::open(asset_path)?;
self.read_cai(&mut f)
@@ -1060,9 +1064,11 @@ impl AssetIO for BmffIO {
_ => (), // todo: handle more patching cases as necessary
}
- std::fs::copy(temp_file.path(), asset_path)?;
-
- Ok(())
+ // copy temp file to asset
+ std::fs::rename(&temp_file.path(), asset_path)
+ // if rename fails, try to copy in case we are on different volumes
+ .or_else(|_| std::fs::copy(&temp_file.path(), asset_path).and(Ok(())))
+ .map_err(Error::IoError)
}
fn get_object_locations(
@@ -1074,9 +1080,93 @@ impl AssetIO for BmffIO {
}
}
+impl AssetPatch for BmffIO {
+ fn patch_cai_store(&self, asset_path: &std::path::Path, store_bytes: &[u8]) -> Result<()> {
+ let mut asset = OpenOptions::new()
+ .write(true)
+ .read(true)
+ .create(false)
+ .open(asset_path)?;
+ let size = asset.seek(SeekFrom::End(0))?;
+ asset.seek(SeekFrom::Start(0))?;
+
+ // create root node
+ let root_box = BoxInfo {
+ path: "".to_string(),
+ offset: 0,
+ size: size as u64,
+ box_type: BoxType::Empty,
+ parent: None,
+ user_type: None,
+ version: None,
+ flags: None,
+ };
+
+ let (mut bmff_tree, root_token) = Arena::with_data(root_box);
+ let mut bmff_map: HashMap<String, Vec<Token>> = HashMap::new();
+
+ // build layout of the BMFF structure
+ build_bmff_tree(
+ &mut asset,
+ size as u64,
+ &mut bmff_tree,
+ &root_token,
+ &mut bmff_map,
+ )?;
+
+ // get position to insert c2pa
+ let (c2pa_start, c2pa_length) = if let Some(uuid_tokens) = bmff_map.get("/uuid") {
+ let uuid_info = &bmff_tree[uuid_tokens[0]].data;
+
+ // is this a C2PA manifest
+ let is_c2pa = if let Some(uuid) = &uuid_info.user_type {
+ // make sure it is a C2PA box
+ vec_compare(&C2PA_UUID, uuid)
+ } else {
+ false
+ };
+
+ if is_c2pa {
+ (uuid_info.offset, Some(uuid_info.size))
+ } else {
+ (0, None)
+ }
+ } else {
+ return Err(Error::BadParam(
+ "patch_cai_store found no manifest store to patch.".to_string(),
+ ));
+ };
+
+ if let Some(manifest_length) = c2pa_length {
+ let mut new_c2pa_box: Vec<u8> = Vec::with_capacity(store_bytes.len() * 2);
+ let merkle_data: &[u8] = &[]; // not yet supported
+ write_c2pa_box(&mut new_c2pa_box, store_bytes, true, merkle_data)?;
+ let new_c2pa_box_size = new_c2pa_box.len();
+
+ if new_c2pa_box_size as u64 == manifest_length {
+ asset.seek(SeekFrom::Start(c2pa_start as u64))?;
+ asset.write_all(&new_c2pa_box)?;
+ Ok(())
+ } else {
+ Err(Error::BadParam(
+ "patch_cai_store store size mismatch.".to_string(),
+ ))
+ }
+ } else {
+ Err(Error::BadParam(
+ "patch_cai_store store size mismatch.".to_string(),
+ ))
+ }
+ }
+}
+
#[cfg(feature = "bmff")]
#[cfg(test)]
pub mod tests {
+ #![allow(clippy::expect_used)]
+ #![allow(clippy::panic)]
+ #![allow(clippy::unwrap_used)]
+
use tempfile::tempdir;
use super::*;
@@ -1150,4 +1240,33 @@ pub mod tests {
}
assert!(success)
}
+
+ #[test]
+ fn test_patch_c2pa_write_mp4() {
+ let test_data = "some test data".as_bytes();
+ let source = fixture_path("video1.mp4");
+
+ let mut success = false;
+ if let Ok(temp_dir) = tempdir() {
+ let output = temp_dir_path(&temp_dir, "mp4_test.mp4");
+
+ if let Ok(_size) = std::fs::copy(&source, &output) {
+ let bmff = BmffIO::new("mp4");
+
+ if let Ok(source_data) = bmff.read_cai_store(&output) {
+ // create replacement data of same size
+ let mut new_data = vec![0u8; source_data.len()];
+ new_data[..test_data.len()].copy_from_slice(test_data);
+ bmff.patch_cai_store(&output, &new_data).unwrap();
+
+ let replaced = bmff.read_cai_store(&output).unwrap();
+
+ assert_eq!(new_data, replaced);
+
+ success = true;
+ }
+ }
+ }
+ assert!(success)
+ }
}
diff --git a/sdk/src/asset_io.rs b/sdk/src/asset_io.rs
@@ -56,9 +56,25 @@ pub trait AssetIO {
// Write the CAI block to an asset
fn save_cai_store(&self, asset_path: &Path, store_bytes: &[u8]) -> Result<()>;
- /// List of standard object offests
+ /// List of standard object offsets
/// If the offsets exist return the start of those locations other it should
/// return the calculated location of when it should start. There may still be a
/// length if the format contains extra header information for example.
fn get_object_locations(&self, asset_path: &Path) -> Result<Vec<HashObjectPositions>>;
+
+ // Returns [`AssetPatch`] trait if this I/O handler supports patching.
+ fn asset_patch_ref(&self) -> Option<&dyn AssetPatch> {
+ None
+ }
+}
+
+// `AssetPatch` optimizes output generation for asset_io handlers that
+// are able to patch blocks of data without changing any other data. The
+// resultant file must still be a valid asset. This saves having to rewrite
+// assets since only the patched bytes are modified.
+pub trait AssetPatch {
+ // Patches an existing manifest store with new manifest store.
+ // Only existing manifest stores of the same size may be patched
+ // since any other changes will invalidate asset hashes.
+ fn patch_cai_store(&self, asset_path: &Path, store_bytes: &[u8]) -> Result<()>;
}
diff --git a/sdk/src/error.rs b/sdk/src/error.rs
@@ -155,6 +155,9 @@ pub enum Error {
#[error("could not create valid JUMBF for claim")]
JumbfCreationError,
+ #[error("thread receive error")]
+ ThreadReceiveError,
+
/// No JUMBF data found.
/// TODO before merging PR: Does this error case need to be part of the public API?
#[error("no JUMBF data found")]
diff --git a/sdk/src/ingredient.rs b/sdk/src/ingredient.rs
@@ -382,7 +382,6 @@ impl Ingredient {
) -> Result<Self> {
Self::from_file_impl(path.as_ref(), options)
}
-
// Internal implementation to avoid code bloat.
#[cfg(feature = "file_io")]
fn from_file_impl(path: &Path, options: &dyn IngredientOptions) -> Result<Self> {
@@ -407,8 +406,6 @@ impl Ingredient {
if let Some(opt_title) = options.title(path) {
ingredient.title = opt_title;
}
- // read the file into a buffer for processing
- let buf = std::fs::read(path).map_err(wrap_io_err)?;
// optionally generate a hash so we know if the file has changed
ingredient.hash = options.hash(path);
@@ -417,7 +414,7 @@ impl Ingredient {
// generate a store from the buffer and then validate from the asset path
// load and verify store in single call - no need to call low level jumbf_io functions
- match Store::load_from_memory(&ingredient.format, &buf, true, &mut report) {
+ match Store::load_from_asset(path, true, &mut report) {
Ok(store) => {
// generate ValidationStatus from ValidationItems filtering for only errors
let statuses = status_for_store(&store, &mut report);
@@ -436,8 +433,7 @@ impl Ingredient {
}
ingredient.active_manifest = Some(claim.label().to_string());
}
- ingredient.manifest_data =
- jumbf_io::load_jumbf_from_memory(&ingredient.format, &buf).ok();
+ ingredient.manifest_data = jumbf_io::load_jumbf_from_file(path).ok();
ingredient.validation_status = if statuses.is_empty() {
None
} else {
diff --git a/sdk/src/jumbf_io.rs b/sdk/src/jumbf_io.rs
@@ -133,7 +133,7 @@ pub fn save_jumbf_to_file(data: &[u8], in_path: &Path, out_path: Option<&Path>)
let ext = get_file_extension(in_path).ok_or(Error::UnsupportedType)?;
// if no output path make a new file based off of source file name
- let img_out_path: PathBuf = match out_path {
+ let asset_out_path: PathBuf = match out_path {
Some(p) => p.to_owned(),
None => {
let filename_osstr = in_path.file_stem().ok_or(Error::UnsupportedType)?;
@@ -145,12 +145,22 @@ pub fn save_jumbf_to_file(data: &[u8], in_path: &Path, out_path: Option<&Path>)
};
// clone output to be overwritten
- if in_path != img_out_path {
- fs::copy(&in_path, &img_out_path).map_err(Error::IoError)?;
+ if in_path != asset_out_path {
+ fs::copy(&in_path, &asset_out_path).map_err(Error::IoError)?;
}
match get_assetio_handler(&ext) {
- Some(asset_handler) => asset_handler.save_cai_store(&img_out_path, data),
+ Some(asset_handler) => {
+ // patch if possible to save time and resources
+ if let Some(patch_handler) = asset_handler.asset_patch_ref() {
+ if patch_handler.patch_cai_store(&asset_out_path, data).is_ok() {
+ return Ok(());
+ }
+ }
+
+ // couldn't patch so just save
+ asset_handler.save_cai_store(&asset_out_path, data)
+ }
_ => Err(Error::UnsupportedType),
}
}
diff --git a/sdk/src/manifest.rs b/sdk/src/manifest.rs
@@ -636,15 +636,23 @@ impl Manifest {
return Err(Error::FileNotFound(path));
}
// we need to copy the source to target before setting the asset info
+ let mut copied = false;
if !dest_path.as_ref().exists() {
std::fs::copy(&source_path, &dest_path)?;
+ copied = true;
}
// first add the information about the target file
self.set_asset_from_path(dest_path.as_ref());
// convert the manifest to a store
let mut store = self.to_store()?;
// sign and write our store to to the output image file
- store.save_to_asset(source_path.as_ref(), signer, dest_path.as_ref())?;
+ if copied {
+ // set source and dest to same path to avoid an unnecessary copy since we already copied above
+ store.save_to_asset(dest_path.as_ref(), signer, dest_path.as_ref())?;
+ } else {
+ // save to asset will do the copy
+ store.save_to_asset(source_path.as_ref(), signer, dest_path.as_ref())?;
+ }
// todo: update xmp
Ok(store)
diff --git a/sdk/src/utils/hash_utils.rs b/sdk/src/utils/hash_utils.rs
@@ -85,6 +85,7 @@ pub fn hash_by_type(hash_type: u8, data: &[u8]) -> Option<Multihash> {
}
}
+#[derive(Clone)]
enum Hasher {
SHA256(Sha256),
SHA384(Sha384),
@@ -232,25 +233,71 @@ pub fn hash_asset_by_alg(
}
};
- // hash the data for ranges
- for r in ranges.into_smallvec() {
- let start = r.start();
- let end = r.end();
- let mut chunk_left = end - start + 1;
+ if cfg!(feature = "no_interleaved_io") {
+ // hash the data for ranges
+ for r in ranges.into_smallvec() {
+ let start = r.start();
+ let end = r.end();
+ let mut chunk_left = end - start + 1;
- // move to start of range
- data.seek(SeekFrom::Start(*start))?;
+ // move to start of range
+ data.seek(SeekFrom::Start(*start))?;
- loop {
- let mut chunk = vec![0u8; std::cmp::min(chunk_left as usize, MAX_HASH_BUF)];
+ loop {
+ let mut chunk = vec![0u8; std::cmp::min(chunk_left as usize, MAX_HASH_BUF)];
+
+ data.read_exact(&mut chunk)?;
+
+ hasher_enum.update(&chunk);
+ chunk_left -= chunk.len() as u64;
+ if chunk_left == 0 {
+ break;
+ }
+ }
+ }
+ } else {
+ // hash the data for ranges
+ for r in ranges.into_smallvec() {
+ let start = r.start();
+ let end = r.end();
+ let mut chunk_left = end - start + 1;
+
+ // move to start of range
+ data.seek(SeekFrom::Start(*start))?;
+
+ let mut chunk = vec![0u8; std::cmp::min(chunk_left as usize, MAX_HASH_BUF)];
data.read_exact(&mut chunk)?;
- hasher_enum.update(&chunk);
+ loop {
+ let (tx, rx) = std::sync::mpsc::channel();
- chunk_left -= chunk.len() as u64;
- if chunk_left == 0 {
- break;
+ chunk_left -= chunk.len() as u64;
+
+ std::thread::spawn(move || {
+ hasher_enum.update(&chunk);
+ tx.send(hasher_enum).unwrap_or_default();
+ });
+
+ // are we done
+ if chunk_left == 0 {
+ hasher_enum = match rx.recv() {
+ Ok(hasher) => hasher,
+ Err(_) => return Err(Error::ThreadReceiveError),
+ };
+ break;
+ }
+
+ // read next chunk while we wait for hash
+ let mut next_chunk = vec![0u8; std::cmp::min(chunk_left as usize, MAX_HASH_BUF)];
+ data.read_exact(&mut next_chunk)?;
+
+ hasher_enum = match rx.recv() {
+ Ok(hasher) => hasher,
+ Err(_) => return Err(Error::ThreadReceiveError),
+ };
+
+ chunk = next_chunk;
}
}
}
@@ -321,6 +368,71 @@ pub fn verify_hash(hash: &str, data: &[u8]) -> bool {
}
}
+// Fast implementation for Blake3 hashing that can handle large assets
+pub fn blake3_from_asset(path: &Path) -> Result<String> {
+ let mut data = File::open(path)?;
+ data.seek(SeekFrom::Start(0))?;
+ let data_len = data.seek(SeekFrom::End(0))?;
+ data.seek(SeekFrom::Start(0))?;
+
+ let mut hasher = blake3::Hasher::new();
+
+ let mut chunk_left = data_len;
+
+ if cfg!(feature = "no_interleaved_io") {
+ loop {
+ let mut chunk = vec![0u8; std::cmp::min(chunk_left as usize, MAX_HASH_BUF)];
+
+ data.read_exact(&mut chunk)?;
+
+ hasher.update(&chunk);
+
+ chunk_left -= chunk.len() as u64;
+ if chunk_left == 0 {
+ break;
+ }
+ }
+ } else {
+ let mut chunk = vec![0u8; std::cmp::min(chunk_left as usize, MAX_HASH_BUF)];
+ data.read_exact(&mut chunk)?;
+
+ loop {
+ let (tx, rx) = std::sync::mpsc::channel();
+
+ chunk_left -= chunk.len() as u64;
+
+ std::thread::spawn(move || {
+ hasher.update(&chunk);
+ tx.send(hasher).unwrap_or_default();
+ });
+
+ // are we done
+ if chunk_left == 0 {
+ hasher = match rx.recv() {
+ Ok(hasher) => hasher,
+ Err(_) => return Err(Error::ThreadReceiveError),
+ };
+ break;
+ }
+
+ // read next chunk while we wait for hash
+ let mut next_chunk = vec![0u8; std::cmp::min(chunk_left as usize, MAX_HASH_BUF)];
+ data.read_exact(&mut next_chunk)?;
+
+ hasher = match rx.recv() {
+ Ok(hasher) => hasher,
+ Err(_) => return Err(Error::ThreadReceiveError),
+ };
+
+ chunk = next_chunk;
+ }
+ }
+
+ let hash = hasher.finalize();
+
+ Ok(hash.to_hex().as_str().to_owned())
+}
+
/// Return the hash of data in the same hash format in_hash
pub fn hash_as_source(in_hash: &str, data: &[u8]) -> Option<String> {
match decode(in_hash) {
