Some code deduplication related to peer db storage

dns_server/Cargo.toml

@@ -17,16 +17,15 @@ storage = { path = "../storage" }
 storage-lmdb = { path = "../storage/lmdb" }
 utils = { path = '../utils' }
 
+async-trait.workspace = true
+clap = { workspace = true, features = ["derive"] }
+directories.workspace = true
 futures = { workspace = true }
+parity-scale-codec.workspace = true
+thiserror.workspace = true
 tokio = { workspace = true, default-features = false }
-
 trust-dns-client.workspace = true
 trust-dns-server.workspace = true
-thiserror.workspace = true
-async-trait.workspace = true
-parity-scale-codec.workspace = true
-clap = { workspace = true, features = ["derive"] }
-directories.workspace = true
 
 [dev-dependencies]
 p2p-test-utils = { path = "../p2p/p2p-test-utils" }

dns_server/src/crawler_p2p/crawler_manager/mod.rs

@@ -33,7 +33,10 @@ use p2p::{
         types::{ConnectivityEvent, SyncingEvent},
         ConnectivityService, NetworkingService, SyncingEventReceiver,
     },
-    peer_manager::ip_or_socket_address_to_peer_address,
+    peer_manager::{
+        ip_or_socket_address_to_peer_address,
+        peerdb_common::{storage::update_db, TransactionRo, TransactionRw},
+    },
     types::{
         ip_or_socket_address::IpOrSocketAddress, peer_address::PeerAddress, peer_id::PeerId,
         socket_address::SocketAddress, IsGlobalIp,
@@ -43,10 +46,7 @@ use tokio::sync::mpsc;
 
 use crate::{dns_server::DnsServerCommand, error::DnsServerError};
 
-use self::storage::{
-    DnsServerStorage, DnsServerStorageRead, DnsServerStorageWrite, DnsServerTransactionRo,
-    DnsServerTransactionRw,
-};
+use self::storage::{DnsServerStorage, DnsServerStorageRead, DnsServerStorageWrite};
 
 use super::crawler::{Crawler, CrawlerCommand, CrawlerEvent};
 
@@ -299,11 +299,11 @@ where
 
                 match (old_state.is_persistent(), new_state.is_persistent()) {
                     (false, true) => {
-                        storage::update_db(storage, |tx| tx.add_address(&address.to_string()))
+                        update_db(storage, |tx| tx.add_address(&address.to_string()))
                             .expect("update_db must succeed (add_address)");
                     }
                     (true, false) => {
-                        storage::update_db(storage, |tx| tx.del_address(&address.to_string()))
+                        update_db(storage, |tx| tx.del_address(&address.to_string()))
                             .expect("update_db must succeed (del_address)");
                     }
                     _ => {}

dns_server/src/crawler_p2p/crawler_manager/storage.rs

@@ -13,6 +13,8 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use p2p::peer_manager::peerdb_common::{TransactionRo, TransactionRw, Transactional};
+
 pub trait DnsServerStorageRead {
     fn get_version(&self) -> Result<Option<u32>, storage::Error>;
 
@@ -27,58 +29,27 @@ pub trait DnsServerStorageWrite {
     fn del_address(&mut self, address: &str) -> Result<(), storage::Error>;
 }
 
-pub trait DnsServerTransactionRo: DnsServerStorageRead {
-    fn close(self);
-}
-
-pub trait DnsServerTransactionRw: DnsServerStorageWrite {
-    fn abort(self);
-
-    fn commit(self) -> Result<(), storage::Error>;
-}
-
-/// Support for transactions over blockchain storage
-pub trait DnsServerTransactional<'t> {
-    /// Associated read-only transaction type.
-    type TransactionRo: DnsServerTransactionRo + 't;
+// Note: here we want to say something like:
+//  pub trait DnsServerStorage: for<'t> Transactional<'t> + Send
+//      where for<'t> <Self as Transactional<'t>>::TransactionRo: DnsServerStorageRead,
+//            for<'t> <Self as Transactional<'t>>::TransactionRw: DnsServerStorageWrite {}
+// But currently Rust would require us to duplicate the "where" constrains in all places
+// where DnsServerStorage is used, so we use this "Helper" approach instead.
+pub trait DnsServerStorage: for<'t> DnsServerStorageHelper<'t> + Send {}
 
-    /// Associated read-write transaction type.
-    type TransactionRw: DnsServerTransactionRw + 't;
-
-    /// Start a read-only transaction.
-    fn transaction_ro<'s: 't>(&'s self) -> Result<Self::TransactionRo, storage::Error>;
-
-    /// Start a read-write transaction.
-    fn transaction_rw<'s: 't>(&'s self) -> Result<Self::TransactionRw, storage::Error>;
+pub trait DnsServerStorageHelper<'t>:
+    Transactional<'t, TransactionRo = Self::TxRo, TransactionRw = Self::TxRw>
+{
+    type TxRo: TransactionRo + DnsServerStorageRead + 't;
+    type TxRw: TransactionRw + DnsServerStorageWrite + 't;
 }
 
-pub trait DnsServerStorage: for<'tx> DnsServerTransactional<'tx> + Send {}
-
-const MAX_RECOVERABLE_ERROR_RETRY_COUNT: u32 = 3;
-
-/// Try update storage, gracefully handle recoverable errors
-pub fn update_db<S, F>(storage: &S, f: F) -> Result<(), storage::Error>
+impl<'t, T> DnsServerStorageHelper<'t> for T
 where
-    S: DnsServerStorage,
-    F: Fn(&mut <S as DnsServerTransactional<'_>>::TransactionRw) -> Result<(), storage::Error>,
+    T: Transactional<'t>,
+    Self::TransactionRo: DnsServerStorageRead + 't,
+    Self::TransactionRw: DnsServerStorageWrite + 't,
 {
-    let mut recoverable_errors = 0;
-    loop {
-        let res = || -> Result<(), storage::Error> {
-            let mut tx = storage.transaction_rw()?;
-            f(&mut tx)?;
-            tx.commit()
-        }();
-
-        match res {
-            Ok(()) => return Ok(()),
-            Err(storage::Error::Recoverable(e)) => {
-                recoverable_errors += 1;
-                if recoverable_errors >= MAX_RECOVERABLE_ERROR_RETRY_COUNT {
-                    return Err(storage::Error::Recoverable(e));
-                }
-            }
-            Err(storage::Error::Fatal(e)) => return Err(storage::Error::Fatal(e)),
-        }
-    }
+    type TxRo = Self::TransactionRo;
+    type TxRw = Self::TransactionRw;
 }

dns_server/src/crawler_p2p/crawler_manager/storage_impl.rs

@@ -13,10 +13,8 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use super::storage::{
-    DnsServerStorage, DnsServerStorageRead, DnsServerStorageWrite, DnsServerTransactionRo,
-    DnsServerTransactionRw, DnsServerTransactional,
-};
+use super::storage::{DnsServerStorage, DnsServerStorageRead, DnsServerStorageWrite};
+use p2p::peer_manager::peerdb_common::storage_impl::{StorageImpl, StorageTxRo, StorageTxRw};
 use serialization::{encoded::Encoded, DecodeAll, Encode};
 
 type ValueId = u32;
@@ -34,76 +32,39 @@ storage::decl_schema! {
 
 const VALUE_ID_VERSION: ValueId = 1;
 
-pub struct DnsServerStoreTxRo<'st, B: storage::Backend>(storage::TransactionRo<'st, B, Schema>);
+type DnsServerStoreTxRo<'st, B> = StorageTxRo<'st, B, Schema>;
+type DnsServerStoreTxRw<'st, B> = StorageTxRw<'st, B, Schema>;
 
-pub struct DnsServerStoreTxRw<'st, B: storage::Backend>(storage::TransactionRw<'st, B, Schema>);
-
-impl<'tx, B: storage::Backend + 'tx> DnsServerTransactional<'tx> for DnsServerStorageImpl<B> {
-    type TransactionRo = DnsServerStoreTxRo<'tx, B>;
-    type TransactionRw = DnsServerStoreTxRw<'tx, B>;
-
-    fn transaction_ro<'st: 'tx>(&'st self) -> Result<Self::TransactionRo, storage::Error> {
-        self.0.transaction_ro().map(DnsServerStoreTxRo)
-    }
-
-    fn transaction_rw<'st: 'tx>(&'st self) -> Result<Self::TransactionRw, storage::Error> {
-        self.0.transaction_rw(None).map(DnsServerStoreTxRw)
-    }
-}
+pub type DnsServerStorageImpl<B> = StorageImpl<B, Schema>;
 
 impl<B: storage::Backend + 'static> DnsServerStorage for DnsServerStorageImpl<B> {}
 
-pub struct DnsServerStorageImpl<T: storage::Backend>(storage::Storage<T, Schema>);
-
-impl<B: storage::Backend> DnsServerStorageImpl<B> {
-    pub fn new(storage: B) -> Result<Self, storage::Error> {
-        let store = storage::Storage::<_, Schema>::new(storage)?;
-        Ok(Self(store))
-    }
-}
-
 impl<'st, B: storage::Backend> DnsServerStorageWrite for DnsServerStoreTxRw<'st, B> {
     fn set_version(&mut self, version: u32) -> Result<(), storage::Error> {
-        self.0.get_mut::<DBValue, _>().put(VALUE_ID_VERSION, version.encode())
+        self.storage().get_mut::<DBValue, _>().put(VALUE_ID_VERSION, version.encode())
     }
 
     fn add_address(&mut self, address: &str) -> Result<(), storage::Error> {
-        self.0.get_mut::<DBAddresses, _>().put(address, ())
+        self.storage().get_mut::<DBAddresses, _>().put(address, ())
     }
 
     fn del_address(&mut self, address: &str) -> Result<(), storage::Error> {
-        self.0.get_mut::<DBAddresses, _>().del(address)
-    }
-}
-
-impl<'st, B: storage::Backend> DnsServerTransactionRw for DnsServerStoreTxRw<'st, B> {
-    fn abort(self) {
-        self.0.abort()
-    }
-
-    fn commit(self) -> Result<(), storage::Error> {
-        self.0.commit()
+        self.storage().get_mut::<DBAddresses, _>().del(address)
     }
 }
 
 impl<'st, B: storage::Backend> DnsServerStorageRead for DnsServerStoreTxRo<'st, B> {
     fn get_version(&self) -> Result<Option<u32>, storage::Error> {
-        let map = self.0.get::<DBValue, _>();
+        let map = self.storage().get::<DBValue, _>();
         let vec_opt = map.get(VALUE_ID_VERSION)?.as_ref().map(Encoded::decode);
         Ok(vec_opt.map(|vec| {
             u32::decode_all(&mut vec.as_ref()).expect("db values to be encoded correctly")
         }))
     }
 
     fn get_addresses(&self) -> Result<Vec<String>, storage::Error> {
-        let map = self.0.get::<DBAddresses, _>();
+        let map = self.storage().get::<DBAddresses, _>();
         let iter = map.prefix_iter_decoded(&())?.map(|(addr, ())| addr);
         Ok(iter.collect::<Vec<_>>())
     }
 }
-
-impl<'st, B: storage::Backend> DnsServerTransactionRo for DnsServerStoreTxRo<'st, B> {
-    fn close(self) {
-        self.0.close()
-    }
-}

dns_server/src/crawler_p2p/crawler_manager/tests/mod.rs

@@ -17,11 +17,11 @@ mod mock_manager;
 
 use std::time::Duration;
 
-use p2p::types::socket_address::SocketAddress;
+use p2p::{peer_manager::peerdb_common::Transactional, types::socket_address::SocketAddress};
 
 use crate::{
     crawler_p2p::crawler_manager::{
-        storage::{DnsServerStorageRead, DnsServerTransactional},
+        storage::DnsServerStorageRead,
         tests::mock_manager::{advance_time, test_crawler},
     },
     dns_server::DnsServerCommand,

p2p/src/peer_manager/mod.rs

@@ -18,6 +18,7 @@
 pub mod address_groups;
 pub mod peer_context;
 pub mod peerdb;
+pub mod peerdb_common;
 mod peers_eviction;
 
 use std::{

p2p/src/peer_manager/peerdb/mod.rs

@@ -49,7 +49,10 @@ use self::{
     storage_load::LoadedStorage,
 };
 
-use super::{address_groups::AddressGroup, ip_or_socket_address_to_peer_address};
+use super::{
+    address_groups::AddressGroup, ip_or_socket_address_to_peer_address,
+    peerdb_common::storage::update_db,
+};
 
 pub struct PeerDb<S> {
     /// P2P configuration
@@ -213,7 +216,7 @@ impl<S: PeerDbStorage> PeerDb<S> {
             let banned = now <= *banned_till;
 
             if !banned {
-                storage::update_db(&self.storage, |tx| {
+                update_db(&self.storage, |tx| {
                     tx.del_banned_address(&address.to_string())
                 })
                 .expect("removing banned address is expected to succeed");
@@ -284,13 +287,13 @@ impl<S: PeerDbStorage> PeerDb<S> {
 
         match (is_persistent_old, is_persistent_new) {
             (false, true) => {
-                storage::update_db(&self.storage, |tx| {
+                update_db(&self.storage, |tx| {
                     tx.add_known_address(&address.to_string())
                 })
                 .expect("adding address expected to succeed (peer_connected)");
             }
             (true, false) => {
-                storage::update_db(&self.storage, |tx| {
+                update_db(&self.storage, |tx| {
                     tx.del_known_address(&address.to_string())
                 })
                 .expect("adding address expected to succeed (peer_connected)");
@@ -326,7 +329,7 @@ impl<S: PeerDbStorage> PeerDb<S> {
     pub fn ban(&mut self, address: BannableAddress) {
         let ban_till = self.time_getter.get_time() + *self.p2p_config.ban_duration;
 
-        storage::update_db(&self.storage, |tx| {
+        update_db(&self.storage, |tx| {
             tx.add_banned_address(&address.to_string(), ban_till)
         })
         .expect("adding banned address is expected to succeed (ban_peer)");
@@ -335,7 +338,7 @@ impl<S: PeerDbStorage> PeerDb<S> {
     }
 
     pub fn unban(&mut self, address: &BannableAddress) {
-        storage::update_db(&self.storage, |tx| {
+        update_db(&self.storage, |tx| {
             tx.del_banned_address(&address.to_string())
         })
         .expect("adding banned address is expected to succeed (ban_peer)");
@@ -351,7 +354,7 @@ impl<S: PeerDbStorage> PeerDb<S> {
         if self.anchor_addresses == anchor_addresses {
             return;
         }
-        storage::update_db(&self.storage, |tx| {
+        update_db(&self.storage, |tx| {
             for address in self.anchor_addresses.difference(&anchor_addresses) {
                 log::debug!("remove anchor peer {address}");
                 tx.del_anchor_address(&address.to_string())?;