From 7149fb6d85be4183c8620596ff133f9d2b9e827c Mon Sep 17 00:00:00 2001 From: atse Date: Wed, 2 Aug 2023 14:43:50 +0800 Subject: [PATCH] Shield pending commands from cancellation --- pytec/pytec/aioclient.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/pytec/pytec/aioclient.py b/pytec/pytec/aioclient.py index ab4fb6a..aefa07e 100644 --- a/pytec/pytec/aioclient.py +++ b/pytec/pytec/aioclient.py @@ -74,12 +74,16 @@ class Client: chunk = await self._reader.readline() return chunk.decode('utf-8', errors='ignore') + async def _read_write(self, command): + self._writer.write(((" ".join(command)).strip() + "\n").encode('utf-8')) + await self._writer.drain() + + return await self._read_line() + async def _command(self, *command): async with self._command_lock: - self._writer.write(((" ".join(command)).strip() + "\n").encode('utf-8')) - await self._writer.drain() - - line = await self._read_line() + # protect the read-write process from being cancelled midway + line = await asyncio.shield(self._read_write(command)) response = json.loads(line) logging.debug(f"{command}: {response}")