From d74bb68f031b955474a9800e4e72a2f93d85bbd8 Mon Sep 17 00:00:00 2001 From: Alexander Bersenev Date: Wed, 27 Jun 2018 11:11:45 +0500 Subject: [PATCH] Revert "Revert "refactor task canceling a bit". The idea with the task cancelation doesn't work" This reverts commit b74079c4337939b695dc68c2a6979112eb368297. --- mtprotoproxy.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/mtprotoproxy.py b/mtprotoproxy.py index bc83a28..a470108 100755 --- a/mtprotoproxy.py +++ b/mtprotoproxy.py @@ -809,18 +809,22 @@ async def handle_client(reader_clt, writer_clt): update_stats(user, octets=len(data)) wr.write(data, extra) await wr.drain() - except (OSError, AttributeError, asyncio.streams.IncompleteReadError) as e: + except (OSError, asyncio.streams.IncompleteReadError) as e: # print_err(e) pass finally: wr.transport.abort() - task_tg_to_clt = connect_reader_to_writer(reader_tg, writer_clt, user) - task_clt_to_tg = connect_reader_to_writer(reader_clt, writer_tg, user) + task_tg_to_clt = asyncio.ensure_future(connect_reader_to_writer(reader_tg, writer_clt, user)) + task_clt_to_tg = asyncio.ensure_future(connect_reader_to_writer(reader_clt, writer_tg, user)) update_stats(user, curr_connects=1) await asyncio.wait([task_tg_to_clt, task_clt_to_tg], return_when=asyncio.FIRST_COMPLETED) update_stats(user, curr_connects=-1) + + task_tg_to_clt.cancel() + task_clt_to_tg.cancel() + writer_tg.transport.abort()