From b0d0fb31d8758602387e4ba1cb712e90dff8b2a7 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Bastian=20K=C3=B6cher?= <git@kchr.de>
Date: Thu, 9 Nov 2023 21:44:23 +0100
Subject: [PATCH] sc-state-db: Keep track of `LAST_PRUNED` after warp syncing
 (#2228)

When warp syncing we import the target block with all its state.
However, we didn't store the `LAST_PRUNED` block which would then lead
to `pruning` to forget about the imported block after a restart of the
node. We just set `LAST_PRUNED` to the parent block of the warp sync
target block to fix this issue.
---
 substrate/client/state-db/src/pruning.rs | 42 ++++++++++++++++++++++--
 1 file changed, 40 insertions(+), 2 deletions(-)

diff --git a/substrate/client/state-db/src/pruning.rs b/substrate/client/state-db/src/pruning.rs
index 7bee2b1d99..623d30b098 100644
--- a/substrate/client/state-db/src/pruning.rs
+++ b/substrate/client/state-db/src/pruning.rs
@@ -406,12 +406,24 @@ impl<BlockHash: Hash, Key: Hash, D: MetaDb> RefWindow<BlockHash, Key, D> {
 		commit: &mut CommitSet<Key>,
 	) -> Result<(), Error<D::Error>> {
 		if self.base == 0 && self.is_empty() && number > 0 {
-			// assume that parent was canonicalized
+			// This branch is taken if the node imports the target block of a warp sync.
+			// assume that the block was canonicalized
 			self.base = number;
+			// The parent of the block was the last block that got pruned.
+			commit
+				.meta
+				.inserted
+				.push((to_meta_key(LAST_PRUNED, &()), (number - 1).encode()));
 		} else if (self.base + self.window_size()) != number {
 			return Err(Error::StateDb(StateDbError::InvalidBlockNumber))
 		}
-		trace!(target: "state-db", "Adding to pruning window: {:?} ({} inserted, {} deleted)", hash, commit.data.inserted.len(), commit.data.deleted.len());
+		trace!(
+			target: "state-db",
+			"Adding to pruning window: {:?} ({} inserted, {} deleted)",
+			hash,
+			commit.data.inserted.len(),
+			commit.data.deleted.len(),
+		);
 		let inserted = if matches!(self.queue, DeathRowQueue::Mem { .. }) {
 			commit.data.inserted.iter().map(|(k, _)| k.clone()).collect()
 		} else {
@@ -869,4 +881,30 @@ mod tests {
 		pruning.prune_one(&mut commit).unwrap();
 		db.commit(&commit);
 	}
+
+	/// Ensure that after warp syncing the state is stored correctly in the db. The warp sync target
+	/// block is imported with all its state at once. This test ensures that after a restart
+	/// `pruning` still knows that this block was imported.
+	#[test]
+	fn store_correct_state_after_warp_syncing() {
+		for count_insertions in [true, false] {
+			let mut db = make_db(&[]);
+			let mut pruning: RefWindow<u64, H256, TestDb> =
+				RefWindow::new(db.clone(), DEFAULT_MAX_BLOCK_CONSTRAINT, count_insertions).unwrap();
+			let block = 10000;
+
+			// import blocks
+			let mut commit = make_commit(&[], &[]);
+			pruning.note_canonical(&block, block, &mut commit).unwrap();
+			push_last_canonicalized(block, &mut commit);
+			db.commit(&commit);
+
+			// load a new queue from db
+			// `cache` should be the same
+			let pruning: RefWindow<u64, H256, TestDb> =
+				RefWindow::new(db, DEFAULT_MAX_BLOCK_CONSTRAINT, count_insertions).unwrap();
+
+			assert_eq!(HaveBlock::Yes, pruning.have_block(&block, block));
+		}
+	}
 }