improve connectedness check

p2p/net/swarm/swarm.go

@@ -384,11 +384,11 @@ func (s *Swarm) addConn(tc transport.CapableConn, dir network.Direction) (*Conn,
 		return nil, ErrSwarmClosed
 	}
 
-	oldState := s.connectednessUnlocked(p, true)
+	oldState := s.connectednessUnlocked(p)
 
 	c.streams.m = make(map[*Stream]struct{})
 	s.conns.m[p] = append(s.conns.m[p], c)
-	newState := s.connectednessUnlocked(p, true)
+	newState := s.connectednessUnlocked(p)
 
 	// Add two swarm refs:
 	// * One will be decremented after the close notifications fire in Conn.doClose
@@ -653,17 +653,38 @@ func (s *Swarm) Connectedness(p peer.ID) network.Connectedness {
 	s.conns.RLock()
 	defer s.conns.RUnlock()
 
-	return s.connectednessUnlocked(p, false)
+	return s.connectednessUnlocked(p)
 }
 
-// connectednessUnlocked returns the connectedness of a peer. For sending peer connectedness
-// changed notifications consider closed connections. When remote closes a connection
-// the underlying transport connection is closed first, so tracking changes to the connectedness
-// state requires considering this recently closed connections impact on Connectedness.
-func (s *Swarm) connectednessUnlocked(p peer.ID, considerClosed bool) network.Connectedness {
+// connectednessUnlocked returns the connectedness of a peer.
+func (s *Swarm) connectednessUnlocked(p peer.ID) network.Connectedness {
 	var haveLimited bool
 	for _, c := range s.conns.m[p] {
-		if !considerClosed && c.IsClosed() {
+		if c.IsClosed() {
+			// These will be garbage collected soon
+			continue
+		}
+		if c.Stat().Limited {
+			haveLimited = true
+		} else {
+			return network.Connected
+		}
+	}
+	if haveLimited {
+		return network.Limited
+	} else {
+		return network.NotConnected
+	}
+}
+
+// connectednessWithClosedConnUnlocked returns connectedness to peer assuming that closedConn is still
+// connected.
+// This is useful when the peer closes the Connection and we are interested in knowing what the old
+// connectedness before closing of the connection was.
+func (s *Swarm) connectednessWithClosedConnUnlocked(p peer.ID, closedConn *Conn) network.Connectedness {
+	var haveLimited bool
+	for _, c := range s.conns.m[p] {
+		if c.IsClosed() && c != closedConn {
 			// These will be garbage collected soon
 			continue
 		}
@@ -776,7 +797,9 @@ func (s *Swarm) removeConn(c *Conn) {
 
 	cs := s.conns.m[p]
 
-	oldState := s.connectednessUnlocked(p, true)
+	// To get an accurate oldState we need to consider the impact of the recently
+	// closed connection on connectedness
+	oldState := s.connectednessWithClosedConnUnlocked(p, c)
 
 	if len(cs) == 1 {
 		delete(s.conns.m, p)
@@ -794,7 +817,7 @@ func (s *Swarm) removeConn(c *Conn) {
 		}
 	}
 
-	newState := s.connectednessUnlocked(p, true)
+	newState := s.connectednessUnlocked(p)
 
 	s.conns.Unlock()