Fix mininode disconnections to work with select

TrezarCoin · Apr 30, 2015 · ef32817 · ef32817
1 parent f026ab6
commit ef32817
Showing 1 changed file with 20 additions and 11 deletions.
diff --git a/qa/rpc-tests/mininode.py b/qa/rpc-tests/mininode.py
@@ -37,6 +37,11 @@
 
 MAX_INV_SZ = 50000
 
+# Keep our own socket map for asyncore, so that we can track disconnects
+# ourselves (to workaround an issue with closing an asyncore socket when 
+# using select)
+mininode_socket_map = dict()
+
 # Serialization/deserialization tools
 def sha256(s):
     return hashlib.new('sha256', s).digest()
@@ -1076,7 +1081,7 @@ class NodeConn(asyncore.dispatcher):
     }
 
     def __init__(self, dstaddr, dstport, rpc, callback, net="regtest"):
-        asyncore.dispatcher.__init__(self)
+        asyncore.dispatcher.__init__(self, map=mininode_socket_map)
         self.log = logging.getLogger("NodeConn(%s:%d)" % (dstaddr, dstport))
         self.dstaddr = dstaddr
         self.dstport = dstport
@@ -1140,14 +1145,10 @@ def readable(self):
         return True
 
     def writable(self):
-        if self.disconnect:
-            self.handle_close()
-            return False
-        else:
-            self.sendbufLock.acquire()
-            length = len(self.sendbuf)
-            self.sendbufLock.release()
-            return (length > 0)
+        self.sendbufLock.acquire()
+        length = len(self.sendbuf)
+        self.sendbufLock.release()
+        return (length > 0)
 
     def handle_write(self):
         self.sendbufLock.acquire()
@@ -1229,12 +1230,20 @@ def got_message(self, message):
 
     def disconnect_node(self):
         self.disconnect = True
-        self.send_message(self.messagemap['ping']())
 
 
 class NetworkThread(Thread):
     def run(self):
-        asyncore.loop(0.1, True)
+        while mininode_socket_map:
+            # We check for whether to disconnect outside of the asyncore
+            # loop to workaround the behavior of asyncore when using
+            # select
+            disconnected = []
+            for fd, obj in mininode_socket_map.items():
+                if obj.disconnect:
+                    disconnected.append(obj)
+            [ obj.handle_close() for obj in disconnected ]
+            asyncore.loop(0.1, use_poll=True, map=mininode_socket_map, count=1)
 
 
 # An exception we can raise if we detect a potential disconnect