From 6dff69bde9da8a68a328a557064b1cb4df341bf8 Mon Sep 17 00:00:00 2001
From: Emilia Hane <emiliaha95@gmail.com>
Date: Tue, 24 Jan 2023 10:27:30 +0100
Subject: [PATCH] Atomically update blob info with pruned blobs

---
 beacon_node/store/src/hot_cold_store.rs | 23 +++++++++++++++--------
 beacon_node/store/src/lib.rs            |  1 +
 2 files changed, 16 insertions(+), 8 deletions(-)

diff --git a/beacon_node/store/src/hot_cold_store.rs b/beacon_node/store/src/hot_cold_store.rs
index 0ea2a95a5..712c2e966 100644
--- a/beacon_node/store/src/hot_cold_store.rs
+++ b/beacon_node/store/src/hot_cold_store.rs
@@ -827,6 +827,10 @@ impl<E: EthSpec, Hot: ItemStore<E>, Cold: ItemStore<E>> HotColdDB<E, Hot, Cold>
                         get_key_for_col(DBColumn::BeaconBlobOrphan.into(), block_root.as_bytes());
                     key_value_batch.push(KeyValueStoreOp::PutKeyValue(db_key, [].into()));
                 }
+
+                StoreOp::PutRawKVStoreOp(kv_store_op) => {
+                    key_value_batch.push(kv_store_op);
+                }
             }
         }
         Ok(key_value_batch)
@@ -869,6 +873,8 @@ impl<E: EthSpec, Hot: ItemStore<E>, Cold: ItemStore<E>> HotColdDB<E, Hot, Cold>
                 StoreOp::DeleteExecutionPayload(_) => (),
 
                 StoreOp::PutOrphanedBlobsKey(_) => (),
+
+                StoreOp::PutRawKVStoreOp(_) => (),
             }
         }
 
@@ -1865,8 +1871,16 @@ impl<E: EthSpec, Hot: ItemStore<E>, Cold: ItemStore<E>> HotColdDB<E, Hot, Cold>
                 break;
             }
         }
-
         let blobs_sidecars_pruned = ops.len();
+
+        let update_blob_info = self.compare_and_set_blob_info(
+            Some(blob_info),
+            Some(BlobInfo {
+                oldest_blob_slot: end_slot + 1,
+            }),
+        )?;
+        ops.push(StoreOp::PutRawKVStoreOp(update_blob_info));
+
         self.do_atomically(ops)?;
         info!(
             self.log,
@@ -1874,13 +1888,6 @@ impl<E: EthSpec, Hot: ItemStore<E>, Cold: ItemStore<E>> HotColdDB<E, Hot, Cold>
             "blobs_sidecars_pruned" => blobs_sidecars_pruned,
         );
 
-        self.compare_and_set_blob_info_with_write(
-            Some(blob_info),
-            Some(BlobInfo {
-                oldest_blob_slot: end_slot + 1,
-            }),
-        )?;
-
         Ok(())
     }
 }
diff --git a/beacon_node/store/src/lib.rs b/beacon_node/store/src/lib.rs
index 7857f256e..18e3b3ca3 100644
--- a/beacon_node/store/src/lib.rs
+++ b/beacon_node/store/src/lib.rs
@@ -166,6 +166,7 @@ pub enum StoreOp<'a, E: EthSpec> {
     DeleteBlobs(Hash256),
     DeleteState(Hash256, Option<Slot>),
     DeleteExecutionPayload(Hash256),
+    PutRawKVStoreOp(KeyValueStoreOp),
 }
 
 /// A unique column identifier.