- Added a mutex to the identity map which mutexes

remove operations against iteration methods,
which now pre-buffer before returning an
iterable.   This because asyncrhonous gc
can remove items via the gc thread at any time.
[ticket:1891]

1 files changed, 49 insertions, 28 deletions

diff --git a/lib/sqlalchemy/orm/identity.py b/lib/sqlalchemy/orm/identity.py
index 4650b066f..30c3a06b7 100644
--- a/lib/sqlalchemy/orm/identity.py
+++ b/lib/sqlalchemy/orm/identity.py
@@ -15,7 +15,7 @@ class IdentityMap(dict):
         self._mutable_attrs = set()
         self._modified = set()
         self._wr = weakref.ref(self)
-
+        
     def replace(self, state):
         raise NotImplementedError()
         
@@ -61,7 +61,7 @@ class IdentityMap(dict):
             
     def has_key(self, key):
         return key in self
-        
+    
     def popitem(self):
         raise NotImplementedError("IdentityMap uses remove() to remove data")
 
@@ -81,6 +81,9 @@ class IdentityMap(dict):
         raise NotImplementedError("IdentityMap uses remove() to remove data")
         
 class WeakInstanceDict(IdentityMap):
+    def __init__(self):
+        IdentityMap.__init__(self)
+        self._remove_mutex = base_util.threading.Lock()
 
     def __getitem__(self, key):
         state = dict.__getitem__(self, key)
@@ -134,8 +137,13 @@ class WeakInstanceDict(IdentityMap):
         self.remove(state)
         
     def remove(self, state):
-        if dict.pop(self, state.key) is not state:
-            raise AssertionError("State %s is not present in this identity map" % state)
+        self._remove_mutex.acquire()
+        try:
+            if dict.pop(self, state.key) is not state:
+                raise AssertionError("State %s is not present in this identity map" % state)
+        finally:
+            self._remove_mutex.release()
+            
         self._manage_removed_state(state)
     
     def discard(self, state):
@@ -153,43 +161,56 @@ class WeakInstanceDict(IdentityMap):
         if o is None:
             return default
         return o
-    
-    # Py2K        
+
+
     def items(self):
+    # Py2K
         return list(self.iteritems())
-
+    
     def iteritems(self):
-        for state in dict.itervalues(self):
     # end Py2K
-    # Py3K
-    #def items(self):
-    #    for state in dict.values(self):
-            value = state.obj()
-            if value is not None:
-                yield state.key, value
+        self._remove_mutex.acquire()
+        try:
+            result = []
+            for state in dict.values(self):
+                value = state.obj()
+                if value is not None:
+                    result.append((state.key, value))
 
-    # Py2K
+            return iter(result)
+        finally:
+            self._remove_mutex.release()
+        
     def values(self):
+    # Py2K
         return list(self.itervalues())
 
     def itervalues(self):
-        for state in dict.itervalues(self):
     # end Py2K
-    # Py3K
-    #def values(self):
-    #    for state in dict.values(self):
-            instance = state.obj()
-            if instance is not None:
-                yield instance
+        self._remove_mutex.acquire()
+        try:
+            result = []
+            for state in dict.values(self):
+                value = state.obj()
+                if value is not None:
+                    result.append(value)
 
+            return iter(result)
+        finally:
+            self._remove_mutex.release()
+    
     def all_states(self):
-        # Py3K
-        # return list(dict.values(self))
+        self._remove_mutex.acquire()
+        try:
+            # Py3K
+            # return list(dict.values(self))
         
-        # Py2K
-        return dict.values(self)
-        # end Py2K
-    
+            # Py2K
+            return dict.values(self)
+            # end Py2K
+        finally:
+            self._remove_mutex.release()
+            
     def prune(self):
         return 0