tag trunk and branch after traversing

Alex6323 · Sep 16, 2020 · 6633eac · 6633eac
1 parent f55cf94
commit 6633eac
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 4 deletions.
diff --git a/bee-protocol/src/worker/milestone_validator.rs b/bee-protocol/src/worker/milestone_validator.rs
@@ -126,10 +126,6 @@ where
                         if let Some(_) = Protocol::get().requested_milestones.remove(&milestone.index) {
                             tangle().update_metadata(&milestone.hash, |meta| meta.flags.set_requested(true));
 
-                            let tx = tangle().get(&milestone.hash).unwrap();
-                            tangle().update_metadata(tx.trunk(), |meta| meta.flags.set_requested(true));
-                            tangle().update_metadata(tx.branch(), |meta| meta.flags.set_requested(true));
-
                             Protocol::trigger_milestone_solidification(milestone.index);
                         }
                     }

diff --git a/bee-protocol/src/worker/solidifier/milestone.rs b/bee-protocol/src/worker/solidifier/milestone.rs
@@ -12,6 +12,7 @@
 use crate::{milestone::MilestoneIndex, protocol::Protocol, tangle::tangle};
 
 use bee_common::{shutdown_stream::ShutdownStream, worker::Error as WorkerError};
+use bee_transaction::Vertex;
 use bee_tangle::traversal;
 
 use futures::{
@@ -57,6 +58,10 @@ impl MilestoneSolidifierWorker {
                     |missing_hash| Protocol::request_transaction(*missing_hash, target_index),
                 );
 
+                let tx = tangle().get(&target_hash).unwrap();
+                tangle().update_metadata(tx.trunk(), |meta| meta.flags.set_requested(true));
+                tangle().update_metadata(tx.branch(), |meta| meta.flags.set_requested(true));
+
                 self.next_ms_index = target_index + MilestoneIndex(1);
             }
         }