Files
utilities/personal/tools/pgbench_benchmark.ps1
2026-04-23 08:26:38 -04:00

269 lines
8.7 KiB
PowerShell

param(
[string]$PgHost = "localhost",
[int]$PgPort = 5432,
[string]$PgUser = "postgres",
[string]$BenchDb = "pgbench",
[int]$Scale = 50,
[int]$Duration = 60,
[string]$Clients = "1 4 8 16 32",
[int]$Jobs = 4,
[int]$WarmupSeconds = 15,
[bool]$RunReadOnly = $true,
[bool]$RunCustom = $false,
[string]$CustomSqlFile = "pgbench_custom_workload.sql",
[string]$CustomModeName = "custom_sql",
[ValidateRange(1, 99)][int]$CustomReadPct = 80,
[ValidateRange(1, 100)][int]$CustomHotPct = 90,
[ValidateRange(1, 4)][int]$CustomTxnSize = 2,
[ValidateRange(100, 100000000)][int]$CustomHotAccounts = 10000,
[bool]$Reinit = $false
)
$ErrorActionPreference = "Stop"
function Require-Command {
param([Parameter(Mandatory = $true)][string]$Name)
if (-not (Get-Command $Name -ErrorAction SilentlyContinue)) {
throw "Missing required command: $Name"
}
}
function Invoke-PsqlScalar {
param(
[Parameter(Mandatory = $true)][string]$Database,
[Parameter(Mandatory = $true)][string]$Sql
)
$output = & psql -X -v ON_ERROR_STOP=1 -h $PgHost -p $PgPort -U $PgUser -d $Database -Atc $Sql
if ($LASTEXITCODE -ne 0) {
throw "psql failed running SQL: $Sql"
}
return ($output | Select-Object -First 1)
}
function Invoke-PsqlNonQuery {
param(
[Parameter(Mandatory = $true)][string]$Database,
[Parameter(Mandatory = $true)][string]$Sql
)
& psql -X -v ON_ERROR_STOP=1 -h $PgHost -p $PgPort -U $PgUser -d $Database -c $Sql | Out-Null
if ($LASTEXITCODE -ne 0) {
throw "psql failed running SQL: $Sql"
}
}
function Ensure-BenchmarkDatabase {
$exists = Invoke-PsqlScalar -Database "postgres" -Sql "SELECT 1 FROM pg_database WHERE datname = '$BenchDb';"
if ($exists -ne "1") {
Write-Host "Creating benchmark database: $BenchDb"
Invoke-PsqlNonQuery -Database "postgres" -Sql "CREATE DATABASE $BenchDb;"
}
}
function Needs-Init {
$hasTable = Invoke-PsqlScalar -Database $BenchDb -Sql "SELECT to_regclass('public.pgbench_accounts') IS NOT NULL;"
return $hasTable -ne "t"
}
function Init-OrReinitData {
if ($Reinit -or (Needs-Init)) {
Write-Host "Initializing pgbench schema/data (scale=$Scale)"
& pgbench -h $PgHost -p $PgPort -U $PgUser -i -s $Scale $BenchDb | Out-Host
if ($LASTEXITCODE -ne 0) {
throw "pgbench initialization failed"
}
}
else {
Write-Host "Using existing pgbench data in $BenchDb"
}
}
function Parse-Tps {
param([Parameter(Mandatory = $true)][string]$Text)
$matches = [regex]::Matches($Text, "(?m)^tps\s*=\s*([0-9]+(?:\.[0-9]+)?)")
if ($matches.Count -gt 0) {
return $matches[$matches.Count - 1].Groups[1].Value
}
return ""
}
function Parse-AvgLatencyMs {
param([Parameter(Mandatory = $true)][string]$Text)
$matches = [regex]::Matches($Text, "(?m)^latency average\s*=\s*([0-9]+(?:\.[0-9]+)?)\s*ms")
if ($matches.Count -gt 0) {
return $matches[$matches.Count - 1].Groups[1].Value
}
return ""
}
function Run-Case {
param(
[Parameter(Mandatory = $true)][string]$Mode,
[Parameter(Mandatory = $true)][int]$ClientCount,
[AllowEmptyCollection()][string[]]$ExtraArgs = @(),
[Parameter(Mandatory = $true)][string]$ResultDir,
[Parameter(Mandatory = $true)][string]$SummaryCsv
)
$logFile = Join-Path $ResultDir ("{0}_c{1}.txt" -f $Mode, $ClientCount)
$txLogPrefix = Join-Path $ResultDir ("{0}_c{1}_txlog" -f $Mode, $ClientCount)
Write-Host "Running mode=$Mode, clients=$ClientCount, duration=${Duration}s"
# Warm cache and avoid measuring first-run effects.
& pgbench -h $PgHost -p $PgPort -U $PgUser -n -M prepared -j $Jobs -c $ClientCount -T $WarmupSeconds @ExtraArgs $BenchDb | Out-Null
if ($LASTEXITCODE -ne 0) {
throw "Warm-up run failed for mode=$Mode clients=$ClientCount"
}
$runOutput = & pgbench -h $PgHost -p $PgPort -U $PgUser -n -M prepared -j $Jobs -c $ClientCount -T $Duration -r -l --log-prefix $txLogPrefix @ExtraArgs $BenchDb 2>&1
if ($LASTEXITCODE -ne 0) {
throw "Measured run failed for mode=$Mode clients=$ClientCount"
}
$runOutput | Set-Content -Path $logFile -Encoding UTF8
$text = ($runOutput -join [Environment]::NewLine)
$tps = Parse-Tps -Text $text
$latencyMs = Parse-AvgLatencyMs -Text $text
$percentiles = Get-PercentileLatency -LogPrefix $txLogPrefix
"$Mode,$ClientCount,$Duration,$Jobs,$tps,$latencyMs,$($percentiles.P95Ms),$($percentiles.P99Ms),$logFile" | Add-Content -Path $SummaryCsv
}
function Get-PercentileLatency {
param([Parameter(Mandatory = $true)][string]$LogPrefix)
$logFiles = Get-ChildItem -Path ($LogPrefix + "*") -File -ErrorAction SilentlyContinue
if (-not $logFiles -or $logFiles.Count -eq 0) {
return @{ P95Ms = ""; P99Ms = "" }
}
# Keep memory bounded even for very large transaction logs.
$maxSampleSize = 200000
$reservoir = New-Object System.Collections.Generic.List[double]
$reservoir.Capacity = $maxSampleSize
$random = [System.Random]::new()
$seen = 0
foreach ($file in $logFiles) {
$reader = [System.IO.File]::OpenText($file.FullName)
try {
while (($line = $reader.ReadLine()) -ne $null) {
if ([string]::IsNullOrWhiteSpace($line) -or $line.StartsWith("#")) {
continue
}
$tokens = $line -split "[\s,]+" | Where-Object { $_ -ne "" }
if ($tokens.Count -ge 3 -and $tokens[2] -match "^[0-9]+(?:\.[0-9]+)?$") {
$latUs = [double]$tokens[2]
$seen++
if ($reservoir.Count -lt $maxSampleSize) {
$reservoir.Add($latUs)
}
else {
$idx = $random.Next(0, $seen)
if ($idx -lt $maxSampleSize) {
$reservoir[$idx] = $latUs
}
}
}
}
}
finally {
$reader.Dispose()
}
}
if ($reservoir.Count -eq 0) {
return @{ P95Ms = ""; P99Ms = "" }
}
$sorted = $reservoir | Sort-Object
$p95Us = Select-PercentileValue -SortedValues $sorted -Percentile 95
$p99Us = Select-PercentileValue -SortedValues $sorted -Percentile 99
return @{
P95Ms = [math]::Round(($p95Us / 1000.0), 3)
P99Ms = [math]::Round(($p99Us / 1000.0), 3)
}
}
function Select-PercentileValue {
param(
[Parameter(Mandatory = $true)][double[]]$SortedValues,
[Parameter(Mandatory = $true)][ValidateRange(1, 99)][int]$Percentile
)
$n = $SortedValues.Count
$rank = [math]::Ceiling(($Percentile / 100.0) * $n)
if ($rank -lt 1) { $rank = 1 }
if ($rank -gt $n) { $rank = $n }
return $SortedValues[$rank - 1]
}
function Resolve-CustomSqlPath {
if ([System.IO.Path]::IsPathRooted($CustomSqlFile)) {
return $CustomSqlFile
}
return (Join-Path (Get-Location) $CustomSqlFile)
}
Require-Command -Name "psql"
Require-Command -Name "pgbench"
$timestamp = Get-Date -Format "yyyyMMdd_HHmmss"
$resultDir = Join-Path (Get-Location) ("pgbench_results_{0}" -f $timestamp)
$summaryCsv = Join-Path $resultDir "summary.csv"
New-Item -Path $resultDir -ItemType Directory -Force | Out-Null
"mode,clients,duration_seconds,jobs,tps,latency_ms,p95_ms,p99_ms,details_file" | Set-Content -Path $summaryCsv -Encoding UTF8
Ensure-BenchmarkDatabase
Init-OrReinitData
$customSqlPath = ""
if ($RunCustom) {
$customSqlPath = Resolve-CustomSqlPath
if (-not (Test-Path -Path $customSqlPath -PathType Leaf)) {
throw "Custom SQL file not found: $customSqlPath"
}
}
$clientValues = $Clients -split "\s+" | Where-Object { $_ -match "^[0-9]+$" }
if ($clientValues.Count -eq 0) {
throw "No valid client values found. Example: -Clients '1 4 8 16 32'"
}
foreach ($c in $clientValues) {
$clientInt = [int]$c
Run-Case -Mode "read_write" -ClientCount $clientInt -ExtraArgs @() -ResultDir $resultDir -SummaryCsv $summaryCsv
if ($RunReadOnly) {
Run-Case -Mode "read_only" -ClientCount $clientInt -ExtraArgs @("-S") -ResultDir $resultDir -SummaryCsv $summaryCsv
}
if ($RunCustom) {
Run-Case -Mode $CustomModeName -ClientCount $clientInt -ExtraArgs @(
"-D", "read_pct=$CustomReadPct",
"-D", "hot_pct=$CustomHotPct",
"-D", "txn_size=$CustomTxnSize",
"-D", "hot_accounts=$CustomHotAccounts",
"-f", $customSqlPath
) -ResultDir $resultDir -SummaryCsv $summaryCsv
}
}
Write-Host ""
Write-Host "Benchmark completed."
Write-Host "Summary: $summaryCsv"
Write-Host "Raw outputs: $resultDir"