add pop_blocks command + fix bug

Author: Boog900

binaries/cuprated/src/blockchain/interface.rs

@@ -170,6 +170,19 @@ pub async fn handle_incoming_block(
         .map_err(IncomingBlockError::InvalidBlock)
 }
 
+pub async fn pop_blocks(numb_blocks: usize) -> Result<(), anyhow::Error> {
+    let Some(incoming_block_tx) = COMMAND_TX.get() else {
+        // We could still be starting up the blockchain manager.
+        return anyhow::bail!("The blockchain manager is not running yet");
+    };
+
+    let (response_tx, response_rx) = oneshot::channel();
+
+    incoming_block_tx.send(BlockchainManagerCommand::PopBlocks {numb_blocks, response_tx }).await?;
+    
+    Ok(response_rx.await?)
+}
+
 /// Check if we have a block with the given hash.
 async fn block_exists(
     block_hash: [u8; 32],

binaries/cuprated/src/blockchain/manager/commands.rs

@@ -17,6 +17,11 @@ pub enum BlockchainManagerCommand {
         /// The channel to send the response down.
         response_tx: oneshot::Sender<Result<IncomingBlockOk, anyhow::Error>>,
     },
+    PopBlocks {
+        numb_blocks: usize,
+        /// The channel to send the response down.
+        response_tx: oneshot::Sender<()>,
+    },
 }
 
 /// The [`Ok`] response for an incoming block.

binaries/cuprated/src/blockchain/manager/handler.rs

@@ -55,6 +55,21 @@ impl super::BlockchainManager {
 
                 drop(response_tx.send(res));
             }
+            BlockchainManagerCommand::PopBlocks {
+                numb_blocks,
+                response_tx,
+            } => {
+                let _guard = REORG_LOCK.write().await;
+                self.pop_blocks(numb_blocks).await;
+                self.blockchain_write_handle
+                    .ready()
+                    .await
+                    .expect(PANIC_CRITICAL_SERVICE_ERROR)
+                    .call(BlockchainWriteRequest::FlushAltBlocks)
+                    .await
+                    .expect(PANIC_CRITICAL_SERVICE_ERROR);
+                let _ = response_tx.send(());
+            }
         }
     }
 

binaries/cuprated/src/commands.rs

@@ -50,6 +50,11 @@ pub enum Command {
 
     /// Print the height of first block not contained in the fast sync hashes.
     FastSyncStopHeight,
+
+    /// Pop blocks from the top of the blockchain.
+    PopBlocks {
+        numb_blocks: usize,
+    }
 }
 
 /// The log output target.
@@ -131,6 +136,14 @@ pub async fn io_loop(
 
                 println!("{stop_height}");
             }
+            Command::PopBlocks {numb_blocks} => {
+                let res = crate::blockchain::interface::pop_blocks(numb_blocks).await;
+
+                match res {
+                    Ok(()) => println!("Popped {numb_blocks} blocks."),
+                    Err(e) => println!("Failed to pop blocks: {e}"),
+                }
+            }
         }
     }
 }

consensus/context/src/task.rs

@@ -222,15 +222,15 @@ impl<D: Database + Clone + Send + 'static> ContextTask<D> {
 
                 self.difficulty_cache
                     .pop_blocks_main_chain(numb_blocks, self.database.clone())
-                    .await?;
+                    .await.unwrap();
                 self.weight_cache
                     .pop_blocks_main_chain(numb_blocks, self.database.clone())
-                    .await?;
+                    .await.unwrap();
                 self.rx_vm_cache
                     .pop_blocks_main_chain(self.chain_height - numb_blocks - 1);
                 self.hardfork_state
                     .pop_blocks_main_chain(numb_blocks, self.database.clone())
-                    .await?;
+                    .await.unwrap();
 
                 self.alt_chain_cache_map.clear();
 

consensus/context/src/weight.rs

@@ -137,31 +137,34 @@ impl BlockWeightsCache {
 
         let chain_height = self.tip_height + 1;
 
-        let new_long_term_start_height = chain_height
-            .saturating_sub(self.config.long_term_window)
-            .saturating_sub(numb_blocks);
-
-        let old_long_term_weights = get_long_term_weight_in_range(
-            new_long_term_start_height
-                // current_chain_height - self.long_term_weights.len() blocks are already in the cache.
-                ..(chain_height - self.long_term_weights.window_len()),
-            database.clone(),
-            Chain::Main,
-        )
-        .await?;
-
-        let new_short_term_start_height = chain_height
-            .saturating_sub(self.config.short_term_window)
-            .saturating_sub(numb_blocks);
+        let old_long_term_weights = if let Some(new_long_term_start_height) = chain_height
+            .checked_sub(self.config.long_term_window + numb_blocks)
+        {
+            get_long_term_weight_in_range(
+                new_long_term_start_height
+                    // current_chain_height - self.long_term_weights.len() blocks are already in the cache.
+                    ..(new_long_term_start_height + numb_blocks),
+                database.clone(),
+                Chain::Main,
+            )
+                .await?
+        } else {
+            vec![]
+        };
 
-        let old_short_term_weights = get_blocks_weight_in_range(
-            new_short_term_start_height
-                // current_chain_height - self.long_term_weights.len() blocks are already in the cache.
-                ..(chain_height - self.short_term_block_weights.window_len()),
-            database,
-            Chain::Main,
-        )
-        .await?;
+        let old_short_term_weights =if let Some(new_short_term_start_height) = chain_height
+            .checked_sub(self.config.short_term_window + numb_blocks) {
+             get_blocks_weight_in_range(
+                new_short_term_start_height
+                    // current_chain_height - self.long_term_weights.len() blocks are already in the cache.
+                    ..(min(numb_blocks, self.short_term_block_weights.window_len()) + new_short_term_start_height),
+                database,
+                Chain::Main,
+            )
+                .await?
+        } else {
+            vec![]
+        };
 
         for _ in 0..numb_blocks {
             self.short_term_block_weights.pop_back();

consensus/src/tests/context/weight.rs

@@ -60,12 +60,12 @@ async fn pop_blocks_greater_than_window() -> Result<(), tower::BoxError> {
 
     let old_cache = weight_cache.clone();
 
-    weight_cache.new_block(5000, 0, 0);
-    weight_cache.new_block(5001, 0, 0);
-    weight_cache.new_block(5002, 0, 0);
+    for i in 0..4999 {
+        weight_cache.new_block(5000 + i, 0, 0);
+    }
 
     weight_cache
-        .pop_blocks_main_chain(3, database)
+        .pop_blocks_main_chain(4999, database)
         .await
         .unwrap();
 

consensus/src/tests/mock_db.rs

@@ -173,13 +173,17 @@ impl Service<BlockchainReadRequest> for DummyDatabase {
                     let mut end = range.end;
                     let mut start = range.start;
 
+                    let block_len = blocks.read().unwrap().len();
                     if let Some(dummy_height) = dummy_height {
-                        let block_len = blocks.read().unwrap().len();
 
                         end -= dummy_height - block_len;
                         start -= dummy_height - block_len;
                     }
 
+                    if block_len < end {
+                        return Err("end block not in database!".into());
+                    }
+
                     BlockchainResponse::BlockExtendedHeaderInRange(
                         blocks
                             .read()