From 9c1ee100b72c8aa3b3ebae3a4b38b89e83ae80d3 Mon Sep 17 00:00:00 2001
From: debris <marek.kotewicz@gmail.com>
Date: Wed, 5 Apr 2017 18:40:48 +0700
Subject: [PATCH] AutoFlushingOverlayDatabase

---
 db/src/block_chain_db.rs | 17 +++++++++++---
 db/src/kv/mod.rs         |  2 +-
 db/src/kv/overlaydb.rs   | 50 ++++++++++++++++++++++++++++++++++++++++
 3 files changed, 65 insertions(+), 4 deletions(-)

diff --git a/db/src/block_chain_db.rs b/db/src/block_chain_db.rs
index 8ce06536..44d6cff0 100644
--- a/db/src/block_chain_db.rs
+++ b/db/src/block_chain_db.rs
@@ -14,7 +14,7 @@ use ser::{
 };
 use kv::{
 	KeyValueDatabase, OverlayDatabase, Transaction as DBTransaction, Location, Value, DiskDatabase,
-	DatabaseConfig, MemoryDatabase
+	DatabaseConfig, MemoryDatabase, AutoFlushingOverlayDatabase
 };
 use best_block::BestBlock;
 use {
@@ -58,7 +58,7 @@ impl<'a, T> ForkChain for ForkChainDatabase<'a, T> where T: KeyValueDatabase {
 	}
 }
 
-impl BlockChainDatabase<DiskDatabase> {
+impl BlockChainDatabase<AutoFlushingOverlayDatabase<DiskDatabase>> {
 	pub fn open_at_path<P>(path: P, total_cache: usize) -> Result<Self, Error> where P: AsRef<Path> {
 		fs::create_dir_all(path.as_ref()).map_err(|err| Error::DatabaseError(err.to_string()))?;
 		let mut cfg = DatabaseConfig::with_columns(Some(COL_COUNT));
@@ -74,7 +74,7 @@ impl BlockChainDatabase<DiskDatabase> {
 		cfg.bloom_filters.insert(Some(COL_TRANSACTIONS_META), 32);
 
 		match DiskDatabase::open(cfg, path) {
-			Ok(db) => Ok(Self::open(db)),
+			Ok(db) => Ok(Self::open_with_cache(db)),
 			Err(err) => Err(Error::DatabaseError(err))
 		}
 	}
@@ -92,6 +92,17 @@ impl BlockChainDatabase<MemoryDatabase> {
 	}
 }
 
+impl<T> BlockChainDatabase<AutoFlushingOverlayDatabase<T>> where T: KeyValueDatabase {
+	pub fn open_with_cache(db: T) -> Self {
+		let db = AutoFlushingOverlayDatabase::new(db, 50);
+		let best_block = Self::read_best_block(&db).unwrap_or_default();
+		BlockChainDatabase {
+			best_block: RwLock::new(best_block),
+			db: db,
+		}
+	}
+}
+
 impl<T> BlockChainDatabase<T> where T: KeyValueDatabase {
 	fn read_best_block(db: &T) -> Option<BestBlock> {
 		let best_number = db.get(COL_META.into(), KEY_BEST_BLOCK_NUMBER);
diff --git a/db/src/kv/mod.rs b/db/src/kv/mod.rs
index 4b1bf3f8..7960a10f 100644
--- a/db/src/kv/mod.rs
+++ b/db/src/kv/mod.rs
@@ -7,5 +7,5 @@ mod transaction;
 pub use self::db::KeyValueDatabase;
 pub use self::diskdb::{Database as DiskDatabase, DatabaseConfig, CompactionProfile};
 pub use self::memorydb::{MemoryDatabase, SharedMemoryDatabase};
-pub use self::overlaydb::OverlayDatabase;
+pub use self::overlaydb::{OverlayDatabase, AutoFlushingOverlayDatabase};
 pub use self::transaction::{Transaction, Operation, Location, Key, Value, KeyState};
diff --git a/db/src/kv/overlaydb.rs b/db/src/kv/overlaydb.rs
index 3df1a2de..c7d54f7e 100644
--- a/db/src/kv/overlaydb.rs
+++ b/db/src/kv/overlaydb.rs
@@ -1,3 +1,4 @@
+use parking_lot::Mutex;
 use kv::{Transaction, Location, Value, KeyValueDatabase, MemoryDatabase};
 
 pub struct OverlayDatabase<'a, T> where T: 'a + KeyValueDatabase {
@@ -31,3 +32,52 @@ impl<'a, T> KeyValueDatabase for OverlayDatabase<'a, T> where T: 'a + KeyValueDa
 		}
 	}
 }
+
+pub struct AutoFlushingOverlayDatabase<T> where T: KeyValueDatabase {
+	db: T,
+	overlay: MemoryDatabase,
+	operations: Mutex<usize>,
+	max_operations: usize,
+}
+
+impl<T> AutoFlushingOverlayDatabase<T> where T: KeyValueDatabase {
+	pub fn new(db: T, max_operations: usize) -> Self {
+		AutoFlushingOverlayDatabase {
+			db: db,
+			overlay: MemoryDatabase::default(),
+			operations: Mutex::default(),
+			max_operations: max_operations,
+		}
+	}
+
+	fn flush(&self) -> Result<(), String> {
+		self.db.write(self.overlay.drain_transaction())
+	}
+}
+
+impl<T> KeyValueDatabase for AutoFlushingOverlayDatabase<T> where T: KeyValueDatabase {
+	fn write(&self, tx: Transaction) -> Result<(), String> {
+		let mut operations = self.operations.lock();
+		*operations += 1;
+		self.overlay.write(tx)?;
+		if *operations == self.max_operations {
+			self.flush()?;
+			*operations = 0;
+		}
+		Ok(())
+	}
+
+	fn get(&self, location: Location, key: &[u8]) -> Result<Option<Value>, String> {
+		if self.overlay.is_known(location, key) {
+			self.overlay.get(location, key)
+		} else {
+			self.db.get(location, key)
+		}
+	}
+}
+
+impl<T> Drop for AutoFlushingOverlayDatabase<T> where T: KeyValueDatabase {
+	fn drop(&mut self) {
+		self.flush().expect("Failed to save database");
+	}
+}