1use std::pin::Pin;
11use std::sync::atomic::{AtomicU64, Ordering};
12use std::sync::{Arc, Mutex};
13
14use futures_channel::mpsc;
15use futures_util::{Stream, StreamExt};
16
17use crate::error::Error;
18use crate::render::{DataPlane, Page, Query, Revision, Row, RowId, RowPatch};
19use crate::table::{filter_rows, sort_rows};
20
21#[derive(Debug, Clone, Default)]
24pub struct Schema {
25 pub column_ids: Vec<String>,
26}
27
28#[derive(Debug, Clone)]
30struct HistoryEntry {
31 revision: Revision,
32 row: Row,
33}
34
35#[derive(Debug)]
36struct State {
37 rows: Vec<Row>,
40 index: std::collections::HashMap<RowId, usize>,
42 history: Vec<HistoryEntry>,
44}
45
46#[derive(Clone)]
48pub struct MemPlane {
49 state: Arc<Mutex<State>>,
50 revision: Arc<AtomicU64>,
51 schema: Schema,
52 live: Arc<Mutex<Vec<mpsc::UnboundedSender<RowPatch>>>>,
55 history_cap: usize,
56}
57
58impl MemPlane {
59 pub fn new(schema: Schema) -> Self {
61 Self {
62 state: Arc::new(Mutex::new(State {
63 rows: Vec::new(),
64 index: std::collections::HashMap::new(),
65 history: Vec::new(),
66 })),
67 revision: Arc::new(AtomicU64::new(0)),
68 schema,
69 live: Arc::new(Mutex::new(Vec::new())),
70 history_cap: 1024,
71 }
72 }
73
74 pub fn revision(&self) -> Revision {
76 Revision(self.revision.load(Ordering::SeqCst))
77 }
78
79 pub fn upsert(&self, row: Row) -> Revision {
81 let id = row.id.clone();
82 let (revision, patch) = {
83 let mut state = self.state.lock().expect("mem plane poisoned");
84 let revision = self.revision.fetch_add(1, Ordering::SeqCst) + 1;
85 match state.index.get(&id) {
86 Some(&i) => {
87 state.rows[i] = row.clone();
88 }
89 None => {
90 let len = state.rows.len();
91 state.index.insert(id.clone(), len);
92 state.rows.push(row.clone());
93 }
94 }
95 state.history.push(HistoryEntry {
96 revision: Revision(revision),
97 row: row.clone(),
98 });
99 if state.history.len() > self.history_cap {
100 state.history.remove(0);
101 }
102 (Revision(revision), RowPatch::Upsert { id, row })
103 };
104 self.broadcast(patch);
105 revision
106 }
107
108 pub fn remove(&self, id: &RowId) -> Revision {
110 let (revision, patch) = {
111 let mut state = self.state.lock().expect("mem plane poisoned");
112 let Some(i) = state.index.remove(id) else {
113 return self.revision();
114 };
115 state.rows.swap_remove(i);
116 if let Some(moved) = state.rows.get(i) {
117 let moved_id = moved.id.clone();
118 state.index.insert(moved_id, i);
119 }
120 let revision = self.revision.fetch_add(1, Ordering::SeqCst) + 1;
121 (Revision(revision), RowPatch::Remove { id: id.clone() })
122 };
123 self.broadcast(patch);
124 revision
125 }
126
127 fn broadcast(&self, patch: RowPatch) {
128 let senders = self.live.lock().expect("mem plane live poisoned");
129 for tx in senders.iter() {
130 let _ = tx.unbounded_send(patch.clone());
131 }
132 }
133
134 pub fn rows(&self) -> Vec<Row> {
136 self.state.lock().expect("mem plane poisoned").rows.clone()
137 }
138}
139
140#[async_trait::async_trait]
141impl DataPlane for MemPlane {
142 async fn query(&self, query: &Query) -> Result<Page, Error> {
143 let revision = self.revision();
144 let mut rows = {
145 let state = self.state.lock().expect("mem plane poisoned");
146 state.rows.clone()
147 };
148 sort_rows(&mut rows, &self.schema.column_ids, query.sort.as_ref());
149 rows = filter_rows(rows, query.filter.as_ref());
150
151 let total = rows.len();
152 let start = query.start.min(total);
153 let end = query.end.min(total).max(start);
154 let page_rows: Vec<Row> = rows.into_iter().skip(start).take(end - start).collect();
155 Ok(Page {
156 rows: page_rows,
157 total,
158 revision,
159 })
160 }
161
162 fn subscribe(&self, from: Revision) -> Pin<Box<dyn Stream<Item = RowPatch> + Send>> {
163 let history: Vec<RowPatch> = {
167 let state = self.state.lock().expect("mem plane poisoned");
168 state
169 .history
170 .iter()
171 .filter(|e| e.revision > from)
172 .map(|e| RowPatch::Upsert {
173 id: e.row.id.clone(),
174 row: e.row.clone(),
175 })
176 .collect()
177 };
178 let (tx, rx) = mpsc::unbounded();
179 self.live.lock().expect("mem plane live poisoned").push(tx);
180 let stream = futures_util::stream::iter(history).chain(rx);
181 Box::pin(stream)
182 }
183}
184
185#[cfg(test)]
186mod tests {
187 use super::*;
188 use crate::render::{Cell, Filter, SortSpec};
189
190 fn text(v: &str) -> Cell {
191 Cell::Text { value: v.into() }
192 }
193 fn row(id: &str, name: &str, n: i64) -> Row {
194 Row {
195 id: RowId(id.into()),
196 cells: vec![text(name), Cell::Number { value: n }],
197 }
198 }
199
200 fn plane() -> MemPlane {
201 MemPlane::new(Schema {
202 column_ids: vec!["name".into(), "count".into()],
203 })
204 }
205
206 #[test]
207 fn query_filters_sorts_and_windows() {
208 let p = plane();
209 p.upsert(row("a", "zebra", 10));
210 p.upsert(row("b", "apple", 2));
211 p.upsert(row("c", "banana", 9));
212
213 let page = tokio_test_block(p.query(&Query {
214 start: 0,
215 end: 100,
216 sort: Some(SortSpec {
217 column: "count".into(),
218 descending: false,
219 }),
220 filter: None,
221 }))
222 .expect("query");
223 let names: Vec<String> = page
224 .rows
225 .iter()
226 .map(|r| crate::table::cell_text(&r.cells[0]))
227 .collect();
228 assert_eq!(names, vec!["apple", "banana", "zebra"]);
229 assert_eq!(page.total, 3);
230
231 let page2 = tokio_test_block(p.query(&Query {
233 start: 1,
234 end: 2,
235 sort: Some(SortSpec {
236 column: "count".into(),
237 descending: false,
238 }),
239 filter: None,
240 }))
241 .expect("query");
242 assert_eq!(page2.rows.len(), 1);
243 assert_eq!(page2.total, 3);
244 assert_eq!(crate::table::cell_text(&page2.rows[0].cells[0]), "banana");
245 }
246
247 #[test]
248 fn filter_by_expression() {
249 let p = plane();
250 p.upsert(row("a", "zebra", 10));
251 p.upsert(row("b", "apple", 2));
252 let page = tokio_test_block(p.query(&Query {
253 start: 0,
254 end: 100,
255 sort: None,
256 filter: Some(Filter {
257 expression: "app".into(),
258 }),
259 }))
260 .expect("query");
261 assert_eq!(page.rows.len(), 1);
262 assert_eq!(page.rows[0].id, RowId("b".into()));
263 }
264
265 #[test]
266 fn subscribe_replays_history_then_live() {
267 let p = plane();
268 p.upsert(row("a", "one", 1));
269 let from = p.revision(); p.upsert(row("b", "two", 2));
271 p.upsert(row("a", "one-v2", 1)); let mut stream = p.subscribe(from);
275 let mut got = Vec::new();
276 let waker = futures_util::task::noop_waker();
278 let mut cx = std::task::Context::from_waker(&waker);
279 let mut pin = Pin::new(&mut stream);
280 while let std::task::Poll::Ready(Some(item)) = pin.as_mut().poll_next(&mut cx) {
281 got.push(item);
282 if got.len() == 2 {
283 break;
284 }
285 }
286 assert_eq!(got.len(), 2);
287 }
288
289 #[test]
290 fn remove_deletes_row() {
291 let p = plane();
292 p.upsert(row("a", "one", 1));
293 p.upsert(row("b", "two", 2));
294 p.remove(&RowId("a".into()));
295 let rows = p.rows();
296 assert_eq!(rows.len(), 1);
297 assert_eq!(rows[0].id, RowId("b".into()));
298 }
299
300 fn tokio_test_block<F: std::future::Future>(f: F) -> F::Output {
302 futures_util::pin_mut!(f);
303 let waker = futures_util::task::noop_waker();
304 let mut cx = std::task::Context::from_waker(&waker);
305 loop {
306 match f.as_mut().poll(&mut cx) {
307 std::task::Poll::Ready(v) => return v,
308 std::task::Poll::Pending => std::thread::yield_now(),
309 }
310 }
311 }
312}