Issue #13390: New function :func:`sys.getallocatedblocks()` returns the number of memory blocks currently allocated.

Also, the ``-R`` option to regrtest uses this function to guard against memory allocation leaks.

9 files changed, 123 insertions, 22 deletions

diff --git a/Doc/library/sys.rst b/Doc/library/sys.rst
index c327f21..ee1f2b2 100644
--- a/Doc/library/sys.rst
+++ b/Doc/library/sys.rst
@@ -393,6 +393,20 @@ always available.
    .. versionadded:: 3.1
 
 
+.. function:: getallocatedblocks()
+
+   Return the number of memory blocks currently allocated by the interpreter,
+   regardless of their size.  This function is mainly useful for debugging
+   small memory leaks.  Because of the interpreter's internal caches, the
+   result can vary from call to call; you may have to call
+   :func:`_clear_type_cache()` to get more predictable results.
+
+   .. versionadded:: 3.4
+
+   .. impl-detail::
+      Not all Python implementations may be able to return this information.
+
+
 .. function:: getcheckinterval()
 
    Return the interpreter's "check interval"; see :func:`setcheckinterval`.
diff --git a/Include/objimpl.h b/Include/objimpl.h
index 3d5f509..6346500 100644
--- a/Include/objimpl.h
+++ b/Include/objimpl.h
@@ -98,6 +98,8 @@ PyAPI_FUNC(void *) PyObject_Malloc(size_t);
 PyAPI_FUNC(void *) PyObject_Realloc(void *, size_t);
 PyAPI_FUNC(void) PyObject_Free(void *);
 
+/* This function returns the number of allocated memory blocks, regardless of size */
+PyAPI_FUNC(Py_ssize_t) _Py_GetAllocatedBlocks(void);
 
 /* Macros */
 #ifdef WITH_PYMALLOC
diff --git a/Lib/test/regrtest.py b/Lib/test/regrtest.py
index 892ff6b..5d9d82d 100755
--- a/Lib/test/regrtest.py
+++ b/Lib/test/regrtest.py
@@ -615,7 +615,7 @@ def main(tests=None, testdir=None, verbose=0, quiet=False,
             sys.exit(2)
         from queue import Queue
         from subprocess import Popen, PIPE
-        debug_output_pat = re.compile(r"\[\d+ refs\]$")
+        debug_output_pat = re.compile(r"\[\d+ refs, \d+ blocks\]$")
         output = Queue()
         pending = MultiprocessTests(tests)
         opt_args = support.args_from_interpreter_flags()
@@ -1320,33 +1320,50 @@ def dash_R(the_module, test, indirect_test, huntrleaks):
             del sys.modules[the_module.__name__]
             exec('import ' + the_module.__name__)
 
-    deltas = []
     nwarmup, ntracked, fname = huntrleaks
     fname = os.path.join(support.SAVEDCWD, fname)
     repcount = nwarmup + ntracked
+    rc_deltas = [0] * repcount
+    alloc_deltas = [0] * repcount
+
     print("beginning", repcount, "repetitions", file=sys.stderr)
     print(("1234567890"*(repcount//10 + 1))[:repcount], file=sys.stderr)
     sys.stderr.flush()
-    dash_R_cleanup(fs, ps, pic, zdc, abcs)
     for i in range(repcount):
-        rc_before = sys.gettotalrefcount()
         run_the_test()
+        alloc_after, rc_after = dash_R_cleanup(fs, ps, pic, zdc, abcs)
         sys.stderr.write('.')
         sys.stderr.flush()
-        dash_R_cleanup(fs, ps, pic, zdc, abcs)
-        rc_after = sys.gettotalrefcount()
         if i >= nwarmup:
-            deltas.append(rc_after - rc_before)
+            rc_deltas[i] = rc_after - rc_before
+            alloc_deltas[i] = alloc_after - alloc_before
+        alloc_before, rc_before = alloc_after, rc_after
     print(file=sys.stderr)
-    if any(deltas):
-        msg = '%s leaked %s references, sum=%s' % (test, deltas, sum(deltas))
-        print(msg, file=sys.stderr)
-        sys.stderr.flush()
-        with open(fname, "a") as refrep:
-            print(msg, file=refrep)
-            refrep.flush()
-        return True
-    return False
+    # These checkers return False on success, True on failure
+    def check_rc_deltas(deltas):
+        return any(deltas)
+    def check_alloc_deltas(deltas):
+        # At least 1/3rd of 0s
+        if 3 * deltas.count(0) < len(deltas):
+            return True
+        # Nothing else than 1s, 0s and -1s
+        if not set(deltas) <= {1,0,-1}:
+            return True
+        return False
+    failed = False
+    for deltas, item_name, checker in [
+        (rc_deltas, 'references', check_rc_deltas),
+        (alloc_deltas, 'memory blocks', check_alloc_deltas)]:
+        if checker(deltas):
+            msg = '%s leaked %s %s, sum=%s' % (
+                test, deltas[nwarmup:], item_name, sum(deltas))
+            print(msg, file=sys.stderr)
+            sys.stderr.flush()
+            with open(fname, "a") as refrep:
+                print(msg, file=refrep)
+                refrep.flush()
+            failed = True
+    return failed
 
 def dash_R_cleanup(fs, ps, pic, zdc, abcs):
     import gc, copyreg
@@ -1412,8 +1429,11 @@ def dash_R_cleanup(fs, ps, pic, zdc, abcs):
     else:
         ctypes._reset_cache()
 
-    # Collect cyclic trash.
+    # Collect cyclic trash and read memory statistics immediately after.
+    func1 = sys.getallocatedblocks
+    func2 = sys.gettotalrefcount
     gc.collect()
+    return func1(), func2()
 
 def warm_caches():
     # char cache
diff --git a/Lib/test/support.py b/Lib/test/support.py
index 4d63904..b98f2bf 100644
--- a/Lib/test/support.py
+++ b/Lib/test/support.py
@@ -1772,7 +1772,7 @@ def strip_python_stderr(stderr):
     This will typically be run on the result of the communicate() method
     of a subprocess.Popen object.
     """
-    stderr = re.sub(br"\[\d+ refs\]\r?\n?", b"", stderr).strip()
+    stderr = re.sub(br"\[\d+ refs, \d+ blocks\]\r?\n?", b"", stderr).strip()
     return stderr
 
 def args_from_interpreter_flags():
diff --git a/Lib/test/test_sys.py b/Lib/test/test_sys.py
index a1074c3..055592b 100644
--- a/Lib/test/test_sys.py
+++ b/Lib/test/test_sys.py
@@ -6,6 +6,7 @@ import textwrap
 import warnings
 import operator
 import codecs
+import gc
 
 # count the number of test runs, used to create unique
 # strings to intern in test_intern()
@@ -611,6 +612,29 @@ class SysModuleTest(unittest.TestCase):
         ret, out, err = assert_python_ok(*args)
         self.assertIn(b"free PyDictObjects", err)
 
+    @unittest.skipUnless(hasattr(sys, "getallocatedblocks"),
+                         "sys.getallocatedblocks unavailable on this build")
+    def test_getallocatedblocks(self):
+        # Some sanity checks
+        a = sys.getallocatedblocks()
+        self.assertIs(type(a), int)
+        self.assertGreater(a, 0)
+        try:
+            # While we could imagine a Python session where the number of
+            # multiple buffer objects would exceed the sharing of references,
+            # it is unlikely to happen in a normal test run.
+            self.assertLess(a, sys.gettotalrefcount())
+        except AttributeError:
+            # gettotalrefcount() not available
+            pass
+        gc.collect()
+        b = sys.getallocatedblocks()
+        self.assertLessEqual(b, a)
+        gc.collect()
+        c = sys.getallocatedblocks()
+        self.assertIn(c, range(b - 50, b + 50))
+
+
 class SizeofTest(unittest.TestCase):
 
     def setUp(self):
diff --git a/Misc/NEWS b/Misc/NEWS
index f69c05c..c305f65 100644
--- a/Misc/NEWS
+++ b/Misc/NEWS
@@ -163,6 +163,9 @@ Core and Builtins
 Library
 -------
 
+- Issue #13390: New function :func:`sys.getallocatedblocks()` returns the
+  number of memory blocks currently allocated.
+
 - Issue #16628: Fix a memory leak in ctypes.resize().
 
 - Issue #13614: Fix setup.py register failure with invalid rst in description.
@@ -433,6 +436,9 @@ Extension Modules
 Tests
 -----
 
+- Issue #13390: The ``-R`` option to regrtest now also checks for memory
+  allocation leaks, using :func:`sys.getallocatedblocks()`.
+
 - Issue #16559: Add more tests for the json module, including some from the
   official test suite at json.org.  Patch by Serhiy Storchaka.
 
diff --git a/Objects/obmalloc.c b/Objects/obmalloc.c
index 6225ebb..c82c978 100644
--- a/Objects/obmalloc.c
+++ b/Objects/obmalloc.c
@@ -525,6 +525,15 @@ static size_t ntimes_arena_allocated = 0;
 /* High water mark (max value ever seen) for narenas_currently_allocated. */
 static size_t narenas_highwater = 0;
 
+static Py_ssize_t _Py_AllocatedBlocks = 0;
+
+Py_ssize_t
+_Py_GetAllocatedBlocks(void)
+{
+    return _Py_AllocatedBlocks;
+}
+
+
 /* Allocate a new arena.  If we run out of memory, return NULL.  Else
  * allocate a new arena, and return the address of an arena_object
  * describing the new arena.  It's expected that the caller will set
@@ -785,6 +794,8 @@ PyObject_Malloc(size_t nbytes)
     if (nbytes > PY_SSIZE_T_MAX)
         return NULL;
 
+    _Py_AllocatedBlocks++;
+
     /*
      * This implicitly redirects malloc(0).
      */
@@ -901,6 +912,7 @@ PyObject_Malloc(size_t nbytes)
                  * and free list are already initialized.
                  */
                 bp = pool->freeblock;
+                assert(bp != NULL);
                 pool->freeblock = *(block **)bp;
                 UNLOCK();
                 return (void *)bp;
@@ -958,7 +970,12 @@ redirect:
      */
     if (nbytes == 0)
         nbytes = 1;
-    return (void *)malloc(nbytes);
+    {
+        void *result = malloc(nbytes);
+        if (!result)
+            _Py_AllocatedBlocks--;
+        return result;
+    }
 }
 
 /* free */
@@ -978,6 +995,8 @@ PyObject_Free(void *p)
     if (p == NULL)      /* free(NULL) has no effect */
         return;
 
+    _Py_AllocatedBlocks--;
+
 #ifdef WITH_VALGRIND
     if (UNLIKELY(running_on_valgrind > 0))
         goto redirect;
diff --git a/Python/pythonrun.c b/Python/pythonrun.c
index dd32017..f0d8550 100644
--- a/Python/pythonrun.c
+++ b/Python/pythonrun.c
@@ -38,9 +38,10 @@
 #ifndef Py_REF_DEBUG
 #define PRINT_TOTAL_REFS()
 #else /* Py_REF_DEBUG */
-#define PRINT_TOTAL_REFS() fprintf(stderr,                              \
-                   "[%" PY_FORMAT_SIZE_T "d refs]\n",                   \
-                   _Py_GetRefTotal())
+#define PRINT_TOTAL_REFS() fprintf(stderr,                   \
+                   "[%" PY_FORMAT_SIZE_T "d refs, "          \
+                   "%" PY_FORMAT_SIZE_T "d blocks]\n",       \
+                   _Py_GetRefTotal(), _Py_GetAllocatedBlocks())
 #endif
 
 #ifdef __cplusplus
diff --git a/Python/sysmodule.c b/Python/sysmodule.c
index 92c5b67..20792c2 100644
--- a/Python/sysmodule.c
+++ b/Python/sysmodule.c
@@ -894,6 +894,19 @@ one higher than you might expect, because it includes the (temporary)\n\
 reference as an argument to getrefcount()."
 );
 
+static PyObject *
+sys_getallocatedblocks(PyObject *self)
+{
+    return PyLong_FromSsize_t(_Py_GetAllocatedBlocks());
+}
+
+PyDoc_STRVAR(getallocatedblocks_doc,
+"getallocatedblocks() -> integer\n\
+\n\
+Return the number of memory blocks currently allocated, regardless of their\n\
+size."
+);
+
 #ifdef COUNT_ALLOCS
 static PyObject *
 sys_getcounts(PyObject *self)
@@ -1062,6 +1075,8 @@ static PyMethodDef sys_methods[] = {
     {"getdlopenflags", (PyCFunction)sys_getdlopenflags, METH_NOARGS,
      getdlopenflags_doc},
 #endif
+    {"getallocatedblocks", (PyCFunction)sys_getallocatedblocks, METH_NOARGS,
+      getallocatedblocks_doc},
 #ifdef COUNT_ALLOCS
     {"getcounts",       (PyCFunction)sys_getcounts, METH_NOARGS},
 #endif