Immediately remove closed websocket clients from weakset

When a client closes the connection, immediately remove the WebSocket
instance from the instances list, rather than waiting for the heartbeat
to time the connection out and garbage collect it.

This is a small optimisation to ensure that the stats on number of
connected clients are as accurate as possible, and that we don't waste
CPU time iterating over sockets that are terminated.

h/streamer/test/websocket_test.py

@@ -13,6 +13,29 @@
 FakeMessage = namedtuple('FakeMessage', ['data'])
 
 
+def test_websocket_stores_instance_list():
+    socket = mock.Mock()
+    clients = [websocket.WebSocket(socket), websocket.WebSocket(socket)]
+
+    for c in clients:
+        assert c in websocket.WebSocket.instances
+
+
+def test_websocket_removes_self_from_instance_list_when_closed():
+    socket = mock.Mock()
+    client1 = websocket.WebSocket(socket)
+    client2 = websocket.WebSocket(socket)
+
+    assert len(websocket.WebSocket.instances) == 2
+    client1.closed(1000)
+    assert client1 not in websocket.WebSocket.instances
+    client2.closed(1000)
+    assert client2 not in websocket.WebSocket.instances
+
+    # A second closure (however unusual) should not raise
+    client1.closed(1000)
+
+
 def test_socket_enqueues_incoming_messages():
     queue = Queue()
     request = testing.DummyRequest()

h/streamer/websocket.py

@@ -49,6 +49,12 @@ def received_message(self, msg):
             log.warn('Streamer work queue full! Unable to queue message from '
                      'WebSocket client having waited 0.1s: giving up.')
 
+    def closed(self, code, reason=None):
+        try:
+            self.instances.remove(self)
+        except KeyError:
+            pass
+
 
 def handle_message(message):
     socket = message.socket