Implement snapshot service

Author: foriequal0

codechain/run_node.rs

@@ -19,6 +19,7 @@ use std::path::Path;
 use std::sync::{Arc, Weak};
 use std::time::{SystemTime, UNIX_EPOCH};
 
+use ccore::snapshot_notify;
 use ccore::{
     AccountProvider, AccountProviderError, BlockId, ChainNotify, Client, ClientConfig, ClientService, EngineInfo,
     EngineType, Miner, MinerService, Scheme, Stratum, StratumConfig, StratumError, NUM_COLUMNS,
@@ -30,6 +31,7 @@ use ckeystore::KeyStore;
 use clap::ArgMatches;
 use clogger::{self, EmailAlarm, LoggerConfig};
 use cnetwork::{Filters, NetworkConfig, NetworkControl, NetworkService, RoutingTable, SocketAddr};
+use csync::snapshot::Service as SnapshotService;
 use csync::{BlockSyncExtension, BlockSyncSender, TransactionSyncExtension};
 use ctimer::TimerLoop;
 use ctrlc::CtrlC;
@@ -363,6 +365,18 @@ pub fn run_node(matches: &ArgMatches) -> Result<(), String> {
         stratum_start(&config.stratum_config(), &miner, client.client())?
     }
 
+    let _snapshot_service = {
+        if !config.snapshot.disable.unwrap() {
+            let client = client.client();
+            let (tx, rx) = snapshot_notify::create();
+            client.engine().register_snapshot_notify_sender(tx);
+            let service = Arc::new(SnapshotService::new(client, rx, config.snapshot.path.unwrap()));
+            Some(service)
+        } else {
+            None
+        }
+    };
+
     // drop the scheme to free up genesis state.
     drop(scheme);
     client.client().engine().complete_register();

core/src/client/mod.rs

@@ -20,6 +20,7 @@ mod client;
 mod config;
 mod error;
 mod importer;
+pub mod snapshot_notify;
 mod test_client;
 
 pub use self::chain_notify::ChainNotify;

core/src/client/snapshot_notify.rs

@@ -0,0 +1,81 @@
+use ctypes::BlockHash;
+
+use parking_lot::RwLock;
+use std::sync::mpsc::{sync_channel, Receiver, RecvError, SyncSender};
+use std::sync::{Arc, Weak};
+
+pub fn create() -> (NotifySender, NotifyReceiverSource) {
+    let (tx, rx) = sync_channel(1);
+    let tx = Arc::new(RwLock::new(Some(tx)));
+    let tx_weak = Arc::downgrade(&tx);
+    (
+        NotifySender {
+            tx,
+        },
+        NotifyReceiverSource(
+            ReceiverCanceller {
+                tx: tx_weak,
+            },
+            NotifyReceiver {
+                rx,
+            },
+        ),
+    )
+}
+
+pub struct NotifySender {
+    tx: Arc<RwLock<Option<SyncSender<BlockHash>>>>,
+}
+
+impl NotifySender {
+    pub fn notify(&self, block_hash: BlockHash) {
+        let guard = self.tx.read();
+        if let Some(tx) = guard.as_ref() {
+            // TODO: Ignore the error. Receiver thread might be terminated or congested.
+            let _ = tx.try_send(block_hash);
+        } else {
+            // ReceiverCanceller is dropped.
+        }
+    }
+}
+
+pub struct NotifyReceiverSource(pub ReceiverCanceller, pub NotifyReceiver);
+
+/// Dropping this makes the receiver stopped.
+///
+/// `recv()` method  of the `Receiver` will stop and return `RecvError` when corresponding `Sender` is dropped.
+/// This is an inherited behaviour of `std::sync::mpsc::{Sender, Receiver}`.
+/// However, we need another way to stop the `Receiver`, since `Sender` is usually shared throughout our codes.
+/// We can't collect them all and destory one by one. We need a kill switch.
+///
+/// `ReceiverCanceller` holds weak reference to the `Sender`, so it doesn't prohibit the default behaviour.
+/// Then, we can upgrade the weak reference and get the shared reference to `Sender` itself, and manually drop it with this.
+pub struct ReceiverCanceller {
+    tx: Weak<RwLock<Option<SyncSender<BlockHash>>>>,
+}
+
+impl Drop for ReceiverCanceller {
+    fn drop(&mut self) {
+        if let Some(tx) = self.tx.upgrade() {
+            let mut guard = tx.write();
+            if let Some(sender) = guard.take() {
+                drop(sender)
+            }
+        } else {
+            // All NotifySender is dropped. No droppable Sender.
+        }
+    }
+}
+
+/// Receiver is dropped when
+/// 1. There are no NotifySenders out there.
+/// 2. ReceiverCanceller is dropped. See the comment of `ReceiverCanceller`.
+pub struct NotifyReceiver {
+    rx: Receiver<BlockHash>,
+}
+
+impl NotifyReceiver {
+    pub fn recv(&self) -> Result<BlockHash, RecvError> {
+        self.rx.recv()
+    }
+}

core/src/consensus/mod.rs

@@ -51,6 +51,7 @@ use primitives::{Bytes, U256};
 use self::bit_set::BitSet;
 use crate::account_provider::AccountProvider;
 use crate::block::{ExecutedBlock, SealedBlock};
+use crate::client::snapshot_notify::NotifySender as SnapshotNotifySender;
 use crate::client::ConsensusClient;
 use crate::codechain_machine::CodeChainMachine;
 use crate::error::Error;
@@ -265,6 +266,8 @@ pub trait ConsensusEngine: Sync + Send {
 
     fn register_chain_notify(&self, _: &Client) {}
 
+    fn register_snapshot_notify_sender(&self, _sender: SnapshotNotifySender) {}
+
     fn complete_register(&self) {}
 
     fn get_best_block_from_best_proposal_header(&self, header: &HeaderView) -> BlockHash {

core/src/lib.rs

@@ -84,6 +84,7 @@ mod tests;
 
 pub use crate::account_provider::{AccountProvider, Error as AccountProviderError};
 pub use crate::block::Block;
+pub use crate::client::snapshot_notify;
 pub use crate::client::{
     AccountData, AssetClient, BlockChainClient, BlockChainTrait, ChainNotify, Client, ClientConfig, DatabaseClient,
     EngineClient, EngineInfo, ExecuteClient, ImportBlock, MiningBlockChainClient, Shard, StateInfo, TermInfo,

sync/src/lib.rs

@@ -46,6 +46,7 @@ extern crate trie_standardmap;
 extern crate util_error;
 
 mod block;
+pub mod snapshot;
 mod transaction;
 
 pub use crate::block::{BlockSyncEvent, BlockSyncExtension, BlockSyncSender};

sync/src/snapshot/mod.rs

@@ -0,0 +1,103 @@
+// Copyright 2019 Kodebox, Inc.
+// This file is part of CodeChain.
+//
+// This program is free software: you can redistribute it and/or modify
+// it under the terms of the GNU Affero General Public License as
+// published by the Free Software Foundation, either version 3 of the
+// License, or (at your option) any later version.
+//
+// This program is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+// GNU Affero General Public License for more details.
+//
+// You should have received a copy of the GNU Affero General Public License
+// along with this program.  If not, see <https://www.gnu.org/licenses/>.
+
+use std::fs::{create_dir_all, File};
+use std::path::PathBuf;
+use std::sync::Arc;
+use std::thread::{spawn, JoinHandle};
+
+use ccore::snapshot_notify::{NotifyReceiverSource, ReceiverCanceller};
+use ccore::{BlockChainTrait, BlockId, Client};
+use cmerkle::snapshot::{ChunkCompressor, Error as SnapshotError, Snapshot};
+use ctypes::BlockHash;
+use hashdb::{AsHashDB, HashDB};
+use primitives::H256;
+use std::ops::Deref;
+
+pub struct Service {
+    join_handle: Option<JoinHandle<()>>,
+    canceller: Option<ReceiverCanceller>,
+}
+
+impl Service {
+    pub fn new(client: Arc<Client>, notify_receiver_source: NotifyReceiverSource, root_dir: String) -> Self {
+        let NotifyReceiverSource(canceller, receiver) = notify_receiver_source;
+        let join_handle = spawn(move || {
+            cinfo!(SYNC, "Snapshot service is on");
+            while let Ok(block_hash) = receiver.recv() {
+                cinfo!(SYNC, "Snapshot is requested for block: {}", block_hash);
+                let state_root = if let Some(header) = client.block_header(&BlockId::Hash(block_hash)) {
+                    header.state_root()
+                } else {
+                    cerror!(SYNC, "There isn't corresponding header for the requested block hash: {}", block_hash,);
+                    continue
+                };
+                let db_lock = client.state_db().read();
+                if let Some(err) = snapshot(db_lock.as_hashdb(), block_hash, state_root, &root_dir).err() {
+                    cerror!(
+                        SYNC,
+                        "Snapshot request failed for block: {}, chunk_root: {}, err: {}",
+                        block_hash,
+                        state_root,
+                        err
+                    );
+                } else {
+                    cinfo!(SYNC, "Snapshot is ready for block: {}", block_hash)
+                }
+            }
+            cinfo!(SYNC, "Snapshot service is stopped")
+        });
+
+        Self {
+            canceller: Some(canceller),
+            join_handle: Some(join_handle),
+        }
+    }
+}
+
+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)?;
+
+    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_file = File::create(chunk_path)?;
+        let compressor = ChunkCompressor::new(chunk_file);
+        compressor.compress_chunk(&chunk)?;
+    }
+
+    Ok(())
+}
+
+impl Drop for Service {
+    fn drop(&mut self) {
+        if let Some(canceller) = self.canceller.take() {
+            // The thread corresponding to the `self.join_handle` waits for the `self.canceller` is dropped.
+            // It must be dropped first not to make deadlock at `handle.join()`.
+            drop(canceller);
+        }
+
+        if let Some(handle) = self.join_handle.take() {
+            handle.join().expect("Snapshot service thread shouldn't panic");
+        }
+    }
+}

util/merkle/src/snapshot/error.rs

@@ -20,6 +20,7 @@ use primitives::H256;
 use rlp::DecoderError as RlpDecoderError;
 
 use crate::TrieError;
+use std::fmt::{Display, Formatter};
 
 #[derive(Debug)]
 pub enum Error {
@@ -53,6 +54,17 @@ impl From<ChunkError> for Error {
     }
 }
 
+impl Display for Error {
+    fn fmt(&self, f: &mut Formatter<'_>) -> Result<(), std::fmt::Error> {
+        match self {
+            Error::IoError(err) => write!(f, "IoError: {}", err),
+            Error::RlpDecoderError(err) => write!(f, "RlpDecoderError: {}", err),
+            Error::TrieError(err) => write!(f, "TrieError: {}", err),
+            Error::ChunkError(err) => write!(f, "ChunkError: {}", err),
+        }
+    }
+}
+
 #[derive(Debug)]
 pub enum ChunkError {
     TooBig,
@@ -63,3 +75,17 @@ pub enum ChunkError {
     },
     InvalidContent,
 }
+
+impl Display for ChunkError {
+    fn fmt(&self, f: &mut Formatter<'_>) -> Result<(), std::fmt::Error> {
+        match self {
+            ChunkError::TooBig => write!(f, "Chunk has too many elements"),
+            ChunkError::InvalidHeight => write!(f, "Chunk height is unexpected height"),
+            ChunkError::ChunkRootMismatch {
+                expected,
+                actual,
+            } => write!(f, "Chunk root is different from expected. expected: {}, actual: {}", expected, actual),
+            ChunkError::InvalidContent => write!(f, "Chunk content is invalid"),
+        }
+    }
+}

util/merkle/src/snapshot/mod.rs

@@ -26,6 +26,8 @@ use hashdb::HashDB;
 use primitives::H256;
 
 use self::chunk::{Chunk, RecoveredChunk, UnresolvedChunk};
+pub use self::compress::{ChunkCompressor, ChunkDecompressor};
+pub use self::error::Error;
 use self::ordered_heap::OrderedHeap;
 use crate::nibbleslice::NibbleSlice;