Skip to content

Commit

Permalink
Check the validity of the env and transaction when opening databases
Browse files Browse the repository at this point in the history
  • Loading branch information
@Kerollmops
Kerollmops committed Jan 11, 2023
1 parent eb2f85e commit a1ed4b1
Show file tree
Hide file tree
Showing 3 changed files with 50 additions and 31 deletions.
54 changes: 27 additions & 27 deletions heed/src/db/polymorph.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -158,7 +158,7 @@ impl PolyDatabase {
KC: BytesEncode<'a>,
DC: BytesDecode<'txn>,
{
assert_matching_env_txn!(self, txn);
assert_eq_env_db_txn!(self, txn);

let key_bytes: Cow<[u8]> = KC::bytes_encode(&key).map_err(Error::Encoding)?;

Expand Down Expand Up @@ -232,7 +232,7 @@ impl PolyDatabase {
KC: BytesEncode<'a> + BytesDecode<'txn>,
DC: BytesDecode<'txn>,
{
assert_matching_env_txn!(self, txn);
assert_eq_env_db_txn!(self, txn);

let mut cursor = RoCursor::new(txn, self.dbi)?;
let key_bytes: Cow<[u8]> = KC::bytes_encode(&key).map_err(Error::Encoding)?;
Expand Down Expand Up @@ -300,7 +300,7 @@ impl PolyDatabase {
KC: BytesEncode<'a> + BytesDecode<'txn>,
DC: BytesDecode<'txn>,
{
assert_matching_env_txn!(self, txn);
assert_eq_env_db_txn!(self, txn);

let mut cursor = RoCursor::new(txn, self.dbi)?;
let key_bytes: Cow<[u8]> = KC::bytes_encode(&key).map_err(Error::Encoding)?;
Expand Down Expand Up @@ -372,7 +372,7 @@ impl PolyDatabase {
KC: BytesEncode<'a> + BytesDecode<'txn>,
DC: BytesDecode<'txn>,
{
assert_matching_env_txn!(self, txn);
assert_eq_env_db_txn!(self, txn);

let mut cursor = RoCursor::new(txn, self.dbi)?;
let key_bytes: Cow<[u8]> = KC::bytes_encode(&key).map_err(Error::Encoding)?;
Expand Down Expand Up @@ -443,7 +443,7 @@ impl PolyDatabase {
KC: BytesEncode<'a> + BytesDecode<'txn>,
DC: BytesDecode<'txn>,
{
assert_matching_env_txn!(self, txn);
assert_eq_env_db_txn!(self, txn);

let mut cursor = RoCursor::new(txn, self.dbi)?;
let key_bytes: Cow<[u8]> = KC::bytes_encode(&key).map_err(Error::Encoding)?;
Expand Down Expand Up @@ -497,7 +497,7 @@ impl PolyDatabase {
KC: BytesDecode<'txn>,
DC: BytesDecode<'txn>,
{
assert_matching_env_txn!(self, txn);
assert_eq_env_db_txn!(self, txn);

let mut cursor = RoCursor::new(txn, self.dbi)?;
match cursor.move_on_first() {
Expand Down Expand Up @@ -550,7 +550,7 @@ impl PolyDatabase {
KC: BytesDecode<'txn>,
DC: BytesDecode<'txn>,
{
assert_matching_env_txn!(self, txn);
assert_eq_env_db_txn!(self, txn);

let mut cursor = RoCursor::new(txn, self.dbi)?;
match cursor.move_on_last() {
Expand Down Expand Up @@ -603,7 +603,7 @@ impl PolyDatabase {
/// # Ok(()) }
/// ```
pub fn len<'txn>(&self, txn: &'txn RoTxn) -> Result<u64> {
assert_matching_env_txn!(self, txn);
assert_eq_env_db_txn!(self, txn);

let mut db_stat = mem::MaybeUninit::uninit();
let result = unsafe { mdb_result(ffi::mdb_stat(txn.txn, self.dbi, db_stat.as_mut_ptr())) };
Expand Down Expand Up @@ -656,7 +656,7 @@ impl PolyDatabase {
/// # Ok(()) }
/// ```
pub fn is_empty<'txn>(&self, txn: &'txn RoTxn) -> Result<bool> {
assert_matching_env_txn!(self, txn);
assert_eq_env_db_txn!(self, txn);

let mut cursor = RoCursor::new(txn, self.dbi)?;
match cursor.move_on_first()? {
Expand Down Expand Up @@ -702,7 +702,7 @@ impl PolyDatabase {
/// # Ok(()) }
/// ```
pub fn iter<'txn, KC, DC>(&self, txn: &'txn RoTxn) -> Result<RoIter<'txn, KC, DC>> {
assert_matching_env_txn!(self, txn);
assert_eq_env_db_txn!(self, txn);

RoCursor::new(txn, self.dbi).map(|cursor| RoIter::new(cursor))
}
Expand Down Expand Up @@ -757,7 +757,7 @@ impl PolyDatabase {
/// # Ok(()) }
/// ```
pub fn iter_mut<'txn, KC, DC>(&self, txn: &'txn mut RwTxn) -> Result<RwIter<'txn, KC, DC>> {
assert_matching_env_txn!(self, txn);
assert_eq_env_db_txn!(self, txn);

RwCursor::new(txn, self.dbi).map(|cursor| RwIter::new(cursor))
}
Expand Down Expand Up @@ -799,7 +799,7 @@ impl PolyDatabase {
/// # Ok(()) }
/// ```
pub fn rev_iter<'txn, KC, DC>(&self, txn: &'txn RoTxn) -> Result<RoRevIter<'txn, KC, DC>> {
assert_matching_env_txn!(self, txn);
assert_eq_env_db_txn!(self, txn);

RoCursor::new(txn, self.dbi).map(|cursor| RoRevIter::new(cursor))
}
Expand Down Expand Up @@ -858,7 +858,7 @@ impl PolyDatabase {
&self,
txn: &'txn mut RwTxn,
) -> Result<RwRevIter<'txn, KC, DC>> {
assert_matching_env_txn!(self, txn);
assert_eq_env_db_txn!(self, txn);

RwCursor::new(txn, self.dbi).map(|cursor| RwRevIter::new(cursor))
}
Expand Down Expand Up @@ -911,7 +911,7 @@ impl PolyDatabase {
KC: BytesEncode<'a>,
R: RangeBounds<KC::EItem>,
{
assert_matching_env_txn!(self, txn);
assert_eq_env_db_txn!(self, txn);

let start_bound = match range.start_bound() {
Bound::Included(bound) => {
Expand Down Expand Up @@ -1002,7 +1002,7 @@ impl PolyDatabase {
KC: BytesEncode<'a>,
R: RangeBounds<KC::EItem>,
{
assert_matching_env_txn!(self, txn);
assert_eq_env_db_txn!(self, txn);

let start_bound = match range.start_bound() {
Bound::Included(bound) => {
Expand Down Expand Up @@ -1080,7 +1080,7 @@ impl PolyDatabase {
KC: BytesEncode<'a>,
R: RangeBounds<KC::EItem>,
{
assert_matching_env_txn!(self, txn);
assert_eq_env_db_txn!(self, txn);

let start_bound = match range.start_bound() {
Bound::Included(bound) => {
Expand Down Expand Up @@ -1171,7 +1171,7 @@ impl PolyDatabase {
KC: BytesEncode<'a>,
R: RangeBounds<KC::EItem>,
{
assert_matching_env_txn!(self, txn);
assert_eq_env_db_txn!(self, txn);

let start_bound = match range.start_bound() {
Bound::Included(bound) => {
Expand Down Expand Up @@ -1249,7 +1249,7 @@ impl PolyDatabase {
where
KC: BytesEncode<'a>,
{
assert_matching_env_txn!(self, txn);
assert_eq_env_db_txn!(self, txn);

let prefix_bytes = KC::bytes_encode(prefix).map_err(Error::Encoding)?;
let prefix_bytes = prefix_bytes.into_owned();
Expand Down Expand Up @@ -1318,7 +1318,7 @@ impl PolyDatabase {
where
KC: BytesEncode<'a>,
{
assert_matching_env_txn!(self, txn);
assert_eq_env_db_txn!(self, txn);

let prefix_bytes = KC::bytes_encode(prefix).map_err(Error::Encoding)?;
let prefix_bytes = prefix_bytes.into_owned();
Expand Down Expand Up @@ -1374,7 +1374,7 @@ impl PolyDatabase {
where
KC: BytesEncode<'a>,
{
assert_matching_env_txn!(self, txn);
assert_eq_env_db_txn!(self, txn);

let prefix_bytes = KC::bytes_encode(prefix).map_err(Error::Encoding)?;
let prefix_bytes = prefix_bytes.into_owned();
Expand Down Expand Up @@ -1443,7 +1443,7 @@ impl PolyDatabase {
where
KC: BytesEncode<'a>,
{
assert_matching_env_txn!(self, txn);
assert_eq_env_db_txn!(self, txn);

let prefix_bytes = KC::bytes_encode(prefix).map_err(Error::Encoding)?;
let prefix_bytes = prefix_bytes.into_owned();
Expand Down Expand Up @@ -1493,7 +1493,7 @@ impl PolyDatabase {
KC: BytesEncode<'a>,
DC: BytesEncode<'a>,
{
assert_matching_env_txn!(self, txn);
assert_eq_env_db_txn!(self, txn);

let key_bytes: Cow<[u8]> = KC::bytes_encode(&key).map_err(Error::Encoding)?;
let data_bytes: Cow<[u8]> = DC::bytes_encode(&data).map_err(Error::Encoding)?;
Expand Down Expand Up @@ -1554,7 +1554,7 @@ impl PolyDatabase {
KC: BytesEncode<'a>,
F: FnMut(&mut ReservedSpace) -> io::Result<()>,
{
assert_matching_env_txn!(self, txn);
assert_eq_env_db_txn!(self, txn);

let key_bytes: Cow<[u8]> = KC::bytes_encode(&key).map_err(Error::Encoding)?;
let mut key_val = unsafe { crate::into_val(&key_bytes) };
Expand Down Expand Up @@ -1620,7 +1620,7 @@ impl PolyDatabase {
KC: BytesEncode<'a>,
DC: BytesEncode<'a>,
{
assert_matching_env_txn!(self, txn);
assert_eq_env_db_txn!(self, txn);

let key_bytes: Cow<[u8]> = KC::bytes_encode(&key).map_err(Error::Encoding)?;
let data_bytes: Cow<[u8]> = DC::bytes_encode(&data).map_err(Error::Encoding)?;
Expand Down Expand Up @@ -1681,7 +1681,7 @@ impl PolyDatabase {
where
KC: BytesEncode<'a>,
{
assert_matching_env_txn!(self, txn);
assert_eq_env_db_txn!(self, txn);

let key_bytes: Cow<[u8]> = KC::bytes_encode(&key).map_err(Error::Encoding)?;
let mut key_val = unsafe { crate::into_val(&key_bytes) };
Expand Down Expand Up @@ -1749,7 +1749,7 @@ impl PolyDatabase {
KC: BytesEncode<'a> + BytesDecode<'txn>,
R: RangeBounds<KC::EItem>,
{
assert_matching_env_txn!(self, txn);
assert_eq_env_db_txn!(self, txn);

let mut count = 0;
let mut iter = self.range_mut::<KC, DecodeIgnore, _>(txn, range)?;
Expand Down Expand Up @@ -1806,7 +1806,7 @@ impl PolyDatabase {
/// # Ok(()) }
/// ```
pub fn clear(&self, txn: &mut RwTxn) -> Result<()> {
assert_matching_env_txn!(self, txn);
assert_eq_env_db_txn!(self, txn);

unsafe { mdb_result(ffi::mdb_drop(txn.txn.txn, self.dbi, 0)).map_err(Into::into) }
}
Expand Down
14 changes: 12 additions & 2 deletions heed/src/env.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,9 @@ use synchronoise::event::SignalEvent;

use crate::mdb::error::mdb_result;
use crate::mdb::ffi;
use crate::{Database, Error, Flags, PolyDatabase, Result, RoCursor, RoTxn, RwTxn};
use crate::{
assert_eq_env_txn, Database, Error, Flags, PolyDatabase, Result, RoCursor, RoTxn, RwTxn,
};

/// The list of opened environments, the value is an optional environment, it is None
/// when someone asks to close the environment, closing is a two-phase step, to make sure
Expand Down Expand Up @@ -297,7 +299,7 @@ impl Drop for EnvInner {
unsafe {
let _ = ffi::mdb_env_close(self.env);
}
// We signal to all the waiters that we have closed the env.
// We signal to all the waiters that the env is closed now.
signal_event.signal();
}
}
Expand Down Expand Up @@ -442,6 +444,8 @@ impl Env {
KC: 'static,
DC: 'static,
{
assert_eq_env_txn!(self, rtxn);

let types = (TypeId::of::<KC>(), TypeId::of::<DC>());
match self.raw_init_database(rtxn.txn, name, Some(types), false) {
Ok(dbi) => Ok(Some(Database::new(self.env_mut_ptr() as _, dbi))),
Expand All @@ -464,6 +468,8 @@ impl Env {
rtxn: &RoTxn,
name: Option<&str>,
) -> Result<Option<PolyDatabase>> {
assert_eq_env_txn!(self, rtxn);

match self.raw_init_database(rtxn.txn, name, None, false) {
Ok(dbi) => Ok(Some(PolyDatabase::new(self.env_mut_ptr() as _, dbi))),
Err(Error::Mdb(e)) if e.not_found() => Ok(None),
Expand All @@ -489,6 +495,8 @@ impl Env {
KC: 'static,
DC: 'static,
{
assert_eq_env_txn!(self, wtxn);

let types = (TypeId::of::<KC>(), TypeId::of::<DC>());
match self.raw_init_database(wtxn.txn.txn, name, Some(types), true) {
Ok(dbi) => Ok(Database::new(self.env_mut_ptr() as _, dbi)),
Expand All @@ -510,6 +518,8 @@ impl Env {
wtxn: &mut RwTxn,
name: Option<&str>,
) -> Result<PolyDatabase> {
assert_eq_env_txn!(self, wtxn);

match self.raw_init_database(wtxn.txn.txn, name, None, true) {
Ok(dbi) => Ok(PolyDatabase::new(self.env_mut_ptr() as _, dbi)),
Err(e) => Err(e),
Expand Down
13 changes: 11 additions & 2 deletions heed/src/lib.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -146,7 +146,7 @@ mod tests {
}
}

macro_rules! assert_matching_env_txn {
macro_rules! assert_eq_env_db_txn {
($database:ident, $txn:ident) => {
assert!(
$database.env_ident == $txn.env_mut_ptr() as usize,
Expand All @@ -155,4 +155,13 @@ macro_rules! assert_matching_env_txn {
};
}

pub(crate) use assert_matching_env_txn;
macro_rules! assert_eq_env_txn {
($env:ident, $txn:ident) => {
assert!(
$env.env_mut_ptr() == $txn.env_mut_ptr(),
"The environment doesn't match the transaction's environment"
);
};
}

pub(crate) use {assert_eq_env_db_txn, assert_eq_env_txn};

0 comments on commit a1ed4b1

Please sign in to comment.