migrate blocks in batches (rework db iterator impl) (#3450)

* rework db iterator more robust block migration based on low level iterator * cleanup * cleanup * fixup
2020-10-07 11:47:25 +01:00
parent eab26b3ae4
commit e7bbda81a0
10 changed files with 187 additions and 128 deletions
@@ -15,7 +15,6 @@
 //! Storage of core types using LMDB.

 use std::fs;
-use std::marker;
 use std::sync::Arc;

 use lmdb_zero as lmdb;
@@ -45,13 +44,13 @@ pub enum Error {
 	#[fail(display = "LMDB error: {} ", _0)]
 	LmdbErr(lmdb::error::Error),
 	/// Wraps a serialization error for Writeable or Readable
-	#[fail(display = "Serialization Error")]
-	SerErr(String),
+	#[fail(display = "Serialization Error: {}", _0)]
+	SerErr(ser::Error),
 	/// File handling error
-	#[fail(display = "File handling Error")]
+	#[fail(display = "File handling Error: {}", _0)]
 	FileErr(String),
 	/// Other error
-	#[fail(display = "Other Error")]
+	#[fail(display = "Other Error: {}", _0)]
 	OtherErr(String),
 }

@@ -61,6 +60,12 @@ impl From<lmdb::error::Error> for Error {
 	}
 }

+impl From<ser::Error> for Error {
+	fn from(e: ser::Error) -> Error {
+		Error::SerErr(e)
+	}
+}
+
 /// unwraps the inner option by converting the none case to a not found error
 pub fn option_to_not_found<T, F>(res: Result<Option<T>, Error>, field_name: F) -> Result<T, Error>
 where
@@ -247,7 +252,7 @@ impl Store {

 	/// Gets a value from the db, provided its key.
 	/// Deserializes the retrieved data using the provided function.
-	pub fn get_with<T, F>(
+	pub fn get_with<F, T>(
 		&self,
 		key: &[u8],
 		access: &lmdb::ConstAccessor<'_>,
@@ -255,12 +260,12 @@ impl Store {
 		deserialize: F,
 	) -> Result<Option<T>, Error>
 	where
-		F: Fn(&[u8]) -> Result<T, Error>,
+		F: Fn(&[u8], &[u8]) -> Result<T, Error>,
 	{
 		let res: Option<&[u8]> = access.get(db, key).to_opt()?;
 		match res {
 			None => Ok(None),
-			Some(res) => deserialize(res).map(|x| Some(x)),
+			Some(res) => deserialize(key, res).map(Some),
 		}
 	}

@@ -274,9 +279,8 @@ impl Store {
 		let txn = lmdb::ReadTransaction::new(self.env.clone())?;
 		let access = txn.access();

-		self.get_with(key, &access, &db, |mut data| {
-			ser::deserialize(&mut data, self.protocol_version())
-				.map_err(|e| Error::SerErr(format!("{}", e)))
+		self.get_with(key, &access, &db, |_, mut data| {
+			ser::deserialize(&mut data, self.protocol_version()).map_err(From::from)
 		})
 	}

@@ -293,23 +297,18 @@ impl Store {
 		Ok(res.is_some())
 	}

-	/// Produces an iterator of (key, value) pairs, where values are `Readable` types
-	/// moving forward from the provided key.
-	pub fn iter<T: ser::Readable>(&self, from: &[u8]) -> Result<SerIterator<T>, Error> {
+	/// Produces an iterator from the provided key prefix.
+	pub fn iter<F, T>(&self, prefix: &[u8], deserialize: F) -> Result<PrefixIterator<F, T>, Error>
+	where
+		F: Fn(&[u8], &[u8]) -> Result<T, Error>,
+	{
 		let lock = self.db.read();
 		let db = lock
 			.as_ref()
 			.ok_or_else(|| Error::NotFoundErr("chain db is None".to_string()))?;
 		let tx = Arc::new(lmdb::ReadTransaction::new(self.env.clone())?);
 		let cursor = Arc::new(tx.cursor(db.clone())?);
-		Ok(SerIterator {
-			tx,
-			cursor,
-			seek: false,
-			prefix: from.to_vec(),
-			version: self.protocol_version(),
-			_marker: marker::PhantomData,
-		})
+		Ok(PrefixIterator::new(tx, cursor, prefix, deserialize))
 	}

 	/// Builds a new batch to be used with this store.
@@ -364,15 +363,15 @@ impl<'a> Batch<'a> {
 		let ser_value = ser::ser_vec(value, version);
 		match ser_value {
 			Ok(data) => self.put(key, &data),
-			Err(err) => Err(Error::SerErr(format!("{}", err))),
+			Err(err) => Err(err.into()),
 		}
 	}

 	/// Low-level access for retrieving data by key.
 	/// Takes a function for flexible deserialization.
-	pub fn get_with<T, F>(&self, key: &[u8], deserialize: F) -> Result<Option<T>, Error>
+	pub fn get_with<F, T>(&self, key: &[u8], deserialize: F) -> Result<Option<T>, Error>
 	where
-		F: Fn(&[u8]) -> Result<T, Error>,
+		F: Fn(&[u8], &[u8]) -> Result<T, Error>,
 	{
 		let access = self.tx.access();
 		let lock = self.store.db.read();
@@ -395,18 +394,20 @@ impl<'a> Batch<'a> {
 		Ok(res.is_some())
 	}

-	/// Produces an iterator of `Readable` types moving forward from the
-	/// provided key.
-	pub fn iter<T: ser::Readable>(&self, from: &[u8]) -> Result<SerIterator<T>, Error> {
-		self.store.iter(from)
+	/// Produces an iterator from the provided key prefix.
+	pub fn iter<F, T>(&self, prefix: &[u8], deserialize: F) -> Result<PrefixIterator<F, T>, Error>
+	where
+		F: Fn(&[u8], &[u8]) -> Result<T, Error>,
+	{
+		self.store.iter(prefix, deserialize)
 	}

 	/// Gets a `Readable` value from the db by provided key and default deserialization strategy.
 	pub fn get_ser<T: ser::Readable>(&self, key: &[u8]) -> Result<Option<T>, Error> {
-		self.get_with(key, |mut data| {
+		self.get_with(key, |_, mut data| {
 			match ser::deserialize(&mut data, self.protocol_version()) {
 				Ok(res) => Ok(res),
-				Err(e) => Err(Error::SerErr(format!("{}", e))),
+				Err(e) => Err(From::from(e)),
 			}
 		})
 	}
@@ -437,57 +438,61 @@ impl<'a> Batch<'a> {
 	}
 }

-/// An iterator that produces Readable instances back. Wraps the lower level
-/// DBIterator and deserializes the returned values.
-pub struct SerIterator<T>
+/// An iterator based on key prefix.
+/// Caller is responsible for deserialization of the data.
+pub struct PrefixIterator<F, T>
 where
-	T: ser::Readable,
+	F: Fn(&[u8], &[u8]) -> Result<T, Error>,
 {
 	tx: Arc<lmdb::ReadTransaction<'static>>,
 	cursor: Arc<lmdb::Cursor<'static, 'static>>,
 	seek: bool,
 	prefix: Vec<u8>,
-	version: ProtocolVersion,
-	_marker: marker::PhantomData<T>,
+	deserialize: F,
 }

-impl<T> Iterator for SerIterator<T>
+impl<F, T> Iterator for PrefixIterator<F, T>
 where
-	T: ser::Readable,
+	F: Fn(&[u8], &[u8]) -> Result<T, Error>,
 {
-	type Item = (Vec<u8>, T);
+	type Item = T;

-	fn next(&mut self) -> Option<(Vec<u8>, T)> {
+	fn next(&mut self) -> Option<Self::Item> {
 		let access = self.tx.access();
-		let kv = if self.seek {
-			Arc::get_mut(&mut self.cursor).unwrap().next(&access)
+		let cursor = Arc::get_mut(&mut self.cursor).expect("failed to get cursor");
+		let kv: Result<(&[u8], &[u8]), _> = if self.seek {
+			cursor.next(&access)
 		} else {
 			self.seek = true;
-			Arc::get_mut(&mut self.cursor)
-				.unwrap()
-				.seek_range_k(&access, &self.prefix[..])
+			cursor.seek_range_k(&access, &self.prefix[..])
 		};
-		match kv {
-			Ok((k, v)) => self.deser_if_prefix_match(k, v),
-			Err(_) => None,
-		}
+		kv.ok()
+			.filter(|(k, _)| k.starts_with(self.prefix.as_slice()))
+			.map(|(k, v)| match (self.deserialize)(k, v) {
+				Ok(v) => Some(v),
+				Err(_) => None,
+			})
+			.flatten()
 	}
 }

-impl<T> SerIterator<T>
+impl<F, T> PrefixIterator<F, T>
 where
-	T: ser::Readable,
+	F: Fn(&[u8], &[u8]) -> Result<T, Error>,
 {
-	fn deser_if_prefix_match(&self, key: &[u8], value: &[u8]) -> Option<(Vec<u8>, T)> {
-		let plen = self.prefix.len();
-		if plen == 0 || (key.len() >= plen && key[0..plen] == self.prefix[..]) {
-			if let Ok(value) = ser::deserialize(&mut &value[..], self.version) {
-				Some((key.to_vec(), value))
-			} else {
-				None
-			}
-		} else {
-			None
+	/// Initialize a new prefix iterator.
+	pub fn new(
+		tx: Arc<lmdb::ReadTransaction<'static>>,
+		cursor: Arc<lmdb::Cursor<'static, 'static>>,
+		prefix: &[u8],
+		deserialize: F,
+	) -> PrefixIterator<F, T> {
+		PrefixIterator {
+			tx,
+			cursor,
+			seek: false,
+			prefix: prefix.to_vec(),
+			deserialize,
 		}
 	}
 }