RateLimitedRestMethod.psm1

function Get-RateLimitedRestMethod {
    <#
    .SYNOPSIS
    Gets the current rate limiting configuration for this module.
 
    .DESCRIPTION
    This function will return the current rate limiting settings that are being used by Invoke-RateLimitedRestMethod.
 
    .EXAMPLE
    Get-RateLimitedRestMethod
 
    .OUTPUTS
    PSCustomObject
    #>

    [cmdletBinding()]
    param(

    )

    $state = Get-RateLimitedRestMethodState
    [pscustomobject]@{
        MinimumMs       = [double]$state.MinimumMs
        MinimumSeconds  = [double]$state.MinimumSeconds
        QueriesPerSecond = [double]$state.QueriesPerSecond
    }
}
function Invoke-RateLimitedRestMethod {
    <#
    .SYNOPSIS
    Runs Invoke-RestMethod with all the supplied parameters according to the rate limiting defined by Set-RateLimitedRestMethod
 
    .DESCRIPTION
    This is a passthru function. It will take all UnboundParameters and pass them straight through to Invoke-RestMethod. This function does not have any parameters so that it can work successfully on future versions of Invoke-RestMethod without issue.
 
    It uses the delay specified in Set-RateLimitedRestMethod to ensure that requests only flow at a specified rate.
 
    .EXAMPLE
    Invoke-RateLimitedRestMethod https://google.com
 
    .EXAMPLE
    irr https://google.com -WebSession $session
 
    .NOTES
    General notes
    #>

    [cmdletBinding()]
    [alias("irr")]
    param()

    dynamicparam {
        $commonParameters = @(
            "Verbose",
            "Debug",
            "ErrorAction",
            "WarningAction",
            "InformationAction",
            "ProgressAction",
            "ErrorVariable",
            "WarningVariable",
            "InformationVariable",
            "OutVariable",
            "OutBuffer",
            "PipelineVariable",
            "WhatIf",
            "Confirm"
        )

        $runtimeParameters = [System.Management.Automation.RuntimeDefinedParameterDictionary]::new()
        $invokeCommand = Get-Command -Name Invoke-RestMethod -CommandType Cmdlet

        foreach ($parameterName in $invokeCommand.Parameters.Keys) {
            if ($parameterName -in $commonParameters) {
                continue
            }

            $parameterMetadata = $invokeCommand.Parameters[$parameterName]
            $attributes = [System.Collections.ObjectModel.Collection[System.Attribute]]::new()
            foreach ($attribute in $parameterMetadata.Attributes) {
                [void]$attributes.Add($attribute)
            }

            $runtimeParameter = [System.Management.Automation.RuntimeDefinedParameter]::new(
                $parameterName,
                $parameterMetadata.ParameterType,
                $attributes
            )

            $runtimeParameters.Add($parameterName, $runtimeParameter)
        }

        $runtimeParameters
    }

    process {
        Wait-RateLimitedRestMethodTurn
        Invoke-RestMethod @PSBoundParameters
    }
}
function Set-RateLimitedRestMethod {
    <#
    .SYNOPSIS
    Sets the rate limiting for this module.
 
    .DESCRIPTION
    This function will accept one of several rate limiting configurations. The rate rate limiting will be used when using Invoke-RateLimitedRestMethod
 
    .PARAMETER MinimumMs
    Number of minimum milliseconds between requests
 
    .PARAMETER MinimumSeconds
    Number of seconds between requests
 
    .PARAMETER QueriesPerSecond
    Number of queries per second that are allowed (accepts decimals)
 
    .EXAMPLE
    Set-RateLimitedRestMethod -MinimumMs 250
 
    .EXAMPLE
    Set-RateLimitedRestMethod -QueriesPerSecond 4
 
    .EXAMPLE
    Set-RatelimitedRestMethod -QueriesPerSecond (240/60) -Passthru
    #>

    [cmdletBinding(DefaultParameterSetName = "QueriesPerSecond")]
    param(
        [Parameter(ParameterSetName = "DelayMs")]
        [double]$MinimumMs,

        [Parameter(ParameterSetName = "DelayS")]
        [double]$MinimumSeconds,

        [Parameter(ParameterSetName = "QueriesPerSecond")]
        [double]$QueriesPerSecond = 0,

        [switch]$Passthru
    )

    $state = Get-RateLimitedRestMethodState

    $minimumMsValue = 0.0
    switch ($PSCmdlet.ParameterSetName) {
        "DelayMs" {
            if ([double]::IsNaN($MinimumMs) -or [double]::IsInfinity($MinimumMs) -or $MinimumMs -lt 0) {
                throw "MinimumMs must be a finite value greater than or equal to 0."
            }
            $minimumMsValue = [double]$MinimumMs
        }
        "DelayS" {
            if ([double]::IsNaN($MinimumSeconds) -or [double]::IsInfinity($MinimumSeconds) -or $MinimumSeconds -lt 0) {
                throw "MinimumSeconds must be a finite value greater than or equal to 0."
            }
            $minimumMsValue = [double]$MinimumSeconds * 1000.0
        }
        "QueriesPerSecond" {
            if ([double]::IsNaN($QueriesPerSecond) -or [double]::IsInfinity($QueriesPerSecond) -or $QueriesPerSecond -lt 0) {
                throw "QueriesPerSecond must be a finite value greater than or equal to 0."
            }

            if ($QueriesPerSecond -eq 0) {
                $minimumMsValue = 0.0
            }
            else {
                $minimumMsValue = 1000.0 / [double]$QueriesPerSecond
            }
        }
    }

    $state.MinimumMs = [double]$minimumMsValue
    $state.MinimumSeconds = [double]($state.MinimumMs / 1000.0)
    if ($state.MinimumMs -le 0) {
        $state.QueriesPerSecond = 0.0
    }
    else {
        $state.QueriesPerSecond = [double](1000.0 / $state.MinimumMs)
    }

    $state.NextAllowedElapsedSeconds = $state.Stopwatch.Elapsed.TotalSeconds

    if ($Passthru) {
        Get-RateLimitedRestMethod
    }


}
function Get-RateLimitedRestMethodState {
    if (-not $script:RateLimitedRestMethodState) {
        $script:RateLimitedRestMethodState = [pscustomobject]@{
            MinimumMs               = 0.0
            MinimumSeconds          = 0.0
            QueriesPerSecond        = 0.0
            Stopwatch               = [System.Diagnostics.Stopwatch]::StartNew()
            NextAllowedElapsedSeconds = 0.0
        }
    }

    if (-not $script:RateLimitedRestMethodState.Stopwatch) {
        $script:RateLimitedRestMethodState.Stopwatch = [System.Diagnostics.Stopwatch]::StartNew()
    }

    $script:RateLimitedRestMethodState
}

function Wait-RateLimitedRestMethodTurn {
    $state = Get-RateLimitedRestMethodState
    if ($state.MinimumMs -le 0) {
        return
    }

    $throttleIntervalSeconds = [double]($state.MinimumMs / 1000.0)
    $elapsedSeconds = $state.Stopwatch.Elapsed.TotalSeconds

    if ($elapsedSeconds -lt $state.NextAllowedElapsedSeconds) {
        $sleepMilliseconds = [int][Math]::Ceiling(($state.NextAllowedElapsedSeconds - $elapsedSeconds) * 1000.0)
        Start-Sleep -Milliseconds ([Math]::Max(0, $sleepMilliseconds))
        $elapsedSeconds = $state.Stopwatch.Elapsed.TotalSeconds
    }

    $state.NextAllowedElapsedSeconds = [Math]::Max($state.NextAllowedElapsedSeconds + $throttleIntervalSeconds, $elapsedSeconds)
}