|
| 1 | +# Usage: .buildkite/scripts/run-script.ps1 <script-or-simple-command> |
| 2 | +# Example: .buildkite/scripts/run-script.ps1 bash .buildkite/scripts/tests.sh |
| 3 | +# Example: .buildkite/scripts/run-script.ps1 .buildkite/scripts/other-tests.ps1 |
| 4 | +# |
| 5 | +# NOTE: Apparently passing arguments in powershell is a nightmare, so you shouldn't do it unless it's really simple. Just use the wrapper to call a script instead. |
| 6 | +# See: https://stackoverflow.com/questions/6714165/powershell-stripping-double-quotes-from-command-line-arguments |
| 7 | +# and: https://github.com/PowerShell/PowerShell/issues/3223#issuecomment-487975049 |
| 8 | +# |
| 9 | +# See here: https://github.com/buildkite/agent/issues/2202 |
| 10 | +# Background processes after the buildkite step script finishes causes the job to hang. |
| 11 | +# So, until this is fixed/changed in buildkite-agent (if ever), we can use this wrapper. |
| 12 | + |
| 13 | +# This wrapper: |
| 14 | +# - Creates a Windows job object (which is like a process group) |
| 15 | +# - Sets JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE, which means that when the job object is closed, all processes in the job object are killed |
| 16 | +# - Starts running your given script, and assigns it to the job object |
| 17 | +# - Now, any child processes created by your script will also end up in the job object |
| 18 | +# - Waits for your script (and only your script, not child processes) to finish |
| 19 | +# - Closes the job object, which kills all processes in the job object (including any leftover child processes) |
| 20 | +# - Exits with the exit status from your script |
| 21 | + |
| 22 | +Add-Type -TypeDefinition @' |
| 23 | +using Microsoft.Win32.SafeHandles; |
| 24 | +using System; |
| 25 | +using System.ComponentModel; |
| 26 | +using System.Runtime.InteropServices; |
| 27 | +
|
| 28 | +public class NativeMethods |
| 29 | +{ |
| 30 | + public enum JOBOBJECTINFOCLASS |
| 31 | + { |
| 32 | + AssociateCompletionPortInformation = 7, |
| 33 | + BasicLimitInformation = 2, |
| 34 | + BasicUIRestrictions = 4, |
| 35 | + EndOfJobTimeInformation = 6, |
| 36 | + ExtendedLimitInformation = 9, |
| 37 | + SecurityLimitInformation = 5, |
| 38 | + GroupInformation = 11 |
| 39 | + } |
| 40 | +
|
| 41 | + [StructLayout(LayoutKind.Sequential)] |
| 42 | + struct JOBOBJECT_BASIC_LIMIT_INFORMATION |
| 43 | + { |
| 44 | + public Int64 PerProcessUserTimeLimit; |
| 45 | + public Int64 PerJobUserTimeLimit; |
| 46 | + public UInt32 LimitFlags; |
| 47 | + public UIntPtr MinimumWorkingSetSize; |
| 48 | + public UIntPtr MaximumWorkingSetSize; |
| 49 | + public UInt32 ActiveProcessLimit; |
| 50 | + public Int64 Affinity; |
| 51 | + public UInt32 PriorityClass; |
| 52 | + public UInt32 SchedulingClass; |
| 53 | + } |
| 54 | +
|
| 55 | + [StructLayout(LayoutKind.Sequential)] |
| 56 | + struct IO_COUNTERS |
| 57 | + { |
| 58 | + public UInt64 ReadOperationCount; |
| 59 | + public UInt64 WriteOperationCount; |
| 60 | + public UInt64 OtherOperationCount; |
| 61 | + public UInt64 ReadTransferCount; |
| 62 | + public UInt64 WriteTransferCount; |
| 63 | + public UInt64 OtherTransferCount; |
| 64 | + } |
| 65 | +
|
| 66 | + [StructLayout(LayoutKind.Sequential)] |
| 67 | + struct JOBOBJECT_EXTENDED_LIMIT_INFORMATION |
| 68 | + { |
| 69 | + public JOBOBJECT_BASIC_LIMIT_INFORMATION BasicLimitInformation; |
| 70 | + public IO_COUNTERS IoInfo; |
| 71 | + public UIntPtr ProcessMemoryLimit; |
| 72 | + public UIntPtr JobMemoryLimit; |
| 73 | + public UIntPtr PeakProcessMemoryUsed; |
| 74 | + public UIntPtr PeakJobMemoryUsed; |
| 75 | + } |
| 76 | +
|
| 77 | + [DllImport("Kernel32.dll", EntryPoint = "AssignProcessToJobObject", SetLastError = true)] |
| 78 | + private static extern bool NativeAssignProcessToJobObject(SafeHandle hJob, SafeHandle hProcess); |
| 79 | +
|
| 80 | + public static void AssignProcessToJobObject(SafeHandle job, SafeHandle process) |
| 81 | + { |
| 82 | + if (!NativeAssignProcessToJobObject(job, process)) |
| 83 | + throw new Win32Exception(); |
| 84 | + } |
| 85 | +
|
| 86 | + [DllImport( |
| 87 | + "Kernel32.dll", |
| 88 | + CharSet = CharSet.Unicode, |
| 89 | + EntryPoint = "CreateJobObjectW", |
| 90 | + SetLastError = true |
| 91 | + )] |
| 92 | + private static extern SafeFileHandle NativeCreateJobObjectW( |
| 93 | + IntPtr lpJobAttributes, |
| 94 | + string lpName |
| 95 | + ); |
| 96 | +
|
| 97 | + [DllImport("Kernel32.dll", EntryPoint = "CloseHandle", SetLastError = true)] |
| 98 | + private static extern bool NativeCloseHandle(SafeHandle hJob); |
| 99 | +
|
| 100 | + [DllImport("kernel32.dll")] |
| 101 | + public static extern bool SetInformationJobObject( |
| 102 | + SafeHandle hJob, |
| 103 | + JOBOBJECTINFOCLASS JobObjectInfoClass, |
| 104 | + IntPtr lpJobObjectInfo, |
| 105 | + uint cbJobObjectInfoLength |
| 106 | + ); |
| 107 | +
|
| 108 | + private const UInt32 JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE = 0x2000; |
| 109 | +
|
| 110 | + public static SafeHandle CreateJobObjectW(string name) |
| 111 | + { |
| 112 | + SafeHandle job = NativeCreateJobObjectW(IntPtr.Zero, name); |
| 113 | + JOBOBJECT_BASIC_LIMIT_INFORMATION info = new JOBOBJECT_BASIC_LIMIT_INFORMATION(); |
| 114 | + info.LimitFlags = JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE; |
| 115 | + JOBOBJECT_EXTENDED_LIMIT_INFORMATION extendedInfo = |
| 116 | + new JOBOBJECT_EXTENDED_LIMIT_INFORMATION(); |
| 117 | + extendedInfo.BasicLimitInformation = info; |
| 118 | + int length = Marshal.SizeOf(typeof(JOBOBJECT_EXTENDED_LIMIT_INFORMATION)); |
| 119 | + IntPtr extendedInfoPtr = Marshal.AllocHGlobal(length); |
| 120 | + Marshal.StructureToPtr(extendedInfo, extendedInfoPtr, false); |
| 121 | + SetInformationJobObject( |
| 122 | + job, |
| 123 | + JOBOBJECTINFOCLASS.ExtendedLimitInformation, |
| 124 | + extendedInfoPtr, |
| 125 | + (uint)length |
| 126 | + ); |
| 127 | + if (job.IsInvalid) |
| 128 | + throw new Win32Exception(); |
| 129 | + return job; |
| 130 | + } |
| 131 | +
|
| 132 | + public static void CloseHandle(SafeHandle job) |
| 133 | + { |
| 134 | + if (!NativeCloseHandle(job)) |
| 135 | + throw new Win32Exception(); |
| 136 | + } |
| 137 | +} |
| 138 | +
|
| 139 | +'@ |
| 140 | + |
| 141 | +$guid = [guid]::NewGuid().Guid |
| 142 | +Write-Output "Creating job object with name $guid" |
| 143 | +$job = [NativeMethods]::CreateJobObjectW($guid) |
| 144 | +$process = Start-Process -PassThru -NoNewWindow powershell -ArgumentList "$args" |
| 145 | +[NativeMethods]::AssignProcessToJobObject($job, $process.SafeHandle) |
| 146 | + |
| 147 | +try { |
| 148 | + Write-Output "Waiting for process $($process.Id) to complete..." |
| 149 | + $process | Wait-Process |
| 150 | + Write-Output "Process finished with exit code $($process.ExitCode), terminating job and exiting..." |
| 151 | +} finally { |
| 152 | + [NativeMethods]::CloseHandle($job) |
| 153 | + exit $process.ExitCode |
| 154 | +} |
0 commit comments