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

3 files changed, 94 insertions, 12 deletions

diff --git a/doc/build/changelog/unreleased_11/4068.rst b/doc/build/changelog/unreleased_11/4068.rst
new file mode 100644
index 000000000..9443b1ba0
--- /dev/null
+++ b/doc/build/changelog/unreleased_11/4068.rst
@@ -0,0 +1,10 @@
+.. change::
+    :tags: bug, orm
+    :tickets: 4068
+
+    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.
\ No newline at end of file
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:
diff --git a/test/orm/test_session.py b/test/orm/test_session.py
index 363d5f782..adfb21497 100644
--- a/test/orm/test_session.py
+++ b/test/orm/test_session.py
@@ -1,5 +1,5 @@
 from sqlalchemy.testing import eq_, assert_raises, \
-    assert_raises_message, assertions, is_true
+    assert_raises_message, assertions, is_true, is_
 from sqlalchemy.testing.util import gc_collect
 from sqlalchemy.testing import pickleable
 from sqlalchemy.util import pickle
@@ -1172,6 +1172,35 @@ class WeakIdentityMapTest(_fixtures.FixtureTest):
         s2.add(u1)
         assert u1 in s2
 
+    def test_fast_discard_race(self):
+        # test issue #4068
+        users, User = self.tables.users, self.classes.User
+
+        mapper(User, users)
+
+        sess = Session()
+
+        u1 = User(name='u1')
+        sess.add(u1)
+        sess.commit()
+
+        u1_state = u1._sa_instance_state
+        ref = u1_state.obj
+        u1_state.obj = lambda: None
+
+        u2 = sess.query(User).first()
+        u1_state._cleanup(ref)
+
+        u3 = sess.query(User).first()
+
+        is_(u2, u3)
+
+        u2_state = u2._sa_instance_state
+        ref = u2_state.obj
+        u2_state.obj = lambda: None
+        u2_state._cleanup(ref)
+        assert not sess.identity_map.contains_state(u2._sa_instance_state)
+
 
 class StrongIdentityMapTest(_fixtures.FixtureTest):
     run_inserts = None
@@ -1305,6 +1334,38 @@ class StrongIdentityMapTest(_fixtures.FixtureTest):
         eq_(prune(), 0)
         self.assert_(len(s.identity_map) == 0)
 
+    @testing.uses_deprecated()
+    def test_fast_discard_race(self):
+        # test issue #4068
+        users, User = self.tables.users, self.classes.User
+
+        mapper(User, users)
+
+        sess = Session(weak_identity_map=False)
+
+        u1 = User(name='u1')
+        sess.add(u1)
+        sess.commit()
+
+        u1_state = u1._sa_instance_state
+        sess.identity_map._dict.pop(u1_state.key)
+        ref = u1_state.obj
+        u1_state.obj = lambda: None
+
+        u2 = sess.query(User).first()
+        u1_state._cleanup(ref)
+
+        u3 = sess.query(User).first()
+
+        is_(u2, u3)
+
+        u2_state = u2._sa_instance_state
+        assert sess.identity_map.contains_state(u2._sa_instance_state)
+        ref = u2_state.obj
+        u2_state.obj = lambda: None
+        u2_state._cleanup(ref)
+        assert not sess.identity_map.contains_state(u2._sa_instance_state)
+
 
 class IsModifiedTest(_fixtures.FixtureTest):
     run_inserts = None