Do addKnownNode recursively for multiple streams

src/knownnodes.py

@@ -8,12 +8,16 @@ import os
 import pickle
 import threading
 import time
+try:
+    from collections.abc import Iterable
+except ImportError:
+    from collections import Iterable
 
 import state
 from bmconfigparser import BMConfigParser
 from network.node import Peer
 
-knownNodesLock = threading.Lock()
+knownNodesLock = threading.RLock()
 """Thread lock for knownnodes modification"""
 knownNodes = {stream: {} for stream in range(1, 4)}
 """The dict of known nodes for each stream"""
@@ -95,9 +99,17 @@ def saveKnownNodes(dirName=None):
 
 def addKnownNode(stream, peer, lastseen=None, is_self=False):
     """
-    Add a new node to the dict or update lastseen if it already exists
+    Add a new node to the dict or update lastseen if it already exists.
+    Do it for each stream number if *stream* is `Iterable`.
+    Returns True if added a new node.
     """
     # pylint: disable=too-many-branches
+    if isinstance(stream, Iterable):
+        with knownNodesLock:
+            for s in stream:
+                addKnownNode(s, peer, lastseen, is_self)
+        return
+
     rating = 0.0
     if not lastseen:
         # FIXME: maybe about 28 days?

src/network/tcp.py

@@ -148,9 +148,8 @@ class TCPConnection(BMProto, TLSDispatcher):
         # The connection having host suitable for knownnodes
         if self.isOutbound or not self.local and not state.socksIP:
             knownnodes.increaseRating(self.destination)
-            with knownnodes.knownNodesLock:
-                for s in self.streams:
-                    knownnodes.addKnownNode(s, self.destination, time.time())
+            knownnodes.addKnownNode(
+                self.streams, self.destination, time.time())
             Dandelion().maybeAddStem(self)
         self.sendAddr()
         self.sendBigInv()
@@ -273,10 +272,8 @@ class TCPConnection(BMProto, TLSDispatcher):
                 (self.isOutbound, False, self.destination)
             ))
             if host_is_global:
-                with knownnodes.knownNodesLock:
-                    for s in self.streams:
                 knownnodes.addKnownNode(
-                            s, self.destination, time.time())
+                    self.streams, self.destination, time.time())
                 Dandelion().maybeRemoveStem(self)
         elif host_is_global:
             knownnodes.decreaseRating(self.destination)