Improve ConditionalWeakTable.Remove

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);
 

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);

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);
 };
 

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()