Read process StandardOutput asynchronously using event handlers

redirectedprocess-spike.ps1

enum StdIOStream {
    StandardOutput
    StandardError
}

class StdIOData {
    [datetime]    $Timestamp
    [int]         $Pid
    [StdIOStream] $Stream
    [string]      $Line
}

class StdIOReceiver {
    hidden static [hashtable] $Data = @{}
    hidden static [hashtable] $LastUpdated = @{}
    hidden static [hashtable] $SourceIdentifiers = @{}

    static [System.Collections.Generic.List[StdIOData]] Read([System.Diagnostics.Process]$process) {
        if ($null -eq ($queue = [StdIOReceiver]::Data[$process.Id])) {
            throw ([System.InvalidOperationException]::new("Process has not been registered with StdIOReceiver."))
        }
        $result = [System.Collections.Generic.List[StdIOData]]::new()
        while ($queue.Count) {
            $result.Add($queue.Dequeue())
        }
        return $result
    }

    static [void] Register([System.Diagnostics.Process]$process) {
        [StdIOReceiver]::Data[$process.Id] = [System.Collections.Generic.Queue[StdIOData]]::new()
        [StdIOReceiver]::LastUpdated[$process.Id] = $null
        [StdIOReceiver]::SourceIdentifiers[$process.Id] = (New-Guid), (New-Guid), (New-Guid)

        Register-ObjectEvent -InputObject $process -EventName OutputDataReceived -Action {
            $data = [StdIOData]@{
                Timestamp = [datetime]::UtcNow
                Pid       = $Sender.Id
                Stream    = [StdIOStream]::StandardOutput
                Line      = $EventArgs.Data
            }
            [StdIOReceiver]::LastUpdated[$Sender.Id] = $data.Timestamp
            [StdIOReceiver]::Data[$Sender.Id].Enqueue($data)
        } -SourceIdentifier ([StdIOReceiver]::SourceIdentifiers[$process.Id][0])

        Register-ObjectEvent -InputObject $process -EventName ErrorDataReceived -Action {
            $data = [StdIOData]@{
                Timestamp = [datetime]::UtcNow
                Pid       = $Sender.Id
                Stream    = [StdIOStream]::StandardError
                Line      = $EventArgs.Data
            }
            [StdIOReceiver]::LastUpdated[$Sender.Id] = $data.Timestamp
            [StdIOReceiver]::Data[$Sender.Id].Enqueue($data)
        } -SourceIdentifier ([StdIOReceiver]::SourceIdentifiers[$process.Id][1])

        Register-ObjectEvent -InputObject $process -EventName Exited -Action {
            if ($process = $Sender -as [System.Diagnostics.Process]) {
                [StdIOReceiver]::Unregister($process)
            }
        } -SourceIdentifier ([StdIOReceiver]::SourceIdentifiers[$process.Id][2])
    }

    static [void] Unregister([System.Diagnostics.Process]$process) {
        try {
            foreach ($id in [StdIOReceiver]::SourceIdentifiers[$process.Id]) {
                Unregister-Event -SourceIdentifier $id.ToString() -Force -ErrorAction Stop
            }
            [StdIOReceiver]::SourceIdentifiers.Remove($process.Id)
            [StdIOReceiver]::LastUpdated.Remove($process.Id)
        } catch {
            throw
        }
    }
}

function Start-RedirectedProcess {
    [CmdletBinding()]
    [OutputType([System.Diagnostics.Process])]
    param(
        [Parameter(Mandatory, Position = 0)]
        [Alias('PSPath', 'Path')]
        [string]
        $FilePath,

        [Parameter(ValueFromRemainingArguments, Position = 1)]
        [Alias('Args')]
        [string[]]
        $ArgumentList
    )

    process {
        try {
            $p = [System.Diagnostics.Process]::new()
            $p.StartInfo = [System.Diagnostics.ProcessStartInfo]@{
                FileName               = $FilePath
                Arguments              = $ArgumentList -join ' '
                CreateNoWindow         = $true
                RedirectStandardInput  = $true
                RedirectStandardOutput = $true
                RedirectStandardError  = $true
                UseShellExecute        = $false
            }
            $p.EnableRaisingEvents = $true
            if (-not $p.Start()) {
                throw "Failed to start process."
            }
            $p.BeginOutputReadLine()
            $p.BeginErrorReadLine()
            [StdIOReceiver]::Register($p)
            $p
        } catch {
            throw
        }
    }
}

function Read-RedirectedProcess {
    [CmdletBinding()]
    [OutputType([StdIOData])]
    param(
        [Parameter(Mandatory, Position = 0, ValueFromPipeline)]
        [System.Diagnostics.Process]
        $Process,

        [Parameter()]
        [switch]
        $Wait
    )

    process {
        do {
            $exited = $Process.HasExited
            [StdIOReceiver]::Read($Process)
            if ($Wait) {
                Start-Sleep -Milliseconds 100
            }
        } while ($Wait -and -not $exited)
    }

    end {
        if ($Process.HasExited) {
            [StdIOReceiver]::Unregister($Process)
        }
    }
}

function Write-RedirectedProcess {
    [CmdletBinding()]
    param(
        [Parameter(Mandatory, ValueFromPipeline)]
        [System.Diagnostics.Process]
        $Process,

        [Parameter(Mandatory, Position = 0)]
        [string[]]
        $StandardInput,

        [Parameter()]
        [switch]
        $PassThru
    )

    process {
        foreach ($s in $StandardInput) {
            if ($Process.HasExited) {
                # No sense in writing to StandardInput if the process isn't listening anymore.
                [StdIOReceiver]::Unregister($Process)
                Write-Error "Process with ID $($Process.Id) exited with ExitCode $($Process.ExitCode)."
                break
            }
            $Process.StandardInput.WriteLine($s)
        }
        if ($PassThru) {
            $Process
        }
    }
}

function Wait-RedirectedProcess {
    [CmdletBinding()]
    param(
        [Parameter(Mandatory, ValueFromPipeline, ParameterSetName = 'Exited')]
        [Parameter(Mandatory, ValueFromPipeline, ParameterSetName = 'Idle')]
        [System.Diagnostics.Process]
        $Process,

        [Parameter(Mandatory, ParameterSetName = 'Exited')]
        [switch]
        $Exited,

        [Parameter(Mandatory, ParameterSetName = 'Idle')]
        [switch]
        $Idle,

        [Parameter(ParameterSetName = 'Idle')]
        [timespan]
        $IdleTime = [timespan]::FromSeconds(2),

        [Parameter()]
        [switch]
        $PassThru
    )

    process {
        if (-not [StdIOReceiver]::LastUpdated.ContainsKey($Process.Id)) {
            Write-Error "Process $($Process.Id) has not been registered with [StdIOReceiver]."
            return
        }

        switch ($PSCmdlet.ParameterSetName) {
            'Exited' {
                while (-not $Process.HasExited) {
                    Start-Sleep -Milliseconds 100
                }
            }

            'Idle' {
                $lastUpdated = [datetime]::UtcNow
                $timeout = $lastUpdated.Add($IdleTime)
                while ($timeout -gt [datetime]::UtcNow) {
                    if ([StdIOReceiver]::LastUpdated[$Process.Id] -gt $lastUpdated) {
                        $lastUpdated = [StdIOReceiver]::LastUpdated[$Process.Id]
                    }
                    $timeout = $lastUpdated.Add($IdleTime)
                    Start-Sleep -Milliseconds 100
                }
            }

            default {
                throw ([System.NotImplementedException]::new("ParameterSetName '$_' not implemented."))
            }
        }

        if ($PassThru) {
            $Process
        }
    }
}


# Start nslookup, wait until stdout and stderr have been idle for default 2 seconds
$process = Start-RedirectedProcess nslookup | Wait-RedirectedProcess -Idle -PassThru

# Grab the initial output from the process, write the records out, and if one
# of them matches "Default Server:", write a URL to query
$process | Read-RedirectedProcess -PipelineVariable record | Foreach-Object {
    $record
    if ($record.Line -match 'Default Server:') {
        $process | Write-RedirectedProcess "www.powershellgallery.com"
    }
}

# Wait until the process output is idle again, then read all the output
$process | Wait-RedirectedProcess -Idle -PassThru | Read-RedirectedProcess

# Write "exit" to tell nslookup we want out, wait until the process has exited
# and then read any output we might have missed at the end.
$process | Write-RedirectedProcess "exit" | Wait-RedirectedProcess -Exited -PassThru | Read-RedirectedProcess