fix(share/discovery): add loop to read events from the channel multip…

…le times (#1684)
celestiaorg · Feb 2, 2023 · 484933e · 484933e
1 parent a60c24f
commit 484933e
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 20 deletions.
diff --git a/nodebuilder/tests/p2p_test.go b/nodebuilder/tests/p2p_test.go
@@ -226,6 +226,7 @@ func TestRestartNodeDiscovery(t *testing.T) {
 		select {
 		case <-ctx.Done():
 			require.True(t, nodes[0].Host.Network().Connectedness(node.Host.ID()) == network.Connected)
+			return
 		case conn := <-connectSub.Out():
 			status := conn.(event.EvtPeerConnectednessChanged)
 			if status.Peer != node.Host.ID() {

diff --git a/share/availability/discovery/discovery.go b/share/availability/discovery/discovery.go
@@ -143,27 +143,29 @@ func (d *Discovery) ensurePeers(ctx context.Context) {
 	// starting to listen to subscriptions async will help us to avoid any blocking
 	// in the case when we will not have the needed amount of FNs and will be blocked in `FindPeers`.
 	go func() {
-		select {
-		case <-ctx.Done():
-			log.Debug("Context canceled. Finish listening for connectedness events.")
-			return
-		case e, ok := <-sub.Out():
-			if !ok {
-				log.Debug("Subscription for connectedness events is closed.")
+		for {
+			select {
+			case <-ctx.Done():
+				log.Debug("Context canceled. Finish listening for connectedness events.")
 				return
-			}
-			// listen to disconnect event to remove peer from set and reset backoff time
-			// reset timer in order to restart the discovery, once stored peer is disconnected
-			connStatus := e.(event.EvtPeerConnectednessChanged)
-			if connStatus.Connectedness == network.NotConnected {
-				if d.set.Contains(connStatus.Peer) {
-					log.Debugw("removing the peer from the peer set",
-						"peer", connStatus.Peer, "status", connStatus.Connectedness.String())
-					d.connector.RestartBackoff(connStatus.Peer)
-					d.set.Remove(connStatus.Peer)
-					d.onUpdatedPeers(connStatus.Peer, false)
-					d.host.ConnManager().UntagPeer(connStatus.Peer, topic)
-					t.Reset(d.discoveryInterval)
+			case e, ok := <-sub.Out():
+				if !ok {
+					log.Debug("Subscription for connectedness events is closed.")
+					return
+				}
+				// listen to disconnect event to remove peer from set and reset backoff time
+				// reset timer in order to restart the discovery, once stored peer is disconnected
+				connStatus := e.(event.EvtPeerConnectednessChanged)
+				if connStatus.Connectedness == network.NotConnected {
+					if d.set.Contains(connStatus.Peer) {
+						log.Debugw("removing the peer from the peer set",
+							"peer", connStatus.Peer, "status", connStatus.Connectedness.String())
+						d.connector.RestartBackoff(connStatus.Peer)
+						d.set.Remove(connStatus.Peer)
+						d.onUpdatedPeers(connStatus.Peer, false)
+						d.host.ConnManager().UntagPeer(connStatus.Peer, topic)
+						t.Reset(d.discoveryInterval)
+					}
 				}
 			}
 		}