Merge pull request #746 from EricRahm/uss_windows

Measure USS on Windows

4 files changed, 95 insertions, 3 deletions

diff --git a/docs/index.rst b/docs/index.rst
index 7f3afcaa..17f20235 100644
--- a/docs/index.rst
+++ b/docs/index.rst
@@ -1012,6 +1012,8 @@ Process class
      +--------+---------+-------+---------+--------------------+
      |        |         |       |         | private            |
      +--------+---------+-------+---------+--------------------+
+     |        |         |       |         | uss                |
+     +--------+---------+-------+---------+--------------------+
 
      Windows metrics are extracted from
      `PROCESS_MEMORY_COUNTERS_EX <http://msdn.microsoft.com/en-us/library/windows/desktop/ms684874(v=vs.85).aspx>`__ structure.
diff --git a/psutil/_psutil_windows.c b/psutil/_psutil_windows.c
index 10ff5dc0..d506660c 100644
--- a/psutil/_psutil_windows.c
+++ b/psutil/_psutil_windows.c
@@ -799,6 +799,89 @@ psutil_proc_memory_info_2(PyObject *self, PyObject *args) {
         pfault_count, m1, m2, m3, m4, m5, m6, m7, m8, private);
 }
 
+/**
+ * Returns the USS of the process.
+ */
+static PyObject *
+psutil_proc_memory_uss(PyObject *self, PyObject *args)
+{
+    DWORD pid;
+    HANDLE proc;
+    PSAPI_WORKING_SET_INFORMATION tmp;
+    DWORD tmp_size = sizeof(tmp);
+    size_t entries;
+    size_t private_pages;
+    size_t i;
+    DWORD info_array_size;
+    PSAPI_WORKING_SET_INFORMATION* info_array;
+    SYSTEM_INFO system_info;
+    PyObject* py_result = NULL;
+    unsigned long long total = 0;
+
+    if (! PyArg_ParseTuple(args, "l", &pid))
+        return NULL;
+
+    proc = psutil_handle_from_pid(pid);
+    if (proc == NULL)
+        return NULL;
+
+    // Determine how many entries we need.
+    memset(&tmp, 0, tmp_size);
+    if (!QueryWorkingSet(proc, &tmp, tmp_size)) {
+        // NB: QueryWorkingSet is expected to fail here due to the
+        // buffer being too small.
+        if (tmp.NumberOfEntries == 0) {
+            PyErr_SetFromWindowsErr(0);
+            goto done;
+        }
+    }
+
+    // Fudge the size in case new entries are added between calls.
+    entries = tmp.NumberOfEntries * 2;
+
+    if (!entries) {
+        goto done;
+    }
+
+    info_array_size = tmp_size + (entries * sizeof(PSAPI_WORKING_SET_BLOCK));
+    info_array = (PSAPI_WORKING_SET_INFORMATION*)malloc(info_array_size);
+    if (!info_array) {
+        PyErr_NoMemory();
+        goto done;
+    }
+
+    if (!QueryWorkingSet(proc, info_array, info_array_size)) {
+        PyErr_SetFromWindowsErr(0);
+        goto done;
+    }
+
+    entries = (size_t)info_array->NumberOfEntries;
+    private_pages = 0;
+    for (i = 0; i < entries; i++) {
+        // Count shared pages that only one process is using as private.
+        if (!info_array->WorkingSetInfo[i].Shared ||
+                info_array->WorkingSetInfo[i].ShareCount <= 1) {
+            private_pages++;
+        }
+    }
+
+    // GetSystemInfo has no return value.
+    GetSystemInfo(&system_info);
+    total = private_pages * system_info.dwPageSize;
+    py_result = Py_BuildValue("K", total);
+
+done:
+    if (proc) {
+        CloseHandle(proc);
+    }
+
+    if (info_array) {
+        free(info_array);
+    }
+
+    return py_result;
+}
+
 
 /*
  * Return a Python integer indicating the total amount of physical memory
@@ -3031,6 +3114,8 @@ PsutilMethods[] = {
      "Return a tuple of process memory information"},
     {"proc_memory_info_2", psutil_proc_memory_info_2, METH_VARARGS,
      "Alternate implementation"},
+    {"proc_memory_uss", psutil_proc_memory_uss, METH_VARARGS,
+     "Return the USS of the process"},
     {"proc_cwd", psutil_proc_cwd, METH_VARARGS,
      "Return process current working directory"},
     {"proc_suspend", psutil_proc_suspend, METH_VARARGS,
diff --git a/psutil/_psutil_windows.h b/psutil/_psutil_windows.h
index 0c22bf7e..76fb09d1 100644
--- a/psutil/_psutil_windows.h
+++ b/psutil/_psutil_windows.h
@@ -24,6 +24,7 @@ static PyObject* psutil_proc_kill(PyObject* self, PyObject* args);
 static PyObject* psutil_proc_memory_info(PyObject* self, PyObject* args);
 static PyObject* psutil_proc_memory_info_2(PyObject* self, PyObject* args);
 static PyObject* psutil_proc_memory_maps(PyObject* self, PyObject* args);
+static PyObject* psutil_proc_memory_uss(PyObject* self, PyObject* args);
 static PyObject* psutil_proc_name(PyObject* self, PyObject* args);
 static PyObject* psutil_proc_num_handles(PyObject* self, PyObject* args);
 static PyObject* psutil_proc_open_files(PyObject* self, PyObject* args);
diff --git a/psutil/_pswindows.py b/psutil/_pswindows.py
index 6392c510..4403e5eb 100644
--- a/psutil/_pswindows.py
+++ b/psutil/_pswindows.py
@@ -87,7 +87,7 @@ svmem = namedtuple('svmem', ['total', 'available', 'percent', 'used', 'free'])
 pextmem = namedtuple(
     'pextmem', ['num_page_faults', 'peak_wset', 'wset', 'peak_paged_pool',
                 'paged_pool', 'peak_nonpaged_pool', 'nonpaged_pool',
-                'pagefile', 'peak_pagefile', 'private'])
+                'pagefile', 'peak_pagefile', 'private', 'uss'])
 pmmap_grouped = namedtuple('pmmap_grouped', ['path', 'rss'])
 pmmap_ext = namedtuple(
     'pmmap_ext', 'addr perms ' + ' '.join(pmmap_grouped._fields))
@@ -359,12 +359,16 @@ class Process(object):
 
     def _get_raw_meminfo(self):
         try:
-            return cext.proc_memory_info(self.pid)
+            info = cext.proc_memory_info(self.pid)
+            uss = cext.proc_memory_uss(self.pid)
+            return info + (uss,)
         except OSError as err:
             if err.errno in ACCESS_DENIED_SET:
                 # TODO: the C ext can probably be refactored in order
                 # to get this from cext.proc_info()
-                return cext.proc_memory_info_2(self.pid)
+                info = cext.proc_memory_info_2(self.pid)
+                uss = cext.proc_memory_uss(self.pid)
+                return info + (uss,)
             raise
 
     @wrap_exceptions