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)