cancel requests for ws and honor timeout if set

Author: joente

thingsdb/client/client.py

@@ -2,6 +2,7 @@
 import logging
 import random
 import ssl
+import time
 from collections import defaultdict
 from ssl import SSLContext, PROTOCOL_TLS
 from typing import Optional, Union, Any
@@ -360,6 +361,9 @@ async def _ensure_write(
     ) -> asyncio.Future:
         if not self._pool:
             raise ConnectionError('no connection')
+
+        start = time.time()
+
         while True:
             if not self.is_connected():
                 logging.info('Wait for a connection')
@@ -369,7 +373,13 @@ async def _ensure_write(
 
             try:
                 res = await self._protocol.write(tp, data, is_bin, timeout)
-            except (CancelledError, NodeError, AuthError) as e:
+            except (asyncio.exceptions.CancelledError,
+                    CancelledError, NodeError, AuthError) as e:
+                if timeout and time.time() - start > timeout:
+                    msg = str(e) or type(e).__name__
+                    raise asyncio.TimeoutError(
+                        f'failed to transmit within timeout (error: {msg})')
+
                 logging.error(
                     f'Failed to transmit package: '
                     f'{e}({e.__class__.__name__}) (will try again)')

thingsdb/client/protocol.py

@@ -222,6 +222,19 @@ def write(
         self._requests[self._pid] = (future, task)
         return future
 
+    def cancel_requests(self):
+        if self._requests:
+            logging.error(
+                f'Canceling {len(self._requests)} requests '
+                'due to a lost connection'
+            )
+            while self._requests:
+                _key, (future, task) = self._requests.popitem()
+                if task is not None:
+                    task.cancel()
+                if not future.cancelled():
+                    future.cancel()
+
     @abstractmethod
     def _write(self, data: Any):
         ...
@@ -269,18 +282,7 @@ def connection_lost(self, exc: Exception) -> None:
         '''
         override asyncio.Protocol
         '''
-        if self._requests:
-            logging.error(
-                f'Canceling {len(self._requests)} requests '
-                'due to a lost connection'
-            )
-            while self._requests:
-                _key, (future, task) = self._requests.popitem()
-                if task is not None:
-                    task.cancel()
-                if not future.cancelled():
-                    future.cancel()
-
+        self.cancel_requests()
         self.close_future.set_result(None)
         self.close_future = None
         self.transport = None
@@ -380,6 +382,7 @@ async def _recv_loop(self):
                     self._handle_package(pkg)
 
         except ConnectionClosed as exc:
+            self.cancel_requests()
             self._proto = None
             self._on_connection_lost(self, exc)
 

thingsdb/version.py

@@ -1 +1 @@
-__version__ = '1.1.4'
+__version__ = '1.1.5'