From e99a12bfee119f2b938a01d2ec023dbe0767f0b7 Mon Sep 17 00:00:00 2001 From: Michal Nazarewicz Date: Wed, 15 Jun 2022 19:38:12 +0200 Subject: [PATCH] Allow passing &CStr arguments (#641) --- src/db.rs | 28 ++++---- src/db_options.rs | 17 ++--- src/ffi_util.rs | 152 +++++++++++++++++++++++++++++++++++++++++ src/lib.rs | 1 + src/slice_transform.rs | 6 +- tests/test_property.rs | 13 ++++ 6 files changed, 191 insertions(+), 26 deletions(-) diff --git a/src/db.rs b/src/db.rs index 86b61de..3e82a2d 100644 --- a/src/db.rs +++ b/src/db.rs @@ -19,7 +19,7 @@ use crate::{ column_family::UnboundColumnFamily, db_options::OptionsMustOutliveDB, ffi, - ffi_util::{from_cstr, opt_bytes_to_ptr, raw_data, to_cpath}, + ffi_util::{from_cstr, opt_bytes_to_ptr, raw_data, to_cpath, CStrLike}, ColumnFamily, ColumnFamilyDescriptor, CompactOptions, DBIteratorWithThreadMode, DBPinnableSlice, DBRawIteratorWithThreadMode, DBWALIterator, Direction, Error, FlushOptions, IngestExternalFileOptions, IteratorMode, Options, ReadOptions, SnapshotWithThreadMode, @@ -1080,16 +1080,14 @@ impl DBWithThreadMode { fn create_inner_cf_handle( &self, - name: &str, + name: impl CStrLike, opts: &Options, ) -> Result<*mut ffi::rocksdb_column_family_handle_t, Error> { - let cf_name = if let Ok(c) = CString::new(name.as_bytes()) { - c - } else { - return Err(Error::new( - "Failed to convert path to CString when creating cf".to_owned(), - )); - }; + let cf_name = name.bake().map_err(|err| { + Error::new(format!( + "Failed to convert path to CString when creating cf: {err}" + )) + })?; Ok(unsafe { ffi_try!(ffi::rocksdb_create_column_family( self.inner, @@ -1567,11 +1565,11 @@ impl DBWithThreadMode { /// the end. That string is parsed using `parse` callback which produces /// the returned result. fn property_value_impl( - name: &str, + name: impl CStrLike, get_property: impl FnOnce(*const c_char) -> *mut c_char, parse: impl FnOnce(&str) -> Result, ) -> Result, Error> { - let value = match CString::new(name) { + let value = match name.bake() { Ok(prop_name) => get_property(prop_name.as_ptr()), Err(e) => { return Err(Error::new(format!( @@ -1600,7 +1598,7 @@ impl DBWithThreadMode { /// /// Full list of properties could be find /// [here](https://github.com/facebook/rocksdb/blob/08809f5e6cd9cc4bc3958dd4d59457ae78c76660/include/rocksdb/db.h#L428-L634). - pub fn property_value(&self, name: &str) -> Result, Error> { + pub fn property_value(&self, name: impl CStrLike) -> Result, Error> { Self::property_value_impl( name, |prop_name| unsafe { ffi::rocksdb_property_value(self.inner, prop_name) }, @@ -1615,7 +1613,7 @@ impl DBWithThreadMode { pub fn property_value_cf( &self, cf: &impl AsColumnFamilyRef, - name: &str, + name: impl CStrLike, ) -> Result, Error> { Self::property_value_impl( name, @@ -1639,7 +1637,7 @@ impl DBWithThreadMode { /// /// Full list of properties that return int values could be find /// [here](https://github.com/facebook/rocksdb/blob/08809f5e6cd9cc4bc3958dd4d59457ae78c76660/include/rocksdb/db.h#L654-L689). - pub fn property_int_value(&self, name: &str) -> Result, Error> { + pub fn property_int_value(&self, name: impl CStrLike) -> Result, Error> { Self::property_value_impl( name, |prop_name| unsafe { ffi::rocksdb_property_value(self.inner, prop_name) }, @@ -1654,7 +1652,7 @@ impl DBWithThreadMode { pub fn property_int_value_cf( &self, cf: &impl AsColumnFamilyRef, - name: &str, + name: impl CStrLike, ) -> Result, Error> { Self::property_value_impl( name, diff --git a/src/db_options.rs b/src/db_options.rs index 636aae8..38608e1 100644 --- a/src/db_options.rs +++ b/src/db_options.rs @@ -24,6 +24,7 @@ use crate::{ comparator::{self, ComparatorCallback, CompareFn}, db::DBAccess, ffi, + ffi_util::CStrLike, merge_operator::{ self, full_merge_callback, partial_merge_callback, MergeFn, MergeOperatorCallback, }, @@ -1274,11 +1275,11 @@ impl Options { pub fn set_merge_operator_associative( &mut self, - name: &str, + name: impl CStrLike, full_merge_fn: F, ) { let cb = Box::new(MergeOperatorCallback { - name: CString::new(name.as_bytes()).unwrap(), + name: name.into_c_string().unwrap(), full_merge_fn: full_merge_fn.clone(), partial_merge_fn: full_merge_fn, }); @@ -1298,12 +1299,12 @@ impl Options { pub fn set_merge_operator( &mut self, - name: &str, + name: impl CStrLike, full_merge_fn: F, partial_merge_fn: PF, ) { let cb = Box::new(MergeOperatorCallback { - name: CString::new(name.as_bytes()).unwrap(), + name: name.into_c_string().unwrap(), full_merge_fn, partial_merge_fn, }); @@ -1339,12 +1340,12 @@ impl Options { /// /// If multi-threaded compaction is used, `filter_fn` may be called multiple times /// simultaneously. - pub fn set_compaction_filter(&mut self, name: &str, filter_fn: F) + pub fn set_compaction_filter(&mut self, name: impl CStrLike, filter_fn: F) where F: CompactionFilterFn + Send + 'static, { let cb = Box::new(CompactionFilterCallback { - name: CString::new(name.as_bytes()).unwrap(), + name: name.into_c_string().unwrap(), filter_fn, }); @@ -1391,9 +1392,9 @@ impl Options { /// The client must ensure that the comparator supplied here has the same /// name and orders keys *exactly* the same as the comparator provided to /// previous open calls on the same DB. - pub fn set_comparator(&mut self, name: &str, compare_fn: CompareFn) { + pub fn set_comparator(&mut self, name: impl CStrLike, compare_fn: CompareFn) { let cb = Box::new(ComparatorCallback { - name: CString::new(name.as_bytes()).unwrap(), + name: name.into_c_string().unwrap(), f: compare_fn, }); diff --git a/src/ffi_util.rs b/src/ffi_util.rs index 6421f46..fc6b459 100644 --- a/src/ffi_util.rs +++ b/src/ffi_util.rs @@ -80,3 +80,155 @@ macro_rules! ffi_try_impl { result }}; } + +/// Value which can be converted into a C string. +/// +/// The trait is used as argument to functions which wish to accept either +/// [`&str`] or [`&CStr`] arguments while internally need to interact with +/// C APIs. Accepting [`&str`] may be more convenient for users but requires +/// conversion into [`CString`] internally which requires allocation. With this +/// trait, latency-conscious users may choose to prepare `CStr` in advance and +/// then pass it directly without having to incur the conversion cost. +/// +/// To use the trait, function should accept `impl CStrLike` and after baking +/// the argument (with [`CStrLike::bake`] method) it can use it as a `&CStr` +/// (since the baked result dereferences into `CStr`). +/// +/// # Example +/// +/// ``` +/// use std::ffi::{CStr, CString}; +/// use rocksdb::CStrLike; +/// +/// fn strlen(arg: impl CStrLike) -> std::result::Result { +/// let baked = arg.bake().map_err(|err| err.to_string())?; +/// Ok(unsafe { libc::strlen(baked.as_ptr()) }) +/// } +/// +/// const FOO: &str = "foo"; +/// const BAR: &CStr = unsafe { CStr::from_bytes_with_nul_unchecked(b"bar\0") }; +/// +/// assert_eq!(Ok(3), strlen(FOO)); +/// assert_eq!(Ok(3), strlen(BAR)); +/// ``` +pub trait CStrLike { + type Baked: std::ops::Deref; + type Error: std::fmt::Debug + std::fmt::Display; + + /// Bakes self into value which can be freely converted into [`&CStr`]. + /// + /// This may require allocation and may fail if `self` has invalid value. + fn bake(self) -> Result; + + /// Consumers and converts value into an owned [`CString`]. + /// + /// If `Self` is already a `CString` simply returns it; if it’s a reference + /// to a `CString` then the value is cloned. In other cases this may + /// require allocation and may fail if `self` has invalid value. + fn into_c_string(self) -> Result; +} + +impl CStrLike for &str { + type Baked = CString; + type Error = std::ffi::NulError; + + fn bake(self) -> Result { + CString::new(self) + } + fn into_c_string(self) -> Result { + CString::new(self) + } +} + +// This is redundant for the most part and exists so that `foo(&string)` (where +// `string: String` works just as if `foo` took `arg: &str` argument. +impl CStrLike for &String { + type Baked = CString; + type Error = std::ffi::NulError; + + fn bake(self) -> Result { + CString::new(self.as_bytes()) + } + fn into_c_string(self) -> Result { + CString::new(self.as_bytes()) + } +} + +impl CStrLike for &CStr { + type Baked = Self; + type Error = std::convert::Infallible; + + fn bake(self) -> Result { + Ok(self) + } + fn into_c_string(self) -> Result { + Ok(self.to_owned()) + } +} + +// This exists so that if caller constructs a `CString` they can pass it into +// the function accepting `CStrLike` argument. Some of such functions may take +// the argument whereas otherwise they would need to allocated a new owned +// object. +impl CStrLike for CString { + type Baked = CString; + type Error = std::convert::Infallible; + + fn bake(self) -> Result { + Ok(self) + } + fn into_c_string(self) -> Result { + Ok(self) + } +} + +// This is redundant for the most part and exists so that `foo(&cstring)` (where +// `string: CString` works just as if `foo` took `arg: &CStr` argument. +impl<'a> CStrLike for &'a CString { + type Baked = &'a CStr; + type Error = std::convert::Infallible; + + fn bake(self) -> Result { + Ok(self) + } + fn into_c_string(self) -> Result { + Ok(self.clone()) + } +} + +#[test] +fn test_c_str_like_bake() { + fn test(value: S) -> Result { + value + .bake() + .map(|value| unsafe { libc::strlen(value.as_ptr()) }) + } + + assert_eq!(Ok(3), test("foo")); // &str + assert_eq!(Ok(3), test(&String::from("foo"))); // String + assert_eq!(Ok(3), test(CString::new("foo").unwrap().as_ref())); // &CStr + assert_eq!(Ok(3), test(&CString::new("foo").unwrap())); // &CString + assert_eq!(Ok(3), test(CString::new("foo").unwrap())); // CString + + assert_eq!(3, test("foo\0bar").err().unwrap().nul_position()); +} + +#[test] +fn test_c_str_like_into() { + fn test(value: S) -> Result { + value.into_c_string() + } + + let want = CString::new("foo").unwrap(); + + assert_eq!(Ok(want.clone()), test("foo")); // &str + assert_eq!(Ok(want.clone()), test(&String::from("foo"))); // &String + assert_eq!( + Ok(want.clone()), + test(CString::new("foo").unwrap().as_ref()) + ); // &CStr + assert_eq!(Ok(want.clone()), test(&CString::new("foo").unwrap())); // &CString + assert_eq!(Ok(want), test(CString::new("foo").unwrap())); // CString + + assert_eq!(3, test("foo\0bar").err().unwrap().nul_position()); +} diff --git a/src/lib.rs b/src/lib.rs index 5ec7de1..3b239ac 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -112,6 +112,7 @@ pub use crate::{ UniversalCompactOptions, UniversalCompactionStopStyle, WriteOptions, }, db_pinnable_slice::DBPinnableSlice, + ffi_util::CStrLike, merge_operator::MergeOperands, perf::{PerfContext, PerfMetric, PerfStatsLevel}, slice_transform::SliceTransform, diff --git a/src/slice_transform.rs b/src/slice_transform.rs index 53e05bd..8d35261 100644 --- a/src/slice_transform.rs +++ b/src/slice_transform.rs @@ -17,7 +17,7 @@ use std::slice; use libc::{c_char, c_void, size_t}; -use crate::ffi; +use crate::{ffi, ffi_util::CStrLike}; /// A `SliceTransform` is a generic pluggable way of transforming one string /// to another. Its primary use-case is in configuring rocksdb @@ -35,12 +35,12 @@ pub struct SliceTransform { impl SliceTransform { pub fn create( - name: &str, + name: impl CStrLike, transform_fn: TransformFn, in_domain_fn: Option, ) -> SliceTransform { let cb = Box::into_raw(Box::new(TransformCallback { - name: CString::new(name.as_bytes()).unwrap(), + name: name.into_c_string().unwrap(), transform_fn, in_domain_fn, })); diff --git a/tests/test_property.rs b/tests/test_property.rs index 5682738..327ee66 100644 --- a/tests/test_property.rs +++ b/tests/test_property.rs @@ -25,7 +25,20 @@ fn property_test() { { let db = DB::open_default(&n).unwrap(); let value = db.property_value(properties::STATS).unwrap().unwrap(); + assert!(value.contains("Stats")); + } + { + let db = DB::open_default(&n).unwrap(); + let prop_name = std::ffi::CString::new(properties::STATS).unwrap(); + let value = db.property_value(prop_name).unwrap().unwrap(); + assert!(value.contains("Stats")); + } + + { + let db = DB::open_default(&n).unwrap(); + let prop_name = std::ffi::CString::new(properties::STATS).unwrap(); + let value = db.property_value(prop_name.as_ref()).unwrap().unwrap(); assert!(value.contains("Stats")); } }