rust_rocksdb/

env.rs

use std::path::Path;
use std::sync::Arc;

use libc::{self, c_int};

use crate::{Error, ffi, ffi_util::to_cpath};

/// An Env is an interface used by the rocksdb implementation to access
/// operating system functionality like the filesystem etc. Callers
/// may wish to provide a custom Env object when opening a database to
/// get fine gain control; e.g., to rate limit file system operations.
///
/// All Env implementations are safe for concurrent access from
/// multiple threads without any external synchronization.
///
/// Note: currently, C API behinds C++ API for various settings.
/// See also: `rocksdb/include/env.h`
#[derive(Clone)]
pub struct Env(pub(crate) Arc<EnvWrapper>);

pub(crate) struct EnvWrapper {
    pub(crate) inner: *mut ffi::rocksdb_env_t,
}

impl Drop for EnvWrapper {
    fn drop(&mut self) {
        unsafe {
            ffi::rocksdb_env_destroy(self.inner);
        }
    }
}

impl Env {
    /// Returns default env
    pub fn new() -> Result<Self, Error> {
        let env = unsafe { ffi::rocksdb_create_default_env() };
        if env.is_null() {
            Err(Error::new("Could not create mem env".to_owned()))
        } else {
            Ok(Self(Arc::new(EnvWrapper { inner: env })))
        }
    }

    /// Returns a new environment that stores its data in memory and delegates
    /// all non-file-storage tasks to base_env.
    pub fn mem_env() -> Result<Self, Error> {
        let env = unsafe { ffi::rocksdb_create_mem_env() };
        if env.is_null() {
            Err(Error::new("Could not create mem env".to_owned()))
        } else {
            Ok(Self(Arc::new(EnvWrapper { inner: env })))
        }
    }

    /// Returns a new environment which wraps and takes ownership of the provided
    /// raw environment.
    ///
    /// # Safety
    ///
    /// Ownership of `env` is transferred to the returned Env, which becomes
    /// responsible for freeing it. The caller should forget the raw pointer
    /// after this call.
    ///
    /// # When would I use this?
    ///
    /// RocksDB's C++ [Env](https://github.com/facebook/rocksdb/blob/main/include/rocksdb/env.h)
    /// class provides many extension points for low-level database subsystems, such as file IO.
    /// These subsystems aren't covered within the scope of the C interface or this crate,
    /// but from_raw() may be used to hand a pre-instrumented Env to this crate for further use.
    ///
    pub unsafe fn from_raw(env: *mut ffi::rocksdb_env_t) -> Self {
        Self(Arc::new(EnvWrapper { inner: env }))
    }

    /// Sets the number of background worker threads of a specific thread pool for this environment.
    /// `LOW` is the default pool.
    ///
    /// Default: 1
    pub fn set_background_threads(&mut self, num_threads: c_int) {
        unsafe {
            ffi::rocksdb_env_set_background_threads(self.0.inner, num_threads);
        }
    }

    /// Sets the size of the high priority thread pool that can be used to
    /// prevent compactions from stalling memtable flushes.
    pub fn set_high_priority_background_threads(&mut self, n: c_int) {
        unsafe {
            ffi::rocksdb_env_set_high_priority_background_threads(self.0.inner, n);
        }
    }

    /// Sets the size of the low priority thread pool that can be used to
    /// prevent compactions from stalling memtable flushes.
    pub fn set_low_priority_background_threads(&mut self, n: c_int) {
        unsafe {
            ffi::rocksdb_env_set_low_priority_background_threads(self.0.inner, n);
        }
    }

    /// Sets the size of the bottom priority thread pool that can be used to
    /// prevent compactions from stalling memtable flushes.
    pub fn set_bottom_priority_background_threads(&mut self, n: c_int) {
        unsafe {
            ffi::rocksdb_env_set_bottom_priority_background_threads(self.0.inner, n);
        }
    }

    /// Wait for all threads started by StartThread to terminate.
    pub fn join_all_threads(&mut self) {
        unsafe {
            ffi::rocksdb_env_join_all_threads(self.0.inner);
        }
    }

    /// Lowering IO priority for threads from the specified pool.
    pub fn lower_thread_pool_io_priority(&mut self) {
        unsafe {
            ffi::rocksdb_env_lower_thread_pool_io_priority(self.0.inner);
        }
    }

    /// Lowering IO priority for high priority thread pool.
    pub fn lower_high_priority_thread_pool_io_priority(&mut self) {
        unsafe {
            ffi::rocksdb_env_lower_high_priority_thread_pool_io_priority(self.0.inner);
        }
    }

    /// Lowering CPU priority for threads from the specified pool.
    pub fn lower_thread_pool_cpu_priority(&mut self) {
        unsafe {
            ffi::rocksdb_env_lower_thread_pool_cpu_priority(self.0.inner);
        }
    }

    /// Lowering CPU priority for high priority thread pool.
    pub fn lower_high_priority_thread_pool_cpu_priority(&mut self) {
        unsafe {
            ffi::rocksdb_env_lower_high_priority_thread_pool_cpu_priority(self.0.inner);
        }
    }

    /// Returns the current `background_threads` setting.
    ///
    /// See [`Self::set_background_threads`] for what this controls.
    pub fn get_background_threads(&self) -> c_int {
        unsafe { ffi::rocksdb_env_get_background_threads(self.0.inner) }
    }

    /// Returns the current `bottom_priority_background_threads` setting.
    ///
    /// See [`Self::set_bottom_priority_background_threads`] for what this controls.
    pub fn get_bottom_priority_background_threads(&self) -> c_int {
        unsafe { ffi::rocksdb_env_get_bottom_priority_background_threads(self.0.inner) }
    }

    /// Returns the current `high_priority_background_threads` setting.
    ///
    /// See [`Self::set_high_priority_background_threads`] for what this controls.
    pub fn get_high_priority_background_threads(&self) -> c_int {
        unsafe { ffi::rocksdb_env_get_high_priority_background_threads(self.0.inner) }
    }

    /// Returns the current `low_priority_background_threads` setting.
    ///
    /// See [`Self::set_low_priority_background_threads`] for what this controls.
    pub fn get_low_priority_background_threads(&self) -> c_int {
        unsafe { ffi::rocksdb_env_get_low_priority_background_threads(self.0.inner) }
    }

    /// Creates `path` through this environment, and returns `Ok` if it is
    /// already there.
    ///
    /// This is `Env::CreateDirIfMissing`, one level only. The POSIX
    /// implementation is a plain `mkdir`, so every parent has to exist
    /// already. It is not the recursive `std::fs::create_dir_all`.
    ///
    /// # Errors
    ///
    /// The POSIX environment fails when a parent directory is missing, when
    /// the process cannot write there, or when `path` exists and is not a
    /// directory. [`Env::mem_env`] always succeeds.
    pub fn create_dir_if_missing<P: AsRef<Path>>(&self, path: P) -> Result<(), Error> {
        let cpath = to_cpath(path)?;
        unsafe {
            ffi_try!(ffi::rocksdb_create_dir_if_missing(
                self.0.inner,
                cpath.as_ptr(),
            ));
        }
        Ok(())
    }
}

unsafe impl Send for EnvWrapper {}
unsafe impl Sync for EnvWrapper {}

/// Priority at which an IO operation is charged to the rate limiter set with
/// [`Options::set_ratelimiter`](crate::Options::set_ratelimiter).
///
/// Mirrors `Env::IOPriority` from `rocksdb/include/env.h`.
#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
#[repr(i32)]
pub enum IoPriority {
    Low = 0,
    Mid = 1,
    High = 2,
    User = 3,
    /// Do not charge the rate limiter at all.
    Total = 4,
}

impl IoPriority {
    pub(crate) fn try_from_raw(raw: c_int) -> Option<Self> {
        match raw {
            0 => Some(Self::Low),
            1 => Some(Self::Mid),
            2 => Some(Self::High),
            3 => Some(Self::User),
            4 => Some(Self::Total),
            _ => None,
        }
    }
}