Serve snapshot responses to peers

Author: remagpie

codechain/run_node.rs

@@ -302,7 +302,10 @@ pub fn run_node(matches: &ArgMatches) -> Result<(), String> {
                         (Some(hash), Some(num)) => Some((hash, num)),
                         _ => None,
                     };
-                    service.register_extension(move |api| BlockSyncExtension::new(client, api, snapshot_target))
+                    let snapshot_dir = config.snapshot.path.clone();
+                    service.register_extension(move |api| {
+                        BlockSyncExtension::new(client, api, snapshot_target, snapshot_dir)
+                    })
                 };
                 let sync = Arc::new(BlockSyncSender::from(sync_sender.clone()));
                 client.client().add_notify(Arc::downgrade(&sync) as Weak<dyn ChainNotify>);

sync/src/block/extension.rs

@@ -16,6 +16,7 @@
 
 use std::collections::hash_map::Entry;
 use std::collections::{HashMap, HashSet};
+use std::fs;
 use std::sync::Arc;
 use std::time::Duration;
 
@@ -43,6 +44,7 @@ use token_generator::TokenGenerator;
 
 use super::downloader::{BodyDownloader, HeaderDownloader};
 use super::message::{Message, RequestMessage, ResponseMessage};
+use crate::snapshot::snapshot_path;
 
 const SYNC_TIMER_TOKEN: TimerToken = 0;
 const SYNC_EXPIRE_TOKEN_BEGIN: TimerToken = SYNC_TIMER_TOKEN + 1;
@@ -81,10 +83,16 @@ pub struct Extension {
     api: Box<dyn Api>,
     last_request: u64,
     nonce: u64,
+    snapshot_dir: Option<String>,
 }
 
 impl Extension {
-    pub fn new(client: Arc<Client>, api: Box<dyn Api>, snapshot_target: Option<(H256, u64)>) -> Extension {
+    pub fn new(
+        client: Arc<Client>,
+        api: Box<dyn Api>,
+        snapshot_target: Option<(H256, 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 {
@@ -137,6 +145,7 @@ impl Extension {
             api,
             last_request: Default::default(),
             nonce: Default::default(),
+            snapshot_dir,
         }
     }
 
@@ -689,8 +698,18 @@ impl Extension {
         ResponseMessage::Bodies(bodies)
     }
 
-    fn create_state_chunk_response(&self, _hash: BlockHash, _tree_root: Vec<H256>) -> ResponseMessage {
-        unimplemented!()
+    fn create_state_chunk_response(&self, hash: BlockHash, chunk_roots: Vec<H256>) -> ResponseMessage {
+        let mut result = Vec::new();
+        for root in chunk_roots {
+            if let Some(dir) = &self.snapshot_dir {
+                let chunk_path = snapshot_path(&dir, &hash, &root);
+                match fs::read(chunk_path) {
+                    Ok(chunk) => result.push(chunk),
+                    _ => result.push(Vec::new()),
+                }
+            }
+        }
+        ResponseMessage::StateChunk(result)
     }
 
     fn on_peer_response(&mut self, from: &NodeId, id: u64, mut response: ResponseMessage) {

sync/src/snapshot/mod.rs

@@ -32,6 +32,19 @@ pub struct Service {
     canceller: Option<ReceiverCanceller>,
 }
 
+pub fn snapshot_dir(root_dir: &str, block: &BlockHash) -> PathBuf {
+    let mut path = PathBuf::new();
+    path.push(root_dir);
+    path.push(format!("{:x}", block.deref()));
+    path
+}
+
+pub fn snapshot_path(root_dir: &str, block: &BlockHash, chunk_root: &H256) -> PathBuf {
+    let mut path = snapshot_dir(root_dir, block);
+    path.push(format!("{:x}", chunk_root));
+    path
+}
+
 impl Service {
     pub fn new(client: Arc<Client>, notify_receiver_source: NotifyReceiverSource, root_dir: String) -> Self {
         let NotifyReceiverSource(canceller, receiver) = notify_receiver_source;
@@ -70,17 +83,11 @@ impl Service {
 }
 
 fn snapshot(db: &dyn HashDB, block_hash: BlockHash, chunk_root: H256, root_dir: &str) -> Result<(), SnapshotError> {
-    let snapshot_dir = {
-        let mut res = PathBuf::new();
-        res.push(root_dir);
-        res.push(format!("{:x}", block_hash.deref()));
-        res
-    };
-    create_dir_all(&snapshot_dir)?;
+    let snapshot_dir = snapshot_dir(root_dir, &block_hash);
+    create_dir_all(snapshot_dir)?;
 
     for chunk in Snapshot::from_hashdb(db, chunk_root) {
-        let mut chunk_path = snapshot_dir.clone();
-        chunk_path.push(format!("{:x}", chunk.root));
+        let chunk_path = snapshot_path(root_dir, &block_hash, &chunk.root);
         let chunk_file = File::create(chunk_path)?;
         let compressor = ChunkCompressor::new(chunk_file);
         compressor.compress_chunk(&chunk)?;