use super::{StorageKey, StorageValue};
use itertools::Itertools;
use std::collections::{HashSet, BTreeMap, BTreeSet};
use smallvec::SmallVec;
use log::warn;
const PROOF_OVERLAY_NON_EMPTY: &str = "\
An OverlayValue is always created with at least one transaction and dropped as soon
as the last transaction is removed; qed";
type DirtyKeysSets = SmallVec<[HashSet<StorageKey>; 5]>;
type Transactions = SmallVec<[InnerValue; 5]>;
#[derive(Debug)]
#[cfg_attr(test, derive(PartialEq))]
pub struct NoOpenTransaction;
#[derive(Debug)]
#[cfg_attr(test, derive(PartialEq))]
pub struct AlreadyInRuntime;
#[derive(Debug)]
#[cfg_attr(test, derive(PartialEq))]
pub struct NotInRuntime;
#[derive(Debug, Clone, Copy)]
pub enum ExecutionMode {
Client,
Runtime,
}
#[derive(Debug, Default, Clone)]
#[cfg_attr(test, derive(PartialEq))]
struct InnerValue {
value: Option<StorageValue>,
extrinsics: BTreeSet<u32>,
}
#[derive(Debug, Default, Clone)]
#[cfg_attr(test, derive(PartialEq))]
pub struct OverlayedValue {
transactions: Transactions,
}
#[derive(Debug, Default, Clone)]
pub struct OverlayedChangeSet {
changes: BTreeMap<StorageKey, OverlayedValue>,
dirty_keys: DirtyKeysSets,
num_client_transactions: usize,
execution_mode: ExecutionMode,
}
impl Default for ExecutionMode {
fn default() -> Self {
Self::Client
}
}
impl OverlayedValue {
pub fn value(&self) -> Option<&StorageValue> {
self.transactions.last().expect(PROOF_OVERLAY_NON_EMPTY).value.as_ref()
}
pub fn extrinsics(&self) -> impl Iterator<Item=&u32> {
self.transactions.iter().flat_map(|t| t.extrinsics.iter()).unique()
}
fn value_mut(&mut self) -> &mut Option<StorageValue> {
&mut self.transactions.last_mut().expect(PROOF_OVERLAY_NON_EMPTY).value
}
fn pop_transaction(&mut self) -> InnerValue {
self.transactions.pop().expect(PROOF_OVERLAY_NON_EMPTY)
}
fn transaction_extrinsics_mut(&mut self) -> &mut BTreeSet<u32> {
&mut self.transactions.last_mut().expect(PROOF_OVERLAY_NON_EMPTY).extrinsics
}
fn set(
&mut self,
value: Option<StorageValue>,
first_write_in_tx: bool,
at_extrinsic: Option<u32>,
) {
if first_write_in_tx || self.transactions.is_empty() {
self.transactions.push(InnerValue {
value,
.. Default::default()
});
} else {
*self.value_mut() = value;
}
if let Some(extrinsic) = at_extrinsic {
self.transaction_extrinsics_mut().insert(extrinsic);
}
}
}
fn insert_dirty(set: &mut DirtyKeysSets, key: StorageKey) -> bool {
set.last_mut().map(|dk| dk.insert(key)).unwrap_or_default()
}
impl OverlayedChangeSet {
pub fn spawn_child(&self) -> Self {
use std::iter::repeat;
Self {
dirty_keys: repeat(HashSet::new()).take(self.transaction_depth()).collect(),
num_client_transactions: self.num_client_transactions,
execution_mode: self.execution_mode,
.. Default::default()
}
}
pub fn is_empty(&self) -> bool {
self.changes.is_empty()
}
pub fn get(&self, key: &[u8]) -> Option<&OverlayedValue> {
self.changes.get(key)
}
pub fn set(
&mut self,
key: StorageKey,
value: Option<StorageValue>,
at_extrinsic: Option<u32>,
) {
let overlayed = self.changes.entry(key.clone()).or_default();
overlayed.set(value, insert_dirty(&mut self.dirty_keys, key), at_extrinsic);
}
#[must_use = "A change was registered, so this value MUST be modified."]
pub fn modify(
&mut self,
key: StorageKey,
init: impl Fn() -> StorageValue,
at_extrinsic: Option<u32>,
) -> &mut Option<StorageValue> {
let overlayed = self.changes.entry(key.clone()).or_default();
let first_write_in_tx = insert_dirty(&mut self.dirty_keys, key);
let clone_into_new_tx = if let Some(tx) = overlayed.transactions.last() {
if first_write_in_tx {
Some(tx.value.clone())
} else {
None
}
} else {
Some(Some(init()))
};
if let Some(cloned) = clone_into_new_tx {
overlayed.set(cloned, first_write_in_tx, at_extrinsic);
}
overlayed.value_mut()
}
pub fn clear_where(
&mut self,
predicate: impl Fn(&[u8], &OverlayedValue) -> bool,
at_extrinsic: Option<u32>,
) {
for (key, val) in self.changes.iter_mut().filter(|(k, v)| predicate(k, v)) {
val.set(None, insert_dirty(&mut self.dirty_keys, key.to_owned()), at_extrinsic);
}
}
pub fn changes(&self) -> impl Iterator<Item=(&StorageKey, &OverlayedValue)> {
self.changes.iter()
}
pub fn next_change(&self, key: &[u8]) -> Option<(&[u8], &OverlayedValue)> {
use std::ops::Bound;
let range = (Bound::Excluded(key), Bound::Unbounded);
self.changes.range::<[u8], _>(range).next().map(|(k, v)| (&k[..], v))
}
pub fn drain_commited(self) -> impl Iterator<Item=(StorageKey, Option<StorageValue>)> {
assert!(self.transaction_depth() == 0, "Drain is not allowed with open transactions.");
self.changes.into_iter().map(|(k, mut v)| (k, v.pop_transaction().value))
}
pub fn transaction_depth(&self) -> usize {
self.dirty_keys.len()
}
pub fn enter_runtime(&mut self) -> Result<(), AlreadyInRuntime> {
if let ExecutionMode::Runtime = self.execution_mode {
return Err(AlreadyInRuntime);
}
self.execution_mode = ExecutionMode::Runtime;
self.num_client_transactions = self.transaction_depth();
Ok(())
}
pub fn exit_runtime(&mut self) -> Result<(), NotInRuntime> {
if let ExecutionMode::Client = self.execution_mode {
return Err(NotInRuntime);
}
self.execution_mode = ExecutionMode::Client;
if self.has_open_runtime_transactions() {
warn!(
"{} storage transactions are left open by the runtime. Those will be rolled back.",
self.transaction_depth() - self.num_client_transactions,
);
}
while self.has_open_runtime_transactions() {
self.rollback_transaction()
.expect("The loop condition checks that the transaction depth is > 0; qed");
}
Ok(())
}
pub fn start_transaction(&mut self) {
self.dirty_keys.push(Default::default());
}
pub fn rollback_transaction(&mut self) -> Result<(), NoOpenTransaction> {
self.close_transaction(true)
}
pub fn commit_transaction(&mut self) -> Result<(), NoOpenTransaction> {
self.close_transaction(false)
}
fn close_transaction(&mut self, rollback: bool) -> Result<(), NoOpenTransaction> {
if let ExecutionMode::Runtime = self.execution_mode {
if !self.has_open_runtime_transactions() {
return Err(NoOpenTransaction)
}
}
for key in self.dirty_keys.pop().ok_or(NoOpenTransaction)? {
let overlayed = self.changes.get_mut(&key).expect("\
A write to an OverlayedValue is recorded in the dirty key set. Before an
OverlayedValue is removed, its containing dirty set is removed. This
function is only called for keys that are in the dirty set. qed\
");
if rollback {
overlayed.pop_transaction();
if overlayed.transactions.is_empty() {
self.changes.remove(&key);
}
} else {
let has_predecessor = if let Some(dirty_keys) = self.dirty_keys.last_mut() {
!dirty_keys.insert(key)
} else {
overlayed.transactions.len() > 1
};
if has_predecessor {
let dropped_tx = overlayed.pop_transaction();
*overlayed.value_mut() = dropped_tx.value;
overlayed.transaction_extrinsics_mut().extend(dropped_tx.extrinsics);
}
}
}
Ok(())
}
fn has_open_runtime_transactions(&self) -> bool {
self.transaction_depth() > self.num_client_transactions
}
}
#[cfg(test)]
mod test {
use super::*;
use pretty_assertions::assert_eq;
type Changes<'a> = Vec<(&'a [u8], (Option<&'a [u8]>, Vec<u32>))>;
type Drained<'a> = Vec<(&'a [u8], Option<&'a [u8]>)>;
fn assert_changes(is: &OverlayedChangeSet, expected: &Changes) {
let is: Changes = is.changes().map(|(k, v)| {
(k.as_ref(), (v.value().map(AsRef::as_ref), v.extrinsics().cloned().collect()))
}).collect();
assert_eq!(&is, expected);
}
fn assert_drained_changes(is: OverlayedChangeSet, expected: Changes) {
let is = is.drain_commited().collect::<Vec<_>>();
let expected = expected
.iter()
.map(|(k, v)| (k.to_vec(), v.0.map(From::from))).collect::<Vec<_>>();
assert_eq!(is, expected);
}
fn assert_drained(is: OverlayedChangeSet, expected: Drained) {
let is = is.drain_commited().collect::<Vec<_>>();
let expected = expected
.iter()
.map(|(k, v)| (k.to_vec(), v.map(From::from))).collect::<Vec<_>>();
assert_eq!(is, expected);
}
#[test]
fn no_transaction_works() {
let mut changeset = OverlayedChangeSet::default();
assert_eq!(changeset.transaction_depth(), 0);
changeset.set(b"key0".to_vec(), Some(b"val0".to_vec()), Some(1));
changeset.set(b"key1".to_vec(), Some(b"val1".to_vec()), Some(2));
changeset.set(b"key0".to_vec(), Some(b"val0-1".to_vec()), Some(9));
assert_drained(changeset, vec![
(b"key0", Some(b"val0-1")),
(b"key1", Some(b"val1")),
]);
}
#[test]
fn transaction_works() {
let mut changeset = OverlayedChangeSet::default();
assert_eq!(changeset.transaction_depth(), 0);
changeset.set(b"key0".to_vec(), Some(b"val0".to_vec()), Some(1));
changeset.set(b"key1".to_vec(), Some(b"val1".to_vec()), Some(1));
changeset.set(b"key0".to_vec(), Some(b"val0-1".to_vec()), Some(10));
changeset.start_transaction();
assert_eq!(changeset.transaction_depth(), 1);
changeset.set(b"key42".to_vec(), Some(b"val42".to_vec()), Some(42));
changeset.set(b"key99".to_vec(), Some(b"val99".to_vec()), Some(99));
changeset.start_transaction();
assert_eq!(changeset.transaction_depth(), 2);
changeset.set(b"key42".to_vec(), Some(b"val42-rolled".to_vec()), Some(421));
changeset.set(b"key7".to_vec(), Some(b"val7-rolled".to_vec()), Some(77));
changeset.set(b"key0".to_vec(), Some(b"val0-rolled".to_vec()), Some(1000));
changeset.set(b"key5".to_vec(), Some(b"val5-rolled".to_vec()), None);
let all_changes: Changes = vec![
(b"key0", (Some(b"val0-rolled"), vec![1, 10, 1000])),
(b"key1", (Some(b"val1"), vec![1])),
(b"key42", (Some(b"val42-rolled"), vec![42, 421])),
(b"key5", (Some(b"val5-rolled"), vec![])),
(b"key7", (Some(b"val7-rolled"), vec![77])),
(b"key99", (Some(b"val99"), vec![99])),
];
assert_changes(&changeset, &all_changes);
changeset.start_transaction();
assert_eq!(changeset.transaction_depth(), 3);
changeset.start_transaction();
assert_eq!(changeset.transaction_depth(), 4);
changeset.rollback_transaction().unwrap();
assert_eq!(changeset.transaction_depth(), 3);
changeset.commit_transaction().unwrap();
assert_eq!(changeset.transaction_depth(), 2);
assert_changes(&changeset, &all_changes);
changeset.rollback_transaction().unwrap();
assert_eq!(changeset.transaction_depth(), 1);
let rolled_back: Changes = vec![
(b"key0", (Some(b"val0-1"), vec![1, 10])),
(b"key1", (Some(b"val1"), vec![1])),
(b"key42", (Some(b"val42"), vec![42])),
(b"key99", (Some(b"val99"), vec![99])),
];
assert_changes(&changeset, &rolled_back);
changeset.commit_transaction().unwrap();
assert_eq!(changeset.transaction_depth(), 0);
assert_changes(&changeset, &rolled_back);
assert_drained_changes(changeset, rolled_back);
}
#[test]
fn transaction_commit_then_rollback_works() {
let mut changeset = OverlayedChangeSet::default();
assert_eq!(changeset.transaction_depth(), 0);
changeset.set(b"key0".to_vec(), Some(b"val0".to_vec()), Some(1));
changeset.set(b"key1".to_vec(), Some(b"val1".to_vec()), Some(1));
changeset.set(b"key0".to_vec(), Some(b"val0-1".to_vec()), Some(10));
changeset.start_transaction();
assert_eq!(changeset.transaction_depth(), 1);
changeset.set(b"key42".to_vec(), Some(b"val42".to_vec()), Some(42));
changeset.set(b"key99".to_vec(), Some(b"val99".to_vec()), Some(99));
changeset.start_transaction();
assert_eq!(changeset.transaction_depth(), 2);
changeset.set(b"key42".to_vec(), Some(b"val42-rolled".to_vec()), Some(421));
changeset.set(b"key7".to_vec(), Some(b"val7-rolled".to_vec()), Some(77));
changeset.set(b"key0".to_vec(), Some(b"val0-rolled".to_vec()), Some(1000));
changeset.set(b"key5".to_vec(), Some(b"val5-rolled".to_vec()), None);
let all_changes: Changes = vec![
(b"key0", (Some(b"val0-rolled"), vec![1, 10, 1000])),
(b"key1", (Some(b"val1"), vec![1])),
(b"key42", (Some(b"val42-rolled"), vec![42, 421])),
(b"key5", (Some(b"val5-rolled"), vec![])),
(b"key7", (Some(b"val7-rolled"), vec![77])),
(b"key99", (Some(b"val99"), vec![99])),
];
assert_changes(&changeset, &all_changes);
changeset.start_transaction();
assert_eq!(changeset.transaction_depth(), 3);
changeset.start_transaction();
assert_eq!(changeset.transaction_depth(), 4);
changeset.rollback_transaction().unwrap();
assert_eq!(changeset.transaction_depth(), 3);
changeset.commit_transaction().unwrap();
assert_eq!(changeset.transaction_depth(), 2);
assert_changes(&changeset, &all_changes);
changeset.commit_transaction().unwrap();
assert_eq!(changeset.transaction_depth(), 1);
assert_changes(&changeset, &all_changes);
changeset.rollback_transaction().unwrap();
assert_eq!(changeset.transaction_depth(), 0);
let rolled_back: Changes = vec![
(b"key0", (Some(b"val0-1"), vec![1, 10])),
(b"key1", (Some(b"val1"), vec![1])),
];
assert_changes(&changeset, &rolled_back);
assert_drained_changes(changeset, rolled_back);
}
#[test]
fn modify_works() {
let mut changeset = OverlayedChangeSet::default();
assert_eq!(changeset.transaction_depth(), 0);
let init = || b"valinit".to_vec();
changeset.set(b"key0".to_vec(), Some(b"val0".to_vec()), Some(0));
changeset.set(b"key1".to_vec(), None, Some(1));
let val = changeset.modify(b"key3".to_vec(), init, Some(3));
assert_eq!(val, &Some(b"valinit".to_vec()));
val.as_mut().unwrap().extend_from_slice(b"-modified");
changeset.start_transaction();
assert_eq!(changeset.transaction_depth(), 1);
changeset.start_transaction();
assert_eq!(changeset.transaction_depth(), 2);
let val = changeset.modify(b"key2".to_vec(), init, Some(2));
assert_eq!(val, &Some(b"valinit".to_vec()));
val.as_mut().unwrap().extend_from_slice(b"-modified");
let val = changeset.modify(b"key0".to_vec(), init, Some(10));
assert_eq!(val, &Some(b"val0".to_vec()));
val.as_mut().unwrap().extend_from_slice(b"-modified");
let val = changeset.modify(b"key1".to_vec(), init, Some(20));
assert_eq!(val, &None);
*val = Some(b"deleted-modified".to_vec());
let all_changes: Changes = vec![
(b"key0", (Some(b"val0-modified"), vec![0, 10])),
(b"key1", (Some(b"deleted-modified"), vec![1, 20])),
(b"key2", (Some(b"valinit-modified"), vec![2])),
(b"key3", (Some(b"valinit-modified"), vec![3])),
];
assert_changes(&changeset, &all_changes);
changeset.commit_transaction().unwrap();
assert_eq!(changeset.transaction_depth(), 1);
assert_changes(&changeset, &all_changes);
changeset.rollback_transaction().unwrap();
assert_eq!(changeset.transaction_depth(), 0);
let rolled_back: Changes = vec![
(b"key0", (Some(b"val0"), vec![0])),
(b"key1", (None, vec![1])),
(b"key3", (Some(b"valinit-modified"), vec![3])),
];
assert_changes(&changeset, &rolled_back);
assert_drained_changes(changeset, rolled_back);
}
#[test]
fn clear_works() {
let mut changeset = OverlayedChangeSet::default();
changeset.set(b"key0".to_vec(), Some(b"val0".to_vec()), Some(1));
changeset.set(b"key1".to_vec(), Some(b"val1".to_vec()), Some(2));
changeset.set(b"del1".to_vec(), Some(b"delval1".to_vec()), Some(3));
changeset.set(b"del2".to_vec(), Some(b"delval2".to_vec()), Some(4));
changeset.start_transaction();
changeset.clear_where(|k, _| k.starts_with(b"del"), Some(5));
assert_changes(&changeset, &vec![
(b"del1", (None, vec![3, 5])),
(b"del2", (None, vec![4, 5])),
(b"key0", (Some(b"val0"), vec![1])),
(b"key1", (Some(b"val1"), vec![2])),
]);
changeset.rollback_transaction().unwrap();
assert_changes(&changeset, &vec![
(b"del1", (Some(b"delval1"), vec![3])),
(b"del2", (Some(b"delval2"), vec![4])),
(b"key0", (Some(b"val0"), vec![1])),
(b"key1", (Some(b"val1"), vec![2])),
]);
}
#[test]
fn next_change_works() {
let mut changeset = OverlayedChangeSet::default();
changeset.set(b"key0".to_vec(), Some(b"val0".to_vec()), Some(0));
changeset.set(b"key1".to_vec(), Some(b"val1".to_vec()), Some(1));
changeset.set(b"key2".to_vec(), Some(b"val2".to_vec()), Some(2));
changeset.start_transaction();
changeset.set(b"key3".to_vec(), Some(b"val3".to_vec()), Some(3));
changeset.set(b"key4".to_vec(), Some(b"val4".to_vec()), Some(4));
changeset.set(b"key11".to_vec(), Some(b"val11".to_vec()), Some(11));
assert_eq!(changeset.next_change(b"key0").unwrap().0, b"key1");
assert_eq!(changeset.next_change(b"key0").unwrap().1.value(), Some(&b"val1".to_vec()));
assert_eq!(changeset.next_change(b"key1").unwrap().0, b"key11");
assert_eq!(changeset.next_change(b"key1").unwrap().1.value(), Some(&b"val11".to_vec()));
assert_eq!(changeset.next_change(b"key11").unwrap().0, b"key2");
assert_eq!(changeset.next_change(b"key11").unwrap().1.value(), Some(&b"val2".to_vec()));
assert_eq!(changeset.next_change(b"key2").unwrap().0, b"key3");
assert_eq!(changeset.next_change(b"key2").unwrap().1.value(), Some(&b"val3".to_vec()));
assert_eq!(changeset.next_change(b"key3").unwrap().0, b"key4");
assert_eq!(changeset.next_change(b"key3").unwrap().1.value(), Some(&b"val4".to_vec()));
assert_eq!(changeset.next_change(b"key4"), None);
changeset.rollback_transaction().unwrap();
assert_eq!(changeset.next_change(b"key0").unwrap().0, b"key1");
assert_eq!(changeset.next_change(b"key0").unwrap().1.value(), Some(&b"val1".to_vec()));
assert_eq!(changeset.next_change(b"key1").unwrap().0, b"key2");
assert_eq!(changeset.next_change(b"key1").unwrap().1.value(), Some(&b"val2".to_vec()));
assert_eq!(changeset.next_change(b"key11").unwrap().0, b"key2");
assert_eq!(changeset.next_change(b"key11").unwrap().1.value(), Some(&b"val2".to_vec()));
assert_eq!(changeset.next_change(b"key2"), None);
assert_eq!(changeset.next_change(b"key3"), None);
assert_eq!(changeset.next_change(b"key4"), None);
}
#[test]
fn no_open_tx_commit_errors() {
let mut changeset = OverlayedChangeSet::default();
assert_eq!(changeset.transaction_depth(), 0);
assert_eq!(changeset.commit_transaction(), Err(NoOpenTransaction));
}
#[test]
fn no_open_tx_rollback_errors() {
let mut changeset = OverlayedChangeSet::default();
assert_eq!(changeset.transaction_depth(), 0);
assert_eq!(changeset.rollback_transaction(), Err(NoOpenTransaction));
}
#[test]
fn unbalanced_transactions_errors() {
let mut changeset = OverlayedChangeSet::default();
changeset.start_transaction();
changeset.commit_transaction().unwrap();
assert_eq!(changeset.commit_transaction(), Err(NoOpenTransaction));
}
#[test]
#[should_panic]
fn drain_with_open_transaction_panics() {
let mut changeset = OverlayedChangeSet::default();
changeset.start_transaction();
let _ = changeset.drain_commited();
}
#[test]
fn runtime_cannot_close_client_tx() {
let mut changeset = OverlayedChangeSet::default();
changeset.start_transaction();
changeset.enter_runtime().unwrap();
changeset.start_transaction();
changeset.commit_transaction().unwrap();
assert_eq!(changeset.commit_transaction(), Err(NoOpenTransaction));
assert_eq!(changeset.rollback_transaction(), Err(NoOpenTransaction));
}
#[test]
fn exit_runtime_closes_runtime_tx() {
let mut changeset = OverlayedChangeSet::default();
changeset.start_transaction();
changeset.set(b"key0".to_vec(), Some(b"val0".to_vec()), Some(1));
changeset.enter_runtime().unwrap();
changeset.start_transaction();
changeset.set(b"key1".to_vec(), Some(b"val1".to_vec()), Some(2));
changeset.exit_runtime().unwrap();
changeset.commit_transaction().unwrap();
assert_eq!(changeset.transaction_depth(), 0);
assert_drained(changeset, vec![
(b"key0", Some(b"val0")),
]);
}
#[test]
fn enter_exit_runtime_fails_when_already_in_requested_mode() {
let mut changeset = OverlayedChangeSet::default();
assert_eq!(changeset.exit_runtime(), Err(NotInRuntime));
assert_eq!(changeset.enter_runtime(), Ok(()));
assert_eq!(changeset.enter_runtime(), Err(AlreadyInRuntime));
assert_eq!(changeset.exit_runtime(), Ok(()));
assert_eq!(changeset.exit_runtime(), Err(NotInRuntime));
}
}