aere-research/parallel-executor/src/mvmemory.rs

//! Multi-version shared memory (MVMemory).
//!
//! For every `Key`, MVMemory keeps an ordered map `txn_index -> cell`, where a
//! cell is either a concrete `Value` written by that transaction's latest
//! incarnation, or an `Estimate` placeholder (a marker that a lower transaction
//! *previously* wrote here but is currently being re-executed, so its value is
//! not trustworthy).
//!
//! A read by transaction `i` for key `k` returns the entry of the HIGHEST
//! transaction `j < i` that wrote `k`. That is what makes speculative parallel
//! execution serial-equivalent: transaction `i` always observes the writes of
//! all lower-indexed transactions that have run, and nothing from higher ones.
//!
//! Storage is sharded across many independent mutexes keyed by a hash of the
//! `Key`, so non-conflicting transactions touch disjoint locks and run without
//! contention. Hot keys (the counter, the AMM pool) serialize on one shard,
//! which is exactly the honest source of the high-conflict slowdown.

use std::collections::BTreeMap;
use std::collections::HashMap;
use std::sync::Mutex;

use crate::types::{Incarnation, Key, TxnIndex, Value};

const SHARDS: usize = 256;

/// One stored cell for a (key, txn_index) pair.
#[derive(Clone, Debug)]
enum Cell {
    /// A concrete value written by this txn's incarnation.
    Value(Incarnation, Value),
    /// Placeholder: this txn wrote here before but is being re-executed.
    Estimate,
}

/// What a read observed. This is recorded into the reader's read set so that
/// validation can later re-check that the read would still return the same
/// thing.
#[derive(Clone, PartialEq, Eq, Debug)]
pub enum ReadDescriptor {
    /// Nothing below the reader wrote this key: it came from base state.
    Base,
    /// It was written by `(txn_index, incarnation)`.
    Version(TxnIndex, Incarnation),
}

/// The result of a live read against MVMemory.
pub enum ReadResult {
    /// Read resolved to a concrete value with the given provenance.
    Ok(ReadDescriptor, Value),
    /// Read hit an ESTIMATE written by `blocking_txn`; the reader must abort and
    /// register a dependency on `blocking_txn`.
    Blocked(TxnIndex),
}

pub struct MvMemory {
    shards: Vec<Mutex<HashMap<Key, BTreeMap<TxnIndex, Cell>>>>,
}

fn shard_of(key: &Key) -> usize {
    // Cheap, deterministic spread. FNV-1a over the key's discriminant + fields.
    let mut h: u64 = 0xcbf29ce484222325;
    let bytes: [u64; 3] = match key {
        Key::Balance(a) => [1, *a, 0],
        Key::Storage(c, s) => [2, *c, *s],
    };
    for b in bytes {
        h ^= b;
        h = h.wrapping_mul(0x100000001b3);
    }
    (h as usize) % SHARDS
}

impl MvMemory {
    pub fn new() -> Self {
        let mut shards = Vec::with_capacity(SHARDS);
        for _ in 0..SHARDS {
            shards.push(Mutex::new(HashMap::new()));
        }
        MvMemory { shards }
    }

    /// Read the value visible to `txn_idx`: the highest entry strictly below it.
    pub fn read(&self, key: &Key, txn_idx: TxnIndex) -> ReadResult {
        let shard = &self.shards[shard_of(key)];
        let guard = shard.lock().unwrap();
        if let Some(versions) = guard.get(key) {
            // highest entry with index < txn_idx
            if let Some((&j, cell)) = versions.range(..txn_idx).next_back() {
                return match cell {
                    Cell::Value(inc, v) => ReadResult::Ok(ReadDescriptor::Version(j, *inc), *v),
                    Cell::Estimate => ReadResult::Blocked(j),
                };
            }
        }
        ReadResult::Ok(ReadDescriptor::Base, 0)
    }

    /// Apply `txn_idx`'s new write set (incarnation `inc`). `prev_keys` is the
    /// set of keys the previous incarnation wrote, so we can (a) delete entries
    /// this incarnation no longer writes and (b) decide whether the *set* of
    /// written keys changed, which downstream validation depends on.
    ///
    /// Returns `true` if the written key-set differs from the previous one
    /// (a key was added or removed) -- Block-STM calls this "wrote a new path"
    /// and uses it to trigger revalidation of the suffix.
    pub fn apply_write_set(
        &self,
        txn_idx: TxnIndex,
        inc: Incarnation,
        writes: &HashMap<Key, Value>,
        prev_keys: &std::collections::HashSet<Key>,
    ) -> bool {
        // Install / overwrite current writes.
        for (k, v) in writes {
            let shard = &self.shards[shard_of(k)];
            let mut guard = shard.lock().unwrap();
            guard
                .entry(*k)
                .or_insert_with(BTreeMap::new)
                .insert(txn_idx, Cell::Value(inc, *v));
        }
        // Remove stale entries: keys written previously but not now.
        let mut removed_any = false;
        for k in prev_keys {
            if !writes.contains_key(k) {
                let shard = &self.shards[shard_of(k)];
                let mut guard = shard.lock().unwrap();
                if let Some(versions) = guard.get_mut(k) {
                    versions.remove(&txn_idx);
                }
                removed_any = true;
            }
        }
        // Added a key not previously written?
        let mut added_any = false;
        for k in writes.keys() {
            if !prev_keys.contains(k) {
                added_any = true;
                break;
            }
        }
        added_any || removed_any
    }

    /// Convert all of `txn_idx`'s current writes (over `keys`) into ESTIMATE
    /// placeholders. Called the moment a transaction is validation-aborted, so
    /// that any concurrent higher transaction that reads one of these keys will
    /// block on `txn_idx` instead of consuming a soon-to-be-wrong value. This is
    /// what protects the suffix during a re-execution window.
    pub fn mark_estimates(&self, txn_idx: TxnIndex, keys: &std::collections::HashSet<Key>) {
        for k in keys {
            let shard = &self.shards[shard_of(k)];
            let mut guard = shard.lock().unwrap();
            if let Some(versions) = guard.get_mut(k) {
                if let Some(cell) = versions.get_mut(&txn_idx) {
                    *cell = Cell::Estimate;
                }
            }
        }
    }

    /// Materialize the final committed state: base overlaid with, for each key,
    /// the value written by the HIGHEST transaction index that wrote a concrete
    /// value to it. Called once after the block is fully executed & validated.
    pub fn snapshot_final(&self, base: &crate::types::StateMap) -> crate::types::StateMap {
        let mut out = base.clone();
        for shard in &self.shards {
            let guard = shard.lock().unwrap();
            for (k, versions) in guard.iter() {
                // highest concrete value wins
                for (_, cell) in versions.iter().rev() {
                    if let Cell::Value(_, v) = cell {
                        out.insert(*k, *v);
                        break;
                    }
                }
            }
        }
        out
    }
}