Private/Collector.ps1

#Requires -Version 5.1

function Get-SMBeatCimInstanceSafe {
    param(
        [Parameter(Mandatory = $true)]
        [string]$ClassName,
        $CimSession
    )

    $params = @{
        ClassName   = $ClassName
        ErrorAction = 'Stop'
    }
    if ($CimSession) {
        $params['CimSession'] = $CimSession
    }

    try {
        Get-CimInstance @params
    }
    catch {
        Write-Verbose ("CIM class {0} not available: {1}" -f $ClassName, $_.Exception.Message)
        @()
    }
}

function Get-SMBeatRawLong {
    param(
        $Object,
        [string[]]$PropertyNames
    )

    foreach ($name in $PropertyNames) {
        $prop = $Object.PSObject.Properties[$name]
        if ($prop -and $null -ne $prop.Value) {
            return [uint64]$prop.Value
        }
    }

    return [uint64]0
}

function New-SMBeatCounterSnapshot {
    param(
        [Parameter(Mandatory = $true)]
        [string]$Kind,
        [Parameter(Mandatory = $true)]
        [string]$Name,
        [uint64]$ReadRaw,
        [uint64]$WriteRaw,
        [uint64]$SentRaw,
        [uint64]$ReceivedRaw,
        [string]$Client,
        [string]$User
    )

    [PSCustomObject]@{
        Kind        = $Kind
        Name        = $Name
        ReadRaw     = $ReadRaw
        WriteRaw    = $WriteRaw
        SentRaw     = $SentRaw
        ReceivedRaw = $ReceivedRaw
        Client      = $Client
        User        = $User
    }
}

function Get-SMBeatSessionMap {
    param(
        $CimSession
    )

    $map = @{}
    $params = @{ ErrorAction = 'SilentlyContinue' }
    if ($CimSession) {
        $params['CimSession'] = $CimSession
    }

    try {
        $sessions = @(Get-SmbSession @params)
    }
    catch {
        return $map
    }

    foreach ($session in $sessions) {
        $id = [string]$session.SessionId
        $map[$id] = [PSCustomObject]@{
            Client = [string]$session.ClientComputerName
            User   = [string]$session.ClientUserName
        }
    }

    return $map
}

function Get-SMBeatCounterSnapshots {
    param(
        $CimSession,
        [switch]$IncludeSessions,
        [switch]$IncludeAdminShares,
        [switch]$IncludeNics
    )

    $items = New-Object System.Collections.Generic.List[object]
    $computer = $env:COMPUTERNAME
    if ($CimSession -and $CimSession.ComputerName) {
        $computer = $CimSession.ComputerName
    }

    $serverRows = @(Get-SMBeatCimInstanceSafe -ClassName 'Win32_PerfRawData_Counters_SMBServer' -CimSession $CimSession)
    foreach ($row in $serverRows) {
        $items.Add((New-SMBeatCounterSnapshot -Kind 'server' -Name $computer `
                -ReadRaw (Get-SMBeatRawLong -Object $row -PropertyNames @('ReadBytesPersec')) `
                -WriteRaw (Get-SMBeatRawLong -Object $row -PropertyNames @('WriteBytesPersec')) `
                -SentRaw (Get-SMBeatRawLong -Object $row -PropertyNames @('SentBytesPersec')) `
                -ReceivedRaw (Get-SMBeatRawLong -Object $row -PropertyNames @('ReceivedBytesPersec')))) | Out-Null
    }

    $shareRows = @(Get-SMBeatCimInstanceSafe -ClassName 'Win32_PerfRawData_Counters_SMBServerShares' -CimSession $CimSession)
    foreach ($row in $shareRows) {
        $name = [string]$row.Name
        if (Test-SMBeatSkipInstance -Name $name -Kind 'share' -IncludeAdminShares:$IncludeAdminShares) {
            continue
        }

        $items.Add((New-SMBeatCounterSnapshot -Kind 'share' -Name $name `
                -ReadRaw (Get-SMBeatRawLong -Object $row -PropertyNames @('ReadBytesPersec')) `
                -WriteRaw (Get-SMBeatRawLong -Object $row -PropertyNames @('WriteBytesPersec')) `
                -SentRaw (Get-SMBeatRawLong -Object $row -PropertyNames @('SentBytesPersec')) `
                -ReceivedRaw (Get-SMBeatRawLong -Object $row -PropertyNames @('ReceivedBytesPersec')))) | Out-Null
    }

    if ($IncludeSessions) {
        $sessionMap = Get-SMBeatSessionMap -CimSession $CimSession
        $sessionRows = @(Get-SMBeatCimInstanceSafe -ClassName 'Win32_PerfRawData_Counters_SMBServerSessions' -CimSession $CimSession)
        foreach ($row in $sessionRows) {
            $name = [string]$row.Name
            if (Test-SMBeatSkipInstance -Name $name -Kind 'session') {
                continue
            }

            $client = $null
            $user = $null
            if ($sessionMap.ContainsKey($name)) {
                $client = $sessionMap[$name].Client
                $user = $sessionMap[$name].User
            }

            $items.Add((New-SMBeatCounterSnapshot -Kind 'session' -Name $name `
                    -ReadRaw (Get-SMBeatRawLong -Object $row -PropertyNames @('ReadBytesPersec')) `
                    -WriteRaw (Get-SMBeatRawLong -Object $row -PropertyNames @('WriteBytesPersec')) `
                    -SentRaw (Get-SMBeatRawLong -Object $row -PropertyNames @('SentBytesPersec')) `
                    -ReceivedRaw (Get-SMBeatRawLong -Object $row -PropertyNames @('ReceivedBytesPersec')) `
                    -Client $client -User $user)) | Out-Null
        }
    }

    if ($IncludeNics) {
        $nicRows = @(Get-SMBeatCimInstanceSafe -ClassName 'Win32_PerfRawData_Tcpip_NetworkInterface' -CimSession $CimSession)
        foreach ($row in $nicRows) {
            $name = [string]$row.Name
            if (Test-SMBeatSkipInstance -Name $name -Kind 'nic') {
                continue
            }
            if ($name -like '*Loopback*' -or $name -like '*isatap*' -or $name -like '*Teredo*') {
                continue
            }

            $items.Add((New-SMBeatCounterSnapshot -Kind 'nic' -Name $name `
                    -ReadRaw 0 -WriteRaw 0 `
                    -SentRaw (Get-SMBeatRawLong -Object $row -PropertyNames @('BytesSentPersec')) `
                    -ReceivedRaw (Get-SMBeatRawLong -Object $row -PropertyNames @('BytesReceivedPersec')))) | Out-Null
        }
    }

    return $items
}

function Read-SMBeatState {
    param(
        [Parameter(Mandatory = $true)]
        [string]$DataRoot
    )

    $path = Get-SMBeatStatePath -DataRoot $DataRoot
    $json = Read-SMBeatTextFile -Path $path
    if ($null -eq $json) {
        return @{}
    }

    $parsed = ConvertFrom-SMBeatJson -Json $json
    $hash = @{}
    $counters = $parsed.Counters
    if ($counters) {
        foreach ($prop in $counters.PSObject.Properties) {
            $hash[$prop.Name] = ConvertTo-SMBeatHashtable -Object $prop.Value
        }
    }
    return $hash
}

function Save-SMBeatState {
    param(
        [Parameter(Mandatory = $true)]
        [string]$DataRoot,
        [Parameter(Mandatory = $true)]
        [hashtable]$Counters,
        [datetime]$Utc
    )

    $obj = [PSCustomObject]@{
        LastSampleUtc = (ConvertTo-SMBeatUtcString -DateTime $Utc)
        Counters      = $Counters
    }
    Write-SMBeatTextFile -Path (Get-SMBeatStatePath -DataRoot $DataRoot) -Text (ConvertTo-SMBeatJson -InputObject $obj)
}

function Save-SMBeatHeartbeat {
    param(
        [Parameter(Mandatory = $true)]
        [string]$DataRoot,
        [datetime]$Utc,
        [int]$IntervalSec,
        [string]$ComputerName,
        [string]$LastError,
        [switch]$Success,
        [int]$EmittedCount
    )

    $now = ConvertTo-SMBeatUtcString -DateTime $Utc
    $existingJson = Read-SMBeatTextFile -Path (Get-SMBeatHeartbeatPath -DataRoot $DataRoot)
    $successCount = 0
    $lastSuccess = $null
    if ($existingJson) {
        $existing = ConvertFrom-SMBeatJson -Json $existingJson
        if ($existing.SuccessCount) {
            $successCount = [int]$existing.SuccessCount
        }
        if ($existing.LastSuccessUtc) {
            $lastSuccess = $existing.LastSuccessUtc
        }
    }

    if ($Success) {
        $successCount++
        $lastSuccess = $now
    }

    $obj = [PSCustomObject]@{
        LastAttemptUtc = $now
        LastSuccessUtc = $lastSuccess
        LastError      = $LastError
        IntervalSec    = $IntervalSec
        ComputerName   = $ComputerName
        SuccessCount   = $successCount
        EmittedCount   = $EmittedCount
        Success        = [bool]$Success
    }

    Write-SMBeatTextFile -Path (Get-SMBeatHeartbeatPath -DataRoot $DataRoot) -Text (ConvertTo-SMBeatJson -InputObject $obj)
}

function Add-SMBeatSampleLines {
    param(
        [Parameter(Mandatory = $true)]
        [string]$DataRoot,
        [Parameter(Mandatory = $true)]
        [datetime]$Utc,
        [Parameter(Mandatory = $true)]
        [object[]]$Records
    )

    if ($null -eq $Records -or $Records.Count -eq 0) {
        return
    }

    $path = Get-SMBeatSampleFilePath -DataRoot $DataRoot -Utc $Utc
    $sb = New-Object System.Text.StringBuilder
    foreach ($record in $Records) {
        $line = (ConvertTo-Json -InputObject $record -Depth 6 -Compress)
        [void]$sb.AppendLine($line)
    }

    Add-SMBeatTextFile -Path $path -Text $sb.ToString()
}

function Invoke-SMBeatRetention {
    param(
        [Parameter(Mandatory = $true)]
        [string]$DataRoot,
        [int]$RetentionDays = 90
    )

    if ($RetentionDays -le 0) {
        return
    }

    $dir = Get-SMBeatSampleDirectory -DataRoot $DataRoot
    if (-not (Test-Path -LiteralPath $dir)) {
        return
    }

    $cutoff = (Get-SMBeatUtcNow).Date.AddDays(-1 * $RetentionDays)
    Get-ChildItem -LiteralPath $dir -Filter '*.jsonl' -ErrorAction SilentlyContinue | ForEach-Object {
        $dayPart = [System.IO.Path]::GetFileNameWithoutExtension($_.Name)
        $parsed = [datetime]::MinValue
        if ([datetime]::TryParseExact($dayPart, 'yyyy-MM-dd', $script:SMBeatInvariant, [System.Globalization.DateTimeStyles]::None, [ref]$parsed)) {
            if ($parsed -lt $cutoff) {
                Remove-Item -LiteralPath $_.FullName -Force -ErrorAction SilentlyContinue
            }
        }
    }
}

function ConvertTo-SMBeatSampleRecord {
    param(
        [Parameter(Mandatory = $true)]
        $Snapshot,
        [Parameter(Mandatory = $true)]
        [string]$Server,
        [Parameter(Mandatory = $true)]
        [datetime]$Utc,
        [Parameter(Mandatory = $true)]
        [int]$IntervalSec,
        [int64]$ReadDelta,
        [int64]$WriteDelta,
        [int64]$SentDelta,
        [int64]$ReceivedDelta
    )

    [PSCustomObject]@{
        ts                = (ConvertTo-SMBeatUtcString -DateTime $Utc)
        server            = $Server
        kind              = $Snapshot.Kind
        name              = $Snapshot.Name
        client            = $Snapshot.Client
        user              = $Snapshot.User
        readBytesDelta    = $ReadDelta
        writeBytesDelta   = $WriteDelta
        sentBytesDelta    = $SentDelta
        receivedBytesDelta = $ReceivedDelta
        sampleIntervalSec = $IntervalSec
    }
}

function Invoke-SMBeatSampleInternal {
    [CmdletBinding()]
    param(
        [string]$Path,
        [string]$ComputerName,
        $CimSession,
        [int]$IntervalSec = 60,
        [int]$RetentionDays = 90,
        [switch]$IncludeSessions,
        [switch]$IncludeAdminShares,
        [switch]$IncludeNics,
        [switch]$AdHocLayout
    )

    $utc = Get-SMBeatUtcNow
    $localComputer = $env:COMPUTERNAME
    $label = $localComputer
    if ($CimSession -and $CimSession.ComputerName) {
        $label = $CimSession.ComputerName
    }
    elseif (-not [string]::IsNullOrWhiteSpace($ComputerName)) {
        $label = $ComputerName
    }

    $dataRoot = Get-SMBeatDataRoot -Path $Path
    if ($AdHocLayout) {
        $dataRoot = Get-SMBeatDataRoot -Path $Path -ComputerName $label
    }

    try {
        $snapshots = @(Get-SMBeatCounterSnapshots -CimSession $CimSession -IncludeSessions:$IncludeSessions -IncludeAdminShares:$IncludeAdminShares -IncludeNics:$IncludeNics)
        $state = Read-SMBeatState -DataRoot $dataRoot
        $newState = @{}
        $emitted = New-Object System.Collections.Generic.List[object]

        foreach ($snap in $snapshots) {
            $key = '{0}:{1}' -f $snap.Kind, $snap.Name
            $newState[$key] = @{
                ReadRaw     = $snap.ReadRaw
                WriteRaw    = $snap.WriteRaw
                SentRaw     = $snap.SentRaw
                ReceivedRaw = $snap.ReceivedRaw
            }

            if (-not $state.ContainsKey($key)) {
                continue
            }

            $prev = $state[$key]
            $prevRead = [uint64]0
            $prevWrite = [uint64]0
            $prevSent = [uint64]0
            $prevRecv = [uint64]0
            if ($prev -is [hashtable]) {
                if ($prev.ContainsKey('ReadRaw')) { $prevRead = [uint64]$prev.ReadRaw }
                if ($prev.ContainsKey('WriteRaw')) { $prevWrite = [uint64]$prev.WriteRaw }
                if ($prev.ContainsKey('SentRaw')) { $prevSent = [uint64]$prev.SentRaw }
                if ($prev.ContainsKey('ReceivedRaw')) { $prevRecv = [uint64]$prev.ReceivedRaw }
            }

            $record = ConvertTo-SMBeatSampleRecord -Snapshot $snap -Server $label -Utc $utc -IntervalSec $IntervalSec `
                -ReadDelta (Get-SMBeatByteDelta -Current $snap.ReadRaw -Previous $prevRead) `
                -WriteDelta (Get-SMBeatByteDelta -Current $snap.WriteRaw -Previous $prevWrite) `
                -SentDelta (Get-SMBeatByteDelta -Current $snap.SentRaw -Previous $prevSent) `
                -ReceivedDelta (Get-SMBeatByteDelta -Current $snap.ReceivedRaw -Previous $prevRecv)

            $emitted.Add($record) | Out-Null
        }

        Save-SMBeatState -DataRoot $dataRoot -Counters $newState -Utc $utc
        Add-SMBeatSampleLines -DataRoot $dataRoot -Utc $utc -Records $emitted.ToArray()
        Invoke-SMBeatRetention -DataRoot $dataRoot -RetentionDays $RetentionDays
        Save-SMBeatHeartbeat -DataRoot $dataRoot -Utc $utc -IntervalSec $IntervalSec -ComputerName $label -LastError '' -Success -EmittedCount $emitted.Count

        [PSCustomObject]@{
            ComputerName = $label
            Utc          = (ConvertTo-SMBeatUtcString -DateTime $utc)
            Path         = $dataRoot
            Emitted      = $emitted.Count
            Success      = $true
        }
    }
    catch {
        $msg = $_.Exception.Message
        Save-SMBeatHeartbeat -DataRoot $dataRoot -Utc $utc -IntervalSec $IntervalSec -ComputerName $label -LastError $msg -EmittedCount 0
        Write-SMBeatLog -DataRoot $dataRoot -Message $msg
        throw
    }
}