def test_make_delayed_pipeline(self): def make_wait_index(idx): async def wait(item): await asyncio.sleep(item[idx] / 10.) return item return wait TEST_CASE = [ (0, 0, 0, 7), (5, 0, 0, 0), (0, 0, 1, 0), (1, 1, 1, 1), (2, 0, 0, 1), (3, 1, 2, 0), ] q = push_aiter() aiter = map_aiter( make_wait_index(0), map_aiter( make_wait_index(1), map_aiter(make_wait_index(2), map_aiter(make_wait_index(3), q, 10), 10), 10), 10) q.push(*TEST_CASE) q.stop() r = run(get_n(aiter)) r1 = sorted(r, key=lambda x: sum(x)) self.assertEqual(r, r1)
async def api_server(rws_aiter, api, workers=1): """ An rws_aiter is an aiter which streams a (StreamReader, StreamWriter, SocketServer) tuples. For a given rws_aiter, create a task which fetches messages from the StreamReader, parses them, and turns them into api calls on api. You can wait forever on this task. If you close the socket, once all connected clients drop off, the task will complete. """ event_aiter = rws_to_event_aiter(rws_aiter, reader_to_cbor_stream) response_writer_for_event = make_response_map_for_api(api) if workers > 1: response_writer_aiter = parallel_map_aiter(response_writer_for_event, workers, event_aiter) else: response_writer_aiter = map_aiter(response_writer_for_event, event_aiter) def to_cbor(response_writer_pair): response, writer = response_writer_pair try: msg = xform_to_cbor_message(response) except Exception as ex: msg = xform_to_cbor_message("problem streaming message: %s" % ex) return writer, msg cbor_msg_aiter = map_aiter(to_cbor, response_writer_aiter) async for writer, cbor_msg in cbor_msg_aiter: writer.write(cbor_msg)
async def main(): server, aiter = await start_server_aiter(7777) line_writer_aiter_aiter = map_aiter( functools.partial(stream_reader_writer_to_line_writer_aiter, server), aiter) line_writer_aiter = join_aiters(line_writer_aiter_aiter) completed_event_aiter = map_aiter(handle_event, line_writer_aiter, worker_count=5) async for line in completed_event_aiter: print(line)
async def start_server( self: "ChiaServer", on_connect: OnConnectFunc = None) -> asyncio.AbstractServer: """ Launches a listening server on host and port specified, to connect to NodeType nodes. On each connection, the on_connect asynchronous generator will be called, and responses will be sent. Whenever a new TCP connection is made, a new srwt tuple is sent through the pipeline. """ require_cert = self._local_type not in (NodeType.FULL_NODE, NodeType.INTRODUCER) ssl_context = ssl_context_for_server(self.root_path, self.config, require_cert=require_cert) server, aiter = await start_server_aiter(self._port, host=None, reuse_address=True, ssl=ssl_context) def add_connection_type( srw: Tuple[asyncio.StreamReader, asyncio.StreamWriter] ) -> Tuple[asyncio.StreamReader, asyncio.StreamWriter, OnConnectFunc]: ssl_object = srw[1].get_extra_info(name="ssl_object") peer_cert = ssl_object.getpeercert() self.log.info(f"Client authed as {peer_cert}") return (srw[0], srw[1], on_connect) srwt_aiter = map_aiter(add_connection_type, aiter) # Push aiters that come from the server into the pipeline if not self._srwt_aiter.is_stopped(): self._srwt_aiter.push(srwt_aiter) self.log.info(f"Server started on port {self._port}") return server
async def start_server( self, host: str, on_connect: OnConnectFunc = None, ) -> bool: """ Launches a listening server on host and port specified, to connect to NodeType nodes. On each connection, the on_connect asynchronous generator will be called, and responses will be sent. Whenever a new TCP connection is made, a new srwt tuple is sent through the pipeline. """ if self._server is not None or self._pipeline_task.done(): return False self._host = host self._server, aiter = await start_server_aiter(self._port, host=None, reuse_address=True) if on_connect is not None: self._on_inbound_connect = on_connect def add_connection_type( srw: Tuple[asyncio.StreamReader, asyncio.StreamWriter] ) -> Tuple[asyncio.StreamReader, asyncio.StreamWriter, None]: return (srw[0], srw[1], None) srwt_aiter = map_aiter(add_connection_type, aiter) # Push all aiters that come from the server, into the pipeline asyncio.create_task(self._add_to_srwt_aiter(srwt_aiter)) log.info(f"Server started on port {self._port}") return True
def test_filter_pipeline(self): async def filter(item_list_of_lists): r = [] for l1 in item_list_of_lists: for item in l1: if item != 0: r.append(item) return r TEST_CASE = [ (0, 0, 0, 7), (5, 0, 0, 0), (0, 0, 1, 0), (1, 1, 1, 1), (2, 0, 0, 1), (3, 1, 2, 0), ] q = push_aiter() aiter = flatten_aiter(map_aiter(filter, q)) q.push(TEST_CASE) q.stop() r = run(get_n(aiter, 12)) r1 = [7, 5, 1, 1, 1, 1, 1, 2, 1, 3, 1, 2] self.assertEqual(r, r1)
def make_client_server(): init_logging() run = asyncio.get_event_loop().run_until_complete path = pathlib.Path(tempfile.mkdtemp(), "port") server, aiter = run(start_unix_server_aiter(path)) rws_aiter = map_aiter(lambda rw: dict( reader=rw[0], writer=rw[1], server=server), aiter) initial_block_hash = bytes(([0] * 31) + [1]) ledger = ledger_api.LedgerAPI(initial_block_hash, RAM_DB()) server_task = asyncio.ensure_future(api_server(rws_aiter, ledger)) remote = run(proxy_for_unix_connection(path)) # make sure server_task isn't garbage collected remote.server_task = server_task return remote
async def main(): server, aiter = await start_server_aiter(7777) line_writer_aiter_aiter = map_aiter( stream_reader_writer_to_line_writer_aiter, aiter) line_writer_aiter = join_aiters(line_writer_aiter_aiter) async for line, sw in line_writer_aiter: print(line) await sw.drain() if line == b"\n": sw.close() sw.write(line) if line == b"quit\n": server.close()
def rws_to_event_aiter(rws_aiter, reader_to_message_stream): def rws_to_reader_event_template_adaptor(rws): return rws, rws["reader"] def reader_event_template_to_event_stream_adaptor(rws_reader): rws, reader = rws_reader return message_stream_to_event_stream(rws, reader_to_message_stream(reader)) def adaptor(rws): return reader_event_template_to_event_stream_adaptor( rws_to_reader_event_template_adaptor(rws)) return join_aiters(map_aiter(adaptor, rws_aiter))
def message_stream_to_event_stream(event_template, message_stream): """ This tweaks each message from message_stream by wrapping it with a dictionary populated with the given template, putting the message is at the top level under "message". """ template = dict(event_template) def adaptor(message): event = dict(template) event.update(message=message) return event return map_aiter(adaptor, message_stream)
def test_make_pipe(self): async def map_f(x): await asyncio.sleep(x / 100.0) return x * x q = push_aiter() aiter = map_aiter(map_f, q) for _ in range(4): q.push(_) for _ in range(3, 9): q.push(_) r = run(get_n(aiter, 10)) q.stop() r.extend(run(get_n(aiter))) r1 = sorted([_*_ for _ in range(4)] + [_ * _ for _ in range(3, 9)]) self.assertEqual(r, r1)
async def start_server(self, on_connect: OnConnectFunc = None) -> bool: """ Launches a listening server on host and port specified, to connect to NodeType nodes. On each connection, the on_connect asynchronous generator will be called, and responses will be sent. Whenever a new TCP connection is made, a new srwt tuple is sent through the pipeline. """ if self._server is not None or self._pipeline_task.done(): return False ssl_context = ssl._create_unverified_context( purpose=ssl.Purpose.CLIENT_AUTH) private_cert, private_key = self.loadSSLConfig("ssl", self.root_path, self.config) ssl_context.load_cert_chain(certfile=private_cert, keyfile=private_key) ssl_context.load_verify_locations(private_cert) if (self._local_type == NodeType.FULL_NODE or self._local_type == NodeType.INTRODUCER): ssl_context.verify_mode = ssl.CERT_NONE else: ssl_context.verify_mode = ssl.CERT_REQUIRED self._server, aiter = await start_server_aiter(self._port, host=None, reuse_address=True, ssl=ssl_context) if on_connect is not None: self._on_inbound_connect = on_connect def add_connection_type( srw: Tuple[asyncio.StreamReader, asyncio.StreamWriter] ) -> Tuple[asyncio.StreamReader, asyncio.StreamWriter, None]: ssl_object = srw[1].get_extra_info(name="ssl_object") peer_cert = ssl_object.getpeercert() self.log.info(f"Client authed as {peer_cert}") return (srw[0], srw[1], None) srwt_aiter = map_aiter(add_connection_type, aiter) # Push all aiters that come from the server, into the pipeline self._tasks.append( asyncio.create_task(self._add_to_srwt_aiter(srwt_aiter))) self.log.info(f"Server started on port {self._port}") return True
def test_client_server(): init_logging() run = asyncio.get_event_loop().run_until_complete path = pathlib.Path(tempfile.mkdtemp(), "port") server, aiter = run(start_unix_server_aiter(path)) rws_aiter = map_aiter( lambda rw: dict(reader=rw[0], writer=rw[1], server=server), aiter) initial_block_hash = bytes(([0] * 31) + [1]) ledger = ledger_api.LedgerAPI(initial_block_hash, RAM_DB()) server_task = asyncio.ensure_future(api_server(rws_aiter, ledger)) run(client_test(path)) server_task.cancel()
def test_syncmap(self): def make_sync_transformation_f(results): def sync_transformation_f(item): results.append(item) return item + 1 return sync_transformation_f results = [] q = push_aiter() q.push(5, 4, 3) q.stop() r = list(q.available_iter()) self.assertEqual(r, [5, 4, 3]) aiter = map_aiter(make_sync_transformation_f(results), q) r = run(get_n(aiter)) self.assertEqual(r, [6, 5, 4]) self.assertEqual(results, [5, 4, 3])
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())
def run_ledger_api(server, aiter): db = RAM_DB() INITIAL_BLOCK_HASH = bytes(([0] * 31) + [1]) ledger = ledger_api.LedgerAPI(INITIAL_BLOCK_HASH, db) rws_aiter = map_aiter(lambda rw: dict(reader=rw[0], writer=rw[1], server=server), aiter) return api_server(rws_aiter, ledger)
def serve_api_on_unix_port(api, path): run = asyncio.get_event_loop().run_until_complete server, aiter = run(start_unix_server_aiter(path)) rws_aiter = map_aiter(lambda rw: dict(reader=rw[0], writer=rw[1], server=server), aiter) return asyncio.ensure_future(api_server(rws_aiter, api, 20))
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)