Public/Runner/Invoke-RetriableCommand.ps1

using module ..\..\Private\Result.psm1
using namespace System.Management.Automation
using namespace System.Collections.ObjectModel

#Requires -Psedition Core
function Invoke-RetriableCommand {
  <#
  .SYNOPSIS
    Runs Retriable Commands
  .DESCRIPTION
    Retries a script or command for a number of times (3 times by default) or until it succeeds.
    This is handy if you know a command or script will fail sometimes and you want to retry it a few times.
    In the return value, example: $result.HasErrors tells you if there were any errors during retries.
    See: Get-Help Invoke-RetriableCommand -Examples
  .EXAMPLE
   # Run a simple ScriptBlock with Attempts
   $result = Invoke-RetriableCommand -ScriptBlock {
       Write-Output "Running retry test..."
       Start-Sleep -Seconds 1
 
       # Simulate success after a few retries using current second
       if ([DateTime]::Now.Second % 5 -eq 0) {
          return [PSCustomObject]@{ StatusCode = 200; Message = "Success" }
       } else {
          throw "Simulated failure. Reason: $([DateTime]::Now.Second)%5 -eq 0 is false"
        }
    } -MaxAttempts 4 -Message "Simulate success after a few attempts"
 
  .EXAMPLE
   # Retry a remote command on a target machine
   $scriptBlock = {
       param($arg1)
       Write-Output "Processing on remote computer with argument: $arg1"
       # Simulate conditional success
       if ($arg1 -eq 'SuccessValue') {
           return 0 # Success return code
       }
       throw "Execution failed on remote machine"
   }
   $result = Invoke-RetriableCommand -ComputerName 'Server01' -ScriptBlock $scriptBlock -ArgumentList 'SuccessValue' -MaxAttempts 3 -Message "Executing remote retry"
 
  .EXAMPLE
   # Retry execution of a local file/script
   $result = Invoke-RetriableCommand -FilePath "C:\Scripts\TestScript.ps1" -ArgumentList "Param1", "Param2" -MaxAttempts 4 -Timeout 60 -Message "Retrying file execution"
 
  .EXAMPLE
   # Customizing success return codes
   $result = Invoke-RetriableCommand -ScriptBlock {
       # Simulate success with custom return code 200
       Write-Output "Custom success return code"
       Start-Sleep -Seconds 1
       return 200
   } -MaxAttempts 3 -SuccessReturnCodes @(0, 200) -Message "Custom success code retry"
 
  .EXAMPLE
   # Retry with a CancellationToken
   $cancellationTokenSource = New-Object System.Threading.CancellationTokenSource
   Start-Job -ScriptBlock {
       Start-Sleep -Seconds 3
       $cancellationTokenSource.Cancel()
   }
   $result = Invoke-RetriableCommand -ScriptBlock {
       Write-Output "Attempting long-running command execution"
       Start-Sleep -Seconds 10 # Simulated long-running task
   } -MaxAttempts 5 -Timeout 15 -CancellationToken $cancellationTokenSource.Token -Message "Retry with cancellation token"
 
  .NOTES
    - Requires Core Psedition due to ternary-operator
      https://learn.microsoft.com/en-us/powershell/module/microsoft.powershell.core/about/about_if?view=powershell-7.4#using-the-ternary-operator-syntax
 
    - All Unnamed arguments will be passed as arguments to the script or command
  .LINK
    Online Version: https://github.com/chadnpc/cliHelper.Core/blob/main/Public/Psrunner/Invoke-RetriableCommand.ps1
  #>

  [CmdletBinding(SupportsShouldProcess, DefaultParameterSetName = 'ScriptBlock')]
  [OutputType([Results])][Alias('Invoke-RtCommand')]
  param (
    [Parameter(Mandatory = $false, Position = 0, ParameterSetName = '__AllParameterSets')]
    [ValidateNotNullOrWhiteSpace()]
    [string]$ComputerName,

    [Parameter(Mandatory = $false, Position = 1, ParameterSetName = 'ScriptBlock')]
    [Alias('Script')][ValidateNotNullOrEmpty()]
    [ScriptBlock]$ScriptBlock,

    [Parameter(Mandatory = $false, Position = 1, ParameterSetName = 'Command')]
    [Alias('Command')][ValidateNotNullOrWhiteSpace()]
    [string]$FilePath,

    [Parameter(Mandatory = $false, Position = 2, ParameterSetName = '__AllParameterSets')]
    [Object[]]$ArgumentList,


    [Parameter(Mandatory = $false, Position = 3, ParameterSetName = '__AllParameterSets')]
    [Alias('Retries', 'MaxRetries')][ValidateNotNullOrEmpty()]
    [int]$MaxAttempts = 3,

    [Parameter(Mandatory = $false, Position = 4, ParameterSetName = '__AllParameterSets')]
    [uint32[]]$SuccessReturnCodes = @(0, 3010),

    [Parameter(Mandatory = $false, Position = 5, ParameterSetName = '__AllParameterSets')]
    [ValidateNotNullOrWhiteSpace()]
    [string]$WorkingDirectory,

    # Timeout in milliseconds
    [Parameter(Mandatory = $false, Position = 6, ParameterSetName = '__AllParameterSets')]
    [Alias('t')][ValidateNotNullOrEmpty()]
    [int]$Timeout = 500,

    [Parameter(Mandatory = $false, Position = 7, ParameterSetName = '__AllParameterSets')]
    [System.Threading.CancellationToken]$CancellationToken = [System.Threading.CancellationToken]::None,

    # The message to display in the verbose stream
    [Parameter(Mandatory = $false, Position = 8, ParameterSetName = '__AllParameterSets')]
    [Alias('msg')][ValidateNotNullOrWhiteSpace()]
    [string]$Message,

    [Parameter(Mandatory = $false, Position = 9, ParameterSetName = '__AllParameterSets')]
    [string]$caller,

    [Parameter(Mandatory = $false, Position = 10, ParameterSetName = 'Command')]
    [switch]$ExpandStrings
  )
  dynamicparam {
    $dynamicParams = [System.Management.Automation.RuntimeDefinedParameterDictionary]::new()
    $attributeCollection = [System.Collections.ObjectModel.Collection[System.Attribute]]::new()
    $attributes = [System.Management.Automation.ParameterAttribute]::new(); $attHash = @{
      Position                        = 11
      ParameterSetName                = '__AllParameterSets'
      Mandatory                       = $false
      ValueFromPipeline               = $false
      ValueFromPipelineByPropertyName = $false
      ValueFromRemainingArguments     = $true
      HelpMessage                     = 'Allows splatting with arguments that do not apply. Do not use directly.'
      DontShow                        = $False
    }; $attHash.Keys | ForEach-Object { $attributes.$_ = $attHash.$_ }
    $attributeCollection.Add($attributes)
    $RuntimeParam = [System.Management.Automation.RuntimeDefinedParameter]::new("IgnoredArguments", [Object[]], $attributeCollection)
    $dynamicParams.Add("IgnoredArguments", $RuntimeParam)
    return $dynamicParams
  }

  begin {
    [ActionPreference]$eap = $ErrorActionPreference; $ErrorActionPreference = "SilentlyContinue"
    $fxn = ![string]::IsNullOrWhiteSpace($caller) ? $caller : [PsRunner]::caller; $PsBoundParameters.GetEnumerator() | ForEach-Object { New-Variable -Name $_.Key -Value $_.Value -ea 'SilentlyContinue' }
    $cmdColors = [PsCustomObject]@{
      Verbose     = 'LemonChiffon'
      Debug       = 'Lavender'
      Error       = 'Salmon'
      Warning     = 'Yellow'
      Progress    = 'SpringGreen'
      Information = 'LawnGreen'
    }
  }

  process {
    $Attempts = 1; $Results = [Results]::new()
    if ($PsBoundParameters.ContainsKey("Message") -and $verbose) {
      Write-Console "$fxn $Message" -f $cmdColors.Information
    }
    while (($Attempts -le $MaxAttempts) -and !$Results.IsSuccess) {
      $AttemptStartTime = Get-Date; $Retries = $MaxAttempts - $Attempts
      if ($cancellationToken.IsCancellationRequested) { $verbose ? (Write-Console "$fxn CancellationRequested when $Retries retries were left." -f $cmdColors.Verbose) : $null; throw }

      $verbose ? (Write-Console "$fxn Attempt # $Attempts/$MaxAttempts ..." -f $cmdColors.Progress) : $null

      $Result = Invoke-Safely {
        if ($PSCmdlet.ParameterSetName -eq 'Command') {
          $verbose ? (Write-Console "Running command line [$FilePath $ArgumentList] on $ComputerName" -f LemonChiffon) : $null
          Invoke-Command -ComputerName $ComputerName -ScriptBlock {
            $VerbosePreference = $using:VerbosePreference
            $WhatIfPreference = $using:WhatIfPreference
            $ps = [System.Diagnostics.Process]::new()
            $ps_startInfo = [System.Diagnostics.ProcessStartInfo]::new()
            $ps_startInfo.FileName = $Using:FilePath;
            if ($Using:ArgumentList) {
              $ps_startInfo.Arguments = $Using:ArgumentList;
              if ($Using:ExpandStrings) {
                $ps_startInfo.Arguments = $ExecutionContext.InvokeCommandWithCred.ExpandString($Using:ArgumentList);
              }
            }
            if ($Using:WorkingDirectory) {
              $ps_startInfo.WorkingDirectory = $Using:WorkingDirectory;
              if ($Using:ExpandStrings) {
                $ps_startInfo.WorkingDirectory = $ExecutionContext.InvokeCommandWithCred.ExpandString($Using:WorkingDirectory);
              }
            }
            $ps_startInfo.UseShellExecute = $false; # This is critical for installs to function on core servers
            $ps.StartInfo = $ps_startInfo;
            $verbose ? (Write-Console "Starting Process path [$($ps_startInfo.FileName)] - Args: [$($ps_startInfo.Arguments)] - Working dir: [$($Using:WorkingDirectory)]" -f LemonChiffon) : $null
            $null = $ps.Start();
            if (!$ps) {
              throw "Error running program: $($ps.ExitCode)"
            } else {
              $ps.WaitForExit()
            }
            # Check the exit code of the process to see if it succeeded.
            if ($ps.ExitCode -notin $Using:SuccessReturnCodes) {
              throw "Error running program: $($ps.ExitCode)"
            }
          }
        } else {
          Invoke-Command -ScriptBlock $ScriptBlock -ArgumentList $ArgumentList
        }
      }

      $Elapsed = [math]::Round(($(Get-Date) - $AttemptStartTime).TotalSeconds, 2)
      $Results.Add($Result, $Elapsed)

      if (!$Result.IsOk()) {
        $verbose ? (Write-Console "$fxn Error after $Elapsed seconds:" -f $cmdColors.Verbose) : $null
        $err = $Result.UnwrapErr()
        if ($err -is [System.Management.Automation.ErrorRecord]) {
          $verbose ? (Write-Console " $($err.CategoryInfo.Category) : $($err.CategoryInfo.Reason) : $($err.Exception.Message)" -f $cmdColors.Error) : $null
        } else {
          $verbose ? (Write-Console " $err" -f $cmdColors.Error) : $null
        }
      }

      if (!$Result.IsOk() -and ($Retries -gt 0) -and !$cancellationToken.IsCancellationRequested) {
        $verbose ? (Write-Console "$fxn Waiting $Timeout ms before retrying. Retries left: $Retries" -f $cmdColors.Verbose) : $null
        Start-Sleep -Milliseconds $Timeout
      }
      $Attempts++
    }
  }

  end {
    if ($verbose) {
      $e = @{
        0 = @{
          c = "Error"
          m = "$Message Completed With Errors. Total time elapsed $($Results.ElapsedTime). Check the log file `$LogPath".Trim()
        };
        1 = @{
          c = "Information"
          m = "$Message Completed Successfully. Total time elapsed $($Results.ElapsedTime)".Trim()
        }
      }[[int]$Results.IsSuccess]
      Write-Console "$fxn $($e.m)" -f $cmdColors.($e.c)
    }
    $ErrorActionPreference = $eap;
    return $Results
  }
}