From 6d337425bf51b086313aa3e7da5936305dd95a28 Mon Sep 17 00:00:00 2001
From: Antioch Peverell <apeverell@protonmail.com>
Date: Tue, 8 Dec 2020 12:42:28 +0000
Subject: [PATCH] [5.0.x] Block migration flag (#3513)

* rework migrate_db_v2_v3

* db flag to track block migration
---
 chain/src/chain.rs | 53 ++++++++++++++++++++++++++++++++--------------
 chain/src/store.rs | 50 ++++++++++++++++++++++++++++++++++++++++++-
 2 files changed, 86 insertions(+), 17 deletions(-)

diff --git a/chain/src/chain.rs b/chain/src/chain.rs
index cd0afd3878..a25df38f7b 100644
--- a/chain/src/chain.rs
+++ b/chain/src/chain.rs
@@ -1484,32 +1484,53 @@ impl Chain {
 	/// Migrate our local db from v2 to v3.
 	/// "commit only" inputs.
 	fn migrate_db_v2_v3(store: &ChainStore) -> Result<(), Error> {
+		if store.batch()?.is_blocks_v3_migrated()? {
+			// Previously migrated so skipping.
+			debug!("migrate_db_v2_v3: previously migrated, skipping");
+			return Ok(());
+		}
+		let mut total = 0;
 		let mut keys_to_migrate = vec![];
 		for (k, v) in store.batch()?.blocks_raw_iter()? {
+			total += 1;
+
 			// We want to migrate all blocks that cannot be read via v3 protocol version.
-			let block_v2: Result<Block, _> =
-				ser::deserialize(&mut Cursor::new(&v), ProtocolVersion(2));
 			let block_v3: Result<Block, _> =
 				ser::deserialize(&mut Cursor::new(&v), ProtocolVersion(3));
-			if let (Ok(_), Err(_)) = (block_v2, block_v3) {
-				keys_to_migrate.push(k);
+			if block_v3.is_err() {
+				let block_v2: Result<Block, _> =
+					ser::deserialize(&mut Cursor::new(&v), ProtocolVersion(2));
+				if block_v2.is_ok() {
+					keys_to_migrate.push(k);
+				}
 			}
 		}
 		debug!(
-			"migrate_db_v2_v3: {} blocks to migrate",
-			keys_to_migrate.len()
+			"migrate_db_v2_v3: {} (of {}) blocks to migrate",
+			keys_to_migrate.len(),
+			total,
 		);
 		let mut count = 0;
-		keys_to_migrate.chunks(100).try_for_each(|keys| {
-			let batch = store.batch()?;
-			for key in keys {
-				batch.migrate_block(&key, ProtocolVersion(2), ProtocolVersion(3))?;
-				count += 1;
-			}
-			batch.commit()?;
-			debug!("migrate_db_v2_v3: successfully migrated {} blocks", count);
-			Ok(())
-		})
+		keys_to_migrate
+			.chunks(100)
+			.try_for_each(|keys| {
+				let batch = store.batch()?;
+				for key in keys {
+					batch.migrate_block(&key, ProtocolVersion(2), ProtocolVersion(3))?;
+					count += 1;
+				}
+				batch.commit()?;
+				debug!("migrate_db_v2_v3: successfully migrated {} blocks", count);
+				Ok(())
+			})
+			.and_then(|_| {
+				// Set flag to indicate we have migrated all blocks in the db.
+				// We will skip migration in the future.
+				let batch = store.batch()?;
+				batch.set_blocks_v3_migrated(true)?;
+				batch.commit()?;
+				Ok(())
+			})
 	}
 
 	/// Gets the block header in which a given output appears in the txhashset.
diff --git a/chain/src/store.rs b/chain/src/store.rs
index 74c520180f..3db41ad4b8 100644
--- a/chain/src/store.rs
+++ b/chain/src/store.rs
@@ -18,7 +18,7 @@ use crate::core::consensus::HeaderInfo;
 use crate::core::core::hash::{Hash, Hashed};
 use crate::core::core::{Block, BlockHeader, BlockSums};
 use crate::core::pow::Difficulty;
-use crate::core::ser::ProtocolVersion;
+use crate::core::ser::{ProtocolVersion, Readable, Writeable};
 use crate::linked_list::MultiIndex;
 use crate::types::{CommitPos, Tip};
 use crate::util::secp::pedersen::Commitment;
@@ -46,6 +46,11 @@ pub const NRD_KERNEL_ENTRY_PREFIX: u8 = b'k';
 const BLOCK_SUMS_PREFIX: u8 = b'M';
 const BLOCK_SPENT_PREFIX: u8 = b'S';
 
+/// Prefix for various boolean flags stored in the db.
+const BOOL_FLAG_PREFIX: u8 = b'B';
+/// Boolean flag for v3 migration.
+const BLOCKS_V3_MIGRATED: &str = "blocks_v3_migrated";
+
 /// All chain-related database operations
 pub struct ChainStore {
 	db: store::Store,
@@ -214,6 +219,27 @@ impl<'a> Batch<'a> {
 		Ok(())
 	}
 
+	/// DB flag representing full migration of blocks to v3 version.
+	/// Default to false if flag not present.
+	pub fn is_blocks_v3_migrated(&self) -> Result<bool, Error> {
+		let migrated: Option<BoolFlag> = self
+			.db
+			.get_ser(&to_key(BOOL_FLAG_PREFIX, BLOCKS_V3_MIGRATED))?;
+		match migrated {
+			None => Ok(false),
+			Some(x) => Ok(x.into()),
+		}
+	}
+
+	/// Set DB flag representing full migration of blocks to v3 version.
+	pub fn set_blocks_v3_migrated(&self, migrated: bool) -> Result<(), Error> {
+		self.db.put_ser(
+			&to_key(BOOL_FLAG_PREFIX, BLOCKS_V3_MIGRATED)[..],
+			&BoolFlag(migrated),
+		)?;
+		Ok(())
+	}
+
 	/// Migrate a block stored in the db reading from one protocol version and writing
 	/// with new protocol version.
 	pub fn migrate_block(
@@ -496,3 +522,25 @@ impl<'a> Iterator for DifficultyIter<'a> {
 pub fn nrd_recent_kernel_index() -> MultiIndex<CommitPos> {
 	MultiIndex::init(NRD_KERNEL_LIST_PREFIX, NRD_KERNEL_ENTRY_PREFIX)
 }
+
+struct BoolFlag(bool);
+
+impl From<BoolFlag> for bool {
+	fn from(b: BoolFlag) -> Self {
+		b.0
+	}
+}
+
+impl Readable for BoolFlag {
+	fn read<R: ser::Reader>(reader: &mut R) -> Result<Self, ser::Error> {
+		let x = reader.read_u8()?;
+		Ok(BoolFlag(1 & x == 1))
+	}
+}
+
+impl Writeable for BoolFlag {
+	fn write<W: ser::Writer>(&self, writer: &mut W) -> Result<(), ser::Error> {
+		writer.write_u8(self.0.into())?;
+		Ok(())
+	}
+}