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    OptimisticTransactionDBOptions, OptimisticTransactionOptions, Options, ThreadMode, Transaction,
24    WriteOptions,
25    db::{DBCommon, DBInner},
26    ffi,
27    ffi_util::to_cpath,
28    write_batch::WriteBatchWithTransaction,
29};
30
31// Default options are kept per-thread to avoid re-allocating on every call while
32// also preventing cross-thread sharing. Some RocksDB option wrappers hold
33// pointers into internal buffers and are not safe to share across threads.
34// Using thread_local allows cheap reuse in the common "default options" path
35// without synchronization overhead. Callers who need non-defaults must pass
36// explicit options.
37thread_local! { static DEFAULT_WRITE_OPTS: WriteOptions = WriteOptions::default(); }
38// See `DEFAULT_TXN_OPTS` in `transaction_db.rs`: building the options per call
39// is a C++ `new`/`delete` pair on the transaction-begin path, and
40// `rocksdb_optimistictransaction_begin` only reads them.
41thread_local! {
42    static DEFAULT_OTXN_OPTS: OptimisticTransactionOptions =
43        OptimisticTransactionOptions::default();
44}
45
46/// A type alias to RocksDB Optimistic Transaction DB.
47///
48/// Please read the official
49/// [guide](https://github.com/facebook/rocksdb/wiki/Transactions#optimistictransactiondb)
50/// to learn more about RocksDB OptimisticTransactionDB.
51///
52/// The default thread mode for [`OptimisticTransactionDB`] is [`SingleThreaded`]
53/// if feature `multi-threaded-cf` is not enabled.
54///
55/// See [`DBCommon`] for full list of methods.
56///
57/// # Examples
58///
59/// ```
60/// use rust_rocksdb::{DB, Options, OptimisticTransactionDB, SingleThreaded};
61/// let tempdir = tempfile::Builder::new()
62///     .prefix("_path_for_optimistic_transaction_db")
63///     .tempdir()
64///     .expect("Failed to create temporary path for the _path_for_optimistic_transaction_db");
65/// let path = tempdir.path();
66/// {
67///     let db: OptimisticTransactionDB = OptimisticTransactionDB::open_default(path).unwrap();
68///     db.put(b"my key", b"my value").unwrap();
69///
70///     // create transaction
71///     let txn = db.transaction();
72///     txn.put(b"key2", b"value2");
73///     txn.put(b"key3", b"value3");
74///     txn.commit().unwrap();
75/// }
76/// let _ = DB::destroy(&Options::default(), path);
77/// ```
78///
79/// [`SingleThreaded`]: crate::SingleThreaded
80#[cfg(not(feature = "multi-threaded-cf"))]
81pub type OptimisticTransactionDB<T = crate::SingleThreaded> =
82    DBCommon<T, OptimisticTransactionDBInner>;
83#[cfg(feature = "multi-threaded-cf")]
84pub type OptimisticTransactionDB<T = crate::MultiThreaded> =
85    DBCommon<T, OptimisticTransactionDBInner>;
86
87pub struct OptimisticTransactionDBInner {
88    base: *mut ffi::rocksdb_t,
89    /// The stacked `OptimisticTransactionDB`, as opposed to `base`, which is the plain DB
90    /// underneath it. Reachable from the crate so `checkpoint` can build a checkpoint on
91    /// the stacked DB rather than the base.
92    pub(crate) db: *mut ffi::rocksdb_optimistictransactiondb_t,
93}
94
95impl DBInner for OptimisticTransactionDBInner {
96    #[inline]
97    fn inner(&self) -> *mut ffi::rocksdb_t {
98        self.base
99    }
100}
101
102impl Drop for OptimisticTransactionDBInner {
103    fn drop(&mut self) {
104        unsafe {
105            ffi::rocksdb_optimistictransactiondb_close_base_db(self.base);
106            ffi::rocksdb_optimistictransactiondb_close(self.db);
107        }
108    }
109}
110
111/// Methods of `OptimisticTransactionDB`.
112impl<T: ThreadMode> OptimisticTransactionDB<T> {
113    /// Opens a database with default options.
114    pub fn open_default<P: AsRef<Path>>(path: P) -> Result<Self, Error> {
115        let mut opts = Options::default();
116        opts.create_if_missing(true);
117        Self::open(&opts, path)
118    }
119
120    /// Opens the database with the specified options.
121    pub fn open<P: AsRef<Path>>(opts: &Options, path: P) -> Result<Self, Error> {
122        Self::open_cf(opts, path, None::<&str>)
123    }
124
125    /// Opens a database with the given database options and column family names.
126    ///
127    /// Column families opened using this function will be created with default `Options`.
128    /// *NOTE*: `default` column family will be opened with the `Options::default()`.
129    /// If you want to open `default` column family with custom options, use `open_cf_descriptors` and
130    /// provide a `ColumnFamilyDescriptor` with the desired options.
131    pub fn open_cf<P, I, N>(opts: &Options, path: P, cfs: I) -> Result<Self, Error>
132    where
133        P: AsRef<Path>,
134        I: IntoIterator<Item = N>,
135        N: AsRef<str>,
136    {
137        let cfs = cfs
138            .into_iter()
139            .map(|name| ColumnFamilyDescriptor::new(name.as_ref(), Options::default()));
140
141        Self::open_cf_descriptors_internal(opts, None, path, cfs)
142    }
143
144    /// Opens a database with the given database options and column family descriptors.
145    pub fn open_cf_descriptors<P, I>(opts: &Options, path: P, cfs: I) -> Result<Self, Error>
146    where
147        P: AsRef<Path>,
148        I: IntoIterator<Item = ColumnFamilyDescriptor>,
149    {
150        Self::open_cf_descriptors_internal(opts, None, path, cfs)
151    }
152
153    /// Opens the database with the specified database and conflict validation options.
154    ///
155    /// Same as [`open`], plus control over how commits are validated. See
156    /// [`OptimisticTransactionDBOptions`].
157    ///
158    /// [`open`]: Self::open
159    pub fn open_with_otxn_db_options<P: AsRef<Path>>(
160        opts: &Options,
161        otxn_db_opts: &OptimisticTransactionDBOptions,
162        path: P,
163    ) -> Result<Self, Error> {
164        Self::open_cf_with_otxn_db_options(opts, otxn_db_opts, path, None::<&str>)
165    }
166
167    /// Opens a database with the given database and conflict validation options and column
168    /// family names.
169    ///
170    /// Same as [`open_cf`], plus control over how commits are validated. Column families
171    /// opened this way are created with default `Options`, including `default`.
172    ///
173    /// [`open_cf`]: Self::open_cf
174    pub fn open_cf_with_otxn_db_options<P, I, N>(
175        opts: &Options,
176        otxn_db_opts: &OptimisticTransactionDBOptions,
177        path: P,
178        cfs: I,
179    ) -> Result<Self, Error>
180    where
181        P: AsRef<Path>,
182        I: IntoIterator<Item = N>,
183        N: AsRef<str>,
184    {
185        let cfs = cfs
186            .into_iter()
187            .map(|name| ColumnFamilyDescriptor::new(name.as_ref(), Options::default()));
188
189        Self::open_cf_descriptors_internal(opts, Some(otxn_db_opts), path, cfs)
190    }
191
192    /// Opens a database with the given database and conflict validation options and column
193    /// family descriptors.
194    ///
195    /// Same as [`open_cf_descriptors`], plus control over how commits are validated.
196    ///
197    /// [`open_cf_descriptors`]: Self::open_cf_descriptors
198    pub fn open_cf_descriptors_with_otxn_db_options<P, I>(
199        opts: &Options,
200        otxn_db_opts: &OptimisticTransactionDBOptions,
201        path: P,
202        cfs: I,
203    ) -> Result<Self, Error>
204    where
205        P: AsRef<Path>,
206        I: IntoIterator<Item = ColumnFamilyDescriptor>,
207    {
208        Self::open_cf_descriptors_internal(opts, Some(otxn_db_opts), path, cfs)
209    }
210
211    /// Internal implementation for opening RocksDB.
212    ///
213    /// `otxn_db_opts` of `None` picks the C entry point that does not take them and lets
214    /// RocksDB supply its own defaults.
215    fn open_cf_descriptors_internal<P, I>(
216        opts: &Options,
217        otxn_db_opts: Option<&OptimisticTransactionDBOptions>,
218        path: P,
219        cfs: I,
220    ) -> Result<Self, Error>
221    where
222        P: AsRef<Path>,
223        I: IntoIterator<Item = ColumnFamilyDescriptor>,
224    {
225        let cfs: Vec<_> = cfs.into_iter().collect();
226        let outlive = iter::once(opts.outlive.clone())
227            .chain(cfs.iter().map(|cf| cf.options.outlive.clone()))
228            .collect();
229
230        let cpath = to_cpath(&path)?;
231
232        if let Err(e) = fs::create_dir_all(&path) {
233            return Err(Error::new(format!(
234                "Failed to create RocksDB directory: `{e:?}`."
235            )));
236        }
237
238        let db: *mut ffi::rocksdb_optimistictransactiondb_t;
239        let mut cf_map = BTreeMap::new();
240
241        if cfs.is_empty() {
242            db = Self::open_raw(opts, otxn_db_opts, &cpath)?;
243        } else {
244            let mut cfs_v = cfs;
245            // Always open the default column family.
246            if !cfs_v.iter().any(|cf| cf.name == DEFAULT_COLUMN_FAMILY_NAME) {
247                cfs_v.push(ColumnFamilyDescriptor {
248                    name: String::from(DEFAULT_COLUMN_FAMILY_NAME),
249                    options: Options::default(),
250                    ttl: ColumnFamilyTtl::SameAsDb,
251                });
252            }
253            // We need to store our CStrings in an intermediate vector
254            // so that their pointers remain valid.
255            let c_cfs: Vec<CString> = cfs_v
256                .iter()
257                .map(|cf| CString::new(cf.name.as_bytes()).unwrap())
258                .collect();
259
260            let cfnames: Vec<_> = c_cfs.iter().map(|cf| cf.as_ptr()).collect();
261
262            // These handles will be populated by DB.
263            let mut cfhandles: Vec<_> = cfs_v.iter().map(|_| ptr::null_mut()).collect();
264
265            let cfopts: Vec<_> = cfs_v
266                .iter()
267                .map(|cf| cf.options.inner.cast_const())
268                .collect();
269
270            db = Self::open_cf_raw(
271                opts,
272                otxn_db_opts,
273                &cpath,
274                &cfs_v,
275                &cfnames,
276                &cfopts,
277                &mut cfhandles,
278            )?;
279
280            for handle in &cfhandles {
281                if handle.is_null() {
282                    return Err(Error::new(
283                        "Received null column family handle from DB.".to_owned(),
284                    ));
285                }
286            }
287
288            for (cf_desc, inner) in cfs_v.iter().zip(cfhandles) {
289                cf_map.insert(cf_desc.name.clone(), inner);
290            }
291        }
292
293        if db.is_null() {
294            return Err(Error::new("Could not initialize database.".to_owned()));
295        }
296
297        let base = unsafe { ffi::rocksdb_optimistictransactiondb_get_base_db(db) };
298        if base.is_null() {
299            unsafe {
300                ffi::rocksdb_optimistictransactiondb_close(db);
301            }
302            return Err(Error::new("Could not initialize database.".to_owned()));
303        }
304        let inner = OptimisticTransactionDBInner { base, db };
305
306        Ok(Self::new(
307            inner,
308            T::new_cf_map_internal(cf_map),
309            path.as_ref().to_path_buf(),
310            outlive,
311        ))
312    }
313
314    fn open_raw(
315        opts: &Options,
316        otxn_db_opts: Option<&OptimisticTransactionDBOptions>,
317        cpath: &CString,
318    ) -> Result<*mut ffi::rocksdb_optimistictransactiondb_t, Error> {
319        unsafe {
320            let db = if let Some(otxn_db_opts) = otxn_db_opts {
321                ffi_try!(
322                    ffi::rocksdb_optimistictransactiondb_open_with_otxn_db_options(
323                        opts.inner,
324                        otxn_db_opts.inner,
325                        cpath.as_ptr()
326                    )
327                )
328            } else {
329                ffi_try!(ffi::rocksdb_optimistictransactiondb_open(
330                    opts.inner,
331                    cpath.as_ptr()
332                ))
333            };
334            Ok(db)
335        }
336    }
337
338    fn open_cf_raw(
339        opts: &Options,
340        otxn_db_opts: Option<&OptimisticTransactionDBOptions>,
341        cpath: &CString,
342        cfs_v: &[ColumnFamilyDescriptor],
343        cfnames: &[*const c_char],
344        cfopts: &[*const ffi::rocksdb_options_t],
345        cfhandles: &mut [*mut ffi::rocksdb_column_family_handle_t],
346    ) -> Result<*mut ffi::rocksdb_optimistictransactiondb_t, Error> {
347        unsafe {
348            let db = if let Some(otxn_db_opts) = otxn_db_opts {
349                ffi_try!(
350                    ffi::rocksdb_optimistictransactiondb_open_column_families_with_otxn_db_options(
351                        opts.inner,
352                        otxn_db_opts.inner,
353                        cpath.as_ptr(),
354                        cfs_v.len() as c_int,
355                        cfnames.as_ptr(),
356                        cfopts.as_ptr(),
357                        cfhandles.as_mut_ptr(),
358                    )
359                )
360            } else {
361                ffi_try!(ffi::rocksdb_optimistictransactiondb_open_column_families(
362                    opts.inner,
363                    cpath.as_ptr(),
364                    cfs_v.len() as c_int,
365                    cfnames.as_ptr(),
366                    cfopts.as_ptr(),
367                    cfhandles.as_mut_ptr(),
368                ))
369            };
370            Ok(db)
371        }
372    }
373
374    /// Creates a transaction with default options.
375    pub fn transaction(&'_ self) -> Transaction<'_, Self> {
376        DEFAULT_WRITE_OPTS.with(|write_opts| {
377            DEFAULT_OTXN_OPTS.with(|otxn_opts| self.transaction_opt(write_opts, otxn_opts))
378        })
379    }
380
381    /// Creates a transaction with default options.
382    pub fn transaction_opt(
383        &'_ self,
384        writeopts: &WriteOptions,
385        otxn_opts: &OptimisticTransactionOptions,
386    ) -> Transaction<'_, Self> {
387        Transaction {
388            inner: unsafe {
389                ffi::rocksdb_optimistictransaction_begin(
390                    self.inner.db,
391                    writeopts.inner,
392                    otxn_opts.inner,
393                    std::ptr::null_mut(),
394                )
395            },
396            _marker: PhantomData,
397        }
398    }
399
400    pub fn write_opt(
401        &self,
402        batch: &WriteBatchWithTransaction<true>,
403        writeopts: &WriteOptions,
404    ) -> Result<(), Error> {
405        unsafe {
406            ffi_try!(ffi::rocksdb_optimistictransactiondb_write(
407                self.inner.db,
408                writeopts.inner,
409                batch.inner
410            ));
411        }
412        Ok(())
413    }
414
415    pub fn write(&self, batch: &WriteBatchWithTransaction<true>) -> Result<(), Error> {
416        DEFAULT_WRITE_OPTS.with(|opts| self.write_opt(batch, opts))
417    }
418
419    pub fn write_without_wal(&self, batch: &WriteBatchWithTransaction<true>) -> Result<(), Error> {
420        let mut wo = WriteOptions::new();
421        wo.disable_wal(true);
422        self.write_opt(batch, &wo)
423    }
424
425    /// Removes the database entries in the range `["from", "to")` using given write options.
426    pub fn delete_range_cf_opt<K: AsRef<[u8]>>(
427        &self,
428        cf: &impl AsColumnFamilyRef,
429        from: K,
430        to: K,
431        writeopts: &WriteOptions,
432    ) -> Result<(), Error> {
433        let from = from.as_ref();
434        let to = to.as_ref();
435
436        unsafe {
437            ffi_try!(ffi::rocksdb_delete_range_cf(
438                self.inner.inner(),
439                writeopts.inner,
440                cf.inner(),
441                from.as_ptr() as *const c_char,
442                from.len() as size_t,
443                to.as_ptr() as *const c_char,
444                to.len() as size_t,
445            ));
446            Ok(())
447        }
448    }
449
450    /// Removes the database entries in the range `["from", "to")` using default write options.
451    pub fn delete_range_cf<K: AsRef<[u8]>>(
452        &self,
453        cf: &impl AsColumnFamilyRef,
454        from: K,
455        to: K,
456    ) -> Result<(), Error> {
457        DEFAULT_WRITE_OPTS.with(|opts| self.delete_range_cf_opt(cf, from, to, opts))
458    }
459}
rust_rocksdb/transactions/optimistic_transaction_db.rs

rust_rocksdb/transactions/
optimistic_transaction_db.rs
