[hermes] Fix concurrency issue (#925)

Author: ali-behjati

hermes/Cargo.lock

@@ -1,6 +1,6 @@
 [package]
 name                   = "hermes"
-version                = "0.1.3"
+version                = "0.1.4"
 edition                = "2021"
 
 [dependencies]
@@ -35,6 +35,7 @@ libp2p                 = { version = "0.42.2", features = [
 ]}
 
 log                    = { version = "0.4.17" }
+prometheus-client      = { version = "0.21.1" }
 pyth-sdk               = { version = "0.7.0" }
 
 # Parse Wormhole attester price attestations.

hermes/Cargo.toml

@@ -19,10 +19,7 @@ use {
             construct_message_states_proofs,
             store_wormhole_merkle_verified_message,
         },
-        storage::{
-            AccumulatorState,
-            CompletedAccumulatorState,
-        },
+        storage::CompletedAccumulatorState,
         types::{
             ProofSet,
             UnixTimestamp,
@@ -150,18 +147,14 @@ impl Store {
             Update::AccumulatorMessages(accumulator_messages) => {
                 let slot = accumulator_messages.slot;
                 log::info!("Storing accumulator messages for slot {:?}.", slot,);
-                let mut accumulator_state = self
-                    .storage
-                    .fetch_accumulator_state(slot)
-                    .await?
-                    .unwrap_or(AccumulatorState {
-                        slot,
-                        accumulator_messages: None,
-                        wormhole_merkle_state: None,
-                    });
-                accumulator_state.accumulator_messages = Some(accumulator_messages);
                 self.storage
-                    .store_accumulator_state(accumulator_state)
+                    .update_accumulator_state(
+                        slot,
+                        Box::new(|mut state| {
+                            state.accumulator_messages = Some(accumulator_messages);
+                            state
+                        }),
+                    )
                     .await?;
                 slot
             }

hermes/src/store.rs

@@ -1,7 +1,6 @@
 use {
     crate::store::{
         storage::{
-            AccumulatorState,
             CompletedAccumulatorState,
             MessageState,
         },
@@ -49,23 +48,18 @@ pub async fn store_wormhole_merkle_verified_message(
     root: WormholeMerkleRoot,
     vaa_bytes: Vec<u8>,
 ) -> Result<()> {
-    let mut accumulator_state = store
-        .storage
-        .fetch_accumulator_state(root.slot)
-        .await?
-        .unwrap_or(AccumulatorState {
-            slot:                  root.slot,
-            accumulator_messages:  None,
-            wormhole_merkle_state: None,
-        });
-
-    accumulator_state.wormhole_merkle_state = Some(WormholeMerkleState {
-        root,
-        vaa: vaa_bytes,
-    });
     store
         .storage
-        .store_accumulator_state(accumulator_state)
+        .update_accumulator_state(
+            root.slot,
+            Box::new(|mut state| {
+                state.wormhole_merkle_state = Some(WormholeMerkleState {
+                    root,
+                    vaa: vaa_bytes,
+                });
+                state
+            }),
+        )
         .await?;
     Ok(())
 }

hermes/src/store/proof/wormhole_merkle.rs

@@ -134,8 +134,20 @@ pub trait Storage: Send + Sync {
         filter: MessageStateFilter,
     ) -> Result<Vec<MessageState>>;
 
+    /// Store the accumulator state. Please note that this call will replace the
+    /// existing accumulator state for the given state's slot. If you wish to
+    /// update the accumulator state, use `update_accumulator_state` instead.
     async fn store_accumulator_state(&self, state: AccumulatorState) -> Result<()>;
     async fn fetch_accumulator_state(&self, slot: Slot) -> Result<Option<AccumulatorState>>;
+
+    /// Update the accumulator state inplace using the provided callback. The callback
+    /// takes the current state and returns the new state. If there is no accumulator
+    /// state for the given slot, the callback will be called with an empty accumulator state.
+    async fn update_accumulator_state(
+        &self,
+        slot: Slot,
+        callback: Box<dyn (FnOnce(AccumulatorState) -> AccumulatorState) + Send>,
+    ) -> Result<()>;
 }
 
 pub type StorageInstance = Box<dyn Storage>;