util: change the way of calculating the approximate distance (#4017)

src/raftstore/store/util.rs

@@ -1360,6 +1360,43 @@ mod tests {
         }
     }
 
+    #[test]
+    fn test_region_maybe_inaccurate_approximate_size() {
+        let path =
+            TempDir::new("_test_raftstore_region_maybe_inaccurate_approximate_size").expect("");
+        let path_str = path.path().to_str().unwrap();
+        let db_opts = DBOptions::new();
+        let mut cf_opts = ColumnFamilyOptions::new();
+        cf_opts.set_disable_auto_compactions(true);
+        let f = Box::new(RangePropertiesCollectorFactory::default());
+        cf_opts.add_table_properties_collector_factory("tikv.range-collector", f);
+        let cfs_opts = LARGE_CFS
+            .iter()
+            .map(|cf| CFOptions::new(cf, cf_opts.clone()))
+            .collect();
+        let db = rocksdb_util::new_engine_opt(path_str, db_opts, cfs_opts).unwrap();
+
+        let mut cf_size = 0;
+        for i in 0..100 {
+            let k1 = keys::data_key(format!("k1{}", i).as_bytes());
+            let k2 = keys::data_key(format!("k9{}", i).as_bytes());
+            let v = vec![0; 4096];
+            cf_size += k1.len() + k2.len() + v.len() * 2;
+            let cf = db.cf_handle("default").unwrap();
+            db.put_cf(cf, &k1, &v).unwrap();
+            db.put_cf(cf, &k2, &v).unwrap();
+            db.flush_cf(cf, true).unwrap();
+        }
+
+        let region = make_region(1, vec![], vec![]);
+        let size = get_region_approximate_size(&db, &region).unwrap();
+        assert_eq!(size, cf_size as u64);
+
+        let region = make_region(1, b"k2".to_vec(), b"k8".to_vec());
+        let size = get_region_approximate_size(&db, &region).unwrap();
+        assert_eq!(size, 0);
+    }
+
     fn check_data(db: &DB, cfs: &[&str], expected: &[(&[u8], &[u8])]) {
         for cf in cfs {
             let handle = get_cf_handle(db, cf).unwrap();

src/util/rocksdb/properties.rs

@@ -555,10 +555,10 @@ impl RangeProperties {
             Some((_, v)) => v.get(kind),
             None => 0,
         };
-        let mut range = self.offsets.range::<[u8], _>((Included(end), Unbounded));
-        let end_offset = match range.next() {
+        let range = self.offsets.range::<[u8], _>((Unbounded, Included(end)));
+        let end_offset = match range.last() {
             Some((_, v)) => v.get(kind),
-            None => self.offsets.iter().last().map_or(0, |(_, v)| v.get(kind)),
+            None => 0,
         };
         if end_offset < start_offset {
             panic!(
@@ -923,8 +923,8 @@ mod tests {
             ("k", "k", k, k),
             ("a", "k", k, a),
             ("a", "i", i, a),
-            ("e", "h", i, e),
-            ("g", "h", i, e),
+            ("e", "h", e, e),
+            ("b", "h", e, a),
             ("g", "g", i, i),
         ];
         for &(start, end, end_idx, start_idx) in &cases {