Add reference-counted TablePropertiesCollection bindings (#425)

* Add reference counted TableProperties bindings
Signed-off-by: Brian Anderson <andersrb@gmail.com>

* Add test for table_properties_rc
Signed-off-by: Brian Anderson <andersrb@gmail.com>

* rustfmt
Signed-off-by: Brian Anderson <andersrb@gmail.com>

* Fix a clippy lint
Signed-off-by: Brian Anderson <andersrb@gmail.com>

* Various changes from review feedback
Signed-off-by: Brian Anderson <andersrb@gmail.com>

src/lib.rs

@@ -70,6 +70,8 @@ mod table_filter;
 mod table_properties;
 mod table_properties_collector;
 mod table_properties_collector_factory;
+pub mod table_properties_rc;
+mod table_properties_rc_handles;
 mod titan;
 
 #[cfg(test)]

src/rocksdb.rs

@@ -15,7 +15,7 @@
 use crocksdb_ffi::{
     self, DBBackupEngine, DBCFHandle, DBCache, DBCompressionType, DBEnv, DBInstance, DBMapProperty,
     DBPinnableSlice, DBSequentialFile, DBStatisticsHistogramType, DBStatisticsTickerType,
-    DBTitanDBOptions, DBWriteBatch,
+    DBTablePropertiesCollection, DBTitanDBOptions, DBWriteBatch,
 };
 use libc::{self, c_char, c_int, c_void, size_t};
 use librocksdb_sys::DBMemoryAllocator;
@@ -40,6 +40,7 @@ use std::sync::Arc;
 use std::{fs, ptr, slice};
 
 use table_properties::{TableProperties, TablePropertiesCollection};
+use table_properties_rc::TablePropertiesCollection as RcTablePropertiesCollection;
 use titan::TitanDBOptions;
 
 pub struct CFHandle {
@@ -1667,6 +1668,13 @@ impl DB {
         }
     }
 
+    pub fn get_properties_of_all_tables_rc(&self) -> Result<RcTablePropertiesCollection, String> {
+        unsafe {
+            let props = ffi_try!(crocksdb_get_properties_of_all_tables(self.inner));
+            Ok(RcTablePropertiesCollection::new(props))
+        }
+    }
+
     pub fn get_properties_of_all_tables_cf(
         &self,
         cf: &CFHandle,
@@ -1684,6 +1692,34 @@ impl DB {
         cf: &CFHandle,
         ranges: &[Range],
     ) -> Result<TablePropertiesCollection, String> {
+        // Safety: transfers ownership of new non-null pointer
+        unsafe {
+            let props = self.get_properties_of_tables_in_range_common(cf, ranges)?;
+            Ok(TablePropertiesCollection::from_raw(props))
+        }
+    }
+
+    /// Like `get_properties_of_table_in_range` but the returned family
+    /// of types don't contain any lifetimes. This is suitable for wrapping
+    /// in further abstractions without needing abstract associated lifetime
+    /// parameters. Used by tikv's `engine_rocks`.
+    pub fn get_properties_of_tables_in_range_rc(
+        &self,
+        cf: &CFHandle,
+        ranges: &[Range],
+    ) -> Result<RcTablePropertiesCollection, String> {
+        // Safety: transfers ownership of new non-null pointer
+        unsafe {
+            let props = self.get_properties_of_tables_in_range_common(cf, ranges)?;
+            Ok(RcTablePropertiesCollection::new(props))
+        }
+    }
+
+    fn get_properties_of_tables_in_range_common(
+        &self,
+        cf: &CFHandle,
+        ranges: &[Range],
+    ) -> Result<*mut DBTablePropertiesCollection, String> {
         let start_keys: Vec<*const u8> = ranges.iter().map(|x| x.start_key.as_ptr()).collect();
         let start_keys_lens: Vec<_> = ranges.iter().map(|x| x.start_key.len()).collect();
         let limit_keys: Vec<*const u8> = ranges.iter().map(|x| x.end_key.as_ptr()).collect();
@@ -1698,7 +1734,7 @@ impl DB {
                 limit_keys.as_ptr(),
                 limit_keys_lens.as_ptr()
             ));
-            Ok(TablePropertiesCollection::from_raw(props))
+            Ok(props)
         }
     }
 

src/table_properties_rc.rs

+//! This provides reference-counted abstractions around table properties
+//! collections. It is used by tikv in its own engine abstractions, to avoid the
+//! complexities of lifetimes in associated types.
+//!
+//! FIXME: Safety - While this does guarantee that all the types created from
+//! the collection stay valid for the lifetime of the collection, it doesn't
+//! guarantee that the _DB_ stays valid for the lifetime of the collection.
+
+use crocksdb_ffi::{DBTablePropertiesCollection, DBTableProperty};
+use libc::size_t;
+use librocksdb_sys as crocksdb_ffi;
+use std::ops::Deref;
+use std::slice;
+use std::str;
+
+use crate::table_properties_rc_handles::{
+    TablePropertiesCollectionHandle, TablePropertiesCollectionIteratorHandle,
+    TablePropertiesHandle, UserCollectedPropertiesHandle,
+};
+
+pub struct TablePropertiesCollection {
+    handle: TablePropertiesCollectionHandle,
+}
+
+impl TablePropertiesCollection {
+    pub unsafe fn new(ptr: *mut DBTablePropertiesCollection) -> TablePropertiesCollection {
+        assert!(!ptr.is_null());
+        TablePropertiesCollection {
+            handle: TablePropertiesCollectionHandle::new(ptr),
+        }
+    }
+
+    pub fn iter(&self) -> TablePropertiesCollectionIter {
+        TablePropertiesCollectionIter::new(self.handle.clone())
+    }
+
+    pub fn len(&self) -> usize {
+        unsafe { crocksdb_ffi::crocksdb_table_properties_collection_len(self.handle.ptr()) }
+    }
+
+    pub fn is_empty(&self) -> bool {
+        self.len() == 0
+    }
+}
+
+pub struct TablePropertiesCollectionIter {
+    handle: TablePropertiesCollectionIteratorHandle,
+}
+
+impl TablePropertiesCollectionIter {
+    fn new(collection: TablePropertiesCollectionHandle) -> TablePropertiesCollectionIter {
+        TablePropertiesCollectionIter {
+            handle: TablePropertiesCollectionIteratorHandle::new(collection),
+        }
+    }
+}
+
+impl Iterator for TablePropertiesCollectionIter {
+    type Item = (TablePropertiesKey, TableProperties);
+
+    fn next(&mut self) -> Option<Self::Item> {
+        unsafe {
+            loop {
+                if !crocksdb_ffi::crocksdb_table_properties_collection_iter_valid(self.handle.ptr())
+                {
+                    return None;
+                }
+
+                let mut keylen: size_t = 0;
+                let key = crocksdb_ffi::crocksdb_table_properties_collection_iter_key(
+                    self.handle.ptr(),
+                    &mut keylen,
+                );
+                let props = crocksdb_ffi::crocksdb_table_properties_collection_iter_value(
+                    self.handle.ptr(),
+                );
+                crocksdb_ffi::crocksdb_table_properties_collection_iter_next(self.handle.ptr());
+                if !props.is_null() {
+                    assert!(!key.is_null() && keylen != 0);
+                    let key = TablePropertiesKey::new(key, keylen, self.handle.clone());
+                    let props_handle = TablePropertiesHandle::new(props, self.handle.clone());
+                    let val = TableProperties::new(props_handle);
+                    return Some((key, val));
+                }
+            }
+        }
+    }
+}
+
+// # Safety
+//
+// The underlying iterator is over an unordered map of heap-allocated strings,
+// so as long as the iterator and collection are alive, the key pointers are
+// valid.
+pub struct TablePropertiesKey {
+    key: *const u8,
+    keylen: size_t,
+    _iter_handle: TablePropertiesCollectionIteratorHandle,
+}
+
+impl TablePropertiesKey {
+    fn new(
+        key: *const u8,
+        keylen: size_t,
+        _iter_handle: TablePropertiesCollectionIteratorHandle,
+    ) -> TablePropertiesKey {
+        // Caller must ensure slice is valid
+        unsafe {
+            let bytes = slice::from_raw_parts(key, keylen);
+            assert!(str::from_utf8(bytes).is_ok());
+        }
+        TablePropertiesKey {
+            key,
+            keylen,
+            _iter_handle,
+        }
+    }
+}
+
+impl Deref for TablePropertiesKey {
+    type Target = str;
+
+    fn deref(&self) -> &str {
+        // Safety: creating slice from values reported by rocksdb, that should
+        // be valid as long is this object is valid. The slice is guaranteed to
+        // be UTF-8 by the constructor.
+        unsafe {
+            let bytes = slice::from_raw_parts(self.key, self.keylen);
+            str::from_utf8_unchecked(bytes)
+        }
+    }
+}
+
+pub struct TableProperties {
+    handle: TablePropertiesHandle,
+}
+
+impl TableProperties {
+    fn new(handle: TablePropertiesHandle) -> TableProperties {
+        TableProperties { handle }
+    }
+
+    fn get_u64(&self, prop: DBTableProperty) -> u64 {
+        unsafe { crocksdb_ffi::crocksdb_table_properties_get_u64(self.handle.ptr(), prop) }
+    }
+
+    pub fn num_entries(&self) -> u64 {
+        self.get_u64(DBTableProperty::NumEntries)
+    }
+
+    pub fn user_collected_properties(&self) -> UserCollectedProperties {
+        UserCollectedProperties::new(self.handle.clone())
+    }
+}
+
+pub struct UserCollectedProperties {
+    handle: UserCollectedPropertiesHandle,
+}
+
+impl UserCollectedProperties {
+    fn new(table_props_handle: TablePropertiesHandle) -> UserCollectedProperties {
+        UserCollectedProperties {
+            handle: UserCollectedPropertiesHandle::new(table_props_handle),
+        }
+    }
+
+    pub fn get<Q: AsRef<[u8]>>(&self, index: Q) -> Option<&[u8]> {
+        let bytes = index.as_ref();
+        let mut size = 0;
+        unsafe {
+            let ptr = crocksdb_ffi::crocksdb_user_collected_properties_get(
+                self.handle.ptr(),
+                bytes.as_ptr(),
+                bytes.len(),
+                &mut size,
+            );
+            if ptr.is_null() {
+                return None;
+            }
+            Some(slice::from_raw_parts(ptr, size))
+        }
+    }
+
+    pub fn len(&self) -> usize {
+        unsafe { crocksdb_ffi::crocksdb_user_collected_properties_len(self.handle.ptr()) }
+    }
+
+    pub fn is_empty(&self) -> bool {
+        self.len() == 0
+    }
+}

src/table_properties_rc_handles.rs

+//! Reference counted handles to various TableProperties* types, that safely
+//! preserve destruction order.
+
+use crocksdb_ffi::{
+    DBTableProperties, DBTablePropertiesCollection, DBTablePropertiesCollectionIterator,
+    DBUserCollectedProperties,
+};
+use librocksdb_sys as crocksdb_ffi;
+use std::rc::Rc;
+
+/// This is a shared wrapper around a DBTablePropertiesCollection w/ dtor
+#[derive(Clone)]
+pub struct TablePropertiesCollectionHandle {
+    shared: Rc<TablePropertiesCollectionHandleWithDrop>,
+}
+
+impl TablePropertiesCollectionHandle {
+    pub unsafe fn new(ptr: *mut DBTablePropertiesCollection) -> TablePropertiesCollectionHandle {
+        assert!(!ptr.is_null());
+        TablePropertiesCollectionHandle {
+            shared: Rc::new(TablePropertiesCollectionHandleWithDrop { ptr }),
+        }
+    }
+
+    pub fn ptr(&self) -> *mut DBTablePropertiesCollection {
+        self.shared.ptr
+    }
+}
+
+struct TablePropertiesCollectionHandleWithDrop {
+    ptr: *mut DBTablePropertiesCollection,
+}
+
+impl Drop for TablePropertiesCollectionHandleWithDrop {
+    fn drop(&mut self) {
+        unsafe {
+            crocksdb_ffi::crocksdb_table_properties_collection_destroy(self.ptr);
+        }
+    }
+}
+
+/// This is a shared wrapper around a DBTablePropertiesCollection w/ dtor.
+//
+// # Safety
+//
+// The safety of this struct depends on drop order, with the iterator
+// needing to drop before the collection.
+#[derive(Clone)]
+pub struct TablePropertiesCollectionIteratorHandle {
+    shared: Rc<TablePropertiesCollectionIteratorHandleWithDrop>,
+    collection: TablePropertiesCollectionHandle,
+}
+
+impl TablePropertiesCollectionIteratorHandle {
+    pub fn new(
+        collection: TablePropertiesCollectionHandle,
+    ) -> TablePropertiesCollectionIteratorHandle {
+        unsafe {
+            let ptr =
+                crocksdb_ffi::crocksdb_table_properties_collection_iter_create(collection.ptr());
+            TablePropertiesCollectionIteratorHandle {
+                shared: Rc::new(TablePropertiesCollectionIteratorHandleWithDrop { ptr }),
+                collection,
+            }
+        }
+    }
+
+    pub fn ptr(&self) -> *mut DBTablePropertiesCollectionIterator {
+        self.shared.ptr
+    }
+}
+
+struct TablePropertiesCollectionIteratorHandleWithDrop {
+    ptr: *mut DBTablePropertiesCollectionIterator,
+}
+
+impl Drop for TablePropertiesCollectionIteratorHandleWithDrop {
+    fn drop(&mut self) {
+        unsafe {
+            crocksdb_ffi::crocksdb_table_properties_collection_iter_destroy(self.ptr);
+        }
+    }
+}
+
+// # Safety
+//
+// `ptr` is valid as long as the iterator is
+#[derive(Clone)]
+pub struct TablePropertiesHandle {
+    ptr: *const DBTableProperties,
+    iter_handle: TablePropertiesCollectionIteratorHandle,
+}
+
+impl TablePropertiesHandle {
+    pub fn new(
+        ptr: *const DBTableProperties,
+        iter_handle: TablePropertiesCollectionIteratorHandle,
+    ) -> TablePropertiesHandle {
+        TablePropertiesHandle { ptr, iter_handle }
+    }
+
+    pub fn ptr(&self) -> *const DBTableProperties {
+        self.ptr
+    }
+}
+
+// # Safety
+//
+// `ptr` is valid as long as the table properties are
+#[derive(Clone)]
+pub struct UserCollectedPropertiesHandle {
+    ptr: *const DBUserCollectedProperties,
+    table_props_handle: TablePropertiesHandle,
+}
+
+impl UserCollectedPropertiesHandle {
+    pub fn new(table_props_handle: TablePropertiesHandle) -> UserCollectedPropertiesHandle {
+        unsafe {
+            let ptr = crocksdb_ffi::crocksdb_table_properties_get_user_properties(
+                table_props_handle.ptr(),
+            );
+            UserCollectedPropertiesHandle {
+                ptr,
+                table_props_handle,
+            }
+        }
+    }
+
+    pub fn ptr(&self) -> *const DBUserCollectedProperties {
+        self.ptr
+    }
+}

tests/cases/mod.rs

@@ -17,6 +17,7 @@ mod test_rocksdb_options;
 mod test_slice_transform;
 mod test_statistics;
 mod test_table_properties;
+mod test_table_properties_rc;
 mod test_titan;
 mod test_ttl;
 

tests/cases/test_table_properties_rc.rs

+use rocksdb::{ColumnFamilyOptions, DBOptions, Writable, DB};
+
+use super::tempdir_with_prefix;
+
+// This is testing that the reference counting prevents use after frees,
+// and can be verified by running under valgrind.
+#[test]
+fn test_lifetimes() {
+    let mut opts = DBOptions::new();
+    opts.create_if_missing(true);
+    let cf_opts = ColumnFamilyOptions::new();
+    let path = tempdir_with_prefix("table_properties_rc");
+    let path = path.path().to_str().unwrap();
+    let db = DB::open_cf(opts, path, vec![("default", cf_opts)]).unwrap();
+
+    let samples = vec![
+        (b"key1".to_vec(), b"value1".to_vec()),
+        (b"key2".to_vec(), b"value2".to_vec()),
+        (b"key3".to_vec(), b"value3".to_vec()),
+        (b"key4".to_vec(), b"value4".to_vec()),
+    ];
+
+    // Put 4 keys.
+    for &(ref k, ref v) in &samples {
+        db.put(k, v).unwrap();
+        assert_eq!(v.as_slice(), &*db.get(k).unwrap().unwrap());
+    }
+    db.flush(true).unwrap();
+
+    let collection = db.get_properties_of_all_tables_rc().unwrap();
+
+    assert_eq!(collection.len(), 1);
+
+    let mut iter = collection.iter();
+
+    drop(collection);
+
+    let (_key, prop) = iter.next().unwrap();
+
+    drop(iter);
+
+    let _ucp = prop.user_collected_properties();
+
+    drop(prop);
+}