BFT-476: Add ctx to PersistentBatchStore methods

node/libs/storage/src/batch_store.rs

@@ -10,7 +10,7 @@ pub struct BatchStoreState {
     /// Stored batch with the lowest number.
     /// If last is `None`, this is the first batch that should be fetched.
     pub first: attester::BatchNumber,
-    /// Stored QC of the latest batch.
+    /// The last stored L1 batch.
     /// None iff store is empty.
     pub last: Option<attester::SyncBatch>,
 }
@@ -48,18 +48,36 @@ impl BatchStoreState {
 /// Trait for the shared state of batches between the consensus and the execution layer.
 #[async_trait::async_trait]
 pub trait PersistentBatchStore: 'static + fmt::Debug + Send + Sync {
+    /// Range of batches persisted in storage.
+    fn persisted(&self) -> sync::watch::Receiver<BatchStoreState>;
+
     /// Get the L1 batch from storage with the highest number.
-    async fn last_batch(&self) -> attester::BatchNumber;
+    ///
+    /// Returns `None` if no batches have been created yet.
+    async fn last_batch(&self, ctx: &ctx::Ctx) -> ctx::Result<Option<attester::BatchNumber>>;
+
     /// Get the L1 batch QC from storage with the highest number.
-    async fn last_batch_qc(&self) -> attester::BatchQC;
+    ///
+    /// Returns `None` if we don't have a QC for any of the batches yet.
+    async fn last_batch_qc(&self, ctx: &ctx::Ctx) -> ctx::Result<Option<attester::BatchQC>>;
+
     /// Returns the batch with the given number.
-    async fn get_batch(&self, number: attester::BatchNumber) -> Option<attester::SyncBatch>;
+    async fn get_batch(
+        &self,
+        ctx: &ctx::Ctx,
+        number: attester::BatchNumber,
+    ) -> ctx::Result<Option<attester::SyncBatch>>;
+
     /// Returns the QC of the batch with the given number.
-    async fn get_batch_qc(&self, number: attester::BatchNumber) -> Option<attester::BatchQC>;
-    /// Store the given QC in the storage.
-    async fn store_qc(&self, qc: attester::BatchQC);
-    /// Range of batches persisted in storage.
-    fn persisted(&self) -> sync::watch::Receiver<BatchStoreState>;
+    async fn get_batch_qc(
+        &self,
+        ctx: &ctx::Ctx,
+        number: attester::BatchNumber,
+    ) -> ctx::Result<Option<attester::BatchQC>>;
+
+    /// Store the given batch QC in the storage.
+    async fn store_qc(&self, ctx: &ctx::Ctx, qc: attester::BatchQC) -> ctx::Result<()>;
+
     /// Queue the batch to be persisted in storage.
     /// `queue_next_batch()` may return BEFORE the batch is actually persisted,
     /// but if the call succeeded the batch is expected to be persisted eventually.
@@ -196,7 +214,7 @@ impl BatchStore {
     /// Fetches a batch (from queue or persistent storage).
     pub async fn batch(
         &self,
-        _ctx: &ctx::Ctx,
+        ctx: &ctx::Ctx,
         number: attester::BatchNumber,
     ) -> ctx::Result<Option<attester::SyncBatch>> {
         {
@@ -210,10 +228,10 @@ impl BatchStore {
         }
         let batch = self
             .persistent
-            .get_batch(number)
+            .get_batch(ctx, number)
             .await
             .context("persistent.batch()")?;
-        Ok(Some(batch))
+        Ok(batch)
     }
 
     /// Append batch to a queue to be persisted eventually.

node/libs/storage/src/testonly/in_memory.rs

@@ -132,36 +132,43 @@ impl PersistentBatchStore for BatchStore {
         self.0.persisted.subscribe()
     }
 
-    async fn last_batch(&self) -> attester::BatchNumber {
-        self.0
-            .persisted
-            .borrow()
-            .last
-            .clone()
-            .map(|qc| qc.number)
-            .unwrap()
+    async fn last_batch(&self, _ctx: &ctx::Ctx) -> ctx::Result<Option<attester::BatchNumber>> {
+        Ok(self.0.persisted.borrow().last.clone().map(|qc| qc.number))
     }
 
-    async fn last_batch_qc(&self) -> attester::BatchQC {
+    async fn last_batch_qc(&self, _ctx: &ctx::Ctx) -> ctx::Result<Option<attester::BatchQC>> {
         let qcs = self.0.qcs.lock().unwrap();
         let last_batch_number = qcs.keys().max().unwrap();
-        qcs.get(last_batch_number).unwrap().clone()
+        Ok(qcs.get(last_batch_number).cloned())
     }
 
-    async fn get_batch_qc(&self, number: attester::BatchNumber) -> Option<attester::BatchQC> {
+    async fn get_batch_qc(
+        &self,
+        _ctx: &ctx::Ctx,
+        number: attester::BatchNumber,
+    ) -> ctx::Result<Option<attester::BatchQC>> {
         let qcs = self.0.qcs.lock().unwrap();
-        qcs.get(&number).cloned()
+        Ok(qcs.get(&number).cloned())
     }
 
-    async fn store_qc(&self, qc: attester::BatchQC) {
+    async fn store_qc(&self, _ctx: &ctx::Ctx, qc: attester::BatchQC) -> ctx::Result<()> {
         self.0.qcs.lock().unwrap().insert(qc.message.number, qc);
+        Ok(())
     }
 
-    async fn get_batch(&self, number: attester::BatchNumber) -> Option<attester::SyncBatch> {
+    async fn get_batch(
+        &self,
+        _ctx: &ctx::Ctx,
+        number: attester::BatchNumber,
+    ) -> ctx::Result<Option<attester::SyncBatch>> {
         let batches = self.0.batches.lock().unwrap();
-        let front = batches.front()?;
-        let idx = number.0.checked_sub(front.number.0)?;
-        batches.get(idx as usize).cloned()
+        let Some(front) = batches.front() else {
+            return Ok(None);
+        };
+        let Some(idx) = number.0.checked_sub(front.number.0) else {
+            return Ok(None);
+        };
+        Ok(batches.get(idx as usize).cloned())
     }
 
     async fn queue_next_batch(

node/libs/storage/src/testonly/mod.rs

@@ -160,7 +160,7 @@ pub async fn dump(ctx: &ctx::Ctx, store: &dyn PersistentBlockStore) -> Vec<valid
 
 /// Dumps all the batches stored in `store`.
 pub async fn dump_batch(
-    _ctx: &ctx::Ctx,
+    ctx: &ctx::Ctx,
     store: &dyn PersistentBatchStore,
 ) -> Vec<attester::SyncBatch> {
     // let genesis = store.genesis(ctx).await.unwrap();
@@ -173,14 +173,14 @@ pub async fn dump_batch(
         .map(|sb| sb.number.next())
         .unwrap_or(state.first);
     for n in (state.first.0..after.0).map(attester::BatchNumber) {
-        let batch = store.get_batch(n).await.unwrap();
+        let batch = store.get_batch(ctx, n).await.unwrap().unwrap();
         assert_eq!(batch.number, n);
         batches.push(batch);
     }
     if let Some(before) = state.first.prev() {
-        assert!(store.get_batch(before).await.is_none());
+        assert!(store.get_batch(ctx, before).await.unwrap().is_none());
     }
-    assert!(store.get_batch(after).await.is_none());
+    assert!(store.get_batch(ctx, after).await.unwrap().is_none());
     batches
 }