release resources correctly (#5287)

prysmaticlabs · Apr 2, 2020 · b47f456 · b47f456
1 parent ca81ef6
commit b47f456
Show file tree

Hide file tree

Showing 4 changed files with 45 additions and 28 deletions.
diff --git a/beacon-chain/p2p/discovery_test.go b/beacon-chain/p2p/discovery_test.go
@@ -106,6 +106,12 @@ func TestStartDiscV5_DiscoverAllPeers(t *testing.T) {
 		}
 		listeners = append(listeners, listener)
 	}
+	defer func() {
+		// Close down all peers.
+		for _, listener := range listeners {
+			listener.Close()
+		}
+	}()
 
 	// Wait for the nodes to have their local routing tables to be populated with the other nodes
 	time.Sleep(discoveryWaitTime)
@@ -116,11 +122,6 @@ func TestStartDiscV5_DiscoverAllPeers(t *testing.T) {
 		t.Errorf("The node's local table doesn't have the expected number of nodes. "+
 			"Expected more than or equal to %d but got %d", 4, len(nodes))
 	}
-
-	// Close all ports
-	for _, listener := range listeners {
-		listener.Close()
-	}
 }
 
 func TestMultiAddrsConversion_InvalidIPAddr(t *testing.T) {
@@ -187,14 +188,15 @@ func TestStaticPeering_PeersAreAdded(t *testing.T) {
 		t.Fatal(err)
 	}
 	s.Start()
+	defer func() {
+		if err := s.Stop(); err != nil {
+			t.Fatal(err)
+		}
+	}()
 
 	time.Sleep(10 * time.Second)
 	peers := s.host.Network().Peers()
 	if len(peers) != 5 {
 		t.Errorf("Not all peers added to peerstore, wanted %d but got %d", 5, len(peers))
 	}
-
-	if err := s.Stop(); err != nil {
-		t.Fatal(err)
-	}
 }
diff --git a/beacon-chain/p2p/fork_test.go b/beacon-chain/p2p/fork_test.go
@@ -66,6 +66,12 @@ func TestStartDiscv5_DifferentForkDigests(t *testing.T) {
 		}
 		listeners = append(listeners, listener)
 	}
+	defer func() {
+		// Close down all peers.
+		for _, listener := range listeners {
+			listener.Close()
+		}
+	}()
 
 	// Wait for the nodes to have their local routing tables to be populated with the other nodes
 	time.Sleep(discoveryWaitTime)
@@ -86,6 +92,7 @@ func TestStartDiscv5_DifferentForkDigests(t *testing.T) {
 	if err != nil {
 		t.Fatal(err)
 	}
+	defer s.Stop()
 	s.genesisTime = genesisTime
 	s.genesisValidatorsRoot = make([]byte, 32)
 	s.dv5Listener = lastListener
@@ -95,11 +102,6 @@ func TestStartDiscv5_DifferentForkDigests(t *testing.T) {
 	if len(multiAddrs) != 0 {
 		t.Errorf("Expected 0 valid peers, got %d", len(multiAddrs))
 	}
-
-	// Close down all peers.
-	for _, listener := range listeners {
-		listener.Close()
-	}
 }
 
 func TestStartDiscv5_SameForkDigests_DifferentNextForkData(t *testing.T) {
@@ -153,6 +155,12 @@ func TestStartDiscv5_SameForkDigests_DifferentNextForkData(t *testing.T) {
 		}
 		listeners = append(listeners, listener)
 	}
+	defer func() {
+		// Close down all peers.
+		for _, listener := range listeners {
+			listener.Close()
+		}
+	}()
 
 	// Wait for the nodes to have their local routing tables to be populated with the other nodes
 	time.Sleep(discoveryWaitTime)
@@ -174,6 +182,8 @@ func TestStartDiscv5_SameForkDigests_DifferentNextForkData(t *testing.T) {
 	if err != nil {
 		t.Fatal(err)
 	}
+	defer s.Stop()
+
 	s.genesisTime = genesisTime
 	s.genesisValidatorsRoot = make([]byte, 32)
 	s.dv5Listener = lastListener
@@ -182,11 +192,6 @@ func TestStartDiscv5_SameForkDigests_DifferentNextForkData(t *testing.T) {
 		t.Error("Expected to have valid peers, got 0")
 	}
 
-	// Close down all peers.
-	for _, listener := range listeners {
-		listener.Close()
-	}
-
 	testutil.AssertLogsContain(t, hook, "Peer matches fork digest but has different next fork epoch")
 }
 

diff --git a/beacon-chain/p2p/service_test.go b/beacon-chain/p2p/service_test.go
@@ -177,6 +177,12 @@ func TestListenForNewNodes(t *testing.T) {
 		listeners = append(listeners, listener)
 		hosts = append(hosts, h)
 	}
+	defer func() {
+		// Close down all peers.
+		for _, listener := range listeners {
+			listener.Close()
+		}
+	}()
 
 	// close peers upon exit of test
 	defer func() {
@@ -195,20 +201,17 @@ func TestListenForNewNodes(t *testing.T) {
 		t.Fatal(err)
 	}
 	s.Start()
+	defer func() {
+		if err := s.Stop(); err != nil {
+			t.Fatal(err)
+		}
+	}()
+
 	time.Sleep(2 * time.Second)
 	peers := s.host.Network().Peers()
 	if len(peers) != 5 {
 		t.Errorf("Not all peers added to peerstore, wanted %d but got %d", 5, len(peers))
 	}
-
-	// close down all peers
-	for _, listener := range listeners {
-		listener.Close()
-	}
-
-	if err := s.Stop(); err != nil {
-		t.Fatal(err)
-	}
 }
 
 func TestPeer_Disconnect(t *testing.T) {

diff --git a/beacon-chain/p2p/subnets_test.go b/beacon-chain/p2p/subnets_test.go
@@ -61,6 +61,12 @@ func TestStartDiscV5_DiscoverPeersWithSubnets(t *testing.T) {
 		listener.LocalNode().Set(entry)
 		listeners = append(listeners, listener)
 	}
+	defer func() {
+		// Close down all peers.
+		for _, listener := range listeners {
+			listener.Close()
+		}
+	}()
 
 	// Make one service on port 3001.
 	port = 4000
@@ -79,6 +85,7 @@ func TestStartDiscV5_DiscoverPeersWithSubnets(t *testing.T) {
 		t.Fatal(err)
 	}
 	s.Start()
+	defer s.Stop()
 
 	// Wait for the nodes to have their local routing tables to be populated with the other nodes
 	time.Sleep(2 * discoveryWaitTime)