def test_aiter_forker(self): q = push_aiter() forker = aiter_forker(q) q.push(1, 2, 3, 4, 5) r0 = run(get_n(forker, 3)) f2 = forker.fork() q.push(*range(7, 14)) q.stop() r1 = run(get_n(forker)) r2 = run(get_n(f2)) self.assertEqual(r0, [1, 2, 3]) self.assertEqual(r1, [4, 5, 7, 8, 9, 10, 11, 12, 13]) self.assertEqual(r2, [4, 5, 7, 8, 9, 10, 11, 12, 13])
def test_aiter_forker_multiple_active(self): """ Multiple forks of an aiter_forker both asking for empty q information at the same time. Make sure the second one doesn't block. """ q = push_aiter() forker = aiter_forker(q) fork_1 = forker.fork(is_active=True) fork_2 = forker.fork(is_active=True) f1 = asyncio.ensure_future(get_n(fork_1, 1)) f2 = asyncio.ensure_future(get_n(fork_2, 1)) run(asyncio.wait([f1, f2], timeout=0.1)) self.assertFalse(f1.done()) self.assertFalse(f2.done()) q.push(1) run(asyncio.wait([f1, f2], timeout=0.1)) self.assertTrue(f1.done()) self.assertTrue(f2.done()) r1 = run(f1) r2 = run(f2) self.assertEqual(r1, [1]) self.assertEqual(r2, [1])
def initialize_pipeline(self, aiter, api: Any, server_port: int) -> asyncio.Task: """ A pipeline that starts with (StreamReader, StreamWriter), maps it though to connections, messages, executes a local API call, and returns responses. """ # Maps a stream reader, writer and NodeType to a Connection object connections_aiter = map_aiter( partial_func.partial_async(self.stream_reader_writer_to_connection, server_port), aiter, ) # Performs a handshake with the peer handshaked_connections_aiter = join_aiters( map_aiter(self.perform_handshake, connections_aiter)) forker = aiter_forker(handshaked_connections_aiter) handshake_finished_1 = forker.fork(is_active=True) handshake_finished_2 = forker.fork(is_active=True) # Reads messages one at a time from the TCP connection messages_aiter = join_aiters( map_aiter(self.connection_to_message, handshake_finished_1, 100)) # Handles each message one at a time, and yields responses to send back or broadcast responses_aiter = join_aiters( map_aiter( partial_func.partial_async_gen(self.handle_message, api), messages_aiter, 100, )) # Uses a forked aiter, and calls the on_connect function to send some initial messages # as soon as the connection is established on_connect_outbound_aiter = join_aiters( map_aiter(self.connection_to_outbound, handshake_finished_2, 100)) # Also uses the instance variable _outbound_aiter, which clients can use to send messages # at any time, not just on_connect. outbound_aiter_mapped = map_aiter(lambda x: (None, x), self._outbound_aiter) responses_aiter = join_aiters( iter_to_aiter([ responses_aiter, on_connect_outbound_aiter, outbound_aiter_mapped ])) # For each outbound message, replicate for each peer that we need to send to expanded_messages_aiter = join_aiters( map_aiter(self.expand_outbound_messages, responses_aiter, 100)) # This will run forever. Sends each message through the TCP connection, using the # length encoding and CBOR serialization async def serve_forever(): async for connection, message in expanded_messages_aiter: log.info( f"-> {message.function} to peer {connection.get_peername()}" ) try: await connection.send(message) except ( ConnectionResetError, BrokenPipeError, RuntimeError, TimeoutError, ) as e: log.error( f"Cannot write to {connection}, already closed. Error {e}." ) self.global_connections.close(connection, True) # We will return a task for this, so user of start_chia_server or start_chia_client can wait until # the server is closed. return asyncio.get_running_loop().create_task(serve_forever())
async def initialize_pipeline( srwt_aiter, api: Any, server_port: int, outbound_aiter: push_aiter, global_connections: PeerConnections, local_type: NodeType, node_id: bytes32, network_id: bytes32, log: logging.Logger, ): """ A pipeline that starts with (StreamReader, StreamWriter), maps it though to connections, messages, executes a local API call, and returns responses. """ # Maps a stream reader, writer and NodeType to a Connection object connections_aiter = map_aiter( partial_func.partial_async( stream_reader_writer_to_connection, server_port, local_type, log, ), join_aiters(srwt_aiter), ) def add_global_connections(connection): return connection, global_connections connections_with_global_connections_aiter = map_aiter( add_global_connections, connections_aiter ) # Performs a handshake with the peer outbound_handshake = Message( "handshake", Handshake( network_id, protocol_version, node_id, uint16(server_port), local_type, ), ) handshaked_connections_aiter = join_aiters( map_aiter( lambda _: perform_handshake(_, srwt_aiter, outbound_handshake), connections_with_global_connections_aiter, ) ) forker = aiter_forker(handshaked_connections_aiter) handshake_finished_1 = forker.fork(is_active=True) handshake_finished_2 = forker.fork(is_active=True) # Reads messages one at a time from the TCP connection messages_aiter = join_aiters( map_aiter(connection_to_message, handshake_finished_1, 100) ) # Handles each message one at a time, and yields responses to send back or broadcast responses_aiter = join_aiters( map_aiter( partial_func.partial_async_gen(handle_message, api), messages_aiter, 100, ) ) # Uses a forked aiter, and calls the on_connect function to send some initial messages # as soon as the connection is established on_connect_outbound_aiter = join_aiters( map_aiter(connection_to_outbound, handshake_finished_2, 100) ) # Also uses the instance variable _outbound_aiter, which clients can use to send messages # at any time, not just on_connect. outbound_aiter_mapped = map_aiter( lambda x: (None, x, global_connections), outbound_aiter ) responses_aiter = join_aiters( iter_to_aiter( [responses_aiter, on_connect_outbound_aiter, outbound_aiter_mapped] ) ) # For each outbound message, replicate for each peer that we need to send to expanded_messages_aiter = join_aiters( map_aiter(expand_outbound_messages, responses_aiter, 100) ) # This will run forever. Sends each message through the TCP connection, using the # length encoding and CBOR serialization async for connection, message in expanded_messages_aiter: if message is None: # Does not ban the peer, this is just a graceful close of connection. global_connections.close(connection, True) continue if connection.is_closing(): connection.log.info( f"Closing, so will not send {message.function} to peer {connection.get_peername()}" ) continue connection.log.info( f"-> {message.function} to peer {connection.get_peername()}" ) try: await connection.send(message) except (RuntimeError, TimeoutError, OSError,) as e: connection.log.warning( f"Cannot write to {connection}, already closed. Error {e}." ) global_connections.close(connection, True)