Fix race condition in ensure_index PYTHON-284

behackett · behackett · commit 2bba32b3d5bf · 2011-11-06T19:37:58.000-08:00
diff --git a/pymongo/collection.py b/pymongo/collection.py
@@ -764,8 +764,8 @@ def ensure_index(self, key_or_list, deprecated_unique=None,
             keys = helpers._index_list(key_or_list)
             name = kwargs["name"] = _gen_index_name(keys)
 
-        if self.__database.connection._cache_index(self.__database.name,
-                                                   self.__name, name, ttl):
+        if not self.__database.connection._cached(self.__database.name,
+                                                  self.__name, name):
             return self.create_index(key_or_list, deprecated_unique,
                                      ttl, **kwargs)
         return None
diff --git a/pymongo/connection.py b/pymongo/connection.py
@@ -407,13 +407,18 @@ def paired(cls, left, right=None, **connection_args):
         return cls([":".join(map(str, left)), ":".join(map(str, right))],
                    **connection_args)
 
+    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])
+
     def _cache_index(self, database, collection, index, ttl):
         """Add an index to the index cache for ensure_index operations.
-
-        Return ``True`` if the index has been newly cached or if the index had
-        expired and is being re-cached.
-
-        Return ``False`` if the index exists and is valid.
         """
         now = datetime.datetime.utcnow()
         expire = datetime.timedelta(seconds=ttl) + now
@@ -422,19 +427,13 @@ def _cache_index(self, database, collection, index, ttl):
             self.__index_cache[database] = {}
             self.__index_cache[database][collection] = {}
             self.__index_cache[database][collection][index] = expire
-            return True
 
-        if collection not in self.__index_cache[database]:
+        elif collection not in self.__index_cache[database]:
             self.__index_cache[database][collection] = {}
             self.__index_cache[database][collection][index] = expire
-            return True
 
-        if index in self.__index_cache[database][collection]:
-            if now < self.__index_cache[database][collection][index]:
-                return False
-
-        self.__index_cache[database][collection][index] = expire
-        return True
+        else:
+            self.__index_cache[database][collection][index] = expire
 
     def _purge_index(self, database_name,
                      collection_name=None, index_name=None):
diff --git a/pymongo/master_slave_connection.py b/pymongo/master_slave_connection.py
@@ -270,6 +270,10 @@ def __iter__(self):
     def next(self):
         raise TypeError("'MasterSlaveConnection' object is not iterable")
 
+    def _cached(self, database_name, collection_name, index_name):
+        return self.__master._cached(database_name,
+                                     collection_name, index_name)
+
     def _cache_index(self, database_name, collection_name, index_name, ttl):
         return self.__master._cache_index(database_name, collection_name,
                                           index_name, ttl)
diff --git a/pymongo/replica_set_connection.py b/pymongo/replica_set_connection.py
@@ -219,13 +219,18 @@ def __init__(self, hosts_or_uri=None, max_pool_size=10,
             if not self[db_name].authenticate(username, password):
                 raise ConfigurationError("authentication failed")
 
+    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])
+
     def _cache_index(self, dbase, collection, index, ttl):
         """Add an index to the index cache for ensure_index operations.
-
-        Return ``True`` if the index has been newly cached or if the index had
-        expired and is being re-cached.
-
-        Return ``False`` if the index exists and is valid.
         """
         now = datetime.datetime.utcnow()
         expire = datetime.timedelta(seconds=ttl) + now
@@ -234,19 +239,13 @@ def _cache_index(self, dbase, collection, index, ttl):
             self.__index_cache[dbase] = {}
             self.__index_cache[dbase][collection] = {}
             self.__index_cache[dbase][collection][index] = expire
-            return True
 
-        if collection not in self.__index_cache[dbase]:
+        elif collection not in self.__index_cache[dbase]:
             self.__index_cache[dbase][collection] = {}
             self.__index_cache[dbase][collection][index] = expire
-            return True
 
-        if index in self.__index_cache[dbase][collection]:
-            if now < self.__index_cache[dbase][collection][index]:
-                return False
-
-        self.__index_cache[dbase][collection][index] = expire
-        return True
+        else:
+            self.__index_cache[dbase][collection][index] = expire
 
     def _purge_index(self, database_name,
                      collection_name=None, index_name=None):
diff --git a/test/test_collection.py b/test/test_collection.py
@@ -19,6 +19,7 @@
 import itertools
 import re
 import sys
+import threading
 import time
 import unittest
 import warnings
@@ -176,9 +177,40 @@ def test_ensure_index(self):
         time.sleep(1.1)
         self.assertEqual("goodbye_1",
                          db.test.ensure_index("goodbye"))
+        # Make sure the expiration time is updated.
+        self.assertEqual(None,
+                         db.test.ensure_index("goodbye"))
         # Clean up indexes for later tests
         db.test.drop_indexes()
 
+    def test_ensure_unique_index_threaded(self):
+        db = self.db
+        db.test.drop()
+        db.test.insert({'foo': i} for i in xrange(10000))
+
+        class Indexer(threading.Thread):
+            def run(self):
+                try:
+                    db.test.ensure_index('foo', unique=True)
+                    db.test.insert({'foo': 'bar'})
+                    db.test.insert({'foo': 'bar'})
+                except OperationFailure:
+                    pass
+
+        threads = []
+        for _ in xrange(10):
+            t = Indexer()
+            threads.append(t)
+
+        for i in xrange(10):
+            threads[i].start()
+
+        for i in xrange(10):
+            threads[i].join()
+
+        self.assertEqual(10001, db.test.count())
+        db.test.drop()
+
     def test_index_on_binary(self):
         db = self.db
         db.drop_collection("test")
