asyncrpc: Set timeout when waiting for reply from server

Signed-off-by: Paul Barker <pbarker@konsulko.com>
Signed-off-by: Richard Purdie <richard.purdie@linuxfoundation.org>

1 files changed, 7 insertions, 2 deletions

diff --git a/lib/bb/asyncrpc/client.py b/lib/bb/asyncrpc/client.py
index d917ab597..3eb4fdde8 100644
--- a/lib/bb/asyncrpc/client.py
+++ b/lib/bb/asyncrpc/client.py
@@ -11,13 +11,14 @@ from . import chunkify, DEFAULT_MAX_CHUNK
 
 
 class AsyncClient(object):
-    def __init__(self, proto_name, proto_version, logger):
+    def __init__(self, proto_name, proto_version, logger, timeout=30):
         self.reader = None
         self.writer = None
         self.max_chunk = DEFAULT_MAX_CHUNK
         self.proto_name = proto_name
         self.proto_version = proto_version
         self.logger = logger
+        self.timeout = timeout
 
     async def connect_tcp(self, address, port):
         async def connect_sock():
@@ -70,7 +71,11 @@ class AsyncClient(object):
 
     async def send_message(self, msg):
         async def get_line():
-            line = await self.reader.readline()
+            try:
+                line = await asyncio.wait_for(self.reader.readline(), self.timeout)
+            except asyncio.TimeoutError:
+                raise ConnectionError("Timed out waiting for server")
+
             if not line:
                 raise ConnectionError("Connection closed")