Cleanup internal ComWrappers cache when object enters Finalization queue

src/coreclr/interop/comwrappers.cpp

@@ -357,7 +357,8 @@ ManagedObjectWrapper* ManagedObjectWrapper::MapFromIUnknown(_In_ IUnknown* pUnk)
     // If the first Vtable entry is part of the ManagedObjectWrapper IUnknown impl,
     // we know how to interpret the IUnknown.
     void** vtable = *reinterpret_cast<void***>(pUnk);
-    if (*vtable != ManagedObjectWrapper_IUnknownImpl.QueryInterface)
+    if (*vtable != ManagedObjectWrapper_IUnknownImpl.QueryInterface
+        && *vtable != ManagedObjectWrapper_IReferenceTrackerTargetImpl.QueryInterface)
         return nullptr;
 
     ABI::ComInterfaceDispatch* disp = reinterpret_cast<ABI::ComInterfaceDispatch*>(pUnk);
@@ -841,7 +842,7 @@ void* NativeObjectWrapperContext::GetRuntimeContext() const noexcept
 
 IReferenceTracker* NativeObjectWrapperContext::GetReferenceTracker() const noexcept
 {
-    return ((_trackerObjectState == TrackerObjectState::NotSet) ? nullptr : _trackerObject);
+    return ((_trackerObjectState == TrackerObjectState::NotSet || _trackerObjectDisconnected) ? nullptr : _trackerObject);
 }
 
 // See TrackerObjectManager::AfterWrapperCreated() for AddRefFromTrackerSource() usage.

src/coreclr/vm/gcenv.ee.cpp

@@ -88,8 +88,7 @@ VOID GCToEEInterface::AfterGcScanRoots (int condemned, int max_gen,
     ::GetAppDomain()->DetachRCWs();
 #endif // FEATURE_COMINTEROP
 
-    if (!sc->concurrent)
-        Interop::OnAfterGCScanRoots();
+    Interop::OnAfterGCScanRoots(sc->concurrent);
 }
 
 /*

src/coreclr/vm/interoplibinterface.cpp

@@ -30,9 +30,15 @@ namespace
         enum
         {
             Flags_None = 0,
+
+            // The EOC has been collected and is no longer visible from managed code.
             Flags_Collected = 1,
+
             Flags_ReferenceTracker = 2,
             Flags_InCache = 4,
+
+            // The EOC is "detached" and no longer used to map between identity and a managed object.
+            Flags_Detached = 8,
         };
         DWORD Flags;
 
@@ -69,7 +75,7 @@ namespace
 
         bool IsActive() const
         {
-            return !IsSet(Flags_Collected)
+            return !IsSet(Flags_Collected | Flags_Detached)
                 && (SyncBlockIndex != InvalidSyncBlockIndex);
         }
 
@@ -80,6 +86,17 @@ namespace
             Flags |= Flags_Collected;
         }
 
+        void MarkDetached()
+        {
+            _ASSERTE(GCHeapUtilities::IsGCInProgress());
+            Flags |= Flags_Detached;
+        }
+
+        void MarkNotInCache()
+        {
+            ::InterlockedAnd((LONG*)&Flags, (~Flags_InCache));
+        }
+
         OBJECTREF GetObjectRef()
         {
             CONTRACTL
@@ -428,6 +445,32 @@ namespace
 
             _hashMap.Remove(cxt->GetKey());
         }
+
+        void DetachNotPromotedEOCs()
+        {
+            CONTRACTL
+            {
+                NOTHROW;
+                GC_NOTRIGGER;
+                MODE_ANY;
+                PRECONDITION(GCHeapUtilities::IsGCInProgress()); // GC is in progress and the runtime is suspended
+            }
+            CONTRACTL_END;
+
+            Iterator curr = _hashMap.Begin();
+            Iterator end = _hashMap.End();
+
+            ExternalObjectContext* cxt;
+            for (; curr != end; ++curr)
+            {
+                cxt = *curr;
+                if (cxt->IsActive()
+                    && !GCHeapUtilities::GetGCHeap()->IsPromoted(OBJECTREFToObject(cxt->GetObjectRef())))
+                {
+                    cxt->MarkDetached();
+                }
+            }
+        }
     };
 
     // Global instance of the external object cache
@@ -727,6 +770,15 @@ namespace
                     handle = handleLocal;
                 }
             }
+            else if (extObjCxt != NULL && extObjCxt->IsSet(ExternalObjectContext::Flags_Detached))
+            {
+                // If an EOC has been found but is marked detached, then we will remove it from the
+                // cache here instead of letting the GC do it later and pretend like it wasn't found.
+                STRESS_LOG1(LF_INTEROP, LL_INFO10, "Detached EOC requested: 0x%p\n", extObjCxt);
+                cache->Remove(extObjCxt);
+                extObjCxt->MarkNotInCache();
+                extObjCxt = NULL;
+            }
         }
 
         STRESS_LOG2(LF_INTEROP, LL_INFO1000, "EOC: 0x%p or Handle: 0x%p\n", extObjCxt, handle);
@@ -1796,4 +1848,19 @@ void ComWrappersNative::OnFullGCFinished()
     }
 }
 
+void ComWrappersNative::AfterRefCountedHandleCallbacks()
+{
+    CONTRACTL
+    {
+        NOTHROW;
+        GC_NOTRIGGER;
+        MODE_ANY;
+    }
+    CONTRACTL_END;
+
+    ExtObjCxtCache* cache = ExtObjCxtCache::GetInstanceNoThrow();
+    if (cache != NULL)
+        cache->DetachNotPromotedEOCs();
+}
+
 #endif // FEATURE_COMWRAPPERS

src/coreclr/vm/interoplibinterface.h

@@ -63,6 +63,7 @@ class ComWrappersNative
 public: // GC interaction
     static void OnFullGCStarted();
     static void OnFullGCFinished();
+    static void AfterRefCountedHandleCallbacks();
 };
 
 class GlobalComWrappersForMarshalling
@@ -187,7 +188,7 @@ class Interop
     // Notify before/after when GC is scanning roots.
     // Present assumption is that calls will never be nested.
     static void OnBeforeGCScanRoots();
-    static void OnAfterGCScanRoots();
+    static void OnAfterGCScanRoots(_In_ bool is_concurrent);
 };
 
 #endif // _INTEROPLIBINTERFACE_H_

src/coreclr/vm/interoplibinterface_shared.cpp

@@ -130,7 +130,7 @@ void Interop::OnBeforeGCScanRoots()
 #endif // FEATURE_OBJCMARSHAL
 }
 
-void Interop::OnAfterGCScanRoots()
+void Interop::OnAfterGCScanRoots(_In_ bool is_concurrent)
 {
     CONTRACTL
     {
@@ -139,7 +139,12 @@ void Interop::OnAfterGCScanRoots()
     }
     CONTRACTL_END;
 
+#ifdef FEATURE_COMWRAPPERS
+    ComWrappersNative::AfterRefCountedHandleCallbacks();
+#endif // FEATURE_COMWRAPPERS
+
 #ifdef FEATURE_OBJCMARSHAL
-    ObjCMarshalNative::AfterRefCountedHandleCallbacks();
+    if (!is_concurrent)
+        ObjCMarshalNative::AfterRefCountedHandleCallbacks();
 #endif // FEATURE_OBJCMARSHAL
 }

src/tests/Interop/COM/ComWrappers/API/Program.cs

@@ -280,6 +280,56 @@ static void ValidatePrecreatedExternalWrapper()
                 });
         }
 
+        static void ValidateExternalWrapperCacheCleanUp()
+        {
+            Console.WriteLine($"Running {nameof(ValidateExternalWrapperCacheCleanUp)}...");
+
+            var cw = new TestComWrappers();
+
+            // Get an object from a tracker runtime.
+            IntPtr trackerObjRaw = MockReferenceTrackerRuntime.CreateTrackerObject();
+
+            // Create a wrapper for the object instance.
+            var weakRef1 = CreateAndRegisterWrapper(cw, trackerObjRaw);
+
+            // Run the GC to have the wrapper marked for collection.
+            ForceGC();
+
+            // Create a new wrapper for the same external object.
+            var weakRef2 = CreateAndRegisterWrapper(cw, trackerObjRaw);
+
+            // We are using a tracking resurrection WeakReference<T> so we should be able
+            // to get back the objects as they are all continually re-registering for Finalization.
+            Assert.IsTrue(weakRef1.TryGetTarget(out ITrackerObjectWrapper wrapper1));
+            Assert.IsTrue(weakRef2.TryGetTarget(out ITrackerObjectWrapper wrapper2));
+
+            // Check that the two wrappers aren't equal, meaning we created a new wrapper since
+            // the first wrapper was removed from the internal cache.
+            Assert.AreNotEqual(wrapper1, wrapper2);
+
+            // Let the wrappers Finalize.
+            wrapper1.ReregisterForFinalize = false;
+            wrapper2.ReregisterForFinalize = false;
+
+            static WeakReference<ITrackerObjectWrapper> CreateAndRegisterWrapper(ComWrappers cw, IntPtr trackerObjRaw)
+            {
+                // Manually create a wrapper
+                var iid = typeof(ITrackerObject).GUID;
+                IntPtr iTestComObject;
+                int hr = Marshal.QueryInterface(trackerObjRaw, ref iid, out iTestComObject);
+                Assert.AreEqual(0, hr);
+                var nativeWrapper = new ITrackerObjectWrapper(iTestComObject);
+
+                nativeWrapper = (ITrackerObjectWrapper)cw.GetOrRegisterObjectForComInstance(trackerObjRaw, CreateObjectFlags.None, nativeWrapper);
+
+                // Set this on the return instead of during creation since the returned wrapper may be the one from
+                // the internal cache and not the one passed in above.
+                nativeWrapper.ReregisterForFinalize = true;
+
+                return new WeakReference<ITrackerObjectWrapper>(nativeWrapper, trackResurrection: true);
+            }
+        }
+
         static void ValidateSuppliedInnerNotAggregation()
         {
             Console.WriteLine($"Running {nameof(ValidateSuppliedInnerNotAggregation)}...");
@@ -566,6 +616,7 @@ static int Main(string[] doNotUse)
                 ValidateCreateObjectCachingScenario();
                 ValidateWrappersInstanceIsolation();
                 ValidatePrecreatedExternalWrapper();
+                ValidateExternalWrapperCacheCleanUp();
                 ValidateSuppliedInnerNotAggregation();
                 ValidateIUnknownImpls();
                 ValidateBadComWrapperImpl();

src/tests/Interop/COM/ComWrappers/Common.cs

@@ -213,9 +213,18 @@ static IntPtr CreateInstance(IntPtr outer, out IntPtr inner)
 
         ~ITrackerObjectWrapper()
         {
-            ComWrappersHelper.Cleanup(ref this.classNative);
+            if (this.ReregisterForFinalize)
+            {
+                GC.ReRegisterForFinalize(this);
+            }
+            else
+            {
+                ComWrappersHelper.Cleanup(ref this.classNative);
+            }
         }
 
+        public bool ReregisterForFinalize { get; set; } = false;
+
         public int AddObjectRef(IntPtr obj)
         {
             int id;