[rm-sled] Make proper use of pinning in LMDB adapter + comment unsafe

This commit is contained in:
Alex 2024-03-08 17:39:17 +01:00
parent b942949940
commit 32aa246300
Signed by untrusted user: lx
GPG key ID: 0E496D15096376BE
2 changed files with 47 additions and 25 deletions

View file

@ -3,6 +3,7 @@ use core::ptr::NonNull;
use std::collections::HashMap; use std::collections::HashMap;
use std::convert::TryInto; use std::convert::TryInto;
use std::pin::Pin;
use std::sync::{Arc, RwLock}; use std::sync::{Arc, RwLock};
use heed::types::ByteSlice; use heed::types::ByteSlice;
@ -319,12 +320,20 @@ where
where where
F: FnOnce(&'a RoTxn<'a>) -> Result<I>, F: FnOnce(&'a RoTxn<'a>) -> Result<I>,
{ {
let mut res = TxAndIterator { tx, iter: None }; let res = TxAndIterator { tx, iter: None };
let mut boxed = Box::pin(res);
let tx = unsafe { NonNull::from(&res.tx).as_ref() }; // This unsafe allows us to bypass lifetime checks
res.iter = Some(iterfun(tx)?); let tx = unsafe { NonNull::from(&boxed.tx).as_ref() };
let iter = iterfun(tx)?;
Ok(Box::new(res)) let mut_ref = Pin::as_mut(&mut boxed);
// This unsafe allows us to write in a field of the pinned struct
unsafe {
Pin::get_unchecked_mut(mut_ref).iter = Some(iter);
}
Ok(Box::new(TxAndIteratorPin(boxed)))
} }
} }
@ -338,14 +347,21 @@ where
} }
} }
impl<'a, I> Iterator for TxAndIterator<'a, I> struct TxAndIteratorPin<'a, I>(Pin<Box<TxAndIterator<'a, I>>>)
where
I: Iterator<Item = IteratorItem<'a>> + 'a;
impl<'a, I> Iterator for TxAndIteratorPin<'a, I>
where where
I: Iterator<Item = IteratorItem<'a>> + 'a, I: Iterator<Item = IteratorItem<'a>> + 'a,
{ {
type Item = Result<(Value, Value)>; type Item = Result<(Value, Value)>;
fn next(&mut self) -> Option<Self::Item> { fn next(&mut self) -> Option<Self::Item> {
match self.iter.as_mut().unwrap().next() { let mut_ref = Pin::as_mut(&mut self.0);
// This unsafe allows us to mutably access the iterator field
let next = unsafe { Pin::get_unchecked_mut(mut_ref).iter.as_mut()?.next() };
match next {
None => None, None => None,
Some(Err(e)) => Some(Err(e.into())), Some(Err(e)) => Some(Err(e.into())),
Some(Ok((k, v))) => Some(Ok((k.to_vec(), v.to_vec()))), Some(Ok((k, v))) => Some(Ok((k.to_vec(), v.to_vec()))),

View file

@ -444,17 +444,23 @@ impl<'a> DbValueIterator<'a> {
let mut boxed = Box::pin(res); let mut boxed = Box::pin(res);
trace!("make iterator with sql: {}", sql); trace!("make iterator with sql: {}", sql);
// This unsafe allows us to bypass lifetime checks
let db = unsafe { NonNull::from(&boxed.db).as_ref() };
let stmt = db.db.prepare(sql)?;
let mut_ref = Pin::as_mut(&mut boxed);
// This unsafe allows us to write in a field of the pinned struct
unsafe { unsafe {
let db = NonNull::from(&boxed.db);
let stmt = db.as_ref().db.prepare(sql)?;
let mut_ref: Pin<&mut DbValueIterator<'a>> = Pin::as_mut(&mut boxed);
Pin::get_unchecked_mut(mut_ref).stmt = Some(stmt); Pin::get_unchecked_mut(mut_ref).stmt = Some(stmt);
}
let mut stmt = NonNull::from(&boxed.stmt); // This unsafe allows us to bypass lifetime checks
let iter = stmt.as_mut().as_mut().unwrap().query(args)?; let stmt = unsafe { NonNull::from(&boxed.stmt).as_mut() };
let iter = stmt.as_mut().unwrap().query(args)?;
let mut_ref: Pin<&mut DbValueIterator<'a>> = Pin::as_mut(&mut boxed); let mut_ref = Pin::as_mut(&mut boxed);
// This unsafe allows us to write in a field of the pinned struct
unsafe {
Pin::get_unchecked_mut(mut_ref).iter = Some(iter); Pin::get_unchecked_mut(mut_ref).iter = Some(iter);
} }
@ -476,10 +482,9 @@ impl<'a> Iterator for DbValueIteratorPin<'a> {
type Item = Result<(Value, Value)>; type Item = Result<(Value, Value)>;
fn next(&mut self) -> Option<Self::Item> { fn next(&mut self) -> Option<Self::Item> {
let next = unsafe { let mut_ref = Pin::as_mut(&mut self.0);
let mut_ref: Pin<&mut DbValueIterator<'a>> = Pin::as_mut(&mut self.0); // This unsafe allows us to mutably access the iterator field
Pin::get_unchecked_mut(mut_ref).iter.as_mut()?.next() let next = unsafe { Pin::get_unchecked_mut(mut_ref).iter.as_mut()?.next() };
};
let row = match next { let row = match next {
Err(e) => return Some(Err(e.into())), Err(e) => return Some(Err(e.into())),
Ok(None) => return None, Ok(None) => return None,
@ -522,11 +527,13 @@ impl<'a> TxValueIterator<'a> {
let mut boxed = Box::pin(res); let mut boxed = Box::pin(res);
trace!("make iterator with sql: {}", sql); trace!("make iterator with sql: {}", sql);
unsafe { // This unsafe allows us to bypass lifetime checks
let mut stmt = NonNull::from(&boxed.stmt); let stmt = unsafe { NonNull::from(&boxed.stmt).as_mut() };
let iter = stmt.as_mut().query(args)?; let iter = stmt.query(args)?;
let mut_ref: Pin<&mut TxValueIterator<'a>> = Pin::as_mut(&mut boxed); let mut_ref = Pin::as_mut(&mut boxed);
// This unsafe allows us to write in a field of the pinned struct
unsafe {
Pin::get_unchecked_mut(mut_ref).iter = Some(iter); Pin::get_unchecked_mut(mut_ref).iter = Some(iter);
} }
@ -547,10 +554,9 @@ impl<'a> Iterator for TxValueIteratorPin<'a> {
type Item = TxOpResult<(Value, Value)>; type Item = TxOpResult<(Value, Value)>;
fn next(&mut self) -> Option<Self::Item> { fn next(&mut self) -> Option<Self::Item> {
let next = unsafe { let mut_ref = Pin::as_mut(&mut self.0);
let mut_ref: Pin<&mut TxValueIterator<'a>> = Pin::as_mut(&mut self.0); // This unsafe allows us to mutably access the iterator field
Pin::get_unchecked_mut(mut_ref).iter.as_mut()?.next() let next = unsafe { Pin::get_unchecked_mut(mut_ref).iter.as_mut()?.next() };
};
let row = match next { let row = match next {
Err(e) => return Some(Err(e.into())), Err(e) => return Some(Err(e.into())),
Ok(None) => return None, Ok(None) => return None,