package transport_integration import ( "context" "fmt" "strings" "sync" "sync/atomic" "testing" "time" "github.com/libp2p/go-libp2p/core/host" "github.com/libp2p/go-libp2p/core/network" mocknetwork "github.com/libp2p/go-libp2p/core/network/mocks" "github.com/libp2p/go-libp2p/core/peer" "github.com/libp2p/go-libp2p/p2p/protocol/identify" "github.com/libp2p/go-libp2p/p2p/protocol/ping" "github.com/stretchr/testify/require" "go.uber.org/mock/gomock" ) func TestResourceManagerIsUsed(t *testing.T) { for _, tc := range transportsToTest { t.Run(tc.Name, func(t *testing.T) { for _, testDialer := range []bool{true, false} { t.Run(tc.Name+fmt.Sprintf(" test_dialer=%v", testDialer), func(t *testing.T) { var reservedMemory, releasedMemory atomic.Int32 defer func() { require.Equal(t, reservedMemory.Load(), releasedMemory.Load()) require.NotEqual(t, 0, reservedMemory.Load()) }() ctrl := gomock.NewController(t) defer ctrl.Finish() rcmgr := mocknetwork.NewMockResourceManager(ctrl) rcmgr.EXPECT().Close() var listener, dialer host.Host var expectedPeer peer.ID var expectedDir network.Direction var expectedAddr interface{} if testDialer { listener = tc.HostGenerator(t, TransportTestCaseOpts{NoRcmgr: true}) dialer = tc.HostGenerator(t, TransportTestCaseOpts{NoListen: true, ResourceManager: rcmgr}) expectedPeer = listener.ID() expectedDir = network.DirOutbound expectedAddr = listener.Addrs()[0] } else { listener = tc.HostGenerator(t, TransportTestCaseOpts{ResourceManager: rcmgr}) dialer = tc.HostGenerator(t, TransportTestCaseOpts{NoListen: true, NoRcmgr: true}) expectedPeer = dialer.ID() expectedDir = network.DirInbound expectedAddr = gomock.Any() } expectFd := true if strings.Contains(tc.Name, "QUIC") || strings.Contains(tc.Name, "WebTransport") || strings.Contains(tc.Name, "WebRTC") { expectFd = false } peerScope := mocknetwork.NewMockPeerScope(ctrl) peerScope.EXPECT().ReserveMemory(gomock.Any(), gomock.Any()).AnyTimes().Do(func(amount int, pri uint8) { reservedMemory.Add(int32(amount)) }) peerScope.EXPECT().ReleaseMemory(gomock.Any()).AnyTimes().Do(func(amount int) { releasedMemory.Add(int32(amount)) }) peerScope.EXPECT().BeginSpan().AnyTimes().DoAndReturn(func() (network.ResourceScopeSpan, error) { s := mocknetwork.NewMockResourceScopeSpan(ctrl) s.EXPECT().BeginSpan().AnyTimes().Return(mocknetwork.NewMockResourceScopeSpan(ctrl), nil) // No need to track these memory reservations since we assert that Done is called s.EXPECT().ReserveMemory(gomock.Any(), gomock.Any()) s.EXPECT().Done() return s, nil }) var calledSetPeer atomic.Bool connScope := mocknetwork.NewMockConnManagementScope(ctrl) connScope.EXPECT().SetPeer(expectedPeer).Do(func(peer.ID) { calledSetPeer.Store(true) }) connScope.EXPECT().PeerScope().AnyTimes().DoAndReturn(func() network.PeerScope { if calledSetPeer.Load() { return peerScope } return nil }) if tc.Name == "WebRTC" { // webrtc receive buffer is a fix sized buffer allocated up front connScope.EXPECT().ReserveMemory(gomock.Any(), gomock.Any()) } connScope.EXPECT().Done().MinTimes(1) var allStreamsDone sync.WaitGroup rcmgr.EXPECT().OpenConnection(expectedDir, expectFd, expectedAddr).Return(connScope, nil) rcmgr.EXPECT().OpenStream(expectedPeer, gomock.Any()).AnyTimes().DoAndReturn(func(id peer.ID, dir network.Direction) (network.StreamManagementScope, error) { allStreamsDone.Add(1) streamScope := mocknetwork.NewMockStreamManagementScope(ctrl) // No need to track these memory reservations since we assert that Done is called streamScope.EXPECT().ReserveMemory(gomock.Any(), gomock.Any()).AnyTimes() streamScope.EXPECT().ReleaseMemory(gomock.Any()).AnyTimes() streamScope.EXPECT().BeginSpan().AnyTimes().DoAndReturn(func() (network.ResourceScopeSpan, error) { s := mocknetwork.NewMockResourceScopeSpan(ctrl) s.EXPECT().BeginSpan().AnyTimes().Return(mocknetwork.NewMockResourceScopeSpan(ctrl), nil) s.EXPECT().Done() return s, nil }) streamScope.EXPECT().SetService(gomock.Any()).MaxTimes(1) streamScope.EXPECT().SetProtocol(gomock.Any()) streamScope.EXPECT().Done().Do(func() { allStreamsDone.Done() }) return streamScope, nil }) require.NoError(t, dialer.Connect(context.Background(), peer.AddrInfo{ ID: listener.ID(), Addrs: listener.Addrs(), })) // Wait for any in progress identifies to finish. // We shouldn't have to do this, but basic host currently // always does an identify. <-dialer.(interface{ IDService() identify.IDService }).IDService().IdentifyWait(dialer.Network().ConnsToPeer(listener.ID())[0]) <-listener.(interface{ IDService() identify.IDService }).IDService().IdentifyWait(listener.Network().ConnsToPeer(dialer.ID())[0]) <-ping.Ping(context.Background(), dialer, listener.ID()) err := dialer.Network().ClosePeer(listener.ID()) require.NoError(t, err) // Wait a bit for any pending .Adds before we call .Wait to avoid a data race. // This shouldn't be necessary since it should be impossible // for an OpenStream to happen *after* a ClosePeer, however // in practice it does and leads to test flakiness. time.Sleep(10 * time.Millisecond) allStreamsDone.Wait() dialer.Close() listener.Close() }) } }) } }