From 89040ca8c8480f8602d302257ba85a31dad340e6 Mon Sep 17 00:00:00 2001
From: Alexey Shekhirin <a.shekhirin@gmail.com>
Date: Mon, 23 Oct 2023 21:55:04 +0100
Subject: [PATCH] DRY close transaction metric

---
 .../db/src/implementation/mdbx/cursor.rs      |  6 +--
 .../storage/db/src/implementation/mdbx/tx.rs  | 50 ++++++++++---------
 2 files changed, 30 insertions(+), 26 deletions(-)

diff --git a/crates/storage/db/src/implementation/mdbx/cursor.rs b/crates/storage/db/src/implementation/mdbx/cursor.rs
index c35fc1d31ff94..4882f16fcffc5 100644
--- a/crates/storage/db/src/implementation/mdbx/cursor.rs
+++ b/crates/storage/db/src/implementation/mdbx/cursor.rs
@@ -43,11 +43,11 @@ impl<'tx, K: TransactionKind, T: Table> Cursor<'tx, K, T> {
         self
     }
 
-    fn execute_with_operation_metric(
+    fn execute_with_operation_metric<T>(
         &mut self,
         operation: Operation,
-        f: impl FnOnce(&mut Self) -> Result<(), DatabaseError>,
-    ) -> Result<(), DatabaseError> {
+        f: impl FnOnce(&mut Self) -> T,
+    ) -> T {
         let start = Instant::now();
         let result = f(self);
 
diff --git a/crates/storage/db/src/implementation/mdbx/tx.rs b/crates/storage/db/src/implementation/mdbx/tx.rs
index 2cf6f18f3c1e7..6f7e185541595 100644
--- a/crates/storage/db/src/implementation/mdbx/tx.rs
+++ b/crates/storage/db/src/implementation/mdbx/tx.rs
@@ -73,11 +73,30 @@ impl<'env, K: TransactionKind, E: EnvironmentKind> Tx<'env, K, E> {
         ))
     }
 
+    fn execute_with_close_transaction_metric<T>(
+        mut self,
+        outcome: TransactionOutcome,
+        f: impl FnOnce(Self) -> T,
+    ) -> T {
+        let start = Instant::now();
+        let (txn_id, metrics_tx) =
+            self.metrics_tx.take().map(|metrics_tx| (self.id(), metrics_tx)).unzip();
+        let result = f(self);
+        if let (Some(txn_id), Some(metrics_tx)) = (txn_id, metrics_tx) {
+            let _ = metrics_tx.send(MetricEvent::CloseTransaction {
+                txn_id,
+                outcome,
+                commit_duration: start.elapsed(),
+            });
+        }
+        result
+    }
+
     fn execute_with_operation_metric<T>(
         &self,
         operation: Operation,
-        f: impl FnOnce(&Transaction<'_, K, E>) -> Result<T, DatabaseError>,
-    ) -> Result<T, DatabaseError> {
+        f: impl FnOnce(&Transaction<'_, K, E>) -> T,
+    ) -> T {
         let start = Instant::now();
         let result = f(&self.inner);
         if let Some(metrics_tx) = &self.metrics_tx {
@@ -111,30 +130,15 @@ impl<K: TransactionKind, E: EnvironmentKind> DbTx for Tx<'_, K, E> {
     }
 
     fn commit(mut self) -> Result<bool, DatabaseError> {
-        let start = Instant::now();
-        let metrics = self.metrics_tx.take().map(|metrics_tx| (self.id(), metrics_tx));
-        let result = self.inner.commit().map_err(|e| DatabaseError::Commit(e.into()));
-        if let Some((txn_id, metrics_tx)) = metrics {
-            let _ = metrics_tx.send(MetricEvent::CloseTransaction {
-                txn_id,
-                outcome: TransactionOutcome::Commit,
-                commit_duration: start.elapsed(),
-            });
-        }
-        result
+        self.execute_with_close_transaction_metric(TransactionOutcome::Commit, |this| {
+            this.inner.commit().map_err(|e| DatabaseError::Commit(e.into()))
+        })
     }
 
     fn drop(mut self) {
-        let start = Instant::now();
-        let metrics = self.metrics_tx.take().map(|metrics_tx| (self.id(), metrics_tx));
-        drop(self.inner);
-        if let Some((txn_id, metrics_tx)) = metrics {
-            let _ = metrics_tx.send(MetricEvent::CloseTransaction {
-                txn_id,
-                outcome: TransactionOutcome::Abort,
-                commit_duration: start.elapsed(),
-            });
-        }
+        self.execute_with_close_transaction_metric(TransactionOutcome::Abort, |this| {
+            drop(this.inner)
+        })
     }
 
     // Iterate over read only values in database.