apollo_batcher: evict the reverted height from the caches after writing the pending results

crates/apollo_batcher/src/batcher.rs

@@ -1557,6 +1557,7 @@ impl Batcher {
                 &mut self.storage_writer,
             )
             .expect("Failed to write commitment results to storage.");
+        self.commitment_manager.evict_reverted_height(height);
 
         info!("Revert task result: {revert_task_result:?}");
         self.validate_revert_task_result(revert_task_result, height).await;

crates/apollo_batcher/src/batcher_test.rs

@@ -1639,9 +1639,26 @@ async fn revert_block() {
         .with(eq(LATEST_BLOCK_IN_STORAGE))
         .returning(|_| ());
 
-    let storage_reader = mock_storage_reader_for_revert();
-    let mock_dependencies =
+    // The commit of the reverted height is still pending when the revert is requested, so its
+    // result is written, and cached, before the storage is reverted.
+    storage_writer.expect_set_global_root_and_block_hash().times(1).returning(|_, _, _, _| Ok(()));
+    let mut storage_reader = mock_storage_reader_for_revert();
+    storage_reader
+        .expect_get_parent_hash_and_partial_block_hash_components()
+        .with(eq(LATEST_BLOCK_IN_STORAGE))
+        .returning(|_| {
+            Ok((
+                Some(BlockHash::default()),
+                Some(PartialBlockHashComponents {
+                    block_number: LATEST_BLOCK_IN_STORAGE,
+                    ..Default::default()
+                }),
+            ))
+        });
+    let mut mock_dependencies =
         MockDependencies { storage_reader, storage_writer, ..Default::default() };
+    // The pending commit's block hash is irrelevant here; don't compare it to the configured one.
+    mock_dependencies.batcher_config.static_config.first_block_with_partial_block_hash = None;
 
     let committer_offset = mock_dependencies.clients.committer_client.get_offset();
 
@@ -1650,16 +1667,27 @@ async fn revert_block() {
     let metrics = recorder.handle().render();
     assert_eq!(BUILDING_HEIGHT.parse_numeric_metric::<u64>(&metrics), Some(INITIAL_HEIGHT.0));
 
-    let revert_input = RevertBlockInput { height: LATEST_BLOCK_IN_STORAGE };
-
-    batcher
-        .commitment_manager
-        .recent_state_commitment_infos_cache
-        .put(LATEST_BLOCK_IN_STORAGE, test_state_commitment_infos());
+    let pending_commit =
+        CommitterTaskInput::ReadPathsAndCommitBlock(ReadPathsAndCommitBlockRequest {
+            commit: CommitBlockRequest {
+                height: LATEST_BLOCK_IN_STORAGE,
+                state_diff: ThinStateDiff::default(),
+                state_diff_commitment: None,
+            },
+            accessed_keys: Default::default(),
+        });
+    batcher.commitment_manager.tasks_sender.send(pending_commit).await.unwrap();
 
+    let revert_input = RevertBlockInput { height: LATEST_BLOCK_IN_STORAGE };
     assert_eq!(*(committer_offset.lock().await), INITIAL_HEIGHT);
     batcher.revert_block(revert_input).await.unwrap();
     assert_eq!(*committer_offset.lock().await, LATEST_BLOCK_IN_STORAGE);
+
+    // The reverted height is evicted from both caches even though its commit result was cached
+    // during the revert.
+    assert!(
+        !batcher.commitment_manager.recent_block_hashes_cache.contains(&LATEST_BLOCK_IN_STORAGE)
+    );
     assert!(
         !batcher
             .commitment_manager

crates/apollo_batcher/src/commitment_manager/commitment_manager_impl.rs

@@ -192,6 +192,14 @@ impl<S: StateCommitterTrait> CommitmentManager<S> {
         }
     }
 
+    /// Evicts the reverted height from the caches. Must run after the commitment results that were
+    /// pending when the revert was requested are written, since they may include the reverted
+    /// height itself.
+    pub(crate) fn evict_reverted_height(&mut self, height: BlockNumber) {
+        self.recent_block_hashes_cache.pop(&height);
+        self.recent_state_commitment_infos_cache.pop(&height);
+    }
+
     /// Fetches all ready commitment results from the state committer. Panics if any task is a
     /// revert.
     pub(crate) fn get_commitment_results(&mut self) -> Vec<CommitmentTaskOutput> {
@@ -240,7 +248,6 @@ impl<S: StateCommitterTrait> CommitmentManager<S> {
                     commitment_results.push(read_path_and_commit_task_result)
                 }
                 CommitterTaskOutput::Revert(revert_task_result) => {
-                    self.recent_state_commitment_infos_cache.pop(&revert_task_result.height);
                     return (commitment_results, revert_task_result);
                 }
             }
@@ -499,9 +506,6 @@ impl<S: StateCommitterTrait> CommitmentManager<S> {
             });
         }
 
-        // Remove the reverted block hash from the cache.
-        self.recent_block_hashes_cache.pop(&height);
-
         let revert_task_input =
             CommitterTaskInput::Revert(RevertBlockRequest { height, reversed_state_diff });
         self.add_task_with_retries(

crates/apollo_batcher/src/commitment_manager/commitment_manager_test.rs

@@ -585,13 +585,9 @@ async fn test_wait_for_revert(mut mock_dependencies: MockDependencies) {
         INITIAL_HEIGHT,
     )
     .await;
-    commitment_manager
-        .recent_state_commitment_infos_cache
-        .put(height, test_state_commitment_infos());
     let (commitment_results, revert_result) = commitment_manager.wait_for_revert_result().await;
     assert_eq!(commitment_results.len(), 2);
     assert_eq!(revert_result.height, height);
-    assert!(!commitment_manager.recent_state_commitment_infos_cache.contains(&height));
 }
 
 #[rstest]