Handle OSError on shutdown & TimeoutError on recv (#974)

* Expect `TimeoutError` during `recv`

* Expect `OSError` during socket shutdown, can happen if other end has already closed the socket

Author: abhinavsingh

proxy/core/base/tcp_server.py

@@ -109,7 +109,11 @@ async def handle_writables(self, writables: Writables) -> bool:
     async def handle_readables(self, readables: Readables) -> bool:
         teardown = False
         if self.work.connection.fileno() in readables:
-            data = self.work.recv(self.flags.client_recvbuf_size)
+            try:
+                data = self.work.recv(self.flags.client_recvbuf_size)
+            except TimeoutError:
+                logger.info('Client recv timeout error')
+                return True
             if data is None:
                 logger.debug(
                     'Connection closed by client {0}'.format(

proxy/core/base/tcp_upstream.py

@@ -81,6 +81,9 @@ async def read_from_descriptors(self, r: Readables) -> bool:
                 else:
                     # Tear down because upstream proxy closed the connection
                     return True
+            except TimeoutError:
+                logger.info('Upstream recv timeout error')
+                return True
             except ssl.SSLWantReadError:
                 logger.info('Upstream SSLWantReadError, will retry')
                 return False

proxy/core/connection/pool.py

@@ -174,7 +174,10 @@ def _remove(self, fileno: int) -> None:
         """Remove a connection by descriptor from the internal data structure."""
         conn = self.connections[fileno]
         logger.debug('Removing conn#{0} from pool'.format(id(conn)))
-        conn.connection.shutdown(socket.SHUT_WR)
+        try:
+            conn.connection.shutdown(socket.SHUT_WR)
+        except OSError:
+            pass
         conn.close()
         self.pools[conn.addr].remove(conn)
         del self.connections[fileno]

proxy/http/websocket/client.py

@@ -115,5 +115,8 @@ def run(self) -> None:
         finally:
             if not self.closed:
                 self.selector.unregister(self.sock)
-                self.sock.shutdown(socket.SHUT_WR)
+                try:
+                    self.sock.shutdown(socket.SHUT_WR)
+                except OSError:
+                    pass
             self.sock.close()