Merge pull request #34 from dlecocq/dan/connection-closed-exception

Connection Closed Exception

nsq/client.py

@@ -10,7 +10,7 @@
 
 from contextlib import contextmanager
 import random
-from select import select
+import select
 import socket
 import time
 import threading
@@ -182,8 +182,15 @@ def read(self):
         # Not all connections need to be written to, so we'll only concern
         # ourselves with those that require writes
         writes = [c for c in connections if c.pending()]
-        readable, writable, exceptable = select(
-            connections, writes, connections, self._timeout)
+        try:
+            readable, writable, exceptable = select.select(
+                connections, writes, connections, self._timeout)
+        except exceptions.ConnectionClosedException:
+            logger.exception('Tried selecting on closed client')
+            return []
+        except select.error:
+            logger.exception('Error running select')
+            return []
 
         # If we returned because the timeout interval passed, log it and return
         if not (readable or writable or exceptable):

nsq/connection.py

@@ -4,7 +4,7 @@
 from . import util
 from . import json
 from . import __version__
-from .exceptions import UnsupportedException
+from .exceptions import UnsupportedException, ConnectionClosedException
 from .sockets import TLSSocket, SnappySocket, DeflateSocket
 from .response import Response, Message
 
@@ -208,7 +208,9 @@ def setblocking(self, blocking):
     def fileno(self):
         '''Returns the socket's fileno. This allows us to select on this'''
         for sock in self.socket():
-            return sock.fileno()
+            if sock:
+                return sock.fileno()
+        raise ConnectionClosedException()
 
     def pending(self):
         '''All of the messages waiting to be sent'''

nsq/exceptions.py

@@ -5,6 +5,10 @@ class NSQException(StandardError):
     '''Base class for all exceptions in this library'''
 
 
+class ConnectionClosedException(NSQException):
+    '''Trying to use a closed connection as if it's alive'''
+
+
 class UnsupportedException(NSQException):
     '''When a requested feature cannot be used'''
 

test/test_client.py

@@ -8,6 +8,8 @@
 
 from common import HttpClientIntegrationTest, MockedConnectionTest
 from contextlib import contextmanager
+import errno
+import select
 import socket
 
 
@@ -61,6 +63,21 @@ def test_conection_checker(self):
             self.assertTrue(checker.is_alive())
         self.assertFalse(checker.is_alive())
 
+    def test_read_closed(self):
+        '''Recovers from reading on a closed connection'''
+        conn = self.client.connections()[0]
+        with mock.patch.object(conn, 'alive', return_value=True):
+            with mock.patch.object(conn, '_socket', None):
+                # This test passes if no exception in raised
+                self.client.read()
+
+    def test_read_select_err(self):
+        '''Recovers from select errors'''
+        with mock.patch('nsq.client.select.select') as mock_select:
+            mock_select.side_effect = select.error(errno.EBADF)
+            # This test passes if no exception is raised
+            self.client.read()
+
 
 class TestClientLookupd(HttpClientIntegrationTest):
     '''Test our client class'''
@@ -119,21 +136,21 @@ class TestClientMultiple(MockedConnectionTest):
     def readable(self, connections):
         '''With all the connections readable'''
         value = (connections, [], [])
-        with mock.patch.object(client, 'select', return_value=value):
+        with mock.patch('nsq.client.select.select', return_value=value):
             yield
 
     @contextmanager
     def writable(self, connections):
         '''With all the connections writable'''
         value = ([], connections, [])
-        with mock.patch.object(client, 'select', return_value=value):
+        with mock.patch('nsq.client.select.select', return_value=value):
             yield
 
     @contextmanager
     def exceptable(self, connections):
         '''With all the connections exceptable'''
         value = ([], [], connections)
-        with mock.patch.object(client, 'select', return_value=value):
+        with mock.patch('nsq.client.select.select', return_value=value):
             yield
 
     def test_multi_read(self):

test/test_connection.py

@@ -219,6 +219,12 @@ def test_fileno(self):
         self.assertEqual(
             self.connection.fileno(), self.socket.fileno())
 
+    def test_fileno_closed(self):
+        '''Raises an exception if the connection's closed'''
+        with mock.patch.object(self.connection, '_socket', None):
+            self.assertRaises(exceptions.ConnectionClosedException,
+                self.connection.fileno)
+
     def test_str_alive(self):
         '''Sane str representation for an alive connection'''
         with mock.patch.object(self.connection, 'alive', return_value=True):