Merge #8882: [qa] Fix race conditions in p2p-compactblocks.py and sendheaders.py

b55d941 [qa] Fix race condition in sendheaders.py (Suhas Daftuar)
6976db2 [qa] Another attempt to fix race condition in p2p-compactblocks.py (Suhas Daftuar)

2 files changed, 48 insertions, 44 deletions

diff --git a/qa/rpc-tests/p2p-compactblocks.py b/qa/rpc-tests/p2p-compactblocks.py
index d91e10d77c..cdcfb36e7b 100755
--- a/qa/rpc-tests/p2p-compactblocks.py
+++ b/qa/rpc-tests/p2p-compactblocks.py
@@ -30,6 +30,10 @@ class TestNode(SingleNodeConnCB):
         self.last_getblocktxn = None
         self.last_block = None
         self.last_blocktxn = None
+        # Store the hashes of blocks we've seen announced.
+        # This is for synchronizing the p2p message traffic,
+        # so we can eg wait until a particular block is announced.
+        self.set_announced_blockhashes = set()
 
     def on_sendcmpct(self, conn, message):
         self.last_sendcmpct.append(message)
@@ -40,14 +44,22 @@ class TestNode(SingleNodeConnCB):
     def on_cmpctblock(self, conn, message):
         self.last_cmpctblock = message
         self.block_announced = True
+        self.last_cmpctblock.header_and_shortids.header.calc_sha256()
+        self.set_announced_blockhashes.add(self.last_cmpctblock.header_and_shortids.header.sha256)
 
     def on_headers(self, conn, message):
         self.last_headers = message
         self.block_announced = True
+        for x in self.last_headers.headers:
+            x.calc_sha256()
+            self.set_announced_blockhashes.add(x.sha256)
 
     def on_inv(self, conn, message):
         self.last_inv = message
-        self.block_announced = True
+        for x in self.last_inv.inv:
+            if x.type == 2:
+                self.block_announced = True
+                self.set_announced_blockhashes.add(x.hash)
 
     def on_getdata(self, conn, message):
         self.last_getdata = message
@@ -87,6 +99,12 @@ class TestNode(SingleNodeConnCB):
         assert(self.received_block_announcement())
         self.clear_block_announcement()
 
+    # Block until a block announcement for a particular block hash is
+    # received.
+    def wait_for_block_announcement(self, block_hash, timeout=30):
+        def received_hash():
+            return (block_hash in self.set_announced_blockhashes)
+        return wait_until(received_hash, timeout=timeout)
 
 class CompactBlocksTest(BitcoinTestFramework):
     def __init__(self):
@@ -278,7 +296,9 @@ class CompactBlocksTest(BitcoinTestFramework):
         if use_witness_address:
             assert(segwit_tx_generated) # check that our test is not broken
 
-        self.test_node.sync_with_ping()
+        # Wait until we've seen the block announcement for the resulting tip
+        tip = int(self.nodes[0].getbestblockhash(), 16)
+        assert(self.test_node.wait_for_block_announcement(tip))
 
         # Now mine a block, and look at the resulting compact block.
         test_node.clear_block_announcement()
diff --git a/qa/rpc-tests/sendheaders.py b/qa/rpc-tests/sendheaders.py
index c3f3180b6b..81b2442e6a 100755
--- a/qa/rpc-tests/sendheaders.py
+++ b/qa/rpc-tests/sendheaders.py
@@ -80,20 +80,19 @@ e. Announce one more that doesn't connect.
    Expect: disconnect.
 '''
 
-class BaseNode(NodeConnCB):
+direct_fetch_response_time = 0.05
+
+class BaseNode(SingleNodeConnCB):
     def __init__(self):
-        NodeConnCB.__init__(self)
-        self.connection = None
+        SingleNodeConnCB.__init__(self)
         self.last_inv = None
         self.last_headers = None
         self.last_block = None
-        self.ping_counter = 1
-        self.last_pong = msg_pong(0)
         self.last_getdata = None
-        self.sleep_time = 0.05
         self.block_announced = False
         self.last_getheaders = None
         self.disconnected = False
+        self.last_blockhash_announced = None
 
     def clear_last_announcement(self):
         with mininode_lock:
@@ -101,9 +100,6 @@ class BaseNode(NodeConnCB):
             self.last_inv = None
             self.last_headers = None
 
-    def add_connection(self, conn):
-        self.connection = conn
-
     # Request data for a list of block hashes
     def get_data(self, block_hashes):
         msg = msg_getdata()
@@ -122,17 +118,17 @@ class BaseNode(NodeConnCB):
         msg.inv = [CInv(2, blockhash)]
         self.connection.send_message(msg)
 
-    # Wrapper for the NodeConn's send_message function
-    def send_message(self, message):
-        self.connection.send_message(message)
-
     def on_inv(self, conn, message):
         self.last_inv = message
         self.block_announced = True
+        self.last_blockhash_announced = message.inv[-1].hash
 
     def on_headers(self, conn, message):
         self.last_headers = message
-        self.block_announced = True
+        if len(message.headers):
+            self.block_announced = True
+            message.headers[-1].calc_sha256()
+            self.last_blockhash_announced = message.headers[-1].sha256
 
     def on_block(self, conn, message):
         self.last_block = message.block
@@ -141,9 +137,6 @@ class BaseNode(NodeConnCB):
     def on_getdata(self, conn, message):
         self.last_getdata = message
 
-    def on_pong(self, conn, message):
-        self.last_pong = message
-
     def on_getheaders(self, conn, message):
         self.last_getheaders = message
 
@@ -157,7 +150,7 @@ class BaseNode(NodeConnCB):
         expect_headers = headers if headers != None else []
         expect_inv = inv if inv != None else []
         test_function = lambda: self.block_announced
-        self.sync(test_function)
+        assert(wait_until(test_function, timeout=60))
         with mininode_lock:
             self.block_announced = False
 
@@ -180,30 +173,14 @@ class BaseNode(NodeConnCB):
         return success
 
     # Syncing helpers
-    def sync(self, test_function, timeout=60):
-        while timeout > 0:
-            with mininode_lock:
-                if test_function():
-                    return
-            time.sleep(self.sleep_time)
-            timeout -= self.sleep_time
-        raise AssertionError("Sync failed to complete")
-        
-    def sync_with_ping(self, timeout=60):
-        self.send_message(msg_ping(nonce=self.ping_counter))
-        test_function = lambda: self.last_pong.nonce == self.ping_counter
-        self.sync(test_function, timeout)
-        self.ping_counter += 1
-        return
-
     def wait_for_block(self, blockhash, timeout=60):
         test_function = lambda: self.last_block != None and self.last_block.sha256 == blockhash
-        self.sync(test_function, timeout)
+        assert(wait_until(test_function, timeout=timeout))
         return
 
     def wait_for_getheaders(self, timeout=60):
         test_function = lambda: self.last_getheaders != None
-        self.sync(test_function, timeout)
+        assert(wait_until(test_function, timeout=timeout))
         return
 
     def wait_for_getdata(self, hash_list, timeout=60):
@@ -211,12 +188,17 @@ class BaseNode(NodeConnCB):
             return
 
         test_function = lambda: self.last_getdata != None and [x.hash for x in self.last_getdata.inv] == hash_list
-        self.sync(test_function, timeout)
+        assert(wait_until(test_function, timeout=timeout))
         return
 
     def wait_for_disconnect(self, timeout=60):
         test_function = lambda: self.disconnected
-        self.sync(test_function, timeout)
+        assert(wait_until(test_function, timeout=timeout))
+        return
+
+    def wait_for_block_announcement(self, block_hash, timeout=60):
+        test_function = lambda: self.last_blockhash_announced == block_hash
+        assert(wait_until(test_function, timeout=timeout))
         return
 
     def send_header_for_blocks(self, new_blocks):
@@ -266,7 +248,9 @@ class SendHeadersTest(BitcoinTestFramework):
     def mine_reorg(self, length):
         self.nodes[0].generate(length) # make sure all invalidated blocks are node0's
         sync_blocks(self.nodes, wait=0.1)
-        [x.clear_last_announcement() for x in self.p2p_connections]
+        for x in self.p2p_connections:
+            x.wait_for_block_announcement(int(self.nodes[0].getbestblockhash(), 16))
+            x.clear_last_announcement()
 
         tip_height = self.nodes[1].getblockcount()
         hash_to_invalidate = self.nodes[1].getblockhash(tip_height-(length-1))
@@ -495,7 +479,7 @@ class SendHeadersTest(BitcoinTestFramework):
 
         test_node.send_header_for_blocks(blocks)
         test_node.sync_with_ping()
-        test_node.wait_for_getdata([x.sha256 for x in blocks], timeout=test_node.sleep_time)
+        test_node.wait_for_getdata([x.sha256 for x in blocks], timeout=direct_fetch_response_time)
 
         [ test_node.send_message(msg_block(x)) for x in blocks ]
 
@@ -526,13 +510,13 @@ class SendHeadersTest(BitcoinTestFramework):
         # both blocks (same work as tip)
         test_node.send_header_for_blocks(blocks[1:2])
         test_node.sync_with_ping()
-        test_node.wait_for_getdata([x.sha256 for x in blocks[0:2]], timeout=test_node.sleep_time)
+        test_node.wait_for_getdata([x.sha256 for x in blocks[0:2]], timeout=direct_fetch_response_time)
 
         # Announcing 16 more headers should trigger direct fetch for 14 more
         # blocks
         test_node.send_header_for_blocks(blocks[2:18])
         test_node.sync_with_ping()
-        test_node.wait_for_getdata([x.sha256 for x in blocks[2:16]], timeout=test_node.sleep_time)
+        test_node.wait_for_getdata([x.sha256 for x in blocks[2:16]], timeout=direct_fetch_response_time)
 
         # Announcing 1 more header should not trigger any response
         test_node.last_getdata = None