lmxopcua/src/Core/ZB.MOM.WW.OtOpcUa.Core.Scripting.OutOfProcess/WindowsJobObject.cs

using System.ComponentModel;
using System.Diagnostics;
using System.Runtime.InteropServices;
using System.Runtime.Versioning;
using Microsoft.Win32.SafeHandles;

namespace ZB.MOM.WW.OtOpcUa.Core.Scripting.OutOfProcess;

/// <summary>
///     A Windows Job Object holding one ScriptWorker (design §3.3, §5.2): <c>KILL_ON_JOB_CLOSE</c> so the kernel kills
///     the worker when the host dies (its handle closes with it), <c>ACTIVE_PROCESS = 1</c> so the worker can start no
///     child process, and a process-memory backstop slightly above the managed heap limit for native and JIT memory.
/// </summary>
[SupportedOSPlatform("windows")]
internal sealed class WindowsJobObject : IDisposable
{
    private const int JobObjectExtendedLimitInformation = 9;
    private const uint JOB_OBJECT_LIMIT_ACTIVE_PROCESS = 0x00000008;
    private const uint JOB_OBJECT_LIMIT_PROCESS_MEMORY = 0x00000100;
    private const uint JOB_OBJECT_LIMIT_DIE_ON_UNHANDLED_EXCEPTION = 0x00000400;
    private const uint JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE = 0x00002000;

    private readonly SafeFileHandle _handle;

    private WindowsJobObject(SafeFileHandle handle) => _handle = handle;

    /// <summary>Creates a job with the worker limits and assigns <paramref name="process"/> to it.</summary>
    /// <param name="process">The freshly started worker.</param>
    /// <param name="processMemoryLimitBytes">The per-process commit limit.</param>
    /// <returns>The job; dispose it to kill the worker.</returns>
    public static WindowsJobObject Assign(Process process, long processMemoryLimitBytes)
        => Assign(process.Handle, processMemoryLimitBytes);

    /// <summary>
    ///     Creates a job with the worker limits and assigns the process behind <paramref name="processHandle"/> to it — the
    ///     restricted launch assigns its worker while it is still suspended (<see cref="RestrictedWorkerLauncher"/>).
    /// </summary>
    /// <param name="processHandle">A handle with <c>PROCESS_SET_QUOTA</c> and <c>PROCESS_TERMINATE</c>.</param>
    /// <param name="processMemoryLimitBytes">The per-process commit limit.</param>
    /// <returns>The job; dispose it to kill the worker.</returns>
    public static WindowsJobObject Assign(IntPtr processHandle, long processMemoryLimitBytes)
    {
        var handle = CreateJobObjectW(IntPtr.Zero, null);
        if (handle.IsInvalid) throw new Win32Exception(Marshal.GetLastWin32Error(), "CreateJobObject failed");
        var job = new WindowsJobObject(handle);
        try
        {
            var info = new JOBOBJECT_EXTENDED_LIMIT_INFORMATION
            {
                BasicLimitInformation = new JOBOBJECT_BASIC_LIMIT_INFORMATION
                {
                    LimitFlags = JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE | JOB_OBJECT_LIMIT_ACTIVE_PROCESS
                                 | JOB_OBJECT_LIMIT_PROCESS_MEMORY | JOB_OBJECT_LIMIT_DIE_ON_UNHANDLED_EXCEPTION,
                    ActiveProcessLimit = 1,
                },
                ProcessMemoryLimit = (UIntPtr)(ulong)processMemoryLimitBytes,
            };
            if (!SetInformationJobObject(handle, JobObjectExtendedLimitInformation, ref info,
                    (uint)Marshal.SizeOf<JOBOBJECT_EXTENDED_LIMIT_INFORMATION>()))
                throw new Win32Exception(Marshal.GetLastWin32Error(), "SetInformationJobObject failed");
            if (!AssignProcessToJobObject(handle, processHandle))
                throw new Win32Exception(Marshal.GetLastWin32Error(), "AssignProcessToJobObject failed");
            return job;
        }
        catch
        {
            job.Dispose();
            throw;
        }
    }

    /// <summary>
    ///     Kills every process in the job at once (lmxopcua#789). The job admits exactly one process, so this is the
    ///     worker's whole tree without <c>Process.Kill(entireProcessTree: true)</c>'s system-wide process snapshot,
    ///     which cost about 100–200 ms of the Windows kill latency.
    /// </summary>
    /// <returns>False when the job could not be terminated (the caller falls back to a process kill).</returns>
    public bool Terminate() => !_handle.IsClosed && TerminateJobObject(_handle, unchecked((uint)-1));

    /// <summary>Closes the job handle, which kills every process still in it.</summary>
    public void Dispose() => _handle.Dispose();

    [StructLayout(LayoutKind.Sequential)]
    private struct JOBOBJECT_BASIC_LIMIT_INFORMATION
    {
        public long PerProcessUserTimeLimit;
        public long PerJobUserTimeLimit;
        public uint LimitFlags;
        public UIntPtr MinimumWorkingSetSize;
        public UIntPtr MaximumWorkingSetSize;
        public uint ActiveProcessLimit;
        public UIntPtr Affinity;
        public uint PriorityClass;
        public uint SchedulingClass;
    }

    [StructLayout(LayoutKind.Sequential)]
    private struct IO_COUNTERS
    {
        public ulong ReadOperationCount;
        public ulong WriteOperationCount;
        public ulong OtherOperationCount;
        public ulong ReadTransferCount;
        public ulong WriteTransferCount;
        public ulong OtherTransferCount;
    }

    [StructLayout(LayoutKind.Sequential)]
    private struct JOBOBJECT_EXTENDED_LIMIT_INFORMATION
    {
        public JOBOBJECT_BASIC_LIMIT_INFORMATION BasicLimitInformation;
        public IO_COUNTERS IoInfo;
        public UIntPtr ProcessMemoryLimit;
        public UIntPtr JobMemoryLimit;
        public UIntPtr PeakProcessMemoryUsed;
        public UIntPtr PeakJobMemoryUsed;
    }

#pragma warning disable SYSLIB1054 // DllImport kept: LibraryImport would need AllowUnsafeBlocks for four calls.
    [DllImport("kernel32.dll", SetLastError = true, CharSet = CharSet.Unicode)]
    private static extern SafeFileHandle CreateJobObjectW(IntPtr securityAttributes, string? name);

    [DllImport("kernel32.dll", SetLastError = true)]
    [return: MarshalAs(UnmanagedType.Bool)]
    private static extern bool SetInformationJobObject(SafeFileHandle job, int infoClass,
        ref JOBOBJECT_EXTENDED_LIMIT_INFORMATION info, uint length);

    [DllImport("kernel32.dll", SetLastError = true)]
    [return: MarshalAs(UnmanagedType.Bool)]
    private static extern bool AssignProcessToJobObject(SafeFileHandle job, IntPtr process);

    [DllImport("kernel32.dll", SetLastError = true)]
    [return: MarshalAs(UnmanagedType.Bool)]
    private static extern bool TerminateJobObject(SafeFileHandle job, uint exitCode);
#pragma warning restore SYSLIB1054
}