From d55ebe8a2b3266d3d7a22b4c1c1a1e11cad7d519 Mon Sep 17 00:00:00 2001
From: Gary Yu <gairy.yu@gmail.com>
Date: Thu, 5 Sep 2019 02:51:28 +0800
Subject: [PATCH] fix: the race condition in compact (#3015)

---
 chain/src/chain.rs               |  6 ++--
 chain/src/txhashset/txhashset.rs | 54 ++++++++++++++++++++++++++++++++
 2 files changed, 56 insertions(+), 4 deletions(-)

diff --git a/chain/src/chain.rs b/chain/src/chain.rs
index 256a2ca38e..2a4cc21dd8 100644
--- a/chain/src/chain.rs
+++ b/chain/src/chain.rs
@@ -1077,7 +1077,7 @@ impl Chain {
 
 			// Rebuild our output_pos index in the db based on current UTXO set.
 			txhashset::extending(&mut txhashset, &mut batch, |extension| {
-				extension.rebuild_index()?;
+				extension.rebuild_height_pos_index()?;
 				Ok(())
 			})?;
 
@@ -1090,8 +1090,6 @@ impl Chain {
 			batch.commit()?;
 		}
 
-		self.rebuild_height_for_pos()?;
-
 		Ok(())
 	}
 
@@ -1223,7 +1221,7 @@ impl Chain {
 	/// Migrate the index 'commitment -> output_pos' to index 'commitment -> (output_pos, block_height)'
 	/// Note: should only be called in two cases:
 	///     - Node start-up. For database migration from the old version.
-	/// 	- After the txhashset 'rebuild_index' when state syncing or compact.
+	/// 	- After the txhashset 'rebuild_index' when state syncing.
 	pub fn rebuild_height_for_pos(&self) -> Result<(), Error> {
 		let txhashset = self.txhashset.read();
 		let mut outputs_pos = txhashset.get_all_output_pos()?;
diff --git a/chain/src/txhashset/txhashset.rs b/chain/src/txhashset/txhashset.rs
index 96da1f55e1..3ffa9f3730 100644
--- a/chain/src/txhashset/txhashset.rs
+++ b/chain/src/txhashset/txhashset.rs
@@ -1311,6 +1311,9 @@ impl<'a> Extension<'a> {
 
 	/// Rebuild the index of MMR positions to the corresponding UTXOs.
 	/// This is a costly operation performed only when we receive a full new chain state.
+	///
+	/// Note: only called by txhashset_write, and should be replaced by 'rebuild_height_pos_index'
+	/// in the future, after a refactoring of 'txhashset_write'.
 	pub fn rebuild_index(&self) -> Result<(), Error> {
 		let now = Instant::now();
 
@@ -1333,6 +1336,57 @@ impl<'a> Extension<'a> {
 		Ok(())
 	}
 
+	/// Rebuild the index of block height & MMR positions to the corresponding UTXOs.
+	/// This is a costly operation performed only when we receive a full new chain state.
+	/// Note: only called by compact.
+	pub fn rebuild_height_pos_index(&self) -> Result<(), Error> {
+		let now = Instant::now();
+
+		// clear it before rebuilding
+		self.batch.clear_output_pos_height()?;
+
+		let mut outputs_pos: Vec<(Commitment, u64)> = vec![];
+		for pos in self.output_pmmr.leaf_pos_iter() {
+			if let Some(out) = self.output_pmmr.get_data(pos) {
+				outputs_pos.push((out.commit, pos));
+			}
+		}
+		let total_outputs = outputs_pos.len();
+		if total_outputs == 0 {
+			debug!("rebuild_height_pos_index: nothing to be rebuilt");
+			return Ok(());
+		} else {
+			debug!(
+				"rebuild_height_pos_index: rebuilding {} outputs position & height...",
+				total_outputs
+			);
+		}
+
+		let max_height = self.head().height;
+
+		let mut i = 0;
+		for search_height in 0..max_height {
+			let h = self.get_header_by_height(search_height + 1)?;
+			while i < total_outputs {
+				let (commit, pos) = outputs_pos[i];
+				if pos > h.output_mmr_size {
+					// Note: MMR position is 1-based and not 0-based, so here must be '>' instead of '>='
+					break;
+				}
+				self.batch.save_output_pos_height(&commit, pos, h.height)?;
+				trace!("rebuild_height_pos_index: {:?}", (commit, pos, h.height));
+				i += 1;
+			}
+		}
+
+		debug!(
+			"txhashset: rebuild_height_pos_index: {} UTXOs, took {}s",
+			total_outputs,
+			now.elapsed().as_secs(),
+		);
+		Ok(())
+	}
+
 	/// Force the rollback of this extension, no matter the result
 	pub fn force_rollback(&mut self) {
 		self.rollback = true;