Shield pending commands from cancellation

This commit is contained in:
atse 2023-08-02 14:43:50 +08:00 committed by Tse Kwok Yan
parent 535954abae
commit 79daa75383
1 changed files with 8 additions and 4 deletions

View File

@ -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}")