translate socket.timeout to TTransportException

Thriftpy · Jan 28, 2021 · e448470 · e448470
1 parent 108cca5
commit e448470
Showing 1 changed file with 12 additions and 1 deletion.
diff --git a/thriftpy2/transport/socket.py b/thriftpy2/transport/socket.py
@@ -108,6 +108,11 @@ def read(self, sz):
         while True:
             try:
                 buff = self.sock.recv(sz)
+            except socket.timeout:
+                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 +138,13 @@ def read(self, sz):
         return buff
 
     def write(self, buff):
-        self.sock.sendall(buff)
+        try:
+            self.sock.sendall(buff)
+        except socket.timeout:
+            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