yamux: pass constructors for peer resource scopes to session construc…

…tor (#1739)
libp2p · Sep 8, 2022 · ff60c5d · ff60c5d
1 parent b52af7d
commit ff60c5d
Show file tree

Hide file tree

Showing 6 changed files with 24 additions and 9 deletions.
diff --git a/go.mod b/go.mod
@@ -29,7 +29,7 @@ require (
 	github.com/libp2p/go-netroute v0.2.0
 	github.com/libp2p/go-openssl v0.1.0
 	github.com/libp2p/go-reuseport v0.2.0
-	github.com/libp2p/go-yamux/v3 v3.1.2
+	github.com/libp2p/go-yamux/v4 v4.0.0
 	github.com/libp2p/zeroconf/v2 v2.2.0
 	github.com/lucas-clemente/quic-go v0.29.0
 	github.com/marten-seemann/tcp v0.0.0-20210406111302-dfbc87cc63fd

diff --git a/go.sum b/go.sum
@@ -312,8 +312,8 @@ github.com/libp2p/go-openssl v0.1.0/go.mod h1:OiOxwPpL3n4xlenjx2h7AwSGaFSC/KZvf6
 github.com/libp2p/go-reuseport v0.2.0 h1:18PRvIMlpY6ZK85nIAicSBuXXvrYoSw3dsBAR7zc560=
 github.com/libp2p/go-reuseport v0.2.0/go.mod h1:bvVho6eLMm6Bz5hmU0LYN3ixd3nPPvtIlaURZZgOY4k=
 github.com/libp2p/go-sockaddr v0.0.2/go.mod h1:syPvOmNs24S3dFVGJA1/mrqdeijPxLV2Le3BRLKd68k=
-github.com/libp2p/go-yamux/v3 v3.1.2 h1:lNEy28MBk1HavUAlzKgShp+F6mn/ea1nDYWftZhFW9Q=
-github.com/libp2p/go-yamux/v3 v3.1.2/go.mod h1:jeLEQgLXqE2YqX1ilAClIfCMDY+0uXQUKmmb/qp0gT4=
+github.com/libp2p/go-yamux/v4 v4.0.0 h1:+Y80dV2Yx/kv7Y7JKu0LECyVdMXm1VUoko+VQ9rBfZQ=
+github.com/libp2p/go-yamux/v4 v4.0.0/go.mod h1:NWjl8ZTLOGlozrXSOZ/HlfG++39iKNnM5wwmtQP1YB4=
 github.com/libp2p/zeroconf/v2 v2.2.0 h1:Cup06Jv6u81HLhIj1KasuNM/RHHrJ8T7wOTS4+Tv53Q=
 github.com/libp2p/zeroconf/v2 v2.2.0/go.mod h1:fuJqLnUwZTshS3U/bMRJ3+ow/v9oid1n0DmyYyNO1Xs=
 github.com/lucas-clemente/quic-go v0.29.0 h1:Vw0mGTfmWqGzh4jx/kMymsIkFK6rErFVmg+t9RLrnZE=

diff --git a/p2p/muxer/testsuite/mux.go b/p2p/muxer/testsuite/mux.go
@@ -71,8 +71,18 @@ func (p *peerScope) Check(t *testing.T) {
 	require.Zero(t, p.memory, "expected all reserved memory to have been released")
 }
 
+type peerScopeSpan struct {
+	peerScope
+}
+
+func (p *peerScopeSpan) Done() {
+	p.mx.Lock()
+	defer p.mx.Unlock()
+	p.memory = 0
+}
+
 func (p *peerScope) Stat() network.ScopeStat                       { return network.ScopeStat{} }
-func (p *peerScope) BeginSpan() (network.ResourceScopeSpan, error) { return nil, nil }
+func (p *peerScope) BeginSpan() (network.ResourceScopeSpan, error) { return &peerScopeSpan{}, nil }
 func (p *peerScope) Peer() peer.ID                                 { panic("implement me") }
 
 var _ network.PeerScope = &peerScope{}

diff --git a/p2p/muxer/yamux/conn.go b/p2p/muxer/yamux/conn.go
@@ -5,7 +5,7 @@ import (
 
 	"github.com/libp2p/go-libp2p/core/network"
 
-	"github.com/libp2p/go-yamux/v3"
+	"github.com/libp2p/go-yamux/v4"
 )
 
 // conn implements mux.MuxedConn over yamux.Session.

diff --git a/p2p/muxer/yamux/stream.go b/p2p/muxer/yamux/stream.go
@@ -5,7 +5,7 @@ import (
 
 	"github.com/libp2p/go-libp2p/core/network"
 
-	"github.com/libp2p/go-yamux/v3"
+	"github.com/libp2p/go-yamux/v4"
 )
 
 // stream implements mux.MuxedStream over yamux.Stream.

diff --git a/p2p/muxer/yamux/transport.go b/p2p/muxer/yamux/transport.go
@@ -7,7 +7,7 @@ import (
 
 	"github.com/libp2p/go-libp2p/core/network"
 
-	"github.com/libp2p/go-yamux/v3"
+	"github.com/libp2p/go-yamux/v4"
 )
 
 var DefaultTransport *Transport
@@ -38,12 +38,17 @@ type Transport yamux.Config
 var _ network.Multiplexer = &Transport{}
 
 func (t *Transport) NewConn(nc net.Conn, isServer bool, scope network.PeerScope) (network.MuxedConn, error) {
+	var newSpan func() (yamux.MemoryManager, error)
+	if scope != nil {
+		newSpan = func() (yamux.MemoryManager, error) { return scope.BeginSpan() }
+	}
+
 	var s *yamux.Session
 	var err error
 	if isServer {
-		s, err = yamux.Server(nc, t.Config(), scope)
+		s, err = yamux.Server(nc, t.Config(), newSpan)
 	} else {
-		s, err = yamux.Client(nc, t.Config(), scope)
+		s, err = yamux.Client(nc, t.Config(), newSpan)
 	}
 	if err != nil {
 		return nil, err