//! Code inspired by [https://github.com/rust-rocksdb/rust-rocksdb][Rust RocksDB] under Apache License 2.0. #![allow(unsafe_code)] use crate::error::invalid_input_error; use libc::{self, c_char, c_void}; use oxrocksdb_sys::*; use std::borrow::Borrow; use std::env::temp_dir; use std::ffi::{CStr, CString}; use std::io::{Error, ErrorKind, Result}; use std::marker::PhantomData; use std::ops::Deref; use std::path::Path; use std::sync::Arc; use std::{ptr, slice}; macro_rules! ffi_result { ( $($function:ident)::*() ) => { ffi_result_impl!($($function)::*()) }; ( $($function:ident)::*( $arg1:expr $(, $arg:expr)* $(,)? ) ) => { ffi_result_impl!($($function)::*($arg1 $(, $arg)* ,)) }; } macro_rules! ffi_result_impl { ( $($function:ident)::*( $($arg:expr,)*) ) => {{ let mut err: *mut ::libc::c_char = ::std::ptr::null_mut(); let result = $($function)::*($($arg,)* &mut err); if err.is_null() { Ok(result) } else { Err(convert_error(err)) } }} } #[derive(Clone)] pub struct Db(Arc); unsafe impl Send for Db {} unsafe impl Sync for Db {} struct DbHandler { db: *mut rocksdb_transactiondb_t, options: *mut rocksdb_options_t, txn_options: *mut rocksdb_transactiondb_options_t, read_options: *mut rocksdb_readoptions_t, write_options: *mut rocksdb_writeoptions_t, flush_options: *mut rocksdb_flushoptions_t, env: Option<*mut rocksdb_env_t>, column_families: Vec<&'static str>, cf_handles: Vec<*mut rocksdb_column_family_handle_t>, } impl Drop for DbHandler { fn drop(&mut self) { unsafe { for cf_handle in &self.cf_handles { rocksdb_column_family_handle_destroy(*cf_handle); } rocksdb_transactiondb_close(self.db); rocksdb_readoptions_destroy(self.read_options); rocksdb_writeoptions_destroy(self.write_options); rocksdb_flushoptions_destroy(self.flush_options); rocksdb_transactiondb_options_destroy(self.txn_options); rocksdb_options_destroy(self.options); if let Some(env) = self.env { rocksdb_env_destroy(env); } } } } impl Db { pub fn new(column_families: &'static [&'static str]) -> Result { let path = if cfg!(target_os = "linux") { "/dev/shm/".into() } else { temp_dir() } .join("oxigraph-temp-rocksdb"); Ok(Self(Arc::new(Self::do_open(&path, column_families, true)?))) } pub fn open(path: &Path, column_families: &'static [&'static str]) -> Result { Ok(Self(Arc::new(Self::do_open(path, column_families, false)?))) } fn do_open( path: &Path, column_families: &'static [&'static str], in_memory: bool, ) -> Result { let c_path = CString::new( path.to_str() .ok_or_else(|| invalid_input_error("The DB path is not valid UTF-8"))?, ) .map_err(invalid_input_error)?; unsafe { let options = rocksdb_options_create(); assert!(!options.is_null(), "rocksdb_options_create returned null"); rocksdb_options_set_create_if_missing(options, 1); rocksdb_options_set_create_missing_column_families(options, 1); let txn_options = rocksdb_transactiondb_options_create(); assert!( !txn_options.is_null(), "rocksdb_transactiondb_options_create returned null" ); let env = if in_memory { let env = rocksdb_create_mem_env(); if env.is_null() { rocksdb_options_destroy(options); rocksdb_transactiondb_options_destroy(txn_options); return Err(other_error("Not able to create an in-memory environment.")); } rocksdb_options_set_env(options, env); Some(env) } else { None }; let mut column_families = column_families.to_vec(); if !column_families.contains(&"default") { column_families.push("default") } let c_column_families = column_families .iter() .map(|cf| CString::new(*cf)) .collect::, _>>() .map_err(invalid_input_error)?; let cf_options: Vec<*const rocksdb_options_t> = vec![options; column_families.len()]; let mut cf_handles: Vec<*mut rocksdb_column_family_handle_t> = vec![ptr::null_mut(); column_families.len()]; let db = ffi_result!(rocksdb_transactiondb_open_column_families( options, txn_options, c_path.as_ptr(), column_families.len().try_into().unwrap(), c_column_families .iter() .map(|cf| cf.as_ptr()) .collect::>() .as_ptr(), cf_options.as_ptr(), cf_handles.as_mut_ptr(), )) .map_err(|e| { rocksdb_options_destroy(options); rocksdb_transactiondb_options_destroy(txn_options); if let Some(env) = env { rocksdb_env_destroy(env); } e })?; assert!(!db.is_null(), "rocksdb_create returned null"); for handle in &cf_handles { if handle.is_null() { rocksdb_transactiondb_close(db); rocksdb_options_destroy(options); rocksdb_transactiondb_options_destroy(txn_options); if let Some(env) = env { rocksdb_env_destroy(env); } return Err(other_error( "Received null column family handle from RocksDB.", )); } } let read_options = rocksdb_readoptions_create(); assert!( !read_options.is_null(), "rocksdb_readoptions_create returned null" ); let write_options = rocksdb_writeoptions_create(); assert!( !read_options.is_null(), "rocksdb_writeoptions_create returned null" ); if in_memory { rocksdb_writeoptions_disable_WAL(write_options, 1); // No need for WAL } let flush_options = rocksdb_flushoptions_create(); assert!( !options.is_null(), "rocksdb_flushoptions_create returned null" ); Ok(DbHandler { db, options, txn_options, read_options, write_options, flush_options, env, column_families, cf_handles, }) } } pub fn column_family(&self, name: &'static str) -> Option { for (cf_name, cf_handle) in self.0.column_families.iter().zip(&self.0.cf_handles) { if *cf_name == name { return Some(ColumnFamily(*cf_handle)); } } None } pub fn flush(&self) -> Result<()> { unsafe { ffi_result!(rocksdb_transactiondb_flush(self.0.db, self.0.flush_options)) } } pub fn get( &self, column_family: &ColumnFamily, key: &[u8], ) -> Result>> { unsafe { let slice = ffi_result!(rocksdb_transactiondb_get_pinned_cf( self.0.db, self.0.read_options, column_family.0, key.as_ptr() as *const c_char, key.len() ))?; Ok(if slice.is_null() { None } else { Some(PinnableSlice { slice, lifetime: PhantomData::default(), }) }) } } pub fn contains_key(&self, column_family: &ColumnFamily, key: &[u8]) -> Result { Ok(self.get(column_family, key)?.is_some()) //TODO: optimize } pub fn insert(&self, column_family: &ColumnFamily, key: &[u8], value: &[u8]) -> Result<()> { unsafe { ffi_result!(rocksdb_transactiondb_put_cf( self.0.db, self.0.write_options, column_family.0, key.as_ptr() as *const c_char, key.len(), value.as_ptr() as *const c_char, value.len(), )) } } pub fn insert_empty(&self, column_family: &ColumnFamily, key: &[u8]) -> Result<()> { self.insert(column_family, key, &[]) } pub fn remove(&self, column_family: &ColumnFamily, key: &[u8]) -> Result<()> { unsafe { ffi_result!(rocksdb_transactiondb_delete_cf( self.0.db, self.0.write_options, column_family.0, key.as_ptr() as *const c_char, key.len() )) } } pub fn iter(&self, column_family: &ColumnFamily) -> Iter { self.scan_prefix(column_family, &[]) } pub fn scan_prefix(&self, column_family: &ColumnFamily, prefix: &[u8]) -> Iter { //We generate the upper bound let upper_bound = { let mut bound = prefix.to_vec(); let mut found = false; for c in bound.iter_mut().rev() { if *c < u8::MAX { *c += 1; found = true; break; } } if found { Some(bound) } else { None } }; unsafe { let options = rocksdb_readoptions_create(); assert!( !options.is_null(), "rocksdb_readoptions_create returned null" ); if let Some(upper_bound) = &upper_bound { rocksdb_readoptions_set_iterate_upper_bound( options, upper_bound.as_ptr() as *const c_char, upper_bound.len(), ); } let iter = rocksdb_transactiondb_create_iterator_cf(self.0.db, options, column_family.0); assert!(!iter.is_null(), "rocksdb_create_iterator returned null"); if prefix.is_empty() { rocksdb_iter_seek_to_first(iter); } else { rocksdb_iter_seek(iter, prefix.as_ptr() as *const c_char, prefix.len()); } let is_currently_valid = rocksdb_iter_valid(iter) != 0; Iter { iter, options, _upper_bound: upper_bound, _db: self.0.clone(), is_currently_valid, } } } pub fn len(&self, column_family: &ColumnFamily) -> Result { let mut count = 0; let mut iter = self.iter(column_family); while iter.is_valid() { count += 1; iter.next(); } iter.status()?; // We makes sure there is no read problem Ok(count) } pub fn is_empty(&self, column_family: &ColumnFamily) -> Result { let iter = self.iter(column_family); iter.status()?; // We makes sure there is no read problem Ok(!iter.is_valid()) } } // It is fine to not keep a lifetime: there is no way to use this type without the database being still in scope. // So, no use after free possible. #[derive(Clone)] pub struct ColumnFamily(*mut rocksdb_column_family_handle_t); unsafe impl Send for ColumnFamily {} unsafe impl Sync for ColumnFamily {} pub struct PinnableSlice<'a> { slice: *mut rocksdb_pinnableslice_t, lifetime: PhantomData<&'a ()>, } impl<'a> Drop for PinnableSlice<'a> { fn drop(&mut self) { unsafe { rocksdb_pinnableslice_destroy(self.slice); } } } impl<'a> Deref for PinnableSlice<'a> { type Target = [u8]; fn deref(&self) -> &[u8] { unsafe { let mut len = 0; let val = rocksdb_pinnableslice_value(self.slice, &mut len); slice::from_raw_parts(val as *const u8, len) } } } impl<'a> AsRef<[u8]> for PinnableSlice<'a> { fn as_ref(&self) -> &[u8] { &*self } } impl<'a> Borrow<[u8]> for PinnableSlice<'a> { fn borrow(&self) -> &[u8] { &*self } } pub struct Iter { iter: *mut rocksdb_iterator_t, is_currently_valid: bool, _upper_bound: Option>, _db: Arc, // needed to ensure that DB still lives while iter is used options: *mut rocksdb_readoptions_t, // needed to ensure that options still lives while iter is used } impl Drop for Iter { fn drop(&mut self) { unsafe { rocksdb_iter_destroy(self.iter); rocksdb_readoptions_destroy(self.options); } } } unsafe impl Send for Iter {} unsafe impl Sync for Iter {} impl Iter { pub fn is_valid(&self) -> bool { self.is_currently_valid } pub fn status(&self) -> Result<()> { unsafe { ffi_result!(rocksdb_iter_get_error(self.iter)) } } pub fn next(&mut self) { unsafe { rocksdb_iter_next(self.iter); self.is_currently_valid = rocksdb_iter_valid(self.iter) != 0; } } pub fn key(&self) -> Option<&[u8]> { if self.is_valid() { unsafe { let mut len = 0; let val = rocksdb_iter_key(self.iter, &mut len); Some(slice::from_raw_parts(val as *const u8, len)) } } else { None } } pub fn value(&self) -> Option<&[u8]> { if self.is_valid() { unsafe { let mut len = 0; let val = rocksdb_iter_value(self.iter, &mut len); Some(slice::from_raw_parts(val as *const u8, len)) } } else { None } } } fn convert_error(ptr: *const c_char) -> Error { let message = unsafe { let s = CStr::from_ptr(ptr).to_string_lossy().into_owned(); libc::free(ptr as *mut c_void); s }; other_error(message) } fn other_error(error: impl Into>) -> Error { Error::new(ErrorKind::InvalidInput, error) }