wasm_dbms_memory/table_registry/
raw_table_reader.rs1use wasm_dbms_api::prelude::{DecodeError, MSize, MemoryError, MemoryResult, Page, PageOffset};
11
12use super::page_ledger::PageLedger;
13use super::raw_record::RAW_RECORD_HEADER_SIZE;
14use super::record_address::RecordAddress;
15use crate::MemoryAccess;
16
17#[derive(Debug, Clone, PartialEq, Eq)]
19pub struct RawRecordBytes {
20 pub address: RecordAddress,
22 pub bytes: Vec<u8>,
24}
25
26pub struct RawTableReader<'a, MA>
30where
31 MA: MemoryAccess,
32{
33 mm: &'a mut MA,
34 page_ledger: &'a PageLedger,
35 page_size: usize,
36 alignment: PageOffset,
37 cursor: Option<Cursor>,
38}
39
40#[derive(Debug, Copy, Clone)]
41struct Cursor {
42 page: Page,
43 offset: PageOffset,
44}
45
46impl<'a, MA> RawTableReader<'a, MA>
47where
48 MA: MemoryAccess,
49{
50 pub fn new(page_ledger: &'a PageLedger, alignment: PageOffset, mm: &'a mut MA) -> Self {
54 let page_size = mm.page_size() as usize;
55 let cursor = page_ledger.pages().first().map(|p| Cursor {
56 page: p.page,
57 offset: 0,
58 });
59 Self {
60 mm,
61 page_ledger,
62 page_size,
63 alignment,
64 cursor,
65 }
66 }
67
68 pub fn try_next(&mut self) -> MemoryResult<Option<RawRecordBytes>> {
70 loop {
71 let Some(Cursor { page, offset }) = self.cursor else {
72 return Ok(None);
73 };
74 let aligned_usize = align_up_usize(offset as usize, self.alignment as usize);
75 if aligned_usize + (RAW_RECORD_HEADER_SIZE as usize) > self.page_size {
77 self.cursor = self.next_page(page);
78 continue;
79 }
80 let aligned = aligned_usize as PageOffset;
81 let mut header = [0u8; RAW_RECORD_HEADER_SIZE as usize];
82 self.mm.read_at_raw(page, aligned, &mut header)?;
83 let length = u16::from_le_bytes(header) as MSize;
84 if length == 0 {
85 let next_offset = aligned_usize + self.alignment as usize;
87 self.cursor = if next_offset >= self.page_size {
88 self.next_page(page)
89 } else {
90 Some(Cursor {
91 page,
92 offset: next_offset as PageOffset,
93 })
94 };
95 continue;
96 }
97 let body_offset_usize = aligned_usize + RAW_RECORD_HEADER_SIZE as usize;
98 if body_offset_usize + (length as usize) > self.page_size {
99 return Err(MemoryError::DecodeError(DecodeError::TooShort));
100 }
101 let body_offset = body_offset_usize as PageOffset;
102 let mut body = vec![0u8; length as usize];
103 self.mm.read_at_raw(page, body_offset, &mut body)?;
104 let address = RecordAddress {
105 page,
106 offset: aligned,
107 };
108 let next_offset =
110 align_up_usize(body_offset_usize + length as usize, self.alignment as usize);
111 self.cursor = if next_offset >= self.page_size {
112 self.next_page(page)
113 } else {
114 Some(Cursor {
115 page,
116 offset: next_offset as PageOffset,
117 })
118 };
119 return Ok(Some(RawRecordBytes {
120 address,
121 bytes: body,
122 }));
123 }
124 }
125
126 fn next_page(&self, current: Page) -> Option<Cursor> {
127 self.page_ledger
128 .next_page_after(current)
129 .map(|page| Cursor { page, offset: 0 })
130 }
131}
132
133fn align_up_usize(offset: usize, alignment: usize) -> usize {
134 if alignment == 0 {
135 return offset;
136 }
137 let rem = offset % alignment;
138 if rem == 0 {
139 offset
140 } else {
141 offset + (alignment - rem)
142 }
143}
144
145#[cfg(test)]
146mod tests {
147 use super::*;
148 use crate::table_registry::test_utils::{User, write_dummy_schema_snapshot};
149 use crate::{HeapMemoryProvider, MemoryManager, TableRegistry, TableRegistryPage};
150
151 #[test]
152 fn test_insert_raw_round_trips_through_raw_reader() {
153 let mut mm = MemoryManager::init(HeapMemoryProvider::default());
154 let schema_snapshot_page = mm.claim_page().unwrap();
155 let pages_list_page = mm.claim_page().unwrap();
156 let free_segments_page = mm.claim_page().unwrap();
157 let index_registry_page = mm.claim_page().unwrap();
158 write_dummy_schema_snapshot(schema_snapshot_page, &mut mm);
159
160 let mut registry = TableRegistry::load(
161 TableRegistryPage {
162 schema_snapshot_page,
163 pages_list_page,
164 free_segments_page,
165 index_registry_page,
166 autoincrement_registry_page: None,
167 },
168 &mut mm,
169 )
170 .unwrap();
171
172 let alignment = 32u16;
173 let payload = vec![1u8, 2, 3, 4, 5, 6, 7];
174 let address = registry.insert_raw(&payload, alignment, &mut mm).unwrap();
175 let read_back = registry.read_raw_at(address, &mut mm).unwrap();
176 assert_eq!(read_back, payload);
177
178 let mut reader = RawTableReader::new(®istry.page_ledger, alignment, &mut mm);
179 let row = reader.try_next().unwrap().expect("missing row");
180 assert_eq!(row.bytes, payload);
181 assert!(reader.try_next().unwrap().is_none());
182 }
183
184 #[test]
185 fn test_delete_raw_frees_segment() {
186 let mut mm = MemoryManager::init(HeapMemoryProvider::default());
187 let schema_snapshot_page = mm.claim_page().unwrap();
188 let pages_list_page = mm.claim_page().unwrap();
189 let free_segments_page = mm.claim_page().unwrap();
190 let index_registry_page = mm.claim_page().unwrap();
191 write_dummy_schema_snapshot(schema_snapshot_page, &mut mm);
192
193 let mut registry = TableRegistry::load(
194 TableRegistryPage {
195 schema_snapshot_page,
196 pages_list_page,
197 free_segments_page,
198 index_registry_page,
199 autoincrement_registry_page: None,
200 },
201 &mut mm,
202 )
203 .unwrap();
204
205 let alignment = 32u16;
206 let bytes = vec![9u8; 10];
207 let addr = registry.insert_raw(&bytes, alignment, &mut mm).unwrap();
208 registry
209 .delete_raw(addr, bytes.len() as MSize, alignment, &mut mm)
210 .unwrap();
211
212 let mut reader = RawTableReader::new(®istry.page_ledger, alignment, &mut mm);
213 assert!(reader.try_next().unwrap().is_none());
214 }
215
216 #[test]
217 fn test_raw_table_reader_reads_all_records_as_bytes() {
218 use wasm_dbms_api::prelude::Encode;
219
220 let mut mm = MemoryManager::init(HeapMemoryProvider::default());
221 let schema_snapshot_page = mm.claim_page().unwrap();
222 let pages_list_page = mm.claim_page().unwrap();
223 let free_segments_page = mm.claim_page().unwrap();
224 let index_registry_page = mm.claim_page().unwrap();
225 write_dummy_schema_snapshot(schema_snapshot_page, &mut mm);
226
227 let mut registry = TableRegistry::load(
228 TableRegistryPage {
229 schema_snapshot_page,
230 pages_list_page,
231 free_segments_page,
232 index_registry_page,
233 autoincrement_registry_page: None,
234 },
235 &mut mm,
236 )
237 .unwrap();
238
239 for id in 0..3u32 {
240 let user = User {
241 id,
242 name: format!("U{id}"),
243 email: "x@x".into(),
244 age: 20,
245 };
246 registry.insert(user, &mut mm).unwrap();
247 }
248
249 let mut reader = RawTableReader::new(®istry.page_ledger, User::ALIGNMENT, &mut mm);
250 let mut count = 0;
251 while reader.try_next().unwrap().is_some() {
252 count += 1;
253 }
254 assert_eq!(count, 3);
255 }
256
257 #[test]
258 fn test_raw_reader_scans_reclaimed_pages_in_ledger_order() {
259 let mut mm = MemoryManager::init(HeapMemoryProvider::default());
260 let schema_snapshot_page = mm.claim_page().unwrap();
261 let pages_list_page = mm.claim_page().unwrap();
262 let free_segments_page = mm.claim_page().unwrap();
263 let index_registry_page = mm.claim_page().unwrap();
264 write_dummy_schema_snapshot(schema_snapshot_page, &mut mm);
265
266 let mut registry = TableRegistry::load(
267 TableRegistryPage {
268 schema_snapshot_page,
269 pages_list_page,
270 free_segments_page,
271 index_registry_page,
272 autoincrement_registry_page: None,
273 },
274 &mut mm,
275 )
276 .unwrap();
277
278 let low = mm.claim_page().unwrap();
280 let high = mm.claim_page().unwrap();
281 mm.unclaim_page(low).unwrap();
282 mm.unclaim_page(high).unwrap();
283
284 let alignment = 32u16;
285 registry
286 .insert_raw(&vec![1u8; 40_000], alignment, &mut mm)
287 .unwrap();
288 registry
289 .insert_raw(&vec![2u8; 40_000], alignment, &mut mm)
290 .unwrap();
291
292 let mut reader = RawTableReader::new(®istry.page_ledger, alignment, &mut mm);
293 let mut first_bytes = Vec::new();
294 while let Some(row) = reader.try_next().unwrap() {
295 first_bytes.push(row.bytes[0]);
296 }
297 assert_eq!(first_bytes, vec![1, 2]);
298 }
299}