diff --git a/kvdb-memorydb/src/lib.rs b/kvdb-memorydb/src/lib.rs index 6a70babea..0ac955a0b 100644 --- a/kvdb-memorydb/src/lib.rs +++ b/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[..]); } diff --git a/kvdb-rocksdb/src/iter.rs b/kvdb-rocksdb/src/iter.rs index a1ef70a53..62f9ba9a2 100644 --- a/kvdb-rocksdb/src/iter.rs +++ b/kvdb-rocksdb/src/iter.rs @@ -107,12 +107,12 @@ where read_lock: RwLockReadGuard<'a, Option>, col: u32, prefix: &[u8], - upper_bound: Box<[u8]>, + upper_bound: Option>, 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, } } diff --git a/kvdb-rocksdb/src/lib.rs b/kvdb-rocksdb/src/lib.rs index 72c96f540..d4e76b4c8 100644 --- a/kvdb-rocksdb/src/lib.rs +++ b/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 { diff --git a/kvdb-shared-tests/src/lib.rs b/kvdb-shared-tests/src/lib.rs index 76fb00c5c..f3352faa0 100644 --- a/kvdb-shared-tests/src/lib.rs +++ b/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; diff --git a/kvdb/src/lib.rs b/kvdb/src/lib.rs index a9c74e949..7cacff666 100644 --- a/kvdb/src/lib.rs +++ b/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 { +/// 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> { 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())); } }