Skip to content

Commit

Permalink
Merge pull request #830 from worldcoin/piohei/improve_logging
Browse files Browse the repository at this point in the history 
Improve logging.
  • Loading branch information
@piohei
piohei authored Nov 29, 2024
2 parents 2c57948 + d16de26 commit 51d7aa1
Show file tree
Hide file tree
Showing 2 changed files with 60 additions and 7 deletions.
45 changes: 40 additions & 5 deletions src/identity_tree/initializer.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -81,6 +81,7 @@ impl TreeInitializer {
initial_root_hash: Hash,
force_cache_purge: bool,
) -> anyhow::Result<TreeState> {
info!("Getting mined commitments from DB");
let mut mined_items = self
.database
.get_commitments_by_status(ProcessedStatus::Mined)
Expand All @@ -90,8 +91,13 @@ impl TreeInitializer {

let mined_items = dedup_tree_updates(mined_items);

info!("Retrieved {} mined commitments from DB", mined_items.len());

if !force_cache_purge {
info!("Attempting to restore tree from cache");
info!(
"Attempting to restore tree from cache (file={})",
&self.config.cache_file
);
if let Some(tree_state) = self
.get_cached_tree_state(&mined_items, initial_root_hash)
.await?
Expand All @@ -118,11 +124,13 @@ impl TreeInitializer {
vec![]
} else {
let max_leaf = mined_items.last().map(|item| item.leaf_index).unwrap();
// if the last index is greater then dense_prefix_depth, 1 << dense_prefix_depth
// if the last index is greater than dense_prefix_depth, 1 << dense_prefix_depth
// should be the last index in restored tree
let last_index = std::cmp::min(max_leaf, (1 << dense_prefix_depth) - 1);
*index = Some(last_index);

info!("calculated last dense tree index: {}", last_index);

if max_leaf - last_index == 0 {
return vec![];
}
Expand All @@ -138,6 +146,11 @@ impl TreeInitializer {
leaves
};

info!(
"number of leaves not included in the dense tree: {}",
leftover_items.len()
);

leftover_items
}

Expand All @@ -146,13 +159,17 @@ impl TreeInitializer {
mined_items: &[TreeUpdate],
initial_root_hash: Hash,
) -> anyhow::Result<Option<TreeState>> {
info!("Restoring tree from cache");

let mut last_mined_index_in_dense: Option<usize> = None;
let leftover_items = Self::get_leftover_leaves_and_update_index(
&mut last_mined_index_in_dense,
self.config.dense_tree_prefix_depth,
mined_items,
);

info!("Restoring canonical mined tree");

let Some(mined_builder) = CanonicalTreeBuilder::restore(
self.config.tree_depth,
self.config.dense_tree_prefix_depth,
Expand All @@ -174,16 +191,20 @@ impl TreeInitializer {
{
Some(root) => {
if !mined.get_root().eq(&root) {
info!("Cached tree root differs from root in database");
return Ok(None);
}
}
None => {
if !mined.get_root().eq(&initial_root_hash) {
info!("Cached tree root differs from init root");
return Ok(None);
}
}
}

info!("Restoring derived processed and batching tree");

let processed_items = self
.database
.get_commitments_by_status(ProcessedStatus::Processed)
Expand All @@ -196,6 +217,8 @@ impl TreeInitializer {
let (processed, batching_builder) = processed_builder.seal_and_continue();
let (batching, mut latest_builder) = batching_builder.seal_and_continue();

info!("Restoring derived latest tree");

let pending_items = self
.database
.get_commitments_by_status(ProcessedStatus::Pending)
Expand All @@ -205,6 +228,8 @@ impl TreeInitializer {
}
let latest = latest_builder.seal();

info!("Batching tree correction");

let batch = self.database.get_latest_batch().await?;
if let Some(batch) = batch {
if batching.get_root() != batch.next_root {
Expand All @@ -213,11 +238,15 @@ impl TreeInitializer {
assert_eq!(batching.get_root(), batch.next_root);
}

info!("Tree restored.");

Ok(Some(TreeState::new(mined, processed, batching, latest)))
}

#[instrument(skip_all)]
async fn initialize_tree(&self, mined_items: Vec<TreeUpdate>) -> anyhow::Result<TreeState> {
info!("Creating tree from the database");

let initial_leaf_value = self.config.initial_leaf_value;

let initial_leaves = if mined_items.is_empty() {
Expand All @@ -233,12 +262,13 @@ impl TreeInitializer {
leaves
};

info!("Creating mined tree");
let tree_depth = self.config.tree_depth;
let dense_tree_prefix_depth = self.config.dense_tree_prefix_depth;
let tree_gc_threshold = self.config.tree_gc_threshold;
let cache_file = self.config.cache_file.clone();

info!("Creating canonical mined tree");

let mined_builder = tokio::task::spawn_blocking(move || {
CanonicalTreeBuilder::new(
tree_depth,
Expand All @@ -253,12 +283,13 @@ impl TreeInitializer {

let (mined, mut processed_builder) = mined_builder.seal();

info!("Creating derived processed and batching tree");

let processed_items = self
.database
.get_commitments_by_status(ProcessedStatus::Processed)
.await?;

info!("Updating processed tree");
let processed_builder = tokio::task::spawn_blocking(move || {
for processed_item in processed_items {
processed_builder.update(&processed_item);
Expand All @@ -271,12 +302,13 @@ impl TreeInitializer {
let (processed, batching_builder) = processed_builder.seal_and_continue();
let (batching, mut latest_builder) = batching_builder.seal_and_continue();

info!("Creating derived latest tree");

let pending_items = self
.database
.get_commitments_by_status(ProcessedStatus::Pending)
.await?;

info!("Updating latest tree");
let latest_builder = tokio::task::spawn_blocking(move || {
for update in pending_items {
latest_builder.update(&update);
Expand All @@ -288,6 +320,7 @@ impl TreeInitializer {

let latest = latest_builder.seal();

info!("Batching tree correction");
let batch = self.database.get_latest_batch().await?;
if let Some(batch) = batch {
if batching.get_root() != batch.next_root {
Expand All @@ -296,6 +329,8 @@ impl TreeInitializer {
assert_eq!(batching.get_root(), batch.next_root);
}

info!("Tree created.");

Ok(TreeState::new(mined, processed, batching, latest))
}
}
Expand Down
22 changes: 20 additions & 2 deletions src/identity_tree/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -644,6 +644,7 @@ impl CanonicalTreeBuilder {
let (initial_leaves_in_dense, leftover_initial_leaves) =
initial_leaves.split_at(initial_leaves_in_dense_count);

info!("Creating mmap dense tree");
let tree =
PoseidonTree::<lazy_merkle_tree::Canonical>::new_mmapped_with_dense_prefix_with_init_values(
tree_depth,
Expand All @@ -652,6 +653,8 @@ impl CanonicalTreeBuilder {
initial_leaves_in_dense,
mmap_file_path
).unwrap();

info!("Applying leaves not in dense tree");
let metadata = CanonicalTreeMetadata {
flatten_threshold: flattening_threshold,
count_since_last_flatten: 0,
Expand All @@ -662,12 +665,19 @@ impl CanonicalTreeBuilder {
metadata,
next: None,
});
let last_leaf_index = leftover_initial_leaves.len() + initial_leaves_in_dense_count;
for (index, leaf) in leftover_initial_leaves.iter().enumerate() {
let leaf_index = index + initial_leaves_in_dense_count;
if leaf_index % 10000 == 0 {
info!("Current index {leaf_index}/{last_leaf_index}");
}
builder.update(&TreeUpdate {
leaf_index: index + initial_leaves_in_dense_count,
leaf_index,
element: *leaf,
});
}

info!("Tree created");
builder
}

Expand All @@ -680,6 +690,7 @@ impl CanonicalTreeBuilder {
flattening_threshold: usize,
mmap_file_path: &str,
) -> Option<Self> {
info!("Restoring tree from file");
let tree: LazyMerkleTree<PoseidonHash, lazy_merkle_tree::Canonical> =
match PoseidonTree::<lazy_merkle_tree::Canonical>::attempt_dense_mmap_restore(
tree_depth,
Expand All @@ -694,6 +705,7 @@ impl CanonicalTreeBuilder {
}
};

info!("Applying leaves not in dense tree");
let metadata = CanonicalTreeMetadata {
flatten_threshold: flattening_threshold,
count_since_last_flatten: 0,
Expand All @@ -706,13 +718,19 @@ impl CanonicalTreeBuilder {
next: None,
});

let last_leaf_index = leftover_items.len() + next_leaf;
for (index, leaf) in leftover_items.iter().enumerate() {
let leaf_index = index + next_leaf;
if index % 10000 == 0 {
info!("Current index {leaf_index}/{last_leaf_index}");
}
builder.update(&TreeUpdate {
leaf_index: next_leaf + index,
leaf_index,
element: *leaf,
});
}

info!("Tree restored");
Some(builder)
}

Expand Down

0 comments on commit 51d7aa1

Please sign in to comment.