Memory-based pruning history size (#4114)

* prune states based on memory param

* pruning memory CLI and usage in sync

* return purged value from memorydb

* calculate memory used incrementally in overlayrecentdb

* refactor shared history pruning code in client

* Fixed usage alignment

* journal_size function for fast memory calculation

util/src/journaldb/overlayrecentdb.rs

@@ -71,6 +71,7 @@ struct JournalOverlay {
 	journal: HashMap<u64, Vec<JournalEntry>>,
 	latest_era: Option<u64>,
 	earliest_era: Option<u64>,
+	cumulative_size: usize, // cumulative size of all entries.
 }
 
 #[derive(PartialEq)]
@@ -127,7 +128,8 @@ impl OverlayRecentDB {
 		journal_overlay.backing_overlay == reconstructed.backing_overlay &&
 		journal_overlay.pending_overlay == reconstructed.pending_overlay &&
 		journal_overlay.journal == reconstructed.journal &&
-		journal_overlay.latest_era == reconstructed.latest_era
+		journal_overlay.latest_era == reconstructed.latest_era &&
+		journal_overlay.cumulative_size == reconstructed.cumulative_size
 	}
 
 	fn payload(&self, key: &H256) -> Option<DBValue> {
@@ -140,6 +142,7 @@ impl OverlayRecentDB {
 		let mut count = 0;
 		let mut latest_era = None;
 		let mut earliest_era = None;
+		let mut cumulative_size = 0;
 		if let Some(val) = db.get(col, &LATEST_ERA_KEY).expect("Low-level database error.") {
 			let mut era = decode::<u64>(&val);
 			latest_era = Some(era);
@@ -161,7 +164,14 @@ impl OverlayRecentDB {
 					for r in insertions.iter() {
 						let k: H256 = r.val_at(0);
 						let v = r.at(1).data();
-						overlay.emplace(to_short_key(&k), DBValue::from_slice(v));
+
+						let short_key = to_short_key(&k);
+
+						if !overlay.contains(&short_key) {
+							cumulative_size += v.len();
+						}
+
+						overlay.emplace(short_key, DBValue::from_slice(v));
 						inserted_keys.push(k);
 						count += 1;
 					}
@@ -186,6 +196,7 @@ impl OverlayRecentDB {
 			journal: journal,
 			latest_era: latest_era,
 			earliest_era: earliest_era,
+			cumulative_size: cumulative_size,
 		}
 	}
 
@@ -207,12 +218,19 @@ impl JournalDB for OverlayRecentDB {
 	fn mem_used(&self) -> usize {
 		let mut mem = self.transaction_overlay.mem_used();
 		let overlay = self.journal_overlay.read();
+
 		mem += overlay.backing_overlay.mem_used();
 		mem += overlay.pending_overlay.heap_size_of_children();
 		mem += overlay.journal.heap_size_of_children();
+
 		mem
 	}
 
+	fn journal_size(&self) -> usize {
+		self.journal_overlay.read().cumulative_size
+
+	}
+
 	fn is_empty(&self) -> bool {
 		self.backing.get(self.column, &LATEST_ERA_KEY).expect("Low level database error").is_none()
 	}
@@ -256,7 +274,13 @@ impl JournalDB for OverlayRecentDB {
 			r.begin_list(2);
 			r.append(&k);
 			r.append(&&*v);
-			journal_overlay.backing_overlay.emplace(to_short_key(&k), v);
+
+			let short_key = to_short_key(&k);
+			if !journal_overlay.backing_overlay.contains(&short_key) {
+				journal_overlay.cumulative_size += v.len();
+			}
+
+			journal_overlay.backing_overlay.emplace(short_key, v);
 		}
 		r.append(&removed_keys);
 
@@ -267,6 +291,7 @@ impl JournalDB for OverlayRecentDB {
 		k.append(&&PADDING[..]);
 		batch.put_vec(self.column, &k.drain(), r.out());
 		if journal_overlay.latest_era.map_or(true, |e| now > e) {
+			trace!(target: "journaldb", "Set latest era to {}", now);
 			batch.put_vec(self.column, &LATEST_ERA_KEY, encode(&now).to_vec());
 			journal_overlay.latest_era = Some(now);
 		}
@@ -322,7 +347,9 @@ impl JournalDB for OverlayRecentDB {
 			}
 			// update the overlay
 			for k in overlay_deletions {
-				journal_overlay.backing_overlay.remove_and_purge(&to_short_key(&k));
+				if let Some(val) = journal_overlay.backing_overlay.remove_and_purge(&to_short_key(&k)) {
+					journal_overlay.cumulative_size -= val.len();
+				}
 			}
 			// apply canon deletions
 			for k in canon_deletions {
@@ -332,6 +359,10 @@ impl JournalDB for OverlayRecentDB {
 			}
 		}
 		journal_overlay.journal.remove(&end_era);
+		if !journal_overlay.journal.is_empty() {
+			trace!(target: "journaldb", "Set earliest_era to {}", end_era + 1);
+			journal_overlay.earliest_era = Some(end_era + 1);
+		}
 
 		Ok(ops as u32)
 	}

util/src/journaldb/traits.rs

@@ -29,6 +29,11 @@ pub trait JournalDB: HashDB {
 	/// Returns heap memory size used
 	fn mem_used(&self) -> usize;
 
+	/// Returns the size of journalled state in memory.
+	/// This function has a considerable speed requirement --
+	/// it must be fast enough to call several times per block imported.
+	fn journal_size(&self) -> usize { 0 }
+
 	/// Check if this database has any commits
 	fn is_empty(&self) -> bool;
 

util/src/memorydb.rs

@@ -133,19 +133,22 @@ impl MemoryDB {
 	}
 
 	/// Remove an element and delete it from storage if reference count reaches zero.
-	pub fn remove_and_purge(&mut self, key: &H256) {
+	/// If the value was purged, return the old value.
+	pub fn remove_and_purge(&mut self, key: &H256) -> Option<DBValue> {
 		if key == &SHA3_NULL_RLP {
-			return;
+			return None;
 		}
 		match self.data.entry(key.clone()) {
 			Entry::Occupied(mut entry) =>
 				if entry.get().1 == 1 {
-					entry.remove();
+					Some(entry.remove().0)
 				} else {
 					entry.get_mut().1 -= 1;
+					None
 				},
 			Entry::Vacant(entry) => {
 				entry.insert((DBValue::new(), -1));
+				None
 			}
 		}
 	}
@@ -265,13 +268,14 @@ fn memorydb_remove_and_purge() {
 	assert_eq!(m.raw(&hello_key), None);
 
 	let mut m = MemoryDB::new();
-	m.remove_and_purge(&hello_key);
+	assert!(m.remove_and_purge(&hello_key).is_none());
 	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.remove_and_purge(&hello_key).unwrap(), hello_bytes);
 	assert_eq!(m.raw(&hello_key), None);
+	assert!(m.remove_and_purge(&hello_key).is_none());
 }
 
 #[test]