fix: try to fix bug mimblewimble#2585.

Signed-off-by: Mike Tang <[email protected]>

chain/src/chain.rs

@@ -34,7 +34,7 @@ use crate::types::{
 use crate::util::secp::pedersen::{Commitment, RangeProof};
 use crate::util::{Mutex, RwLock, StopState};
 use grin_store::Error::NotFoundErr;
-use std::collections::HashMap;
+use std::collections::{HashMap, HashSet};
 use std::fs::{self, File};
 use std::path::PathBuf;
 use std::sync::atomic::{AtomicUsize, Ordering};
@@ -1065,6 +1065,49 @@ impl Chain {
 			count, tail.height
 		);
 
+		// Remove short live fork blocks
+		// Get all block headers in db
+		let all_block_headers = self.store.all_block_headers()?;
+
+		let mut on_chain_block_set: HashSet<Hash> = HashSet::new();
+
+		let current_hash = txhashset.get_header_hash_by_height(head.height)?;
+		let mut current = batch.get_block_header(&current_hash)?;
+
+		loop {
+			// Go to the store directly so we can handle NotFoundErr robustly.
+			match self.store.get_block(&current.hash()) {
+				Ok(_b) => {
+					on_chain_block_set.insert(current.hash());
+				}
+				Err(NotFoundErr(_)) => {
+					break;
+				}
+				Err(e) => {
+					return Err(
+						ErrorKind::StoreErr(e, "retrieving block to compact".to_owned()).into(),
+					);
+				}
+			}
+			if current.height <= 1 {
+				break;
+			}
+			match batch.get_previous_header(&current) {
+				Ok(h) => current = h,
+				Err(NotFoundErr(_)) => break,
+				Err(e) => return Err(From::from(e)),
+			}
+		}
+
+		// clean all not in current chain fork
+		for header in all_block_headers {
+			let header_hash = header.hash();
+			if !on_chain_block_set.contains(&header_hash) {
+				let block = self.store.get_block(&header_hash)?;
+				batch.delete_block(&block.hash())?;
+			}
+		}
+
 		Ok(())
 	}
 

chain/src/store.rs

@@ -120,6 +120,16 @@ impl ChainStore {
 		)
 	}
 
+	/// Retrieve all block headers
+	pub fn all_block_headers(&self) -> Result<Vec<BlockHeader>, Error> {
+		let key = to_key(BLOCK_HEADER_PREFIX, &mut "".to_string().into_bytes());
+		Ok(self
+			.db
+			.iter::<BlockHeader>(&key)?
+			.map(|(_, v)| v)
+			.collect::<Vec<_>>())
+	}
+
 	/// Builds a new batch to be used with this store.
 	pub fn batch(&self) -> Result<Batch<'_>, Error> {
 		Ok(Batch {