PYTHON-1070 - Make index cache thread safe

Author: behackett

pymongo/collection.py

@@ -1402,6 +1402,12 @@ def ensure_index(self, key_or_list, cache_for=300, **kwargs):
  keys = helpers._index_list(key_or_list)
  name = kwargs.setdefault("name", helpers._gen_index_name(keys))
 
+ # Note that there is a race condition here. One thread could
+ # check if the index is cached and be preempted before creating
+ # and caching the index. This means multiple threads attempting
+ # to create the same index concurrently could send the index
+ # to the server two or more times. This has no practical impact
+ # other than wasted round trips.
  if not self.__database.client._cached(self.__database.name,
  self.__name, name):
  self.__create_index(keys, kwargs)

pymongo/mongo_client.py

@@ -353,6 +353,7 @@ def __init__(
 
  # Cache of existing indexes used by ensure_index ops.
  self.__index_cache = {}
+ self.__index_cache_lock = threading.Lock()
 
  super(MongoClient, self).__init__(options.codec_options,
  options.read_preference,
@@ -433,27 +434,29 @@ def _cached(self, dbname, coll, index):
  """Test if `index` is cached."""
  cache = self.__index_cache
  now = datetime.datetime.utcnow()
- return (dbname in cache and
- coll in cache[dbname] and
- index in cache[dbname][coll] and
- now < cache[dbname][coll][index])
+ with self.__index_cache_lock:
+ return (dbname in cache and
+ coll in cache[dbname] and
+ index in cache[dbname][coll] and
+ now < cache[dbname][coll][index])
 
  def _cache_index(self, dbname, collection, index, cache_for):
  """Add an index to the index cache for ensure_index operations."""
  now = datetime.datetime.utcnow()
  expire = datetime.timedelta(seconds=cache_for) + now
 
- if database not in self.__index_cache:
- self.__index_cache[dbname] = {}
- self.__index_cache[dbname][collection] = {}
- self.__index_cache[dbname][collection][index] = expire
+ with self.__index_cache_lock:
+ if database not in self.__index_cache:
+ self.__index_cache[dbname] = {}
+ self.__index_cache[dbname][collection] = {}
+ self.__index_cache[dbname][collection][index] = expire
 
- elif collection not in self.__index_cache[dbname]:
- self.__index_cache[dbname][collection] = {}
- self.__index_cache[dbname][collection][index] = expire
+  elif collection not in self.__index_cache[dbname]:
+  self.__index_cache[dbname][collection] = {}
+  self.__index_cache[dbname][collection][index] = expire
 
- else:
- self.__index_cache[dbname][collection][index] = expire
+  else:
+  self.__index_cache[dbname][collection][index] = expire
 
  def _purge_index(self, database_name,
  collection_name=None, index_name=None):
@@ -463,22 +466,23 @@ def _purge_index(self, database_name,
 
  If `collection_name` is None purge an entire database.
  """
- if not database_name in self.__index_cache:
- return
+ with self.__index_cache_lock:
+ if not database_name in self.__index_cache:
+ return
 
- if collection_name is None:
- del self.__index_cache[database_name]
- return
+  if collection_name is None:
+  del self.__index_cache[database_name]
+  return
 
- if not collection_name in self.__index_cache[database_name]:
- return
+  if not collection_name in self.__index_cache[database_name]:
+  return
 
- if index_name is None:
- del self.__index_cache[database_name][collection_name]
- return
+  if index_name is None:
+  del self.__index_cache[database_name][collection_name]
+  return
 
- if index_name in self.__index_cache[database_name][collection_name]:
- del self.__index_cache[database_name][collection_name][index_name]
+  if index_name in self.__index_cache[database_name][collection_name]:
+  del self.__index_cache[database_name][collection_name][index_name]
 
  def _server_property(self, attr_name):
  """An attribute of the current server's description.

test/test_legacy_api.py

@@ -1118,6 +1118,66 @@ def test_ensure_index(self):
  # Clean up indexes for later tests
  db.test.drop_indexes()
 
+ def test_ensure_index_threaded(self):
+ coll = self.db.threaded_index_creation
+ index_docs = []
+
+ class Indexer(threading.Thread):
+ def run(self):
+ coll.ensure_index('foo0')
+ coll.ensure_index('foo1')
+ coll.ensure_index('foo2')
+ index_docs.append(coll.index_information())
+
+ try:
+ threads = []
+ for _ in range(10):
+ t = Indexer()
+ t.setDaemon(True)
+ threads.append(t)
+
+ for thread in threads:
+ thread.start()
+
+ joinall(threads)
+
+ first = index_docs[0]
+ for index_doc in index_docs[1:]:
+ self.assertEqual(index_doc, first)
+ finally:
+ coll.drop()
+
+ def test_ensure_purge_index_threaded(self):
+ coll = self.db.threaded_index_creation
+
+ class Indexer(threading.Thread):
+ def run(self):
+ coll.ensure_index('foo')
+ try:
+ coll.drop_index('foo')
+ except OperationFailure:
+ # The index may have already been dropped.
+ pass
+ coll.ensure_index('foo')
+ coll.drop_indexes()
+ coll.ensure_index('foo')
+
+ try:
+ threads = []
+ for _ in range(10):
+ t = Indexer()
+ t.setDaemon(True)
+ threads.append(t)
+
+ for thread in threads:
+ thread.start()
+
+ joinall(threads)
+
+ self.assertTrue('foo_1' in coll.index_information())
+ finally:
+ coll.drop()
+
  def test_ensure_unique_index_threaded(self):
  coll = self.db.test_unique_threaded
  coll.drop()