[release/6.0] Notify Reference Tracker runtime of disconnect at the right time

src/coreclr/interop/comwrappers.cpp

@@ -791,17 +791,6 @@ void NativeObjectWrapperContext::Destroy(_In_ NativeObjectWrapperContext* wrappe
 {
     _ASSERTE(wrapper != nullptr);
 
-    // Check if the tracker object manager should be informed prior to being destroyed.
-    IReferenceTracker* trackerMaybe = wrapper->GetReferenceTracker();
-    if (trackerMaybe != nullptr)
-    {
-        // We only call this during a GC so ignore the failure as
-        // there is no way we can handle it at this point.
-        HRESULT hr = TrackerObjectManager::BeforeWrapperDestroyed(trackerMaybe);
-        _ASSERTE(SUCCEEDED(hr));
-        (void)hr;
-    }
-
     // Manually trigger the destructor since placement
     // new was used to allocate the object.
     wrapper->~NativeObjectWrapperContext();

src/coreclr/interop/comwrappers.hpp

@@ -212,8 +212,8 @@ class TrackerObjectManager
     // Called after wrapper has been created.
     static HRESULT AfterWrapperCreated(_In_ IReferenceTracker* obj);
 
-    // Called before wrapper is about to be destroyed (the same lifetime as short weak handle).
-    static HRESULT BeforeWrapperDestroyed(_In_ IReferenceTracker* obj);
+    // Called before wrapper is about to be finalized (the same lifetime as short weak handle).
+    static HRESULT BeforeWrapperFinalized(_In_ IReferenceTracker* obj);
 
 public:
     // Begin the reference tracking process for external objects.

src/coreclr/interop/inc/interoplib.h

@@ -87,8 +87,12 @@ namespace InteropLib
             _In_ size_t contextSize,
             _Out_ ExternalWrapperResult* result) noexcept;
 
-        // Destroy the supplied wrapper.
-        void DestroyWrapperForExternal(_In_ void* context) noexcept;
+        // Inform the wrapper it is being collected.
+        void NotifyWrapperForExternalIsBeingCollected(_In_ void* context) noexcept;
+
+         // Destroy the supplied wrapper.
+        // Optionally notify the wrapper of collection at the same time.
+        void DestroyWrapperForExternal(_In_ void* context, _In_ bool notifyIsBeingCollected = false) noexcept;
 
         // Separate the supplied wrapper from the tracker runtime.
         void SeparateWrapperFromTrackerRuntime(_In_ void* context) noexcept;

src/coreclr/interop/interoplib.cpp

@@ -166,15 +166,37 @@ namespace InteropLib
             return S_OK;
         }
 
-        void DestroyWrapperForExternal(_In_ void* contextMaybe) noexcept
+        void NotifyWrapperForExternalIsBeingCollected(_In_ void* contextMaybe) noexcept
+         {
+             NativeObjectWrapperContext* context = NativeObjectWrapperContext::MapFromRuntimeContext(contextMaybe);
+
+             // A caller should not be destroying a context without knowing if the context is valid.
+             _ASSERTE(context != nullptr);
+
+            // Check if the tracker object manager should be informed of collection.
+            IReferenceTracker* trackerMaybe = context->GetReferenceTracker();
+            if (trackerMaybe != nullptr)
+            {
+                // We only call this during a GC so ignore the failure as
+                // there is no way we can handle it at this point.
+                HRESULT hr = TrackerObjectManager::BeforeWrapperFinalized(trackerMaybe);
+                _ASSERTE(SUCCEEDED(hr));
+                (void)hr;
+            }
+        }
+
+        void DestroyWrapperForExternal(_In_ void* contextMaybe, _In_ bool notifyIsBeingCollected) noexcept
         {
             NativeObjectWrapperContext* context = NativeObjectWrapperContext::MapFromRuntimeContext(contextMaybe);
 
             // A caller should not be destroying a context without knowing if the context is valid.
             _ASSERTE(context != nullptr);
 
-            NativeObjectWrapperContext::Destroy(context);
-        }
+            if (notifyIsBeingCollected)
+                NotifyWrapperForExternalIsBeingCollected(contextMaybe);
+
+             NativeObjectWrapperContext::Destroy(context);
+         }
 
         void SeparateWrapperFromTrackerRuntime(_In_ void* contextMaybe) noexcept
         {

src/coreclr/interop/trackerobjectmanager.cpp

@@ -305,13 +305,13 @@ HRESULT TrackerObjectManager::AfterWrapperCreated(_In_ IReferenceTracker* obj)
     return S_OK;
 }
 
-HRESULT TrackerObjectManager::BeforeWrapperDestroyed(_In_ IReferenceTracker* obj)
+HRESULT TrackerObjectManager::BeforeWrapperFinalized(_In_ IReferenceTracker* obj)
 {
     _ASSERTE(obj != nullptr);
 
     HRESULT hr;
 
-    // Notify tracker runtime that we are about to destroy a wrapper
+    // Notify tracker runtime that we are about to finalize a wrapper
     // (same timing as short weak handle) for this object.
     // They need this information to disconnect weak refs and stop firing events,
     // so that they can avoid resurrecting the object.

src/coreclr/vm/interoplibinterface_comwrappers.cpp

@@ -176,7 +176,10 @@ namespace
             if (Result.Context != NULL)
             {
                 GCX_PREEMP();
-                InteropLib::Com::DestroyWrapperForExternal(Result.Context);
+                // We also request collection notification since this holder is presently
+                // only used for new activation of wrappers therefore the notification won't occur
+                // at the typical time of before finalization.
+                InteropLib::Com::DestroyWrapperForExternal(Result.Context, /* notifyIsBeingCollected */ true);
             }
         }
         InteropLib::Com::ExternalWrapperResult* operator&()
@@ -478,7 +481,11 @@ namespace
                 if (!cxt->IsSet(ExternalObjectContext::Flags_Detached)
                     && !GCHeapUtilities::GetGCHeap()->IsPromoted(OBJECTREFToObject(cxt->GetObjectRef())))
                 {
+                    // Indicate the wrapper entry has been detached.
                     cxt->MarkDetached();
+
+                    // Notify the wrapper it was not promoted and is being collected.
+                    InteropLib::Com::NotifyWrapperForExternalIsBeingCollected(cxt);
                 }
             }
         }

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

@@ -624,6 +624,9 @@ static int Main(string[] doNotUse)
                 ValidateQueryInterfaceAfterManagedObjectCollected();
                 ValidateAggregationWithComObject();
                 ValidateAggregationWithReferenceTrackerObject();
+
+                // Ensure all objects have been cleaned up.
+                ForceGC();
             }
             catch (Exception e)
             {

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

@@ -156,6 +156,12 @@ public static AllocationCountResult CountTrackerObjectAllocations()
 
         [DllImport(nameof(MockReferenceTrackerRuntime))]
         extern public static int TrackerTarget_ReleaseFromReferenceTracker(IntPtr ptr);
+
+        // Suppressing the GC transition here as we want to make sure we are in-sync
+        // with the GC which is setting the connected value.
+        [SuppressGCTransition]
+        [DllImport(nameof(MockReferenceTrackerRuntime))]
+        extern public static byte IsTrackerObjectConnected(IntPtr instance);
     }
 
     [Guid("42951130-245C-485E-B60B-4ED4254256F8")]
@@ -219,6 +225,12 @@ static IntPtr CreateInstance(IntPtr outer, out IntPtr inner)
             }
             else
             {
+                byte isConnected = MockReferenceTrackerRuntime.IsTrackerObjectConnected(this.classNative.Instance);
+                if (isConnected != 0)
+                {
+                    throw new Exception("TrackerObject should be disconnected prior to finalization");
+                }
+
                 ComWrappersHelper.Cleanup(ref this.classNative);
             }
         }

src/tests/Interop/COM/ComWrappers/MockReferenceTrackerRuntime/ReferenceTrackerRuntime.cpp

@@ -218,6 +218,11 @@ namespace
                 }
             }
 
+            bool IsConnected()
+            {
+                return _connected;
+            }
+
             STDMETHOD(AddObjectRef)(_In_ IUnknown* c, _Out_ int* id)
             {
                 assert(c != nullptr && id != nullptr);
@@ -546,6 +551,12 @@ extern "C" DLL_EXPORT int STDMETHODCALLTYPE Trigger_NotifyEndOfReferenceTracking
     return TrackerRuntimeManager.NotifyEndOfReferenceTrackingOnThread();
 }
 
+extern "C" DLL_EXPORT bool STDMETHODCALLTYPE IsTrackerObjectConnected(IUnknown* inst)
+{
+    auto trackerObject = reinterpret_cast<TrackerObject::TrackerObjectImpl*>(inst);
+    return trackerObject->IsConnected();
+}
+
 extern "C" DLL_EXPORT void* STDMETHODCALLTYPE TrackerTarget_AddRefFromReferenceTrackerAndReturn(IUnknown *obj)
 {
     assert(obj != nullptr);