1use std::fmt;
9use std::sync::{Arc, Mutex, MutexGuard};
10
11use crate::error::{Error, Result};
12use crate::kv::KVTransaction;
13use crate::txn::{OwnedLeaseOutcome, OwnedReadSessionStatus, OwnedTransactionSessionStatus};
14use crate::types::{Key, TxnId, TxnMode, Value};
15
16pub trait OwnedKVScan: Send {
22 fn next_entry(&mut self) -> Result<Option<(Key, Value)>>;
24
25 fn close(&mut self) -> Result<()> {
27 Ok(())
28 }
29}
30
31pub trait OwnedKVTransaction: Send {
36 fn id(&self) -> TxnId;
38
39 fn mode(&self) -> TxnMode;
41
42 fn get(&mut self, key: &Key) -> Result<Option<Value>>;
44
45 fn put(&mut self, key: Key, value: Value) -> Result<()>;
47
48 fn delete(&mut self, key: Key) -> Result<()>;
50
51 fn scan_prefix(&mut self, prefix: &[u8]) -> Result<Box<dyn OwnedKVScan>>;
53
54 fn scan_range(&mut self, start: &[u8], end: &[u8]) -> Result<Box<dyn OwnedKVScan>>;
56
57 fn commit(self: Box<Self>) -> Result<()>;
59
60 fn rollback(self: Box<Self>) -> Result<()>;
62}
63
64pub struct OwnedKVTransactionAdapter<'a> {
74 transaction: &'a mut dyn OwnedKVTransaction,
75}
76
77impl<'a> OwnedKVTransactionAdapter<'a> {
78 pub fn new(transaction: &'a mut dyn OwnedKVTransaction) -> Self {
80 Self { transaction }
81 }
82
83 fn collect_cursor(cursor: Box<dyn OwnedKVScan>) -> Result<Vec<(Key, Value)>> {
84 let mut cursor = cursor;
85 let mut entries = Vec::new();
86 while let Some(entry) = cursor.next_entry()? {
87 entries.push(entry);
88 }
89 cursor.close()?;
90 Ok(entries)
91 }
92}
93
94impl<'a> KVTransaction<'a> for OwnedKVTransactionAdapter<'a> {
95 fn id(&self) -> TxnId {
96 self.transaction.id()
97 }
98
99 fn mode(&self) -> TxnMode {
100 self.transaction.mode()
101 }
102
103 fn get(&mut self, key: &Key) -> Result<Option<Value>> {
104 self.transaction.get(key)
105 }
106
107 fn put(&mut self, key: Key, value: Value) -> Result<()> {
108 self.transaction.put(key, value)
109 }
110
111 fn delete(&mut self, key: Key) -> Result<()> {
112 self.transaction.delete(key)
113 }
114
115 fn scan_prefix(
116 &mut self,
117 prefix: &[u8],
118 ) -> Result<Box<dyn Iterator<Item = (Key, Value)> + '_>> {
119 let entries = Self::collect_cursor(self.transaction.scan_prefix(prefix)?)?;
120 Ok(Box::new(entries.into_iter()))
121 }
122
123 fn scan_range(
124 &mut self,
125 start: &[u8],
126 end: &[u8],
127 ) -> Result<Box<dyn Iterator<Item = (Key, Value)> + '_>> {
128 let entries = Self::collect_cursor(self.transaction.scan_range(start, end)?)?;
129 Ok(Box::new(entries.into_iter()))
130 }
131
132 fn commit_self(self) -> Result<()> {
133 Err(Error::InvalidParameter {
134 param: "owned_transaction_adapter".to_owned(),
135 reason: "terminal ownership belongs to OwnedTransactionSession".to_owned(),
136 })
137 }
138
139 fn rollback_self(self) -> Result<()> {
140 Err(Error::InvalidParameter {
141 param: "owned_transaction_adapter".to_owned(),
142 reason: "terminal ownership belongs to OwnedTransactionSession".to_owned(),
143 })
144 }
145}
146
147pub trait OwnedKVStore: Send + Sync + 'static {
152 fn begin_owned_kv_transaction(
154 self: Arc<Self>,
155 mode: TxnMode,
156 ) -> Result<Box<dyn OwnedKVTransaction>>;
157}
158
159#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
164#[non_exhaustive]
165pub struct OwnedReadOptions {}
166
167pub trait OwnedSessionFactory: OwnedKVStore {
169 fn begin_owned_read(self: Arc<Self>, _options: OwnedReadOptions) -> Result<OwnedReadSession> {
171 let transaction = self.begin_owned_kv_transaction(TxnMode::ReadOnly)?;
172 OwnedReadSession::new(transaction)
173 }
174
175 fn begin_owned_transaction(self: Arc<Self>, mode: TxnMode) -> Result<OwnedTransactionSession> {
177 let transaction = self.begin_owned_kv_transaction(mode)?;
178 Ok(OwnedTransactionSession::new(transaction))
179 }
180}
181
182impl<T> OwnedSessionFactory for T where T: OwnedKVStore {}
183
184pub trait OwnedReadSessionApi: Send + Sync {
186 fn status(&self) -> OwnedReadSessionStatus;
188
189 fn acquire_lease(&self) -> Result<OwnedReadLease>;
191
192 fn close(&self) -> Result<OwnedReadSessionStatus>;
194}
195
196pub trait OwnedTransactionSessionApi: Send + Sync {
198 fn status(&self) -> OwnedTransactionSessionStatus;
200
201 fn acquire_lease(&self) -> Result<OwnedTransactionLease>;
203
204 fn commit(&self) -> Result<OwnedTransactionSessionStatus>;
206
207 fn rollback(&self) -> Result<OwnedTransactionSessionStatus>;
209}
210
211#[derive(Clone)]
213pub struct OwnedReadSession {
214 inner: Arc<Mutex<OwnedReadInner>>,
215}
216
217struct OwnedReadInner {
218 transaction: Option<Box<dyn OwnedKVTransaction>>,
219 status: OwnedReadSessionStatus,
220}
221
222impl Drop for OwnedReadInner {
223 fn drop(&mut self) {
224 if let Some(transaction) = self.transaction.take() {
225 let _ = transaction.rollback();
226 }
227 }
228}
229
230impl OwnedReadSession {
231 pub fn new(transaction: Box<dyn OwnedKVTransaction>) -> Result<Self> {
233 if transaction.mode() != TxnMode::ReadOnly {
234 return Err(Error::InvalidParameter {
235 param: "owned_read_session".to_owned(),
236 reason: "requires a read-only transaction".to_owned(),
237 });
238 }
239 Ok(Self {
240 inner: Arc::new(Mutex::new(OwnedReadInner {
241 transaction: Some(transaction),
242 status: OwnedReadSessionStatus::Open,
243 })),
244 })
245 }
246
247 pub fn status(&self) -> OwnedReadSessionStatus {
249 self.lock().status
250 }
251
252 pub fn acquire_lease(&self) -> Result<OwnedReadLease> {
254 let mut inner = self.lock();
255 if inner.status != OwnedReadSessionStatus::Open || inner.transaction.is_none() {
256 return Err(Error::TxnClosed);
257 }
258 inner.status = OwnedReadSessionStatus::LeaseActive;
259 Ok(OwnedReadLease {
260 inner: self.inner.clone(),
261 released: false,
262 })
263 }
264
265 pub fn close(&self) -> Result<OwnedReadSessionStatus> {
267 finish_read_session(&self.inner, OwnedLeaseOutcome::Closed)
268 }
269
270 fn lock(&self) -> MutexGuard<'_, OwnedReadInner> {
271 self.inner
272 .lock()
273 .expect("owned read session mutex poisoned")
274 }
275}
276
277impl OwnedReadSessionApi for OwnedReadSession {
278 fn status(&self) -> OwnedReadSessionStatus {
279 Self::status(self)
280 }
281
282 fn acquire_lease(&self) -> Result<OwnedReadLease> {
283 Self::acquire_lease(self)
284 }
285
286 fn close(&self) -> Result<OwnedReadSessionStatus> {
287 Self::close(self)
288 }
289}
290
291impl fmt::Debug for OwnedReadSession {
292 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
293 formatter
294 .debug_struct("OwnedReadSession")
295 .field("status", &self.status())
296 .finish_non_exhaustive()
297 }
298}
299
300pub struct OwnedReadLease {
302 inner: Arc<Mutex<OwnedReadInner>>,
303 released: bool,
304}
305
306impl OwnedReadLease {
307 pub fn with_transaction<T>(
312 &self,
313 operation: impl FnOnce(&mut dyn OwnedKVTransaction) -> Result<T>,
314 ) -> Result<T> {
315 let mut inner = self
316 .inner
317 .lock()
318 .expect("owned read session mutex poisoned");
319 if inner.status != OwnedReadSessionStatus::LeaseActive {
320 return Err(Error::TxnClosed);
321 }
322 operation(inner.transaction.as_deref_mut().ok_or(Error::TxnClosed)?)
323 }
324
325 pub fn finish(mut self, outcome: OwnedLeaseOutcome) -> Result<OwnedReadSessionStatus> {
327 let result = finish_read_session(&self.inner, outcome);
328 self.released = true;
329 result
330 }
331}
332
333impl Drop for OwnedReadLease {
334 fn drop(&mut self) {
335 if !self.released {
336 let _ = finish_read_session(&self.inner, OwnedLeaseOutcome::Closed);
337 self.released = true;
338 }
339 }
340}
341
342fn finish_read_session(
343 inner: &Arc<Mutex<OwnedReadInner>>,
344 outcome: OwnedLeaseOutcome,
345) -> Result<OwnedReadSessionStatus> {
346 let (transaction, target) = {
347 let mut inner = inner.lock().expect("owned read session mutex poisoned");
348 if inner.status.is_terminal() {
349 return Ok(inner.status);
350 }
351 let target = OwnedReadSessionStatus::from(outcome);
352 let transaction = inner.transaction.take();
353 inner.status = target;
354 (transaction, target)
355 };
356
357 if let Some(transaction) = transaction {
358 if transaction.rollback().is_err() {
359 let mut inner = inner.lock().expect("owned read session mutex poisoned");
360 inner.status = OwnedReadSessionStatus::Failed;
361 return Err(Error::TxnClosed);
362 }
363 }
364 Ok(target)
365}
366
367#[derive(Clone)]
369pub struct OwnedTransactionSession {
370 inner: Arc<Mutex<OwnedTransactionInner>>,
371}
372
373struct OwnedTransactionInner {
374 transaction: Option<Box<dyn OwnedKVTransaction>>,
375 status: OwnedTransactionSessionStatus,
376}
377
378impl Drop for OwnedTransactionInner {
379 fn drop(&mut self) {
380 if let Some(transaction) = self.transaction.take() {
381 let _ = transaction.rollback();
382 }
383 }
384}
385
386impl OwnedTransactionSession {
387 pub fn new(transaction: Box<dyn OwnedKVTransaction>) -> Self {
389 Self {
390 inner: Arc::new(Mutex::new(OwnedTransactionInner {
391 transaction: Some(transaction),
392 status: OwnedTransactionSessionStatus::Open,
393 })),
394 }
395 }
396
397 pub fn status(&self) -> OwnedTransactionSessionStatus {
399 self.lock().status
400 }
401
402 pub fn acquire_lease(&self) -> Result<OwnedTransactionLease> {
404 let mut inner = self.lock();
405 if !inner.status.can_acquire_lease() || inner.transaction.is_none() {
406 return Err(Error::TxnClosed);
407 }
408 inner.status = OwnedTransactionSessionStatus::LeaseActive;
409 Ok(OwnedTransactionLease {
410 inner: self.inner.clone(),
411 released: false,
412 })
413 }
414
415 pub fn with_transaction<T>(
422 &self,
423 operation: impl FnOnce(&mut dyn OwnedKVTransaction) -> Result<T>,
424 ) -> Result<T> {
425 let lease = self.acquire_lease()?;
426 let result = lease.with_transaction(operation);
427 let release = lease.finish(OwnedLeaseOutcome::Exhausted);
428 match (result, release) {
429 (Ok(value), Ok(_)) => Ok(value),
430 (Err(error), Ok(_)) => Err(error),
431 (_, Err(error)) => Err(error),
432 }
433 }
434
435 pub fn commit(&self) -> Result<OwnedTransactionSessionStatus> {
437 self.finish_transaction(true)
438 }
439
440 pub fn rollback(&self) -> Result<OwnedTransactionSessionStatus> {
442 self.finish_transaction(false)
443 }
444
445 fn finish_transaction(&self, commit: bool) -> Result<OwnedTransactionSessionStatus> {
446 let transaction = {
447 let mut inner = self.lock();
448 let allowed = if commit {
449 inner.status.can_commit()
450 } else {
451 inner.status.can_rollback()
452 };
453 if !allowed {
454 return Err(Error::TxnClosed);
455 }
456 inner.status = OwnedTransactionSessionStatus::Closed;
457 inner.transaction.take().ok_or(Error::TxnClosed)?
458 };
459
460 let result = if commit {
461 transaction.commit()
462 } else {
463 transaction.rollback()
464 };
465 let mut inner = self.lock();
466 inner.status = if result.is_ok() {
467 if commit {
468 OwnedTransactionSessionStatus::Committed
469 } else {
470 OwnedTransactionSessionStatus::RolledBack
471 }
472 } else {
473 OwnedTransactionSessionStatus::Closed
474 };
475 match result {
476 Ok(()) => Ok(inner.status),
477 Err(error) => Err(error),
478 }
479 }
480
481 fn lock(&self) -> MutexGuard<'_, OwnedTransactionInner> {
482 self.inner
483 .lock()
484 .expect("owned transaction session mutex poisoned")
485 }
486}
487
488impl OwnedTransactionSessionApi for OwnedTransactionSession {
489 fn status(&self) -> OwnedTransactionSessionStatus {
490 Self::status(self)
491 }
492
493 fn acquire_lease(&self) -> Result<OwnedTransactionLease> {
494 Self::acquire_lease(self)
495 }
496
497 fn commit(&self) -> Result<OwnedTransactionSessionStatus> {
498 Self::commit(self)
499 }
500
501 fn rollback(&self) -> Result<OwnedTransactionSessionStatus> {
502 Self::rollback(self)
503 }
504}
505
506impl fmt::Debug for OwnedTransactionSession {
507 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
508 formatter
509 .debug_struct("OwnedTransactionSession")
510 .field("status", &self.status())
511 .finish_non_exhaustive()
512 }
513}
514
515pub struct OwnedTransactionLease {
517 inner: Arc<Mutex<OwnedTransactionInner>>,
518 released: bool,
519}
520
521impl OwnedTransactionLease {
522 pub fn with_transaction<T>(
524 &self,
525 operation: impl FnOnce(&mut dyn OwnedKVTransaction) -> Result<T>,
526 ) -> Result<T> {
527 let mut inner = self
528 .inner
529 .lock()
530 .expect("owned transaction session mutex poisoned");
531 if inner.status != OwnedTransactionSessionStatus::LeaseActive {
532 return Err(Error::TxnClosed);
533 }
534 operation(inner.transaction.as_deref_mut().ok_or(Error::TxnClosed)?)
535 }
536
537 pub fn finish(mut self, outcome: OwnedLeaseOutcome) -> Result<OwnedTransactionSessionStatus> {
539 let mut inner = self
540 .inner
541 .lock()
542 .expect("owned transaction session mutex poisoned");
543 if inner.status != OwnedTransactionSessionStatus::LeaseActive {
544 self.released = true;
545 return Err(Error::TxnClosed);
546 }
547 inner.status = OwnedTransactionSessionStatus::from(outcome);
548 self.released = true;
549 Ok(inner.status)
550 }
551}
552
553impl Drop for OwnedTransactionLease {
554 fn drop(&mut self) {
555 if !self.released {
556 let mut inner = self
557 .inner
558 .lock()
559 .expect("owned transaction session mutex poisoned");
560 if inner.status == OwnedTransactionSessionStatus::LeaseActive {
561 inner.status = OwnedTransactionSessionStatus::MustAbort;
562 }
563 self.released = true;
564 }
565 }
566}
567
568#[cfg(test)]
569mod tests {
570 use std::sync::atomic::{AtomicUsize, Ordering};
571 use std::sync::Arc;
572
573 use super::{
574 OwnedKVScan, OwnedKVTransaction, OwnedKVTransactionAdapter, OwnedReadSession,
575 OwnedSessionFactory, OwnedTransactionSession,
576 };
577 use crate::error::Result;
578 use crate::kv::{memory::MemoryKV, KVTransaction};
579 use crate::txn::{OwnedLeaseOutcome, OwnedReadSessionStatus, OwnedTransactionSessionStatus};
580 use crate::types::{Key, TxnId, TxnMode, Value};
581
582 #[derive(Default)]
583 struct Calls {
584 commits: AtomicUsize,
585 rollbacks: AtomicUsize,
586 }
587
588 struct TestCursor;
589
590 impl OwnedKVScan for TestCursor {
591 fn next_entry(&mut self) -> Result<Option<(Key, Value)>> {
592 Ok(None)
593 }
594 }
595
596 struct TestTransaction {
597 calls: Arc<Calls>,
598 mode: TxnMode,
599 }
600
601 impl TestTransaction {
602 fn new(calls: Arc<Calls>, mode: TxnMode) -> Self {
603 Self { calls, mode }
604 }
605 }
606
607 impl OwnedKVTransaction for TestTransaction {
608 fn id(&self) -> TxnId {
609 TxnId(7)
610 }
611
612 fn mode(&self) -> TxnMode {
613 self.mode
614 }
615
616 fn get(&mut self, _key: &Key) -> Result<Option<Value>> {
617 Ok(None)
618 }
619
620 fn put(&mut self, _key: Key, _value: Value) -> Result<()> {
621 Ok(())
622 }
623
624 fn delete(&mut self, _key: Key) -> Result<()> {
625 Ok(())
626 }
627
628 fn scan_prefix(&mut self, _prefix: &[u8]) -> Result<Box<dyn OwnedKVScan>> {
629 Ok(Box::new(TestCursor))
630 }
631
632 fn scan_range(&mut self, _start: &[u8], _end: &[u8]) -> Result<Box<dyn OwnedKVScan>> {
633 Ok(Box::new(TestCursor))
634 }
635
636 fn commit(self: Box<Self>) -> Result<()> {
637 self.calls.commits.fetch_add(1, Ordering::SeqCst);
638 Ok(())
639 }
640
641 fn rollback(self: Box<Self>) -> Result<()> {
642 self.calls.rollbacks.fetch_add(1, Ordering::SeqCst);
643 Ok(())
644 }
645 }
646
647 #[test]
648 fn read_session_releases_its_transaction_once_at_terminal_outcome() {
649 let calls = Arc::new(Calls::default());
650 let session = OwnedReadSession::new(Box::new(TestTransaction::new(
651 calls.clone(),
652 TxnMode::ReadOnly,
653 )))
654 .unwrap();
655
656 let lease = session.acquire_lease().unwrap();
657 assert_eq!(session.status(), OwnedReadSessionStatus::LeaseActive);
658 assert!(session.acquire_lease().is_err());
659 assert_eq!(
660 lease.finish(OwnedLeaseOutcome::Exhausted).unwrap(),
661 OwnedReadSessionStatus::Exhausted
662 );
663 assert_eq!(calls.rollbacks.load(Ordering::SeqCst), 1);
664 assert_eq!(session.close().unwrap(), OwnedReadSessionStatus::Exhausted);
665 assert_eq!(calls.rollbacks.load(Ordering::SeqCst), 1);
666 }
667
668 #[test]
669 fn transaction_session_blocks_commit_while_leased_then_commits_once() {
670 let calls = Arc::new(Calls::default());
671 let session = OwnedTransactionSession::new(Box::new(TestTransaction::new(
672 calls.clone(),
673 TxnMode::ReadWrite,
674 )));
675
676 let lease = session.acquire_lease().unwrap();
677 assert!(session.commit().is_err());
678 assert_eq!(
679 lease.finish(OwnedLeaseOutcome::Exhausted).unwrap(),
680 OwnedTransactionSessionStatus::Committable
681 );
682 assert_eq!(
683 session.commit().unwrap(),
684 OwnedTransactionSessionStatus::Committed
685 );
686 assert_eq!(calls.commits.load(Ordering::SeqCst), 1);
687 assert!(session.rollback().is_err());
688 assert_eq!(calls.rollbacks.load(Ordering::SeqCst), 0);
689 }
690
691 #[test]
692 fn cancelled_or_dropped_lease_requires_a_single_rollback() {
693 let calls = Arc::new(Calls::default());
694 let session = OwnedTransactionSession::new(Box::new(TestTransaction::new(
695 calls.clone(),
696 TxnMode::ReadWrite,
697 )));
698
699 let lease = session.acquire_lease().unwrap();
700 assert_eq!(
701 lease.finish(OwnedLeaseOutcome::Cancelled).unwrap(),
702 OwnedTransactionSessionStatus::MustAbort
703 );
704 assert!(session.commit().is_err());
705 assert_eq!(
706 session.rollback().unwrap(),
707 OwnedTransactionSessionStatus::RolledBack
708 );
709 assert_eq!(calls.rollbacks.load(Ordering::SeqCst), 1);
710
711 let drop_calls = Arc::new(Calls::default());
712 let dropped = OwnedTransactionSession::new(Box::new(TestTransaction::new(
713 drop_calls.clone(),
714 TxnMode::ReadWrite,
715 )));
716 drop(dropped.acquire_lease().unwrap());
717 assert_eq!(dropped.status(), OwnedTransactionSessionStatus::MustAbort);
718 drop(dropped);
719 assert_eq!(drop_calls.rollbacks.load(Ordering::SeqCst), 1);
720 }
721
722 #[test]
723 fn compatibility_adapter_never_consumes_the_owned_terminal_transition() {
724 let store = Arc::new(MemoryKV::new());
725 let session = store
726 .clone()
727 .begin_owned_transaction(TxnMode::ReadWrite)
728 .unwrap();
729
730 session
731 .with_transaction(|owned| {
732 let mut compatibility = OwnedKVTransactionAdapter::new(owned);
733 compatibility.put(b"adapter".to_vec(), b"value".to_vec())?;
734 let rows = compatibility.scan_prefix(b"adapter")?.collect::<Vec<_>>();
735 assert_eq!(rows, vec![(b"adapter".to_vec(), b"value".to_vec())]);
736 assert!(compatibility.commit_self().is_err());
737 Ok(())
738 })
739 .unwrap();
740
741 assert_eq!(session.status(), OwnedTransactionSessionStatus::Committable);
742 session.commit().unwrap();
743
744 let read = store.begin_owned_read(Default::default()).unwrap();
745 let lease = read.acquire_lease().unwrap();
746 let value = lease
747 .with_transaction(|transaction| transaction.get(&b"adapter".to_vec()))
748 .unwrap();
749 assert_eq!(value, Some(b"value".to_vec()));
750 lease.finish(OwnedLeaseOutcome::Exhausted).unwrap();
751 }
752}