From 493d703f5dd24f6e9bde42e749e31f0e0e042a9f Mon Sep 17 00:00:00 2001
From: Niklas Adolfsson <niklasadolfsson1@gmail.com>
Date: Tue, 9 Jan 2024 10:22:42 +0100
Subject: [PATCH] fix: runtime upgrades in unstable backend

This commit changes the following:

1. FollowEvent::Finalized: just remove the finalized and pruned blocks from runtimes
2. FollowEvent::Initalized: just remove the finalized block from runtimes
---
 subxt/src/backend/unstable/mod.rs | 33 +++++++++++++++++++++++++------
 1 file changed, 27 insertions(+), 6 deletions(-)

diff --git a/subxt/src/backend/unstable/mod.rs b/subxt/src/backend/unstable/mod.rs
index fb1437352d..9067d7b6c7 100644
--- a/subxt/src/backend/unstable/mod.rs
+++ b/subxt/src/backend/unstable/mod.rs
@@ -353,7 +353,7 @@ impl<T: Config + Send + Sync + 'static> Backend<T> for UnstableBackend<T> {
             .filter_map(move |ev| {
                 let output = match ev {
                     FollowEvent::Initialized(ev) => {
-                        runtimes.clear();
+                        runtimes.remove(&ev.finalized_block_hash.hash());
                         ev.finalized_block_runtime
                     }
                     FollowEvent::NewBlock(ev) => {
@@ -363,14 +363,35 @@ impl<T: Config + Send + Sync + 'static> Backend<T> for UnstableBackend<T> {
                         None
                     }
                     FollowEvent::Finalized(ev) => {
-                        let next_runtime = ev
+                        let next_runtime = {
+                            let mut it = ev
+                                .finalized_block_hashes
+                                .iter()
+                                .rev()
+                                .filter_map(|h| runtimes.get(&h.hash()).cloned())
+                                .peekable();
+
+                            let next = it.next();
+
+                            if it.peek().is_some() {
+                                tracing::warn!(
+                                    target: "subxt",
+                                    "Several runtime upgrades in the finalized blocks but only the latest runtime upgrade is returned"
+                                );
+                            }
+
+                            next
+                        };
+
+                        // Remove finalized and pruned blocks as valid runtime upgrades.
+                        for block in ev
                             .finalized_block_hashes
                             .iter()
-                            .rev()
-                            .filter_map(|h| runtimes.get(&h.hash()).cloned())
-                            .next();
+                            .chain(ev.pruned_block_hashes.iter())
+                        {
+                            runtimes.remove(&block.hash());
+                        }
 
-                        runtimes.clear();
                         next_runtime
                     }
                     _ => None,