1use std::collections::VecDeque;
11use std::pin::Pin;
12use std::sync::atomic::{AtomicU64, Ordering};
13use std::sync::{Arc, Mutex};
14
15use futures_channel::mpsc;
16use futures_util::{Stream, StreamExt};
17
18use crate::error::Error;
19use crate::render::{DataPlane, Page, Query, Revision, Row, RowId, RowPatch};
20use crate::table::{filter_rows, sort_rows};
21
22#[derive(Debug, Clone, Default)]
25pub struct Schema {
26 pub column_ids: Vec<String>,
27}
28
29#[derive(Debug, Clone)]
33struct HistoryEntry {
34 revision: Revision,
35 patch: RowPatch,
36}
37
38#[derive(Debug)]
39struct State {
40 rows: Vec<Row>,
43 index: std::collections::HashMap<RowId, usize>,
45 history: VecDeque<HistoryEntry>,
49}
50
51#[derive(Clone)]
53pub struct MemPlane {
54 state: Arc<Mutex<State>>,
55 revision: Arc<AtomicU64>,
56 schema: Schema,
57 live: Arc<Mutex<Vec<mpsc::UnboundedSender<RowPatch>>>>,
60 history_cap: usize,
61}
62
63impl MemPlane {
64 pub fn new(schema: Schema) -> Self {
66 Self {
67 state: Arc::new(Mutex::new(State {
68 rows: Vec::new(),
69 index: std::collections::HashMap::new(),
70 history: VecDeque::new(),
71 })),
72 revision: Arc::new(AtomicU64::new(0)),
73 schema,
74 live: Arc::new(Mutex::new(Vec::new())),
75 history_cap: 1024,
76 }
77 }
78
79 pub fn revision(&self) -> Revision {
81 Revision(self.revision.load(Ordering::SeqCst))
82 }
83
84 pub fn upsert(&self, row: Row) -> Revision {
86 let id = row.id.clone();
87 let (revision, patch) = {
88 let mut state = self.state.lock().expect("mem plane poisoned");
89 let revision = self.revision.fetch_add(1, Ordering::SeqCst) + 1;
90 match state.index.get(&id) {
91 Some(&i) => {
92 state.rows[i] = row.clone();
93 }
94 None => {
95 let len = state.rows.len();
96 state.index.insert(id.clone(), len);
97 state.rows.push(row.clone());
98 }
99 }
100 let patch = RowPatch::Upsert { id, row };
101 state.history.push_back(HistoryEntry {
102 revision: Revision(revision),
103 patch: patch.clone(),
104 });
105 if state.history.len() > self.history_cap {
106 state.history.pop_front();
107 }
108 (Revision(revision), patch)
109 };
110 self.broadcast(patch);
111 revision
112 }
113
114 pub fn remove(&self, id: &RowId) -> Revision {
116 let (revision, patch) = {
117 let mut state = self.state.lock().expect("mem plane poisoned");
118 let Some(i) = state.index.remove(id) else {
119 return self.revision();
120 };
121 state.rows.swap_remove(i);
122 if let Some(moved) = state.rows.get(i) {
123 let moved_id = moved.id.clone();
124 state.index.insert(moved_id, i);
125 }
126 let revision = self.revision.fetch_add(1, Ordering::SeqCst) + 1;
127 let patch = RowPatch::Remove { id: id.clone() };
128 state.history.push_back(HistoryEntry {
129 revision: Revision(revision),
130 patch: patch.clone(),
131 });
132 if state.history.len() > self.history_cap {
133 state.history.pop_front();
134 }
135 (Revision(revision), patch)
136 };
137 self.broadcast(patch);
138 revision
139 }
140
141 fn broadcast(&self, patch: RowPatch) {
142 let mut senders = self.live.lock().expect("mem plane live poisoned");
143 senders.retain(|tx| !tx.is_closed());
146 for tx in senders.iter() {
147 let _ = tx.unbounded_send(patch.clone());
148 }
149 }
150
151 pub fn rows(&self) -> Vec<Row> {
153 self.state.lock().expect("mem plane poisoned").rows.clone()
154 }
155
156 pub fn schema_column_ids(&self) -> Vec<String> {
158 self.schema.column_ids.clone()
159 }
160}
161
162#[async_trait::async_trait]
163impl DataPlane for MemPlane {
164 async fn query(&self, query: &Query) -> Result<Page, Error> {
165 let revision = self.revision();
166 let mut rows = {
167 let state = self.state.lock().expect("mem plane poisoned");
168 state.rows.clone()
169 };
170 sort_rows(&mut rows, &self.schema.column_ids, query.sort.as_ref());
171 rows = filter_rows(rows, query.filter.as_ref());
172
173 let total = rows.len();
174 let start = query.start.min(total);
175 let end = query.end.min(total).max(start);
176 let page_rows: Vec<Row> = rows.into_iter().skip(start).take(end - start).collect();
177 Ok(Page {
178 rows: page_rows,
179 total,
180 revision,
181 })
182 }
183
184 fn subscribe(&self, from: Revision) -> Pin<Box<dyn Stream<Item = RowPatch> + Send>> {
185 let history: Vec<RowPatch> = {
190 let state = self.state.lock().expect("mem plane poisoned");
191 state
192 .history
193 .iter()
194 .filter(|e| e.revision > from)
195 .map(|e| e.patch.clone())
196 .collect()
197 };
198 let (tx, rx) = mpsc::unbounded();
199 self.live.lock().expect("mem plane live poisoned").push(tx);
200 let stream = futures_util::stream::iter(history).chain(rx);
201 Box::pin(stream)
202 }
203}
204
205#[cfg(test)]
206mod tests {
207 use super::*;
208 use crate::render::{Cell, Filter, SortSpec};
209
210 fn text(v: &str) -> Cell {
211 Cell::Text { value: v.into() }
212 }
213 fn row(id: &str, name: &str, n: i64) -> Row {
214 Row {
215 id: RowId(id.into()),
216 cells: vec![text(name), Cell::Number { value: n }],
217 }
218 }
219
220 fn plane() -> MemPlane {
221 MemPlane::new(Schema {
222 column_ids: vec!["name".into(), "count".into()],
223 })
224 }
225
226 #[test]
227 fn query_filters_sorts_and_windows() {
228 let p = plane();
229 p.upsert(row("a", "zebra", 10));
230 p.upsert(row("b", "apple", 2));
231 p.upsert(row("c", "banana", 9));
232
233 let page = tokio_test_block(p.query(&Query {
234 start: 0,
235 end: 100,
236 sort: Some(SortSpec {
237 column: "count".into(),
238 descending: false,
239 }),
240 filter: None,
241 }))
242 .expect("query");
243 let names: Vec<String> = page
244 .rows
245 .iter()
246 .map(|r| crate::table::cell_text(&r.cells[0]))
247 .collect();
248 assert_eq!(names, vec!["apple", "banana", "zebra"]);
249 assert_eq!(page.total, 3);
250
251 let page2 = tokio_test_block(p.query(&Query {
253 start: 1,
254 end: 2,
255 sort: Some(SortSpec {
256 column: "count".into(),
257 descending: false,
258 }),
259 filter: None,
260 }))
261 .expect("query");
262 assert_eq!(page2.rows.len(), 1);
263 assert_eq!(page2.total, 3);
264 assert_eq!(crate::table::cell_text(&page2.rows[0].cells[0]), "banana");
265 }
266
267 #[test]
268 fn filter_by_expression() {
269 let p = plane();
270 p.upsert(row("a", "zebra", 10));
271 p.upsert(row("b", "apple", 2));
272 let page = tokio_test_block(p.query(&Query {
273 start: 0,
274 end: 100,
275 sort: None,
276 filter: Some(Filter {
277 expression: "app".into(),
278 }),
279 }))
280 .expect("query");
281 assert_eq!(page.rows.len(), 1);
282 assert_eq!(page.rows[0].id, RowId("b".into()));
283 }
284
285 #[test]
286 fn subscribe_replays_history_then_live() {
287 let p = plane();
288 p.upsert(row("a", "one", 1));
289 let from = p.revision(); p.upsert(row("b", "two", 2));
291 p.upsert(row("a", "one-v2", 1)); let mut stream = p.subscribe(from);
295 let mut got = Vec::new();
296 let waker = futures_util::task::noop_waker();
298 let mut cx = std::task::Context::from_waker(&waker);
299 let mut pin = Pin::new(&mut stream);
300 while let std::task::Poll::Ready(Some(item)) = pin.as_mut().poll_next(&mut cx) {
301 got.push(item);
302 if got.len() == 2 {
303 break;
304 }
305 }
306 assert_eq!(got.len(), 2);
307 }
308
309 #[test]
310 fn remove_deletes_row() {
311 let p = plane();
312 p.upsert(row("a", "one", 1));
313 p.upsert(row("b", "two", 2));
314 p.remove(&RowId("a".into()));
315 let rows = p.rows();
316 assert_eq!(rows.len(), 1);
317 assert_eq!(rows[0].id, RowId("b".into()));
318 }
319
320 #[test]
321 fn history_replay_includes_removals() {
322 let p = plane();
323 p.upsert(row("a", "one", 1));
324 let from = p.revision(); p.upsert(row("b", "two", 2));
326 p.remove(&RowId("a".into())); let mut stream = p.subscribe(from);
329 let waker = futures_util::task::noop_waker();
330 let mut cx = std::task::Context::from_waker(&waker);
331 let mut pin = Pin::new(&mut stream);
332 let mut got = Vec::new();
333 while let std::task::Poll::Ready(Some(item)) = pin.as_mut().poll_next(&mut cx) {
334 got.push(item);
335 if got.len() == 2 {
336 break;
337 }
338 }
339 assert_eq!(got.len(), 2);
341 assert!(matches!(&got[0], RowPatch::Upsert { id, .. } if id.0 == "b"));
342 assert!(matches!(&got[1], RowPatch::Remove { id } if id.0 == "a"));
343 }
344
345 #[test]
346 fn broadcast_drops_closed_senders() {
347 let p = plane();
348 {
350 let _stream = p.subscribe(Revision(0));
351 }
352 p.upsert(row("a", "one", 1));
353 let live_len = p.live.lock().expect("live").len();
356 assert!(
357 live_len <= 1,
358 "closed sender should be dropped, got {live_len}"
359 );
360 }
361
362 #[test]
368 fn query_50k_rows_is_not_quadratic() {
369 let p = plane();
370 for i in 0..50_000 {
371 p.upsert(row(&format!("row-{i}"), &format!("name-{i}"), i));
372 }
373 let start = std::time::Instant::now();
374 for _ in 0..10 {
375 let page = tokio_test_block(p.query(&Query {
376 start: 0,
377 end: 50,
378 sort: Some(SortSpec {
379 column: "count".into(),
380 descending: true,
381 }),
382 filter: Some(Filter {
383 expression: "name-4".into(),
384 }),
385 }))
386 .expect("query");
387 assert_eq!(
388 page.total, 11_111,
389 "filter should match name-4*, name-40*.."
390 );
391 }
392 let elapsed = start.elapsed();
393 assert!(
396 elapsed.as_secs() < 5,
397 "50k-row query regressed: {elapsed:?} for 10 queries"
398 );
399 }
400
401 fn tokio_test_block<F: std::future::Future>(f: F) -> F::Output {
403 futures_util::pin_mut!(f);
404 let waker = futures_util::task::noop_waker();
405 let mut cx = std::task::Context::from_waker(&waker);
406 loop {
407 match f.as_mut().poll(&mut cx) {
408 std::task::Poll::Ready(v) => return v,
409 std::task::Poll::Pending => std::thread::yield_now(),
410 }
411 }
412 }
413}