Cancel all background work on last reference drop (#98)

This cancels all background work before dropping the last reference to
prevent background work from keeping the database open.

src/db_reference.rs

@@ -0,0 +1,39 @@
+use rocksdb::DB;
+use std::cell::RefCell;
+use std::sync::Arc;
+
+/// The type of a reference to a [rocksdb::DB] that is passed around the library.
+pub(crate) type DbReference = Arc<RefCell<DB>>;
+
+/// A wrapper around [DbReference] that cancels all background work when dropped.
+///
+/// All users of [rocksdb::DB] should use this wrapper instead to avoid keeping background threads
+/// alive after the database is dropped.
+#[derive(Clone)]
+pub(crate) struct DbReferenceHolder {
+ inner: Option<DbReference>,
+}
+
+impl DbReferenceHolder {
+ pub fn new(db: DB) -> Self {
+ Self {
+ inner: Some(Arc::new(RefCell::new(db))),
+ }
+ }
+
+ pub fn get(&self) -> Option<&DbReference> {
+ self.inner.as_ref()
+ }
+
+ pub fn close(&mut self) {
+ if let Some(db) = self.inner.take().and_then(Arc::into_inner) {
+ db.borrow_mut().cancel_all_background_work(true);
+ }
+ }
+}
+
+impl Drop for DbReferenceHolder {
+ fn drop(&mut self) {
+ self.close();
+ }
+}

src/iter.rs

@@ -1,20 +1,21 @@
+use crate::db_reference::DbReferenceHolder;
 use crate::encoder::{decode_value, encode_key};
+use crate::exceptions::DbClosedError;
 use crate::util::error_message;
 use crate::{ReadOpt, ReadOptionsPy};
 use core::slice;
 use libc::{c_char, c_uchar, size_t};
 use pyo3::exceptions::PyException;
 use pyo3::prelude::*;
-use rocksdb::{AsColumnFamilyRef, ColumnFamily, DB};
-use std::cell::RefCell;
+use rocksdb::{AsColumnFamilyRef, ColumnFamily};
 use std::ptr::null_mut;
 use std::sync::Arc;
 
 #[pyclass]
 #[allow(dead_code)]
 pub(crate) struct RdictIter {
  /// iterator must keep a reference count of DB to keep DB alive.
- pub(crate) db: Arc<RefCell<DB>>,
+ pub(crate) db: DbReferenceHolder,
 
  pub(crate) inner: *mut librocksdb_sys::rocksdb_iterator_t,
 
@@ -49,26 +50,29 @@ pub(crate) struct RdictValues {
 
 impl RdictIter {
  pub(crate) fn new(
- db: &Arc<RefCell<DB>>,
+ db: &DbReferenceHolder,
  cf: &Option<Arc<ColumnFamily>>,
  readopts: ReadOptionsPy,
  pickle_loads: &PyObject,
  raw_mode: bool,
  py: Python,
  ) -> PyResult<Self> {
  let readopts = readopts.to_read_opt(raw_mode, py)?;
+
+ let db_inner = db
+ .get()
+ .ok_or_else(|| DbClosedError::new_err("DB instance already closed"))?
+ .borrow()
+ .inner();
+
  Ok(RdictIter {
  db: db.clone(),
  inner: unsafe {
  match cf {
- None => {
- librocksdb_sys::rocksdb_create_iterator(db.borrow().inner(), readopts.0)
+ None => librocksdb_sys::rocksdb_create_iterator(db_inner, readopts.0),
+ Some(cf) => {
+ librocksdb_sys::rocksdb_create_iterator_cf(db_inner, readopts.0, cf.inner())
  }
- Some(cf) => librocksdb_sys::rocksdb_create_iterator_cf(
- db.borrow().inner(),
- readopts.0,
- cf.inner(),
- ),
  }
  },
  readopts,

src/lib.rs

@@ -1,4 +1,5 @@
 // #![feature(core_intrinsics)]
+mod db_reference;
 mod encoder;
 mod exceptions;
 mod iter;

src/rdict.rs

@@ -1,3 +1,4 @@
+use crate::db_reference::{DbReference, DbReferenceHolder};
 use crate::encoder::{decode_value, encode_key, encode_value};
 use crate::exceptions::DbClosedError;
 use crate::iter::{RdictItems, RdictKeys, RdictValues};
@@ -71,7 +72,7 @@ pub(crate) struct Rdict {
  pub(crate) access_type: AccessType,
  pub(crate) slice_transforms: Arc<RwLock<HashMap<String, SliceTransformType>>>,
  // drop DB last
- pub(crate) db: Option<Arc<RefCell<DB>>>,
+ pub(crate) db: DbReferenceHolder,
 }
 
 /// Define DB Access Types.
@@ -136,9 +137,9 @@ impl Rdict {
  .save(config_path)
  }
 
- fn get_db(&self) -> PyResult<&Arc<RefCell<DB>>> {
+ fn get_db(&self) -> PyResult<&DbReference> {
  self.db
- .as_ref()
+ .get()
  .ok_or_else(|| DbClosedError::new_err("DB instance already closed"))
  }
 }
@@ -266,7 +267,7 @@ impl Rdict {
  // save rocksdict config
  rocksdict_config.save(config_path)?;
  Ok(Rdict {
- db: Some(Arc::new(RefCell::new(db))),
+ db: DbReferenceHolder::new(db),
  write_opt: (&w_opt).into(),
  flush_opt: FlushOptionsPy::new(),
  read_opt: r_opt.to_read_options(options.raw_mode, py)?,
@@ -640,7 +641,7 @@ impl Rdict {
  };
 
  RdictIter::new(
- self.get_db()?,
+ &self.db,
  &self.column_family,
  read_opt,
  &self.loads,
@@ -818,7 +819,7 @@ impl Rdict {
  "column name `{name}` does not exist, use `create_cf` to creat it",
  ))),
  Some(cf) => Ok(Self {
- db: Some(db.clone()),
+ db: self.db.clone(),
  write_opt: (&self.write_opt_py).into(),
  flush_opt: self.flush_opt,
  read_opt: self.read_opt_py.to_read_options(self.opt_py.raw_mode, py)?,
@@ -855,7 +856,10 @@ impl Rdict {
  None => Err(PyException::new_err(format!(
  "column name `{name}` does not exist, use `create_cf` to creat it",
  ))),
- Some(cf) => Ok(ColumnFamilyPy { cf, db: db.clone() }),
+ Some(cf) => Ok(ColumnFamilyPy {
+ cf,
+ db: self.db.clone(),
+ }),
  }
  }
 
@@ -1014,8 +1018,8 @@ impl Rdict {
  /// Other Column Family `Rdict` instances, `ColumnFamily`
  /// (cf handle) instances, iterator instances such as`RdictIter`,
  /// `RdictItems`, `RdictKeys`, `RdictValues` can all keep RocksDB
- /// alive. `del` all associated instances mentioned above
- /// to actually shut down RocksDB.
+ /// alive. `del` or `close` all associated instances mentioned
+ /// above to actually shut down RocksDB.
  ///
  fn close(&mut self) -> PyResult<()> {
  let f_opt = &self.flush_opt;
@@ -1024,7 +1028,7 @@ impl Rdict {
  AccessTypeInner::ReadOnly { .. } | AccessTypeInner::Secondary { .. } => {
  drop(db);
  drop(self.column_family.take());
- drop(self.db.take());
+ self.db.close();
  return Ok(());
  }
  _ => (),
@@ -1037,7 +1041,7 @@ impl Rdict {
  };
  drop(db);
  drop(self.column_family.take());
- drop(self.db.take());
+ self.db.close();
  match (flush_result, flush_wal_result) {
  (Ok(_), Ok(_)) => Ok(()),
  (Err(e), Ok(_)) => Err(PyException::new_err(e.to_string())),
@@ -1257,7 +1261,7 @@ fn get_batch_inner<'a>(
 impl Drop for Rdict {
  // flush
  fn drop(&mut self) {
- if let Some(db) = &self.db {
+ if let Some(db) = self.db.get() {
  let f_opt = &self.flush_opt;
  let db = db.borrow();
  let _ = if let Some(cf) = &self.column_family {
@@ -1269,7 +1273,7 @@ impl Drop for Rdict {
  // important, always drop column families first
  // to ensure that CF handles have shorter life than DB.
  drop(self.column_family.take());
- drop(self.db.take());
+ self.db.close();
  }
 }
 
@@ -1283,7 +1287,7 @@ pub(crate) struct ColumnFamilyPy {
  // must follow this drop order
  pub(crate) cf: Arc<ColumnFamily>,
  // must keep db alive
- db: Arc<RefCell<DB>>,
+ db: DbReferenceHolder,
 }
 
 unsafe impl Send for ColumnFamilyPy {}

src/snapshot.rs

@@ -1,9 +1,10 @@
+use crate::db_reference::{DbReference, DbReferenceHolder};
 use crate::encoder::{decode_value, encode_key};
+use crate::exceptions::DbClosedError;
 use crate::{Rdict, RdictItems, RdictIter, RdictKeys, RdictValues, ReadOptionsPy};
 use pyo3::exceptions::PyException;
 use pyo3::prelude::*;
-use rocksdb::{ColumnFamily, ReadOptions, DB};
-use std::cell::RefCell;
+use rocksdb::{ColumnFamily, ReadOptions};
 use std::ops::Deref;
 use std::sync::Arc;
 
@@ -43,7 +44,7 @@ pub struct Snapshot {
  pub(crate) pickle_loads: PyObject,
  pub(crate) read_opt: ReadOptions,
  // decrease db Rc last
- pub(crate) db: Arc<RefCell<DB>>,
+ pub(crate) db: DbReferenceHolder,
  pub(crate) raw_mode: bool,
 }
 
@@ -133,7 +134,7 @@ impl Snapshot {
 
  /// read from snapshot
  fn __getitem__(&self, key: &PyAny, py: Python) -> PyResult<PyObject> {
- let db = self.db.borrow();
+ let db = self.get_db().borrow();
  let key = encode_key(key, self.raw_mode)?;
  let value_result = if let Some(cf) = &self.column_family {
  db.get_pinned_cf_opt(cf.deref(), &key[..], &self.read_opt)
@@ -152,33 +153,40 @@ impl Snapshot {
 
 impl Snapshot {
  pub(crate) fn new(rdict: &Rdict, py: Python) -> PyResult<Self> {
- if let Some(db) = &rdict.db {
- let db_borrow = db.borrow();
- let snapshot = unsafe { librocksdb_sys::rocksdb_create_snapshot(db_borrow.inner()) };
- let r_opt: ReadOptions = rdict
- .read_opt_py
- .to_read_options(rdict.opt_py.raw_mode, py)?;
- unsafe {
- set_snapshot(r_opt.inner(), snapshot);
- }
- Ok(Snapshot {
- inner: snapshot,
- column_family: rdict.column_family.clone(),
- pickle_loads: rdict.loads.clone(),
- read_opt: r_opt,
- db: db.clone(),
- raw_mode: rdict.opt_py.raw_mode,
- })
- } else {
- Err(PyException::new_err("DB already closed"))
+ let db_inner = rdict
+ .db
+ .get()
+ .ok_or_else(|| DbClosedError::new_err("DB instance already closed"))?
+ .borrow()
+ .inner();
+ let snapshot = unsafe { librocksdb_sys::rocksdb_create_snapshot(db_inner) };
+ let r_opt: ReadOptions = rdict
+ .read_opt_py
+ .to_read_options(rdict.opt_py.raw_mode, py)?;
+ unsafe {
+ set_snapshot(r_opt.inner(), snapshot);
  }
+ Ok(Snapshot {
+ inner: snapshot,
+ column_family: rdict.column_family.clone(),
+ pickle_loads: rdict.loads.clone(),
+ read_opt: r_opt,
+ db: rdict.db.clone(),
+ raw_mode: rdict.opt_py.raw_mode,
+ })
+ }
+
+ fn get_db(&self) -> &DbReference {
+ self.db
+ .get()
+ .expect("Snapshot should never close its DbReference")
  }
 }
 
 impl Drop for Snapshot {
  fn drop(&mut self) {
  unsafe {
- librocksdb_sys::rocksdb_release_snapshot(self.db.borrow().inner(), self.inner);
+ librocksdb_sys::rocksdb_release_snapshot(self.get_db().borrow().inner(), self.inner);
  }
  }
 }