From 85f643023fed3d4e2fb8e399f9ad57f3a65ef237 Mon Sep 17 00:00:00 2001 From: xdegaye Date: Sat, 1 Jul 2017 14:14:45 +0200 Subject: bpo-30695: Add set_nomemory(start, stop) to _testcapi (GH-2406) --- Lib/test/test_capi.py | 34 +++++- .../Tests/2017-06-30-11-20-20.bpo-30695.lo7FQX.rst | 2 + Modules/_testcapimodule.c | 128 +++++++++++++++++++++ 3 files changed, 163 insertions(+), 1 deletion(-) create mode 100644 Misc/NEWS.d/next/Tests/2017-06-30-11-20-20.bpo-30695.lo7FQX.rst diff --git a/Lib/test/test_capi.py b/Lib/test/test_capi.py index 1cf5cd7..c3a04b4 100644 --- a/Lib/test/test_capi.py +++ b/Lib/test/test_capi.py @@ -15,7 +15,7 @@ import time import unittest from test import support from test.support import MISSING_C_DOCSTRINGS -from test.support.script_helper import assert_python_failure +from test.support.script_helper import assert_python_failure, assert_python_ok try: import _posixsubprocess except ImportError: @@ -243,6 +243,38 @@ class CAPITest(unittest.TestCase): def test_buildvalue_N(self): _testcapi.test_buildvalue_N() + def test_set_nomemory(self): + code = """if 1: + import _testcapi + + class C(): pass + + # The first loop tests both functions and that remove_mem_hooks() + # can be called twice in a row. The second loop checks a call to + # set_nomemory() after a call to remove_mem_hooks(). The third + # loop checks the start and stop arguments of set_nomemory(). + for outer_cnt in range(1, 4): + start = 10 * outer_cnt + for j in range(100): + if j == 0: + if outer_cnt != 3: + _testcapi.set_nomemory(start) + else: + _testcapi.set_nomemory(start, start + 1) + try: + C() + except MemoryError as e: + if outer_cnt != 3: + _testcapi.remove_mem_hooks() + print('MemoryError', outer_cnt, j) + _testcapi.remove_mem_hooks() + break + """ + rc, out, err = assert_python_ok('-c', code) + self.assertIn(b'MemoryError 1 10', out) + self.assertIn(b'MemoryError 2 20', out) + self.assertIn(b'MemoryError 3 30', out) + @unittest.skipUnless(threading, 'Threading required for this test.') class TestPendingCalls(unittest.TestCase): diff --git a/Misc/NEWS.d/next/Tests/2017-06-30-11-20-20.bpo-30695.lo7FQX.rst b/Misc/NEWS.d/next/Tests/2017-06-30-11-20-20.bpo-30695.lo7FQX.rst new file mode 100644 index 0000000..a57bbe7 --- /dev/null +++ b/Misc/NEWS.d/next/Tests/2017-06-30-11-20-20.bpo-30695.lo7FQX.rst @@ -0,0 +1,2 @@ +Add the `set_nomemory(start, stop)` and `remove_mem_hooks()` functions to +the _testcapi module. diff --git a/Modules/_testcapimodule.c b/Modules/_testcapimodule.c index 83ba33d..c9c2798 100644 --- a/Modules/_testcapimodule.c +++ b/Modules/_testcapimodule.c @@ -3420,6 +3420,130 @@ test_pyobject_setallocators(PyObject *self) return test_setallocators(PYMEM_DOMAIN_OBJ); } +/* Most part of the following code is inherited from the pyfailmalloc project + * written by Victor Stinner. */ +static struct { + int installed; + PyMemAllocatorEx raw; + PyMemAllocatorEx mem; + PyMemAllocatorEx obj; +} FmHook; + +static struct { + int start; + int stop; + Py_ssize_t count; +} FmData; + +static int +fm_nomemory(void) +{ + FmData.count++; + if (FmData.count > FmData.start && + (FmData.stop <= 0 || FmData.count <= FmData.stop)) { + return 1; + } + return 0; +} + +static void * +hook_fmalloc(void *ctx, size_t size) +{ + PyMemAllocatorEx *alloc = (PyMemAllocatorEx *)ctx; + if (fm_nomemory()) { + return NULL; + } + return alloc->malloc(alloc->ctx, size); +} + +static void * +hook_fcalloc(void *ctx, size_t nelem, size_t elsize) +{ + PyMemAllocatorEx *alloc = (PyMemAllocatorEx *)ctx; + if (fm_nomemory()) { + return NULL; + } + return alloc->calloc(alloc->ctx, nelem, elsize); +} + +static void * +hook_frealloc(void *ctx, void *ptr, size_t new_size) +{ + PyMemAllocatorEx *alloc = (PyMemAllocatorEx *)ctx; + if (fm_nomemory()) { + return NULL; + } + return alloc->realloc(alloc->ctx, ptr, new_size); +} + +static void +hook_ffree(void *ctx, void *ptr) +{ + PyMemAllocatorEx *alloc = (PyMemAllocatorEx *)ctx; + alloc->free(alloc->ctx, ptr); +} + +static void +fm_setup_hooks(void) +{ + PyMemAllocatorEx alloc; + + if (FmHook.installed) { + return; + } + FmHook.installed = 1; + + alloc.malloc = hook_fmalloc; + alloc.calloc = hook_fcalloc; + alloc.realloc = hook_frealloc; + alloc.free = hook_ffree; + PyMem_GetAllocator(PYMEM_DOMAIN_RAW, &FmHook.raw); + PyMem_GetAllocator(PYMEM_DOMAIN_MEM, &FmHook.mem); + PyMem_GetAllocator(PYMEM_DOMAIN_OBJ, &FmHook.obj); + + alloc.ctx = &FmHook.raw; + PyMem_SetAllocator(PYMEM_DOMAIN_RAW, &alloc); + + alloc.ctx = &FmHook.mem; + PyMem_SetAllocator(PYMEM_DOMAIN_MEM, &alloc); + + alloc.ctx = &FmHook.obj; + PyMem_SetAllocator(PYMEM_DOMAIN_OBJ, &alloc); +} + +static void +fm_remove_hooks(void) +{ + if (FmHook.installed) { + FmHook.installed = 0; + PyMem_SetAllocator(PYMEM_DOMAIN_RAW, &FmHook.raw); + PyMem_SetAllocator(PYMEM_DOMAIN_MEM, &FmHook.mem); + PyMem_SetAllocator(PYMEM_DOMAIN_OBJ, &FmHook.obj); + } +} + +static PyObject* +set_nomemory(PyObject *self, PyObject *args) +{ + /* Memory allocation fails after 'start' allocation requests, and until + * 'stop' allocation requests except when 'stop' is negative or equal + * to 0 (default) in which case allocation failures never stop. */ + FmData.count = 0; + FmData.stop = 0; + if (!PyArg_ParseTuple(args, "i|i", &FmData.start, &FmData.stop)) { + return NULL; + } + fm_setup_hooks(); + Py_RETURN_NONE; +} + +static PyObject* +remove_mem_hooks(PyObject *self) +{ + fm_remove_hooks(); + Py_RETURN_NONE; +} + PyDoc_STRVAR(docstring_empty, "" ); @@ -4287,6 +4411,10 @@ static PyMethodDef TestMethods[] = { (PyCFunction)test_pymem_setallocators, METH_NOARGS}, {"test_pyobject_setallocators", (PyCFunction)test_pyobject_setallocators, METH_NOARGS}, + {"set_nomemory", (PyCFunction)set_nomemory, METH_VARARGS, + PyDoc_STR("set_nomemory(start:int, stop:int = 0)")}, + {"remove_mem_hooks", (PyCFunction)remove_mem_hooks, METH_NOARGS, + PyDoc_STR("Remove memory hooks.")}, {"no_docstring", (PyCFunction)test_with_docstring, METH_NOARGS}, {"docstring_empty", -- cgit v0.12