From c86861f0bc8d8339622ffd7f41747053a92f36e2 Mon Sep 17 00:00:00 2001
From: Niven <sieniven@gmail.com>
Date: Wed, 24 Jan 2024 15:59:04 +0800
Subject: [PATCH] Shift evm db flush in DVM db flush pipeline

---
 lib/ain-evm/src/core.rs        |  7 +------
 lib/ain-evm/src/evm.rs         | 14 +++++++++++---
 lib/ain-evm/src/storage/mod.rs |  3 +--
 lib/ain-rs-exports/src/evm.rs  |  5 +++++
 lib/ain-rs-exports/src/lib.rs  |  2 ++
 src/validation.cpp             |  7 +++++++
 6 files changed, 27 insertions(+), 11 deletions(-)

diff --git a/lib/ain-evm/src/core.rs b/lib/ain-evm/src/core.rs
index 7ef263b21c..cb97d3dc1d 100644
--- a/lib/ain-evm/src/core.rs
+++ b/lib/ain-evm/src/core.rs
@@ -31,10 +31,7 @@ use crate::{
     gas::check_tx_intrinsic_gas,
     precompiles::MetachainPrecompiles,
     receipt::ReceiptService,
-    storage::{
-        traits::{BlockStorage, FlushableStorage},
-        Storage,
-    },
+    storage::{traits::BlockStorage, Storage},
     transaction::{
         cache::{TransactionCache, ValidateTxInfo},
         SignedTx,
@@ -146,8 +143,6 @@ impl EVMCoreService {
         );
         storage.put_latest_block(Some(&block))?;
         storage.put_block(&block)?;
-        handler.flush()?;
-        storage.flush()?;
 
         Ok(handler)
     }
diff --git a/lib/ain-evm/src/evm.rs b/lib/ain-evm/src/evm.rs
index ae725fcf7d..91c4926bac 100644
--- a/lib/ain-evm/src/evm.rs
+++ b/lib/ain-evm/src/evm.rs
@@ -272,10 +272,7 @@ impl EVMServices {
             .sender
             .send(Notification::Block(block.header.hash()))
             .map_err(|e| format_err!(e.to_string()))?;
-
         self.core.clear_account_nonce();
-        self.core.flush()?;
-        self.storage.flush()?;
         Ok(())
     }
 
@@ -467,6 +464,17 @@ impl EVMServices {
         template.backend.increase_tx_count();
         Ok(())
     }
+
+    ///
+    /// # Safety
+    ///
+    /// Result cannot be used safety unless `cs_main` lock is taken on C++ side
+    /// across all usages. Note: To be replaced with a proper lock flow later.
+    ///
+    pub unsafe fn flush_state_to_db(&self) -> Result<()> {
+        self.core.flush()?;
+        self.storage.flush()
+    }
 }
 
 // Block template methods
diff --git a/lib/ain-evm/src/storage/mod.rs b/lib/ain-evm/src/storage/mod.rs
index b01ad8dcde..00f5eda18d 100644
--- a/lib/ain-evm/src/storage/mod.rs
+++ b/lib/ain-evm/src/storage/mod.rs
@@ -227,7 +227,6 @@ impl Storage {
 impl Rollback for Storage {
     fn disconnect_latest_block(&self) -> Result<()> {
         self.cache.disconnect_latest_block()?;
-        self.blockstore.disconnect_latest_block()?;
-        self.flush()
+        self.blockstore.disconnect_latest_block()
     }
 }
diff --git a/lib/ain-rs-exports/src/evm.rs b/lib/ain-rs-exports/src/evm.rs
index c9a5f8d977..082a76903d 100644
--- a/lib/ain-rs-exports/src/evm.rs
+++ b/lib/ain-rs-exports/src/evm.rs
@@ -858,3 +858,8 @@ fn evm_try_dispatch_pending_transactions_event(raw_tx: &str) -> Result<()> {
         .send(Notification::Transaction(signed_tx.hash()))
         .map_err(|e| format_err!(e.to_string()))?)
 }
+
+#[ffi_fallible]
+fn evm_try_flush_db() -> Result<()> {
+    unsafe { SERVICES.evm.flush_state_to_db() }
+}
diff --git a/lib/ain-rs-exports/src/lib.rs b/lib/ain-rs-exports/src/lib.rs
index 07ff3ef055..e97f7e9a53 100644
--- a/lib/ain-rs-exports/src/lib.rs
+++ b/lib/ain-rs-exports/src/lib.rs
@@ -341,6 +341,8 @@ pub mod ffi {
             result: &mut CrossBoundaryResult,
             raw_tx: &str,
         );
+
+        fn evm_try_flush_db(result: &mut CrossBoundaryResult);
     }
 
     // =========  Debug ==========
diff --git a/src/validation.cpp b/src/validation.cpp
index 733348709d..1682c65191 100644
--- a/src/validation.cpp
+++ b/src/validation.cpp
@@ -3583,6 +3583,13 @@ bool CChainState::FlushStateToDisk(const CChainParams &chainparams,
                 if (!CoinsTip().Flush() || !pcustomcsDB->Flush()) {
                     return AbortNode(state, "Failed to write to coin or masternode db to disk");
                 }
+                // Flush the EVM chainstate
+                if (IsEVMEnabled(pcustomcsview->GetAttributes())) {
+                    auto res = XResultStatusLogged(evm_try_flush_db(result));
+                    if (!res) {
+                        return AbortNode(state, "Failed to write to EVM db to disk");
+                    }
+                }
                 if (!compactBegin.empty() && !compactEnd.empty()) {
                     auto time = GetTimeMillis();
                     pcustomcsDB->Compact(compactBegin, compactEnd);