Measure USS on OSX

6 files changed, 151 insertions, 7 deletions

diff --git a/docs/index.rst b/docs/index.rst
index 566b9996..03139f58 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..e3582b8b 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',
@@ -278,8 +278,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)
+        rss, vms, pfaults, pageins, uss = cext.proc_memory_info(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..ab3ac53a 100644
--- a/psutil/_psutil_osx.c
+++ b/psutil/_psutil_osx.c
@@ -40,6 +40,7 @@
 #include <IOKit/IOBSD.h>
 
 #include "_psutil_osx.h"
+#include "_psutil_osx_uss.h"
 #include "_psutil_common.h"
 #include "arch/osx/process_info.h"
 
@@ -515,6 +516,9 @@ static PyObject *
 psutil_proc_memory_info(PyObject *self, PyObject *args) {
     long pid;
     struct proc_taskinfo pti;
+    int err;
+    mach_port_t task = MACH_PORT_NULL;
+    int64_t uss = 0;
 
     if (! PyArg_ParseTuple(args, "l", &pid))
         return NULL;
@@ -525,13 +529,31 @@ psutil_proc_memory_info(PyObject *self, PyObject *args) {
     // I just give up...
     // struct proc_regioninfo pri;
     // psutil_proc_pidinfo(pid, PROC_PIDREGIONINFO, &pri, sizeof(pri))
+
+    err = task_for_pid(mach_task_self(), pid, &task);
+    if (err != KERN_SUCCESS) {
+        psutil_raise_ad_or_nsp(pid);
+        goto error;
+    }
+
+    calc_uss(task, &uss);
+
+    if (task != MACH_PORT_NULL)
+        mach_port_deallocate(mach_task_self(), task);
+
     return Py_BuildValue(
-        "(KKkk)",
+        "(KKkkK)",
         pti.pti_resident_size,  // resident memory size (rss)
         pti.pti_virtual_size,   // virtual memory size (vms)
         pti.pti_faults,         // number of page faults (pages)
-        pti.pti_pageins         // number of actual pageins (pages)
+        pti.pti_pageins,        // number of actual pageins (pages)
+        (long long)uss          // unique memory size (uss)
     );
+
+error:
+    if (task != MACH_PORT_NULL)
+        mach_port_deallocate(mach_task_self(), task);
+    return NULL;
 }
 
 
diff --git a/psutil/_psutil_osx_uss.c b/psutil/_psutil_osx_uss.c
new file mode 100644
index 00000000..72912f52
--- /dev/null
+++ b/psutil/_psutil_osx_uss.c
@@ -0,0 +1,105 @@
+/* -*- Mode: C++; tab-width: 8; indent-tabs-mode: nil; c-basic-offset: 2 -*- */
+/* vim: set ts=8 sts=2 et sw=2 tw=80: */
+/* This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at http://mozilla.org/MPL/2.0/. */
+
+#include <mach/mach_init.h>
+#include <mach/mach_vm.h>
+#include <mach/shared_region.h>
+#include <mach/task.h>
+#include <sys/sysctl.h>
+
+#include "_psutil_osx_uss.h"
+
+bool
+InSharedRegion(mach_vm_address_t aAddr, cpu_type_t aType)
+{
+  mach_vm_address_t base;
+  mach_vm_address_t size;
+
+  switch (aType) {
+    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 <= aAddr && aAddr < (base + size);
+}
+
+bool
+calc_uss(mach_port_t target, int64_t* aN)
+{
+  if (!aN) {
+    return false;
+  }
+
+  cpu_type_t cpu_type;
+  size_t len = sizeof(cpu_type);
+  if (sysctlbyname("sysctl.proc_cputype", &cpu_type, &len, NULL, 0) != 0) {
+    return false;
+  }
+
+  /* Roughly based on libtop_update_vm_regions in
+    http://www.opensource.apple.com/source/top/top-100.1.2/libtop.c */
+  size_t privatePages = 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 infoCount = VM_REGION_TOP_INFO_COUNT;
+    mach_port_t objectName;
+
+    kern_return_t kr =
+        mach_vm_region(target, &addr, &size, VM_REGION_TOP_INFO,
+                       (vm_region_info_t)&info,
+                       &infoCount, &objectName);
+    if (kr == KERN_INVALID_ADDRESS) {
+      /* Done iterating VM regions. */
+      break;
+    } else if (kr != KERN_SUCCESS) {
+      return false;
+    }
+
+    if (InSharedRegion(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:
+        privatePages += info.private_pages_resident;
+        privatePages += info.shared_pages_resident;
+        break;
+      case SM_COW:
+        privatePages += info.private_pages_resident;
+        if (info.ref_count == 1) {
+          /* Treat copy-on-write pages as private if they only have one reference. */
+          privatePages += info.shared_pages_resident;
+        }
+        break;
+      case SM_SHARED:
+      default:
+        break;
+    }
+  }
+
+  vm_size_t pageSize;
+  if (host_page_size(mach_host_self(), &pageSize) != KERN_SUCCESS) {
+    pageSize = PAGE_SIZE;
+  }
+
+  *aN = privatePages * pageSize;
+  return true;
+}
diff --git a/psutil/_psutil_osx_uss.h b/psutil/_psutil_osx_uss.h
new file mode 100644
index 00000000..87e75f6f
--- /dev/null
+++ b/psutil/_psutil_osx_uss.h
@@ -0,0 +1,16 @@
+/* -*- Mode: C++; tab-width: 8; indent-tabs-mode: nil; c-basic-offset: 2 -*- */
+/* vim: set ts=8 sts=2 et sw=2 tw=80: */
+/* This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at http://mozilla.org/MPL/2.0/. */
+
+#ifndef _psutil_osx_uss_h
+#define _psutil_osx_uss_h
+
+#include <stdbool.h>
+#include <stdint.h>
+#include <mach/port.h>
+
+bool calc_uss(mach_port_t target, int64_t* aN);
+
+#endif
diff --git a/setup.py b/setup.py
index 4d4d8650..aef6abdd 100644
--- a/setup.py
+++ b/setup.py
@@ -119,7 +119,8 @@ elif sys.platform.startswith("darwin"):
         sources=[
             'psutil/_psutil_osx.c',
             'psutil/_psutil_common.c',
-            'psutil/arch/osx/process_info.c'
+            'psutil/arch/osx/process_info.c',
+            'psutil/_psutil_osx_uss.c'
         ],
         define_macros=[VERSION_MACRO],
         extra_link_args=[