1use std::sync::atomic::{AtomicU64, Ordering};
10use std::sync::Arc;
11
12use parking_lot::Mutex;
13use uqa_core::memory::{BudgetedSharedMap, BudgetedVec, MemoryBudget, MemoryReservation};
14
15use crate::read_control::StorageReadControl;
16use crate::StorageSavepointId;
17
18use super::key::RecordKey;
19use super::{PreparedRecordCommit, PreparedRecordWrite, RecordWrite, VersionError, VersionResult};
20
21struct Change {
22 write: PreparedRecordWrite,
23 identity: PrivateRecordRevision,
24}
25
26type Records = BudgetedSharedMap<RecordKey, Change>;
27
28#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord)]
30pub struct PrivateRecordRevision(u64);
31
32impl PrivateRecordRevision {
33 fn allocate() -> VersionResult<Self> {
34 static NEXT: AtomicU64 = AtomicU64::new(1);
35 NEXT.fetch_update(Ordering::Relaxed, Ordering::Relaxed, |id| id.checked_add(1))
36 .map(Self)
37 .map_err(|_| VersionError::PrivateRevisionExhausted)
38 }
39
40 pub fn as_u64(self) -> u64 {
41 self.0
42 }
43}
44
45pub struct PrivateRecordKey {
47 key: RecordKey,
48 revision: PrivateRecordRevision,
49}
50
51impl PrivateRecordKey {
52 pub fn key(&self) -> &[u8] {
53 self.key.bytes()
54 }
55
56 pub fn revision(&self) -> PrivateRecordRevision {
57 self.revision
58 }
59}
60
61struct Savepoint {
62 id: StorageSavepointId,
63 records: Records,
64}
65
66struct State {
67 records: Records,
68 savepoints: BudgetedVec<Savepoint>,
69}
70
71impl State {
72 fn savepoint_position(&self, id: StorageSavepointId) -> VersionResult<usize> {
73 self.savepoints
74 .iter()
75 .rposition(|savepoint| savepoint.id == id)
76 .ok_or(VersionError::SavepointMissing(id))
77 }
78
79 fn truncate_savepoints(&mut self, len: usize) {
80 if len == 0 {
81 self.savepoints = BudgetedVec::new(self.records.budget());
82 } else {
83 self.savepoints.truncate(len);
84 }
85 }
86}
87
88struct Owner {
89 state: Mutex<State>,
90 memory: MemoryBudget,
91}
92
93pub struct PrivateRecordChanges {
97 owner: Arc<Owner>,
98}
99
100impl PrivateRecordChanges {
101 pub(super) fn share_owner(&self) -> Self {
102 Self {
103 owner: Arc::clone(&self.owner),
104 }
105 }
106
107 pub(super) fn write_kind(
108 &self,
109 key: &[u8],
110 control: &StorageReadControl,
111 ) -> VersionResult<Option<super::commit::RecordWriteKind>> {
112 control.check()?;
113 Ok(self
114 .owner
115 .state
116 .lock()
117 .records
118 .get(key)
119 .map(|change| change.write.kind()))
120 }
121
122 pub fn new(memory: &MemoryBudget) -> Self {
123 Self {
124 owner: Arc::new(Owner {
125 state: Mutex::new(State {
126 records: Records::new(memory),
127 savepoints: BudgetedVec::new(memory),
128 }),
129 memory: memory.clone(),
130 }),
131 }
132 }
133
134 pub fn apply(
136 &self,
137 writes: &[RecordWrite<'_>],
138 control: &StorageReadControl,
139 ) -> VersionResult<()> {
140 let owned_control = StorageReadControl::new(&self.owner.memory, control.cancellation());
141 let prepared = PreparedRecordCommit::new(writes, &owned_control)?;
142 self.apply_owned(prepared.records(), control)
143 }
144
145 pub(crate) fn apply_owned(
146 &self,
147 writes: &[PreparedRecordWrite],
148 control: &StorageReadControl,
149 ) -> VersionResult<()> {
150 if writes.is_empty() {
151 return Ok(());
152 }
153 let mut state = self.owner.state.lock();
154 for (index, write) in writes.iter().enumerate() {
155 control.cancellation().check()?;
156 if let Some(previous) = state.records.get(write.key()) {
157 if previous.write.expected() != write.expected() {
158 return Err(VersionError::WriteConflict {
159 mutation: index,
160 expected: write.expected(),
161 actual: previous.write.expected(),
162 });
163 }
164 }
165 }
166 let identity = PrivateRecordRevision::allocate()?;
167 if let [write] = writes {
168 control.cancellation().check()?;
169 state.records.try_insert(
170 write.shared_key(),
171 Change {
172 write: write.clone(),
173 identity,
174 },
175 )?;
176 return Ok(());
177 }
178 let mut records = state.records.clone();
179 for write in writes {
180 control.cancellation().check()?;
181 records.try_insert(
182 write.shared_key(),
183 Change {
184 write: write.clone(),
185 identity,
186 },
187 )?;
188 }
189 control.cancellation().check()?;
190 state.records = records;
192 Ok(())
193 }
194
195 pub fn has_written(&self) -> bool {
196 !self.owner.state.lock().records.is_empty()
197 }
198
199 pub fn snapshot(&self) -> VersionResult<PrivateRecordSnapshot> {
200 let memory = self
201 .owner
202 .memory
203 .reserve(std::mem::size_of::<PrivateRecordSnapshot>())?;
204 Ok(PrivateRecordSnapshot {
205 records: self.owner.state.lock().records.clone(),
206 _memory: memory,
207 })
208 }
209
210 pub fn savepoint(&self, id: StorageSavepointId) -> VersionResult<()> {
211 let mut state = self.owner.state.lock();
212 let records = state.records.clone();
213 state.savepoints.push(Savepoint { id, records })?;
214 Ok(())
215 }
216
217 pub fn release_savepoint(&self, id: StorageSavepointId) -> VersionResult<()> {
219 let mut state = self.owner.state.lock();
220 let position = state.savepoint_position(id)?;
221 state.truncate_savepoints(position);
222 Ok(())
223 }
224
225 pub fn rollback_to_savepoint(&self, id: StorageSavepointId) -> VersionResult<()> {
227 let mut state = self.owner.state.lock();
228 let position = state.savepoint_position(id)?;
229 state.records = state.savepoints[position].records.clone();
230 state.truncate_savepoints(position + 1);
231 Ok(())
232 }
233
234 pub fn rollback(&self) -> VersionResult<()> {
235 let mut state = self.owner.state.lock();
236 state.records = Records::new(&self.owner.memory);
237 state.truncate_savepoints(0);
238 Ok(())
239 }
240
241 pub fn prepare(&self, control: &StorageReadControl) -> VersionResult<PreparedRecordCommit> {
243 control.cancellation().check()?;
244 let state = self.owner.state.lock();
245 let mut writes = BudgetedVec::new(control.memory());
246 for (_, change) in &state.records {
247 control.cancellation().check()?;
248 writes.push(change.write.clone())?;
249 }
250 PreparedRecordCommit::from_unique_owned(writes, control)
251 }
252}
253
254pub struct PrivateRecordSnapshot {
256 records: Records,
257 _memory: MemoryReservation,
258}
259
260impl PrivateRecordSnapshot {
261 pub(super) fn try_clone(&self) -> VersionResult<Self> {
263 let memory = self.records.budget().reserve(std::mem::size_of::<Self>())?;
264 Ok(Self {
265 records: self.records.clone(),
266 _memory: memory,
267 })
268 }
269
270 pub fn scan_keys(
272 &self,
273 prefix: &[u8],
274 after: Option<&[u8]>,
275 limit: usize,
276 control: &StorageReadControl,
277 ) -> VersionResult<BudgetedVec<PrivateRecordKey>> {
278 control.check()?;
279 let mut result = BudgetedVec::new(control.memory());
280 if limit == 0 {
281 return Ok(result);
282 }
283 let start = after.filter(|after| *after >= prefix).unwrap_or(prefix);
284 for (key, change) in self.records.range_from(std::ops::Bound::Included(start)) {
285 control.check()?;
286 let key = key.bytes();
287 if !key.starts_with(prefix) {
288 break;
289 }
290 if after.is_some_and(|after| key <= after) {
291 continue;
292 }
293 result.push(PrivateRecordKey {
294 key: change.write.shared_key(),
295 revision: change.identity,
296 })?;
297 if result.len() == limit {
298 break;
299 }
300 }
301 Ok(result)
302 }
303
304 pub(super) fn visit_merged<P: super::projection::Projection>(
305 &self,
306 committed: &dyn super::CommittedRecordSnapshot,
307 prefix: &[u8],
308 after: Option<&[u8]>,
309 limit: usize,
310 control: &StorageReadControl,
311 visit: &mut super::projection::Visitor<'_, P>,
312 ) -> VersionResult<()> {
313 control.cancellation().check()?;
314 if limit == 0 {
315 return Ok(());
316 }
317 let start = after.filter(|after| *after >= prefix).unwrap_or(prefix);
318 let mut entries = self.records.range_from(std::ops::Bound::Included(start));
319 let mut next_private = || -> VersionResult<Option<&PreparedRecordWrite>> {
320 for (key, change) in entries.by_ref() {
321 control.cancellation().check()?;
322 if !key.bytes().starts_with(prefix) {
323 return Ok(None);
324 }
325 if after.is_some_and(|after| key.bytes() <= after) {
326 continue;
327 }
328 return Ok(Some(&change.write));
329 }
330 Ok(None)
331 };
332 let mut pending = next_private()?;
333 let mut count = 0;
334 let mut running = true;
335 let mut emit = |key: &[u8], record: P::Record<'_>| -> VersionResult<bool> {
336 control.cancellation().check()?;
337 count += 1;
338 let more = visit(key, record)?;
339 control.cancellation().check()?;
340 Ok(more && count < limit)
341 };
342 P::visit(committed, prefix, after, control, &mut |key, record| {
343 while let Some(write) = pending.filter(|write| write.key() < key) {
344 running = emit(write.key(), P::private(write))?;
345 if !running {
346 return Ok(false);
347 }
348 pending = next_private()?;
349 }
350 if let Some(write) = pending.filter(|write| write.key() == key) {
351 running = emit(key, P::private(write))?;
352 pending = next_private()?;
353 } else {
354 running = emit(key, record)?;
355 }
356 Ok(running)
357 })?;
358 while running {
359 let Some(write) = pending else {
360 break;
361 };
362 running = emit(write.key(), P::private(write))?;
363 pending = next_private()?;
364 }
365 Ok(())
366 }
367
368 pub fn get(
369 &self,
370 key: &[u8],
371 control: &StorageReadControl,
372 ) -> VersionResult<Option<PreparedRecordWrite>> {
373 control.cancellation().check()?;
374 Ok(self.records.get(key).map(|change| change.write.clone()))
375 }
376
377 pub fn scan(
379 &self,
380 prefix: &[u8],
381 after: Option<&[u8]>,
382 limit: usize,
383 control: &StorageReadControl,
384 ) -> VersionResult<BudgetedVec<PreparedRecordWrite>> {
385 control.cancellation().check()?;
386 let mut result = BudgetedVec::new(control.memory());
387 if limit == 0 {
388 return Ok(result);
389 }
390 let start = after.filter(|after| *after >= prefix).unwrap_or(prefix);
391 for (key, change) in self.records.range_from(std::ops::Bound::Included(start)) {
392 control.cancellation().check()?;
393 let key = key.bytes();
394 if !key.starts_with(prefix) {
395 break;
396 }
397 if after.is_some_and(|after| key <= after) {
398 continue;
399 }
400 result.push(change.write.clone())?;
401 if result.len() == limit {
402 break;
403 }
404 }
405 Ok(result)
406 }
407}