Wrap conn with metrics (#2131)

libp2p · Feb 23, 2023 · 90aedbd · 90aedbd
1 parent 7c762f1
commit 90aedbd
Show file tree

Hide file tree

Showing 4 changed files with 26 additions and 8 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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())