From 79ac26b6be9f1dcd535ff2f6d7c657a6d5725a62 Mon Sep 17 00:00:00 2001
From: Jan Bujak <jan@parity.io>
Date: Mon, 26 Jul 2021 10:51:32 +0000
Subject: [PATCH] Do not leak active head data in statement distribution

---
 .../network/statement-distribution/src/lib.rs | 22 +++++++++----------
 1 file changed, 10 insertions(+), 12 deletions(-)

diff --git a/node/network/statement-distribution/src/lib.rs b/node/network/statement-distribution/src/lib.rs
index 8d0cecdd63b4..b8974497955c 100644
--- a/node/network/statement-distribution/src/lib.rs
+++ b/node/network/statement-distribution/src/lib.rs
@@ -1847,6 +1847,16 @@ impl StatementDistribution {
 			FromOverseer::Signal(OverseerSignal::ActiveLeaves(ActiveLeavesUpdate { activated, deactivated })) => {
 				let _timer = metrics.time_active_leaves_update();
 
+				for deactivated in deactivated {
+					if active_heads.remove(&deactivated).is_some() {
+						tracing::trace!(
+							target: LOG_TARGET,
+							hash = ?deactivated,
+							"Deactivating leaf",
+						);
+					}
+				}
+
 				for activated in activated {
 					let relay_parent = activated.hash;
 					let span = PerLeafSpan::new(activated.span, "statement-distribution");
@@ -1862,18 +1872,6 @@ impl StatementDistribution {
 
 					active_heads.entry(relay_parent)
 						.or_insert(ActiveHeadData::new(session_info.validators.clone(), session_index, span));
-
-					active_heads.retain(|h, _| {
-						let live = !deactivated.contains(h);
-						if !live {
-							tracing::trace!(
-								target: LOG_TARGET,
-								hash = ?h,
-								"Deactivating leaf",
-							);
-						}
-						live
-					});
 				}
 			}
 			FromOverseer::Signal(OverseerSignal::BlockFinalized(..)) => {