Windows版本smem_通过进程名统计对应内存占用

文章目录

[Windows 版 smem（支持 RSS/PSS/USS、排序、GROUP BY、--help）](#Windows 版 smem（支持 RSS/PSS/USS、排序、GROUP BY、--help）)
使用方法
查看帮助
[按 PSS 排序](#按 PSS 排序)
[按 RSS 倒序](#按 RSS 倒序)
[按 USS 排序](#按 USS 排序)
[仅显示 GROUP BY](#仅显示 GROUP BY)
[仅显示 DETAIL](#仅显示 DETAIL)
使用之前需要安装依赖
完整代码
支持的排序字段
示例
- [按 PSS 倒序](#按 PSS 倒序)
- [按 RSS 升序](#按 RSS 升序)
- 按进程名排序
- [仅显示 GROUP](#仅显示 GROUP)
- [仅显示 DETAIL](#仅显示 DETAIL)

Windows 版 smem（支持 RSS/PSS/USS、排序、GROUP BY、--help）

下面是增强版 Python3 实现。

支持功能：

RSS / PSS / USS
进程明细
GROUP BY 进程名
TOTAL 汇总
支持排序
支持 --help
支持按字段排序
支持升序/降序
支持仅显示 GROUP BY
支持仅显示 DETAIL

使用方法

bash 复制代码

python wsmem.py

默认：

显示 DETAIL
显示 GROUP
按进程名排序

查看帮助

bash 复制代码

python wsmem.py --help

按 PSS 排序

bash 复制代码

python wsmem.py --sort pss

按 RSS 倒序

bash 复制代码

python wsmem.py --sort rss --desc

按 USS 排序

bash 复制代码

python wsmem.py --sort uss

仅显示 GROUP BY

bash 复制代码

python wsmem.py --group-only

仅显示 DETAIL

bash 复制代码

python wsmem.py --detail-only

使用之前需要安装依赖

shell 复制代码

pip install psutil

完整代码

python 复制代码

import ctypes
import argparse
import psutil

from ctypes import wintypes
from collections import defaultdict

PAGE_SIZE = 4096

PROCESS_QUERY_INFORMATION = 0x0400
PROCESS_VM_READ = 0x0010

kernel32 = ctypes.WinDLL("kernel32")
psapi = ctypes.WinDLL("psapi")

# =========================================================
# Structures
# =========================================================

class MEMORY_BASIC_INFORMATION(ctypes.Structure):
    _fields_ = [
        ("BaseAddress", ctypes.c_void_p),
        ("AllocationBase", ctypes.c_void_p),
        ("AllocationProtect", wintypes.DWORD),
        ("RegionSize", ctypes.c_size_t),
        ("State", wintypes.DWORD),
        ("Protect", wintypes.DWORD),
        ("Type", wintypes.DWORD),
    ]


class PSAPI_WORKING_SET_EX_BLOCK(ctypes.Structure):
    _fields_ = [
        ("Flags", ctypes.c_ulonglong)
    ]


class PSAPI_WORKING_SET_EX_INFORMATION(ctypes.Structure):
    _fields_ = [
        ("VirtualAddress", ctypes.c_void_p),
        ("VirtualAttributes",
         PSAPI_WORKING_SET_EX_BLOCK)
    ]


VirtualQueryEx = kernel32.VirtualQueryEx
QueryWorkingSetEx = psapi.QueryWorkingSetEx

OpenProcess = kernel32.OpenProcess
CloseHandle = kernel32.CloseHnadle

# =========================================================
# Argument
# =========================================================

parser = argparse.ArgumentParser(
    description="Windows smem-like tool (RSS/PSS/USS)")

parser.add_argument(
    "--sort",
    choices=["name", "rss", "pss", "uss"],
    default="name",
    help="sort field")

parser.add_argument(
    "--desc",
    action="store_true",
    help="descending sort")

parser.add_argument(
    "--group-only",
    action="store_true",
    help="show group summary only")

parser.add_argument(
    "--detail-only",
    action="store_true",
    help="show detail only")

args = parser.parse_args()

# =========================================================
# Helper
# =========================================================

def mb(x):
    return x / 1024 / 1024


def parse_ws_flags(flags):

    valid = flags & 0x1

    share_count = (flags >> 1) & 0x7

    shared = (flags >> 15) & 0x1

    return valid, share_count, shared


# =========================================================
# Calculate Process Memory
# =========================================================

def calculate_process_memory(pid):

    try:

        hProcess = OpenProcess(
            PROCESS_QUERY_INFORMATION |
            PROCESS_VM_READ,
            False,
            pid)

        if not hProcess:
            return None

        proc = psutil.Process(pid)

        rss = proc.memory_info().rss

        uss = 0
        pss = 0

        addr = 0

        mbi = MEMORY_BASIC_INFORMATION()

        while VirtualQueryEx(
                hProcess,
                ctypes.c_void_p(addr),
                ctypes.byref(mbi),
                ctypes.sizeof(mbi)):

            MEM_COMMIT = 0x1000

            if mbi.State == MEM_COMMIT:

                region_size = mbi.RegionSize

                page_count = region_size // PAGE_SIZE

                if page_count == 0:
                    addr += mbi.RegionSize
                    continue

                ws_array = (
                    PSAPI_WORKING_SET_EX_INFORMATION
                    * page_count)()

                for i in range(page_count):

                    ws_array[i].VirtualAddress = \
                        addr + i * PAGE_SIZE

                if QueryWorkingSetEx(
                        hProcess,
                        ws_array,
                        ctypes.sizeof(ws_array)):

                    for i in range(page_count):

                        flags = \
                            ws_array[i]\
                            .VirtualAttributes\
                            .Flags

                        valid, share_count, is_shared = \
                            parse_ws_flags(flags)

                        if not valid:
                            continue

                        # USS
                        if share_count <= 1:

                            uss += PAGE_SIZE
                            pss += PAGE_SIZE

                        # Shared
                        else:

                            pss += \
                                PAGE_SIZE / share_count

            addr += mbi.RegionSize

        CloseHandle(hProcess)

        return {
            "pid": pid,
            "name": proc.name(),
            "rss": rss,
            "pss": pss,
            "uss": uss
        }

    except:
        return None


# =========================================================
# Scan Process
# =========================================================


results = [ ]


print("Scanning processes...\n")

for proc in psutil.process_iter(['pid']):

    r = calculate_process_memory(
        proc.info['pid'])

    if r:
        results.append(r)

# =========================================================
# Sort Detail
# =========================================================

if args.sort == "name":

    results.sort(
        key=lambda x: x['name'].lower(),
        reverse=args.desc)

else:

    results.sort(
        key=lambda x: x[args.sort],
        reverse=args.desc)

# =========================================================
# DETAIL
# =========================================================

if not args.group_only:

    print("=" * 100)
    print("PROCESS DETAIL")
    print("=" * 100)

    print(
        f"{'PID':<8}"
        f"{'PROCESS':<35}"
        f"{'RSS(MB)':>15}"
        f"{'PSS(MB)':>15}"
        f"{'USS(MB)':>15}"
    )

    print("-" * 100)

    for r in results:

        print(
            f"{r['pid']:<8}"
            f"{r['name']:<35}"
            f"{mb(r['rss']):>15.2f}"
            f"{mb(r['pss']):>15.2f}"
            f"{mb(r['uss']):>15.2f}"
        )

# =========================================================
# GROUP BY
# =========================================================

group_stats = defaultdict(lambda: {
    "count": 0,
    "rss": 0,
    "pss": 0,
    "uss": 0
})

total_count = 0
total_rss = 0
total_pss = 0
total_uss = 0

for r in results:

    name = r["name"]

    group_stats[name]["count"] += 1
    group_stats[name]["rss"] += r["rss"]
    group_stats[name]["pss"] += r["pss"]
    group_stats[name]["uss"] += r["uss"]

    total_count += 1
    total_rss += r["rss"]
    total_pss += r["pss"]
    total_uss += r["uss"]

# =========================================================
# Sort Group
# =========================================================

group_items = list(group_stats.items())

if args.sort == "name":

    group_items.sort(
        key=lambda x: x[0].lower(),
        reverse=args.desc)

else:

    group_items.sort(
        key=lambda x: x[1][args.sort],
        reverse=args.desc)

# =========================================================
# GROUP OUTPUT
# =========================================================

if not args.detail_only:

    print()
    print("=" * 100)
    print("GROUP BY PROCESS NAME")
    print("=" * 100)

    print(
        f"{'PROCESS':<35}"
        f"{'COUNT':>10}"
        f"{'SUM_RSS(MB)':>18}"
        f"{'SUM_PSS(MB)':>18}"
        f"{'SUM_USS(MB)':>18}"
    )

    print("-" * 100)

    for name, stat in group_items:

        print(
            f"{name:<35}"
            f"{stat['count']:>10}"
            f"{mb(stat['rss']):>18.2f}"
            f"{mb(stat['pss']):>18.2f}"
            f"{mb(stat['uss']):>18.2f}"
        )

    print("-" * 100)

    print(
        f"{'TOTAL':<35}"
        f"{total_count:>10}"
        f"{mb(total_rss):>18.2f}"
        f"{mb(total_pss):>18.2f}"
        f"{mb(total_uss):>18.2f}"
    )

支持的排序字段

字段	含义
name	进程名
rss	RSS
pss	PSS
uss	USS

示例

按 PSS 倒序

bash 复制代码

python wsmem.py --sort pss --desc

按 RSS 升序

bash 复制代码

python wsmem.py --sort rss

按进程名排序

bash 复制代码

python wsmem.py --sort name

仅显示 GROUP

bash 复制代码

python wsmem.py --group-only

仅显示 DETAIL

bash 复制代码

python wsmem.py --detail-only