Collect the count and memory size of the .NET objects used by a process

Get-RuntimeDiagnostics.ps1

function Get-RuntimeDiagnostics {
<#
.SYNOPSIS
    Collect the count and memory size of the .net objects used by a process
.DESCRIPTION
    Collect the count and memory size of the .net objects used by a process

    This function is using the Microsoft.Diagnostics.Runtime .NET library from nuget:
    https://www.nuget.org/packages/Microsoft.Diagnostics.Runtime

    The number of CLR objects that are loaded by a .NET process can easily count
    up to hundreds of thousands or even more than a million.

    For that reason, in order to speed-up execution:
    - I'm using LINQ to group, sort and calculate the sum, instead of
      the native functions (Group-Object, Sort-Object and Measure-Object).
    - And I'm also using the .Where() and .ForEach() methods instead of
      sending to the pipeline through the Where-Object and ForEach-Object.

    Moreover, I'm executing the code of the main logic in a single streaming
    fashion (as-in I don't save any variables for the CLR Objects or the Groupings)
    to reduce the memory requirements of the function.

    Understandbly the code is harder to read as opposed to using the native
    PS functions, or using some variables to shorten the length of the .NET commands
    but it's needed in this case due to the increased data load.
.EXAMPLE
    Get-RuntimeDiagnostics -Verbose
.EXAMPLE
    cd (md C:\RuntimeDiagnostics -Force)
    nuget install Microsoft.Diagnostics.Runtime | Out-Null
    Add-Type -Path (dir '*\lib\netstandard2.0\*.dll').FullName

    $diag = Get-RuntimeDiagnostics -Verbose
    
    # download and load the required libraries
    # and then collect the runtime diagnostics
.EXAMPLE
    Add-Type -Path (dir 'C:\RuntimeDiagnostics\*\lib\netstandard2.0\*.dll').FullName

    $diag = Get-RuntimeDiagnostics
    # load the required libraries and then collect the diagnostics
.NOTES
    Inspiration was taken from Adam Driscoll's Get-ClrObject here:
    https://github.com/ironmansoftware/RuntimeDiagnostics/blob/main/Commands/GetObjectCommand.cs
    And from Microsoft's clrMD repo here:
    https://github.com/microsoft/clrmd/blob/main/doc/GettingStarted.md

    Make sure to use the CreateSnapshotAndAttach() method as per the Getting Started documentation
    and not the AttachToProcess()
#>
[CmdletBinding(DefaultParameterSetName = 'PID')]
[OutputType([PSCustomObject])]
param (
    [Parameter(ParameterSetName = 'Process')]
    [System.Diagnostics.Process]$Process = [System.Diagnostics.Process]::GetCurrentProcess(),

    [Parameter(ParameterSetName = 'PID')]
    [int]$ProcessID = $PID,

    [int]$NumberOfItems = 20
)

# make sure we have the required libraries
if (-not ('Microsoft.Diagnostics.Runtime.ClrObject' -as [type])) {
    Write-Warning 'Could not find the Microsoft.Diagnostics.Runtime namespace'
    Write-Warning 'Please load all the required assemblies for Microsoft.Diagnostics.Runtime'
    return
}

if ($PSCmdlet.ParameterSetName -eq 'PID') {
    $ThisProcess = [System.Diagnostics.Process]::GetProcesses().Where({$_.Id -eq $ProcessID})
    if (-not [bool]$ThisProcess) {
        Write-Warning "Cannot find process with ID $ProcessID"
        return
    }
}
elseif ($PSCmdlet.ParameterSetName -eq 'Process') {
    $ProcessID = $Process.Id
}

try {
    # save some stats to show memory usage and runtime duration of this function
    $Timer      = [System.Diagnostics.Stopwatch]::StartNew()
    $MemBefore  = [System.Diagnostics.Process]::GetCurrentProcess().PrivateMemorySize64

    # get the CLR Runtime
    $Target  = [Microsoft.Diagnostics.Runtime.DataTarget]::CreateSnapshotAndAttach($ProcessID)
    $Runtime = $Target.ClrVersions[0].CreateRuntime()
    # Note: do NOT use the [Microsoft.Diagnostics.Runtime.DataTarget]::AttachToProcess() method
    #       it breaks the memory size values, and the results are way off.

    # get the count and memory size from the objects
    $Results = [System.Linq.Enumerable]::OrderByDescending(
        [PSObject[]](
            [System.Linq.Enumerable]::GroupBy(
                [Microsoft.Diagnostics.Runtime.ClrObject[]]$Runtime.Heap.EnumerateObjects(),
                [Func[Microsoft.Diagnostics.Runtime.ClrObject,System.String]]{$args[0].Type.Name}
            ).Where({$_.Count -gt 1}).ForEach({
                [PSCustomObject]@{
                    Type   = $_.Key
                    Count  = $_.Count
                    Size   = [System.Linq.Enumerable]::Sum([System.UInt64[]]$_.Size)
                }
            }) # foreach group
        ),
        [Func[PSObject,System.UInt64]]{$args[0].Size}
    ).Where({$_.Type -ne 'Free'},'First',$NumberOfItems)

    # add some properties to the output object
    $Def  = [string[]]('Memory','Count','Type')
    $DDPS = [Management.Automation.PSPropertySet]::new('DefaultDisplayPropertySet',$Def)
    $Std  = [Management.Automation.PSMemberInfo[]]@($DDPS)
    $Labels = ('b', 'KB',  'MB',  'GB',  'TB',  'PB',  'EB',  'ZB',  'YB')
    $Results.ForEach({
        $Order  = [Math]::Floor([Math]::Log($_.Size, 1000))
        $Number = $_.Size / [Math]::Pow(1000, $Order)
        if     ($Number -lt 2)  {$RoundTo = 2}
        elseif ($Number -lt 20) {$RoundTo = 1}
        else                    {$RoundTo = 0}
        $Rounded  = [Math]::Round($Number,$RoundTo)
        $_.psobject.Properties.Add(
            [PSNoteProperty]::new('Memory',('{0}{1}' -f $Rounded,$Labels[$Order]))
        )
        $_ | Add-Member MemberSet PSStandardMembers $Std
    })
}

catch {throw $_}

finally {
    # get memory usage before garbage collection (that's the peak)
    $MemMid = [System.Diagnostics.Process]::GetCurrentProcess().PrivateMemorySize64

    # clean up
    if ($Runtime) {$Runtime.Dispose()}
    if ($Target)  {$Target.Dispose()}
    Remove-Variable Target,Runtime -Verbose:$false -ErrorAction Ignore
    [System.GC]::Collect()
    [System.GC]::WaitForPendingFinalizers()
    [System.GC]::Collect()
}

# show some verbose info
$ThisPID  = [System.Diagnostics.Process]::GetCurrentProcess().Id
$MemAfter = [System.Diagnostics.Process]::GetCurrentProcess().PrivateMemorySize64
$MemDiff  = $MemAfter - $MemBefore
$MemUse   = '{0}MB' -f [Math]::Floor($MemDiff/1000000)
$MemPeak  = '{0}MB' -f [Math]::Floor($MemMid/1000000)
$MemStart = '{0}MB' -f [Math]::Floor($MemBefore/1000000)
$MemEnd   = '{0}MB' -f [Math]::Floor($MemAfter/1000000)
$Timer.Stop()
$Duration = $Timer.Elapsed.TotalSeconds.ToString('#"sec" .###"ms"').Replace('.','')
Write-Verbose "Execution runtime was $Duration"
Write-Verbose "During execution, the memory of this process (PID:$ThisPID) fluctuated like so"
if ($MemDiff -gt 0) {$Suffix = ", thus used $MemUse"}
else                {$Suffix = ''}
Write-Verbose ("It started at $MemStart, peaked at $MemPeak, and ended at $MemEnd" + $Suffix)

# show the output
if ($Results.Count -ge 1) {Write-Output $Results}
}


## for clarity, here's another take but this time with native PS functions
## the following way takes up approx. ~20 seconds to finish, and uses ~300MB of memory
<#
# helper function
function ConvertTo-PrettyCapacity {
    [OutputType([string])]
    param ([Parameter(Mandatory,ValueFromPipeline)][UInt64]$Bytes)
    $Labels = ('b', 'KB',  'MB',  'GB',  'TB',  'PB',  'EB',  'ZB',  'YB')
    $Order  = [Math]::Floor([Math]::Log($Bytes, 1000))   # <-- orders of magnitude
    $Number = $Bytes / [Math]::Pow(1000, $Order)
    if     ($Number -lt 2)  {$RoundTo = 2}
    elseif ($Number -lt 20) {$RoundTo = 1}
    else                    {$RoundTo = 0}
    '{0}{1}' -f [Math]::Round($Number,$RoundTo),$Labels[$Order]
}

Add-Type -Path (dir 'C:\RuntimeDiagnostics\*\lib\netstandard2.0\*.dll').FullName

$Target  = [Microsoft.Diagnostics.Runtime.DataTarget]::CreateSnapshotAndAttach($PID)
$Runtime = $Target.ClrVersions[0].CreateRuntime()
$Results = $Runtime.Heap.EnumerateObjects() |
    Group-Object -Property {$_.Type.Name} | foreach {
        [PSCustomObject]@{
            Type   = $_.Name
            Count  = $_.Count
            Size   = ($_.Group | Measure-Object -Property Size -Sum).Sum
        }
    } |
    Sort-Object Size -Descending | 
    Select-Object -First 20 -Property @{n='Memory';e={ConvertTo-PrettyCapacity $_.Size}},Count,Type
#>