From d8a4a79fe63da80160d36379c72e542efb692bf6 Mon Sep 17 00:00:00 2001 From: Rushikesh Sakharle Date: Tue, 19 May 2026 22:50:24 +0530 Subject: [PATCH] Add Windows process memory support --- README.md | 28 +++++ ps_mem.py | 363 +++++++++++++++++++++++++++++++++++++++++++++++++++--- 2 files changed, 376 insertions(+), 15 deletions(-) diff --git a/README.md b/README.md index 732893e..2848653 100644 --- a/README.md +++ b/README.md @@ -12,6 +12,34 @@ Install: are available for most distros. Also the ps_mem.py script can be run directly. +Windows: + +The script can also be run directly on Windows with Python. The Windows +implementation uses native PSAPI/kernel32 calls through `ctypes`, so no +third-party runtime package is required: + +```bat +python ps_mem.py +python ps_mem.py -p 1234 -d +``` + +Windows reports the current working set per process. Shared resident pages are +apportioned by the Windows page share count when the process can be queried. +If Windows denies that working-set query for a process, ps_mem falls back to +standard process memory counters and prints a best-effort total with a warning +because shared memory may be over-counted. The `-S/--swap` option is accepted +for compatibility on Windows, but it is ignored and no fake zero-valued swap +column is printed. The `-s/--split-args` option shows the full executable path +on Windows. + +To build a standalone Windows executable: + +```bat +python -m pip install pyinstaller +python -m PyInstaller --onefile ps_mem.py +dist\ps_mem.exe +``` + Usage: ``` diff --git a/ps_mem.py b/ps_mem.py index d03a987..c5c4c8a 100755 --- a/ps_mem.py +++ b/ps_mem.py @@ -74,6 +74,7 @@ # FreeBSD 8.0 supports up to a level of Linux 2.6.16 import argparse +import ctypes import errno import os import sys @@ -96,11 +97,98 @@ def std_exceptions(etype, value, tb): # Define some global variables # -PAGESIZE = os.sysconf("SC_PAGE_SIZE") / 1024 #KiB +IS_WINDOWS = os.name == 'nt' + +def page_size_kib(): + try: + return os.sysconf("SC_PAGE_SIZE") / 1024.0 #KiB + except (AttributeError, OSError, ValueError): + import mmap + return mmap.PAGESIZE / 1024.0 #KiB + +PAGESIZE = page_size_kib() our_pid = os.getpid() have_pss = 0 have_swap_pss = 0 +windows_used_fallback = False + +if IS_WINDOWS: + from ctypes import wintypes + + DWORD = wintypes.DWORD + BOOL = wintypes.BOOL + HANDLE = wintypes.HANDLE + SIZE_T = ctypes.c_size_t + ULONG_PTR = ctypes.c_size_t + + PROCESS_QUERY_INFORMATION = 0x0400 + PROCESS_QUERY_LIMITED_INFORMATION = 0x1000 + PROCESS_VM_READ = 0x0010 + ERROR_BAD_LENGTH = 24 + ERROR_INSUFFICIENT_BUFFER = 122 + + def load_windows_dll(name): + try: + return ctypes.WinDLL(name, use_last_error=True) + except TypeError: + return ctypes.WinDLL(name) + + def get_last_windows_error(): + if hasattr(ctypes, 'get_last_error'): + return ctypes.get_last_error() + return ctypes.GetLastError() + + kernel32 = load_windows_dll('kernel32') + psapi = load_windows_dll('psapi') + + kernel32.OpenProcess.argtypes = [DWORD, BOOL, DWORD] + kernel32.OpenProcess.restype = HANDLE + kernel32.CloseHandle.argtypes = [HANDLE] + kernel32.CloseHandle.restype = BOOL + + try: + QueryFullProcessImageNameW = kernel32.QueryFullProcessImageNameW + except AttributeError: + QueryFullProcessImageNameW = None + else: + QueryFullProcessImageNameW.argtypes = [ + HANDLE, + DWORD, + wintypes.LPWSTR, + ctypes.POINTER(DWORD), + ] + QueryFullProcessImageNameW.restype = BOOL + + class PROCESS_MEMORY_COUNTERS_EX(ctypes.Structure): + _fields_ = [ + ('cb', DWORD), + ('PageFaultCount', DWORD), + ('PeakWorkingSetSize', SIZE_T), + ('WorkingSetSize', SIZE_T), + ('QuotaPeakPagedPoolUsage', SIZE_T), + ('QuotaPagedPoolUsage', SIZE_T), + ('QuotaPeakNonPagedPoolUsage', SIZE_T), + ('QuotaNonPagedPoolUsage', SIZE_T), + ('PagefileUsage', SIZE_T), + ('PeakPagefileUsage', SIZE_T), + ('PrivateUsage', SIZE_T), + ] + + psapi.EnumProcesses.argtypes = [ + ctypes.POINTER(DWORD), + DWORD, + ctypes.POINTER(DWORD), + ] + psapi.EnumProcesses.restype = BOOL + psapi.GetProcessMemoryInfo.argtypes = [ + HANDLE, + ctypes.POINTER(PROCESS_MEMORY_COUNTERS_EX), + DWORD, + ] + psapi.GetProcessMemoryInfo.restype = BOOL + psapi.QueryWorkingSet.argtypes = [HANDLE, ctypes.c_void_p, DWORD] + psapi.QueryWorkingSet.restype = BOOL class Unbuffered(io.TextIOBase): def __init__(self, stream): @@ -114,6 +202,9 @@ def close(self): class Proc: def __init__(self): + if IS_WINDOWS: + self.proc = None + return uname = os.uname() if uname[0] == "FreeBSD": self.proc = '/compat/linux/proc' @@ -121,6 +212,8 @@ def __init__(self): self.proc = '/proc' def path(self, *args): + if self.proc is None: + raise OSError(errno.ENOENT, "procfs is not available on Windows") return os.path.join(self.proc, *(str(a) for a in args)) def open(self, *args): @@ -146,6 +239,131 @@ def open(self, *args): # Functions # +def close_windows_handle(handle): + if handle: + kernel32.CloseHandle(handle) + + +def open_windows_process(pid, access): + handle = kernel32.OpenProcess(access, False, pid) + if not handle: + raise LookupError + return handle + + +def enum_windows_pids(): + process_count = 2048 + while True: + processes = (DWORD * process_count)() + bytes_returned = DWORD() + if not psapi.EnumProcesses(processes, ctypes.sizeof(processes), + ctypes.byref(bytes_returned)): + raise ctypes.WinError(get_last_windows_error()) + + returned = bytes_returned.value // ctypes.sizeof(DWORD) + if returned < process_count: + return [int(pid) for pid in processes[:returned] if int(pid)] + process_count *= 2 + + +def query_windows_image_path(handle): + if QueryFullProcessImageNameW is None: + return '' + + size = DWORD(32768) + buf = ctypes.create_unicode_buffer(size.value) + if QueryFullProcessImageNameW(handle, 0, buf, ctypes.byref(size)): + return buf.value + return '' + + +def get_windows_cmd_name(pid, handle, split_args, discriminate_by_pid): + path = query_windows_image_path(handle) + if split_args and path: + cmd = path + elif path: + cmd = os.path.basename(path) + else: + cmd = 'pid-%d' % pid + + if sys.version_info >= (3,): + cmd = cmd.encode(errors='replace').decode() + if discriminate_by_pid: + cmd = '%s [%d]' % (cmd, pid) + return cmd + + +def get_windows_memory_counters(handle): + counters = PROCESS_MEMORY_COUNTERS_EX() + counters.cb = ctypes.sizeof(PROCESS_MEMORY_COUNTERS_EX) + if not psapi.GetProcessMemoryInfo(handle, ctypes.byref(counters), + counters.cb): + raise LookupError + return counters + + +def get_windows_working_set(pid, handle, counters): + global have_pss + + page_bytes = int(PAGESIZE * 1024) + process_pages = int(counters.WorkingSetSize // page_bytes) + 1024 + if process_pages < 1024: + process_pages = 1024 + + for _ in range(12): + working_set = (ULONG_PTR * (process_pages + 1))() + if psapi.QueryWorkingSet(handle, working_set, + ctypes.sizeof(working_set)): + have_pss = 1 + private = 0.0 + shared = 0.0 + entry_count = min(int(working_set[0]), process_pages) + for idx in range(1, entry_count + 1): + flags = int(working_set[idx]) + is_shared = (flags >> 8) & 1 + if is_shared: + share_count = (flags >> 5) & 7 + if share_count < 1: + share_count = 1 + shared += PAGESIZE / share_count + else: + private += PAGESIZE + return private, shared, 0, 0, pid + + error = get_last_windows_error() + if error not in (ERROR_BAD_LENGTH, ERROR_INSUFFICIENT_BUFFER): + raise LookupError + process_pages *= 2 + + raise LookupError + + +def get_windows_fallback_memory(pid, counters): + global windows_used_fallback + + windows_used_fallback = True + working_set = counters.WorkingSetSize / 1024.0 + private_commit = counters.PrivateUsage / 1024.0 + private = min(working_set, private_commit) + shared = max(working_set - private, 0) + return private, shared, 0, 0, pid + + +def get_windows_mem_stats(pid, handle, can_query_working_set): + counters = get_windows_memory_counters(handle) + if can_query_working_set: + try: + return get_windows_working_set(pid, handle, counters) + except LookupError: + pass + return get_windows_fallback_memory(pid, counters) + + +def totals_available(): + if IS_WINDOWS: + return True + return have_pss and not windows_used_fallback + def parse_options(): help_msg = 'Show program core memory usage.' parser = argparse.ArgumentParser(prog='ps_mem', description=help_msg) @@ -210,6 +428,8 @@ def parse_options(): # (major,minor,release) def kernel_ver(): + if IS_WINDOWS: + return (0, 0, 0) kv = proc.open('sys/kernel/osrelease').readline().split(".")[:3] last = len(kv) if last == 2: @@ -400,6 +620,10 @@ def cmd_with_count(cmd, count): #-1= not available def val_accuracy(show_swap): """http://wiki.apache.org/spamassassin/TopSharedMemoryBug""" + if IS_WINDOWS: + ram_accuracy = (2, 1)[windows_used_fallback] + return ram_accuracy, -1 + kv = kernel_ver() pid = os.getpid() swap_accuracy = -1 @@ -425,11 +649,17 @@ def val_accuracy(show_swap): else: return 1, swap_accuracy -def show_val_accuracy( ram_inacc, swap_inacc, only_total, show_swap ): +def show_val_accuracy( ram_inacc, swap_inacc, only_total, show_swap, + swap_requested=None ): + if swap_requested is None: + swap_requested = show_swap + level = ("Warning","Error")[only_total] + if IS_WINDOWS and not show_swap and (ram_inacc == 1 or swap_requested): + level = "Warning" # Only show significant warnings - if not show_swap: + if not swap_requested: swap_inacc = 2 elif only_total: ram_inacc = 2 @@ -449,10 +679,16 @@ def show_val_accuracy( ram_inacc, swap_inacc, only_total, show_swap ): "Values reported could be too large, and totals are not reported\n" ) elif ram_inacc == 1: - sys.stderr.write( - "%s: Shared memory is slightly over-estimated by this system\n" - "for each program, so totals are not reported.\n" % level - ) + if IS_WINDOWS: + sys.stderr.write( + "%s: Shared memory is slightly over-estimated by this system\n" + "for some processes, so totals may be too large.\n" % level + ) + else: + sys.stderr.write( + "%s: Shared memory is slightly over-estimated by this system\n" + "for each program, so totals are not reported.\n" % level + ) if swap_inacc == -1: sys.stderr.write( @@ -470,12 +706,101 @@ def show_val_accuracy( ram_inacc, swap_inacc, only_total, show_swap ): accuracy = swap_inacc else: accuracy = ram_inacc + if IS_WINDOWS and not show_swap and accuracy == 1: + return if accuracy != 2: sys.exit(1) +def get_windows_memory_usage(pids_to_show, split_args, discriminate_by_pid, + include_self=False, only_self=False): + global have_pss + global windows_used_fallback + + have_pss = 0 + windows_used_fallback = False + cmds = {} + shareds = {} + shared_huges = {} + mem_ids = {} + count = {} + swaps = {} + + for pid in enum_windows_pids(): + if only_self and pid != our_pid: + continue + if pid == our_pid and not include_self: + continue + if pids_to_show and pid not in pids_to_show: + continue + + handle = None + can_query_working_set = True + try: + try: + handle = open_windows_process( + pid, PROCESS_QUERY_INFORMATION | PROCESS_VM_READ) + except LookupError: + handle = open_windows_process( + pid, PROCESS_QUERY_LIMITED_INFORMATION) + can_query_working_set = False + + cmd = get_windows_cmd_name(pid, handle, split_args, + discriminate_by_pid) + private, shared, shared_huge, swap, mem_id = \ + get_windows_mem_stats(pid, handle, can_query_working_set) + except LookupError: + continue + finally: + close_windows_handle(handle) + + if shareds.get(cmd): + if have_pss: + shareds[cmd] += shared + elif shareds[cmd] < shared: + shareds[cmd] = shared + else: + shareds[cmd] = shared + if shared_huges.get(cmd): + if shared_huges[cmd] < shared_huge: + shared_huges[cmd] = shared_huge + else: + shared_huges[cmd] = shared_huge + cmds[cmd] = cmds.setdefault(cmd, 0) + private + if cmd in count: + count[cmd] += 1 + else: + count[cmd] = 1 + mem_ids.setdefault(cmd, {}).update({mem_id: None}) + swaps[cmd] = swaps.setdefault(cmd, 0) + swap + + total_swap = 0 + total = 0 + + for cmd in cmds: + cmd_count = count[cmd] + if len(mem_ids[cmd]) == 1 and cmd_count > 1: + cmds[cmd] /= cmd_count + if have_pss: + shareds[cmd] /= cmd_count + shareds[cmd] += shared_huges[cmd] + cmds[cmd] = cmds[cmd] + shareds[cmd] + total += cmds[cmd] + total_swap += swaps[cmd] + + sorted_cmds = sorted(cmds.items(), key=lambda x:x[1]) + sorted_cmds = [x for x in sorted_cmds if x[1]] + + return sorted_cmds, shareds, count, total, swaps, total_swap + + def get_memory_usage(pids_to_show, split_args, discriminate_by_pid, include_self=False, only_self=False): + if IS_WINDOWS: + return get_windows_memory_usage(pids_to_show, split_args, + discriminate_by_pid, include_self, + only_self) + cmds = {} shareds = {} shared_huges = {} @@ -581,16 +906,19 @@ def print_memory_usage(sorted_cmds, shareds, count, total, swaps, total_swap, sys.stdout.write(output_string % output_data) # Only show totals if appropriate - if have_swap_pss and show_swap: # kernel will have_pss + if have_swap_pss and show_swap and totals_available(): # kernel will have_pss sys.stdout.write("%s\n%s%9s%s%9s\n%s\n" % ("-" * 45, " " * 24, human(total), " " * 3, human(total_swap), "=" * 45)) - elif have_pss: + elif totals_available(): sys.stdout.write("%s\n%s%9s\n%s\n" % ("-" * 33, " " * 24, human(total), "=" * 33)) def verify_environment(pids_to_show): + if IS_WINDOWS: + return + if os.geteuid() != 0 and not pids_to_show: sys.stderr.write("Sorry, root permission required, or specify pids with -p\n") sys.stderr.close() @@ -615,6 +943,9 @@ def main(): split_args, pids_to_show, watch, only_total, discriminate_by_pid, \ show_swap = parse_options() + swap_requested = show_swap + if IS_WINDOWS: + show_swap = False verify_environment(pids_to_show) @@ -628,9 +959,10 @@ def main(): sorted_cmds, shareds, count, total, swaps, total_swap = \ get_memory_usage(pids_to_show, split_args, discriminate_by_pid) - if only_total and show_swap and have_swap_pss: + if only_total and show_swap and have_swap_pss and \ + totals_available(): sys.stdout.write(human(total_swap, units=1)+'\n') - elif only_total and not show_swap and have_pss: + elif only_total and not show_swap and totals_available(): sys.stdout.write(human(total, units=1)+'\n') elif not only_total: print_memory_usage(sorted_cmds, shareds, count, total, @@ -647,9 +979,9 @@ def main(): sorted_cmds, shareds, count, total, swaps, total_swap = \ get_memory_usage(pids_to_show, split_args, discriminate_by_pid) - if only_total and show_swap and have_swap_pss: + if only_total and show_swap and have_swap_pss and totals_available(): sys.stdout.write(human(total_swap, units=1)+'\n') - elif only_total and not show_swap and have_pss: + elif only_total and not show_swap and totals_available(): sys.stdout.write(human(total, units=1)+'\n') elif not only_total: print_memory_usage(sorted_cmds, shareds, count, total, swaps, @@ -660,7 +992,8 @@ def main(): # one which is reenabled after this script finishes. sys.stdout.close() - ram_accuracy, swap_accuracy = val_accuracy( show_swap ) - show_val_accuracy( ram_accuracy, swap_accuracy, only_total, show_swap ) + ram_accuracy, swap_accuracy = val_accuracy( swap_requested ) + show_val_accuracy( ram_accuracy, swap_accuracy, only_total, show_swap, + swap_requested ) if __name__ == '__main__': main()