From be1c808fcad201adc4d5d6cca52ddb24aeb5e367 Mon Sep 17 00:00:00 2001 From: vxiiduu <73044267+vxiiduu@users.noreply.github.com> Date: Fri, 15 Mar 2024 09:09:36 +1000 Subject: gh-116195: Implements a fast path for nt.getppid (GH-116205) Use the NtQueryInformationProcess system call to efficiently retrieve the parent process ID in a single step, rather than using the process snapshots API which retrieves large amounts of unnecessary information and is more prone to failure (since it makes heap allocations). Includes a fallback to the original win32_getppid implementation in case the unstable API appears to return strange results. --- .../2024-03-14-20-46-23.gh-issue-116195.Cu_rYs.rst | 1 + Modules/posixmodule.c | 86 +++++++++++++++++++++- 2 files changed, 85 insertions(+), 2 deletions(-) create mode 100644 Misc/NEWS.d/next/Windows/2024-03-14-20-46-23.gh-issue-116195.Cu_rYs.rst diff --git a/Misc/NEWS.d/next/Windows/2024-03-14-20-46-23.gh-issue-116195.Cu_rYs.rst b/Misc/NEWS.d/next/Windows/2024-03-14-20-46-23.gh-issue-116195.Cu_rYs.rst new file mode 100644 index 0000000..32122d7 --- /dev/null +++ b/Misc/NEWS.d/next/Windows/2024-03-14-20-46-23.gh-issue-116195.Cu_rYs.rst @@ -0,0 +1 @@ +Improves performance of :func:`os.getppid` by using an alternate system API when available. Contributed by vxiiduu. diff --git a/Modules/posixmodule.c b/Modules/posixmodule.c index 19e9257..7b2d366 100644 --- a/Modules/posixmodule.c +++ b/Modules/posixmodule.c @@ -9115,7 +9115,81 @@ os_setpgrp_impl(PyObject *module) #ifdef HAVE_GETPPID #ifdef MS_WINDOWS -#include +#include +#include + +// The structure definition in winternl.h may be incomplete. +// This structure is the full version from the MSDN documentation. +typedef struct _PROCESS_BASIC_INFORMATION_FULL { + NTSTATUS ExitStatus; + PVOID PebBaseAddress; + ULONG_PTR AffinityMask; + LONG BasePriority; + ULONG_PTR UniqueProcessId; + ULONG_PTR InheritedFromUniqueProcessId; +} PROCESS_BASIC_INFORMATION_FULL; + +typedef NTSTATUS (NTAPI *PNT_QUERY_INFORMATION_PROCESS) ( + IN HANDLE ProcessHandle, + IN PROCESSINFOCLASS ProcessInformationClass, + OUT PVOID ProcessInformation, + IN ULONG ProcessInformationLength, + OUT PULONG ReturnLength OPTIONAL); + +// This function returns the process ID of the parent process. +// Returns 0 on failure. +static ULONG +win32_getppid_fast(void) +{ + NTSTATUS status; + HMODULE ntdll; + PNT_QUERY_INFORMATION_PROCESS pNtQueryInformationProcess; + PROCESS_BASIC_INFORMATION_FULL basic_information; + static ULONG cached_ppid = 0; + + if (cached_ppid) { + // No need to query the kernel again. + return cached_ppid; + } + + ntdll = GetModuleHandleW(L"ntdll.dll"); + if (!ntdll) { + return 0; + } + + pNtQueryInformationProcess = (PNT_QUERY_INFORMATION_PROCESS) GetProcAddress(ntdll, "NtQueryInformationProcess"); + if (!pNtQueryInformationProcess) { + return 0; + } + + status = pNtQueryInformationProcess(GetCurrentProcess(), + ProcessBasicInformation, + &basic_information, + sizeof(basic_information), + NULL); + + if (!NT_SUCCESS(status)) { + return 0; + } + + // Perform sanity check on the parent process ID we received from NtQueryInformationProcess. + // The check covers values which exceed the 32-bit range (if running on x64) as well as + // zero and (ULONG) -1. + + if (basic_information.InheritedFromUniqueProcessId == 0 || + basic_information.InheritedFromUniqueProcessId >= ULONG_MAX) + { + return 0; + } + + // Now that we have reached this point, the BasicInformation.InheritedFromUniqueProcessId + // structure member contains a ULONG_PTR which represents the process ID of our parent + // process. This process ID will be correctly returned even if the parent process has + // exited or been terminated. + + cached_ppid = (ULONG) basic_information.InheritedFromUniqueProcessId; + return cached_ppid; +} static PyObject* win32_getppid(void) @@ -9123,8 +9197,16 @@ win32_getppid(void) DWORD error; PyObject* result = NULL; HANDLE process = GetCurrentProcess(); - HPSS snapshot = NULL; + ULONG pid; + + pid = win32_getppid_fast(); + if (pid != 0) { + return PyLong_FromUnsignedLong(pid); + } + + // If failure occurs in win32_getppid_fast(), fall back to using the PSS API. + error = PssCaptureSnapshot(process, PSS_CAPTURE_NONE, 0, &snapshot); if (error != ERROR_SUCCESS) { return PyErr_SetFromWindowsErr(error); -- cgit v0.12