Bring processing onto asyncio

proxy/common_neon/data.py

@@ -26,5 +26,5 @@ def add_instruction(self, sol_tx_hash: str, sol_spent: int, steps: int, bpf: int
 
 @dataclass
 class NeonTxData:
-    tx_signed: bytes
+    tx_signed: str
 

proxy/mempool/__init__.py

@@ -1,4 +1,4 @@
-from .mempool_server import MemPoolServer
+from .mempool_server import PickableDataServer
 from .mempool_client import MemPoolClient
 from .mempool_service import MemPoolService
 from .mempool_client import MemPoolClient

proxy/mempool/mem_pool.py

@@ -1,60 +1,40 @@
+import asyncio
 from logged_groups import logged_group
-from multiprocessing import Pool, Queue
-import queue
+from concurrent.futures import ProcessPoolExecutor
+import time
 
 from ..common_neon.data import NeonTxData
 
 
 @logged_group("neon.MemPool")
 class MemPool:
 
-    POOL_PROC_COUNT = 3
-    TX_QUEUE_TIMEOUT_SEC = 0.04
-    BREAK_PROC_INVOCATION = 0
+    POOL_PROC_COUNT = 8
 
     def __init__(self):
         self._pool = None
-        self._tx_queue = Queue()
-
-    def _process_init(self):
-        self._pool = Pool(processes=self.POOL_PROC_COUNT)
+        self._tx_queue = asyncio.Queue()
+        self._pool = ProcessPoolExecutor(self.POOL_PROC_COUNT)
+        self._neon_tx_futures = set()
+        self._event_loop = asyncio.get_event_loop()
+        self._event_loop.create_task(self.do_mempool_stuff())
+        self._send_tx_futures = list()
 
     def on_eth_send_raw_transaction(self, neon_tx_data: NeonTxData):
-        self._tx_queue.put(neon_tx_data)
-
-    def error_callback(self, error):
-        self.error("Failed to invoke on worker process: ", error)
-
-    def on_eth_send_raw_transaction_callback(self, result):
-        self.debug(f"Processing result: {result}")
+        self._tx_queue.put_nowait(neon_tx_data)
 
     @staticmethod
-    def _on_eth_send_raw_transaction_impl(neon_tx_data: NeonTxData) -> bool:
+    def on_eth_send_raw_transaction_impl(neon_tx_data: NeonTxData) -> bool:
+        print(f"neon_tx_data: {neon_tx_data}")
+        time.sleep(0.1)
         return True
 
-    def do_extras(self):
-        pass
-
-    def __getstate__(self):
-        self_dict = self.__dict__.copy()
-        del self_dict['_pool']
-        return self_dict
-
-    def __setstate__(self, state):
-        self.__dict__.update(state)
-
-    def run(self):
-        self._process_init()
-        while True:
-            try:
-                if not self._process_queue():
-                    break
-            except queue.Empty:
-                self.do_extras()
-
-    def _process_queue(self) -> bool:
-        neon_tx_data = self._tx_queue.get(block=True, timeout=self.TX_QUEUE_TIMEOUT_SEC)
-        if neon_tx_data == self.BREAK_PROC_INVOCATION:
-            return False
-        self._pool.apply_async(MemPool._on_eth_send_raw_transaction_impl, (neon_tx_data, ), callback=self.on_eth_send_raw_transaction_callback, error_callback=self.error_callback)
-        return True
+    async def do_mempool_stuff(self):
+        with self._pool as pool:
+            while True:
+                try:
+                    neon_tx_data = await self._tx_queue.get()
+                    pool.submit(MemPool.on_eth_send_raw_transaction_impl, neon_tx_data)
+                except Exception as err:
+                    print(f"Failed to submit neon tx onto the mempool worker: {err}")
+

proxy/mempool/mempool_client.py

@@ -1,26 +1,32 @@
-from multiprocessing.managers import BaseManager
+import socket
+import pickle
+import struct
 from logged_groups import logged_group
 
 
+from ..common_neon.data import NeonTxData
+
+
 @logged_group("neon.Proxy")
 class MemPoolClient:
 
     def __init__(self, host: str, port: int):
 
         self.info(f"Initialize MemPoolClient connecting to: {port} at: {host}")
+        self._connection = socket.create_connection((host, port))
 
-        class MemPoolManager(BaseManager):
-            def __init__(self):
-                super(MemPoolManager, self).__init__(address=(host, port), authkey=b'abracadabra')
-                self.register("MemPool")
-
-        self._mempool_manager = MemPoolManager()
-        self._mempool_manager.connect()
-        self._mempool = self._mempool_manager.MemPool()
-
-    def on_eth_send_raw_transaction(self, neon_tx_data):
+    def on_eth_send_raw_transaction(self, neon_tx_data: NeonTxData):
         try:
-            self._mempool.on_eth_send_raw_transaction(neon_tx_data)
+            self.debug(f"Send transaction: {neon_tx_data}")
+            payload = self.decode_neon_tx_data(neon_tx_data)
+            self._connection.send(payload)
         except BaseException as err:
             self.error(f"Failed to send raw transaction onto mempool: {err}")
             raise Exception("Failed to send to the mempool")
+
+    def decode_neon_tx_data(self, neon_tx_data: NeonTxData):
+        data = pickle.dumps(neon_tx_data)
+        data_len = len(data)
+        packed_len = struct.pack("!I", data_len)
+        payload = packed_len + data
+        return payload

proxy/mempool/mempool_server.py

@@ -1,47 +1,67 @@
-import multiprocessing as mp
-import os
-import signal
-from multiprocessing.managers import BaseManager
 from abc import ABC, abstractmethod
+import asyncio
+import socket
+import pickle
+import struct
+from typing import Any
 from logged_groups import logged_group
 
-from ..common_neon.data import NeonTxData
 
-
-class MemPoolServerUser(ABC):
+class PickableDataServerUser(ABC):
 
     @abstractmethod
-    def on_eth_send_raw_transaction(self, neon_tx_data: NeonTxData):
+    def on_data_received(self, data: Any):
         """Gets neon_tx_data from the neon rpc api service worker"""
 
 
 @logged_group("neon.MemPool")
-class MemPoolServer(ABC):
+class PickableDataServer(ABC):
 
     QUEUE_TIMEOUT_SEC = 0.4
     BREAK_PROC_INVOCATION = 0
     JOIN_PROC_TIMEOUT_SEC = 5
 
-    def __init__(self, *, user: MemPoolServerUser, host: str, port: int):
+    def __init__(self, *, user: PickableDataServerUser, host: str, port: int):
         self._user = user
         self._port = port
         self._host = host
-        self._timeout = self.QUEUE_TIMEOUT_SEC
-
-        class MemPoolManager(BaseManager):
-            pass
-
-        MemPoolManager.register("MemPool", lambda: self._user)
-        self._mempool_manager = MemPoolManager(address=(host, port), authkey=b'abracadabra')
-        self._mempool_server = self._mempool_manager.get_server()
-        self._mempool_server_process = mp.Process(target=self._mempool_server.serve_forever, name="mempool_listen_proc")
-
-        pid = os.getpid()
-        signal.signal(signal.SIGINT, lambda sif, frame: self.finish() if os.getpid() == pid else 0)
 
     def start(self):
         self.info(f"Start listen on: {self._port} at: {self._host}")
         self._mempool_server_process.start()
 
     def finish(self):
         self._mempool_server_process.terminate()
+
+    async def handle_client(self, client):
+        loop = asyncio.get_event_loop()
+        peer_name = client.getpeername()
+        self.debug(f"Got new incoming connection: {peer_name}")
+        while True:
+            try:
+                len_packed: bytes = await loop.sock_recv(client, 4)
+                if len(len_packed) == 0:
+                    break
+                # TODO: all the data can be received by parts, handle it
+                payload_len_data = struct.unpack("!I", len_packed[:4])[0]
+                payload = await loop.sock_recv(client, payload_len_data)
+                data = pickle.loads(payload)
+                self._user.on_data_received(data)
+                response = pickle.dumps({"data": data, "status": "ok"})
+                await loop.sock_sendall(client, response)
+            except ConnectionResetError:
+                self.error(f"Client connection: {peer_name} - has been interrupted")
+                break
+        client.close()
+
+    async def run_server(self):
+        self.info(f"Listen port: {self._port} on: {self._host}")
+        server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+        server.bind((self._host, self._port))
+        server.listen(8)
+        server.setblocking(False)
+
+        loop = asyncio.get_event_loop()
+        while True:
+            client, _ = await loop.sock_accept(server)
+            loop.create_task(self.handle_client(client))

proxy/mempool/mempool_service.py

@@ -1,29 +1,27 @@
 from logged_groups import logged_group
-from multiprocessing import Process
+import asyncio
 
-from .mempool_server import MemPoolServer, MemPoolServerUser
+from .mempool_server import PickableDataServer, PickableDataServerUser
 from .mem_pool import MemPool
 
-from ..common_neon.data import NeonTxData
-
+from typing import Any
 
 @logged_group("neon.MemPool")
-class MemPoolService(MemPoolServerUser):
+class MemPoolService(PickableDataServerUser):
 
     MEMPOOL_SERVICE_PORT = 9091
-    MEMPOOL_SERVICE_HOST = "127.0.0.1"
+    MEMPOOL_SERVICE_HOST = "0.0.0.0"
 
     def __init__(self):
-        self._mempool_server = MemPoolServer(user=self, host=self.MEMPOOL_SERVICE_HOST, port=self.MEMPOOL_SERVICE_PORT)
+        self.event_loop = asyncio.new_event_loop()
+        asyncio.set_event_loop(self.event_loop)
+
+        self._mempool_server = PickableDataServer(user=self, host=self.MEMPOOL_SERVICE_HOST, port=self.MEMPOOL_SERVICE_PORT)
         self._mempool = MemPool()
-        self._mempool_proc = Process(target=self.run_mempool)
 
     def start(self):
-        self._mempool_server.start()
-        self._mempool_proc.start()
-
-    def on_eth_send_raw_transaction(self, neon_tx_data: NeonTxData):
-        self._mempool.on_eth_send_raw_transaction(neon_tx_data)
+        self.info("Run until complete")
+        self.event_loop.run_until_complete(self._mempool_server.run_server())
 
-    def run_mempool(self):
-        self._mempool.run()
+    def on_data_received(self, data: Any):
+        self._mempool.on_eth_send_raw_transaction(data)

proxy/neon_rpc_api_model/neon_rpc_api_model.py

@@ -8,14 +8,14 @@
 from web3.auto import w3
 
 from ..common_neon.address import EthereumAddress
-from ..common_neon.emulator_interactor import call_emulated, call_trx_emulated
+from ..common_neon.emulator_interactor import call_emulated
 from ..common_neon.errors import EthereumError, InvalidParamError, PendingTxError
 from ..common_neon.estimate import GasEstimate
 from ..common_neon.eth_proto import Trx as EthTrx
 from ..common_neon.keys_storage import KeyStorage
 from ..common_neon.solana_interactor import SolanaInteractor
 from ..common_neon.utils import SolanaBlockInfo
-from ..common_neon.data import NeonTxPrecheckResult
+from ..common_neon.data import NeonTxPrecheckResult, NeonTxData
 from ..environment import SOLANA_URL, PP_SOLANA_URL, PYTH_MAPPING_ACCOUNT, NEON_EVM_VERSION, NEON_EVM_REVISION, \
                           CHAIN_ID, neon_cli, EVM_STEP_COUNT
 from ..memdb.memdb import MemDB
@@ -451,7 +451,8 @@ def eth_sendRawTransaction(self, rawTrx: str) -> str:
                 tx_sender.execute(neon_tx_precheck_result)
 
             self._stat_tx_success()
-            self._mempool_client.on_eth_send_raw_transaction(eth_signature)
+            neon_tx_data = NeonTxData(tx_signed=rawTrx)
+            self._mempool_client.on_eth_send_raw_transaction(neon_tx_data)
             return eth_signature
 
         except PendingTxError as err: