Always check that discarded state is the expected one

Fixed race condition in ORM identity map which would cause objects
to be inappropriately removed during a load operation, causing
duplicate object identities to occur, particularly under joined eager
loading which involves deduplication of objects.  The issue is specific
to garbage collection of weak references and is observed only under the
Pypy interpreter.

Change-Id: I9f6ae3fe5b078f26146af82b15d16f3a549a9032
Fixes: #4068

1 files changed, 22 insertions, 11 deletions

diff --git a/lib/sqlalchemy/orm/identity.py b/lib/sqlalchemy/orm/identity.py
index 8f4304ad2..e1bc9bb98 100644
--- a/lib/sqlalchemy/orm/identity.py
+++ b/lib/sqlalchemy/orm/identity.py
@@ -209,13 +209,19 @@ class WeakInstanceDict(IdentityMap):
             return list(self._dict.values())
 
     def _fast_discard(self, state):
-        self._dict.pop(state.key, None)
+        # used by InstanceState for state being
+        # GC'ed, inlines _managed_removed_state
+        try:
+            st = self._dict[state.key]
+        except KeyError:
+            # catch gc removed the key after we just checked for it
+            pass
+        else:
+            if st is state:
+                self._dict.pop(state.key, None)
 
     def discard(self, state):
-        st = self._dict.pop(state.key, None)
-        if st:
-            assert st is state
-            self._manage_removed_state(state)
+        self.safe_discard(state)
 
     def safe_discard(self, state):
         if state.key in self._dict:
@@ -311,14 +317,19 @@ class StrongInstanceDict(IdentityMap):
         state._instance_dict = self._wr
 
     def _fast_discard(self, state):
-        self._dict.pop(state.key, None)
+        # used by InstanceState for state being
+        # GC'ed, inlines _managed_removed_state
+        try:
+            obj = self._dict[state.key]
+        except KeyError:
+            # catch gc removed the key after we just checked for it
+            pass
+        else:
+            if attributes.instance_state(obj) is state:
+                self._dict.pop(state.key, None)
 
     def discard(self, state):
-        obj = self._dict.pop(state.key, None)
-        if obj is not None:
-            self._manage_removed_state(state)
-            st = attributes.instance_state(obj)
-            assert st is state
+        self.safe_discard(state)
 
     def safe_discard(self, state):
         if state.key in self._dict: