From bfdc097538862feb1bc17582b782525701a302d5 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Andr=C3=A9=20Silva?= <andre.beat@gmail.com>
Date: Tue, 20 Feb 2018 16:20:30 +0000
Subject: [PATCH] kvdb-rocksdb: remove buffered operations when committing
 transaction (#7950)

---
 util/kvdb-rocksdb/src/lib.rs | 20 ++++++++++++++++++++
 util/kvdb/src/lib.rs         | 20 ++++++++++++++++++++
 2 files changed, 40 insertions(+)

diff --git a/util/kvdb-rocksdb/src/lib.rs b/util/kvdb-rocksdb/src/lib.rs
index f30ceba20..aef448554 100644
--- a/util/kvdb-rocksdb/src/lib.rs
+++ b/util/kvdb-rocksdb/src/lib.rs
@@ -491,6 +491,9 @@ impl Database {
 				let batch = WriteBatch::new();
 				let ops = tr.ops;
 				for op in ops {
+					// remove any buffered operation for this key
+					self.overlay.write()[Self::to_overlay_column(op.col())].remove(op.key());
+
 					match op {
 						DBOp::Insert { col, key, value } => {
 							col.map_or_else(|| batch.put(&key, &value), |c| batch.put_cf(cfs[c as usize], &key, &value))?
@@ -857,4 +860,21 @@ mod tests {
 			assert_eq!(db.num_columns(), 0);
 		}
 	}
+
+	#[test]
+	fn write_clears_buffered_ops() {
+		let tempdir = TempDir::new("").unwrap();
+		let config = DatabaseConfig::default();
+		let db = Database::open(&config, tempdir.path().to_str().unwrap()).unwrap();
+
+		let mut batch = db.transaction();
+		batch.put(None, b"foo", b"bar");
+		db.write_buffered(batch);
+
+		let mut batch = db.transaction();
+		batch.put(None, b"foo", b"baz");
+		db.write(batch).unwrap();
+
+		assert_eq!(db.get(None, b"foo").unwrap().unwrap().as_ref(), b"baz");
+	}
 }
diff --git a/util/kvdb/src/lib.rs b/util/kvdb/src/lib.rs
index c10a392d5..b5acc539a 100644
--- a/util/kvdb/src/lib.rs
+++ b/util/kvdb/src/lib.rs
@@ -67,6 +67,26 @@ pub enum DBOp {
 	}
 }
 
+impl DBOp {
+	/// Returns the key associated with this operation.
+	pub fn key(&self) -> &[u8] {
+		match *self {
+			DBOp::Insert { ref key, .. } => key,
+			DBOp::InsertCompressed { ref key, .. } => key,
+			DBOp::Delete { ref key, .. } => key,
+		}
+	}
+
+	/// Returns the column associated with this operation.
+	pub fn col(&self) -> Option<u32> {
+		match *self {
+			DBOp::Insert { col, .. } => col,
+			DBOp::InsertCompressed { col, .. } => col,
+			DBOp::Delete { col, .. } => col,
+		}
+	}
+}
+
 impl DBTransaction {
 	/// Create new transaction.
 	pub fn new() -> DBTransaction {