Merge pull request #745 from EricRahm/uss_osx

Measure USS on OSX

4 files changed, 117 insertions, 3 deletions

diff --git a/docs/index.rst b/docs/index.rst
index aa1bd36b..7f3afcaa 100644
--- a/docs/index.rst
+++ b/docs/index.rst
@@ -1000,7 +1000,7 @@ Process class
      +--------+---------+-------+---------+--------------------+
      | text   | pageins | data  |         | peak_paged_pool    |
      +--------+---------+-------+---------+--------------------+
-     | lib    |         | stack |         | paged_pool         |
+     | lib    | uss     | stack |         | paged_pool         |
      +--------+---------+-------+---------+--------------------+
      | data   |         |       |         | peak_nonpaged_pool |
      +--------+---------+-------+---------+--------------------+
diff --git a/psutil/_psosx.py b/psutil/_psosx.py
index 271e31f9..52994847 100644
--- a/psutil/_psosx.py
+++ b/psutil/_psosx.py
@@ -58,7 +58,7 @@ svmem = namedtuple(
     'svmem', ['total', 'available', 'percent', 'used', 'free',
               'active', 'inactive', 'wired'])
 
-pextmem = namedtuple('pextmem', ['rss', 'vms', 'pfaults', 'pageins'])
+pextmem = namedtuple('pextmem', ['rss', 'vms', 'pfaults', 'pageins', 'uss'])
 
 pmmap_grouped = namedtuple(
     'pmmap_grouped',
@@ -279,7 +279,8 @@ class Process(object):
     @wrap_exceptions
     def memory_info_ex(self):
         rss, vms, pfaults, pageins = cext.proc_memory_info(self.pid)
-        return pextmem(rss, vms, pfaults * PAGESIZE, pageins * PAGESIZE)
+        uss = cext.proc_memory_uss(self.pid)
+        return pextmem(rss, vms, pfaults * PAGESIZE, pageins * PAGESIZE, uss)
 
     @wrap_exceptions
     def cpu_times(self):
diff --git a/psutil/_psutil_osx.c b/psutil/_psutil_osx.c
index 4ffae898..0cf0684b 100644
--- a/psutil/_psutil_osx.c
+++ b/psutil/_psutil_osx.c
@@ -534,6 +534,116 @@ psutil_proc_memory_info(PyObject *self, PyObject *args) {
     );
 }
 
+/**
+ * Indicates if the given virtual address on the given architecture is in the
+ * shared VM region.
+ */
+bool
+in_shared_region(mach_vm_address_t addr, cpu_type_t type)
+{
+    mach_vm_address_t base;
+    mach_vm_address_t size;
+
+    switch (type) {
+        case CPU_TYPE_ARM:
+            base = SHARED_REGION_BASE_ARM;
+            size = SHARED_REGION_SIZE_ARM;
+            break;
+        case CPU_TYPE_I386:
+            base = SHARED_REGION_BASE_I386;
+            size = SHARED_REGION_SIZE_I386;
+            break;
+        case CPU_TYPE_X86_64:
+            base = SHARED_REGION_BASE_X86_64;
+            size = SHARED_REGION_SIZE_X86_64;
+            break;
+        default:
+            return false;
+    }
+
+    return base <= addr && addr < (base + size);
+}
+
+
+/**
+ * Returns the USS (unique set size) of the process.
+ */
+static PyObject *
+psutil_proc_memory_uss(PyObject *self, PyObject *args)
+{
+    long pid;
+    if (! PyArg_ParseTuple(args, "l", &pid))
+        return NULL;
+
+    mach_port_t task = MACH_PORT_NULL;
+    int err = task_for_pid(mach_task_self(), pid, &task);
+    if (err != KERN_SUCCESS) {
+        psutil_raise_ad_or_nsp(pid);
+        return NULL;
+    }
+
+    cpu_type_t cpu_type;
+    size_t len = sizeof(cpu_type);
+    if (sysctlbyname("sysctl.proc_cputype", &cpu_type, &len, NULL, 0) != 0) {
+        PyErr_SetFromErrno(PyExc_OSError);
+        return NULL;
+    }
+
+    // Roughly based on libtop_update_vm_regions in
+    // http://www.opensource.apple.com/source/top/top-100.1.2/libtop.c
+    size_t private_pages = 0;
+    mach_vm_size_t size = 0;
+    for (mach_vm_address_t addr = MACH_VM_MIN_ADDRESS; ; addr += size) {
+        vm_region_top_info_data_t info;
+        mach_msg_type_number_t info_count = VM_REGION_TOP_INFO_COUNT;
+        mach_port_t object_name;
+
+        kern_return_t kr =
+            mach_vm_region(task, &addr, &size, VM_REGION_TOP_INFO,
+                    (vm_region_info_t)&info,
+                    &info_count, &object_name);
+        if (kr == KERN_INVALID_ADDRESS) {
+            // Done iterating VM regions.
+            break;
+        }
+        else if (kr != KERN_SUCCESS) {
+            return false;
+        }
+
+        if (in_shared_region(addr, cpu_type) && info.share_mode != SM_PRIVATE) {
+            continue;
+        }
+
+        switch (info.share_mode) {
+            case SM_LARGE_PAGE:
+                // NB: Large pages are not shareable and always resident.
+            case SM_PRIVATE:
+                private_pages += info.private_pages_resident;
+                private_pages += info.shared_pages_resident;
+                break;
+            case SM_COW:
+                private_pages += info.private_pages_resident;
+                if (info.ref_count == 1) {
+                    // Treat copy-on-write pages as private if they only have one reference.
+                    private_pages += info.shared_pages_resident;
+                }
+                break;
+            case SM_SHARED:
+            default:
+                break;
+        }
+    }
+
+    mach_port_deallocate(mach_task_self(), task);
+
+    vm_size_t page_size;
+    if (host_page_size(mach_host_self(), &page_size) != KERN_SUCCESS) {
+        page_size = PAGE_SIZE;
+    }
+
+    return Py_BuildValue("K", private_pages * page_size);
+}
+
 
 /*
  * Return number of threads used by process as a Python integer.
@@ -1700,6 +1810,8 @@ PsutilMethods[] = {
      "seconds since the epoch"},
     {"proc_memory_info", psutil_proc_memory_info, METH_VARARGS,
      "Return memory information about a process"},
+    {"proc_memory_uss", psutil_proc_memory_uss, METH_VARARGS,
+     "Return USS a process"},
     {"proc_num_threads", psutil_proc_num_threads, METH_VARARGS,
      "Return number of threads used by process"},
     {"proc_status", psutil_proc_status, METH_VARARGS,
diff --git a/psutil/_psutil_osx.h b/psutil/_psutil_osx.h
index 2220ec3b..a6a2957f 100644
--- a/psutil/_psutil_osx.h
+++ b/psutil/_psutil_osx.h
@@ -17,6 +17,7 @@ static PyObject* psutil_proc_exe(PyObject* self, PyObject* args);
 static PyObject* psutil_proc_gids(PyObject* self, PyObject* args);
 static PyObject* psutil_proc_memory_info(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_fds(PyObject* self, PyObject* args);
 static PyObject* psutil_proc_num_threads(PyObject* self, PyObject* args);