Discard a socket after an exception to ensure we don't have inconsistent state PYTHON-294

Author: A. Jesse Jiryu Davis

pymongo/connection.py

@@ -165,6 +165,12 @@ def get_socket(self, host, port):
             self.sock = (pid, self.connect(host, port))
             return (self.sock[1], False)
 
+    def discard_socket(self):
+        """Close and discard the active socket.
+        """
+        if self.sock:
+            self.sock[1].close()
+            self.sock = None
 
     def return_socket(self):
         if self.sock is not None and self.sock[0] == os.getpid():
@@ -827,7 +833,18 @@ def __receive_data_on_socket(self, length, sock):
         """
         message = ""
         while len(message) < length:
-            chunk = sock.recv(length - len(message))
+            try:
+                chunk = sock.recv(length - len(message))
+            except BaseException, e:
+                # PYTHON-294: must close socket here, because if it remains open
+                # after this then the next time we read from it we may get stale
+                # data from a previous operation.
+                try:
+                    self.__pool.discard_socket()
+                except:
+                    pass
+                raise e
+
             if chunk == "":
                 raise ConnectionFailure("connection closed")
             message += chunk
@@ -877,7 +894,11 @@ def _send_message_with_response(self, message,
                 raise AutoReconnect(str(e))
         finally:
             if "network_timeout" in kwargs:
-                sock.settimeout(self.__net_timeout)
+                try:
+                    sock.settimeout(self.__net_timeout)
+                except socket.error:
+                    # There was an exception and we've closed the socket
+                    pass
 
     def start_request(self):
         """DEPRECATED all operations will start a request.

test/test_connection.py

@@ -476,6 +476,57 @@ def test_contextlib(self):
         self.assertEqual(None, connection._Connection__pool.sock)
         self.assertEqual(0, len(connection._Connection__pool.sockets))
 
+    def test_interrupt_signal(self):
+        """
+        Test fix for PYTHON-294 -- make sure Connection closes its socket if it
+        gets an interrupt while waiting to recv() from it.
+        """
+        import pymongo
+        import signal
+
+        c = get_connection()
+        db = c.pymongo_test
+
+        # A $where clause which takes 1 sec to execute
+        where = '''function() {
+            var d = new Date((new Date()).getTime() + 1*1000);
+            while (d > (new Date())) { }; return true;
+        }'''
+
+        # Need exactly 1 document so find() will execute its $where clause once
+        db.drop_collection('foo')
+        db.foo.insert({'_id': 1}, safe=True)
+
+        # Convert SIGALRM to SIGINT -- it's hard to schedule a SIGINT for 1/2 a
+        # second in the future, but easy to schedule SIGALRM.
+        def sigalarm(num, frame):
+            import sys
+            raise KeyboardInterrupt
+
+        old_signal_handler = signal.signal(signal.SIGALRM, sigalarm)
+        signal.setitimer(signal.ITIMER_REAL, 0.5, 0)
+        raised = False
+        try:
+            # Need list() to actually iterate the cursor and create the
+            # cursor on the server; find is lazily evaluated. The find() will
+            # be interrupted by sigalarm() raising a KeyboardInterrupt.
+            list(db.foo.find({'$where': where}))
+        except KeyboardInterrupt:
+            raised = True
+        finally:
+            signal.signal(signal.SIGALRM, old_signal_handler)
+
+        # Can't use self.assertRaises() because it doesn't catch system
+        # exceptions
+        self.assert_(raised, "Didn't raise expected KeyboardInterrupt")
+
+        # Raises AssertionError due to PYTHON-294 -- Mongo's response to the
+        # previous find() is still waiting to be read on the socket, so the
+        # request id's don't match.
+        self.assertEqual(
+            [{'_id': 1}],
+            list(db.foo.find())
+        )
 
 if __name__ == "__main__":
     unittest.main()