input_set_best_block now properly accepts an Option (#497)

* `input_set_best_block` now properly accepts an `Option` * PR link
smol-dot · May 2, 2023 · 0737e37 · 0737e37
1 parent 4759b7b
commit 0737e37
Show file tree

Hide file tree

Showing 4 changed files with 32 additions and 14 deletions.
diff --git a/lib/src/chain/async_tree.rs b/lib/src/chain/async_tree.rs
@@ -816,27 +816,35 @@ where
     /// Updates the state machine to take into account that the best block of the input has been
     /// modified.
     ///
+    /// Pass `None` if the input best block is now the same as the output finalized block.
+    ///
     /// # Panic
     ///
     /// Panics if `new_best_block` isn't a valid node.
     /// Panics if `new_best_block` isn't equal or a descendant of the input finalized block.
     ///
-    pub fn input_set_best_block(&mut self, new_best_block: NodeIndex) {
+    pub fn input_set_best_block(&mut self, new_best_block: Option<NodeIndex>) {
         // Make sure that `new_best_block` is a descendant of the current input finalized block,
         // otherwise the state of the tree will be corrupted.
         // This is checked with an `assert!` rather than a `debug_assert!`, as this constraint
         // is part of the public API of this method.
-        assert!(self.input_finalized_index.map_or(true, |fin_idx| self
-            .non_finalized_blocks
-            .is_ancestor(fin_idx, new_best_block)));
+        assert!(match (self.input_finalized_index, new_best_block) {
+            (Some(f), Some(b)) => self.non_finalized_blocks.is_ancestor(f, b),
+            (Some(_), None) => false,
+            (None, Some(_)) => true,
+            (None, None) => true,
+        });
 
         // If necessary, update the weight of the block.
-        // TODO: this will panic if the new best block is equal to the output finalized block?
-        match &mut self
-            .non_finalized_blocks
-            .get_mut(new_best_block)
-            .unwrap()
-            .input_best_block_weight
+        match new_best_block
+            .map(|new_best_block| {
+                &mut self
+                    .non_finalized_blocks
+                    .get_mut(new_best_block)
+                    .unwrap()
+                    .input_best_block_weight
+            })
+            .unwrap_or(&mut self.finalized_block_weight)
         {
             w if *w == self.input_best_block_next_weight - 1 => {}
             w => {

diff --git a/light-base/src/runtime_service.rs b/light-base/src/runtime_service.rs
@@ -1398,14 +1398,19 @@ async fn run_background<TPlat: PlatformRef>(
 
                             match &mut guarded.tree {
                                 GuardedInner::FinalizedBlockRuntimeKnown {
+                                    finalized_block,
                                     tree, ..
                                 } => {
-                                    let idx = tree.input_iter_unordered().find(|block| block.user_data.hash == hash).unwrap().id;
+                                    let idx = if hash == finalized_block.hash {
+                                        None
+                                    } else {
+                                        Some(tree.input_iter_unordered().find(|block| block.user_data.hash == hash).unwrap().id)
+                                    };
                                     tree.input_set_best_block(idx);
                                 }
                                 GuardedInner::FinalizedBlockRuntimeUnknown { tree, .. } => {
                                     let idx = tree.input_iter_unordered().find(|block| block.user_data.hash == hash).unwrap().id;
-                                    tree.input_set_best_block(idx);
+                                    tree.input_set_best_block(Some(idx));
                                 }
                             }
 

diff --git a/light-base/src/sync_service/parachain.rs b/light-base/src/sync_service/parachain.rs
@@ -1031,12 +1031,13 @@ impl<TPlat: PlatformRef> ParachainBackgroundTask<TPlat> {
                     HashDisplay(&hash)
                 );
 
+                // If the block isn't found in `async_tree`, assume that it is equal to the
+                // finalized block (that has left the tree already).
                 let node_idx = runtime_subscription
                     .async_tree
                     .input_iter_unordered()
                     .find(|b| *b.user_data == hash)
-                    .unwrap()
-                    .id;
+                    .map(|b| b.id);
                 runtime_subscription
                     .async_tree
                     .input_set_best_block(node_idx);

diff --git a/wasm-node/CHANGELOG.md b/wasm-node/CHANGELOG.md
@@ -2,6 +2,10 @@
 
 ## Unreleased
 
+### Fixed
+
+- Fix panic when the best block of a chain switches to being equal to the current finalized block. This can occasionally happen for parachains in case of a reorg on the relay chain. ([#497](https://github.com/smol-dot/smoldot/pull/497))
+
 ## 1.0.3 - 2023-04-27
 
 ### Changed