Fix graceful shutdown on linux (#58)

dungeon-master-666 · web-flow · commit 7a29df7bdf25 · 2022-06-26T13:52:21.000+02:00
diff --git a/ton-http-api/pyTON/manager.py b/ton-http-api/pyTON/manager.py
@@ -2,11 +2,12 @@
 import time
 import traceback
 import random
-import aioprocessing as ap
+import queue
 
 from collections import defaultdict
 from collections.abc import Mapping
 from copy import deepcopy
+from concurrent.futures import ThreadPoolExecutor
 
 from pyTON.worker import TonlibWorker
 from pyTON.models import TonlibWorkerMsgType, TonlibClientResult, ConsensusBlock
@@ -38,6 +39,8 @@ def __init__(self,
         # cache setup
         self.setup_cache()
 
+        self.threadpool_executor = ThreadPoolExecutor(max_workers=len(self.tonlib_settings.liteserver_config['liteservers']) * 2)
+
         # workers spawn
         self.loop = loop or asyncio.get_running_loop()
         for ls_index in range(len(self.tonlib_settings.liteserver_config['liteservers'])):
@@ -60,6 +63,8 @@ async def shutdown(self):
         for i in self.workers:
             await self.worker_control(i, enabled=False)
 
+        self.threadpool_executor.shutdown()
+
     def setup_cache(self):
         self.raw_get_transactions = self.cache_manager.cached(expire=5)(self.raw_get_transactions)
         self.get_transactions = self.cache_manager.cached(expire=15, check_error=False)(self.get_transactions)
@@ -84,14 +89,12 @@ def spawn_worker(self, ls_index, force_restart=False):
                 return
             try:
                 worker_info['reader'].cancel()  
+                worker_info['worker'].exit_event.set()
+                worker_info['worker'].output_queue.cancel_join_thread()
+                worker_info['worker'].input_queue.cancel_join_thread()
+                worker_info['worker'].output_queue.close()
+                worker_info['worker'].input_queue.close()
                 worker_info['worker'].join(timeout=3)
-                if worker_info['worker'].is_alive():
-                    worker_info['worker'].terminate()
-                    worker_info['worker'].join()
-                self.workers[ls_index]['worker'].output_queue.close()
-                self.workers[ls_index]['worker'].output_queue.join_thread()
-                self.workers[ls_index]['worker'].input_queue.close()
-                self.workers[ls_index]['worker'].input_queue.join_thread()
             except Exception as ee:
                 logger.error('Failed to delete existing process: {exc}', exc=ee)
         # running new worker
@@ -112,15 +115,17 @@ def spawn_worker(self, ls_index, force_restart=False):
 
     async def worker_control(self, ls_index, enabled):
         if enabled == False:
-            self.workers[ls_index]['worker'].terminate()
-            self.workers[ls_index]['worker'].join()
             self.workers[ls_index]['reader'].cancel()
-            await self.workers[ls_index]['reader']
+            self.workers[ls_index]['worker'].exit_event.set()
 
+            self.workers[ls_index]['worker'].output_queue.cancel_join_thread()
+            self.workers[ls_index]['worker'].input_queue.cancel_join_thread()
             self.workers[ls_index]['worker'].output_queue.close()
-            self.workers[ls_index]['worker'].output_queue.join_thread()
             self.workers[ls_index]['worker'].input_queue.close()
-            self.workers[ls_index]['worker'].input_queue.join_thread()
+
+            self.workers[ls_index]['worker'].join()
+            
+            await self.workers[ls_index]['reader']
 
         self.workers[ls_index]['is_enabled'] = enabled
 
@@ -152,7 +157,10 @@ async def read_results(self, ls_index):
         worker = self.workers[ls_index]['worker']
         while True:
             try:
-                msg_type, msg_content = await worker.output_queue.coro_get()
+                try:
+                    msg_type, msg_content = await self.loop.run_in_executor(self.threadpool_executor, worker.output_queue.get, True, 1)
+                except queue.Empty:
+                    continue
                 if msg_type == TonlibWorkerMsgType.TASK_RESULT:
                     task_id = msg_content.task_id
                     result = msg_content.result
@@ -173,9 +181,6 @@ async def read_results(self, ls_index):
 
                 if msg_type == TonlibWorkerMsgType.ARCHIVAL_UPDATE:
                     worker.is_archival = msg_content
-
-                if msg_type == TonlibWorkerMsgType.DEAD_REPORT:
-                    self.spawn_worker(ls_index, force_restart=True)
             except asyncio.CancelledError:
                 logger.info("Task read_results was cancelled")
                 return
@@ -269,7 +274,7 @@ async def dispatch_request_to_worker(self, method, ls_index, *args, **kwargs):
 
         logger.info("Sending request method: {method}, task_id: {task_id}, ls_index: {ls_index}", 
             method=method, task_id=task_id, ls_index=ls_index)
-        await self.workers[ls_index]['worker'].input_queue.coro_put((task_id, timeout, method, args, kwargs))
+        await self.loop.run_in_executor(self.threadpool_executor, self.workers[ls_index]['worker'].input_queue.put, (task_id, timeout, method, args, kwargs))
 
         try:
             self.futures[task_id] = self.loop.create_future()
@@ -332,7 +337,7 @@ async def raw_send_message(self, serialized_boc):
             for ls_index in ls_index_list:
                 task_id = "{}:{}".format(time.time(), random.random())
                 timeout = time.time() + self.tonlib_settings.request_timeout
-                await self.workers[ls_index]['worker'].input_queue.coro_put((task_id, timeout, 'raw_send_message', [serialized_boc], {}))
+                await self.loop.run_in_executor(self.threadpool_executor, self.workers[ls_index]['worker'].input_queue.put, (task_id, timeout, 'raw_send_message', [serialized_boc], {}))
 
                 self.futures[task_id] = self.loop.create_future()
                 task_ids.append(task_id)
diff --git a/ton-http-api/pyTON/models.py b/ton-http-api/pyTON/models.py
@@ -17,11 +17,9 @@ class TonlibClientResult:
 
 
 class TonlibWorkerMsgType(Enum):
-    DEAD_REPORT = 0
-    TASK_RESULT = 1
-    LAST_BLOCK_UPDATE = 2
-    ARCHIVAL_UPDATE = 3
-
+    TASK_RESULT = 0
+    LAST_BLOCK_UPDATE = 1
+    ARCHIVAL_UPDATE = 2
 
 
 @dataclass
diff --git a/ton-http-api/pyTON/worker.py b/ton-http-api/pyTON/worker.py
@@ -2,13 +2,14 @@
 import traceback
 import random
 import time
-import aioprocessing as ap
+import queue
 import multiprocessing as mp
 
 from pyTON.settings import TonlibSettings
 from pyTON.models import TonlibWorkerMsgType, TonlibClientResult
 from pytonlib import TonlibClient
 from datetime import datetime
+from concurrent.futures import ThreadPoolExecutor
 
 from enum import Enum
 from dataclasses import dataclass
@@ -21,14 +22,13 @@ class TonlibWorker(mp.Process):
     def __init__(self, 
                  ls_index: int, 
                  tonlib_settings: TonlibSettings,
-                 input_queue: Optional[ap.AioQueue]=None,
-                 output_queue: Optional[ap.AioQueue]=None):
+                 input_queue: Optional[mp.Queue]=None,
+                 output_queue: Optional[mp.Queue]=None):
         super(TonlibWorker, self).__init__(daemon=True)
 
-        self.input_queue = input_queue or ap.AioQueue()
-        self.output_queue = output_queue or ap.AioQueue()
-        self.input_queue.cancel_join_thread()
-        self.output_queue.cancel_join_thread()
+        self.input_queue = input_queue or mp.Queue()
+        self.output_queue = output_queue or mp.Queue()
+        self.exit_event = mp.Event()
 
         self.ls_index = ls_index
         self.tonlib_settings = tonlib_settings
@@ -39,11 +39,14 @@ def __init__(self,
         self.loop = None
         self.tasks = {}
         self.tonlib = None
+        self.threadpool_executor = None
 
         self.timeout_count = 0
         self.is_dead = False
 
     def run(self):
+        self.threadpool_executor = ThreadPoolExecutor(max_workers=6)
+
         policy = asyncio.get_event_loop_policy()
         policy.set_event_loop(policy.new_event_loop())
         self.loop = asyncio.new_event_loop()
@@ -61,7 +64,27 @@ def run(self):
         self.tasks['report_last_block'] = self.loop.create_task(self.report_last_block())
         self.tasks['report_archival'] = self.loop.create_task(self.report_archival())
         self.tasks['main_loop'] = self.loop.create_task(self.main_loop())
-        self.loop.run_until_complete(self.idle_loop())
+        self.tasks['idle_loop'] = self.loop.create_task(self.idle_loop())
+
+        finished, unfinished = self.loop.run_until_complete(asyncio.wait([
+            self.tasks['report_last_block'], self.tasks['report_archival'], self.tasks['main_loop'], self.tasks['idle_loop']], 
+            return_when=asyncio.FIRST_COMPLETED))
+        
+        self.exit_event.set()        
+        
+        for to_cancel in unfinished:
+            to_cancel.cancel()
+            try:
+                self.loop.run_until_complete(to_cancel)
+            except:
+                pass
+
+        self.threadpool_executor.shutdown()
+
+        self.output_queue.cancel_join_thread()
+        self.input_queue.cancel_join_thread()
+        self.output_queue.close()
+        self.input_queue.close()
 
     @property
     def info(self):
@@ -73,97 +96,81 @@ def info(self):
             'number': self.ls_index,
         }
 
-    async def report_dead(self):
-        if not self.is_dead:
-            self.is_dead = True
-            
-            format_exc = traceback.format_exc()
-            logger.error('Dead report: {format_exc}', format_exc=format_exc)
-            await self.output_queue.coro_put((TonlibWorkerMsgType.DEAD_REPORT, format_exc))
-
     async def report_last_block(self):
-        try:
-            while not self.is_dead:
-                last_block = -1
-                try:
-                    masterchain_info = await self.tonlib.get_masterchain_info()
-                    last_block = masterchain_info["last"]["seqno"]
-                    self.timeout_count = 0
-                except asyncio.CancelledError:
-                    logger.warning('Client #{ls_index:03d} report_last_block timeout', ls_index=self.ls_index)
-                    self.timeout_count += 1
-                except Exception as e:
-                    logger.error("Client #{ls_index:03d} report_last_block exception: {exc}", ls_index=self.ls_index, exc=e)
-                    self.timeout_count += 1
-
-                if self.timeout_count >= 10:
-                    raise RuntimeError(f'Client #{self.ls_index:03d} got {self.timeout_count} timeouts in report_last_block')
-                
-                self.last_block = last_block
-                await self.output_queue.coro_put((TonlibWorkerMsgType.LAST_BLOCK_UPDATE, self.last_block))
-                await asyncio.sleep(1)
-        except:
-            await self.report_dead()
+        while not self.exit_event.is_set():
+            last_block = -1
+            try:
+                masterchain_info = await self.tonlib.get_masterchain_info()
+                last_block = masterchain_info["last"]["seqno"]
+                self.timeout_count = 0
+            except asyncio.CancelledError:
+                logger.warning('Client #{ls_index:03d} report_last_block timeout', ls_index=self.ls_index)
+                self.timeout_count += 1
+            except Exception as e:
+                logger.error("Client #{ls_index:03d} report_last_block exception: {exc}", ls_index=self.ls_index, exc=e)
+                self.timeout_count += 1
+
+            if self.timeout_count >= 10:
+                raise RuntimeError(f'Client #{self.ls_index:03d} got {self.timeout_count} timeouts in report_last_block')
+            
+            self.last_block = last_block
+            await self.loop.run_in_executor(self.threadpool_executor, self.output_queue.put, (TonlibWorkerMsgType.LAST_BLOCK_UPDATE, self.last_block))
+            await asyncio.sleep(1)
 
     async def report_archival(self):
-        try:
-            while not self.is_dead:
-                is_archival = False
-                try:
-                    block_transactions = await self.tonlib.get_block_transactions(-1, -9223372036854775808, random.randint(2, 4096), count=10)
-                    is_archival = block_transactions.get("@type", "") == "blocks.transactions"
-                except asyncio.CancelledError:
-                    logger.warning('Client #{ls_index:03d} report_archival timeout', ls_index=self.ls_index)
-                except Exception as e:
-                    logger.error("Client #{ls_index:03d} report_archival exception: {exc}", ls_index=self.ls_index, exc=e)
-                self.is_archival = is_archival
-                await self.output_queue.coro_put((TonlibWorkerMsgType.ARCHIVAL_UPDATE, self.is_archival))
-                await asyncio.sleep(600)
-        except:
-            await self.report_dead()
+        while not self.exit_event.is_set():
+            is_archival = False
+            try:
+                block_transactions = await self.tonlib.get_block_transactions(-1, -9223372036854775808, random.randint(2, 4096), count=10)
+                is_archival = block_transactions.get("@type", "") == "blocks.transactions"
+            except asyncio.CancelledError:
+                logger.warning('Client #{ls_index:03d} report_archival timeout', ls_index=self.ls_index)
+            except Exception as e:
+                logger.error("Client #{ls_index:03d} report_archival exception: {exc}", ls_index=self.ls_index, exc=e)
+            self.is_archival = is_archival
+            await self.loop.run_in_executor(self.threadpool_executor, self.output_queue.put, (TonlibWorkerMsgType.ARCHIVAL_UPDATE, self.is_archival))
+            await asyncio.sleep(600)
         
     async def main_loop(self):
-        try:
-            while not self.is_dead:
+        while not self.exit_event.is_set():
+            try:
+                task_id, timeout, method, args, kwargs = await self.loop.run_in_executor(self.threadpool_executor, self.input_queue.get, True, 1)
+            except queue.Empty:
+                continue
+
+            result = None
+            exception = None
+
+            start_time = datetime.now()
+            if time.time() < timeout:
                 try:
-                    task_id, timeout, method, args, kwargs = await self.input_queue.coro_get(timeout=3)
-                except:
-                    continue
-
-                result = None
-                exception = None
-
-                start_time = datetime.now()
-                if time.time() < timeout:
-                    try:
-                        result = await self.tonlib.__getattribute__(method)(*args, **kwargs)
-                    except asyncio.CancelledError:
-                        exception = Exception("Liteserver timeout")
-                        logger.warning("Client #{ls_index:03d} did not get response from liteserver before timeout", ls_index=self.ls_index)
-                    except Exception as e:
-                        exception = e
-                        logger.warning("Client #{ls_index:03d} raised exception while executing task. Method: {method}, args: {args}, kwargs: {kwargs}, exception: {exc}", 
-                            ls_index=self.ls_index, method=method, args=args, kwargs=kwargs, exc=e)
-                    else:
-                        logger.debug("Client #{ls_index:03d} got result {method}", ls_index=self.ls_index, method=method)
+                    result = await self.tonlib.__getattribute__(method)(*args, **kwargs)
+                except asyncio.CancelledError:
+                    exception = Exception("Liteserver timeout")
+                    logger.warning("Client #{ls_index:03d} did not get response from liteserver before timeout", ls_index=self.ls_index)
+                except Exception as e:
+                    exception = e
+                    logger.warning("Client #{ls_index:03d} raised exception while executing task. Method: {method}, args: {args}, kwargs: {kwargs}, exception: {exc}", 
+                        ls_index=self.ls_index, method=method, args=args, kwargs=kwargs, exc=e)
                 else:
-                    exception = asyncio.TimeoutError()
-                    logger.warning("Client #{ls_index:03d} received task after timeout", ls_index=self.ls_index)
-                end_time = datetime.now()
-                elapsed_time = (end_time - start_time).total_seconds()
-
-                # result
-                tonlib_task_result = TonlibClientResult(task_id,
-                                                        method,
-                                                        elapsed_time=elapsed_time,
-                                                        params=[args, kwargs],
-                                                        result=result,
-                                                        exception=exception,
-                                                        liteserver_info=self.info)
-                await self.output_queue.coro_put((TonlibWorkerMsgType.TASK_RESULT, tonlib_task_result))
-        except:
-            await self.report_dead()
-        
+                    logger.debug("Client #{ls_index:03d} got result {method}", ls_index=self.ls_index, method=method)
+            else:
+                exception = asyncio.TimeoutError()
+                logger.warning("Client #{ls_index:03d} received task after timeout", ls_index=self.ls_index)
+            end_time = datetime.now()
+            elapsed_time = (end_time - start_time).total_seconds()
+
+            # result
+            tonlib_task_result = TonlibClientResult(task_id,
+                                                    method,
+                                                    elapsed_time=elapsed_time,
+                                                    params=[args, kwargs],
+                                                    result=result,
+                                                    exception=exception,
+                                                    liteserver_info=self.info)
+            await self.loop.run_in_executor(self.threadpool_executor, self.output_queue.put, (TonlibWorkerMsgType.TASK_RESULT, tonlib_task_result))
+
     async def idle_loop(self):
-        while not self.is_dead:
+        while not self.exit_event.is_set():
             await asyncio.sleep(0.5)
+        raise Exception("exit_event set")
diff --git a/ton-http-api/requirements.txt b/ton-http-api/requirements.txt
@@ -1,4 +1,3 @@
-aioprocessing>=2.0.0
 aioredis>=2.0.1
 loguru>=0.6.0
 fastapi>=0.78.0
diff --git a/ton-http-api/setup.py b/ton-http-api/setup.py
@@ -16,7 +16,6 @@
     version=version,
     packages=find_packages('.', exclude=['tests']),
     install_requires=[
-        'aioprocessing>=2.0.0',
         'aioredis>=2.0.1',
         'loguru>=0.6.0',
         'fastapi>=0.78.0',

Original file line number	Diff line number	Diff line change
`@@ -1,4 +1,3 @@`
`1`		`-aioprocessing>=2.0.0`
`2`	`1`	`aioredis>=2.0.1`
`3`	`2`	`loguru>=0.6.0`
`4`	`3`	`fastapi>=0.78.0`