Add sync state SnapshotBody

Author: remagpie

codechain/run_node.rs

@@ -35,6 +35,7 @@ use csync::snapshot::Service as SnapshotService;
 use csync::{BlockSyncExtension, BlockSyncSender, TransactionSyncExtension};
 use ctimer::TimerLoop;
 use ctrlc::CtrlC;
+use ctypes::BlockHash;
 use fdlimit::raise_fd_limit;
 use kvdb::KeyValueDB;
 use kvdb_rocksdb::{Database, DatabaseConfig};
@@ -303,7 +304,7 @@ pub fn run_node(matches: &ArgMatches) -> Result<(), String> {
                 let sync_sender = {
                     let client = client.client();
                     let snapshot_target = match (config.network.snapshot_hash, config.network.snapshot_number) {
-                        (Some(hash), Some(num)) => Some((hash, num)),
+                        (Some(hash), Some(num)) => Some((BlockHash::from(hash), num)),
                         _ => None,
                     };
                     let snapshot_dir = config.snapshot.path.clone();

sync/src/block/extension.rs

@@ -63,6 +63,10 @@ pub struct TokenInfo {
 #[derive(Debug)]
 enum State {
     SnapshotHeader(BlockHash, u64),
+    SnapshotBody {
+        block: BlockHash,
+        prev_root: H256,
+    },
     SnapshotChunk {
         block: BlockHash,
         restore: SnapshotRestore,
@@ -90,28 +94,12 @@ impl Extension {
     pub fn new(
         client: Arc<Client>,
         api: Box<dyn Api>,
-        snapshot_target: Option<(H256, u64)>,
+        snapshot_target: Option<(BlockHash, u64)>,
         snapshot_dir: Option<String>,
     ) -> Extension {
         api.set_timer(SYNC_TIMER_TOKEN, Duration::from_millis(SYNC_TIMER_INTERVAL)).expect("Timer set succeeds");
 
-        let state = match snapshot_target {
-            Some((hash, num)) => match client.block_header(&BlockId::Number(num)) {
-                Some(ref header) if *header.hash() == hash => {
-                    let state_db = client.state_db().read();
-                    let state_root = header.state_root();
-                    match TrieFactory::readonly(state_db.as_hashdb(), &state_root) {
-                        Ok(ref trie) if trie.is_complete() => State::Full,
-                        _ => State::SnapshotChunk {
-                            block: hash.into(),
-                            restore: SnapshotRestore::new(state_root),
-                        },
-                    }
-                }
-                _ => State::SnapshotHeader(hash.into(), num),
-            },
-            None => State::Full,
-        };
+        let state = Extension::initial_state(client.clone(), snapshot_target);
         cdebug!(SYNC, "Initial state is {:?}", state);
         let mut header = client.best_header();
         let mut hollow_headers = vec![header.decode()];
@@ -149,6 +137,36 @@ impl Extension {
         }
     }
 
+    fn initial_state(client: Arc<Client>, snapshot_target: Option<(BlockHash, u64)>) -> State {
+        let (hash, num) = match snapshot_target {
+            Some(target) => target,
+            None => return State::Full,
+        };
+        let header = match client.block_header(&num.into()) {
+            Some(ref h) if h.hash() == hash => h.clone(),
+            _ => return State::SnapshotHeader(hash, num),
+        };
+        if client.block_body(&hash.into()).is_none() {
+            let parent_hash = header.parent_hash();
+            let parent =
+                client.block_header(&parent_hash.into()).expect("Parent header of the snapshot header must exist");
+            return State::SnapshotBody {
+                block: hash,
+                prev_root: parent.state_root(),
+            }
+        }
+
+        let state_db = client.state_db().read();
+        let state_root = header.state_root();
+        match TrieFactory::readonly(state_db.as_hashdb(), &state_root) {
+            Ok(ref trie) if trie.is_complete() => State::Full,
+            _ => State::SnapshotChunk {
+                block: hash,
+                restore: SnapshotRestore::new(state_root),
+            },
+        }
+    }
+
     fn dismiss_request(&mut self, id: &NodeId, request_id: u64) {
         if let Some(requests) = self.requests.get_mut(id) {
             requests.retain(|(i, _)| *i != request_id);
@@ -395,6 +413,9 @@ impl NetworkExtension<Event> for Extension {
                             });
                         }
                     }
+                    State::SnapshotBody {
+                        ..
+                    } => unimplemented!(),
                     State::SnapshotChunk {
                         block,
                         ref mut restore,
@@ -509,6 +530,9 @@ impl Extension {
                     None
                 }
             }
+            State::SnapshotBody {
+                ..
+            } => None,
             State::SnapshotChunk {
                 ..
             } => None,
@@ -562,6 +586,9 @@ impl Extension {
                     None
                 }
             }
+            State::SnapshotBody {
+                ..
+            } => unimplemented!(),
             State::SnapshotChunk {
                 ..
             } => None,
@@ -855,6 +882,9 @@ impl Extension {
                     headers.len()
                 ),
             },
+            State::SnapshotBody {
+                ..
+            } => {}
             State::SnapshotChunk {
                 ..
             } => {}