notify: make locking clearer

notif.go

@@ -45,13 +45,16 @@ func (nn *netNotifiee) Connected(n inet.Network, v inet.Conn) {
 		cancel:   cancel,
 	}
 
-	go func() {
-
 	// Note: We *could* just check the peerstore to see if the remote side supports the dht
 	// protocol, but its not clear that that information will make it into the peerstore
 	// by the time this notification is sent. So just to be very careful, we brute force this
 	// and open a new stream
+	go nn.testConnection(ctx, v)
+
+}
 
+func (nn *netNotifiee) testConnection(ctx context.Context, v inet.Conn) {
+	dht := nn.DHT()
 	for {
 		s, err := dht.host.NewStream(ctx, v.RemotePeer(), ProtocolDHT, ProtocolDHTOld)
 
@@ -59,12 +62,13 @@ func (nn *netNotifiee) Connected(n inet.Network, v inet.Conn) {
 		case nil:
 			s.Close()
 			dht.plk.Lock()
-				defer dht.plk.Unlock()
 
 			// Check if canceled under the lock.
 			if ctx.Err() == nil {
 				dht.Update(dht.Context(), v.RemotePeer())
 			}
+
+			dht.plk.Unlock()
 		case io.EOF:
 			if ctx.Err() == nil {
 				// Connection died but we may still have *an* open connection (context not canceled) so try again.
@@ -78,7 +82,6 @@ func (nn *netNotifiee) Connected(n inet.Network, v inet.Conn) {
 		}
 		return
 	}
-	}()
 }
 
 func (nn *netNotifiee) Disconnected(n inet.Network, v inet.Conn) {