Add support for collectible types to SOS (#19842)

* Add support for collectible types to SOS

Collectible types indirectly reference managed LoaderAllocator via
pointer to native AssemblyLoaderAllocator stored in their MethodTable.
GC uses this relation when scanning object graph to determine which
objects are rooted and which ones are not.
The gcroot command in SOS doesn't understand this relation and so it
is unable to find all roots for LoaderAllocator.
This change fixes it.

* PR feedback

Make the failure to get the collectible info non-fatal to make it
compatible with older runtimes.

5 files changed, 153 insertions, 52 deletions

diff --git a/src/ToolBox/SOS/Strike/eeheap.cpp b/src/ToolBox/SOS/Strike/eeheap.cpp
index 5a5680fcfb..c28c9641f5 100644
--- a/src/ToolBox/SOS/Strike/eeheap.cpp
+++ b/src/ToolBox/SOS/Strike/eeheap.cpp
@@ -868,8 +868,7 @@ DWORD GetNumComponents(TADDR obj)
     return Value;
 }
 
-BOOL GetSizeEfficient(DWORD_PTR dwAddrCurrObj, 
-    DWORD_PTR dwAddrMethTable, BOOL bLarge, size_t& s, BOOL& bContainsPointers)
+static MethodTableInfo* GetMethodTableInfo(DWORD_PTR dwAddrMethTable)
 {
     // Remove lower bits in case we are in mark phase
     dwAddrMethTable = dwAddrMethTable & ~3;
@@ -880,12 +879,34 @@ BOOL GetSizeEfficient(DWORD_PTR dwAddrCurrObj,
         // from the target
         DacpMethodTableData dmtd;
         // see code:ClrDataAccess::RequestMethodTableData for details
-        if (dmtd.Request(g_sos,dwAddrMethTable) != S_OK)
-            return FALSE;
+        if (dmtd.Request(g_sos, dwAddrMethTable) != S_OK)
+            return NULL;
+
 
         info->BaseSize = dmtd.BaseSize;
         info->ComponentSize = dmtd.ComponentSize;
         info->bContainsPointers = dmtd.bContainsPointers;
+
+        // The following request doesn't work on older runtimes. For those, the
+        // objects would just look like non-collectible, which is acceptable.
+        DacpMethodTableCollectibleData dmtcd;
+        if (SUCCEEDED(dmtcd.Request(g_sos, dwAddrMethTable)))
+        {
+            info->bCollectible = dmtcd.bCollectible;
+            info->LoaderAllocatorObjectHandle = TO_TADDR(dmtcd.LoaderAllocatorObjectHandle);
+        }
+    }
+
+    return info;
+}
+
+BOOL GetSizeEfficient(DWORD_PTR dwAddrCurrObj, 
+    DWORD_PTR dwAddrMethTable, BOOL bLarge, size_t& s, BOOL& bContainsPointers)
+{
+    MethodTableInfo* info = GetMethodTableInfo(dwAddrMethTable);
+    if (info == NULL)
+    {
+        return FALSE;
     }
         
     bContainsPointers = info->bContainsPointers;
@@ -911,6 +932,20 @@ BOOL GetSizeEfficient(DWORD_PTR dwAddrCurrObj,
     return TRUE;
 }
 
+BOOL GetCollectibleDataEfficient(DWORD_PTR dwAddrMethTable, BOOL& bCollectible, TADDR& loaderAllocatorObjectHandle)
+{
+    MethodTableInfo* info = GetMethodTableInfo(dwAddrMethTable);
+    if (info == NULL)
+    {
+        return FALSE;
+    }
+
+    bCollectible = info->bCollectible;
+    loaderAllocatorObjectHandle = info->LoaderAllocatorObjectHandle;
+
+    return TRUE;
+}
+
 // This function expects stat to be valid, and ready to get statistics.
 void GatherOneHeapFinalization(DacpGcHeapDetails& heapDetails, HeapStat *stat, BOOL bAllReady, BOOL bShort)
 {
diff --git a/src/ToolBox/SOS/Strike/gcroot.cpp b/src/ToolBox/SOS/Strike/gcroot.cpp
index e4262208e4..cd13719094 100644
--- a/src/ToolBox/SOS/Strike/gcroot.cpp
+++ b/src/ToolBox/SOS/Strike/gcroot.cpp
@@ -1009,7 +1009,7 @@ GCRootImpl::RootNode *GCRootImpl::GetGCRefs(RootNode *path, RootNode *node)
 
     // Only calculate the size if we need it.
     size_t objSize = 0;
-    if (mSize || node->MTInfo->ContainsPointers)
+    if (mSize || node->MTInfo->ContainsPointers || node->MTInfo->Collectible)
     {
         objSize = GetSizeOfObject(obj, node->MTInfo);
         
@@ -1027,7 +1027,7 @@ GCRootImpl::RootNode *GCRootImpl::GetGCRefs(RootNode *path, RootNode *node)
     }
     
     // Early out:  If the object doesn't contain any pointers, return.
-    if (!node->MTInfo->ContainsPointers)
+    if (!node->MTInfo->ContainsPointers && !node->MTInfo->Collectible)
         return NULL;
     
     // Make sure we have the object's data in the cache.
@@ -1139,6 +1139,15 @@ GCRootImpl::MTInfo *GCRootImpl::GetMTInfo(TADDR mt)
     curr->ComponentSize = (size_t)dmtd.ComponentSize;
     curr->ContainsPointers = dmtd.bContainsPointers ? true : false;
 
+    // The following request doesn't work on older runtimes. For those, the
+    // objects would just look like non-collectible, which is acceptable.
+    DacpMethodTableCollectibleData dmtcd;
+    if (SUCCEEDED(dmtcd.Request(g_sos, mt)))
+    {
+        curr->Collectible = dmtcd.bCollectible ? true : false;
+        curr->LoaderAllocatorObjectHandle = TO_TADDR(dmtcd.LoaderAllocatorObjectHandle);
+    }
+
     // If this method table contains pointers, fill out and cache the GCDesc.
     if (curr->ContainsPointers)
     {
diff --git a/src/ToolBox/SOS/Strike/sos.cpp b/src/ToolBox/SOS/Strike/sos.cpp
index 64ee4b9a0c..5ff77a905a 100644
--- a/src/ToolBox/SOS/Strike/sos.cpp
+++ b/src/ToolBox/SOS/Strike/sos.cpp
@@ -180,6 +180,15 @@ namespace sos
             info->BaseSize = mMTData->BaseSize;
             info->ComponentSize = mMTData->ComponentSize;
             info->bContainsPointers = mMTData->bContainsPointers;
+
+            // The following request doesn't work on older runtimes. For those, the
+            // objects would just look like non-collectible, which is acceptable.
+            DacpMethodTableCollectibleData mtcd;
+            if (SUCCEEDED(mtcd.Request(g_sos, GetMT())))
+            {
+                info->bCollectible = mtcd.bCollectible;
+                info->LoaderAllocatorObjectHandle = TO_TADDR(mtcd.LoaderAllocatorObjectHandle);
+            }
         }
         
         if (mSize == (size_t)~0)
@@ -380,14 +389,14 @@ namespace sos
 
 
     RefIterator::RefIterator(TADDR obj, LinearReadCache *cache)
-        : mCache(cache), mGCDesc(0), mArrayOfVC(false), mDone(false), mBuffer(0), mCurrSeries(0),
+        : mCache(cache), mGCDesc(0), mArrayOfVC(false), mDone(false), mBuffer(0), mCurrSeries(0), mLoaderAllocatorObjectHandle(0),
           i(0), mCount(0), mCurr(0), mStop(0), mObject(obj), mObjSize(0)
     {
         Init();
     }
 
     RefIterator::RefIterator(TADDR obj, CGCDesc *desc, bool arrayOfVC, LinearReadCache *cache)
-        : mCache(cache), mGCDesc(desc), mArrayOfVC(arrayOfVC), mDone(false), mBuffer(0), mCurrSeries(0),
+        : mCache(cache), mGCDesc(desc), mArrayOfVC(arrayOfVC), mDone(false), mBuffer(0), mCurrSeries(0), mLoaderAllocatorObjectHandle(0),
           i(0), mCount(0), mCurr(0), mStop(0), mObject(obj), mObjSize(0)
     {
         Init();
@@ -403,6 +412,13 @@ namespace sos
     {
         if (mDone)
             Throw<Exception>("Attempt to move past the end of the iterator.");
+
+        if (mCurr == mLoaderAllocatorObjectHandle)
+        {
+            // The mLoaderAllocatorObjectHandle is always the last reference returned
+            mDone = true;
+            return *this;
+        }
         
         if (!mArrayOfVC)
         {
@@ -440,6 +456,14 @@ namespace sos
                 mDone = true;
         }
         
+        if (mDone && mLoaderAllocatorObjectHandle != NULL)
+        {
+            // The iteration over all regular object references is done, but there is one more
+            // reference for collectible types - the LoaderAllocator for GC
+            mCurr = mLoaderAllocatorObjectHandle;
+            mDone = false;
+        }
+
         return *this;
     }
     
@@ -457,66 +481,89 @@ namespace sos
     {
         TADDR mt = ReadPointer(mObject);
         BOOL bContainsPointers = FALSE;
-        
+        BOOL bCollectible = FALSE;
+        TADDR loaderAllocatorObjectHandle;
+
         if (!GetSizeEfficient(mObject, mt, FALSE, mObjSize, bContainsPointers))
             Throw<DataRead>("Failed to get size of object.");
-        
-        if (!bContainsPointers)
+
+        if (!GetCollectibleDataEfficient(mt, bCollectible, loaderAllocatorObjectHandle))
+            Throw<DataRead>("Failed to get collectible info of object.");
+
+        if (!bContainsPointers && !bCollectible)
         {
             mDone = true;
             return;
         }
-        
-        if (!mGCDesc)
+
+        if (bContainsPointers)
         {
-            int entries = 0;
-            
-            if (FAILED(MOVE(entries, mt-sizeof(TADDR))))
-                Throw<DataRead>("Failed to request number of entries.");
-            
-            // array of vc?
-            if (entries < 0)
+            if (!mGCDesc)
             {
-                entries = -entries;
-                mArrayOfVC = true;
+                int entries = 0;
+
+                if (FAILED(MOVE(entries, mt-sizeof(TADDR))))
+                    Throw<DataRead>("Failed to request number of entries.");
+
+                // array of vc?
+                if (entries < 0)
+                {
+                    entries = -entries;
+                    mArrayOfVC = true;
+                }
+                else
+                {
+                    mArrayOfVC = false;
+                }
+
+                size_t slots = 1 + entries * sizeof(CGCDescSeries)/sizeof(TADDR);
+
+                ArrayHolder<TADDR> buffer = new TADDR[slots];
+
+                ULONG fetched = 0;
+                CLRDATA_ADDRESS address = TO_CDADDR(mt - slots*sizeof(TADDR));
+                if (FAILED(g_ExtData->ReadVirtual(address, buffer, (ULONG)(slots*sizeof(TADDR)), &fetched)))
+                    Throw<DataRead>("Failed to request GCDesc.");
+
+                mBuffer = buffer.Detach();
+                mGCDesc = (CGCDesc*)(mBuffer + slots);
+            }
+
+            mCurrSeries = mGCDesc->GetHighestSeries();
+
+            if (!mArrayOfVC)
+            {
+                mCurr = mObject + mCurrSeries->GetSeriesOffset();
+                mStop = mCurr + mCurrSeries->GetSeriesSize() + mObjSize;
             }
             else
             {
-                mArrayOfVC = false;
+                i = 0;
+                mCurr = mObject + mCurrSeries->startoffset;
+                mStop = mCurr + mCurrSeries->val_serie[i].nptrs * sizeof(TADDR);
+                mCount = (int)mGCDesc->GetNumSeries();
             }
-            
-            size_t slots = 1 + entries * sizeof(CGCDescSeries)/sizeof(TADDR);
-            
-            ArrayHolder<TADDR> buffer = new TADDR[slots];
-            
-            ULONG fetched = 0;
-            CLRDATA_ADDRESS address = TO_CDADDR(mt - slots*sizeof(TADDR));
-            if (FAILED(g_ExtData->ReadVirtual(address, buffer, (ULONG)(slots*sizeof(TADDR)), &fetched)))
-                Throw<DataRead>("Failed to request GCDesc.");
-            
-            mBuffer = buffer.Detach();
-            mGCDesc = (CGCDesc*)(mBuffer + slots);
+
+            if (mCurr == mStop)
+                operator++();
+            else if (mCurr >= mObject + mObjSize - plug_skew)
+                mDone = true;
         }
-        
-        mCurrSeries = mGCDesc->GetHighestSeries();
-        
-        if (!mArrayOfVC)
+        else
         {
-            mCurr = mObject + mCurrSeries->GetSeriesOffset();
-            mStop = mCurr + mCurrSeries->GetSeriesSize() + mObjSize;
+            mDone = true;
         }
-        else
+
+        if (bCollectible)
         {
-            i = 0;
-            mCurr = mObject + mCurrSeries->startoffset;
-            mStop = mCurr + mCurrSeries->val_serie[i].nptrs * sizeof(TADDR);
-            mCount = (int)mGCDesc->GetNumSeries();
+            mLoaderAllocatorObjectHandle = loaderAllocatorObjectHandle;
+            if (mDone)
+            {
+                // There are no object references, but there is still a reference for 
+                // collectible types - the LoaderAllocator for GC
+                mCurr = mLoaderAllocatorObjectHandle;
+            }
         }
-        
-        if (mCurr == mStop)
-          operator++();
-        else if (mCurr >= mObject + mObjSize - plug_skew)
-          mDone = true;
     }
 
 
diff --git a/src/ToolBox/SOS/Strike/sos.h b/src/ToolBox/SOS/Strike/sos.h
index 3778235964..ff5b53a9c7 100644
--- a/src/ToolBox/SOS/Strike/sos.h
+++ b/src/ToolBox/SOS/Strike/sos.h
@@ -501,6 +501,8 @@ namespace sos
         TADDR *mBuffer;
         CGCDescSeries *mCurrSeries;
         
+        TADDR mLoaderAllocatorObjectHandle;
+
         int i, mCount;
         
         TADDR mCurr, mStop, mObject;
diff --git a/src/ToolBox/SOS/Strike/util.h b/src/ToolBox/SOS/Strike/util.h
index 78516546ac..ebad2e49ff 100644
--- a/src/ToolBox/SOS/Strike/util.h
+++ b/src/ToolBox/SOS/Strike/util.h
@@ -1660,9 +1660,11 @@ struct MethodTableInfo
     DWORD BaseSize;           // Caching BaseSize and ComponentSize for a MethodTable
     DWORD ComponentSize;      // here has HUGE perf benefits in heap traversals.
     BOOL  bContainsPointers;
+    BOOL  bCollectible;
     DWORD_PTR* GCInfoBuffer;  // Start of memory of GC info
     CGCDesc* GCInfo;    // Just past GC info (which is how it is stored)
     bool  ArrayOfVC;
+    TADDR LoaderAllocatorObjectHandle;
 };
 
 class MethodTableCache
@@ -1680,9 +1682,11 @@ protected:
             info.BaseSize = 0;
             info.ComponentSize = 0;
             info.bContainsPointers = false;
+            info.bCollectible = false;
             info.GCInfo = NULL;
             info.ArrayOfVC = false;
             info.GCInfoBuffer = NULL;
+            info.LoaderAllocatorObjectHandle = NULL;
         }
     };
     Node *head;
@@ -1948,6 +1952,8 @@ size_t NextOSPageAddress (size_t addr);
 BOOL GetSizeEfficient(DWORD_PTR dwAddrCurrObj, 
     DWORD_PTR dwAddrMethTable, BOOL bLarge, size_t& s, BOOL& bContainsPointers);
 
+BOOL GetCollectibleDataEfficient(DWORD_PTR dwAddrMethTable, BOOL& bCollectible, TADDR& loaderAllocatorObjectHandle);
+
 // ObjSize now uses the methodtable cache for its work too.
 size_t ObjectSize (DWORD_PTR obj, BOOL fIsLargeObject=FALSE);
 size_t ObjectSize(DWORD_PTR obj, DWORD_PTR mt, BOOL fIsValueClass, BOOL fIsLargeObject=FALSE);
@@ -2856,8 +2862,10 @@ private:
         TADDR *Buffer;
         CGCDesc *GCDesc;
 
+        TADDR LoaderAllocatorObjectHandle;
         bool ArrayOfVC;
         bool ContainsPointers;
+        bool Collectible;
         size_t BaseSize;
         size_t ComponentSize;
         
@@ -2874,7 +2882,7 @@ private:
 
         MTInfo()
             : MethodTable(0), TypeName(0), Buffer(0), GCDesc(0),
-              ArrayOfVC(false), ContainsPointers(false), BaseSize(0), ComponentSize(0)
+              ArrayOfVC(false), ContainsPointers(false), Collectible(false), BaseSize(0), ComponentSize(0)
         {
         }