Merge pull request #109 from libp2p/rcmgr

use the resource manager
libp2p · Jan 18, 2022 · 2d29093 · 2d29093
2 parents 2e12799 + abde7f4
commit 2d29093
Show file tree

Hide file tree

Showing 5 changed files with 21 additions and 9 deletions.
diff --git a/p2p/transport/websocket/browser_integration_browser_test.go b/p2p/transport/websocket/browser_integration_browser_test.go
@@ -30,7 +30,7 @@ func TestInBrowser(t *testing.T) {
 	if err != nil {
 		t.Fatal(err)
 	}
-	tpt := New(u)
+	tpt := New(u, nil)
 	addr, err := ma.NewMultiaddr("/ip4/127.0.0.1/tcp/5555/ws")
 	if err != nil {
 		t.Fatal("could not parse multiaddress:" + err.Error())

diff --git a/p2p/transport/websocket/browser_integration_native_test.go b/p2p/transport/websocket/browser_integration_native_test.go
@@ -56,7 +56,7 @@ func TestInBrowser(t *testing.T) {
 		if err != nil {
 			t.Error("SERVER:", err)
 		}
-		tpt := New(u)
+		tpt := New(u, nil)
 		addr, err := ma.NewMultiaddr("/ip4/127.0.0.1/tcp/5555/ws")
 		if err != nil {
 			t.Error("SERVER:", err)

diff --git a/p2p/transport/websocket/websocket.go b/p2p/transport/websocket/websocket.go
@@ -29,11 +29,18 @@ var _ transport.Transport = (*WebsocketTransport)(nil)
 
 // WebsocketTransport is the actual go-libp2p transport
 type WebsocketTransport struct {
-	Upgrader transport.Upgrader
+	upgrader transport.Upgrader
+	rcmgr    network.ResourceManager
 }
 
-func New(u transport.Upgrader) *WebsocketTransport {
-	return &WebsocketTransport{u}
+func New(u transport.Upgrader, rcmgr network.ResourceManager) *WebsocketTransport {
+	if rcmgr == nil {
+		rcmgr = network.NullResourceManager
+	}
+	return &WebsocketTransport{
+		upgrader: u,
+		rcmgr:    rcmgr,
+	}
 }
 
 func (t *WebsocketTransport) CanDial(a ma.Multiaddr) bool {
@@ -49,9 +56,14 @@ func (t *WebsocketTransport) Proxy() bool {
 }
 
 func (t *WebsocketTransport) Dial(ctx context.Context, raddr ma.Multiaddr, p peer.ID) (transport.CapableConn, error) {
+	connScope, err := t.rcmgr.OpenConnection(network.DirOutbound, true)
+	if err != nil {
+		return nil, err
+	}
 	macon, err := t.maDial(ctx, raddr)
 	if err != nil {
+		connScope.Done()
 		return nil, err
 	}
-	return t.Upgrader.Upgrade(ctx, t, macon, network.DirOutbound, p)
+	return t.upgrader.Upgrade(ctx, t, macon, network.DirOutbound, p, connScope)
 }
diff --git a/p2p/transport/websocket/websocket_native.go b/p2p/transport/websocket/websocket_native.go
@@ -75,7 +75,7 @@ func (t *WebsocketTransport) Listen(a ma.Multiaddr) (transport.Listener, error)
 	if err != nil {
 		return nil, err
 	}
-	return t.Upgrader.UpgradeListener(t, malist), nil
+	return t.upgrader.UpgradeListener(t, malist), nil
 }
 
 func (t *WebsocketTransport) wrapListener(l net.Listener, origin *url.URL) (*listener, error) {

diff --git a/p2p/transport/websocket/websocket_test.go b/p2p/transport/websocket/websocket_test.go
@@ -47,12 +47,12 @@ func TestWebsocketTransport(t *testing.T) {
 	if err != nil {
 		t.Fatal(err)
 	}
-	ta := New(ua)
+	ta := New(ua, nil)
 	ub, err := tptu.New(newSecureMuxer(t, "peerB"), new(mplex.Transport))
 	if err != nil {
 		t.Fatal(err)
 	}
-	tb := New(ub)
+	tb := New(ub, nil)
 
 	zero := "/ip4/127.0.0.1/tcp/0/ws"
 	ttransport.SubtestTransport(t, ta, tb, zero, "peerA")