From 6f1590b253ea13ef1e991614ee89f048133300ab Mon Sep 17 00:00:00 2001
From: Andronik Ordian <write@reusable.software>
Date: Wed, 9 Dec 2020 16:33:58 +0100
Subject: [PATCH] runtime/session_info: avoid heavy loop when introduced on a
 live chain  (#2099)

* session_info: a heavy loop test

* session_info: fix a typo

* session_info: fix heavy loop

* session_info: crank the iterations all the way up
---
 .../runtime/parachains/src/session_info.rs    | 31 +++++++++++++++++--
 1 file changed, 28 insertions(+), 3 deletions(-)

diff --git a/polkadot/runtime/parachains/src/session_info.rs b/polkadot/runtime/parachains/src/session_info.rs
index c3faec7eeb1..decc25a20d0 100644
--- a/polkadot/runtime/parachains/src/session_info.rs
+++ b/polkadot/runtime/parachains/src/session_info.rs
@@ -101,8 +101,11 @@ impl<T: Config> Module<T> {
 		// update `EarliestStoredSession` based on `config.dispute_period`
 		EarliestStoredSession::set(new_earliest_stored_session);
 		// remove all entries from `Sessions` from the previous value up to the new value
-		for idx in old_earliest_stored_session..new_earliest_stored_session {
-			Sessions::remove(&idx);
+		// avoid a potentially heavy loop when introduced on a live chain
+		if old_earliest_stored_session != 0 || Sessions::get(0).is_some() {
+			for idx in old_earliest_stored_session..new_earliest_stored_session {
+				Sessions::remove(&idx);
+			}
 		}
 		// create a new entry in `Sessions` with information about the current session
 		let new_session_info = SessionInfo {
@@ -216,7 +219,7 @@ mod tests {
 	}
 
 	#[test]
-	fn session_pruning_is_based_on_dispute_deriod() {
+	fn session_pruning_is_based_on_dispute_period() {
 		new_test_ext(genesis_config()).execute_with(|| {
 			run_to_block(100, session_changes);
 			assert_eq!(EarliestStoredSession::get(), 9);
@@ -253,4 +256,26 @@ mod tests {
 			assert_eq!(session.needed_approvals, 42);
 		})
 	}
+
+	#[test]
+	fn session_pruning_avoids_heavy_loop() {
+		new_test_ext(genesis_config()).execute_with(|| {
+			let start = 1_000_000_000;
+			System::on_initialize(start);
+			System::set_block_number(start);
+
+			if let Some(notification) = new_session_every_block(start) {
+				Configuration::initializer_on_new_session(&notification.validators, &notification.queued);
+				SessionInfo::initializer_on_new_session(&notification);
+			}
+
+			Configuration::initializer_initialize(start);
+			SessionInfo::initializer_initialize(start);
+
+			assert_eq!(EarliestStoredSession::get(), start - 1);
+
+			run_to_block(start + 1, new_session_every_block);
+			assert_eq!(EarliestStoredSession::get(), start);
+		})
+	}
 }
-- 
GitLab