From 775d9fc6b416da1cb9aae7736485c105c1793d34 Mon Sep 17 00:00:00 2001 From: Gusto Date: Mon, 8 Jul 2024 16:21:14 +0300 Subject: [PATCH] WIP: Direct DA connection protocol --- da/executor/__init__.py | 0 da/executor/executor.py | 47 ++++++++++++++++++++++++++++++++++ da/executor/mock_network.py | 33 ++++++++++++++++++++++++ da/executor/node.py | 36 ++++++++++++++++++++++++++ da/executor/proto.py | 50 +++++++++++++++++++++++++++++++++++++ da/executor/transport.py | 30 ++++++++++++++++++++++ 6 files changed, 196 insertions(+) create mode 100644 da/executor/__init__.py create mode 100644 da/executor/executor.py create mode 100644 da/executor/mock_network.py create mode 100644 da/executor/node.py create mode 100644 da/executor/proto.py create mode 100644 da/executor/transport.py diff --git a/da/executor/__init__.py b/da/executor/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/da/executor/executor.py b/da/executor/executor.py new file mode 100644 index 0000000..da965a8 --- /dev/null +++ b/da/executor/executor.py @@ -0,0 +1,47 @@ +import asyncio +import proto +from transport import Transport + +class Executor: + def __init__(self, addr, port, col_num): + self.addr = addr + self.port = port + self.col_num = col_num + self.connections = [] + self.interval = 10 + + async def execute(self): + data = "TEST DATA" + while True: + try: + for transport in self.connections: + message = proto.new_dispersal_put_msg(data) + await transport.write(message) + await asyncio.sleep(self.interval) + except asyncio.CancelledError: + break + except Exception as e: + print(f"Error during message sending: {e}") + + async def connect(self): + try: + reader, writer = await asyncio.open_connection(self.addr, self.port) + conn_id = len(self.connections) + transport = Transport(conn_id, reader, writer, self._handle) + self.connections.append(transport) + print(f"Connected to {self.addr}:{self.port}, ID: {conn_id}") + asyncio.create_task(transport.read_and_process()) + except Exception as e: + print(f"Failed to connect or lost connection: {e}") + + async def _handle(self, conn_id, writer, message): + msg_type, msg_id, data = message + if msg_type == proto.DISPERSAL_OK: + print(f"Executor: Dispersal OK from connection {conn_id}, message ID {msg_id}.") + elif msg_type == proto.SAMPLE_OK: + print(f"Executor: Sample OK from connection {conn_id}, message ID {msg_id}.") + + async def run(self): + await asyncio.gather(*(self.connect() for _ in range(self.col_num))) + await self.execute() + diff --git a/da/executor/mock_network.py b/da/executor/mock_network.py new file mode 100644 index 0000000..8f57229 --- /dev/null +++ b/da/executor/mock_network.py @@ -0,0 +1,33 @@ +import asyncio +import argparse +from node import Node +from executor import Executor + +class App: + def __init__(self, addr='localhost'): + self.addr = addr + + async def run_nodes(self, start_port, num_nodes): + nodes = [Node(self.addr, start_port + i) for i in range(num_nodes)] + await asyncio.gather(*(node.run() for node in nodes)) + + async def run_nodes_with_executor(self, col_number): + node = Node(self.addr, 8888) + executor = Executor(self.addr, 8888, col_number) + await asyncio.gather(node.run(), executor.run()) + + async def run_executor(self, remote_addr, start_port, col_number): + executor = Executor(remote_addr, start_port, col_number) + await asyncio.gather(executor.run()) + +def main(): + # TODO: Add args parser. + app = App() + + # asyncio.run(app.run_nodes(10000, 4096)) + # asyncio.run(app.run_executor('localhost', 8888, 1)) + asyncio.run(app.run_nodes_with_executor(10)) + +if __name__ == '__main__': + main() + diff --git a/da/executor/node.py b/da/executor/node.py new file mode 100644 index 0000000..93b3749 --- /dev/null +++ b/da/executor/node.py @@ -0,0 +1,36 @@ +import asyncio +import struct +import proto +from itertools import count +from transport import Transport + +conn_id_counter = count(start=1) + +class Node: + def __init__(self, addr, port): + self.addr = addr + self.port = port + + async def _on_conn(self, reader, writer): + conn_id = next(conn_id_counter) + transport = Transport(conn_id, reader, writer, self._handle) + await transport.read_and_process() + + async def listen(self): + server = await asyncio.start_server( + self._on_conn, self.addr, self.port + ) + print(f"Server started at {self.addr}:{self.port}") + async with server: + await server.serve_forever() + + async def _handle(self, conn_id, writer, message): + msg_type, msg_id, data = message + if msg_type == proto.DISPERSAL_PUT: + response = proto.new_dispersal_ok_msg(msg_id) + writer.write(response) + elif msg_type == proto.SAMPLE_PUT: + pass + + async def run(self): + await self.listen() diff --git a/da/executor/proto.py b/da/executor/proto.py new file mode 100644 index 0000000..dc6ed8b --- /dev/null +++ b/da/executor/proto.py @@ -0,0 +1,50 @@ +import struct +from itertools import count + +DISPERSAL_PUT = 0x01 +DISPERSAL_OK = 0x02 +SAMPLE_PUT = 0x03 +SAMPLE_OK = 0x04 + +HEADER_SIZE = 9 +HEADER_FORMAT = "!B I I" + +DISPERSAL_HASH_COL_SIZE = 6 +# First 4 bytes for the hash and the next 2 bytes for the column index. +DISPERSAL_HASH_COL_FORMAT = "!I H" + +msg_id_counter = count(start=1) + +def pack_header(msg_type, msg_id, data): + encoded_data = data.encode() + data_length = len(encoded_data) + header = struct.pack(HEADER_FORMAT, msg_type, msg_id, data_length) + return header + encoded_data + +def unpack_header(data): + if len(data) < HEADER_SIZE: + return None, None, None, None + msg_type, msg_id, data_length = struct.unpack(HEADER_FORMAT, data[:HEADER_SIZE]) + return msg_type, msg_id, data_length + +def new_dispersal_put_msg(data): + msg_id = next(msg_id_counter) + return pack_header(DISPERSAL_PUT, msg_id, data) + +def new_dispersal_ok_msg(msg_id): + return pack_header(DISPERSAL_OK, msg_id, "") + +def new_sample_put_msg(data): + msg_id = next(msg_id_counter) + return pack_header(SAMPLE_PUT, msg_id, data) + +def new_sample_ok_msg(msg_id, response_data): + return pack_header(SAMPLE_OK, msg_id, response_data) + +def parse_dispersal_data(data): + if len(data) >= DISPERSAL_HASH_COL_SIZE: + hash_value, col_index = struct.unpack(DISPERSAL_HASH_COL_FORMAT, data[:DISPERSAL_HASH_COL_SIZE]) + remaining_data = data[DISPERSAL_HASH_COL_SIZE:] + return hash_value, col_index, remaining_data + else: + raise ValueError("Data is too short to unpack hash and col.") diff --git a/da/executor/transport.py b/da/executor/transport.py new file mode 100644 index 0000000..de09fdb --- /dev/null +++ b/da/executor/transport.py @@ -0,0 +1,30 @@ +import asyncio +import struct +import proto + +class Transport: + def __init__(self, conn_id, reader, writer, handler): + self.conn_id = conn_id + self.reader = reader + self.writer = writer + self.handler = handler + + async def read_and_process(self): + try: + while True: + header = await self.reader.readexactly(9) # Assuming the header is 9 bytes long + msg_type, msg_id, data_length = proto.unpack_header(header) + data = await self.reader.readexactly(data_length) + await self.handler(self.conn_id, self.writer, (msg_type, msg_id, data)) + except asyncio.IncompleteReadError: + print("Transport: Connection closed by the peer.") + except Exception as e: + print(f"Transport: An error occurred: {e}") + finally: + self.writer.close() + await self.writer.wait_closed() + + async def write(self, data): + self.writer.write(data) + await self.writer.drain() +