@lix-js/sdk 0.6.0-preview.3 → 0.6.0-preview.5

package/dist-engine-src/src/storage/context.rs

@@ -1,356 +1,620 @@
-use std::sync::Arc;
+use std::ops::Bound;
 
-use async_trait::async_trait;
-
-use crate::backend::{Backend, BackendReadTransaction, BackendWriteTransaction};
-use crate::storage::types::{KvWriteBatch, StorageWriter};
+use crate::backend::{
+    Backend, BackendError, BackendWrite, CommitResult, InMemoryBackend, KeyRange, Prefix,
+    ReadOptions, WriteOptions,
+};
 use crate::storage::{
-    KvEntryPage, KvExistsBatch, KvGetRequest, KvKeyPage, KvScanRequest, KvValueBatch, KvValuePage,
-    KvWriteStats, StorageReadTransaction, StorageReader, StorageWriteTransaction,
+    StorageRead, StorageReadScope, StorageSpace, StorageWriteSet, StorageWriteSetError,
+    StorageWriteSetStats,
 };
-use crate::LixError;
 
-#[derive(Clone)]
-pub(crate) struct StorageContext {
-    backend: Arc<dyn Backend + Send + Sync>,
+#[derive(Clone, Debug)]
+pub struct StorageContext<B = InMemoryBackend> {
+    backend: B,
 }
 
-impl StorageContext {
-    pub(crate) fn new(backend: Arc<dyn Backend + Send + Sync>) -> Self {
+impl<B> StorageContext<B>
+where
+    B: Backend,
+{
+    pub fn new(backend: B) -> Self {
         Self { backend }
     }
 
-    pub(crate) async fn begin_read_transaction(
+    pub fn begin_read(
         &self,
-    ) -> Result<Box<dyn StorageReadTransaction + Send + Sync + 'static>, LixError> {
-        let transaction = self.backend.begin_read_transaction().await?;
-        Ok(Box::new(StorageContextReadTransaction { transaction }))
+        opts: ReadOptions,
+    ) -> Result<StorageReadScope<B::Read<'_>>, BackendError> {
+        self.backend.begin_read(opts).map(StorageReadScope::new)
     }
 
-    pub(crate) async fn begin_write_transaction(
-        &self,
-    ) -> Result<Box<dyn StorageWriteTransaction + Send + Sync + 'static>, LixError> {
-        let transaction = self.backend.begin_write_transaction().await?;
-        Ok(Box::new(StorageContextWriteTransaction { transaction }))
+    pub fn new_write_set(&self) -> StorageWriteSet {
+        StorageWriteSet::new()
     }
 
-    pub(crate) async fn close(&self) -> Result<(), LixError> {
-        self.backend.close().await
+    pub async fn begin_read_transaction(
+        &self,
+    ) -> Result<Box<StorageReadTransaction<B::Read<'_>>>, crate::LixError> {
+        Ok(Box::new(StorageReadTransaction {
+            read: self.begin_read(ReadOptions::default())?,
+        }))
     }
 
-    pub(crate) async fn destroy(&self) -> Result<(), LixError> {
-        self.backend.destroy().await
+    pub async fn begin_write_transaction(
+        &self,
+    ) -> Result<Box<StorageWriteTransaction<'_, B>>, crate::LixError> {
+        Ok(Box::new(StorageWriteTransaction {
+            storage: self,
+            read: self.begin_read(ReadOptions::default())?,
+        }))
     }
-}
 
-#[cfg(any(test, feature = "storage-benches"))]
-#[async_trait]
-impl StorageReader for StorageContext {
-    async fn get_values(&mut self, request: KvGetRequest) -> Result<KvValueBatch, LixError> {
-        let mut transaction = self.begin_read_transaction().await?;
-        let result = transaction.get_values(request).await;
-        match result {
-            Ok(result) => {
-                transaction.rollback().await?;
-                Ok(result)
-            }
-            Err(error) => {
-                let _ = transaction.rollback().await;
-                Err(error)
-            }
-        }
+    pub fn commit_write_set(
+        &self,
+        write_set: StorageWriteSet,
+        opts: WriteOptions,
+    ) -> Result<(CommitResult, StorageWriteSetStats), StorageWriteSetError> {
+        write_set.commit(&self.backend, opts)
     }
 
-    async fn exists_many(&mut self, request: KvGetRequest) -> Result<KvExistsBatch, LixError> {
-        let mut transaction = self.begin_read_transaction().await?;
-        let result = transaction.exists_many(request).await;
-        match result {
-            Ok(result) => {
-                transaction.rollback().await?;
-                Ok(result)
-            }
-            Err(error) => {
-                let _ = transaction.rollback().await;
-                Err(error)
-            }
-        }
+    pub fn durable_write_lock(&self) -> crate::backend::DurableWriteLock {
+        self.backend.durable_write_lock()
     }
 
-    async fn scan_keys(&mut self, request: KvScanRequest) -> Result<KvKeyPage, LixError> {
-        let mut transaction = self.begin_read_transaction().await?;
-        let result = transaction.scan_keys(request).await;
-        match result {
-            Ok(result) => {
-                transaction.rollback().await?;
-                Ok(result)
-            }
-            Err(error) => {
-                let _ = transaction.rollback().await;
-                Err(error)
-            }
+    pub fn delete_range(
+        &self,
+        space: StorageSpace,
+        range: KeyRange,
+        opts: WriteOptions,
+    ) -> Result<CommitResult, BackendError> {
+        let mut write = self.backend.begin_write(opts)?;
+        let physical_range = space.encode_range(range, None);
+        if let Err(error) = write.delete_range(physical_range) {
+            let _ = write.rollback();
+            return Err(error);
         }
+        write.commit()
     }
 
-    async fn scan_values(&mut self, request: KvScanRequest) -> Result<KvValuePage, LixError> {
-        let mut transaction = self.begin_read_transaction().await?;
-        let result = transaction.scan_values(request).await;
-        match result {
-            Ok(result) => {
-                transaction.rollback().await?;
-                Ok(result)
-            }
-            Err(error) => {
-                let _ = transaction.rollback().await;
-                Err(error)
-            }
-        }
+    pub fn delete_prefix(
+        &self,
+        space: StorageSpace,
+        prefix: Prefix,
+        opts: WriteOptions,
+    ) -> Result<CommitResult, BackendError> {
+        self.delete_range(space, prefix.to_range()?, opts)
     }
 
-    async fn scan_entries(&mut self, request: KvScanRequest) -> Result<KvEntryPage, LixError> {
-        let mut transaction = self.begin_read_transaction().await?;
-        let result = transaction.scan_entries(request).await;
-        match result {
-            Ok(result) => {
-                transaction.rollback().await?;
-                Ok(result)
-            }
-            Err(error) => {
-                let _ = transaction.rollback().await;
-                Err(error)
-            }
-        }
+    pub fn clear_space(
+        &self,
+        space: StorageSpace,
+        opts: WriteOptions,
+    ) -> Result<CommitResult, BackendError> {
+        self.delete_range(
+            space,
+            KeyRange {
+                lower: Bound::Unbounded,
+                upper: Bound::Unbounded,
+            },
+            opts,
+        )
     }
 }
 
-struct StorageContextReadTransaction {
-    transaction: Box<dyn BackendReadTransaction + Send + Sync + 'static>,
+pub struct StorageReadTransaction<R>
+where
+    R: crate::backend::BackendRead,
+{
+    read: StorageReadScope<R>,
 }
 
-struct StorageContextWriteTransaction {
-    transaction: Box<dyn BackendWriteTransaction + Send + Sync + 'static>,
+impl<R> StorageReadTransaction<R>
+where
+    R: crate::backend::BackendRead,
+{
+    pub async fn rollback(self: Box<Self>) -> Result<(), crate::LixError> {
+        Ok(())
+    }
 }
 
-#[async_trait]
-impl StorageReader for StorageContextReadTransaction {
-    async fn get_values(&mut self, request: KvGetRequest) -> Result<KvValueBatch, LixError> {
-        self.transaction
-            .get_values(request.into())
-            .await
-            .map(Into::into)
-    }
+impl<R> StorageRead for StorageReadTransaction<R>
+where
+    R: crate::backend::BackendRead,
+{
+    type BackendRead = R;
 
-    async fn exists_many(&mut self, request: KvGetRequest) -> Result<KvExistsBatch, LixError> {
-        self.transaction
-            .exists_many(request.into())
-            .await
-            .map(Into::into)
+    fn backend_read(&self) -> &Self::BackendRead {
+        self.read.backend_read()
     }
+}
+
+pub struct StorageWriteTransaction<'a, B>
+where
+    B: Backend,
+{
+    storage: &'a StorageContext<B>,
+    read: StorageReadScope<B::Read<'a>>,
+}
 
-    async fn scan_keys(&mut self, request: KvScanRequest) -> Result<KvKeyPage, LixError> {
-        self.transaction
-            .scan_keys(request.into())
-            .await
-            .map(Into::into)
+impl<B> StorageWriteTransaction<'_, B>
+where
+    B: Backend,
+{
+    pub async fn commit(self: Box<Self>) -> Result<(), crate::LixError> {
+        Ok(())
     }
 
-    async fn scan_values(&mut self, request: KvScanRequest) -> Result<KvValuePage, LixError> {
-        self.transaction
-            .scan_values(request.into())
-            .await
-            .map(Into::into)
+    pub async fn rollback(self: Box<Self>) -> Result<(), crate::LixError> {
+        Ok(())
     }
 
-    async fn scan_entries(&mut self, request: KvScanRequest) -> Result<KvEntryPage, LixError> {
-        self.transaction
-            .scan_entries(request.into())
-            .await
-            .map(Into::into)
+    pub fn write_set(
+        &mut self,
+        write_set: StorageWriteSet,
+    ) -> Result<StorageWriteSetStats, crate::LixError> {
+        let (_commit, stats) = self
+            .storage
+            .commit_write_set(write_set, WriteOptions::default())?;
+        Ok(stats)
     }
 }
 
-#[async_trait]
-impl StorageReadTransaction for StorageContextReadTransaction {
-    async fn rollback(self: Box<Self>) -> Result<(), LixError> {
-        self.transaction.rollback().await
+impl<'a, B> StorageRead for StorageWriteTransaction<'a, B>
+where
+    B: Backend,
+{
+    type BackendRead = B::Read<'a>;
+
+    fn backend_read(&self) -> &Self::BackendRead {
+        self.read.backend_read()
     }
 }
 
-#[async_trait]
-impl StorageReader for StorageContextWriteTransaction {
-    async fn get_values(&mut self, request: KvGetRequest) -> Result<KvValueBatch, LixError> {
-        self.transaction
-            .get_values(request.into())
-            .await
-            .map(Into::into)
-    }
+#[cfg(test)]
+mod tests {
+    use bytes::Bytes;
+
+    use crate::backend::{
+        DurableWriteLock, GetOptions, InMemoryBackend, Key, ProjectedValue, ReadOptions, SpaceId,
+        StoredValue, WriteOptions,
+    };
+    use crate::storage::{PointReadPlan, StorageContext, StorageSpace};
 
-    async fn exists_many(&mut self, request: KvGetRequest) -> Result<KvExistsBatch, LixError> {
-        self.transaction
-            .exists_many(request.into())
-            .await
-            .map(Into::into)
+    fn key(bytes: &'static str) -> Key {
+        Key(Bytes::from_static(bytes.as_bytes()))
     }
 
-    async fn scan_keys(&mut self, request: KvScanRequest) -> Result<KvKeyPage, LixError> {
-        self.transaction
-            .scan_keys(request.into())
-            .await
-            .map(Into::into)
+    fn value(bytes: &'static str) -> StoredValue {
+        StoredValue {
+            bytes: Bytes::from_static(bytes.as_bytes()),
+        }
     }
 
-    async fn scan_values(&mut self, request: KvScanRequest) -> Result<KvValuePage, LixError> {
-        self.transaction
-            .scan_values(request.into())
-            .await
-            .map(Into::into)
+    fn space(id: u32) -> StorageSpace {
+        match id {
+            1 => StorageSpace::new(SpaceId(1), "test.space.one"),
+            2 => StorageSpace::new(SpaceId(2), "test.space.two"),
+            _ => StorageSpace::new(SpaceId(id), "test.space.other"),
+        }
     }
 
-    async fn scan_entries(&mut self, request: KvScanRequest) -> Result<KvEntryPage, LixError> {
-        self.transaction
-            .scan_entries(request.into())
-            .await
-            .map(Into::into)
+    #[test]
+    fn context_commits_write_set_and_reads_through_storage_contract() {
+        let storage = StorageContext::new(InMemoryBackend::new());
+        let mut writes = storage.new_write_set();
+        writes.put(space(1), key("a"), value("A"));
+        writes.put(space(1), key("b"), value("B"));
+        writes.put(space(2), key("a"), value("other"));
+        writes.delete(space(2), key("missing"));
+
+        let (_commit, stats) = storage
+            .commit_write_set(writes, WriteOptions::default())
+            .expect("commit write set");
+
+        assert_eq!(stats.staged_puts, 3);
+        assert_eq!(stats.staged_deletes, 1);
+        assert_eq!(stats.touched_spaces, 2);
+        assert_eq!(stats.put_batches, 2);
+        assert_eq!(stats.delete_batches, 1);
+        assert_eq!(stats.backend_calls, 3);
+
+        let read = storage
+            .begin_read(ReadOptions::default())
+            .expect("begin read");
+        let result = PointReadPlan::new(space(1), &[key("a"), key("b")])
+            .materialize(&read, GetOptions::default())
+            .expect("read back values");
+        assert_eq!(
+            result.value,
+            vec![
+                Some(ProjectedValue::FullValue(Bytes::from_static(b"A"))),
+                Some(ProjectedValue::FullValue(Bytes::from_static(b"B"))),
+            ]
+        );
     }
-}
 
-#[async_trait]
-impl StorageWriter for StorageContextWriteTransaction {
-    async fn write_kv_batch(&mut self, batch: KvWriteBatch) -> Result<KvWriteStats, LixError> {
-        self.transaction
-            .write_kv_batch(batch.into())
-            .await
-            .map(Into::into)
+    #[test]
+    fn context_read_scope_pins_snapshot_across_later_commits() {
+        let storage = StorageContext::new(InMemoryBackend::new());
+        let mut writes = storage.new_write_set();
+        writes.put(space(1), key("a"), value("A"));
+        storage
+            .commit_write_set(writes, WriteOptions::default())
+            .expect("seed");
+
+        let read = storage
+            .begin_read(ReadOptions::default())
+            .expect("begin read");
+
+        let mut later = storage.new_write_set();
+        later.put(space(1), key("a"), value("B"));
+        storage
+            .commit_write_set(later, WriteOptions::default())
+            .expect("later commit");
+
+        let result = PointReadPlan::new(space(1), &[key("a")])
+            .materialize(&read, GetOptions::default())
+            .expect("read old scope");
+
+        assert_eq!(
+            result.value[0],
+            Some(ProjectedValue::FullValue(Bytes::from_static(b"A")))
+        );
     }
 }
 
-#[async_trait]
-impl StorageReadTransaction for StorageContextWriteTransaction {
-    async fn rollback(self: Box<Self>) -> Result<(), LixError> {
-        self.transaction.rollback().await
+#[cfg(test)]
+mod shape_tests {
+    use std::cell::{Cell, RefCell};
+    use std::ops::Bound;
+    use std::rc::Rc;
+
+    use bytes::Bytes;
+
+    use crate::backend::{
+        Backend, BackendCapabilities, BackendError, BackendRangeScan, BackendRead, BackendWrite,
+        BufferedRangeScan, CommitResult, DurableWriteLock, GetOptions, Key, KeyRange, PointVisitor,
+        ProjectedValue, ProjectedValueRef, PutBatch, ReadOptions, ScanOptions, ScanResult,
+        ScanVisitor, SpaceId, StoredValue, WriteConcurrency, WriteOptions, WriteStats,
+    };
+    use crate::storage::{PointReadPlan, ScanPlan, StorageContext, StorageReadScope, StorageSpace};
+
+    #[test]
+    fn write_set_across_g_spaces_lowers_to_g_put_many_calls_and_one_commit() {
+        let backend = CountingBackend::default();
+        let storage = StorageContext::new(backend.clone());
+        let mut writes = storage.new_write_set();
+        writes.put(space(1), key("a"), value("A"));
+        writes.put(space(2), key("a"), value("B"));
+        writes.put(space(3), key("a"), value("C"));
+
+        storage
+            .commit_write_set(writes, WriteOptions::default())
+            .expect("commit write set");
+
+        assert_eq!(backend.state.begin_write_calls.get(), 1);
+        assert_eq!(backend.state.commit_calls.get(), 1);
+        assert_eq!(backend.state.put_batches.borrow().len(), 3);
+        assert_eq!(
+            backend
+                .state
+                .put_batches
+                .borrow()
+                .iter()
+                .map(|(space, keys)| (*space, keys.clone()))
+                .collect::<Vec<_>>(),
+            vec![
+                (SpaceId(1), vec![key("a")]),
+                (SpaceId(2), vec![key("a")]),
+                (SpaceId(3), vec![key("a")]),
+            ]
+        );
     }
-}
 
-#[async_trait]
-impl StorageWriteTransaction for StorageContextWriteTransaction {
-    async fn commit(self: Box<Self>) -> Result<(), LixError> {
-        self.transaction.commit().await
+    #[test]
+    fn point_read_m_requested_u_unique_sends_u_backend_keys() {
+        let read = SpyRead::default();
+        let scope = StorageReadScope::new(read.clone());
+
+        let result = PointReadPlan::new(
+            space(1),
+            &[key("b"), key("a"), key("b"), key("missing"), key("a")],
+        )
+        .materialize(&scope, GetOptions::default())
+        .expect("point read");
+
+        assert_eq!(
+            read.get_many_keys.borrow().clone(),
+            vec![key("b"), key("a"), key("missing")]
+                .into_iter()
+                .map(|key| space(1).encode_key(&key))
+                .collect::<Vec<_>>()
+        );
+        assert_eq!(
+            result.value,
+            vec![
+                Some(ProjectedValue::FullValue(space(1).encode_key(&key("b")).0)),
+                Some(ProjectedValue::FullValue(space(1).encode_key(&key("a")).0)),
+                Some(ProjectedValue::FullValue(space(1).encode_key(&key("b")).0)),
+                Some(ProjectedValue::FullValue(
+                    space(1).encode_key(&key("missing")).0
+                )),
+                Some(ProjectedValue::FullValue(space(1).encode_key(&key("a")).0)),
+            ]
+        );
     }
-}
 
-#[cfg(test)]
-mod tests {
-    use std::sync::Arc;
-
-    use crate::backend::testing::UnitTestBackend;
-    use crate::storage::types::KvWriteBatch;
-    use crate::storage::{KvGetGroup, KvScanRange, StorageWriteSet};
-
-    use super::*;
-
-    #[tokio::test]
-    async fn storage_context_roundtrips_batched_writes_and_reads() {
-        let backend: Arc<dyn Backend + Send + Sync> = Arc::new(UnitTestBackend::new());
-        let storage = StorageContext::new(backend);
-        let mut tx = storage
-            .begin_write_transaction()
-            .await
-            .expect("transaction opens");
-
-        let mut batch = KvWriteBatch::new();
-        batch.put("ns", b"a".to_vec(), b"1".to_vec());
-        batch.put("ns", b"b".to_vec(), b"2".to_vec());
-        let stats = tx.write_kv_batch(batch).await.expect("batch writes");
-        assert_eq!(stats.puts, 2);
-        tx.commit().await.expect("commit succeeds");
-
-        let mut tx = storage
-            .begin_read_transaction()
-            .await
-            .expect("read transaction opens");
-        let result = tx
-            .get_values(KvGetRequest {
-                groups: vec![KvGetGroup {
-                    namespace: "ns".to_string(),
-                    keys: vec![b"a".to_vec(), b"b".to_vec()],
-                }],
-            })
-            .await
-            .expect("batch reads");
-        assert_eq!(result.groups[0].value(0), Some(Some(b"1".as_slice())));
-        assert_eq!(result.groups[0].value(1), Some(Some(b"2".as_slice())));
-
-        let exists = tx
-            .exists_many(KvGetRequest {
-                groups: vec![KvGetGroup {
-                    namespace: "ns".to_string(),
-                    keys: vec![b"a".to_vec(), b"missing".to_vec()],
-                }],
+    #[test]
+    fn prefix_scan_calls_scan_range_once() {
+        let read = SpyRead::default();
+        let scope = StorageReadScope::new(read.clone());
+
+        ScanPlan::prefix(
+            space(1),
+            crate::backend::Prefix {
+                bytes: Bytes::from_static(b"a"),
+            },
+        )
+        .collect(&scope, ScanOptions::default())
+        .expect("prefix scan");
+
+        assert_eq!(read.scan_range_calls.get(), 1);
+        assert_eq!(
+            read.scan_range.borrow().clone(),
+            Some(KeyRange {
+                lower: Bound::Included(space(1).encode_key(&key("a"))),
+                upper: Bound::Excluded(space(1).encode_key(&key("b"))),
             })
-            .await
-            .expect("existence reads");
-        assert_eq!(exists.groups[0].exists, vec![true, false]);
-
-        let result = tx
-            .scan_entries(KvScanRequest {
-                namespace: "ns".to_string(),
-                range: KvScanRange::prefix(Vec::new()),
-                after: Some(b"a".to_vec()),
-                limit: 1,
-            })
-            .await
-            .expect("scan reads");
-        assert_eq!(result.key(0).expect("key exists"), b"b");
-        assert_eq!(result.value(0).expect("value exists"), b"2");
-
-        let key_only = tx
-            .scan_keys(KvScanRequest {
-                namespace: "ns".to_string(),
-                range: KvScanRange::prefix(Vec::new()),
-                after: None,
-                limit: 2,
+        );
+    }
+
+    #[test]
+    fn delete_range_lowers_to_one_backend_delete_range() {
+        let backend = CountingBackend::default();
+        let storage = StorageContext::new(backend.clone());
+
+        storage
+            .delete_range(
+                space(7),
+                KeyRange {
+                    lower: Bound::Included(key("a")),
+                    upper: Bound::Excluded(key("c")),
+                },
+                WriteOptions::default(),
+            )
+            .expect("delete range");
+
+        assert_eq!(backend.state.begin_write_calls.get(), 1);
+        assert_eq!(backend.state.commit_calls.get(), 1);
+        assert_eq!(backend.state.delete_many_calls.get(), 0);
+        assert_eq!(
+            backend.state.delete_ranges.borrow().as_slice(),
+            &[KeyRange {
+                lower: Bound::Included(space(7).encode_key(&key("a"))),
+                upper: Bound::Excluded(space(7).encode_key(&key("c"))),
+            }]
+        );
+    }
+
+    #[test]
+    fn delete_prefix_lowers_to_one_backend_delete_range() {
+        let backend = CountingBackend::default();
+        let storage = StorageContext::new(backend.clone());
+
+        storage
+            .delete_prefix(
+                space(7),
+                crate::backend::Prefix {
+                    bytes: Bytes::from_static(b"ab"),
+                },
+                WriteOptions::default(),
+            )
+            .expect("delete prefix");
+
+        assert_eq!(backend.state.begin_write_calls.get(), 1);
+        assert_eq!(backend.state.commit_calls.get(), 1);
+        assert_eq!(backend.state.delete_many_calls.get(), 0);
+        assert_eq!(
+            backend.state.delete_ranges.borrow().as_slice(),
+            &[KeyRange {
+                lower: Bound::Included(space(7).encode_key(&key("ab"))),
+                upper: Bound::Excluded(space(7).encode_key(&key("ac"))),
+            }]
+        );
+    }
+
+    #[test]
+    fn clear_space_lowers_to_one_backend_delete_range() {
+        let backend = CountingBackend::default();
+        let storage = StorageContext::new(backend.clone());
+
+        storage
+            .clear_space(space(7), WriteOptions::default())
+            .expect("clear space");
+
+        assert_eq!(backend.state.begin_write_calls.get(), 1);
+        assert_eq!(backend.state.commit_calls.get(), 1);
+        assert_eq!(backend.state.delete_many_calls.get(), 0);
+        assert_eq!(
+            backend.state.delete_ranges.borrow().as_slice(),
+            &[KeyRange {
+                lower: Bound::Included(Key(Bytes::from_static(b"\0\0\0\x07"))),
+                upper: Bound::Excluded(Key(Bytes::from_static(b"\0\0\0\x08"))),
+            }]
+        );
+    }
+
+    #[derive(Clone, Default)]
+    struct CountingBackend {
+        state: Rc<CountingState>,
+    }
+
+    #[derive(Default)]
+    struct CountingState {
+        begin_write_calls: Cell<u64>,
+        commit_calls: Cell<u64>,
+        delete_many_calls: Cell<u64>,
+        put_batches: RefCell<Vec<(SpaceId, Vec<Key>)>>,
+        delete_ranges: RefCell<Vec<KeyRange>>,
+    }
+
+    struct CountingWrite {
+        state: Rc<CountingState>,
+        put_batches: Vec<(SpaceId, Vec<Key>)>,
+        delete_ranges: Vec<KeyRange>,
+    }
+
+    impl Backend for CountingBackend {
+        type Read<'a>
+            = SpyRead
+        where
+            Self: 'a;
+
+        type Write<'a>
+            = CountingWrite
+        where
+            Self: 'a;
+
+        fn capabilities(&self) -> BackendCapabilities {
+            BackendCapabilities::v0(WriteConcurrency::SingleWriter)
+        }
+
+        fn begin_read(&self, _opts: ReadOptions) -> Result<Self::Read<'_>, BackendError> {
+            Ok(SpyRead::default())
+        }
+
+        fn begin_write(&self, _opts: WriteOptions) -> Result<Self::Write<'_>, BackendError> {
+            self.state
+                .begin_write_calls
+                .set(self.state.begin_write_calls.get() + 1);
+            Ok(CountingWrite {
+                state: Rc::clone(&self.state),
+                put_batches: Vec::new(),
+                delete_ranges: Vec::new(),
             })
-            .await
-            .expect("key-only scan reads");
-        assert_eq!(key_only.keys.iter().collect::<Vec<_>>(), vec![b"a", b"b"]);
-        tx.rollback().await.expect("rollback succeeds");
-    }
-
-    #[tokio::test]
-    async fn storage_write_set_applies_as_one_batch() {
-        let backend: Arc<dyn Backend + Send + Sync> = Arc::new(UnitTestBackend::new());
-        let storage = StorageContext::new(backend);
-        let mut tx = storage
-            .begin_write_transaction()
-            .await
-            .expect("transaction opens");
-
-        let mut writes = StorageWriteSet::new();
-        assert!(writes.is_empty());
-        writes.put("ns", b"a".to_vec(), b"1".to_vec());
-        writes.put("ns", b"b".to_vec(), b"2".to_vec());
-        writes.delete("ns", b"missing".to_vec());
-        assert!(!writes.is_empty());
-
-        let stats = writes.apply(tx.as_mut()).await.expect("write set applies");
-        assert_eq!(stats.puts, 2);
-        assert_eq!(stats.deletes, 1);
-        tx.commit().await.expect("commit succeeds");
-
-        let mut tx = storage
-            .begin_read_transaction()
-            .await
-            .expect("read transaction opens");
-        let result = tx
-            .get_values(KvGetRequest {
-                groups: vec![KvGetGroup {
-                    namespace: "ns".to_string(),
-                    keys: vec![b"a".to_vec(), b"b".to_vec()],
-                }],
+        }
+
+        fn durable_write_lock(&self) -> DurableWriteLock {
+            DurableWriteLock::new()
+        }
+    }
+
+    impl BackendWrite for CountingWrite {
+        fn put_many(&mut self, entries: PutBatch) -> Result<(), BackendError> {
+            let space = entries
+                .entries
+                .first()
+                .map(|entry| physical_space(&entry.key))
+                .unwrap_or(SpaceId(0));
+            self.put_batches.push((
+                space,
+                entries
+                    .entries
+                    .into_iter()
+                    .map(|entry| logical_key(entry.key))
+                    .collect(),
+            ));
+            Ok(())
+        }
+
+        fn delete_many(&mut self, _keys: &[Key]) -> Result<(), BackendError> {
+            self.state
+                .delete_many_calls
+                .set(self.state.delete_many_calls.get() + 1);
+            Ok(())
+        }
+
+        fn delete_range(&mut self, range: KeyRange) -> Result<(), BackendError> {
+            self.delete_ranges.push(range);
+            Ok(())
+        }
+
+        fn commit(self) -> Result<CommitResult, BackendError> {
+            self.state
+                .commit_calls
+                .set(self.state.commit_calls.get() + 1);
+            self.state.put_batches.borrow_mut().extend(self.put_batches);
+            self.state
+                .delete_ranges
+                .borrow_mut()
+                .extend(self.delete_ranges);
+            Ok(CommitResult {
+                commit_id: None,
+                stats: WriteStats::default(),
             })
-            .await
-            .expect("batch reads");
-        assert_eq!(result.groups[0].value(0), Some(Some(&b"1"[..])));
-        assert_eq!(result.groups[0].value(1), Some(Some(&b"2"[..])));
-        tx.rollback().await.expect("rollback succeeds");
+        }
+
+        fn rollback(self) -> Result<(), BackendError> {
+            Ok(())
+        }
+    }
+
+    #[derive(Clone, Default)]
+    struct SpyRead {
+        get_many_keys: Rc<RefCell<Vec<Key>>>,
+        scan_range_calls: Rc<Cell<u64>>,
+        scan_range: Rc<RefCell<Option<KeyRange>>>,
+    }
+
+    impl BackendRead for SpyRead {
+        type RangeScan<'a> = BufferedRangeScan;
+
+        fn visit_keys<V>(
+            &self,
+            keys: &[Key],
+            _opts: GetOptions<'_>,
+            visitor: &mut V,
+        ) -> Result<(), BackendError>
+        where
+            V: PointVisitor + ?Sized,
+        {
+            self.get_many_keys.replace(keys.to_vec());
+            for (index, key) in keys.iter().enumerate() {
+                let value = (key.0.as_ref() != b"missing")
+                    .then_some(ProjectedValueRef::FullValue(key.0.as_ref()));
+                visitor.visit(index, key, value)?;
+            }
+            Ok(())
+        }
+
+        fn with_range_scan<T, F>(
+            &self,
+            range: KeyRange,
+            _opts: ScanOptions<'_>,
+            f: F,
+        ) -> Result<T, BackendError>
+        where
+            F: FnOnce(&mut Self::RangeScan<'_>) -> Result<T, BackendError>,
+        {
+            self.scan_range_calls.set(self.scan_range_calls.get() + 1);
+            self.scan_range.replace(Some(range));
+            let mut cursor = BufferedRangeScan::default();
+            f(&mut cursor)
+        }
+    }
+
+    fn space(id: u32) -> StorageSpace {
+        StorageSpace::new(SpaceId(id), "shape.test.space")
+    }
+
+    fn physical_space(key: &Key) -> SpaceId {
+        SpaceId(u32::from_be_bytes(
+            key.0[..4].try_into().expect("space prefix"),
+        ))
+    }
+
+    fn logical_key(key: Key) -> Key {
+        Key(key.0.slice(4..))
+    }
+
+    fn key(bytes: &'static str) -> Key {
+        Key(Bytes::from_static(bytes.as_bytes()))
+    }
+
+    fn value(bytes: &'static str) -> StoredValue {
+        StoredValue {
+            bytes: Bytes::from_static(bytes.as_bytes()),
+        }
     }
 }