baiyue-firstleg/scripts/local-check.ps1

[CmdletBinding()]
param(
    [string[]]$Command,
    [string]$WorkingDirectory = (Get-Location).Path,
    [string]$Label = 'local check',
    [ValidateRange(0, 86400)][int]$WaitSeconds = 900,
    [ValidateRange(1, 86400)][int]$TimeoutSeconds = 1200,
    [switch]$Status
)

$ErrorActionPreference = 'Stop'
if (-not $IsWindows -or $PSVersionTable.PSVersion.Major -lt 7) {
    throw 'local-check requires Windows and PowerShell 7 (pwsh).'
}
# 所有工作目录共用两个用户级锁；首槽保留原锁名，运行中的旧入口仍计入容量。
$stateDirectory = Join-Path $env:LOCALAPPDATA 'Chongchong/LocalChecks'
[IO.Directory]::CreateDirectory($stateDirectory) | Out-Null
$slotDefinitions = @(
    [pscustomobject]@{ number = 1; lock = (Join-Path $stateDirectory 'heavy.lock'); owner = (Join-Path $stateDirectory 'owner.json'); result = (Join-Path $stateDirectory 'last-result.json') },
    [pscustomobject]@{ number = 2; lock = (Join-Path $stateDirectory 'heavy-2.lock'); owner = (Join-Path $stateDirectory 'owner-2.json'); result = (Join-Path $stateDirectory 'last-result-2.json') }
)
function Read-Owner($definition) {
    try { return [IO.File]::ReadAllText($definition.owner) } catch { return 'Owner information is being updated.' }
}
function Open-Slot($definition) {
    try { return [IO.File]::Open($definition.lock, 'OpenOrCreate', 'ReadWrite', 'None') }
    catch [IO.IOException] {
        if (($_.Exception.HResult -band 0xffff) -notin @(32, 33)) { throw }
        return $null
    }
}
function Open-AvailableSlot {
    foreach ($definition in $slotDefinitions) {
        $handle = Open-Slot $definition
        if ($handle) { return [pscustomobject]@{ definition = $definition; handle = $handle } }
    }
    return $null
}
if ($Status) {
    $busyCount = 0
    foreach ($definition in $slotDefinitions) {
        $probe = Open-Slot $definition
        if ($probe) { $probe.Dispose(); Write-Host "IDLE: slot $($definition.number)/2 is free." }
        else { $busyCount++; Write-Host "BUSY: slot $($definition.number)/2"; Write-Host (Read-Owner $definition) }
        if (Test-Path -LiteralPath $definition.result) {
            Write-Host "Last result for slot $($definition.number):"
            try { Write-Host ([IO.File]::ReadAllText($definition.result)) } catch { Write-Host 'Result is being updated.' }
        }
    }
    Write-Host "Capacity: $busyCount/2 slots occupied."
    exit 0
}
if (-not $Command -or -not $Command[0]) { throw 'Pass -Command as an array, for example @("pnpm", "typecheck").' }
$directory = (Resolve-Path -LiteralPath $WorkingDirectory).Path
if (-not (Test-Path -LiteralPath $directory -PathType Container)) { throw 'WorkingDirectory must be a directory.' }
if ($env:CHONGCHONG_LOCAL_CHECK_OWNER) {
    # 拒绝嵌套申请，避免外层等待内层、内层等待同一锁。
    throw 'Already inside local-check. Run the original command directly inside this checked batch.'
}
$timer = [Diagnostics.Stopwatch]::StartNew()
$slot = $null
$job = $null
$child = $null
$gatePath = Join-Path $stateDirectory ("gate-{0}.tmp" -f [guid]::NewGuid())
$owner = $null
$exitCode = 1
try {
    $nextNotice = 0
    while (-not ($lease = Open-AvailableSlot)) {
        if ($timer.Elapsed.TotalSeconds -ge $WaitSeconds) {
            Write-Host "QUEUE TIMEOUT after $WaitSeconds seconds. No check was started."
            exit 75
        }
        if ($timer.Elapsed.TotalSeconds -ge $nextNotice) {
            $owners = @($slotDefinitions | ForEach-Object { "slot $($_.number)/2: $(Read-Owner $_)" }) -join ' | '
            Write-Host "WAITING for either shared heavy check slot: $owners"
            $nextNotice = $timer.Elapsed.TotalSeconds + 15
        }
        Start-Sleep -Milliseconds 250
    }
    $slot = $lease.handle
    $ownerPath = $lease.definition.owner
    $resultPath = $lease.definition.result
    $owner = [ordered]@{ label = $Label; pid = $PID; slot = $lease.definition.number; directory = $directory; tool = $Command[0]; startedAt = [DateTimeOffset]::Now.ToString('o') }
    [IO.File]::WriteAllText($ownerPath, ($owner | ConvertTo-Json -Compress))
    Write-Host "RUNNING [$Label] in $directory (shared slot $($lease.definition.number)/2; timeout ${TimeoutSeconds}s)"

    # Job Object 在包装器被强制结束时也回收本次子进程树，不影响其他 Node/Java。
    if (-not ('LocalCheckJob' -as [type])) {
        Add-Type -TypeDefinition @'
using System;
using System.ComponentModel;
using System.Runtime.InteropServices;
public static class LocalCheckJob {
    [StructLayout(LayoutKind.Sequential)] struct BasicLimits {
        public long PerProcessUserTimeLimit, PerJobUserTimeLimit;
        public uint LimitFlags;
        public UIntPtr MinimumWorkingSetSize, MaximumWorkingSetSize;
        public uint ActiveProcessLimit;
        public UIntPtr Affinity;
        public uint PriorityClass, SchedulingClass;
    }
    [StructLayout(LayoutKind.Sequential)] struct IoCounters {
        public ulong ReadOperationCount, WriteOperationCount, OtherOperationCount;
        public ulong ReadTransferCount, WriteTransferCount, OtherTransferCount;
    }
    [StructLayout(LayoutKind.Sequential)] struct ExtendedLimits {
        public BasicLimits BasicLimitInformation;
        public IoCounters IoInfo;
        public UIntPtr ProcessMemoryLimit, JobMemoryLimit, PeakProcessMemoryUsed, PeakJobMemoryUsed;
    }
    [StructLayout(LayoutKind.Sequential)] struct Accounting {
        public long UserTime, KernelTime, PeriodUserTime, PeriodKernelTime;
        public uint PageFaults, TotalProcesses, ActiveProcesses, TerminatedProcesses;
    }
    [DllImport("kernel32.dll", CharSet=CharSet.Unicode, SetLastError=true)] static extern IntPtr CreateJobObject(IntPtr attributes, string name);
    [DllImport("kernel32.dll", SetLastError=true)] static extern bool SetInformationJobObject(IntPtr job, int infoClass, ref ExtendedLimits info, uint length);
    [DllImport("kernel32.dll", SetLastError=true)] static extern bool AssignProcessToJobObject(IntPtr job, IntPtr process);
    [DllImport("kernel32.dll", SetLastError=true)] static extern bool TerminateJobObject(IntPtr job, uint code);
    [DllImport("kernel32.dll", SetLastError=true)] static extern bool QueryInformationJobObject(IntPtr job, int infoClass, out Accounting info, uint length, IntPtr returnedLength);
    [DllImport("kernel32.dll")] public static extern bool CloseHandle(IntPtr handle);
    public static IntPtr Create() {
        var job = CreateJobObject(IntPtr.Zero, null);
        if (job == IntPtr.Zero) throw new Win32Exception();
        var limits = new ExtendedLimits();
        limits.BasicLimitInformation.LimitFlags = 0x2000; // KILL_ON_JOB_CLOSE
        if (!SetInformationJobObject(job, 9, ref limits, (uint)Marshal.SizeOf<ExtendedLimits>())) {
            int error = Marshal.GetLastWin32Error(); CloseHandle(job); throw new Win32Exception(error);
        }
        return job;
    }
    public static void Attach(IntPtr job, IntPtr process) {
        if (!AssignProcessToJobObject(job, process)) throw new Win32Exception();
    }
    public static void Stop(IntPtr job) {
        if (!TerminateJobObject(job, 124)) throw new Win32Exception();
        // 后代全部退出再归还执行名额，避免退出瞬间下一任务与遗留 worker 重叠。
        Accounting info;
        do {
            if (!QueryInformationJobObject(job, 1, out info, (uint)Marshal.SizeOf<Accounting>(), IntPtr.Zero)) throw new Win32Exception();
            if (info.ActiveProcesses != 0) System.Threading.Thread.Sleep(50);
        } while (info.ActiveProcesses != 0);
    }
}
'@
    }
    $job = [LocalCheckJob]::Create()
    $payload = @{ command = $Command; gate = $gatePath } | ConvertTo-Json -Compress
    $payload64 = [Convert]::ToBase64String([Text.Encoding]::UTF8.GetBytes($payload))
    $bootstrap = @'
$ErrorActionPreference = 'Stop'
$payload = [Text.Encoding]::UTF8.GetString([Convert]::FromBase64String($env:CHONGCHONG_LOCAL_CHECK_PAYLOAD)) | ConvertFrom-Json
while (-not (Test-Path -LiteralPath $payload.gate)) { Start-Sleep -Milliseconds 50 }
$env:CHONGCHONG_LOCAL_CHECK_PAYLOAD = $null
$commandArgs = @($payload.command | Select-Object -Skip 1)
$executable = [string]$payload.command[0]
if ([IO.Path]::GetFileNameWithoutExtension($executable) -in @('mvn', 'mvnw')) {
    if (@($commandArgs | Where-Object { $_ -match '^-T|^--threads|^-D(forkCount|reuseForks|parallel)=' }).Count) {
        throw 'Do not override Maven concurrency in local-check; -T 1 and forkCount=1 are supplied automatically.'
    }
    $commandArgs = @('-T', '1', '-DforkCount=1', '-DreuseForks=true') + $commandArgs
}
if ([IO.Path]::GetFileNameWithoutExtension($executable) -eq 'node' -and $commandArgs -contains '--test') {
    if (@($commandArgs | Where-Object { $_ -match '^--test-concurrency' }).Count) { throw 'local-check supplies --test-concurrency=2.' }
    $commandArgs = @('--test-concurrency=2') + $commandArgs
}
$global:LASTEXITCODE = 0
& $executable @commandArgs
if (-not $?) { if ($LASTEXITCODE) { exit $LASTEXITCODE }; exit 1 }
exit $LASTEXITCODE
'@
    $startInfo = [Diagnostics.ProcessStartInfo]::new()
    $startInfo.FileName = (Get-Process -Id $PID).Path
    foreach ($arg in @('-NoLogo', '-NoProfile', '-NonInteractive', '-OutputFormat', 'Text', '-EncodedCommand', [Convert]::ToBase64String([Text.Encoding]::Unicode.GetBytes($bootstrap)))) { $startInfo.ArgumentList.Add($arg) }
    $startInfo.WorkingDirectory = $directory
    $startInfo.UseShellExecute = $false
    $startInfo.Environment['CHONGCHONG_LOCAL_CHECK_OWNER'] = [string]$PID
    $startInfo.Environment['CHONGCHONG_LOCAL_CHECK_PAYLOAD'] = $payload64
    $startInfo.Environment['VITEST_MAX_WORKERS'] = '2'
    $startInfo.Environment['NODE_OPTIONS'] = (($env:NODE_OPTIONS -replace '--max-old-space-size(?:=|\s+)\S+', '') + ' --max-old-space-size=2048').Trim()
    $startInfo.Environment['JAVA_TOOL_OPTIONS'] = (($env:JAVA_TOOL_OPTIONS -replace '-XX:ActiveProcessorCount=\S+|-Xmx\S+', '') + ' -XX:ActiveProcessorCount=2 -Xmx1536m').Trim()
    # Maven 启动器可能在 JAVA_TOOL_OPTIONS 之后应用 MAVEN_OPTS，末尾再限定一次。
    $startInfo.Environment['MAVEN_OPTS'] = (($env:MAVEN_OPTS -replace '-XX:ActiveProcessorCount=\S+|-Xmx\S+', '') + ' -XX:ActiveProcessorCount=2 -Xmx1536m').Trim()
    $child = [Diagnostics.Process]::Start($startInfo)
    [LocalCheckJob]::Attach($job, $child.Handle)
    # 先纳入 Job 再放行，防止短命包装器提前产生未受管理的后代。
    [IO.File]::WriteAllText($gatePath, '')
    $running = [Diagnostics.Stopwatch]::StartNew()
    while (-not $child.WaitForExit(250)) {
        if ($running.Elapsed.TotalSeconds -ge $TimeoutSeconds) {
            Write-Host "EXECUTION TIMEOUT [$Label]; stopping only this check's process tree."
            $exitCode = 124
            break
        }
    }
    if ($exitCode -ne 124) { $exitCode = $child.ExitCode }
} finally {
    if ($job) {
        try { [LocalCheckJob]::Stop($job) }
        finally { [LocalCheckJob]::CloseHandle($job) | Out-Null }
    }
    if ($child) {
        if (-not $child.HasExited) { $child.Kill($true); $child.WaitForExit() }
        $child.Dispose()
    }
    if (Test-Path -LiteralPath $gatePath) { Remove-Item -LiteralPath $gatePath -Force }
    if ($slot) {
        try {
            if ($owner) {
                $owner['finishedAt'] = [DateTimeOffset]::Now.ToString('o')
                $owner['exitCode'] = $exitCode
                [IO.File]::WriteAllText($resultPath, ($owner | ConvertTo-Json -Compress))
            }
        } finally { $slot.Dispose() }
    }
}
Write-Host "FINISHED [$Label] exit=$exitCode"
exit $exitCode