Clean-up runtimes management in runtime_service (#2781)

Fix #2467

We now clean up old runtimes when it makes sense, rather than after a
download where it doesn't make sense.

Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

bin/light-base/src/runtime_service.rs

@@ -1528,19 +1528,20 @@ impl<TPlat: Platform> Background<TPlat> {
     ) {
         let mut guarded = self.guarded.lock().await;
 
-        guarded.runtimes.retain(|_, rt| rt.strong_count() > 0);
-
-        // Try to find an existing identical runtime.
+        // Try to find an existing runtime identical to the one that has just been downloaded.
+        // This loop is `O(n)`, but given that we expect this list to very small (at most 1 or
+        // 2 elements), this is not a problem.
         let existing_runtime = guarded
             .runtimes
             .iter()
             .filter_map(|(_, rt)| rt.upgrade())
             .find(|rt| rt.runtime_code == storage_code && rt.heap_pages == storage_heap_pages);
 
+        // If no identical runtime was found, try compiling the runtime.
+        // TODO: use a let-else construct here once stable
         let runtime = if let Some(existing_runtime) = existing_runtime {
             existing_runtime
         } else {
-            // No identical runtime was found. Try compiling the new runtime.
             let runtime = SuccessfulRuntime::from_storage(&storage_code, &storage_heap_pages).await;
             match &runtime {
                 Ok(runtime) => {
@@ -1572,6 +1573,7 @@ impl<TPlat: Platform> Background<TPlat> {
             runtime
         };
 
+        // Insert the runtime into the tree.
         match &mut guarded.tree {
             GuardedInner::FinalizedBlockRuntimeKnown { tree, .. } => {
                 tree.async_op_finished(async_op_id, runtime);
@@ -1598,6 +1600,7 @@ impl<TPlat: Platform> Background<TPlat> {
                         user_data: new_finalized,
                         best_block_index,
                         pruned_blocks,
+                        former_finalized_async_op_user_data: former_finalized_runtime,
                         ..
                     }) => {
                         *finalized_block = new_finalized;
@@ -1610,6 +1613,13 @@ impl<TPlat: Platform> Background<TPlat> {
                             HashDisplay(&finalized_block.hash), HashDisplay(&best_block_hash)
                         );
 
+                        // The finalization might cause some runtimes in the list of runtimes
+                        // to have become unused. Clean them up.
+                        drop(former_finalized_runtime);
+                        guarded
+                            .runtimes
+                            .retain(|_, runtime| runtime.strong_count() > 0);
+
                         let all_blocks_notif = Notification::Finalized {
                             best_block_hash,
                             hash: finalized_block.hash,
@@ -1965,7 +1975,7 @@ impl<TPlat: Platform> Background<TPlat> {
         // Clean up unused runtimes to free up resources.
         guarded
             .runtimes
-            .retain(|_, runtime| runtime.strong_count() == 0);
+            .retain(|_, runtime| runtime.strong_count() > 0);
     }
 }
 

bin/wasm-node/CHANGELOG.md

@@ -6,6 +6,10 @@
 
 - Removed the `version` field of the struct returned by the `rpc_methods` function. This is technically a breaking change, but it has been introduced in a minor version bump because it is very insubstantial. ([#2756](https://github.com/paritytech/smoldot/pull/2756))
 
+### Fixed
+
+- Fix old runtimes not being cleaned up properly and runtimes being downloaded multiple times after an on-chain runtime upgrade. ([#2781](https://github.com/paritytech/smoldot/pull/2781))
+
 ## 0.6.34 - 2022-09-20
 
 ### Added