rust_rocksdb/transactions/
optimistic_transaction_db.rs

1// Copyright 2021 Yiyuan Liu
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7// http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14//
15
16use std::{collections::BTreeMap, ffi::CString, fs, iter, marker::PhantomData, path::Path, ptr};
17
18use libc::{c_char, c_int, size_t};
19
20use crate::column_family::ColumnFamilyTtl;
21use crate::{
22    AsColumnFamilyRef, ColumnFamilyDescriptor, DEFAULT_COLUMN_FAMILY_NAME, Error,
23    OptimisticTransactionOptions, Options, ThreadMode, Transaction, WriteOptions,
24    db::{DBCommon, DBInner},
25    ffi,
26    ffi_util::to_cpath,
27    write_batch::WriteBatchWithTransaction,
28};
29
30// Default options are kept per-thread to avoid re-allocating on every call while
31// also preventing cross-thread sharing. Some RocksDB option wrappers hold
32// pointers into internal buffers and are not safe to share across threads.
33// Using thread_local allows cheap reuse in the common "default options" path
34// without synchronization overhead. Callers who need non-defaults must pass
35// explicit options.
36thread_local! { static DEFAULT_WRITE_OPTS: WriteOptions = WriteOptions::default(); }
37// See `DEFAULT_TXN_OPTS` in `transaction_db.rs`: building the options per call
38// is a C++ `new`/`delete` pair on the transaction-begin path, and
39// `rocksdb_optimistictransaction_begin` only reads them.
40thread_local! {
41    static DEFAULT_OTXN_OPTS: OptimisticTransactionOptions =
42        OptimisticTransactionOptions::default();
43}
44
45/// A type alias to RocksDB Optimistic Transaction DB.
46///
47/// Please read the official
48/// [guide](https://github.com/facebook/rocksdb/wiki/Transactions#optimistictransactiondb)
49/// to learn more about RocksDB OptimisticTransactionDB.
50///
51/// The default thread mode for [`OptimisticTransactionDB`] is [`SingleThreaded`]
52/// if feature `multi-threaded-cf` is not enabled.
53///
54/// See [`DBCommon`] for full list of methods.
55///
56/// # Examples
57///
58/// ```
59/// use rust_rocksdb::{DB, Options, OptimisticTransactionDB, SingleThreaded};
60/// let tempdir = tempfile::Builder::new()
61///     .prefix("_path_for_optimistic_transaction_db")
62///     .tempdir()
63///     .expect("Failed to create temporary path for the _path_for_optimistic_transaction_db");
64/// let path = tempdir.path();
65/// {
66///     let db: OptimisticTransactionDB = OptimisticTransactionDB::open_default(path).unwrap();
67///     db.put(b"my key", b"my value").unwrap();
68///
69///     // create transaction
70///     let txn = db.transaction();
71///     txn.put(b"key2", b"value2");
72///     txn.put(b"key3", b"value3");
73///     txn.commit().unwrap();
74/// }
75/// let _ = DB::destroy(&Options::default(), path);
76/// ```
77///
78/// [`SingleThreaded`]: crate::SingleThreaded
79#[cfg(not(feature = "multi-threaded-cf"))]
80pub type OptimisticTransactionDB<T = crate::SingleThreaded> =
81    DBCommon<T, OptimisticTransactionDBInner>;
82#[cfg(feature = "multi-threaded-cf")]
83pub type OptimisticTransactionDB<T = crate::MultiThreaded> =
84    DBCommon<T, OptimisticTransactionDBInner>;
85
86pub struct OptimisticTransactionDBInner {
87    base: *mut ffi::rocksdb_t,
88    db: *mut ffi::rocksdb_optimistictransactiondb_t,
89}
90
91impl DBInner for OptimisticTransactionDBInner {
92    #[inline]
93    fn inner(&self) -> *mut ffi::rocksdb_t {
94        self.base
95    }
96}
97
98impl Drop for OptimisticTransactionDBInner {
99    fn drop(&mut self) {
100        unsafe {
101            ffi::rocksdb_optimistictransactiondb_close_base_db(self.base);
102            ffi::rocksdb_optimistictransactiondb_close(self.db);
103        }
104    }
105}
106
107/// Methods of `OptimisticTransactionDB`.
108impl<T: ThreadMode> OptimisticTransactionDB<T> {
109    /// Opens a database with default options.
110    pub fn open_default<P: AsRef<Path>>(path: P) -> Result<Self, Error> {
111        let mut opts = Options::default();
112        opts.create_if_missing(true);
113        Self::open(&opts, path)
114    }
115
116    /// Opens the database with the specified options.
117    pub fn open<P: AsRef<Path>>(opts: &Options, path: P) -> Result<Self, Error> {
118        Self::open_cf(opts, path, None::<&str>)
119    }
120
121    /// Opens a database with the given database options and column family names.
122    ///
123    /// Column families opened using this function will be created with default `Options`.
124    /// *NOTE*: `default` column family will be opened with the `Options::default()`.
125    /// If you want to open `default` column family with custom options, use `open_cf_descriptors` and
126    /// provide a `ColumnFamilyDescriptor` with the desired options.
127    pub fn open_cf<P, I, N>(opts: &Options, path: P, cfs: I) -> Result<Self, Error>
128    where
129        P: AsRef<Path>,
130        I: IntoIterator<Item = N>,
131        N: AsRef<str>,
132    {
133        let cfs = cfs
134            .into_iter()
135            .map(|name| ColumnFamilyDescriptor::new(name.as_ref(), Options::default()));
136
137        Self::open_cf_descriptors_internal(opts, path, cfs)
138    }
139
140    /// Opens a database with the given database options and column family descriptors.
141    pub fn open_cf_descriptors<P, I>(opts: &Options, path: P, cfs: I) -> Result<Self, Error>
142    where
143        P: AsRef<Path>,
144        I: IntoIterator<Item = ColumnFamilyDescriptor>,
145    {
146        Self::open_cf_descriptors_internal(opts, path, cfs)
147    }
148
149    /// Internal implementation for opening RocksDB.
150    fn open_cf_descriptors_internal<P, I>(opts: &Options, path: P, cfs: I) -> Result<Self, Error>
151    where
152        P: AsRef<Path>,
153        I: IntoIterator<Item = ColumnFamilyDescriptor>,
154    {
155        let cfs: Vec<_> = cfs.into_iter().collect();
156        let outlive = iter::once(opts.outlive.clone())
157            .chain(cfs.iter().map(|cf| cf.options.outlive.clone()))
158            .collect();
159
160        let cpath = to_cpath(&path)?;
161
162        if let Err(e) = fs::create_dir_all(&path) {
163            return Err(Error::new(format!(
164                "Failed to create RocksDB directory: `{e:?}`."
165            )));
166        }
167
168        let db: *mut ffi::rocksdb_optimistictransactiondb_t;
169        let mut cf_map = BTreeMap::new();
170
171        if cfs.is_empty() {
172            db = Self::open_raw(opts, &cpath)?;
173        } else {
174            let mut cfs_v = cfs;
175            // Always open the default column family.
176            if !cfs_v.iter().any(|cf| cf.name == DEFAULT_COLUMN_FAMILY_NAME) {
177                cfs_v.push(ColumnFamilyDescriptor {
178                    name: String::from(DEFAULT_COLUMN_FAMILY_NAME),
179                    options: Options::default(),
180                    ttl: ColumnFamilyTtl::SameAsDb,
181                });
182            }
183            // We need to store our CStrings in an intermediate vector
184            // so that their pointers remain valid.
185            let c_cfs: Vec<CString> = cfs_v
186                .iter()
187                .map(|cf| CString::new(cf.name.as_bytes()).unwrap())
188                .collect();
189
190            let cfnames: Vec<_> = c_cfs.iter().map(|cf| cf.as_ptr()).collect();
191
192            // These handles will be populated by DB.
193            let mut cfhandles: Vec<_> = cfs_v.iter().map(|_| ptr::null_mut()).collect();
194
195            let cfopts: Vec<_> = cfs_v
196                .iter()
197                .map(|cf| cf.options.inner.cast_const())
198                .collect();
199
200            db = Self::open_cf_raw(opts, &cpath, &cfs_v, &cfnames, &cfopts, &mut cfhandles)?;
201
202            for handle in &cfhandles {
203                if handle.is_null() {
204                    return Err(Error::new(
205                        "Received null column family handle from DB.".to_owned(),
206                    ));
207                }
208            }
209
210            for (cf_desc, inner) in cfs_v.iter().zip(cfhandles) {
211                cf_map.insert(cf_desc.name.clone(), inner);
212            }
213        }
214
215        if db.is_null() {
216            return Err(Error::new("Could not initialize database.".to_owned()));
217        }
218
219        let base = unsafe { ffi::rocksdb_optimistictransactiondb_get_base_db(db) };
220        if base.is_null() {
221            unsafe {
222                ffi::rocksdb_optimistictransactiondb_close(db);
223            }
224            return Err(Error::new("Could not initialize database.".to_owned()));
225        }
226        let inner = OptimisticTransactionDBInner { base, db };
227
228        Ok(Self::new(
229            inner,
230            T::new_cf_map_internal(cf_map),
231            path.as_ref().to_path_buf(),
232            outlive,
233        ))
234    }
235
236    fn open_raw(
237        opts: &Options,
238        cpath: &CString,
239    ) -> Result<*mut ffi::rocksdb_optimistictransactiondb_t, Error> {
240        unsafe {
241            let db = ffi_try!(ffi::rocksdb_optimistictransactiondb_open(
242                opts.inner,
243                cpath.as_ptr()
244            ));
245            Ok(db)
246        }
247    }
248
249    fn open_cf_raw(
250        opts: &Options,
251        cpath: &CString,
252        cfs_v: &[ColumnFamilyDescriptor],
253        cfnames: &[*const c_char],
254        cfopts: &[*const ffi::rocksdb_options_t],
255        cfhandles: &mut [*mut ffi::rocksdb_column_family_handle_t],
256    ) -> Result<*mut ffi::rocksdb_optimistictransactiondb_t, Error> {
257        unsafe {
258            let db = ffi_try!(ffi::rocksdb_optimistictransactiondb_open_column_families(
259                opts.inner,
260                cpath.as_ptr(),
261                cfs_v.len() as c_int,
262                cfnames.as_ptr(),
263                cfopts.as_ptr(),
264                cfhandles.as_mut_ptr(),
265            ));
266            Ok(db)
267        }
268    }
269
270    /// Creates a transaction with default options.
271    pub fn transaction(&'_ self) -> Transaction<'_, Self> {
272        DEFAULT_WRITE_OPTS.with(|write_opts| {
273            DEFAULT_OTXN_OPTS.with(|otxn_opts| self.transaction_opt(write_opts, otxn_opts))
274        })
275    }
276
277    /// Creates a transaction with default options.
278    pub fn transaction_opt(
279        &'_ self,
280        writeopts: &WriteOptions,
281        otxn_opts: &OptimisticTransactionOptions,
282    ) -> Transaction<'_, Self> {
283        Transaction {
284            inner: unsafe {
285                ffi::rocksdb_optimistictransaction_begin(
286                    self.inner.db,
287                    writeopts.inner,
288                    otxn_opts.inner,
289                    std::ptr::null_mut(),
290                )
291            },
292            _marker: PhantomData,
293        }
294    }
295
296    pub fn write_opt(
297        &self,
298        batch: &WriteBatchWithTransaction<true>,
299        writeopts: &WriteOptions,
300    ) -> Result<(), Error> {
301        unsafe {
302            ffi_try!(ffi::rocksdb_optimistictransactiondb_write(
303                self.inner.db,
304                writeopts.inner,
305                batch.inner
306            ));
307        }
308        Ok(())
309    }
310
311    pub fn write(&self, batch: &WriteBatchWithTransaction<true>) -> Result<(), Error> {
312        DEFAULT_WRITE_OPTS.with(|opts| self.write_opt(batch, opts))
313    }
314
315    pub fn write_without_wal(&self, batch: &WriteBatchWithTransaction<true>) -> Result<(), Error> {
316        let mut wo = WriteOptions::new();
317        wo.disable_wal(true);
318        self.write_opt(batch, &wo)
319    }
320
321    /// Removes the database entries in the range `["from", "to")` using given write options.
322    pub fn delete_range_cf_opt<K: AsRef<[u8]>>(
323        &self,
324        cf: &impl AsColumnFamilyRef,
325        from: K,
326        to: K,
327        writeopts: &WriteOptions,
328    ) -> Result<(), Error> {
329        let from = from.as_ref();
330        let to = to.as_ref();
331
332        unsafe {
333            ffi_try!(ffi::rocksdb_delete_range_cf(
334                self.inner.inner(),
335                writeopts.inner,
336                cf.inner(),
337                from.as_ptr() as *const c_char,
338                from.len() as size_t,
339                to.as_ptr() as *const c_char,
340                to.len() as size_t,
341            ));
342            Ok(())
343        }
344    }
345
346    /// Removes the database entries in the range `["from", "to")` using default write options.
347    pub fn delete_range_cf<K: AsRef<[u8]>>(
348        &self,
349        cf: &impl AsColumnFamilyRef,
350        from: K,
351        to: K,
352    ) -> Result<(), Error> {
353        DEFAULT_WRITE_OPTS.with(|opts| self.delete_range_cf_opt(cf, from, to, opts))
354    }
355}
rust_rocksdb/transactions/optimistic_transaction_db.rs

rust_rocksdb/transactions/
optimistic_transaction_db.rs
