From 90aedbddaa06f77af010daaac16a0ccfc81922fa Mon Sep 17 00:00:00 2001
From: Marco Munizaga <git@marcopolo.io>
Date: Wed, 22 Feb 2023 18:21:29 -0800
Subject: [PATCH] Wrap conn with metrics (#2131)

---
 p2p/net/swarm/swarm.go        | 22 ++++++++++++++++++++++
 p2p/net/swarm/swarm_conn.go   |  4 ----
 p2p/net/swarm/swarm_dial.go   |  6 +++---
 p2p/net/swarm/swarm_listen.go |  2 +-
 4 files changed, 26 insertions(+), 8 deletions(-)

diff --git a/p2p/net/swarm/swarm.go b/p2p/net/swarm/swarm.go
index 327f5d6e1..8f317f46f 100644
--- a/p2p/net/swarm/swarm.go
+++ b/p2p/net/swarm/swarm.go
@@ -641,3 +641,25 @@ func (s *Swarm) ResourceManager() network.ResourceManager {
 // Swarm is a Network.
 var _ network.Network = (*Swarm)(nil)
 var _ transport.TransportNetwork = (*Swarm)(nil)
+
+type connWithMetrics struct {
+	transport.CapableConn
+	opened        time.Time
+	dir           network.Direction
+	metricsTracer MetricsTracer
+}
+
+func wrapWithMetrics(capableConn transport.CapableConn, metricsTracer MetricsTracer, opened time.Time, dir network.Direction) connWithMetrics {
+	c := connWithMetrics{CapableConn: capableConn, opened: opened, dir: dir, metricsTracer: metricsTracer}
+	c.metricsTracer.OpenedConnection(c.dir, capableConn.RemotePublicKey(), capableConn.ConnState(), capableConn.LocalMultiaddr())
+	return c
+}
+
+func (c connWithMetrics) completedHandshake() {
+	c.metricsTracer.CompletedHandshake(time.Since(c.opened), c.ConnState(), c.LocalMultiaddr())
+}
+
+func (c connWithMetrics) Close() error {
+	c.metricsTracer.ClosedConnection(c.dir, time.Since(c.opened), c.ConnState(), c.LocalMultiaddr())
+	return c.CapableConn.Close()
+}
diff --git a/p2p/net/swarm/swarm_conn.go b/p2p/net/swarm/swarm_conn.go
index d9202fafa..406b100a5 100644
--- a/p2p/net/swarm/swarm_conn.go
+++ b/p2p/net/swarm/swarm_conn.go
@@ -60,10 +60,6 @@ func (c *Conn) Close() error {
 }
 
 func (c *Conn) doClose() {
-	if c.swarm.metricsTracer != nil {
-		c.swarm.metricsTracer.ClosedConnection(c.stat.Direction, time.Since(c.stat.Stats.Opened), c.ConnState(), c.LocalMultiaddr())
-	}
-
 	c.swarm.removeConn(c)
 
 	// Prevent new streams from opening.
diff --git a/p2p/net/swarm/swarm_dial.go b/p2p/net/swarm/swarm_dial.go
index 4f58a9b0e..e6e7a8e92 100644
--- a/p2p/net/swarm/swarm_dial.go
+++ b/p2p/net/swarm/swarm_dial.go
@@ -499,9 +499,9 @@ func (s *Swarm) dialAddr(ctx context.Context, p peer.ID, addr ma.Multiaddr) (tra
 	}
 	canonicallog.LogPeerStatus(100, connC.RemotePeer(), connC.RemoteMultiaddr(), "connection_status", "established", "dir", "outbound")
 	if s.metricsTracer != nil {
-		connState := connC.ConnState()
-		s.metricsTracer.OpenedConnection(network.DirOutbound, connC.RemotePublicKey(), connState, connC.LocalMultiaddr())
-		s.metricsTracer.CompletedHandshake(time.Since(start), connState, connC.LocalMultiaddr())
+		connWithMetrics := wrapWithMetrics(connC, s.metricsTracer, start, network.DirOutbound)
+		connWithMetrics.completedHandshake()
+		connC = connWithMetrics
 	}
 
 	// Trust the transport? Yeah... right.
diff --git a/p2p/net/swarm/swarm_listen.go b/p2p/net/swarm/swarm_listen.go
index 5ad2b08c3..d3e61aab3 100644
--- a/p2p/net/swarm/swarm_listen.go
+++ b/p2p/net/swarm/swarm_listen.go
@@ -131,7 +131,7 @@ func (s *Swarm) AddListenAddr(a ma.Multiaddr) error {
 			}
 			canonicallog.LogPeerStatus(100, c.RemotePeer(), c.RemoteMultiaddr(), "connection_status", "established", "dir", "inbound")
 			if s.metricsTracer != nil {
-				s.metricsTracer.OpenedConnection(network.DirInbound, c.RemotePublicKey(), c.ConnState(), c.LocalMultiaddr())
+				c = wrapWithMetrics(c, s.metricsTracer, time.Now(), network.DirOutbound)
 			}
 
 			log.Debugf("swarm listener accepted connection: %s <-> %s", c.LocalMultiaddr(), c.RemoteMultiaddr())