use std::pin::Pin;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, Mutex};
use futures_channel::mpsc;
use futures_util::{Stream, StreamExt};
use crate::error::Error;
use crate::render::{DataPlane, Page, Query, Revision, Row, RowId, RowPatch};
use crate::table::{filter_rows, sort_rows};
#[derive(Debug, Clone, Default)]
pub struct Schema {
pub column_ids: Vec<String>,
}
#[derive(Debug, Clone)]
struct HistoryEntry {
revision: Revision,
row: Row,
}
#[derive(Debug)]
struct State {
rows: Vec<Row>,
index: std::collections::HashMap<RowId, usize>,
history: Vec<HistoryEntry>,
}
#[derive(Clone)]
pub struct MemPlane {
state: Arc<Mutex<State>>,
revision: Arc<AtomicU64>,
schema: Schema,
live: Arc<Mutex<Vec<mpsc::UnboundedSender<RowPatch>>>>,
history_cap: usize,
}
impl MemPlane {
pub fn new(schema: Schema) -> Self {
Self {
state: Arc::new(Mutex::new(State {
rows: Vec::new(),
index: std::collections::HashMap::new(),
history: Vec::new(),
})),
revision: Arc::new(AtomicU64::new(0)),
schema,
live: Arc::new(Mutex::new(Vec::new())),
history_cap: 1024,
}
}
pub fn revision(&self) -> Revision {
Revision(self.revision.load(Ordering::SeqCst))
}
pub fn upsert(&self, row: Row) -> Revision {
let id = row.id.clone();
let (revision, patch) = {
let mut state = self.state.lock().expect("mem plane poisoned");
let revision = self.revision.fetch_add(1, Ordering::SeqCst) + 1;
match state.index.get(&id) {
Some(&i) => {
state.rows[i] = row.clone();
}
None => {
let len = state.rows.len();
state.index.insert(id.clone(), len);
state.rows.push(row.clone());
}
}
state.history.push(HistoryEntry {
revision: Revision(revision),
row: row.clone(),
});
if state.history.len() > self.history_cap {
state.history.remove(0);
}
(Revision(revision), RowPatch::Upsert { id, row })
};
self.broadcast(patch);
revision
}
pub fn remove(&self, id: &RowId) -> Revision {
let (revision, patch) = {
let mut state = self.state.lock().expect("mem plane poisoned");
let Some(i) = state.index.remove(id) else {
return self.revision();
};
state.rows.swap_remove(i);
if let Some(moved) = state.rows.get(i) {
let moved_id = moved.id.clone();
state.index.insert(moved_id, i);
}
let revision = self.revision.fetch_add(1, Ordering::SeqCst) + 1;
(Revision(revision), RowPatch::Remove { id: id.clone() })
};
self.broadcast(patch);
revision
}
fn broadcast(&self, patch: RowPatch) {
let senders = self.live.lock().expect("mem plane live poisoned");
for tx in senders.iter() {
let _ = tx.unbounded_send(patch.clone());
}
}
pub fn rows(&self) -> Vec<Row> {
self.state.lock().expect("mem plane poisoned").rows.clone()
}
}
#[async_trait::async_trait]
impl DataPlane for MemPlane {
async fn query(&self, query: &Query) -> Result<Page, Error> {
let revision = self.revision();
let mut rows = {
let state = self.state.lock().expect("mem plane poisoned");
state.rows.clone()
};
sort_rows(&mut rows, &self.schema.column_ids, query.sort.as_ref());
rows = filter_rows(rows, query.filter.as_ref());
let total = rows.len();
let start = query.start.min(total);
let end = query.end.min(total).max(start);
let page_rows: Vec<Row> = rows.into_iter().skip(start).take(end - start).collect();
Ok(Page {
rows: page_rows,
total,
revision,
})
}
fn subscribe(&self, from: Revision) -> Pin<Box<dyn Stream<Item = RowPatch> + Send>> {
let history: Vec<RowPatch> = {
let state = self.state.lock().expect("mem plane poisoned");
state
.history
.iter()
.filter(|e| e.revision > from)
.map(|e| RowPatch::Upsert {
id: e.row.id.clone(),
row: e.row.clone(),
})
.collect()
};
let (tx, rx) = mpsc::unbounded();
self.live.lock().expect("mem plane live poisoned").push(tx);
let stream = futures_util::stream::iter(history).chain(rx);
Box::pin(stream)
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::render::{Cell, Filter, SortSpec};
fn text(v: &str) -> Cell {
Cell::Text { value: v.into() }
}
fn row(id: &str, name: &str, n: i64) -> Row {
Row {
id: RowId(id.into()),
cells: vec![text(name), Cell::Number { value: n }],
}
}
fn plane() -> MemPlane {
MemPlane::new(Schema {
column_ids: vec!["name".into(), "count".into()],
})
}
#[test]
fn query_filters_sorts_and_windows() {
let p = plane();
p.upsert(row("a", "zebra", 10));
p.upsert(row("b", "apple", 2));
p.upsert(row("c", "banana", 9));
let page = tokio_test_block(p.query(&Query {
start: 0,
end: 100,
sort: Some(SortSpec {
column: "count".into(),
descending: false,
}),
filter: None,
}))
.expect("query");
let names: Vec<String> = page
.rows
.iter()
.map(|r| crate::table::cell_text(&r.cells[0]))
.collect();
assert_eq!(names, vec!["apple", "banana", "zebra"]);
assert_eq!(page.total, 3);
let page2 = tokio_test_block(p.query(&Query {
start: 1,
end: 2,
sort: Some(SortSpec {
column: "count".into(),
descending: false,
}),
filter: None,
}))
.expect("query");
assert_eq!(page2.rows.len(), 1);
assert_eq!(page2.total, 3);
assert_eq!(crate::table::cell_text(&page2.rows[0].cells[0]), "banana");
}
#[test]
fn filter_by_expression() {
let p = plane();
p.upsert(row("a", "zebra", 10));
p.upsert(row("b", "apple", 2));
let page = tokio_test_block(p.query(&Query {
start: 0,
end: 100,
sort: None,
filter: Some(Filter {
expression: "app".into(),
}),
}))
.expect("query");
assert_eq!(page.rows.len(), 1);
assert_eq!(page.rows[0].id, RowId("b".into()));
}
#[test]
fn subscribe_replays_history_then_live() {
let p = plane();
p.upsert(row("a", "one", 1));
let from = p.revision(); p.upsert(row("b", "two", 2));
p.upsert(row("a", "one-v2", 1));
let mut stream = p.subscribe(from);
let mut got = Vec::new();
let waker = futures_util::task::noop_waker();
let mut cx = std::task::Context::from_waker(&waker);
let mut pin = Pin::new(&mut stream);
while let std::task::Poll::Ready(Some(item)) = pin.as_mut().poll_next(&mut cx) {
got.push(item);
if got.len() == 2 {
break;
}
}
assert_eq!(got.len(), 2);
}
#[test]
fn remove_deletes_row() {
let p = plane();
p.upsert(row("a", "one", 1));
p.upsert(row("b", "two", 2));
p.remove(&RowId("a".into()));
let rows = p.rows();
assert_eq!(rows.len(), 1);
assert_eq!(rows[0].id, RowId("b".into()));
}
fn tokio_test_block<F: std::future::Future>(f: F) -> F::Output {
futures_util::pin_mut!(f);
let waker = futures_util::task::noop_waker();
let mut cx = std::task::Context::from_waker(&waker);
loop {
match f.as_mut().poll(&mut cx) {
std::task::Poll::Ready(v) => return v,
std::task::Poll::Pending => std::thread::yield_now(),
}
}
}
}