translate socket.timeout to TTransportException

thriftpy2/rpc.py

@@ -31,13 +31,15 @@ def make_client(service, host="localhost", port=9090, unix_socket=None,
                 proto_factory=TBinaryProtocolFactory(),
                 trans_factory=TBufferedTransportFactory(),
                 timeout=3000, cafile=None, ssl_context=None, certfile=None,
-                keyfile=None, url="", socket_family=socket.AF_INET):
+                keyfile=None, url="", socket_family=socket.AF_INET,
+                handle_timeout_error=False):
     if url:
         parsed_url = urllib.parse.urlparse(url)
         host = parsed_url.hostname or host
         port = parsed_url.port or port
     if unix_socket:
-        socket = TSocket(unix_socket=unix_socket, socket_timeout=timeout)
+        socket = TSocket(unix_socket=unix_socket, socket_timeout=timeout,
+                         handle_timeout_error=handle_timeout_error)
         if certfile:
             warnings.warn("SSL only works with host:port, not unix_socket.")
     elif host and port:

thriftpy2/transport/socket.py

@@ -16,7 +16,7 @@ class TSocket(object):
 
     def __init__(self, host=None, port=None, unix_socket=None,
                  sock=None, socket_family=socket.AF_INET,
-                 socket_timeout=3000, connect_timeout=None):
+                 socket_timeout=3000, connect_timeout=None, handle_timeout_error=False):
         """Initialize a TSocket
 
         TSocket can be initialized in 3 ways:
@@ -35,6 +35,8 @@ def __init__(self, host=None, port=None, unix_socket=None,
         @param socket_timeout   socket timeout in ms
         @param connect_timeout  connect timeout in ms, only used in
             connection, will be set to socket_timeout if not set.
+        @param handle_timeout_error(bool)    Whether translate socket.timeout
+            error to TTransportException. Default is False for compalibility.
         """
         if sock:
             self.sock = sock
@@ -54,6 +56,8 @@ def __init__(self, host=None, port=None, unix_socket=None,
         self.connect_timeout = connect_timeout / 1000 if connect_timeout \
             else self.socket_timeout
 
+        self.handle_timeout_error = handle_timeout_error
+
     def _init_sock(self):
         if self.unix_socket:
             _sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
@@ -108,6 +112,13 @@ def read(self, sz):
         while True:
             try:
                 buff = self.sock.recv(sz)
+            except socket.timeout:
+                if not self.handle_timeout_error:
+                    raise
+                addr = self.sock.getsockname()
+                typ = TTransportException.TIMED_OUT
+                msg = "Timeouted when read from %s" % str(addr)
+                raise TTransportException(type=typ, message=msg)
             except socket.error as e:
                 if e.errno == errno.EINTR:
                     continue
@@ -133,7 +144,15 @@ def read(self, sz):
         return buff
 
     def write(self, buff):
-        self.sock.sendall(buff)
+        try:
+            self.sock.sendall(buff)
+        except socket.timeout:
+            if not self.handle_timeout_error:
+                raise
+            addr = self.sock.getsockname()
+            typ = TTransportException.TIMED_OUT
+            msg = "Timeouted when write to %s" % str(addr)
+            raise TTransportException(type=typ, message=msg)
 
     def flush(self):
         pass