Remove extra data allocation used for MT generation (#837)

Make (a small) part of the replication the replication process sequential to avoid creating an extra `mmap` the size of a sector size. Decrease memory usage by 1x. Increase replication time (for parallel setups which will waste part of that parallelism) by a small fraction (encoding time still dominates).

Author: schomatis

README.md

@@ -176,15 +176,7 @@ cargo build                                                                   \
     --size 1048576
 ```
 
-To optimize even more for memory there's another option (used in addition to the `disk-trees` feature) to generate all MTs in sequential order, to make sure we can offload them to disk before we start buildding the next one,
-
-```
-FIL_PROOFS_GENERATE_MERKLE_TREES_IN_PARALLEL=0
-```
-
-You can inspect that it's working also in the replication log, where you'll see that the MTs are all generated in order without any layer index out of place.
-
-All these optimizations (`disk-trees` with a directory to offload MTs plus sequential generation) should reduce the maximum RSS, in the `zigzag` example, to 3 times the sector size used (so in the above command that tested ZigZag with a 1 GiB sector the maximum RSS reported by commands like `/usr/bin/time -v` should not exceed 3 GiB, please submit an issue if you observe otherwise).
+This optimization should reduce the maximum RSS, in the `zigzag` example, to 1-2 times the sector size used (so in the above command that tested ZigZag with a 1 GiB sector the maximum RSS reported by commands like `/usr/bin/time -v` should not exceed 2 GiB, please submit an issue if you observe otherwise).
 
 ## Generate Documentation
 

storage-proofs/Cargo.toml

@@ -15,7 +15,7 @@ bench = false
 bitvec = "0.5"
 rand = "0.4"
 libc = "0.2"
-merkletree = "=0.9"
+merkletree = "=0.10"
 failure = "0.1"
 byteorder = "1"
 config = "0.9.3"

storage-proofs/src/drgraph.rs

@@ -7,7 +7,7 @@ use rayon::prelude::*;
 use crate::error::*;
 use crate::hasher::pedersen::PedersenHasher;
 use crate::hasher::{Domain, Hasher};
-use crate::merkle::MerkleTree;
+use crate::merkle::{MerkleStore, MerkleTree};
 use crate::parameter_cache::ParameterSetMetadata;
 use crate::util::{data_at_node, NODE_SIZE};
 use merkletree::merkle::FromIndexedParallelIterator;
@@ -62,6 +62,16 @@ pub trait Graph<H: Hasher>: ::std::fmt::Debug + Clone + PartialEq + Eq {
         }
     }
 
+    /// Builds a merkle tree based on the given leaves store.
+    // FIXME: Add the parallel case (if it's worth it).
+    fn merkle_tree_from_leaves(
+        &self,
+        leaves: MerkleStore<H::Domain>,
+        leafs_number: usize,
+    ) -> Result<MerkleTree<H::Domain, H::Function>> {
+        Ok(MerkleTree::from_leaves_store(leaves, leafs_number))
+    }
+
     /// Returns the merkle tree depth.
     fn merkle_tree_depth(&self) -> u64 {
         graph_height(self.size()) as u64

storage-proofs/src/hasher/pedersen.rs

@@ -319,7 +319,7 @@ mod tests {
 
     use merkletree::hash::Hashable;
 
-    use crate::merkle::{MerkleTree, VecMerkleTree};
+    use crate::merkle::MerkleTree;
 
     #[test]
     fn test_path() {
@@ -335,8 +335,7 @@ mod tests {
     fn test_pedersen_hasher() {
         let values = ["hello", "world", "you", "two"];
 
-        let t = VecMerkleTree::<PedersenDomain, PedersenFunction>::from_data(values.iter());
-        // Using `VecMerkleTree` since the `MmapStore` of `MerkleTree` doesn't support `Deref` (`as_slice`).
+        let t = MerkleTree::<PedersenDomain, PedersenFunction>::from_data(values.iter());
 
         assert_eq!(t.leafs(), 4);
 

storage-proofs/src/hasher/sha256.rs

@@ -17,7 +17,7 @@ mod tests {
     use std::fmt;
     use std::iter::FromIterator;
 
-    use crate::merkle::VecMerkleTree;
+    use crate::merkle::MerkleTree;
     use merkletree::hash::{Algorithm, Hashable};
 
     use super::super::{DigestDomain, Hasher};
@@ -109,8 +109,8 @@ mod tests {
 
         let v: Vec<DigestDomain> = vec![h1.into(), h2.into(), h3.into()];
         let v2: Vec<DigestDomain> = vec![h1.into(), h2.into()];
-        let t = VecMerkleTree::<<Sha256Hasher as Hasher>::Domain, <Sha256Hasher as Hasher>::Function>::from_iter(v);
-        let t2 = VecMerkleTree::<
+        let t = MerkleTree::<<Sha256Hasher as Hasher>::Domain, <Sha256Hasher as Hasher>::Function>::from_iter(v);
+        let t2 = MerkleTree::<
             <Sha256Hasher as Hasher>::Domain,
             <Sha256Hasher as Hasher>::Function,
         >::from_iter(v2);

storage-proofs/src/layered_drgporep.rs

@@ -3,8 +3,6 @@ use std::marker::PhantomData;
 use std::sync::mpsc::channel;
 
 use crossbeam_utils::thread;
-use memmap::MmapMut;
-use memmap::MmapOptions;
 use rayon::prelude::*;
 use serde::de::Deserialize;
 use serde::ser::Serialize;
@@ -14,22 +12,15 @@ use crate::drgporep::{self, DrgPoRep};
 use crate::drgraph::Graph;
 use crate::error::{Error, Result};
 use crate::hasher::{Domain, HashFunction, Hasher};
-use crate::merkle::MerkleTree;
+use crate::merkle::{next_pow2, populate_leaves, MerkleStore, MerkleTree, Store};
 use crate::parameter_cache::ParameterSetMetadata;
 use crate::porep::{self, PoRep};
 use crate::proof::ProofScheme;
-use crate::settings;
+use crate::util::{data_at_node, NODE_SIZE};
 use crate::vde;
 
 type Tree<H> = MerkleTree<<H as Hasher>::Domain, <H as Hasher>::Function>;
 
-fn anonymous_mmap(len: usize) -> MmapMut {
-    MmapOptions::new()
-        .len(len)
-        .map_anon()
-        .expect("Failed to create memory map")
-}
-
 #[derive(Debug, Clone, Serialize)]
 pub enum LayerChallenges {
     Fixed {
@@ -372,140 +363,110 @@ pub trait Layers {
         let mut taus = Vec::with_capacity(layers);
         let mut auxs: Vec<Tree<Self::Hasher>> = Vec::with_capacity(layers);
 
-        if !&settings::SETTINGS
-            .lock()
-            .unwrap()
-            .generate_merkle_trees_in_parallel
-        {
-            let mut sorted_trees: Vec<_> = Vec::new();
-
-            (0..=layers).fold(graph.clone(), |current_graph, layer| {
-                let tree_d = current_graph.merkle_tree(&data).unwrap();
-
-                info!("returning tree (layer: {})", layer);
-
-                sorted_trees.push(tree_d);
-
-                if layer < layers {
-                    info!("encoding (layer: {})", layer);
-                    vde::encode(&current_graph, replica_id, data)
-                        .expect("encoding failed in thread");
-                }
-
-                Self::transform(&current_graph)
-            });
-
-            sorted_trees
-                .into_iter()
-                .fold(None, |previous_comm_r: Option<_>, replica_tree| {
-                    let comm_r = replica_tree.root();
-                    // Each iteration's replica_tree becomes the next iteration's previous_tree (data_tree).
-                    // The first iteration has no previous_tree.
-                    if let Some(comm_d) = previous_comm_r {
-                        let tau = porep::Tau { comm_r, comm_d };
-                        // info!("setting tau/aux (layer: {})", i - 1);
-                        // FIXME: Use `enumerate` if this log is worth it.
-                        taus.push(tau);
-                    };
-
-                    auxs.push(replica_tree);
-
-                    Some(comm_r)
-                });
-        } else {
-            // The parallel case is more complicated but should produce the same results as the
-            // serial case. Note that to make lifetimes work out, we have to inline and tease apart
-            // the definition of DrgPoRep::replicate. This is because as implemented, it entangles
-            // encoding and merkle tree generation too tightly to be used as a subcomponent.
-            // Instead, we need to create a scope which encloses all the work, spawning threads
-            // for merkle tree generation and sending the results back to a channel.
-            // The received results need to be sorted by layer because ordering of the completed results
-            // is not guaranteed. Misordered results will be seen in practice when trees are small.
-
-            // The outer scope ensure that `tx` is dropped and closed before we read from `outer_rx`.
-            // Otherwise, the read loop will block forever waiting for more input.
-            let outer_rx = {
-                let (tx, rx) = channel();
-
-                let errf = |e| {
-                    let err_string = format!("{:?}", e);
-                    error!(
-                        "MerkleTreeGenerationError: {} - {:?}",
-                        &err_string,
-                        failure::Backtrace::new()
-                    );
-                    Error::MerkleTreeGenerationError(err_string)
-                };
-
-                let _ = thread::scope(|scope| -> Result<()> {
-                    let mut threads = Vec::with_capacity(layers + 1);
-                    (0..=layers).fold(graph.clone(), |current_graph, layer| {
-                        let mut data_copy = anonymous_mmap(data.len());
-                        data_copy[0..data.len()].clone_from_slice(data);
-
-                        let return_channel = tx.clone();
-                        let (transfer_tx, transfer_rx) = channel::<Self::Graph>();
+        // We need to create a scope which encloses all the work, spawning threads
+        // for merkle tree generation and sending the results back to a channel.
+        // The received results need to be sorted by layer because ordering of the completed results
+        // is not guaranteed. Misordered results will be seen in practice when trees are small.
+
+        // The outer scope ensure that `tx` is dropped and closed before we read from `outer_rx`.
+        // Otherwise, the read loop will block forever waiting for more input.
+        let outer_rx = {
+            let (tx, rx) = channel();
+
+            let errf = |e| {
+                let err_string = format!("{:?}", e);
+                error!(
+                    "MerkleTreeGenerationError: {} - {:?}",
+                    &err_string,
+                    failure::Backtrace::new()
+                );
+                Error::MerkleTreeGenerationError(err_string)
+            };
 
-                        transfer_tx
-                            .send(current_graph.clone())
+            let _ = thread::scope(|scope| -> Result<()> {
+                let mut threads = Vec::with_capacity(layers + 1);
+                (0..=layers).fold(graph.clone(), |current_graph, layer| {
+                    let leafs = data.len() / NODE_SIZE;
+                    assert_eq!(data.len() % NODE_SIZE, 0);
+                    let pow = next_pow2(leafs);
+                    // FIXME: Who's actually responsible for ensuring power of 2
+                    //  sector sizes?
+
+                    let mut leaves_store = MerkleStore::new(pow);
+
+                    populate_leaves::<
+                        _,
+                        <Self::Hasher as Hasher>::Function,
+                        _,
+                        std::iter::Map<_, _>,
+                    >(&mut leaves_store, (0..leafs).map(|i| {
+                        let d = data_at_node(&data, i).expect("data_at_node math failed");
+                        <Self::Hasher as Hasher>::Domain::try_from_bytes(d)
+                            .expect("failed to convert node data to domain element")
+                    }));
+                    let return_channel = tx.clone();
+                    let (transfer_tx, transfer_rx) = channel::<Self::Graph>();
+
+                    transfer_tx
+                        .send(current_graph.clone())
+                        .expect("Failed to send value through channel");
+
+                    let thread = scope.spawn(move |_| {
+                        // If we panic anywhere in this closure, thread.join() below will receive an error —
+                        // so it is safe to unwrap.
+                        let graph = transfer_rx
+                            .recv()
+                            .expect("Failed to receive value through channel");
+
+                        let tree_d =
+                            graph.merkle_tree_from_leaves(leaves_store, leafs).unwrap();
+
+                        info!("returning tree (layer: {})", layer);
+                        return_channel
+                            .send((layer, tree_d))
                             .expect("Failed to send value through channel");
-
-                        let thread = scope.spawn(move |_| {
-                            // If we panic anywhere in this closure, thread.join() below will receive an error —
-                            // so it is safe to unwrap.
-                            let graph = transfer_rx
-                                .recv()
-                                .expect("Failed to receive value through channel");
-
-                            let tree_d = graph.merkle_tree(&data_copy).unwrap();
-
-                            info!("returning tree (layer: {})", layer);
-                            return_channel
-                                .send((layer, tree_d))
-                                .expect("Failed to send value through channel");
-                        });
-
-                        threads.push(thread);
-
-                        if layer < layers {
-                            info!("encoding (layer: {})", layer);
-                            vde::encode(&current_graph, replica_id, data)
-                                .expect("encoding failed in thread");
-                        }
-                        Self::transform(&current_graph)
                     });
 
-                    for thread in threads {
-                        thread.join().map_err(errf)?;
+                    threads.push(thread);
+
+                    if layer < layers {
+                        info!("encoding (layer: {})", layer);
+                        vde::encode(&current_graph, replica_id, data)
+                            .expect("encoding failed in thread");
                     }
+                    Self::transform(&current_graph)
+                });
 
-                    Ok(())
-                })
-                .map_err(errf)?;
+                for thread in threads {
+                    thread.join().map_err(errf)?;
+                }
 
-                rx
-            };
+                Ok(())
+            })
+            .map_err(errf)?;
 
-            let mut sorted_trees = outer_rx.iter().collect::<Vec<_>>();
-            sorted_trees.sort_unstable_by_key(|k| k.0);
+            rx
+        };
 
-            sorted_trees
-                .into_iter()
-                .fold(None, |previous_comm_r: Option<_>, (i, replica_tree)| {
-                    let comm_r = replica_tree.root();
-                    // Each iteration's replica_tree becomes the next iteration's previous_tree (data_tree).
-                    // The first iteration has no previous_tree.
-                    if let Some(comm_d) = previous_comm_r {
-                        let tau = porep::Tau { comm_r, comm_d };
-                        info!("setting tau/aux (layer: {})", i - 1);
-                        taus.push(tau);
-                    };
+        let mut sorted_trees = outer_rx.iter().collect::<Vec<_>>();
+        sorted_trees.sort_unstable_by_key(|k| k.0);
 
-                    auxs.push(replica_tree);
+        sorted_trees
+            .into_iter()
+            .fold(None, |previous_comm_r: Option<_>, (i, replica_tree)| {
+                let comm_r = replica_tree.root();
+                // Each iteration's replica_tree becomes the next iteration's previous_tree (data_tree).
+                // The first iteration has no previous_tree.
+                if let Some(comm_d) = previous_comm_r {
+                    let tau = porep::Tau { comm_r, comm_d };
+                    info!("setting tau/aux (layer: {})", i - 1);
+                    taus.push(tau);
+                };
 
-                    Some(comm_r)
-                });
-        }
+                auxs.push(replica_tree);
+
+                Some(comm_r)
+            });
 
         Ok((taus, auxs))
     }

storage-proofs/src/merkle.rs

@@ -5,23 +5,22 @@ use std::marker::PhantomData;
 // Reexport here, so we don't depend on merkletree directly in other places.
 use merkletree::hash::Algorithm;
 use merkletree::merkle;
-#[cfg(not(feature = "disk-trees"))]
-use merkletree::merkle::MmapStore;
-use merkletree::merkle::VecStore;
 use merkletree::proof;
 use paired::bls12_381::Fr;
 
 use crate::hasher::{Domain, Hasher};
 
-// `mmap`ed `MerkleTree` (replacing the previously `Vec`-backed
-// `MerkleTree`, now encapsulated in `merkle::VecStore` and exposed
-// as `VecMerkleTree`).
+pub use merkletree::merkle::{next_pow2, populate_leaves, Store};
+pub type MmapStore<E> = merkletree::merkle::MmapStore<E>;
 pub type DiskStore<E> = merkletree::merkle::DiskStore<E>;
-pub type VecMerkleTree<T, A> = merkle::MerkleTree<T, A, VecStore<T>>;
 #[cfg(feature = "disk-trees")]
 pub type MerkleTree<T, A> = merkle::MerkleTree<T, A, DiskStore<T>>;
 #[cfg(not(feature = "disk-trees"))]
 pub type MerkleTree<T, A> = merkle::MerkleTree<T, A, MmapStore<T>>;
+#[cfg(feature = "disk-trees")]
+pub type MerkleStore<T> = DiskStore<T>;
+#[cfg(not(feature = "disk-trees"))]
+pub type MerkleStore<T> = MmapStore<T>;
 
 /// Representation of a merkle proof.
 /// Each element in the `path` vector consists of a tuple `(hash, is_right)`, with `hash` being the the hash of the node at the current level and `is_right` a boolean indicating if the path is taking the right path.