tools/blackout/blackout-target/src/static_tree.rs - fuchsia - Git at Google

 // Copyright 2019 The Fuchsia Authors. All rights reserved.
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.

 //! Generate a random directory tree that can be commited to disk all at once. The general formula
 //! goes like this -
 //!
 //! ```
 //! use {rand::{Rng, SeedableRng, rngs::StdRng}, static_tree::{EntryDistribution, DirectoryEntry}};
 //! let mut rng = StdRng::seed_from_u64(seed);
 //! let dist = EntryDistribution::new(depth);
 //! let tree: DirectoryEntry = rng.sample(&dist);
 //! tree.write_tree_at(root).expect("failed to write tree");
 //! ```

 use {
     anyhow::Error,
     rand::{distributions, Rng},
     std::{fmt, fs::File, io::Write},
 };

 /// A random distribution specialized to generation of random directory trees. This distribution
 /// decreases the likelyhood of a directory being generated linearly relative to the depth of the
 /// subset of the tree being generated, until it's 0% at the maximum depth.
 #[derive(Clone, Copy, Debug, Eq, PartialEq)]
 pub struct EntryDistribution {
     depth: u32,
     max_depth: u32,
 }

 impl EntryDistribution {
     /// Create a new `EntryDistribution` with a maximum depth of `max_depth`. This distribution is
     /// used for generating [`DirectoryEntry`]s and [`Entry`]s.
     pub fn new(max_depth: u32) -> EntryDistribution {
         EntryDistribution { depth: 1, max_depth }
     }

     /// get the entry distribution for the next level down on the directory tree.
     fn next_level(&self) -> EntryDistribution {
         debug_assert!(
             self.depth <= self.max_depth,
             "directory tree exceeded max depth ({} vs max of {}). programming error.",
             self.depth,
             self.max_depth
         );
         EntryDistribution { depth: self.depth + 1, max_depth: self.max_depth }
     }

     /// creates a bernoulli distribution where the likelyhood is progressively decreased at further
     /// depths until it's 0% at max_depth.
     fn directory_distribution(&self) -> distributions::Bernoulli {
         distributions::Bernoulli::from_ratio(self.max_depth - self.depth, self.max_depth)
     }
 }

 /// A file entry in the generated tree. Contains a randomly generated u64 for a name and randomly
 /// generated byte contents. The file size is random, in a range of 1 byte to 2^16 bytes (~65kb).
 #[derive(Clone, Eq, PartialEq)]
 pub struct FileEntry {
     name: u64,
     contents: Vec<u8>,
 }

 impl fmt::Debug for FileEntry {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         write!(f, "FileEntry {{ name: {:?} }}", self.name)
     }
 }

 impl distributions::Distribution<FileEntry> for distributions::Standard {
     fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> FileEntry {
         let size = rng.gen_range(1, 1 << 16);
         // unfortunately, we can't use sample_iter to generate the content here. the trait
         // definition for distribution requires the Rng type parameter to have ?Sized (or "maybe
         // sized"), but the sample_iter function requires the provided Rng be Sized, either through
         // the explicit Self: Sized on the Rng sample_iter function, or the implicit Sized present
         // on type parameters for the equivalent Distribution function.
         let mut contents = vec![0; size];
         rng.fill(contents.as_mut_slice());
         FileEntry { name: rng.gen(), contents }
     }
 }

 impl FileEntry {
     fn write_file_at(self, root: &str) -> Result<(), Error> {
         let file_path = format!("{}/{}", root, self.name);
         let mut file = File::create(&file_path)?;
         file.write_all(&self.contents)?;
         Ok(())
     }
 }

 /// A directory entry in the generated tree. Contains a randomly generated u64 for a name and a
 /// vector of randomly generated directory entries, with a number of entries somewhere in the range
 /// of [0, 6). The sample function generates the entire subtree depth-first before returning, and is
 /// mutually recursive with the [`Entry`] sample function.
 #[derive(Clone, Debug, Eq, PartialEq)]
 pub struct DirectoryEntry {
     name: u64,
     entries: Vec<Entry>,
 }

 impl distributions::Distribution<DirectoryEntry> for EntryDistribution {
     fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> DirectoryEntry {
         // each directory has a random number of entries in the range [0, 6)
         let num_entries = rng.gen_range(0, 6);
         let mut entries = vec![];
         let entry_dist = self.next_level();
         for _ in 0..num_entries {
             entries.push(rng.sample(&entry_dist));
         }
         DirectoryEntry { name: rng.gen(), entries }
     }
 }

 impl DirectoryEntry {
     /// get the path the directory entry will be written to relative to a given root.
     pub fn get_name(&self) -> String {
         self.name.to_string()
     }

     /// Take the entire randomly generated tree and commit it to disk.
     pub fn write_tree_at(self, root: &str) -> Result<(), Error> {
         let path = format!("{}/{}", root, self.get_name());
         std::fs::create_dir_all(&path)?;
         for entry in self.entries {
             match entry {
                 Entry::File(file_entry) => file_entry.write_file_at(&path)?,
                 Entry::Directory(dir_entry) => dir_entry.write_tree_at(&path)?,
             }
         }
         Ok(())
     }
 }

 /// An entry of either File or Directory. The chance of it being a directory is relative to the depth
 /// recorded in EntryDistribution. The associated entry is also then generated.
 #[derive(Clone, Debug, Eq, PartialEq)]
 pub enum Entry {
     /// This entry is a file.
     File(FileEntry),
     /// This entry is a directory
     Directory(DirectoryEntry),
 }

 impl distributions::Distribution<Entry> for EntryDistribution {
     fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> Entry {
         if rng.sample(self.directory_distribution()) {
             Entry::Directory(rng.sample(self))
         } else {
             Entry::File(rng.sample(distributions::Standard))
         }
     }
 }

 #[cfg(test)]
 mod tests {
     use {
         super::{DirectoryEntry, Entry, EntryDistribution, FileEntry},
         fs_management::Minfs,
         ramdevice_client::RamdiskClient,
         rand::{rngs::StdRng, Rng, SeedableRng},
     };

     // this fixture will have to get updated any time the generation logic changes. depending on the
     // nature of the change, the comparison logic may have to be changed as well.
     fn get_fixture(seed: u64, depth: u32) -> DirectoryEntry {
         // confirm that the seed and depth used to generate this tree are the same.
         assert_eq!(seed, 1234, "seed value changed - update fixture");
         assert_eq!(depth, 2, "depth value changed - update fixture");
         DirectoryEntry {
             name: 6998163126151612998,
             entries: vec![
                 Entry::File(FileEntry { name: 15265247264153552834, contents: vec![] }),
                 Entry::File(FileEntry { name: 8405537650559459291, contents: vec![] }),
                 Entry::File(FileEntry { name: 12776892210595361501, contents: vec![] }),
             ],
         }
     }

     #[test]
     fn gen_tree() {
         let seed = 1234;
         let depth = 2;
         // make sure we are generating a tree at all.
         let mut rng = StdRng::seed_from_u64(seed);
         let dist = EntryDistribution::new(depth);
         let tree: DirectoryEntry = rng.sample(dist);
         println!("generated tree: {:?}", tree);
         // this skips the contents for the files, because they are huge, so we do our own manual
         // comparison.
         let fixture = get_fixture(seed, depth);
         assert_eq!(fixture.name, tree.name);
         assert_eq!(fixture.entries.len(), tree.entries.len());
         for i in 0..fixture.entries.len() {
             match &fixture.entries[i] {
                 Entry::File(fixture_file_entry) => match &tree.entries[i] {
                     Entry::File(tree_file_entry) => {
                         assert_eq!(fixture_file_entry.name, tree_file_entry.name)
                     }
                     _ => panic!("expected a file in generated tree"),
                 },
                 _ => panic!("expected a file in fixture tree"),
             }
         }
     }

     #[test]
     fn same_seed_same_tree() {
         // this confirms an assumption about the properties of tree generation that we rely on for
         // consistency checking.
         let seed = 1234;
         let gen_tree: fn(u64) -> DirectoryEntry = |seed| {
             let mut rng = StdRng::seed_from_u64(seed);
             let dist = EntryDistribution::new(5);
             rng.sample(dist)
         };

         let tree1 = gen_tree(seed);
         let tree2 = gen_tree(seed);

         assert_eq!(tree1, tree2);
     }

     #[test]
     fn write_tree() {
         let root = "/test-root";
         let seed = 1234;
         let depth = 2;

         let mut rng = StdRng::seed_from_u64(seed);
         let dist = EntryDistribution::new(depth);
         let tree: DirectoryEntry = rng.sample(dist);

         let ramdisk = RamdiskClient::create(512, 1 << 16).expect("failed to make ramdisk");
         let device_path = ramdisk.get_path();

         let mut minfs = Minfs::new(device_path).expect("failed to make new minfs");
         minfs.format().expect("failed to format minfs");
         minfs.mount(root).expect("failed to mount minfs");

         tree.write_tree_at(root).expect("failed to write tree");

         let fixture = get_fixture(seed, depth);
         let path = std::path::PathBuf::from(format!("{}/{}", root, fixture.name));
         assert!(path.is_dir());
         for (i, entry) in std::fs::read_dir(&path).expect("failed to read directory").enumerate() {
             let entry = entry.expect("failed to read entry");
             let file_type = entry.file_type().expect("failed to get file type");
             assert!(file_type.is_file());
             let file_name =
                 entry.file_name().into_string().expect("failed to convert file name to string");
             let expected_name = match &fixture.entries[i] {
                 Entry::File(f) => f.name,
                 Entry::Directory(_) => panic!("expected a file in the fixture tree"),
             };
             assert_eq!(file_name, expected_name.to_string());
         }

         minfs.unmount().expect("failed to unmount minfs");
         ramdisk.destroy().expect("failed to destroy ramdisk");
     }
 }
	// Copyright 2019 The Fuchsia Authors. All rights reserved.
	// Use of this source code is governed by a BSD-style license that can be
	// found in the LICENSE file.

	//! Generate a random directory tree that can be commited to disk all at once. The general formula
	//! goes like this -
	//!
	//! ```
	//! use {rand::{Rng, SeedableRng, rngs::StdRng}, static_tree::{EntryDistribution, DirectoryEntry}};
	//! let mut rng = StdRng::seed_from_u64(seed);
	//! let dist = EntryDistribution::new(depth);
	//! let tree: DirectoryEntry = rng.sample(&dist);
	//! tree.write_tree_at(root).expect("failed to write tree");
	//! ```

	use {
	anyhow::Error,
	rand::{distributions, Rng},
	std::{fmt, fs::File, io::Write},
	};

	/// A random distribution specialized to generation of random directory trees. This distribution
	/// decreases the likelyhood of a directory being generated linearly relative to the depth of the
	/// subset of the tree being generated, until it's 0% at the maximum depth.
	#[derive(Clone, Copy, Debug, Eq, PartialEq)]
	pub struct EntryDistribution {
	depth: u32,
	max_depth: u32,
	}

	impl EntryDistribution {
	/// Create a new `EntryDistribution` with a maximum depth of `max_depth`. This distribution is
	/// used for generating [`DirectoryEntry`]s and [`Entry`]s.
	pub fn new(max_depth: u32) -> EntryDistribution {
	EntryDistribution { depth: 1, max_depth }
	}

	/// get the entry distribution for the next level down on the directory tree.
	fn next_level(&self) -> EntryDistribution {
	debug_assert!(
	self.depth <= self.max_depth,
	"directory tree exceeded max depth ({} vs max of {}). programming error.",
	self.depth,
	self.max_depth
	);
	EntryDistribution { depth: self.depth + 1, max_depth: self.max_depth }
	}

	/// creates a bernoulli distribution where the likelyhood is progressively decreased at further
	/// depths until it's 0% at max_depth.
	fn directory_distribution(&self) -> distributions::Bernoulli {
	distributions::Bernoulli::from_ratio(self.max_depth - self.depth, self.max_depth)
	}
	}

	/// A file entry in the generated tree. Contains a randomly generated u64 for a name and randomly
	/// generated byte contents. The file size is random, in a range of 1 byte to 2^16 bytes (~65kb).
	#[derive(Clone, Eq, PartialEq)]
	pub struct FileEntry {
	name: u64,
	contents: Vec<u8>,
	}

	impl fmt::Debug for FileEntry {
	fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
	write!(f, "FileEntry {{ name: {:?} }}", self.name)
	}
	}

	impl distributions::Distribution<FileEntry> for distributions::Standard {
	fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> FileEntry {
	let size = rng.gen_range(1, 1 << 16);
	// unfortunately, we can't use sample_iter to generate the content here. the trait
	// definition for distribution requires the Rng type parameter to have ?Sized (or "maybe
	// sized"), but the sample_iter function requires the provided Rng be Sized, either through
	// the explicit Self: Sized on the Rng sample_iter function, or the implicit Sized present
	// on type parameters for the equivalent Distribution function.
	let mut contents = vec![0; size];
	rng.fill(contents.as_mut_slice());
	FileEntry { name: rng.gen(), contents }
	}
	}

	impl FileEntry {
	fn write_file_at(self, root: &str) -> Result<(), Error> {
	let file_path = format!("{}/{}", root, self.name);
	let mut file = File::create(&file_path)?;
	file.write_all(&self.contents)?;
	Ok(())
	}
	}

	/// A directory entry in the generated tree. Contains a randomly generated u64 for a name and a
	/// vector of randomly generated directory entries, with a number of entries somewhere in the range
	/// of [0, 6). The sample function generates the entire subtree depth-first before returning, and is
	/// mutually recursive with the [`Entry`] sample function.
	#[derive(Clone, Debug, Eq, PartialEq)]
	pub struct DirectoryEntry {
	name: u64,
	entries: Vec<Entry>,
	}

	impl distributions::Distribution<DirectoryEntry> for EntryDistribution {
	fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> DirectoryEntry {
	// each directory has a random number of entries in the range [0, 6)
	let num_entries = rng.gen_range(0, 6);
	let mut entries = vec![];
	let entry_dist = self.next_level();
	for _ in 0..num_entries {
	entries.push(rng.sample(&entry_dist));
	}
	DirectoryEntry { name: rng.gen(), entries }
	}
	}

	impl DirectoryEntry {
	/// get the path the directory entry will be written to relative to a given root.
	pub fn get_name(&self) -> String {
	self.name.to_string()
	}

	/// Take the entire randomly generated tree and commit it to disk.
	pub fn write_tree_at(self, root: &str) -> Result<(), Error> {
	let path = format!("{}/{}", root, self.get_name());
	std::fs::create_dir_all(&path)?;
	for entry in self.entries {
	match entry {
	Entry::File(file_entry) => file_entry.write_file_at(&path)?,
	Entry::Directory(dir_entry) => dir_entry.write_tree_at(&path)?,
	}
	}
	Ok(())
	}
	}

	/// An entry of either File or Directory. The chance of it being a directory is relative to the depth
	/// recorded in EntryDistribution. The associated entry is also then generated.
	#[derive(Clone, Debug, Eq, PartialEq)]
	pub enum Entry {
	/// This entry is a file.
	File(FileEntry),
	/// This entry is a directory
	Directory(DirectoryEntry),
	}

	impl distributions::Distribution<Entry> for EntryDistribution {
	fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> Entry {
	if rng.sample(self.directory_distribution()) {
	Entry::Directory(rng.sample(self))
	} else {
	Entry::File(rng.sample(distributions::Standard))
	}
	}
	}

	#[cfg(test)]
	mod tests {
	use {
	super::{DirectoryEntry, Entry, EntryDistribution, FileEntry},
	fs_management::Minfs,
	ramdevice_client::RamdiskClient,
	rand::{rngs::StdRng, Rng, SeedableRng},
	};

	// this fixture will have to get updated any time the generation logic changes. depending on the
	// nature of the change, the comparison logic may have to be changed as well.
	fn get_fixture(seed: u64, depth: u32) -> DirectoryEntry {
	// confirm that the seed and depth used to generate this tree are the same.
	assert_eq!(seed, 1234, "seed value changed - update fixture");
	assert_eq!(depth, 2, "depth value changed - update fixture");
	DirectoryEntry {
	name: 6998163126151612998,
	entries: vec![
	Entry::File(FileEntry { name: 15265247264153552834, contents: vec![] }),
	Entry::File(FileEntry { name: 8405537650559459291, contents: vec![] }),
	Entry::File(FileEntry { name: 12776892210595361501, contents: vec![] }),
	],
	}
	}

	#[test]
	fn gen_tree() {
	let seed = 1234;
	let depth = 2;
	// make sure we are generating a tree at all.
	let mut rng = StdRng::seed_from_u64(seed);
	let dist = EntryDistribution::new(depth);
	let tree: DirectoryEntry = rng.sample(dist);
	println!("generated tree: {:?}", tree);
	// this skips the contents for the files, because they are huge, so we do our own manual
	// comparison.
	let fixture = get_fixture(seed, depth);
	assert_eq!(fixture.name, tree.name);
	assert_eq!(fixture.entries.len(), tree.entries.len());
	for i in 0..fixture.entries.len() {
	match &fixture.entries[i] {
	Entry::File(fixture_file_entry) => match &tree.entries[i] {
	Entry::File(tree_file_entry) => {
	assert_eq!(fixture_file_entry.name, tree_file_entry.name)
	}
	_ => panic!("expected a file in generated tree"),
	},
	_ => panic!("expected a file in fixture tree"),
	}
	}
	}

	#[test]
	fn same_seed_same_tree() {
	// this confirms an assumption about the properties of tree generation that we rely on for
	// consistency checking.
	let seed = 1234;
	let gen_tree: fn(u64) -> DirectoryEntry = \|seed\| {
	let mut rng = StdRng::seed_from_u64(seed);
	let dist = EntryDistribution::new(5);
	rng.sample(dist)
	};

	let tree1 = gen_tree(seed);
	let tree2 = gen_tree(seed);

	assert_eq!(tree1, tree2);
	}

	#[test]
	fn write_tree() {
	let root = "/test-root";
	let seed = 1234;
	let depth = 2;

	let mut rng = StdRng::seed_from_u64(seed);
	let dist = EntryDistribution::new(depth);
	let tree: DirectoryEntry = rng.sample(dist);

	let ramdisk = RamdiskClient::create(512, 1 << 16).expect("failed to make ramdisk");
	let device_path = ramdisk.get_path();

	let mut minfs = Minfs::new(device_path).expect("failed to make new minfs");
	minfs.format().expect("failed to format minfs");
	minfs.mount(root).expect("failed to mount minfs");

	tree.write_tree_at(root).expect("failed to write tree");

	let fixture = get_fixture(seed, depth);
	let path = std::path::PathBuf::from(format!("{}/{}", root, fixture.name));
	assert!(path.is_dir());
	for (i, entry) in std::fs::read_dir(&path).expect("failed to read directory").enumerate() {
	let entry = entry.expect("failed to read entry");
	let file_type = entry.file_type().expect("failed to get file type");
	assert!(file_type.is_file());
	let file_name =
	entry.file_name().into_string().expect("failed to convert file name to string");
	let expected_name = match &fixture.entries[i] {
	Entry::File(f) => f.name,
	Entry::Directory(_) => panic!("expected a file in the fixture tree"),
	};
	assert_eq!(file_name, expected_name.to_string());
	}

	minfs.unmount().expect("failed to unmount minfs");
	ramdisk.destroy().expect("failed to destroy ramdisk");
	}
	}