BGProcess

0.1.1

BGProcess.psm1


using namespace System.Management.Automation
using namespace System.Management.Automation.Runspaces
using namespace System.Threading.Tasks

class AsyncShellOutput {
    [PSDataCollection[psobject]]    $Result
    [PSDataCollection[ErrorRecord]] $ErrorStream
}

class AsyncShell {
    [powershell]   $Shell
    [IAsyncResult] $AsyncResult

    AsyncShell() {}

    AsyncShell([powershell]$powershell) {
        $this.Shell = $powershell
    }

    AsyncShell([runspace]$runspace) {
        $this.Shell = [powershell]::Create($runspace)
    }

    AsyncShell([RunspacePool]$runspacepool) {
        $this.Shell = [powershell]::Create()
        $this.Shell.RunspacePool = $runspacepool
    }

    [AsyncShellOutput] CollectResults() {
        if ($this.AsyncResult.IsCompleted) {
            try {
                return ([AsyncShellOutput]@{
                    Result      = $this.Shell.EndInvoke($this.AsyncResult)
                    ErrorStream = $this.Shell.Streams.Error
                })
            } finally {
                $this.Shell.Streams.ClearStreams()
                $this.AsyncResult = $null
            }
        }
        return $null
    }
}

class BGProcessData {
    [string]        $StandardOutput
    [string]        $StandardError
    [ErrorRecord[]] $ErrorRecords
    [bool]          $HasData
}

class BGProcessDataBuilder {
    hidden [text.stringbuilder]                    $stdOut
    hidden [text.stringbuilder]                    $stdErr
    hidden [collections.generic.list[errorrecord]] $errors
    hidden [bool]                                  $hasData
    BGProcessDataBuilder() {
        $this.stdOut = [text.stringbuilder]::new()
        $this.stdErr = [text.stringbuilder]::new()
        $this.errors = [collections.generic.list[errorrecord]]::new()
    }

    [void] AddStandardOutput([string]$data) {
        $this.stdOut.Append($data)
        $this.hasData = $true
    }

    [void] AddStandardError([string]$data) {
        $this.stdErr.Append($data)
        $this.hasData = $true
    }

    [void] AddErrorRecord([errorrecord]$errorrecord) {
        $this.errors.Add($errorrecord)
        $this.hasData = $true
    }

    [BGProcessData] ToBGProcessData() {
        return ([BGProcessData]@{
            StandardOutput = $this.stdOut.ToString()
            StandardError  = $this.stdErr.ToString()
            ErrorRecords   = $this.errors.ToArray()
            HasData        = $this.hasData
        })
    }
}


class BGProcess {
    [System.Diagnostics.Process] $Process
    [int]                        $Id
    [string]                     $Name
    [nullable[int]]              $ExitCode
    [bool]                       $HasExited

    hidden [AsyncShell]$stdOut
    hidden [AsyncShell]$StdErr

    hidden static [hashtable] $Instances = @{}

    BGProcess ([diagnostics.process]$process) {
        $this.Process = $process
        $this.Id = $process.Id
        $this.Name = $process.Name
        [BGProcess]::Instances[$process.Id] = $this

        $pool = [runspacefactory]::CreateRunspacePool(2, 2)
        $pool.Open()
        $readStream = (Get-Command -Name ReadStream).ScriptBlock
        $stdOutShell = [powershell]::Create().AddScript($readStream).AddArgument($process.StandardOutput.BaseStream)
        $stdErrShell = [powershell]::Create().AddScript($readStream).AddArgument($process.StandardError.BaseStream)
        $stdOutShell.RunspacePool = $pool
        $stdErrShell.RunspacePool = $pool

        $this.stdOut = [AsyncShell]::new($stdOutShell)
        $this.stdErr = [AsyncShell]::new($stdErrShell)
        $this.stdOut.AsyncResult = $this.stdOut.Shell.BeginInvoke()
        $this.StdErr.AsyncResult = $this.StdErr.Shell.BeginInvoke()
        Register-ObjectEvent -InputObject $this.Process -EventName Exited -Action {
            $process = $Sender -as [diagnostics.process]
            [BGProcess]::Instances[$process.Id].ExitCode = $process.ExitCode
            [BGProcess]::Instances[$process.Id].HasExited = $true
        }
    }

    [void] Write([string] $data) {
        $this.Process.StandardInput.Write($data)
    }

    [BGProcessData] Read() {
        $dataBuilder = [BGProcessDataBuilder]::new()
        while ($this.stdOut.AsyncResult.IsCompleted) {
            $results = $this.stdOut.CollectResults()
            foreach ($record in $results.Result) {
                $dataBuilder.AddStandardOutput($record)
            }
            foreach ($record in $results.ErrorStream) {
                $dataBuilder.AddErrorRecord($record)
            }
            $this.stdOut.AsyncResult = $this.stdOut.Shell.BeginInvoke()
        }
        while ($this.StdErr.AsyncResult.IsCompleted) {
            $results = $this.StdErr.CollectResults()
            foreach ($record in $results.Result) {
                $dataBuilder.AddStandardError($record)
            }
            foreach ($record in $results.ErrorStream) {
                $dataBuilder.AddErrorRecord($record)
            }
            $this.StdErr.AsyncResult = $this.StdErr.Shell.BeginInvoke()
        }
        return $dataBuilder.ToBGProcessData()
    }
}
function ReadStream {
    [CmdletBinding()]
    [OutputType([string])]
    param(
        [Parameter(Mandatory, Position = 0)]
        [ValidateNotNull()]
        [System.IO.Stream]
        $Stream,

        [Parameter(Position = 1)]
        [System.Text.Encoding]
        $Encoding = [System.Text.Encoding]::UTF8
    )

    process {
        $sb = [text.stringbuilder]::new()
        $buffer = [byte[]]::new(1KB)
        $read = 0
        do {
            $read = $Stream.Read($buffer, 0, $buffer.Length)
            if ($read -gt 0) {
                $null = $sb.Append($Encoding.GetString($buffer, 0, $read))
            }
        } while ($read -eq $buffer.Length)
        $sb.ToString()
    }
}
function Read-BGProcess {
    [CmdletBinding()]
    param(
        [Parameter(Mandatory, ValueFromPipeline)]
        [BGProcess]
        $Process,

        [Parameter()]
        [switch]
        $MapErrorsToStdOut,

        [Parameter()]
        [switch]
        $Wait,

        [Parameter()]
        [timespan]
        $Timeout = (New-TimeSpan -Seconds 1)
    )

    process {
        $stopwatch = [diagnostics.stopwatch]::StartNew()
        do {
            $data = $Process.Read()

            foreach ($record in $data.ErrorRecords) {
                $stopwatch.Restart()
                Write-Error -ErrorRecord $record
            }

            if (-not [string]::IsNullOrEmpty($data.StandardError)) {
                $stopwatch.Restart()
                if ($MapErrorsToStdOut) {
                    $data.StandardError
                } else {
                    Write-Error -Message $data.StandardError
                }
            }

            if (-not [string]::IsNullOrEmpty($data.StandardOutput)) {
                $stopwatch.Restart()
                $data.StandardOutput
            }

            if ($Wait -and $stopwatch.Elapsed -lt $Timeout) {
                Start-Sleep -Milliseconds 100
            }

        } while ($Wait -and $stopwatch.Elapsed -lt $Timeout)
    }
}
function Start-BGProcess {
    [CmdletBinding(SupportsShouldProcess)]
    [OutputType([BGProcess])]
    param(
        [Parameter(Mandatory, Position = 0)]
        [Alias('Path')]
        [string]
        $FileName,

        [Parameter(ValueFromRemainingArguments, Position = 1)]
        [string[]]
        $Arguments
    )

    process {
        try {
            $startInfo = [System.Diagnostics.ProcessStartInfo]@{
                FileName               = $FileName
                Arguments              = $Arguments -join ' '
                CreateNoWindow         = $true
                RedirectStandardInput  = $true
                RedirectStandardOutput = $true
                RedirectStandardError  = $true
                UseShellExecute        = $false
            }
            if ($PSCmdlet.ShouldProcess($FileName, "Start process with arguments: $($startInfo.Arguments)")) {
                $p = [System.Diagnostics.Process]::Start($startInfo)
                $p.EnableRaisingEvents = $true
                [BGProcess]::new($p)
            }
        } catch {
            throw
        }
    }
}
function Stop-BGProcess {
    [CmdletBinding(SupportsShouldProcess)]
    [OutputType([BGProcess])]
    param (
        [Parameter(Mandatory, ValueFromPipeline)]
        [BGProcess[]]
        $Process,

        [Parameter()]
        [switch]
        $PassThru
    )

    process {
        foreach ($p in $Process) {
            if ($PSCmdlet.ShouldProcess(("{0} ({1})" -f $p.Process.Name, $p.Id), "Stop-BGProcess")) {
                $p.process | Stop-Process
                if ($PassThru) {
                    $p
                }
            }
        }
    }
}
function Wait-BGProcess {
    [CmdletBinding()]
    [OutputType([BGProcess])]
    param (
        [Parameter(Mandatory, ValueFromPipeline)]
        [BGProcess[]]
        $Process,

        [Parameter()]
        [switch]
        $PassThru
    )

    process {
        foreach ($p in $Process) {
            while (-not $p.HasExited) {
                Start-Sleep -Milliseconds 100
            }
            if ($PassThru) {
                $p
            }
        }
    }
}
function Write-BGProcess {
    [CmdletBinding()]
    [OutputType([BGProcess])]
    param(
        [Parameter(Mandatory, ValueFromPipeline)]
        [BGProcess]
        $Process,

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

        [Parameter()]
        [string]
        $LineTerminator = ([System.Environment]::NewLine),

        [Parameter()]
        [switch]
        $NoLineTerminator,

        [Parameter()]
        [switch]
        $PassThru
    )

    process {
        $Process.Write($Text)
        if (-not $NoLineTerminator) {
            $Process.Write($LineTerminator)
        }
        if ($Passthru) {
            $Process
        }
    }
}