diff --git a/util/src/journaldb/overlayrecentdb.rs b/util/src/journaldb/overlayrecentdb.rs index fcc537d5364..3a33ea9b224 100644 --- a/util/src/journaldb/overlayrecentdb.rs +++ b/util/src/journaldb/overlayrecentdb.rs @@ -295,7 +295,7 @@ impl JournalDB for OverlayRecentDB { } // update the overlay for k in overlay_deletions { - journal_overlay.backing_overlay.remove(&OverlayRecentDB::to_short_key(&k)); + journal_overlay.backing_overlay.remove_and_purge(&OverlayRecentDB::to_short_key(&k)); } // apply canon deletions for k in canon_deletions { @@ -303,7 +303,6 @@ impl JournalDB for OverlayRecentDB { try!(batch.delete(&k)); } } - journal_overlay.backing_overlay.purge(); } journal_overlay.journal.remove(&end_era); } diff --git a/util/src/memorydb.rs b/util/src/memorydb.rs index 73435e40a3a..f15394583f3 100644 --- a/util/src/memorydb.rs +++ b/util/src/memorydb.rs @@ -24,6 +24,7 @@ use hashdb::*; use heapsize::*; use std::mem; use std::collections::HashMap; +use std::collections::hash_map::Entry; use std::default::Default; #[derive(Debug,Clone)] @@ -162,6 +163,24 @@ impl MemoryDB { pub fn mem_used(&self) -> usize { self.data.heap_size_of_children() } + + /// Remove an element and delete it from storage if reference count reaches zero. + pub fn remove_and_purge(&mut self, key: &H256) { + if key == &SHA3_NULL_RLP { + return; + } + match self.data.entry(key.clone()) { + Entry::Occupied(mut entry) => + if entry.get().1 == 1 { + entry.remove(); + } else { + entry.get_mut().1 -= 1; + }, + Entry::Vacant(entry) => { + entry.insert((Bytes::new(), -1)); + } + } + } } static NULL_RLP_STATIC: [u8; 1] = [0x80; 1]; @@ -269,3 +288,28 @@ fn memorydb_denote() { assert_eq!(m.get(&hash).unwrap(), b"Hello world!"); } + +#[test] +fn memorydb_remove_and_purge() { + let hello_bytes = b"Hello world!"; + let hello_key = hello_bytes.sha3(); + + let mut m = MemoryDB::new(); + m.remove(&hello_key); + assert_eq!(m.raw(&hello_key).unwrap().1, -1); + m.purge(); + assert_eq!(m.raw(&hello_key).unwrap().1, -1); + m.insert(hello_bytes); + assert_eq!(m.raw(&hello_key).unwrap().1, 0); + m.purge(); + assert_eq!(m.raw(&hello_key), None); + + let mut m = MemoryDB::new(); + m.remove_and_purge(&hello_key); + assert_eq!(m.raw(&hello_key).unwrap().1, -1); + m.insert(hello_bytes); + m.insert(hello_bytes); + assert_eq!(m.raw(&hello_key).unwrap().1, 1); + m.remove_and_purge(&hello_key); + assert_eq!(m.raw(&hello_key), None); +}