WIP: Direct DA connection protocol

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()
+

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()
+

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()

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.")

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()
+