Don't reuse socket after an exception PYTHON-398

Prevent sockets from being returned to the pool if they're in an inconsistent state due to an async exception from PyThreadState_SetAsyncExc(). Connection and ReplicaSetConnection always use sockets the same way: first sendall() with a request, then __receive_message_on_socket() to get the response. (Named __recv_msg in RSC.) A socket is vulnerable to async exceptions starting at sendall() and ending after __receive_message_on_socket(). If both are completed, then the socket can be returned to the pool. If an exception's raised during this period, the socket must be closed. I wrap each section comprising sendall() and __receive_message_on_socket() like: try: sock_info.sock.sendall(msg) response = self.__receive_message_on_socket(1, rqst_id, sock_info) except: sock_info.close() raise Since the socket's pool of origin isn't easily available in every function where the socket must be closed, I close it directly instead of calling pool.discard_socket(sock_info). This is safe, because Pool checks sock_info.closed everywhere that matters, so the socket won't be reused after it's closed. Calling close() instead of discard_socket() also keeps Connection and RSC from further divergence, which is good in case we ever refactor them.

Author: A. Jesse Jiryu Davis

pymongo/connection.py

@@ -511,8 +511,13 @@ def __simple_command(self, sock_info, dbname, spec):
  """
  rqst_id, msg, _ = message.query(0, dbname + '.$cmd', 0, -1, spec)
  start = time.time()
- sock_info.sock.sendall(msg)
- response = self.__receive_message_on_socket(1, rqst_id, sock_info)
+ try:
+ sock_info.sock.sendall(msg)
+ response = self.__receive_message_on_socket(1, rqst_id, sock_info)
+ except:
+ sock_info.close()
+ raise
+
  end = time.time()
  response = helpers._unpack_response(response)['data'][0]
  msg = "command %r failed: %%s" % spec
@@ -838,6 +843,9 @@ def _send_message(self, message, with_last_error=False):
  except (ConnectionFailure, socket.error), e:
  self.disconnect()
  raise AutoReconnect(str(e))
+ except:
+ sock_info.close()
+ raise
 
  def __receive_data_on_socket(self, length, sock_info):
  """Lowest level receive operation.
@@ -847,12 +855,7 @@ def __receive_data_on_socket(self, length, sock_info):
  """
  chunks = []
  while length:
- try:
- chunk = sock_info.sock.recv(length)
- except:
- # recv was interrupted
- self.__pool.discard_socket(sock_info)
- raise
+ chunk = sock_info.sock.recv(length)
  if chunk == EMPTY:
  raise ConnectionFailure("connection closed")
  length -= len(chunk)
@@ -877,8 +880,12 @@ def __send_and_receive(self, message, sock_info):
  """Send a message on the given socket and return the response data.
  """
  (request_id, data) = self.__check_bson_size(message)
- sock_info.sock.sendall(data)
- return self.__receive_message_on_socket(1, request_id, sock_info)
+ try:
+ sock_info.sock.sendall(data)
+ return self.__receive_message_on_socket(1, request_id, sock_info)
+ except:
+ sock_info.close()
+ raise
 
  # we just ignore _must_use_master here: it's only relevant for
  # MasterSlaveConnection instances.

pymongo/replica_set_connection.py

@@ -636,8 +636,13 @@ def __simple_command(self, sock_info, dbname, spec):
  """
  rqst_id, msg, _ = message.query(0, dbname + '.$cmd', 0, -1, spec)
  start = time.time()
- sock_info.sock.sendall(msg)
- response = self.__recv_msg(1, rqst_id, sock_info)
+ try:
+ sock_info.sock.sendall(msg)
+ response = self.__recv_msg(1, rqst_id, sock_info)
+ except:
+ sock_info.close()
+ raise
+
  end = time.time()
  response = helpers._unpack_response(response)['data'][0]
  msg = "command %r failed: %%s" % spec
@@ -1011,7 +1016,7 @@ def _send_message(self, msg, safe=False, _connection_to_use=None):
  self.disconnect()
  raise AutoReconnect(str(why))
  except:
- member.pool.discard_socket(sock_info)
+ sock_info.close()
  raise
 
  def __send_and_receive(self, member, msg, **kwargs):
@@ -1038,7 +1043,7 @@ def __send_and_receive(self, member, msg, **kwargs):
  member.pool.discard_socket(sock_info)
  raise AutoReconnect("%s:%d: %s" % (host, port, str(why)))
  except:
- member.pool.discard_socket(sock_info)
+ sock_info.close()
  raise
 
  def __try_read(self, member, msg, **kwargs):