Better client disconnection handling.

sjrpc/server/simple.py

@@ -41,8 +41,12 @@ class SimpleRpcServer(ConnectionManager):
             connection.shutdown(self._on_disconnect)
         self._listening_sock.close()
 
-    def shutdown_client(self, connection):
-        connection.shutdown(callback=self._on_disconnect)
+    def shutdown_client(self, fd):
+        conn = self._clients.get(fd)
+        self._poll.unregister(fd)
+        if fd is not None:
+            del self._clients[fd]
+            conn.shutdown(callback=self._on_disconnect)
 
     def all_connections(self):
         return set(self._clients.values())
@@ -69,13 +73,11 @@ class SimpleRpcServer(ConnectionManager):
                 except socket.error as err:
                     logging.error('Socket error while receiving from client '
                                   'fd/%s: %s' % (fd, err))
-                    connection.shutdown(self._on_disconnect)
-                    del self._clients[fd]
+                    self.shutdown_client(fd)
                 except Exception as err:
                     logging.error('Unknown error while receiving from client '
                                   'fd/%s: %s' % (fd, err))
-                    connection.shutdown(self._on_disconnect)
-                    del self._clients[fd]
+                    self.shutdown_client(fd)
                 
             if event & select.EPOLLOUT:
                 # Data are ready to be written on socket
@@ -84,18 +86,15 @@ class SimpleRpcServer(ConnectionManager):
                 except socket.error as err:
                     logging.error('Socket error while sending to the client '
                                   'fd/%s: %s' % (fd, err))
-                    connection.shutdown(self._on_disconnect)
-                    del self._clients[fd]
+                    self.shutdown_client(fd)
                 except Exception as err:
                     logging.error('Unknown error while sending to the client '
                                   'fd/%s: %s' % (fd, err))
-                    connection.shutdown(self._on_disconnect)
-                    del self._clients[fd]
+                    self.shutdown_client(fd)
 
             if event & select.EPOLLHUP:
                 logging.error('Socket HUP fd/%s: %s' % (fd, err))
-                connection.shutdown(self._on_disconnect)
-                del self._clients[fd]
+                self.shutdown_client(fd)
 
 
 class SimpleSslRpcServer(SimpleRpcServer):