[pypy-commit] pypy default: Implement _PyTraceMalloc_Track() and define a macro that can be

[arigo]

Author: Armin Rigo <arigo at tunes.org>
Branch: 
Changeset: r90637:f0290a07535c
Date: 2017-03-12 09:11 +0100
http://bitbucket.org/pypy/pypy/changeset/f0290a07535c/

Log:	Implement _PyTraceMalloc_Track() and define a macro that can be
	tested for in CPython C extensions. See pymem.h for the exact way to
	say that...

diff --git a/pypy/module/cpyext/include/pymem.h b/pypy/module/cpyext/include/pymem.h
--- a/pypy/module/cpyext/include/pymem.h
+++ b/pypy/module/cpyext/include/pymem.h
@@ -51,6 +51,25 @@
 #define PyMem_Del               PyMem_Free
 #define PyMem_DEL               PyMem_FREE
 
+
+/* From CPython 3.6, with a different goal.  _PyTraceMalloc_Track()
+ * is equivalent to __pypy__.add_memory_pressure(size); it works with
+ * or without the GIL.  _PyTraceMalloc_Untrack() is an empty stub.
+ * You can check if these functions are available by using:
+ *
+ *    #if defined(PYPY_TRACEMALLOC) || \
+ *         (PY_VERSION_HEX >= 0x03060000 && !defined(Py_LIMITED_API))
+ */
+#define PYPY_TRACEMALLOC        1
+
+typedef unsigned int _PyTraceMalloc_domain_t;
+
+PyAPI_FUNC(int) _PyTraceMalloc_Track(_PyTraceMalloc_domain_t domain,
+                                     uintptr_t ptr, size_t size);
+PyAPI_FUNC(int) _PyTraceMalloc_Untrack(_PyTraceMalloc_domain_t domain,
+                                       uintptr_t ptr);
+
+
 #ifdef __cplusplus
 }
 #endif
diff --git a/pypy/module/cpyext/object.py b/pypy/module/cpyext/object.py
--- a/pypy/module/cpyext/object.py
+++ b/pypy/module/cpyext/object.py
@@ -475,3 +475,8 @@
     with rffi.scoped_nonmovingbuffer(data) as buf:
         fwrite(buf, 1, count, fp)
     return 0
+
+ at cpython_api([lltype.Signed], lltype.Void)
+def _PyPyGC_AddMemoryPressure(space, report):
+    from rpython.rlib import rgc
+    rgc.add_memory_pressure(report)
diff --git a/pypy/module/cpyext/src/pymem.c b/pypy/module/cpyext/src/pymem.c
--- a/pypy/module/cpyext/src/pymem.c
+++ b/pypy/module/cpyext/src/pymem.c
@@ -4,3 +4,46 @@
 {
     return malloc((n) ? (n) : 1);
 }
+
+int _PyTraceMalloc_Track(_PyTraceMalloc_domain_t domain,
+                         uintptr_t ptr, size_t size)
+{
+    /* to avoid acquiring/releasing the GIL too often, only do it
+       if the total reported size exceeds 64KB. */
+    static volatile long unreported_size = 0;
+    long prev, next, report;
+
+    size += sizeof(long);
+    /* ^^^ to account for some alignment.  Important, otherwise we'd
+     * collect sizes of, say, 1-bytes mallocs in 1-bytes increment */
+
+ retry:
+    report = 0;
+    prev = unreported_size;
+    next = prev + size;
+    if (next >= 65536) {
+        report = next;
+        next = 0;
+    }
+    if (prev != next) {
+#ifdef _WIN32
+        if (InterlockedCompareExchange(&unreported_size, next, prev) != prev)
+            goto retry;
+#else
+        if (!__sync_bool_compare_and_swap(&unreported_size, prev, next))
+            goto retry;
+#endif
+    }
+
+    if (report) {
+        PyGILState_STATE state = PyGILState_Ensure();
+        _PyPyGC_AddMemoryPressure(report);
+        PyGILState_Release(state);
+    }
+}
+
+int _PyTraceMalloc_Untrack(_PyTraceMalloc_domain_t domain,
+                           uintptr_t ptr)
+{
+    /* nothing */
+}
diff --git a/pypy/module/cpyext/test/test_object.py b/pypy/module/cpyext/test/test_object.py
--- a/pypy/module/cpyext/test/test_object.py
+++ b/pypy/module/cpyext/test/test_object.py
@@ -214,6 +214,9 @@
 
 class AppTestObject(AppTestCpythonExtensionBase):
     def setup_class(cls):
+        from rpython.rlib import rgc
+        from pypy.interpreter import gateway
+
         AppTestCpythonExtensionBase.setup_class.im_func(cls)
         tmpname = str(py.test.ensuretemp('out', dir=0))
         if cls.runappdirect:
@@ -221,6 +224,28 @@
         else:
             cls.w_tmpname = cls.space.wrap(tmpname)
 
+        cls.total_mem = 0
+        def add_memory_pressure(estimate):
+            assert estimate >= 0
+            cls.total_mem += estimate
+        cls.orig_add_memory_pressure = [rgc.add_memory_pressure]
+        rgc.add_memory_pressure = add_memory_pressure
+
+        def _reset_memory_pressure(space):
+            cls.total_mem = 0
+        cls.w_reset_memory_pressure = cls.space.wrap(
+            gateway.interp2app(_reset_memory_pressure))
+
+        def _cur_memory_pressure(space):
+            return space.newint(cls.total_mem)
+        cls.w_cur_memory_pressure = cls.space.wrap(
+            gateway.interp2app(_cur_memory_pressure))
+
+    def teardown_class(cls):
+        from rpython.rlib import rgc
+        if hasattr(cls, 'orig_add_memory_pressure'):
+            [rgc.add_memory_pressure] = cls.orig_add_memory_pressure
+
     def test_object_malloc(self):
         module = self.import_extension('foo', [
             ("malloctest", "METH_NOARGS",
@@ -323,6 +348,30 @@
         a = module.empty_format('hello')
         assert isinstance(a, unicode)
 
+    def test_add_memory_pressure(self):
+        module = self.import_extension('foo', [
+            ("foo", "METH_O",
+            """
+                _PyTraceMalloc_Track(0, 0, PyInt_AsLong(args) - sizeof(long));
+                Py_INCREF(Py_None);
+                return Py_None;
+            """)])
+        self.reset_memory_pressure()
+        module.foo(42)
+        assert self.cur_memory_pressure() == 0
+        module.foo(65000 - 42)
+        assert self.cur_memory_pressure() == 0
+        module.foo(536)
+        assert self.cur_memory_pressure() == 65536
+        module.foo(40000)
+        assert self.cur_memory_pressure() == 65536
+        module.foo(40000)
+        assert self.cur_memory_pressure() == 65536 + 80000
+        module.foo(35000)
+        assert self.cur_memory_pressure() == 65536 + 80000
+        module.foo(35000)
+        assert self.cur_memory_pressure() == 65536 + 80000 + 70000
+
 class AppTestPyBuffer_FillInfo(AppTestCpythonExtensionBase):
     """
     PyBuffer_FillInfo populates the fields of a Py_buffer from its arguments.