Fix limit prefix delete case (#368)

* Fix limit prefix delete case. Warning this uses iter by prefix for
it in rocksdb.

* Comment

* Changes needed for merging, end_prefix returning option in new code and
renaming of iter with prefix.

* reduce cornercase iteration windows size.

* extract common iter prefix delete corner case code.

* Update kvdb-rocksdb/src/lib.rs

Co-Authored-By: Andronik Ordian <[email protected]>

* Remove unused code, there can be call to 'delete_range_cf' on empty range,
but that is fine.

Co-authored-by: Andronik Ordian <[email protected]>

kvdb-memorydb/src/lib.rs

@@ -74,8 +74,11 @@ impl KeyValueDB for InMemory {
 							col.clear();
 						} else {
 							let start_range = Bound::Included(prefix.to_vec());
-							let end_range = Bound::Excluded(kvdb::end_prefix(&prefix[..]));
-							let keys: Vec<_> = col.range((start_range, end_range)).map(|(k, _)| k.clone()).collect();
+							let keys: Vec<_> = if let Some(end_range) = kvdb::end_prefix(&prefix[..]) {
+								col.range((start_range, Bound::Excluded(end_range))).map(|(k, _)| k.clone()).collect()
+							} else {
+								col.range((start_range, Bound::Unbounded)).map(|(k, _)| k.clone()).collect()
+							};
 							for key in keys.into_iter() {
 								col.remove(&key[..]);
 							}

kvdb-rocksdb/src/iter.rs

@@ -107,12 +107,12 @@ where
 		read_lock: RwLockReadGuard<'a, Option<T>>,
 		col: u32,
 		prefix: &[u8],
-		upper_bound: Box<[u8]>,
+		upper_bound: Option<Box<[u8]>>,
 		read_opts: &ReadOptions,
 	) -> Self {
 		Self {
 			inner: Self::new_inner(read_lock, |db| db.iter_with_prefix(col, prefix, read_opts)),
-			upper_bound_prefix: Some(upper_bound),
+			upper_bound_prefix: upper_bound,
 		}
 	}
 

kvdb-rocksdb/src/lib.rs

@@ -443,15 +443,16 @@ impl Database {
 							stats_total_bytes += key.len();
 							batch.delete_cf(cf, &key).map_err(other_io_err)?
 						}
-						DBOp::DeletePrefix { col: _, prefix } => {
-							if !prefix.is_empty() {
-								let end_range = kvdb::end_prefix(&prefix[..]);
-								batch.delete_range_cf(cf, &prefix[..], &end_range[..]).map_err(other_io_err)?;
-							} else {
-								// Deletes all values in the column.
-								let end_range = &[u8::max_value()];
-								batch.delete_range_cf(cf, &[][..], &end_range[..]).map_err(other_io_err)?;
-								batch.delete_cf(cf, &end_range[..]).map_err(other_io_err)?;
+						DBOp::DeletePrefix { col, prefix } => {
+							let end_prefix = kvdb::end_prefix(&prefix[..]);
+							let no_end = end_prefix.is_none();
+							let end_range = end_prefix.unwrap_or_else(|| vec![u8::max_value(); 16]);
+							batch.delete_range_cf(cf, &prefix[..], &end_range[..]).map_err(other_io_err)?;
+							if no_end {
+								let prefix = if prefix.len() > end_range.len() { &prefix[..] } else { &end_range[..] };
+								for (key, _) in self.iter_with_prefix(col, prefix) {
+									batch.delete_cf(cf, &key[..]).map_err(other_io_err)?;
+								}
 							}
 						}
 					};
@@ -517,15 +518,16 @@ impl Database {
 		let optional = if read_lock.is_some() {
 			let mut read_opts = ReadOptions::default();
 			read_opts.set_verify_checksums(false);
-			let end_prefix = kvdb::end_prefix(prefix).into_boxed_slice();
 			// rocksdb doesn't work with an empty upper bound
-			if !end_prefix.is_empty() {
+			let end_prefix = kvdb::end_prefix(prefix).map(|end_prefix| {
+				let end_prefix = end_prefix.into_boxed_slice();
 				// SAFETY: the end_prefix lives as long as the iterator
 				// See `ReadGuardedIterator` definition for more details.
 				unsafe {
 					read_opts.set_iterate_upper_bound(&end_prefix);
 				}
-			}
+				end_prefix
+			});
 			let guarded = iter::ReadGuardedIterator::new_with_prefix(read_lock, col, prefix, end_prefix, &read_opts);
 			Some(guarded)
 		} else {

kvdb-shared-tests/src/lib.rs

@@ -175,7 +175,7 @@ pub fn test_io_stats(db: &dyn KeyValueDB) -> io::Result<()> {
 }
 
 /// The number of columns required to run `test_delete_prefix`.
-pub const DELETE_PREFIX_NUM_COLUMNS: u32 = 5;
+pub const DELETE_PREFIX_NUM_COLUMNS: u32 = 7;
 
 /// A test for `KeyValueDB::delete_prefix`.
 pub fn test_delete_prefix(db: &dyn KeyValueDB) -> io::Result<()> {
@@ -190,6 +190,7 @@ pub fn test_delete_prefix(db: &dyn KeyValueDB) -> io::Result<()> {
 		&[2][..],
 		&[2, 0][..],
 		&[2, 255][..],
+		&[255; 16][..],
 	];
 	let init_db = |ix: u32| -> io::Result<()> {
 		let mut batch = db.transaction();
@@ -199,8 +200,8 @@ pub fn test_delete_prefix(db: &dyn KeyValueDB) -> io::Result<()> {
 		db.write(batch)?;
 		Ok(())
 	};
-	let check_db = |ix: u32, content: [bool; 10]| -> io::Result<()> {
-		let mut state = [true; 10];
+	let check_db = |ix: u32, content: [bool; 11]| -> io::Result<()> {
+		let mut state = [true; 11];
 		for (c, key) in keys.iter().enumerate() {
 			state[c] = db.get(ix, key)?.is_some();
 		}
@@ -209,15 +210,19 @@ pub fn test_delete_prefix(db: &dyn KeyValueDB) -> io::Result<()> {
 	};
 	let tests: [_; DELETE_PREFIX_NUM_COLUMNS as usize] = [
 		// standard
-		(&[1u8][..], [true, true, true, false, false, false, false, true, true, true]),
+		(&[1u8][..], [true, true, true, false, false, false, false, true, true, true, true]),
 		// edge
-		(&[1u8, 255, 255][..], [true, true, true, true, true, true, false, true, true, true]),
+		(&[1u8, 255, 255][..], [true, true, true, true, true, true, false, true, true, true, true]),
 		// none 1
-		(&[1, 2][..], [true, true, true, true, true, true, true, true, true, true]),
+		(&[1, 2][..], [true, true, true, true, true, true, true, true, true, true, true]),
 		// none 2
-		(&[8][..], [true, true, true, true, true, true, true, true, true, true]),
+		(&[8][..], [true, true, true, true, true, true, true, true, true, true, true]),
+		// last value
+		(&[255, 255][..], [true, true, true, true, true, true, true, true, true, true, false]),
+		// last value, limit prefix
+		(&[255][..], [true, true, true, true, true, true, true, true, true, true, false]),
 		// all
-		(&[][..], [false, false, false, false, false, false, false, false, false, false]),
+		(&[][..], [false, false, false, false, false, false, false, false, false, false, false]),
 	];
 	for (ix, test) in tests.iter().enumerate() {
 		let ix = ix as u32;

kvdb/src/lib.rs

@@ -143,15 +143,19 @@ pub trait KeyValueDB: Sync + Send + parity_util_mem::MallocSizeOf {
 
 /// For a given start prefix (inclusive), returns the correct end prefix (non-inclusive).
 /// This assumes the key bytes are ordered in lexicographical order.
-pub fn end_prefix(prefix: &[u8]) -> Vec<u8> {
+/// Since key length is not limited, for some case we return `None` because there is
+/// no bounded limit (every keys in the serie `[]`, `[255]`, `[255, 255]` ...).
+pub fn end_prefix(prefix: &[u8]) -> Option<Vec<u8>> {
 	let mut end_range = prefix.to_vec();
 	while let Some(0xff) = end_range.last() {
 		end_range.pop();
 	}
 	if let Some(byte) = end_range.last_mut() {
 		*byte += 1;
+		Some(end_range)
+	} else {
+		None
 	}
-	end_range
 }
 
 #[cfg(test)]
@@ -160,18 +164,18 @@ mod test {
 
 	#[test]
 	fn end_prefix_test() {
-		assert_eq!(end_prefix(&[5, 6, 7]), vec![5, 6, 8]);
-		assert_eq!(end_prefix(&[5, 6, 255]), vec![5, 7]);
+		assert_eq!(end_prefix(&[5, 6, 7]), Some(vec![5, 6, 8]));
+		assert_eq!(end_prefix(&[5, 6, 255]), Some(vec![5, 7]));
 		// This is not equal as the result is before start.
-		assert_ne!(end_prefix(&[5, 255, 255]), vec![5, 255]);
+		assert_ne!(end_prefix(&[5, 255, 255]), Some(vec![5, 255]));
 		// This is equal ([5, 255] will not be deleted because
 		// it is before start).
-		assert_eq!(end_prefix(&[5, 255, 255]), vec![6]);
-		assert_eq!(end_prefix(&[255, 255, 255]), vec![]);
+		assert_eq!(end_prefix(&[5, 255, 255]), Some(vec![6]));
+		assert_eq!(end_prefix(&[255, 255, 255]), None);
 
-		assert_eq!(end_prefix(&[0x00, 0xff]), vec![0x01]);
-		assert_eq!(end_prefix(&[0xff]), vec![]);
-		assert_eq!(end_prefix(b"0"), b"1".to_vec());
-		assert_eq!(end_prefix(&[]), vec![]);
+		assert_eq!(end_prefix(&[0x00, 0xff]), Some(vec![0x01]));
+		assert_eq!(end_prefix(&[0xff]), None);
+		assert_eq!(end_prefix(&[]), None);
+		assert_eq!(end_prefix(b"0"), Some(b"1".to_vec()));
 	}
 }