Skip to content

Commit e99a12b

Browse files
authored
Allow passing &CStr arguments (tikv#641)
1 parent ef87246 commit e99a12b

File tree

6 files changed

+191
-26
lines changed

6 files changed

+191
-26
lines changed

src/db.rs

Lines changed: 13 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@ use crate::{
1919
column_family::UnboundColumnFamily,
2020
db_options::OptionsMustOutliveDB,
2121
ffi,
22-
ffi_util::{from_cstr, opt_bytes_to_ptr, raw_data, to_cpath},
22+
ffi_util::{from_cstr, opt_bytes_to_ptr, raw_data, to_cpath, CStrLike},
2323
ColumnFamily, ColumnFamilyDescriptor, CompactOptions, DBIteratorWithThreadMode,
2424
DBPinnableSlice, DBRawIteratorWithThreadMode, DBWALIterator, Direction, Error, FlushOptions,
2525
IngestExternalFileOptions, IteratorMode, Options, ReadOptions, SnapshotWithThreadMode,
@@ -1080,16 +1080,14 @@ impl<T: ThreadMode> DBWithThreadMode<T> {
10801080

10811081
fn create_inner_cf_handle(
10821082
&self,
1083-
name: &str,
1083+
name: impl CStrLike,
10841084
opts: &Options,
10851085
) -> Result<*mut ffi::rocksdb_column_family_handle_t, Error> {
1086-
let cf_name = if let Ok(c) = CString::new(name.as_bytes()) {
1087-
c
1088-
} else {
1089-
return Err(Error::new(
1090-
"Failed to convert path to CString when creating cf".to_owned(),
1091-
));
1092-
};
1086+
let cf_name = name.bake().map_err(|err| {
1087+
Error::new(format!(
1088+
"Failed to convert path to CString when creating cf: {err}"
1089+
))
1090+
})?;
10931091
Ok(unsafe {
10941092
ffi_try!(ffi::rocksdb_create_column_family(
10951093
self.inner,
@@ -1567,11 +1565,11 @@ impl<T: ThreadMode> DBWithThreadMode<T> {
15671565
/// the end. That string is parsed using `parse` callback which produces
15681566
/// the returned result.
15691567
fn property_value_impl<R>(
1570-
name: &str,
1568+
name: impl CStrLike,
15711569
get_property: impl FnOnce(*const c_char) -> *mut c_char,
15721570
parse: impl FnOnce(&str) -> Result<R, Error>,
15731571
) -> Result<Option<R>, Error> {
1574-
let value = match CString::new(name) {
1572+
let value = match name.bake() {
15751573
Ok(prop_name) => get_property(prop_name.as_ptr()),
15761574
Err(e) => {
15771575
return Err(Error::new(format!(
@@ -1600,7 +1598,7 @@ impl<T: ThreadMode> DBWithThreadMode<T> {
16001598
///
16011599
/// Full list of properties could be find
16021600
/// [here](https://github.com/facebook/rocksdb/blob/08809f5e6cd9cc4bc3958dd4d59457ae78c76660/include/rocksdb/db.h#L428-L634).
1603-
pub fn property_value(&self, name: &str) -> Result<Option<String>, Error> {
1601+
pub fn property_value(&self, name: impl CStrLike) -> Result<Option<String>, Error> {
16041602
Self::property_value_impl(
16051603
name,
16061604
|prop_name| unsafe { ffi::rocksdb_property_value(self.inner, prop_name) },
@@ -1615,7 +1613,7 @@ impl<T: ThreadMode> DBWithThreadMode<T> {
16151613
pub fn property_value_cf(
16161614
&self,
16171615
cf: &impl AsColumnFamilyRef,
1618-
name: &str,
1616+
name: impl CStrLike,
16191617
) -> Result<Option<String>, Error> {
16201618
Self::property_value_impl(
16211619
name,
@@ -1639,7 +1637,7 @@ impl<T: ThreadMode> DBWithThreadMode<T> {
16391637
///
16401638
/// Full list of properties that return int values could be find
16411639
/// [here](https://github.com/facebook/rocksdb/blob/08809f5e6cd9cc4bc3958dd4d59457ae78c76660/include/rocksdb/db.h#L654-L689).
1642-
pub fn property_int_value(&self, name: &str) -> Result<Option<u64>, Error> {
1640+
pub fn property_int_value(&self, name: impl CStrLike) -> Result<Option<u64>, Error> {
16431641
Self::property_value_impl(
16441642
name,
16451643
|prop_name| unsafe { ffi::rocksdb_property_value(self.inner, prop_name) },
@@ -1654,7 +1652,7 @@ impl<T: ThreadMode> DBWithThreadMode<T> {
16541652
pub fn property_int_value_cf(
16551653
&self,
16561654
cf: &impl AsColumnFamilyRef,
1657-
name: &str,
1655+
name: impl CStrLike,
16581656
) -> Result<Option<u64>, Error> {
16591657
Self::property_value_impl(
16601658
name,

src/db_options.rs

Lines changed: 9 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@ use crate::{
2424
comparator::{self, ComparatorCallback, CompareFn},
2525
db::DBAccess,
2626
ffi,
27+
ffi_util::CStrLike,
2728
merge_operator::{
2829
self, full_merge_callback, partial_merge_callback, MergeFn, MergeOperatorCallback,
2930
},
@@ -1274,11 +1275,11 @@ impl Options {
12741275

12751276
pub fn set_merge_operator_associative<F: MergeFn + Clone>(
12761277
&mut self,
1277-
name: &str,
1278+
name: impl CStrLike,
12781279
full_merge_fn: F,
12791280
) {
12801281
let cb = Box::new(MergeOperatorCallback {
1281-
name: CString::new(name.as_bytes()).unwrap(),
1282+
name: name.into_c_string().unwrap(),
12821283
full_merge_fn: full_merge_fn.clone(),
12831284
partial_merge_fn: full_merge_fn,
12841285
});
@@ -1298,12 +1299,12 @@ impl Options {
12981299

12991300
pub fn set_merge_operator<F: MergeFn, PF: MergeFn>(
13001301
&mut self,
1301-
name: &str,
1302+
name: impl CStrLike,
13021303
full_merge_fn: F,
13031304
partial_merge_fn: PF,
13041305
) {
13051306
let cb = Box::new(MergeOperatorCallback {
1306-
name: CString::new(name.as_bytes()).unwrap(),
1307+
name: name.into_c_string().unwrap(),
13071308
full_merge_fn,
13081309
partial_merge_fn,
13091310
});
@@ -1339,12 +1340,12 @@ impl Options {
13391340
///
13401341
/// If multi-threaded compaction is used, `filter_fn` may be called multiple times
13411342
/// simultaneously.
1342-
pub fn set_compaction_filter<F>(&mut self, name: &str, filter_fn: F)
1343+
pub fn set_compaction_filter<F>(&mut self, name: impl CStrLike, filter_fn: F)
13431344
where
13441345
F: CompactionFilterFn + Send + 'static,
13451346
{
13461347
let cb = Box::new(CompactionFilterCallback {
1347-
name: CString::new(name.as_bytes()).unwrap(),
1348+
name: name.into_c_string().unwrap(),
13481349
filter_fn,
13491350
});
13501351

@@ -1391,9 +1392,9 @@ impl Options {
13911392
/// The client must ensure that the comparator supplied here has the same
13921393
/// name and orders keys *exactly* the same as the comparator provided to
13931394
/// previous open calls on the same DB.
1394-
pub fn set_comparator(&mut self, name: &str, compare_fn: CompareFn) {
1395+
pub fn set_comparator(&mut self, name: impl CStrLike, compare_fn: CompareFn) {
13951396
let cb = Box::new(ComparatorCallback {
1396-
name: CString::new(name.as_bytes()).unwrap(),
1397+
name: name.into_c_string().unwrap(),
13971398
f: compare_fn,
13981399
});
13991400

src/ffi_util.rs

Lines changed: 152 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -80,3 +80,155 @@ macro_rules! ffi_try_impl {
8080
result
8181
}};
8282
}
83+
84+
/// Value which can be converted into a C string.
85+
///
86+
/// The trait is used as argument to functions which wish to accept either
87+
/// [`&str`] or [`&CStr`] arguments while internally need to interact with
88+
/// C APIs. Accepting [`&str`] may be more convenient for users but requires
89+
/// conversion into [`CString`] internally which requires allocation. With this
90+
/// trait, latency-conscious users may choose to prepare `CStr` in advance and
91+
/// then pass it directly without having to incur the conversion cost.
92+
///
93+
/// To use the trait, function should accept `impl CStrLike` and after baking
94+
/// the argument (with [`CStrLike::bake`] method) it can use it as a `&CStr`
95+
/// (since the baked result dereferences into `CStr`).
96+
///
97+
/// # Example
98+
///
99+
/// ```
100+
/// use std::ffi::{CStr, CString};
101+
/// use rocksdb::CStrLike;
102+
///
103+
/// fn strlen(arg: impl CStrLike) -> std::result::Result<usize, String> {
104+
/// let baked = arg.bake().map_err(|err| err.to_string())?;
105+
/// Ok(unsafe { libc::strlen(baked.as_ptr()) })
106+
/// }
107+
///
108+
/// const FOO: &str = "foo";
109+
/// const BAR: &CStr = unsafe { CStr::from_bytes_with_nul_unchecked(b"bar\0") };
110+
///
111+
/// assert_eq!(Ok(3), strlen(FOO));
112+
/// assert_eq!(Ok(3), strlen(BAR));
113+
/// ```
114+
pub trait CStrLike {
115+
type Baked: std::ops::Deref<Target = CStr>;
116+
type Error: std::fmt::Debug + std::fmt::Display;
117+
118+
/// Bakes self into value which can be freely converted into [`&CStr`].
119+
///
120+
/// This may require allocation and may fail if `self` has invalid value.
121+
fn bake(self) -> Result<Self::Baked, Self::Error>;
122+
123+
/// Consumers and converts value into an owned [`CString`].
124+
///
125+
/// If `Self` is already a `CString` simply returns it; if it’s a reference
126+
/// to a `CString` then the value is cloned. In other cases this may
127+
/// require allocation and may fail if `self` has invalid value.
128+
fn into_c_string(self) -> Result<CString, Self::Error>;
129+
}
130+
131+
impl CStrLike for &str {
132+
type Baked = CString;
133+
type Error = std::ffi::NulError;
134+
135+
fn bake(self) -> Result<Self::Baked, Self::Error> {
136+
CString::new(self)
137+
}
138+
fn into_c_string(self) -> Result<CString, Self::Error> {
139+
CString::new(self)
140+
}
141+
}
142+
143+
// This is redundant for the most part and exists so that `foo(&string)` (where
144+
// `string: String` works just as if `foo` took `arg: &str` argument.
145+
impl CStrLike for &String {
146+
type Baked = CString;
147+
type Error = std::ffi::NulError;
148+
149+
fn bake(self) -> Result<Self::Baked, Self::Error> {
150+
CString::new(self.as_bytes())
151+
}
152+
fn into_c_string(self) -> Result<CString, Self::Error> {
153+
CString::new(self.as_bytes())
154+
}
155+
}
156+
157+
impl CStrLike for &CStr {
158+
type Baked = Self;
159+
type Error = std::convert::Infallible;
160+
161+
fn bake(self) -> Result<Self::Baked, Self::Error> {
162+
Ok(self)
163+
}
164+
fn into_c_string(self) -> Result<CString, Self::Error> {
165+
Ok(self.to_owned())
166+
}
167+
}
168+
169+
// This exists so that if caller constructs a `CString` they can pass it into
170+
// the function accepting `CStrLike` argument. Some of such functions may take
171+
// the argument whereas otherwise they would need to allocated a new owned
172+
// object.
173+
impl CStrLike for CString {
174+
type Baked = CString;
175+
type Error = std::convert::Infallible;
176+
177+
fn bake(self) -> Result<Self::Baked, Self::Error> {
178+
Ok(self)
179+
}
180+
fn into_c_string(self) -> Result<CString, Self::Error> {
181+
Ok(self)
182+
}
183+
}
184+
185+
// This is redundant for the most part and exists so that `foo(&cstring)` (where
186+
// `string: CString` works just as if `foo` took `arg: &CStr` argument.
187+
impl<'a> CStrLike for &'a CString {
188+
type Baked = &'a CStr;
189+
type Error = std::convert::Infallible;
190+
191+
fn bake(self) -> Result<Self::Baked, Self::Error> {
192+
Ok(self)
193+
}
194+
fn into_c_string(self) -> Result<CString, Self::Error> {
195+
Ok(self.clone())
196+
}
197+
}
198+
199+
#[test]
200+
fn test_c_str_like_bake() {
201+
fn test<S: CStrLike>(value: S) -> Result<usize, S::Error> {
202+
value
203+
.bake()
204+
.map(|value| unsafe { libc::strlen(value.as_ptr()) })
205+
}
206+
207+
assert_eq!(Ok(3), test("foo")); // &str
208+
assert_eq!(Ok(3), test(&String::from("foo"))); // String
209+
assert_eq!(Ok(3), test(CString::new("foo").unwrap().as_ref())); // &CStr
210+
assert_eq!(Ok(3), test(&CString::new("foo").unwrap())); // &CString
211+
assert_eq!(Ok(3), test(CString::new("foo").unwrap())); // CString
212+
213+
assert_eq!(3, test("foo\0bar").err().unwrap().nul_position());
214+
}
215+
216+
#[test]
217+
fn test_c_str_like_into() {
218+
fn test<S: CStrLike>(value: S) -> Result<CString, S::Error> {
219+
value.into_c_string()
220+
}
221+
222+
let want = CString::new("foo").unwrap();
223+
224+
assert_eq!(Ok(want.clone()), test("foo")); // &str
225+
assert_eq!(Ok(want.clone()), test(&String::from("foo"))); // &String
226+
assert_eq!(
227+
Ok(want.clone()),
228+
test(CString::new("foo").unwrap().as_ref())
229+
); // &CStr
230+
assert_eq!(Ok(want.clone()), test(&CString::new("foo").unwrap())); // &CString
231+
assert_eq!(Ok(want), test(CString::new("foo").unwrap())); // CString
232+
233+
assert_eq!(3, test("foo\0bar").err().unwrap().nul_position());
234+
}

src/lib.rs

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -112,6 +112,7 @@ pub use crate::{
112112
UniversalCompactOptions, UniversalCompactionStopStyle, WriteOptions,
113113
},
114114
db_pinnable_slice::DBPinnableSlice,
115+
ffi_util::CStrLike,
115116
merge_operator::MergeOperands,
116117
perf::{PerfContext, PerfMetric, PerfStatsLevel},
117118
slice_transform::SliceTransform,

src/slice_transform.rs

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@ use std::slice;
1717

1818
use libc::{c_char, c_void, size_t};
1919

20-
use crate::ffi;
20+
use crate::{ffi, ffi_util::CStrLike};
2121

2222
/// A `SliceTransform` is a generic pluggable way of transforming one string
2323
/// to another. Its primary use-case is in configuring rocksdb
@@ -35,12 +35,12 @@ pub struct SliceTransform {
3535

3636
impl SliceTransform {
3737
pub fn create(
38-
name: &str,
38+
name: impl CStrLike,
3939
transform_fn: TransformFn,
4040
in_domain_fn: Option<InDomainFn>,
4141
) -> SliceTransform {
4242
let cb = Box::into_raw(Box::new(TransformCallback {
43-
name: CString::new(name.as_bytes()).unwrap(),
43+
name: name.into_c_string().unwrap(),
4444
transform_fn,
4545
in_domain_fn,
4646
}));

tests/test_property.rs

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,20 @@ fn property_test() {
2525
{
2626
let db = DB::open_default(&n).unwrap();
2727
let value = db.property_value(properties::STATS).unwrap().unwrap();
28+
assert!(value.contains("Stats"));
29+
}
2830

31+
{
32+
let db = DB::open_default(&n).unwrap();
33+
let prop_name = std::ffi::CString::new(properties::STATS).unwrap();
34+
let value = db.property_value(prop_name).unwrap().unwrap();
35+
assert!(value.contains("Stats"));
36+
}
37+
38+
{
39+
let db = DB::open_default(&n).unwrap();
40+
let prop_name = std::ffi::CString::new(properties::STATS).unwrap();
41+
let value = db.property_value(prop_name.as_ref()).unwrap().unwrap();
2942
assert!(value.contains("Stats"));
3043
}
3144
}

0 commit comments

Comments
 (0)