From 6ef31918fce2a9790e01f863e06575a9ca2ec77a Mon Sep 17 00:00:00 2001
From: Jan Kotas <jkotas@microsoft.com>
Date: Sat, 10 Dec 2016 00:31:20 -0800
Subject: [PATCH] Improve ConditionalWeakTable.Remove (#8580)

Clear the key of the deleted entry to allow GC collect objects pointed to by it

Fixes #8577
---
 .../CompilerServices/ConditionalWeakTable.cs     | 16 +++++++++++++++-
 src/vm/comdependenthandle.cpp                    | 13 ++++++++++++-
 src/vm/comdependenthandle.h                      |  1 +
 src/vm/ecalllist.h                               |  1 +
 4 files changed, 29 insertions(+), 2 deletions(-)

diff --git a/src/mscorlib/src/System/Runtime/CompilerServices/ConditionalWeakTable.cs b/src/mscorlib/src/System/Runtime/CompilerServices/ConditionalWeakTable.cs
index 5e6a3f4bb937..9d9b61cc218b 100644
--- a/src/mscorlib/src/System/Runtime/CompilerServices/ConditionalWeakTable.cs
+++ b/src/mscorlib/src/System/Runtime/CompilerServices/ConditionalWeakTable.cs
@@ -516,10 +516,16 @@ internal bool Remove(TKey key)
                 int entryIndex = FindEntry(key, out value);
                 if (entryIndex != -1)
                 {
+                    ref Entry entry = ref _entries[entryIndex];
+
                     // We do not free the handle here, as we may be racing with readers who already saw the hash code.
                     // Instead, we simply overwrite the entry's hash code, so subsequent reads will ignore it.
                     // The handle will be free'd in Container's finalizer, after the table is resized or discarded.
-                    Volatile.Write(ref _entries[entryIndex].HashCode, -1);
+                    Volatile.Write(ref entry.HashCode, -1);
+
+                    // Also, clear the key to allow GC to collect objects pointed to by the entry
+                    entry.depHnd.SetPrimary(null);
+
                     return true;
                 }
 
@@ -840,6 +846,11 @@ public void GetPrimaryAndSecondary(out object primary, out object secondary)
             nGetPrimaryAndSecondary(_handle, out primary, out secondary);
         }
 
+        public void SetPrimary(object primary)
+        {
+            nSetPrimary(_handle, primary);
+        }
+
         public void SetSecondary(object secondary)
         {
             nSetSecondary(_handle, secondary);
@@ -870,6 +881,9 @@ public void Free()
         [MethodImpl(MethodImplOptions.InternalCall)]
         private static extern void nGetPrimaryAndSecondary(IntPtr dependentHandle, out object primary, out object secondary);
 
+        [MethodImpl(MethodImplOptions.InternalCall)]
+        private static extern void nSetPrimary(IntPtr dependentHandle, object primary);
+
         [MethodImpl(MethodImplOptions.InternalCall)]
         private static extern void nSetSecondary(IntPtr dependentHandle, object secondary);
 
diff --git a/src/vm/comdependenthandle.cpp b/src/vm/comdependenthandle.cpp
index 7a46f69e61ee..6535a804aee0 100644
--- a/src/vm/comdependenthandle.cpp
+++ b/src/vm/comdependenthandle.cpp
@@ -74,11 +74,22 @@ FCIMPL3(VOID, DependentHandle::nGetPrimaryAndSecondary, OBJECTHANDLE handle, Obj
 }
 FCIMPLEND
 
+FCIMPL2(VOID, DependentHandle::nSetPrimary, OBJECTHANDLE handle, Object *_primary)
+{
+    FCALL_CONTRACT;
+
+    _ASSERTE(handle != NULL);
+
+    OBJECTREF primary(_primary);
+    StoreObjectInHandle(handle, primary);
+}
+FCIMPLEND
+
 FCIMPL2(VOID, DependentHandle::nSetSecondary, OBJECTHANDLE handle, Object *_secondary)
 {
     FCALL_CONTRACT;
 
-    _ASSERTE(handle != NULL && _secondary != NULL);
+    _ASSERTE(handle != NULL);
 
     OBJECTREF secondary(_secondary);
     SetDependentHandleSecondary(handle, secondary);
diff --git a/src/vm/comdependenthandle.h b/src/vm/comdependenthandle.h
index 29110c108707..7192a4bbc3ab 100644
--- a/src/vm/comdependenthandle.h
+++ b/src/vm/comdependenthandle.h
@@ -45,6 +45,7 @@ class DependentHandle
     static FCDECL2(VOID,   nGetPrimary, OBJECTHANDLE handle, Object **outPrimary);
     static FCDECL3(VOID,   nGetPrimaryAndSecondary, OBJECTHANDLE handle, Object **outPrimary, Object **outSecondary);
     static FCDECL1(VOID,   nFree, OBJECTHANDLE handle);
+    static FCDECL2(VOID,   nSetPrimary, OBJECTHANDLE handle, Object *primary);
     static FCDECL2(VOID,   nSetSecondary, OBJECTHANDLE handle, Object *secondary);
 };
 
diff --git a/src/vm/ecalllist.h b/src/vm/ecalllist.h
index 03b675fa1793..b110d0eea650 100644
--- a/src/vm/ecalllist.h
+++ b/src/vm/ecalllist.h
@@ -108,6 +108,7 @@ FCFuncStart(gDependentHandleFuncs)
     FCFuncElement("nGetPrimary",             DependentHandle::nGetPrimary)
     FCFuncElement("nGetPrimaryAndSecondary", DependentHandle::nGetPrimaryAndSecondary)
     FCFuncElement("nFree",                   DependentHandle::nFree)
+    FCFuncElement("nSetPrimary",             DependentHandle::nSetPrimary)
     FCFuncElement("nSetSecondary",           DependentHandle::nSetSecondary)
 FCFuncEnd()