How can a Windows process detect that it will soon hit its memory limit?

Question

If a process' memory is being limited by JOBOBJECT_EXTENDED_LIMIT_INFORMATION. ProcessMemoryLimit ("...specifies the limit for the virtual memory that can be committed by a process..."), how can the process detect when it is getting close to this limit? The most obvious (to me) is to periodically check Process.VirtualMemorySize64 ("The amount of virtual memory, in bytes, allocated for the associated process"). Is this the right measure? There is also WorkingSet64, PrivateMemorySize64, PagedSystemMemorySize64, PagedMemorySize64 and NonPagedSystemMemorySize64.

Answer 1

@Ben really answered this in a comment: you can use SetInformationJobObject (overview in Job Objects). Here is my code, which extends this question to provide a class that can be used to limit a process' memory, and to provide a callback at a nominated threshold:

#nullable enable
using System;
using System.Collections.Generic;
using System.ComponentModel;
using System.Runtime.InteropServices;
using System.Threading;
using CommonLib;
using QWeb.QWebLib;

namespace QWeb.QhostLib
{
    /// <summary>
    /// A Class to enable the use of Win32 [Job Objects](https://learn.microsoft.com/en-us/windows/win32/procthread/job-objects)
    /// taken from: https://stackoverflow.com/a/9164742/5626740
    ///
    /// 1. Create a MemoryLimitedJobObject, specifying the memory limit (CreateJobObject)
    /// 2. Add the process to be limited to the job (AddProcess)
    /// 3. [optional] TellMeWhenMemoryUsedExceeds() to register a callback for when memory
    ///    use exceeds some threshold.
    /// </summary>
    public unsafe class MemoryLimitedJobObject : IDisposable
    {
        [DllImport("kernel32.dll", CharSet = CharSet.Unicode)]
        static extern IntPtr CreateJobObject(IntPtr a, string? lp_name);

        [DllImport("kernel32.dll", SetLastError = true)]
        static extern bool SetInformationJobObject(IntPtr h_job, JobObjectInfoType info_type, void *lp_job_object_info, int cb_job_object_info_length);

        [DllImport("kernel32.dll", SetLastError = true)]
        static extern bool AssignProcessToJobObject(IntPtr job, IntPtr process);

        static readonly ListenToJobLimitViolations listenToJobLimitViolations = new ListenToJobLimitViolations();

        IntPtr jobHandle;
        bool disposed;

        /// <summary>Creates a Job Object with limited memory, processes can later be added to this Job Object</summary>
        /// <param name="memory_limit_bytes">The limit to be set on the Job Object</param>
        public MemoryLimitedJobObject(ulong memory_limit_bytes)
        {
            this.jobHandle = CreateJobObject(IntPtr.Zero, null);

            var extended_info = new JOBOBJECT_EXTENDED_LIMIT_INFORMATION {
                BasicLimitInformation = new JOBOBJECT_BASIC_LIMIT_INFORMATION {
                    LimitFlags = 0x0100 // JOB_OBJECT_LIMIT_PROCESS_MEMORY
                },
                ProcessMemoryLimit = new UIntPtr(memory_limit_bytes)
            };
            if (!SetInformationJobObject(jobHandle, JobObjectInfoType.ExtendedLimitInformation, &extended_info, sizeof(JOBOBJECT_EXTENDED_LIMIT_INFORMATION)))
                throw new ApplicationException($"Unable to set information.  Error: {Marshal.GetLastWin32Error()}");
        }

        public void Dispose()
        {
            Dispose(true);
            GC.SuppressFinalize(this);
        }

        private void Dispose(bool disposing)
        {
            if (disposed)
                return;

            if (disposing) { }

            Close();
            disposed = true;
        }

        public void Close()
        {
            listenToJobLimitViolations.Forget(jobHandle);
            Win32.CloseHandle(jobHandle);
            this.jobHandle = IntPtr.Zero;
        }

        /// <summary>As soon as the process hits memory_report_trigger_bytes of RAM we will
        /// call memory_trigger_exceeded with the number of bytes in use.  This will not be called
        /// again.</summary>
        public void TellMeWhenMemoryUsedExceeds(ulong memory_report_trigger_bytes, Action<ulong> memory_trigger_exceeded)
        {
            // Tell Windows to notify us when this job exceeds its limit.
            var limit_info = new JOBOBJECT_NOTIFICATION_LIMIT_INFORMATION {
                JobMemoryLimit = memory_report_trigger_bytes,
                LimitFlags = 0x00000200 // JOB_OBJECT_LIMIT_JOB_MEMORY
            };
            if (!SetInformationJobObject(jobHandle, JobObjectInfoType.NotificationLimitInformation, &limit_info, sizeof(JOBOBJECT_NOTIFICATION_LIMIT_INFORMATION)))
                throw new Win32Exception(Marshal.GetLastWin32Error());

            // Tell Windows that any asynchronous notifications from this job should come to the completion
            // port we created earlier.
            var async_completion_port_handle = listenToJobLimitViolations.KnowAbout(jobHandle, memory_trigger_exceeded);
            var port_info = new JOBOBJECT_ASSOCIATE_COMPLETION_PORT {
                CompletionKey = jobHandle,  // so we can tell 
                CompletionPort = async_completion_port_handle
            };
            if (!SetInformationJobObject(jobHandle, JobObjectInfoType.AssociateCompletionPortInformation, &port_info, sizeof(JOBOBJECT_ASSOCIATE_COMPLETION_PORT)))
                throw new Win32Exception(Marshal.GetLastWin32Error());
        }

        public bool AddProcess(IntPtr process_handle)
        {
            return AssignProcessToJobObject(jobHandle, process_handle);
        }
    }

    /// <summary>
    /// Runs a Thread that is shared between all MemoryLimitedJobObjects.  It gets notifications from Windows
    /// when a process exceeds its memory threshold, which lets us warn people that the process probably
    /// needs more space.
    /// </summary>
    public class ListenToJobLimitViolations
    {
        [DllImport("kernel32.dll", SetLastError = true)]
        static extern IntPtr CreateIoCompletionPort(IntPtr file_handle, IntPtr existing_completion_port, UIntPtr completion_key, uint number_of_concurrent_threads);

        [DllImport("kernel32.dll", SetLastError = true)]
        static extern unsafe bool GetQueuedCompletionStatus(IntPtr completion_port, out uint number_of_bytes, out IntPtr completion_key, void* overlapped, uint milliseconds);

        [DllImport("kernel32.dll", SetLastError = true)]
        static extern bool QueryInformationJobObject(IntPtr job_handle, JobObjectInfoType job_object_information_class, out JOBOBJECT_LIMIT_VIOLATION_INFORMATION job_object_information, int job_object_information_length, out uint return_length);

        // Extracted from winnt.h in Windows SDK
        public const int JOB_OBJECT_MSG_NOTIFICATION_LIMIT = 11;

        readonly IntPtr asyncCompletionPortHandle;
        readonly Dictionary<IntPtr, Action<ulong>> jobsRegisteredForNotification = new();

        internal ListenToJobLimitViolations()
        {
            // Create a asynchronous completion port to receive events for this job, and start
            // monitoring it.
            this.asyncCompletionPortHandle = CreateIoCompletionPort(Win32.INVALID_HANDLE_VALUE, IntPtr.Zero, UIntPtr.Zero, 0);
            if (asyncCompletionPortHandle == IntPtr.Zero)
                throw new Win32Exception(Marshal.GetLastWin32Error());
            new Thread(() => {
                try {
                    MonitorCompletionPort(asyncCompletionPortHandle);
                } catch (Exception ex) {
                    Error.Report(nameof(MonitorCompletionPort), ex);
                }
            }) {
                IsBackground = true,
                Name = "ListenToJobLimitViolations"
            }.Start();
        }

        unsafe void MonitorCompletionPort(IntPtr async_completion_port_handle)
        {
            while (true) {
                NativeOverlapped* native_overlapped;  // documented as being set to whatever was provided when the async operation was
                                                      // started, which doesn't apply here.  So I figure it's not our job to dispose.
                if (!GetQueuedCompletionStatus(async_completion_port_handle, out var bytes_read, out var job_handle, &native_overlapped, uint.MaxValue))
                    throw new Win32Exception(Marshal.GetLastWin32Error());
                // Identify what sort of message we're getting.  winnt.h says
                // "These values are returned via the lpNumberOfBytesTransferred parameter"!
                if (bytes_read == JOB_OBJECT_MSG_NOTIFICATION_LIMIT) {
                    if (!QueryInformationJobObject(job_handle, JobObjectInfoType.LimitViolationInformation, out var limit_violation, sizeof(JOBOBJECT_LIMIT_VIOLATION_INFORMATION), out var _))
                        throw new Win32Exception(Marshal.GetLastWin32Error());
                    Action<ulong> trigger_action;
                    lock (jobsRegisteredForNotification) {
                        trigger_action = jobsRegisteredForNotification[job_handle];
                    }
                    trigger_action(limit_violation.JobMemory);
                }
            }
        }

        internal IntPtr KnowAbout(IntPtr job_handle, Action<ulong> memory_trigger_exceeded)
        {
            lock (jobsRegisteredForNotification) {
                jobsRegisteredForNotification.Add(job_handle, memory_trigger_exceeded);
                return asyncCompletionPortHandle;  // caller must register their Job with this, to direct notifications here
            }
        }

        internal void Forget(IntPtr job_handle)
        {
            lock (jobsRegisteredForNotification) {
                jobsRegisteredForNotification.Remove(job_handle);
            }
        }
    }

    #region Helper classes

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

    [StructLayout(LayoutKind.Sequential)]
    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)]
    public struct SECURITY_ATTRIBUTES
    {
#pragma warning disable IDE1006 // Naming Styles
        public uint nLength;
        public IntPtr lpSecurityDescriptor;
        public int bInheritHandle;
#pragma warning restore IDE1006 // Naming Styles
    }

    [StructLayout(LayoutKind.Sequential)]
    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;
    }

    [StructLayout(LayoutKind.Sequential)]
    struct JOBOBJECT_NOTIFICATION_LIMIT_INFORMATION
    {
        public ulong IoReadBytesLimit;
        public ulong IoWriteBytesLimit;
        public long PerJobUserTimeLimit;
        public ulong JobMemoryLimit;
        public int RateControlTolerance;
        public int RateControlToleranceInterval;
        public uint LimitFlags;
    }

    [StructLayout(LayoutKind.Sequential)]
    struct JOBOBJECT_ASSOCIATE_COMPLETION_PORT
    {
        public IntPtr CompletionKey;
        public IntPtr CompletionPort;
    }

    [StructLayout(LayoutKind.Sequential)]
    struct JOBOBJECT_LIMIT_VIOLATION_INFORMATION
    {
        public uint LimitFlags;
        public uint ViolationLimitFlags;
        public ulong IoReadBytes;
        public ulong IoReadBytesLimit;
        public ulong IoWriteBytes;
        public ulong IoWriteBytesLimit;
        public long PerJobUserTime;
        public long PerJobUserTimeLimit;
        public ulong JobMemory;
        public ulong JobMemoryLimit;
        public int RateControlTolerance;
        public int RateControlToleranceLimit;
    }

    public enum JobObjectInfoType
    {
        AssociateCompletionPortInformation = 7,
        ExtendedLimitInformation = 9,
        NotificationLimitInformation = 12,
        LimitViolationInformation = 13
    }
    #endregion
}