Repository: paritytech/parity-db (language: rust)

Issue:
Values do not seems to be persisted on ref-counted BTree
Here is a failing test I believe should pass:

```rust
	#[test]
	fn fuzz_test() {
		let dir = tempdir().unwrap();
		let options = Options {
			path: dir.path().to_owned(),
			columns: vec![ColumnOptions {
				ref_counted: true,
				btree_index: true,
				..ColumnOptions::default()
			}],
			sync_wal: true,
			sync_data: true,
			stats: false,
			salt: None,
		};

		let db = Db::open_or_create(&options).unwrap();
		db.commit_changes(vec![(0, Operation::Set(vec![0], vec![0]))]).unwrap();

		// The value is there before closing and opening again
		assert_eq!(db.get(0, &[0]).unwrap(), Some(vec![0]));

		// Ensures we should properly join on shutdown
		assert!(db.join_on_shutdown);

		// The value is not there anymore
		drop(db);
		let db = Db::open_read_only(&options).unwrap();
		assert_eq!(db.get(0, &[0]).unwrap(), Some(vec![0]));
	}
```


Discussion / hints:

Fix patch:
diff --git a/src/btree/mod.rs b/src/btree/mod.rs
index a3a5077d..248173b6 100644
--- a/src/btree/mod.rs
+++ b/src/btree/mod.rs
@@ -256,23 +256,6 @@ impl BTreeTable {
 		Ok(())
 	}
 
-	pub fn write_plan(
-		tables: TablesRef,
-		btree: &mut BTree,
-		writer: &mut LogWriter,
-		record_id: u64,
-		btree_header: &mut BTreeHeader,
-	) -> Result<()> {
-		let root = BTree::fetch_root(btree.root_index.unwrap_or(NULL_ADDRESS), tables, writer)?;
-		if let Some(ix) = Self::write_node_plan(tables, root, writer, btree.root_index)? {
-			btree.root_index = Some(ix);
-		}
-		btree.record_id = record_id;
-		btree_header.root = btree.root_index.unwrap_or(NULL_ADDRESS);
-		btree_header.depth = btree.depth;
-		Ok(())
-	}
-
 	fn write_plan_remove_node(
 		tables: TablesRef,
 		writer: &mut LogWriter,
@@ -285,6 +268,7 @@ impl BTreeTable {
 			&Operation::Dereference(()),
 			writer,
 			None,
+			false,
 		)?;
 		Ok(())
 	}
@@ -347,6 +331,7 @@ impl BTreeTable {
 				&Operation::Set((), entry.encoded),
 				writer,
 				None,
+				false,
 			)? {
 				Some(new_index)
 			} else {
@@ -456,7 +441,9 @@ pub mod commit_overlay {
 			self.changes.sort();
 			tree.write_sorted_changes(self.changes.as_slice(), locked, writer)?;
 			*ops += self.changes.len() as u64;
-			BTreeTable::write_plan(locked, &mut tree, writer, record_id, &mut btree_header)?;
+
+			btree_header.root = tree.root_index.unwrap_or(NULL_ADDRESS);
+			btree_header.depth = tree.depth;
 
 			if old_btree_header != btree_header {
 				let mut entry = Entry::empty();
@@ -468,8 +455,11 @@ pub mod commit_overlay {
 					&Operation::Set((), &entry.encoded.as_ref()[..HEADER_SIZE as usize]),
 					writer,
 					None,
+					false,
 				)?;
 			}
+			#[cfg(test)]
+			tree.is_balanced(locked, writer)?;
 			Ok(())
 		}
 	}
diff --git a/src/btree/node.rs b/src/btree/node.rs
index 64f05a8e..c5b34c2d 100644
--- a/src/btree/node.rs
+++ b/src/btree/node.rs
@@ -64,23 +64,13 @@ impl Node {
 		log: &mut LogWriter,
 	) -> Result<(Option<(Separator, Child)>, bool)> {
 		loop {
-			if changes.len() > 1 && !changes[0].is_reference_ops() {
-				let mut skip = false;
-				for next in changes[1..].iter() {
-					if changes[0].key() == next.key() {
-						if !next.is_reference_ops() {
-							skip = true;
-							break
-						}
-					} else {
-						break
-					}
-				}
-				if skip {
-					// TODO only when advancing (here rec call useless)
-					*changes = &changes[1..];
-					continue
-				}
+			if !btree.ref_counted &&
+				changes.len() > 1 &&
+				!matches!(changes[1], Operation::Reference(..)) &&
+				changes[0].key() == changes[1].key()
+			{
+				*changes = &changes[1..];
+				continue
 			}
 			let r = match &changes[0] {
 				Operation::Set(key, value) =>
@@ -260,30 +250,35 @@ impl Node {
 		if at {
 			let existing = self.separator_address(i);
 			if let Some(existing) = existing {
-				Column::write_existing_value_plan::<_, Vec<u8>>(
+				if Column::write_existing_value_plan::<_, Vec<u8>>(
 					&TableKey::NoHash,
 					values,
 					existing,
 					change,
 					log,
 					None,
-				)?;
-			}
-			let _ = self.remove_separator(i);
-			if depth != 0 {
-				// replace by bigger value in left child.
-				if let Some(mut child) = self.fetch_child(i, values, log)? {
-					let (need_balance, sep) = child.remove_last(depth - 1, values, log)?;
-					self.write_child(i, child, values, log)?;
-					if let Some(sep) = sep {
-						self.set_separator(i, sep);
-					}
-					if need_balance {
-						self.rebalance(depth, i, values, log)?;
+					values.ref_counted,
+				)?
+				.0
+				.is_none()
+				{
+					let _ = self.remove_separator(i);
+					if depth != 0 {
+						// replace by bigger value in left child.
+						if let Some(mut child) = self.fetch_child(i, values, log)? {
+							let (need_balance, sep) = child.remove_last(depth - 1, values, log)?;
+							self.write_child(i, child, values, log)?;
+							if let Some(sep) = sep {
+								self.set_separator(i, sep);
+							}
+							if need_balance {
+								self.rebalance(depth, i, values, log)?;
+							}
+						}
+					} else {
+						self.remove_from(i, false, true);
 					}
 				}
-			} else {
-				self.remove_from(i, false, true);
 			}
 		} else {
 			if !has_child {
@@ -665,12 +660,6 @@ pub struct Child {
 }
 
 impl Node {
-	pub fn clear(&mut self) {
-		self.separators = Default::default();
-		self.children = Default::default();
-		self.changed = true;
-	}
-
 	pub fn from_encoded(enc: Vec<u8>) -> Self {
 		let mut entry = Entry::from_encoded(enc);
 		let mut node =
@@ -757,13 +746,15 @@ impl Node {
 				&Operation::Set((), value),
 				log,
 				None,
+				btree.ref_counted,
 			)?
 			.1
 			.unwrap_or(address)
 		} else {
 			Column::write_new_value_plan(&TableKey::NoHash, btree, value, log, None)?
 		};
-		Ok(Separator { modified: true, separator: Some(SeparatorInner { key, value }) })
+		let modified = Some(value) != existing;
+		Ok(Separator { modified, separator: Some(SeparatorInner { key, value }) })
 	}
 
 	fn split(
diff --git a/src/column.rs b/src/column.rs
index 4ce41186..5348c92c 100644
--- a/src/column.rs
+++ b/src/column.rs
@@ -520,6 +520,7 @@ impl HashColumn {
 			change,
 			log,
 			stats,
+			self.ref_counted,
 		)? {
 			(Some(outcome), _) => Ok(outcome),
 			(None, Some(value_address)) => {
@@ -886,6 +887,7 @@ impl Column {
 		change: &Operation<K, V>,
 		log: &mut LogWriter,
 		stats: Option<&ColumnStats>,
+		ref_counted: bool,
 	) -> Result<(Option<PlanOutcome>, Option<Address>)> {
 		let tier = address.size_tier() as usize;
 
@@ -908,7 +910,7 @@ impl Column {
 
 		match change {
 			Operation::Reference(_) =>
-				if tables.ref_counted {
+				if ref_counted {
 					log::trace!(target: "parity-db", "{}: Increment ref {}", tables.col, key);
 					tables.tables[tier].write_inc_ref(address.offset(), log)?;
 					if let Some(stats) = stats {
@@ -919,7 +921,7 @@ impl Column {
 					Ok((Some(PlanOutcome::Skipped), None))
 				},
 			Operation::Set(_, val) => {
-				if tables.ref_counted {
+				if ref_counted {
 					log::trace!(target: "parity-db", "{}: Increment ref {}", tables.col, key);
 					tables.tables[tier].write_inc_ref(address.offset(), log)?;
 					return Ok((Some(PlanOutcome::Written), None))
@@ -967,7 +969,7 @@ impl Column {
 			Operation::Dereference(_) => {
 				// Deletion
 				let cur_size = if stats.is_some() { Some(fetch_size()?) } else { None };
-				let remove = if tables.ref_counted {
+				let remove = if ref_counted {
 					let removed = !tables.tables[tier].write_dec_ref(address.offset(), log)?;
 					log::trace!(target: "parity-db", "{}: Dereference {}, deleted={}", tables.col, key, removed);
 					removed
diff --git a/src/db.rs b/src/db.rs
index 651ed0b9..eae116de 100644
--- a/src/db.rs
+++ b/src/db.rs
@@ -1169,13 +1169,6 @@ impl<Key, Value> Operation<Key, Value> {
 			Operation::Set(k, _) | Operation::Dereference(k) | Operation::Reference(k) => k,
 		}
 	}
-
-	pub fn is_reference_ops(&self) -> bool {
-		match self {
-			Operation::Set(..) | Operation::Dereference(..) => false,
-			Operation::Reference(..) => true,
-		}
-	}
 }
 
 impl<K: AsRef<[u8]>, Value> Operation<K, Value> {
@@ -1824,25 +1817,36 @@ mod tests {
 	}
 
 	fn test_basic(change_set: &[(Vec<u8>, Option<Vec<u8>>)]) {
+		test_basic_inner(change_set, false, false);
+		test_basic_inner(change_set, false, true);
+		test_basic_inner(change_set, true, false);
+		test_basic_inner(change_set, true, true);
+	}
+
+	fn test_basic_inner(
+		change_set: &[(Vec<u8>, Option<Vec<u8>>)],
+		btree_index: bool,
+		ref_counted: bool,
+	) {
 		let tmp = tempdir().unwrap();
 		let col_nb = 1u8;
 		let mut options = Options::with_columns(tmp.path(), 2);
-		options.columns[col_nb as usize].btree_index = true;
+		options.columns[col_nb as usize].btree_index = btree_index;
+		options.columns[col_nb as usize].ref_counted = ref_counted;
 		let db_test = EnableCommitPipelineStages::DbFile;
+		// ref counted and commit overlay currently don't support removal
+		assert!(!(ref_counted && matches!(db_test, EnableCommitPipelineStages::CommitOverlay)));
 		let inner_options =
 			InternalOptions { create: true, commit_stages: db_test, ..Default::default() };
 		let db = Db::open_inner(&options, &inner_options).unwrap();
 
-		let mut iter = db.iter(col_nb).unwrap();
-		assert_eq!(iter.next().unwrap(), None);
+		let iter = btree_index.then(|| db.iter(col_nb).unwrap());
+		assert_eq!(iter.and_then(|mut i| i.next().unwrap()), None);
 
-		db.commit(change_set.iter().map(|(k, v)| (0, k.clone(), v.clone()))).unwrap();
 		db.commit(change_set.iter().map(|(k, v)| (col_nb, k.clone(), v.clone())))
 			.unwrap();
 		db_test.run_stages(&db);
 
-		let stats = db.stats();
-		let col_stats = stats.columns.into_iter().next().unwrap().unwrap();
 		let mut keys = HashSet::new();
 		let mut expected_count: u64 = 0;
 		for (k, v) in change_set.iter() {
@@ -1854,12 +1858,47 @@ mod tests {
 				expected_count -= 1;
 			}
 		}
-		assert_eq!(col_stats.total_values, expected_count);
+		if ref_counted {
+			let mut state: BTreeMap<Vec<u8>, Option<(Vec<u8>, usize)>> = Default::default();
+			for (k, v) in change_set.iter() {
+				let mut remove = false;
+				let mut insert = false;
+				match state.get_mut(k) {
+					Some(Some((_, counter))) =>
+						if v.is_some() {
+							*counter += 1;
+						} else if *counter == 1 {
+							remove = true;
+						} else {
+							*counter -= 1;
+						},
+					Some(None) | None =>
+						if v.is_some() {
+							insert = true;
+						},
+				}
+				if insert {
+					state.insert(k.clone(), v.clone().map(|v| (v, 1)));
+				}
+				if remove {
+					state.remove(k);
+				}
+			}
+			for (key, value) in state {
+				assert_eq!(db.get(col_nb, &key).unwrap(), value.map(|v| v.0));
+			}
+		} else {
+			let stats = db.stats();
+			// btree do not have stats implemented
+			if let Some(stats) = stats.columns[col_nb as usize].as_ref() {
+				assert_eq!(stats.total_values, expected_count);
+			}
 
-		let state: BTreeMap<Vec<u8>, Option<Vec<u8>>> =
-			change_set.iter().map(|(k, v)| (k.clone(), v.clone())).collect();
-		for (key, value) in state.iter() {
-			assert_eq!(&db.get(col_nb, key).unwrap(), value);
+			let state: BTreeMap<Vec<u8>, Option<Vec<u8>>> =
+				change_set.iter().map(|(k, v)| (k.clone(), v.clone())).collect();
+			for (key, value) in state.iter() {
+				assert_eq!(&db.get(col_nb, key).unwrap(), value);
+			}
 		}
 	}
 
@@ -1904,6 +1943,22 @@ mod tests {
 
 	#[test]
 	fn test_simple() {
+		test_basic(&[
+			(b"key1".to_vec(), Some(b"value1".to_vec())),
+			(b"key1".to_vec(), Some(b"value1".to_vec())),
+			(b"key1".to_vec(), None),
+		]);
+		test_basic(&[
+			(b"key1".to_vec(), Some(b"value1".to_vec())),
+			(b"key1".to_vec(), Some(b"value1".to_vec())),
+			(b"key1".to_vec(), None),
+			(b"key1".to_vec(), None),
+		]);
+		test_basic(&[
+			(b"key1".to_vec(), Some(b"value1".to_vec())),
+			(b"key1".to_vec(), Some(b"value2".to_vec())),
+		]);
+		test_basic(&[(b"key1".to_vec(), Some(b"value1".to_vec()))]);
 		test_basic(&[
 			(b"key1".to_vec(), Some(b"value1".to_vec())),
 			(b"key2".to_vec(), Some(b"value2".to_vec())),


Test patch: