use crate::types::*;
use slab::Slab;
use std::fmt;
struct TableAndEpoch<T> {
table: AtomicPtr<T>,
epoch: AtomicUsize,
}
type ReadersList<T> = Arc<Mutex<Slab<Arc<TableAndEpoch<T>>>>>;
struct Reader<T> {
sync_state: Arc<TableAndEpoch<T>>,
key_in_readers: usize,
readers: ReadersList<T>,
}
pub struct AsLockReadGuard<'r, T> {
active_table: &'r T,
epoch: &'r AtomicUsize,
}
struct Writer<T> {
active_table: Box<T>,
standby_table: Box<T>,
ops_to_replay: Vec<Box<dyn FnOnce(&mut T) + Send>>,
readers: ReadersList<T>,
blocking_readers: std::collections::HashMap<usize, usize>,
}
pub struct AsLockHandle<T> {
writer: Arc<Mutex<Writer<T>>>,
reader: Reader<T>,
_not_sync: std::cell::UnsafeCell<fn(&T)>,
}
pub struct AsLockWriteGuard<'w, T> {
writer: MutexGuard<'w, Writer<T>>,
}
impl<T> Reader<T> {
pub fn read(&self) -> AsLockReadGuard<'_, T> {
let TableAndEpoch { table, epoch } = &*self.sync_state;
let old_epoch = epoch.load(Ordering::Acquire);
assert_eq!(old_epoch % 2, 0, "Reader is not reentrant");
epoch.store(old_epoch + 1, Ordering::Release);
fence(Ordering::SeqCst);
let active_table = unsafe { &*table.load(Ordering::SeqCst) };
AsLockReadGuard {
active_table,
epoch,
}
}
}
impl<T> Drop for Reader<T> {
fn drop(&mut self) {
self.readers.lock().remove(self.key_in_readers);
}
}
impl<T: fmt::Debug> fmt::Debug for Reader<T> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("Reader")
.field("num_readers", &self.readers.lock().len())
.field("active_table", &*self.read())
.finish()
}
}
impl<'r, T> Drop for AsLockReadGuard<'r, T> {
fn drop(&mut self) {
let old_epoch = self.epoch.load(Ordering::Acquire);
debug_assert_eq!(old_epoch % 2, 1);
self.epoch.store(old_epoch + 1, Ordering::Release);
}
}
impl<'r, T> std::ops::Deref for AsLockReadGuard<'r, T> {
type Target = T;
fn deref(&self) -> &Self::Target {
self.active_table
}
}
impl<'r, T: fmt::Debug> fmt::Debug for AsLockReadGuard<'r, T> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
self.active_table.fmt(f)
}
}
impl<T> Writer<T> {
pub fn from_identical(t1: T, t2: T) -> Writer<T> {
Writer {
active_table: Box::new(t1),
standby_table: Box::new(t2),
ops_to_replay: vec![],
readers: Arc::new(Mutex::new(Slab::with_capacity(1024))),
blocking_readers: std::collections::HashMap::new(),
}
}
pub fn new_reader(&mut self) -> Reader<T> {
let readers = Arc::clone(&self.readers);
let sync_state = Arc::new(TableAndEpoch {
table: AtomicPtr::new(self.active_table.as_mut() as *mut T),
epoch: AtomicUsize::new(0),
});
let key_in_readers = readers.lock().insert(Arc::clone(&sync_state));
Reader {
sync_state,
key_in_readers,
readers,
}
}
fn await_standby_table_free(&mut self) {
while !self.blocking_readers.is_empty() {
let readers = self.readers.lock();
self.blocking_readers.retain(|key, first_epoch_after_swap| {
let epoch = match readers.get(*key) {
None => {
return false;
}
Some(table_and_epoch) => table_and_epoch.epoch.load(Ordering::Acquire),
};
epoch <= *first_epoch_after_swap && *first_epoch_after_swap % 2 != 0
});
if !self.blocking_readers.is_empty() {
spin_loop();
}
}
}
}
impl<T> std::fmt::Debug for Writer<T>
where
T: std::fmt::Debug,
{
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("Writer")
.field("num_readers", &self.readers.lock().len())
.field("ops_to_replay", &self.ops_to_replay.len())
.field("standby_table", &self.standby_table)
.finish()
}
}
impl<T> AsLockHandle<T> {
pub fn from_identical(t1: T, t2: T) -> AsLockHandle<T> {
let mut writer = Writer::from_identical(t1, t2);
let reader = writer.new_reader();
AsLockHandle {
writer: Arc::new(Mutex::new(writer)),
reader,
_not_sync: std::cell::UnsafeCell::new(|_| {}),
}
}
pub fn read(&self) -> AsLockReadGuard<'_, T> {
self.reader.read()
}
pub fn write(&self) -> AsLockWriteGuard<'_, T> {
let mut mg = self.writer.lock();
let writer: &mut Writer<_> = &mut mg;
writer.await_standby_table_free();
std::sync::atomic::compiler_fence(Ordering::SeqCst);
for op in writer.ops_to_replay.drain(..) {
op(&mut writer.standby_table);
}
writer.ops_to_replay.clear();
AsLockWriteGuard { writer: mg }
}
}
impl<T> AsLockHandle<T>
where
T: Clone,
{
pub fn new(t: T) -> AsLockHandle<T> {
Self::from_identical(t.clone(), t)
}
}
impl<T> Default for AsLockHandle<T>
where
T: Default,
{
fn default() -> AsLockHandle<T> {
Self::from_identical(T::default(), T::default())
}
}
impl<T> Clone for AsLockHandle<T> {
fn clone(&self) -> AsLockHandle<T> {
let writer = Arc::clone(&self.writer);
let reader = writer.lock().new_reader();
AsLockHandle {
writer,
reader,
_not_sync: std::cell::UnsafeCell::new(|_| {}),
}
}
}
impl<T> std::fmt::Debug for AsLockHandle<T>
where
T: std::fmt::Debug,
{
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
let writer = self.writer.lock();
let num_readers = writer.readers.lock().len();
let num_ops_to_replay = writer.ops_to_replay.len();
f.debug_struct("AsLockHandle")
.field("num_readers", &num_readers)
.field("num_ops_to_replay", &num_ops_to_replay)
.field("standby_table", &writer.standby_table)
.field("active_table", &*self.read())
.finish()
}
}
impl<'w, T> AsLockWriteGuard<'w, T> {
pub fn update_tables<'a, R>(
&'a mut self,
mut update: impl UpdateTables<'a, T, R> + 'static + Sized + Send,
) -> R {
let writer: &mut Writer<_> = &mut self.writer;
let res = update.apply_first(&mut writer.standby_table);
writer.ops_to_replay.push(Box::new(move |table| {
update.apply_second(table);
}));
res
}
pub fn update_tables_closure<R>(
&mut self,
update: impl Fn(&mut T) -> R + 'static + Sized + Send,
) -> R {
let res = update(&mut self.writer.standby_table);
self.writer.ops_to_replay.push(Box::new(move |table| {
update(table);
}));
res
}
}
impl<'w, T> Drop for AsLockWriteGuard<'w, T> {
fn drop(&mut self) {
let writer: &mut Writer<_> = &mut self.writer;
assert!(writer.blocking_readers.is_empty());
std::mem::swap(&mut writer.active_table, &mut writer.standby_table);
for (key, table_and_epoch) in writer.readers.lock().iter_mut() {
let res = table_and_epoch.table.compare_exchange(
writer.standby_table.as_mut() as *mut T,
writer.active_table.as_mut() as *mut T,
Ordering::SeqCst,
Ordering::SeqCst,
);
assert_eq!(res, Ok(writer.standby_table.as_mut() as *mut T));
fence(Ordering::SeqCst);
let first_epoch_after_swap = table_and_epoch.epoch.load(Ordering::Acquire);
if first_epoch_after_swap % 2 != 0 {
writer.blocking_readers.insert(key, first_epoch_after_swap);
}
}
}
}
impl<'w, T> std::ops::Deref for AsLockWriteGuard<'w, T> {
type Target = T;
fn deref(&self) -> &Self::Target {
&self.writer.standby_table
}
}
impl<'w, T> std::fmt::Debug for AsLockWriteGuard<'w, T>
where
T: std::fmt::Debug,
{
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("AsLockWriteGuard")
.field("num_readers", &self.writer.readers.lock().len())
.field("ops_to_replay", &self.writer.ops_to_replay.len())
.field("standby_table", &self.writer.standby_table)
.finish()
}
}
#[cfg(test)]
mod test {
use super::*;
use std::thread;
struct PushVec<T> {
value: T,
}
impl<'a, T> UpdateTables<'a, Vec<T>, ()> for PushVec<T>
where
T: Clone + std::fmt::Debug,
{
fn apply_first(&mut self, table: &'a mut Vec<T>) {
dbg!(&table);
table.push(self.value.clone());
dbg!(&table);
}
fn apply_second(self, table: &mut Vec<T>) {
table.push(self.value); }
}
struct PopVec {}
impl PopVec {
fn apply<T>(&mut self, table: &mut Vec<T>) -> Option<T> {
table.pop()
}
}
impl<'a, T> UpdateTables<'a, Vec<T>, Option<T>> for PopVec {
fn apply_first(&mut self, table: &'a mut Vec<T>) -> Option<T> {
self.apply(table)
}
fn apply_second(mut self, table: &mut Vec<T>) {
(&mut self).apply(table);
}
}
struct MutableRef {}
impl<'a, T> UpdateTables<'a, Vec<T>, &'a mut T> for MutableRef {
fn apply_first(&mut self, table: &'a mut Vec<T>) -> &'a mut T {
&mut table[0]
}
fn apply_second(self, table: &mut Vec<T>) {
let _ = &mut table[0];
}
}
#[test]
#[should_panic(expected = "Reader is not reentrant")]
fn reader_not_reentrant() {
let table = AsLockHandle::<Vec<i32>>::default();
let _rg1 = table.read();
let _rg2 = table.read();
}
#[test]
fn writer_not_reentrant() {
let table = AsLockHandle::<Vec<i32>>::from_identical(vec![], vec![]);
let _wg = table.write();
}
#[test]
fn publish_update() {
let table = AsLockHandle::<Vec<i32>>::new(vec![]);
assert_eq!(table.read().len(), 0);
{
let mut wg = table.write();
wg.update_tables(PushVec { value: 2 });
assert_eq!(wg.len(), 1);
assert_eq!(table.read().len(), 0);
}
assert_eq!(*table.read(), vec![2]);
}
#[test]
fn update_tables_closure() {
let table = AsLockHandle::<Vec<i32>>::default();
assert_eq!(table.read().len(), 0);
{
let mut wg = table.write();
wg.update_tables_closure(|vec| vec.push(2));
assert_eq!(wg.len(), 1);
assert_eq!(table.read().len(), 0);
}
assert_eq!(*table.read(), vec![2]);
}
#[test]
fn multi_apply() {
let table = AsLockHandle::<Vec<i32>>::default();
{
let mut wg = table.write();
wg.update_tables(PushVec { value: 2 });
wg.update_tables(PushVec { value: 3 });
wg.update_tables(PushVec { value: 4 });
wg.update_tables(PopVec {});
wg.update_tables(PushVec { value: 5 });
}
assert_eq!(*table.read(), vec![2, 3, 5]);
}
#[test]
fn multi_publish() {
let table = AsLockHandle::<Vec<Box<i32>>>::default();
{
let mut wg = table.write();
wg.update_tables(PushVec { value: Box::new(2) });
wg.update_tables(PushVec { value: Box::new(3) });
wg.update_tables(PopVec {});
wg.update_tables(PushVec { value: Box::new(5) });
}
assert_eq!(*table.read(), vec![Box::new(2), Box::new(5)]);
{
let mut wg = table.write();
wg.update_tables(PushVec { value: Box::new(9) });
wg.update_tables(PushVec { value: Box::new(8) });
wg.update_tables(PopVec {});
wg.update_tables(PushVec { value: Box::new(7) });
}
assert_eq!(
*table.read(),
vec![Box::new(2), Box::new(5), Box::new(9), Box::new(7)]
);
table.write().update_tables(PopVec {});
assert_eq!(*table.read(), vec![Box::new(2), Box::new(5), Box::new(9)]);
}
#[test]
fn multi_thread() {
let table = AsLockHandle::<Vec<i32>>::default();
let handler = {
let table = table.clone();
thread::spawn(move || {
while *table.read() != vec![2, 3, 5] {
assert_ne!(*table.read(), vec![2, 3, 4]);
}
let handler = {
let table = table;
thread::spawn(move || while *table.read() != vec![2, 3, 5] {})
};
assert!(handler.join().is_ok());
})
};
{
let mut wg = table.write();
wg.update_tables(PushVec { value: 2 });
wg.update_tables(PushVec { value: 3 });
wg.update_tables(PushVec { value: 4 });
wg.update_tables(PopVec {});
wg.update_tables(PushVec { value: 5 });
}
assert!(handler.join().is_ok());
}
#[test]
fn writer_dropped() {
let table;
{
table = AsLockHandle::<Vec<i32>>::default();
{
let mut wg = table.write();
wg.update_tables(PushVec { value: 2 });
wg.update_tables(PushVec { value: 3 });
wg.update_tables(PushVec { value: 4 });
wg.update_tables(PopVec {});
wg.update_tables(PushVec { value: 5 });
}
}
assert_eq!(*table.read(), vec![2, 3, 5]);
}
#[test]
fn mutable_ref() {
let table = AsLockHandle::<Vec<i32>>::default();
{
let mut wg = table.write();
wg.update_tables(PushVec { value: 2 });
let mr = wg.update_tables(MutableRef {});
*mr = 10;
}
assert_eq!(*table.read(), vec![10]);
assert_ne!(*table.read(), *table.write());
}
#[test]
fn debug_str() {
let table = AsLockHandle::<Vec<i32>>::default();
assert_eq!(
format!("{:?}", table),
"AsLockHandle { num_readers: 1, num_ops_to_replay: 0, standby_table: [], active_table: [] }"
);
{
let mut wg = table.write();
wg.update_tables(PushVec { value: 2 });
assert_eq!(
format!("{:?}", wg),
"AsLockWriteGuard { num_readers: 1, ops_to_replay: 1, standby_table: [2] }"
);
}
assert_eq!(
format!("{:?}", table),
"AsLockHandle { num_readers: 1, num_ops_to_replay: 1, standby_table: [], active_table: [2] }"
);
assert_eq!(format!("{:?}", table.read()), "[2]");
}
}