Merge pull request #150 from rrichardson/open-cfs-with-options

Open cfs with options - fixes merge conflict in #136

librocksdb-sys/Cargo.toml

@@ -22,6 +22,6 @@ libc = "0.2"
 const-cstr = "0.2"
 
 [build-dependencies]
-cc = { version = "1.0", features = ["parallel"] }
+cc = { version = "^1.0", features = ["parallel"] }
 make-cmd = "0.1"
 bindgen = "0.29"

src/db.rs

@@ -14,7 +14,7 @@
 //
 
 
-use {DB, Error, Options, WriteOptions, ColumnFamily};
+use {DB, Error, Options, WriteOptions, ColumnFamily, ColumnFamilyDescriptor};
 use ffi;
 use ffi_util::opt_bytes_to_ptr;
 
@@ -572,6 +572,16 @@ impl<'a> Drop for Snapshot<'a> {
     }
 }
 
+impl ColumnFamilyDescriptor {
+    // Create a new column family descriptor with the specified name and options.
+    pub fn new<S>(name: S, options: Options) -> Self where S: Into<String> {
+        ColumnFamilyDescriptor {
+            name: name.into(),
+            options
+        }
+    }
+}
+
 impl DB {
     /// Open a database with default options.
     pub fn open_default<P: AsRef<Path>>(path: P) -> Result<DB, Error> {
@@ -585,14 +595,17 @@ impl DB {
         DB::open_cf(opts, path, &[])
     }
 
-    /// Open a database with specified options and column family.
-    ///
-    /// A column family must be created first by calling `DB::create_cf`.
-    ///
-    /// # Panics
+    /// Open a database with the given database options and column family names.
     ///
-    /// * Panics if the column family doesn't exist.
+    /// Column families opened using this function will be created with default `Options`.
     pub fn open_cf<P: AsRef<Path>>(opts: &Options, path: P, cfs: &[&str]) -> Result<DB, Error> {
+        let cfs_v = cfs.to_vec().iter().map(|name| ColumnFamilyDescriptor::new(*name, Options::default())).collect();
+
+        DB::open_cf_descriptors(opts, path, cfs_v)
+    }
+
+    /// Open a database with the given database options and column family names/options.
+    pub fn open_cf_descriptors<P: AsRef<Path>>(opts: &Options, path: P, cfs: Vec<ColumnFamilyDescriptor>) -> Result<DB, Error> {
         let path = path.as_ref();
         let cpath = match CString::new(path.to_string_lossy().as_bytes()) {
             Ok(c) => c,
@@ -621,17 +634,19 @@ impl DB {
                 db = ffi_try!(ffi::rocksdb_open(opts.inner, cpath.as_ptr() as *const _,));
             }
         } else {
-            let mut cfs_v = cfs.to_vec();
+            let mut cfs_v = cfs;
             // Always open the default column family.
-            if !cfs_v.contains(&"default") {
-                cfs_v.push("default");
+            if !cfs_v.iter().any(|cf| cf.name == "default") {
+                cfs_v.push(ColumnFamilyDescriptor {
+                    name: String::from("default"),
+                    options: Options::default()
+                });
             }
-
             // We need to store our CStrings in an intermediate vector
             // so that their pointers remain valid.
             let c_cfs: Vec<CString> = cfs_v
                 .iter()
-                .map(|cf| CString::new(cf.as_bytes()).unwrap())
+                .map(|cf| CString::new(cf.name.as_bytes()).unwrap())
                 .collect();
 
             let mut cfnames: Vec<_> = c_cfs.iter().map(|cf| cf.as_ptr()).collect();
@@ -639,10 +654,8 @@ impl DB {
             // These handles will be populated by DB.
             let mut cfhandles: Vec<_> = cfs_v.iter().map(|_| ptr::null_mut()).collect();
 
-            // TODO(tyler) allow options to be passed in.
-            let mut cfopts: Vec<_> = cfs_v
-                .iter()
-                .map(|_| unsafe { ffi::rocksdb_options_create() as *const _ })
+            let mut cfopts: Vec<_> = cfs_v.iter()
+                .map(|cf| cf.options.inner as *const _)
                 .collect();
 
             unsafe {
@@ -666,7 +679,7 @@ impl DB {
             }
 
             for (n, h) in cfs_v.iter().zip(cfhandles) {
-                cf_map.insert(n.to_string(), ColumnFamily { inner: h });
+                cf_map.insert(n.name.clone(), ColumnFamily { inner: h });
             }
         }
 

src/db_options.rs

@@ -149,6 +149,25 @@ impl Options {
         }
     }
 
+    /// If true, any column families that didn't exist when opening the database
+    /// will be created.
+    ///
+    /// Default: `false`
+    ///
+    /// # Example
+    ///
+    /// ```
+    /// use rocksdb::Options;
+    ///
+    /// let mut opts = Options::default();
+    /// opts.create_missing_column_families(true);
+    /// ```
+    pub fn create_missing_column_families(&mut self, create_missing_cfs: bool) {
+        unsafe {
+            ffi::rocksdb_options_set_create_missing_column_families(self.inner, create_missing_cfs as c_uchar);
+        }
+    }
+
     /// Sets the compression algorithm that will be used for the bottommost level that
     /// contain files. If level-compaction is used, this option will only affect
     /// levels after base level.
@@ -202,10 +221,11 @@ impl Options {
         }
     }
 
-    pub fn set_merge_operator(&mut self, name: &str, merge_fn: MergeFn) {
+    pub fn set_merge_operator(&mut self, name: &str, full_merge_fn: MergeFn, partial_merge_fn: Option<MergeFn>) {
         let cb = Box::new(MergeOperatorCallback {
             name: CString::new(name.as_bytes()).unwrap(),
-            merge_fn: merge_fn,
+            full_merge_fn: full_merge_fn,
+            partial_merge_fn: partial_merge_fn.unwrap_or(full_merge_fn),
         });
 
         unsafe {
@@ -224,7 +244,7 @@ impl Options {
     #[deprecated(since = "0.5.0",
                  note = "add_merge_operator has been renamed to set_merge_operator")]
     pub fn add_merge_operator(&mut self, name: &str, merge_fn: MergeFn) {
-        self.set_merge_operator(name, merge_fn);
+        self.set_merge_operator(name, merge_fn, None);
     }
 
     /// Sets a compaction filter used to determine if entries should be kept, changed,

src/lib.rs

@@ -30,6 +30,21 @@
 //!  db.delete(b"my key").unwrap();
 //! ```
 //!
+//! Opening a database and a single column family with custom options:
+//!
+//! ```
+//! use rocksdb::{DB, ColumnFamilyDescriptor, Options};
+//! let mut cf_opts = Options::default();
+//! cf_opts.set_max_write_buffer_number(16);
+//! let cf = ColumnFamilyDescriptor::new("cf1", cf_opts);
+//!
+//! let mut db_opts = Options::default();
+//! db_opts.create_missing_column_families(true);
+//! db_opts.create_if_missing(true);
+//!
+//! let db = DB::open_cf_descriptors(&db_opts, "path/for/rocksdb/storage_with_cfs", vec![cf]).unwrap();
+//! ```
+//!
 
 extern crate libc;
 extern crate librocksdb_sys as ffi;
@@ -64,6 +79,14 @@ pub struct DB {
     path: PathBuf,
 }
 
+/// A descriptor for a RocksDB column family.
+///
+/// A description of the column family, containing the name and `Options`.
+pub struct ColumnFamilyDescriptor {
+    name: String,
+    options: Options,
+}
+
 /// A simple wrapper round a string, used for errors reported from
 /// ffi calls.
 #[derive(Debug, Clone, PartialEq)]
@@ -173,6 +196,7 @@ pub struct WriteOptions {
     inner: *mut ffi::rocksdb_writeoptions_t,
 }
 
+
 /// An opaque type used to represent a column family. Returned from some functions, and used
 /// in others
 #[derive(Copy, Clone)]

src/merge_operator.rs

@@ -21,7 +21,7 @@
 //! fn concat_merge(new_key: &[u8],
 //!                 existing_val: Option<&[u8]>,
 //!                 operands: &mut MergeOperands)
-//!                 -> Vec<u8> {
+//!                 -> Option<Vec<u8>> {
 //!
 //!    let mut result: Vec<u8> = Vec::with_capacity(operands.size_hint().0);
 //!    existing_val.map(|v| {
@@ -34,14 +34,14 @@
 //!            result.push(*e)
 //!        }
 //!    }
-//!    result
+//!    Some(result)
 //! }
 //!
 //! fn main() {
 //!    let path = "path/to/rocksdb";
 //!    let mut opts = Options::default();
 //!    opts.create_if_missing(true);
-//!    opts.add_merge_operator("test operator", concat_merge);
+//!    opts.set_merge_operator("test operator", concat_merge, None);
 //!    let db = DB::open(&opts, path).unwrap();
 //!    let p = db.put(b"k1", b"a");
 //!    db.merge(b"k1", b"b");
@@ -60,11 +60,12 @@ use std::mem;
 use std::ptr;
 use std::slice;
 
-pub type MergeFn = fn(&[u8], Option<&[u8]>, &mut MergeOperands) -> Vec<u8>;
+pub type MergeFn = fn(&[u8], Option<&[u8]>, &mut MergeOperands) -> Option<Vec<u8>>;
 
 pub struct MergeOperatorCallback {
 	pub name: CString,
-    pub merge_fn: MergeFn,
+	pub full_merge_fn: MergeFn,
+	pub partial_merge_fn: MergeFn,
 }
 
 pub unsafe extern "C" fn destructor_callback(raw_cb: *mut c_void) {
@@ -91,8 +92,13 @@ pub unsafe extern "C" fn full_merge_callback(
 	let cb = &mut *(raw_cb as *mut MergeOperatorCallback);
 	let operands = &mut MergeOperands::new(operands_list, operands_list_len, num_operands);
 	let key = slice::from_raw_parts(raw_key as *const u8, key_len as usize);
-    let oldval = slice::from_raw_parts(existing_value as *const u8, existing_value_len as usize);
-    let mut result = (cb.merge_fn)(key, Some(oldval), operands);
+	let oldval = 
+		if existing_value == ptr::null() {
+			None
+		} else {
+			Some(slice::from_raw_parts(existing_value as *const u8, existing_value_len as usize))
+		};
+	if let Some(mut result) = (cb.full_merge_fn)(key, oldval, operands) {
 		result.shrink_to_fit();
 		// TODO(tan) investigate zero-copy techniques to improve performance
 		let buf = libc::malloc(result.len() as size_t);
@@ -101,6 +107,10 @@ pub unsafe extern "C" fn full_merge_callback(
 		*success = 1 as u8;
 		ptr::copy(result.as_ptr() as *mut c_void, &mut *buf, result.len());
 		buf as *mut c_char
+	} else {
+		*success = 0 as u8;
+		ptr::null_mut() as *mut c_char
+	}
 }
 
 pub unsafe extern "C" fn partial_merge_callback(
@@ -116,7 +126,7 @@ pub unsafe extern "C" fn partial_merge_callback(
 	let cb = &mut *(raw_cb as *mut MergeOperatorCallback);
 	let operands = &mut MergeOperands::new(operands_list, operands_list_len, num_operands);
 	let key = slice::from_raw_parts(raw_key as *const u8, key_len as usize);
-    let mut result = (cb.merge_fn)(key, None, operands);
+	if let Some(mut result) = (cb.partial_merge_fn)(key, None, operands) {
 		result.shrink_to_fit();
 		// TODO(tan) investigate zero-copy techniques to improve performance
 		let buf = libc::malloc(result.len() as size_t);
@@ -125,6 +135,10 @@ pub unsafe extern "C" fn partial_merge_callback(
 		*success = 1 as u8;
 		ptr::copy(result.as_ptr() as *mut c_void, &mut *buf, result.len());
 		buf as *mut c_char
+	} else {
+		*success = 0 as u8;
+		ptr::null_mut::<c_char>()
+	}
 }
 
 
@@ -182,12 +196,15 @@ impl<'a> Iterator for &'a mut MergeOperands {
 }
 
 #[cfg(test)]
-#[allow(unused_variables)]
+mod test {
+
+	use super::*;
+
 	fn test_provided_merge(
-    new_key: &[u8],
+		_new_key: &[u8],
 		existing_val: Option<&[u8]>,
 		operands: &mut MergeOperands,
-) -> Vec<u8> {
+		) -> Option<Vec<u8>> {
 		let nops = operands.size_hint().0;
 		let mut result: Vec<u8> = Vec::with_capacity(nops);
 		if let Some(v) = existing_val {
@@ -200,7 +217,7 @@ fn test_provided_merge(
 				result.push(*e);
 			}
 		}
-    result
+		Some(result)
 	}
 
 #[test]
@@ -210,7 +227,7 @@ fn mergetest() {
 		let path = "_rust_rocksdb_mergetest";
 		let mut opts = Options::default();
 		opts.create_if_missing(true);
-    opts.set_merge_operator("test operator", test_provided_merge);
+		opts.set_merge_operator("test operator", test_provided_merge, None);
 		{
 			let db = DB::open(&opts, path).unwrap();
 			let p = db.put(b"k1", b"a");
@@ -240,3 +257,202 @@ fn mergetest() {
 		}
 		assert!(DB::destroy(&opts, path).is_ok());
 	}
+
+	unsafe fn to_slice<T: Sized>(p: &T) -> &[u8] {
+		::std::slice::from_raw_parts(
+			(p as *const T) as *const u8,
+			::std::mem::size_of::<T>(),
+			)
+	}
+
+	fn from_slice<T: Sized>(s: &[u8]) -> Option<&T> {
+		if ::std::mem::size_of::<T>() != s.len() {
+			println!("slice {:?} is len {}, but T is size {}", s, s.len(), ::std::mem::size_of::<T>());
+			None
+		} else {
+			unsafe {
+				Some(::std::mem::transmute(s.as_ptr()))
+			}
+		}
+	}
+
+#[repr(packed)]
+
+	#[derive(Clone, Debug)]
+	struct ValueCounts {
+		num_a: u32,
+		num_b: u32,
+		num_c: u32,
+		num_d: u32,
+	}
+
+	fn test_counting_partial_merge(
+		_new_key: &[u8],
+		existing_val: Option<&[u8]>,
+		operands: &mut MergeOperands,
+		) -> Option<Vec<u8>> {
+		let nops = operands.size_hint().0;
+		let mut result: Vec<u8> = Vec::with_capacity(nops);
+		for op in operands {
+			for e in op {
+				result.push(*e);
+			}
+		}
+		Some(result)
+	}
+
+	fn test_counting_full_merge(
+		_new_key: &[u8],
+		existing_val: Option<&[u8]>,
+		operands: &mut MergeOperands,
+		) -> Option<Vec<u8>> {
+
+		let nops = operands.size_hint().0;
+		let mut counts : ValueCounts = 
+			if let Some(v) = existing_val {
+				from_slice::<ValueCounts>(v).unwrap().clone()
+			} else {
+				ValueCounts {
+					num_a: 0,
+					num_b: 0,
+					num_c: 0,
+					num_d: 0 }
+			};
+
+		for op in operands {
+			for e in op {
+				match *e {
+					b'a' => counts.num_a += 1,
+					b'b' => counts.num_b += 1,
+					b'c' => counts.num_c += 1,
+					b'd' => counts.num_d += 1,
+					_ => {}
+				}
+			}
+		}
+		let slc = unsafe { to_slice(&counts) };
+		Some(slc.to_vec())
+	}
+
+#[test]
+	fn counting_mergetest() {
+		use std::thread;
+		use std::time::Duration;
+		use std::sync::Arc;
+		use {DB, Options, DBCompactionStyle};
+
+		let path = "_rust_rocksdb_partial_mergetest";
+		let mut opts = Options::default();
+		opts.create_if_missing(true);
+		opts.set_compaction_style(DBCompactionStyle::Universal);
+		opts.set_min_write_buffer_number_to_merge(10);
+
+		opts.set_merge_operator("sort operator", test_counting_full_merge, Some(test_counting_partial_merge));
+		{
+			let db = Arc::new(DB::open(&opts, path).unwrap());
+			let _ = db.delete(b"k1");
+			let _ = db.delete(b"k2");
+			let _ = db.merge(b"k1", b"a");
+			let _ = db.merge(b"k1", b"b");
+			let _ = db.merge(b"k1", b"d");
+			let _ = db.merge(b"k1", b"a");
+			let _ = db.merge(b"k1", b"a");
+			let _ = db.merge(b"k1", b"efg");
+			for i in 0..500 {
+				let _ = db.merge(b"k2", b"c");
+				if i % 20 == 0 {
+                    let _ = db.get(b"k2");
+				}
+			}
+			for i in 0..500 {
+				let _ = db.merge(b"k2", b"c");
+				if i % 20 == 0 {
+                    let _ = db.get(b"k2");
+				}
+			}
+			db.compact_range(None, None);
+			let d1 = db.clone();
+			let d2 = db.clone();
+			let d3 = db.clone();
+			let h1 = thread::spawn(move || {
+				for i in 0..500 {
+					let _ = d1.merge(b"k2", b"c");
+					if i % 20 == 0 {
+						let _ = d1.get(b"k2");
+					}
+				}
+				for i in 0..500 {
+					let _ = d1.merge(b"k2", b"a");
+					if i % 20 == 0 {
+						let _ = d1.get(b"k2");
+					}
+				}
+			});
+			let h2 = thread::spawn(move || {
+				for i in 0..500 {
+					let _ = d2.merge(b"k2", b"b");
+					if i % 20 == 0 {
+						let _ = d2.get(b"k2");
+					}
+				}
+				for i in 0..500 {
+					let _ = d2.merge(b"k2", b"d");
+					if i % 20 == 0 {
+						let _ = d2.get(b"k2");
+					}
+				}
+				d2.compact_range(None, None);
+			});
+			h2.join();
+			let h3 = thread::spawn(move || {
+				for i in 0..500 {
+					let _ = d3.merge(b"k2", b"a");
+					if i % 20 == 0 {
+						let _ = d3.get(b"k2");
+					}
+				}
+				for i in 0..500 {
+					let _ = d3.merge(b"k2", b"c");
+					if i % 20 == 0 {
+						let _ = d3.get(b"k2");
+					}
+				}
+			});
+			let m = db.merge(b"k1", b"b");
+			assert!(m.is_ok());
+			h3.join();
+			h1.join();
+			match db.get(b"k2") {
+				Ok(Some(value)) => {
+					match from_slice::<ValueCounts>(&*value) {
+						Some(v) => {
+							assert_eq!(v.num_a, 1000);
+							assert_eq!(v.num_b, 500);
+							assert_eq!(v.num_c, 2000);
+							assert_eq!(v.num_d, 500);
+						},
+						None => panic!("Failed to get ValueCounts from db"),
+					}
+				}
+				Err(e) => panic!("error reading value {:?}", e),
+				_ => panic!("value not present"),
+			}
+			match db.get(b"k1") {
+				Ok(Some(value)) => {
+					match from_slice::<ValueCounts>(&*value) {
+						Some(v) => {
+							assert_eq!(v.num_a, 3);
+							assert_eq!(v.num_b, 2);
+							assert_eq!(v.num_c, 0);
+							assert_eq!(v.num_d, 1);
+						},
+						None => panic!("Failed to get ValueCounts from db"),
+					}
+				}
+				Err(e) => panic!("error reading value {:?}", e),
+				_ => panic!("value not present"),
+			}
+		}
+		assert!(DB::destroy(&opts, path).is_ok());
+	}
+}

tests/test_column_family.rs

@@ -14,7 +14,7 @@
 //
 extern crate rocksdb;
 
-use rocksdb::{DB, MergeOperands, Options};
+use rocksdb::{DB, MergeOperands, Options, ColumnFamilyDescriptor};
 
 #[test]
 pub fn test_column_family() {
@@ -24,7 +24,7 @@ pub fn test_column_family() {
     {
         let mut opts = Options::default();
         opts.create_if_missing(true);
-        opts.set_merge_operator("test operator", test_provided_merge);
+        opts.set_merge_operator("test operator", test_provided_merge, None);
         let mut db = DB::open(&opts, path).unwrap();
         let opts = Options::default();
         match db.create_cf("cf1", &opts) {
@@ -38,7 +38,7 @@ pub fn test_column_family() {
     // should fail to open db without specifying same column families
     {
         let mut opts = Options::default();
-        opts.set_merge_operator("test operator", test_provided_merge);
+        opts.set_merge_operator("test operator", test_provided_merge, None);
         match DB::open(&opts, path) {
             Ok(_) => {
                 panic!("should not have opened DB successfully without \
@@ -56,7 +56,7 @@ pub fn test_column_family() {
     // should properly open db when specyfing all column families
     {
         let mut opts = Options::default();
-        opts.set_merge_operator("test operator", test_provided_merge);
+        opts.set_merge_operator("test operator", test_provided_merge, None);
         match DB::open_cf(&opts, path, &["cf1"]) {
             Ok(_) => println!("successfully opened db with column family"),
             Err(e) => panic!("failed to open db with column family: {}", e),
@@ -91,6 +91,25 @@ pub fn test_column_family() {
     assert!(DB::destroy(&Options::default(), path).is_ok());
 }
 
+#[test]
+fn test_create_missing_column_family() {
+    let path = "_rust_rocksdb_missing_cftest";
+
+    // should be able to create new column families when opening a new database
+    {
+        let mut opts = Options::default();
+        opts.create_if_missing(true);
+        opts.create_missing_column_families(true);
+
+        match DB::open_cf(&opts, path, &["cf1"]) {
+            Ok(_) => println!("successfully created new column family"),
+            Err(e) => panic!("failed to create new column family: {}", e),
+        }
+    }
+
+    assert!(DB::destroy(&Options::default(), path).is_ok());
+}
+
 #[test]
 #[ignore]
 fn test_merge_operator() {
@@ -98,7 +117,7 @@ fn test_merge_operator() {
     // TODO should be able to write, read, merge, batch, and iterate over a cf
     {
         let mut opts = Options::default();
-        opts.set_merge_operator("test operator", test_provided_merge);
+        opts.set_merge_operator("test operator", test_provided_merge, None);
         let db = match DB::open_cf(&opts, path, &["cf1"]) {
             Ok(db) => {
                 println!("successfully opened db with column family");
@@ -140,7 +159,7 @@ fn test_merge_operator() {
 fn test_provided_merge(_: &[u8],
                        existing_val: Option<&[u8]>,
                        operands: &mut MergeOperands)
-                       -> Vec<u8> {
+                       -> Option<Vec<u8>> {
     let nops = operands.size_hint().0;
     let mut result: Vec<u8> = Vec::with_capacity(nops);
     match existing_val {
@@ -156,5 +175,45 @@ fn test_provided_merge(_: &[u8],
             result.push(*e);
         }
     }
-    result
+    Some(result)
+}
+
+#[test]
+pub fn test_column_family_with_options() {
+    let path = "_rust_rocksdb_cf_with_optionstest";
+    {
+        let mut cfopts = Options::default();
+        cfopts.set_max_write_buffer_number(16);
+        let cf_descriptor = ColumnFamilyDescriptor::new("cf1", cfopts);
+
+        let mut opts = Options::default();
+        opts.create_if_missing(true);
+        opts.create_missing_column_families(true);
+
+        let cfs = vec![cf_descriptor];
+        match DB::open_cf_descriptors(&opts, path, cfs) {
+            Ok(_) => println!("created db with column family descriptors succesfully"),
+            Err(e) => {
+                panic!("could not create new database with column family descriptors: {}", e);
+            }
+        }
+    }
+
+    {
+        let mut cfopts = Options::default();
+        cfopts.set_max_write_buffer_number(16);
+        let cf_descriptor = ColumnFamilyDescriptor::new("cf1", cfopts);
+
+        let opts = Options::default();
+        let cfs = vec![cf_descriptor];
+
+        match DB::open_cf_descriptors(&opts, path, cfs) {
+            Ok(_) => println!("succesfully re-opened database with column family descriptorrs"),
+            Err(e) => {
+                panic!("unable to re-open database with column family descriptors: {}", e);
+            }
+        }
+    }
+
+    assert!(DB::destroy(&Options::default(), path).is_ok());
 }