From fd9c10c2e23247fc14f5e58a03e9626639b091f0 Mon Sep 17 00:00:00 2001
From: behzad nouri <behzadnouri@gmail.com>
Date: Sat, 10 Jul 2021 14:13:52 +0000
Subject: [PATCH] adds a generic implementation of Gossip{Read,Write}Lock
 (#18559)

---
 gossip/src/cluster_info.rs         | 86 +++---------------------------
 gossip/src/cluster_info_metrics.rs | 36 +++++++++++++
 2 files changed, 44 insertions(+), 78 deletions(-)

diff --git a/gossip/src/cluster_info.rs b/gossip/src/cluster_info.rs
index 03a8224a9..ee338a66f 100644
--- a/gossip/src/cluster_info.rs
+++ b/gossip/src/cluster_info.rs
@@ -14,7 +14,9 @@
 //! Bank needs to provide an interface for us to query the stake weight
 use {
     crate::{
-        cluster_info_metrics::{submit_gossip_stats, Counter, GossipStats, ScopedTimer},
+        cluster_info_metrics::{
+            submit_gossip_stats, Counter, GossipStats, ScopedTimer, TimedGuard,
+        },
         contact_info::ContactInfo,
         crds::{Crds, Cursor},
         crds_gossip::CrdsGossip,
@@ -73,7 +75,7 @@ use {
         io::BufReader,
         iter::repeat,
         net::{IpAddr, Ipv4Addr, SocketAddr, TcpListener, UdpSocket},
-        ops::{Deref, DerefMut, Div},
+        ops::{Deref, Div},
         path::{Path, PathBuf},
         result::Result,
         sync::{
@@ -137,78 +139,6 @@ pub enum ClusterInfoError {
     BadGossipAddress,
 }
 
-struct GossipWriteLock<'a> {
-    gossip: RwLockWriteGuard<'a, CrdsGossip>,
-    timer: Measure,
-    counter: &'a Counter,
-}
-
-impl<'a> GossipWriteLock<'a> {
-    fn new(
-        gossip: RwLockWriteGuard<'a, CrdsGossip>,
-        label: &'static str,
-        counter: &'a Counter,
-    ) -> Self {
-        Self {
-            gossip,
-            timer: Measure::start(label),
-            counter,
-        }
-    }
-}
-
-impl<'a> Deref for GossipWriteLock<'a> {
-    type Target = RwLockWriteGuard<'a, CrdsGossip>;
-    fn deref(&self) -> &Self::Target {
-        &self.gossip
-    }
-}
-
-impl<'a> DerefMut for GossipWriteLock<'a> {
-    fn deref_mut(&mut self) -> &mut Self::Target {
-        &mut self.gossip
-    }
-}
-
-impl<'a> Drop for GossipWriteLock<'a> {
-    fn drop(&mut self) {
-        self.counter.add_measure(&mut self.timer);
-    }
-}
-
-struct GossipReadLock<'a> {
-    gossip: RwLockReadGuard<'a, CrdsGossip>,
-    timer: Measure,
-    counter: &'a Counter,
-}
-
-impl<'a> GossipReadLock<'a> {
-    fn new(
-        gossip: RwLockReadGuard<'a, CrdsGossip>,
-        label: &'static str,
-        counter: &'a Counter,
-    ) -> Self {
-        Self {
-            gossip,
-            timer: Measure::start(label),
-            counter,
-        }
-    }
-}
-
-impl<'a> Deref for GossipReadLock<'a> {
-    type Target = RwLockReadGuard<'a, CrdsGossip>;
-    fn deref(&self) -> &Self::Target {
-        &self.gossip
-    }
-}
-
-impl<'a> Drop for GossipReadLock<'a> {
-    fn drop(&mut self) {
-        self.counter.add_measure(&mut self.timer);
-    }
-}
-
 pub struct ClusterInfo {
     /// The network
     pub gossip: RwLock<CrdsGossip>,
@@ -986,16 +916,16 @@ impl ClusterInfo {
         &'a self,
         label: &'static str,
         counter: &'a Counter,
-    ) -> GossipReadLock<'a> {
-        GossipReadLock::new(self.gossip.read().unwrap(), label, counter)
+    ) -> TimedGuard<'a, RwLockReadGuard<CrdsGossip>> {
+        TimedGuard::new(self.gossip.read().unwrap(), label, counter)
     }
 
     fn time_gossip_write_lock<'a>(
         &'a self,
         label: &'static str,
         counter: &'a Counter,
-    ) -> GossipWriteLock<'a> {
-        GossipWriteLock::new(self.gossip.write().unwrap(), label, counter)
+    ) -> TimedGuard<'a, RwLockWriteGuard<CrdsGossip>> {
+        TimedGuard::new(self.gossip.write().unwrap(), label, counter)
     }
 
     pub fn push_message(&self, message: CrdsValue) {
diff --git a/gossip/src/cluster_info_metrics.rs b/gossip/src/cluster_info_metrics.rs
index 0401b17a3..b3c228ed6 100644
--- a/gossip/src/cluster_info_metrics.rs
+++ b/gossip/src/cluster_info_metrics.rs
@@ -4,6 +4,7 @@ use {
     solana_sdk::pubkey::Pubkey,
     std::{
         collections::HashMap,
+        ops::{Deref, DerefMut},
         sync::{
             atomic::{AtomicU64, Ordering},
             RwLock,
@@ -28,6 +29,12 @@ impl Counter {
     }
 }
 
+pub(crate) struct TimedGuard<'a, T> {
+    guard: T,
+    timer: Measure,
+    counter: &'a Counter,
+}
+
 pub(crate) struct ScopedTimer<'a> {
     clock: Instant,
     metric: &'a AtomicU64,
@@ -52,6 +59,35 @@ impl Drop for ScopedTimer<'_> {
     }
 }
 
+impl<'a, T> TimedGuard<'a, T> {
+    pub(crate) fn new(guard: T, label: &'static str, counter: &'a Counter) -> Self {
+        Self {
+            guard,
+            timer: Measure::start(label),
+            counter,
+        }
+    }
+}
+
+impl<'a, T> Deref for TimedGuard<'a, T> {
+    type Target = T;
+    fn deref(&self) -> &Self::Target {
+        &self.guard
+    }
+}
+
+impl<'a, T> DerefMut for TimedGuard<'a, T> {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        &mut self.guard
+    }
+}
+
+impl<'a, T> Drop for TimedGuard<'a, T> {
+    fn drop(&mut self) {
+        self.counter.add_measure(&mut self.timer);
+    }
+}
+
 #[derive(Default)]
 pub(crate) struct GossipStats {
     pub(crate) all_tvu_peers: Counter,