feat: persistent tx-pool data into a file when it has been shutdown

ckb-bin/src/subcommand/replay.rs

@@ -19,7 +19,7 @@ pub fn replay(args: ReplayArgs, async_handle: Handle) -> Result<(), ExitCode> {
     )?;
     let (shared, _) = shared_builder
         .consensus(args.consensus.clone())
-        .tx_pool_config(args.config.tx_pool)
+        .tx_pool_config(args.config.tx_pool.clone())
         .build()?;
 
     if !args.tmp_target.is_dir() {

ckb-bin/src/subcommand/run.rs

@@ -63,7 +63,10 @@ pub fn run(args: RunArgs, version: Version, async_handle: Handle) -> Result<(),
     exit_handler.wait_for_exit();
 
     info!("Finishing work, please wait...");
+    shared.tx_pool_controller().save_pool().map_err(|err| {
+        eprintln!("TxPool Error: {}", err);
+        ExitCode::Failure
+    })?;
     drop(rpc_server);
-
     Ok(())
 }

test/Cargo.toml

@@ -39,6 +39,9 @@ lazy_static = "1.4.0"
 byteorder = "1.3.1"
 jsonrpc-core = "18.0"
 
+[target.'cfg(not(target_os="windows"))'.dependencies]
+nix = "0.20.0"
+
 # Prevent this from interfering with workspaces
 [workspace]
 members = ["."]

test/src/main.rs

@@ -382,6 +382,8 @@ fn all_specs() -> Vec<Box<dyn Spec>> {
         Box::new(TemplateSizeLimit),
         Box::new(PoolReconcile),
         Box::new(PoolResurrect),
+        #[cfg(not(target_os = "windows"))]
+        Box::new(PoolPersisted),
         Box::new(TransactionRelayBasic),
         Box::new(TransactionRelayLowFeeRate),
         // TODO failed on poor CI server

test/src/node.rs

@@ -26,15 +26,20 @@ use std::process::{self, Child, Command, Stdio};
 use std::thread::sleep;
 use std::time::{Duration, Instant};
 
-struct ProcessGuard(pub Child);
+struct ProcessGuard {
+    pub child: Child,
+    pub killed: bool,
+}
 
 impl Drop for ProcessGuard {
     fn drop(&mut self) {
-        match self.0.kill() {
-            Err(e) => error!("Could not kill ckb process: {}", e),
-            Ok(_) => debug!("Successfully killed ckb process"),
+        if !self.killed {
+            match self.child.kill() {
+                Err(e) => error!("Could not kill ckb process: {}", e),
+                Ok(_) => debug!("Successfully killed ckb process"),
+            }
+            let _ = self.child.wait();
         }
-        let _ = self.0.wait();
     }
 }
 
@@ -532,14 +537,33 @@ impl Node {
             }
         };
 
-        self.guard = Some(ProcessGuard(child_process));
+        self.guard = Some(ProcessGuard {
+            child: child_process,
+            killed: false,
+        });
         self.node_id = Some(node_info.node_id);
     }
 
     pub fn stop(&mut self) {
         drop(self.guard.take())
     }
 
+    #[cfg(not(target_os = "windows"))]
+    pub fn stop_gracefully(&mut self) {
+        if let Some(mut guard) = self.guard.take() {
+            if !guard.killed {
+                // send SIGINT to the child
+                nix::sys::signal::kill(
+                    nix::unistd::Pid::from_raw(guard.child.id() as i32),
+                    nix::sys::signal::Signal::SIGINT,
+                )
+                .expect("cannot send ctrl-c");
+                let _ = guard.child.wait();
+                guard.killed = true;
+            }
+        }
+    }
+
     pub fn export(&self, target: String) {
         Command::new(binary())
             .args(&[

test/src/specs/tx_pool/mod.rs

@@ -5,6 +5,8 @@ mod depend_tx_in_same_block;
 mod descendant;
 mod different_txs_with_same_input;
 mod limit;
+#[cfg(not(target_os = "windows"))]
+mod pool_persisted;
 mod pool_reconcile;
 mod pool_resurrect;
 mod proposal_expire_rule;
@@ -27,6 +29,8 @@ pub use depend_tx_in_same_block::*;
 pub use descendant::*;
 pub use different_txs_with_same_input::*;
 pub use limit::*;
+#[cfg(not(target_os = "windows"))]
+pub use pool_persisted::*;
 pub use pool_reconcile::*;
 pub use pool_resurrect::*;
 pub use proposal_expire_rule::*;

test/src/specs/tx_pool/pool_persisted.rs

@@ -0,0 +1,65 @@
+use crate::util::mining::{mine, mine_until_out_bootstrap_period};
+use crate::{Node, Spec};
+use ckb_logger::info;
+
+pub struct PoolPersisted;
+
+impl Spec for PoolPersisted {
+    crate::setup!(num_nodes: 1);
+
+    fn run(&self, nodes: &mut Vec<Node>) {
+        let node0 = &mut nodes[0];
+
+        info!("Generate 1 block on node0");
+        mine_until_out_bootstrap_period(node0);
+
+        info!("Generate 6 txs on node0");
+        let mut hash = node0.generate_transaction();
+
+        (0..5).for_each(|_| {
+            let tx = node0.new_transaction(hash.clone());
+            hash = node0.rpc_client().send_transaction(tx.data().into());
+        });
+
+        info!("Generate 1 more blocks on node0");
+        mine(node0, 1);
+
+        info!("Generate 5 more txs on node0");
+        (0..5).for_each(|_| {
+            let tx = node0.new_transaction(hash.clone());
+            hash = node0.rpc_client().send_transaction(tx.data().into());
+        });
+
+        info!("Generate 1 more blocks on node0");
+        mine(node0, 1);
+
+        node0.wait_for_tx_pool();
+
+        let tx_pool_info_original = node0.get_tip_tx_pool_info();
+
+        info!("Stop node0 gracefully");
+        node0.stop_gracefully();
+
+        info!("Start node0");
+        node0.start();
+
+        let tx_pool_info_reloaded = node0.get_tip_tx_pool_info();
+        info!("TxPool should be same as before");
+        info!("tx_pool_info_original: {:?}", tx_pool_info_original);
+        info!("tx_pool_info_reloaded: {:?}", tx_pool_info_reloaded);
+        assert_eq!(
+            tx_pool_info_original.proposed,
+            tx_pool_info_reloaded.proposed
+        );
+        assert_eq!(tx_pool_info_original.orphan, tx_pool_info_reloaded.orphan);
+        assert_eq!(tx_pool_info_original.pending, tx_pool_info_reloaded.pending);
+        assert_eq!(
+            tx_pool_info_original.total_tx_size,
+            tx_pool_info_reloaded.total_tx_size
+        );
+        assert_eq!(
+            tx_pool_info_original.total_tx_cycles,
+            tx_pool_info_reloaded.total_tx_cycles
+        );
+    }
+}

tx-pool/src/lib.rs

@@ -6,6 +6,7 @@ mod callback;
 mod chunk_process;
 mod component;
 pub mod error;
+mod persisted;
 pub mod pool;
 mod process;
 pub mod service;

tx-pool/src/persisted.rs

@@ -0,0 +1,97 @@
+use crate::TxPool;
+use ckb_error::{AnyError, OtherError};
+use ckb_types::{
+    core::TransactionView,
+    packed::{TransactionVec, TransactionVecReader},
+    prelude::*,
+};
+use std::{
+    fs::OpenOptions,
+    io::{Read as _, Write as _},
+};
+
+/// The version of the persisted tx-pool data.
+pub(crate) const VERSION: u32 = 1;
+
+impl TxPool {
+    pub(crate) fn load_from_file(&self) -> Result<Vec<TransactionView>, AnyError> {
+        let mut persisted_data_file = self.config.persisted_data.clone();
+        persisted_data_file.set_extension(format!("v{}", VERSION));
+
+        if persisted_data_file.exists() {
+            let mut file = OpenOptions::new()
+                .read(true)
+                .open(&persisted_data_file)
+                .map_err(|err| {
+                    let errmsg = format!(
+                        "Failed to open the tx-pool persisted data file [{:?}], cause: {}",
+                        persisted_data_file, err
+                    );
+                    OtherError::new(errmsg)
+                })?;
+            let mut buffer = Vec::new();
+            file.read_to_end(&mut buffer).map_err(|err| {
+                let errmsg = format!(
+                    "Failed to read the tx-pool persisted data file [{:?}], cause: {}",
+                    persisted_data_file, err
+                );
+                OtherError::new(errmsg)
+            })?;
+
+            let persisted_data = TransactionVecReader::from_slice(&buffer)
+                .map_err(|err| {
+                    let errmsg = format!(
+                        "The tx-pool persisted data file [{:?}] is broken, cause: {}",
+                        persisted_data_file, err
+                    );
+                    OtherError::new(errmsg)
+                })?
+                .to_entity();
+
+            Ok(persisted_data
+                .into_iter()
+                .map(|tx| tx.into_view())
+                .collect())
+        } else {
+            Ok(Vec::new())
+        }
+    }
+
+    pub(crate) fn save_into_file(&mut self) -> Result<(), AnyError> {
+        let mut persisted_data_file = self.config.persisted_data.clone();
+        persisted_data_file.set_extension(format!("v{}", VERSION));
+
+        let mut file = OpenOptions::new()
+            .create(true)
+            .write(true)
+            .truncate(true)
+            .open(&persisted_data_file)
+            .map_err(|err| {
+                let errmsg = format!(
+                    "Failed to open the tx-pool persisted data file [{:?}], cause: {}",
+                    persisted_data_file, err
+                );
+                OtherError::new(errmsg)
+            })?;
+
+        let txs = TransactionVec::new_builder()
+            .extend(self.drain_all_transactions().iter().map(|tx| tx.data()))
+            .build();
+
+        file.write_all(txs.as_slice()).map_err(|err| {
+            let errmsg = format!(
+                "Failed to write the tx-pool persisted data into file [{:?}], cause: {}",
+                persisted_data_file, err
+            );
+            OtherError::new(errmsg)
+        })?;
+        file.sync_all().map_err(|err| {
+            let errmsg = format!(
+                "Failed to sync the tx-pool persisted data file [{:?}], cause: {}",
+                persisted_data_file, err
+            );
+            OtherError::new(errmsg)
+        })?;
+        Ok(())
+    }
+}

tx-pool/src/pool.rs

@@ -93,14 +93,14 @@ impl TxPool {
         const COMMITTED_HASH_CACHE_SIZE: usize = 100_000;
 
         TxPool {
-            config,
             pending: PendingQueue::new(),
             gap: PendingQueue::new(),
             proposed: ProposedPool::new(config.max_ancestors_count),
             committed_txs_hash_cache: LruCache::new(COMMITTED_HASH_CACHE_SIZE),
             last_txs_updated_at,
             total_tx_size: 0,
             total_tx_cycles: 0,
+            config,
             snapshot,
         }
     }

tx-pool/src/process.rs

@@ -982,10 +982,17 @@ impl TxPoolService {
 
     pub(crate) async fn clear_pool(&mut self, new_snapshot: Arc<Snapshot>) {
         let mut tx_pool = self.tx_pool.write().await;
-        let config = tx_pool.config;
+        let config = tx_pool.config.clone();
         self.last_txs_updated_at = Arc::new(AtomicU64::new(0));
         *tx_pool = TxPool::new(config, new_snapshot, Arc::clone(&self.last_txs_updated_at));
     }
+
+    pub(crate) async fn save_pool(&mut self) {
+        let mut tx_pool = self.tx_pool.write().await;
+        if let Err(err) = tx_pool.save_into_file() {
+            error!("failed to save pool, error: {:?}", err)
+        }
+    }
 }
 
 type PreCheckedTx = (